diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c4930a5a..7f41b525 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -36,6 +36,19 @@ jobs: - name: Run cargo fmt run: make fmt + # 2.1 + clippy: + name: Rust clippy + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: clippy + - uses: Swatinem/rust-cache@v2 + + - name: Run cargo clippy + run: make clippy # 3 e2e: name: Rust e2e sqllogictest @@ -47,6 +60,19 @@ jobs: - name: Run sqllogictest suite run: make test-slt + # 3.1 + sqlite-slt: + name: SQLite sqllogictest corpus + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + submodules: true + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + + - name: Run SQLite sqllogictest corpus + run: make test-slt-sqlite # 4 coverage: name: Rust coverage @@ -165,3 +191,22 @@ jobs: with: name: fuzz-artifacts path: fuzz/artifacts + # 10 + miri: + name: Miri (unsafe code UB check) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@master + with: + toolchain: nightly-2026-10-04 + components: miri, rust-src + - uses: Swatinem/rust-cache@v2 + with: + key: miri + + - name: Set up Miri sysroot + run: cargo +nightly-2026-10-04 miri setup + + - name: Run Miri tests + run: make miri diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 00000000..13412ef0 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "tests/kitesql-testing"] + path = tests/kitesql-testing + url = https://github.com/KipData/kitesql-testing.git diff --git a/AGENT.md b/AGENT.md index 8073a5ff..e59d87c8 100644 --- a/AGENT.md +++ b/AGENT.md @@ -77,6 +77,18 @@ PRs that modify logic but leave obvious test gaps untouched may be rejected. --- +### 2.4 Write Only Necessary Tests + +Coverage means each behavior is tested once, not that every case gets its own test. + +- Test each distinct branch or behavior once, at the layer that owns it +- Do not repeat in a driver or caller what the rule or module already tests +- Group related cases into one table-driven test instead of near-identical tests +- Skip tests for trivial control flow (simple loops, early returns, plain delegation) +- When touching a test module, delete tests made redundant by others + +--- + ## 3. Code Simplicity Rules ### 3.1 Prefer Simple Code @@ -134,6 +146,13 @@ If code looks similar in multiple places, it probably wants to be reused. - Avoid exposing internal details - Keep modules small, focused, and cohesive +### 4.3 Rust Imports + +- Declare `use` imports at module scope, not inside functions or methods. +- Import functions and types before using them; do not write fully qualified paths + such as `crate::optimizer::rule::normalization::reorder_joins(...)` in function bodies. +- Use import aliases when names conflict; keep test imports at the test module scope. + --- ## 5. Dependency Policy diff --git a/Cargo.lock b/Cargo.lock index cf0640c2..15984401 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1135,7 +1135,7 @@ dependencies = [ [[package]] name = "kite_sql" -version = "0.5.0" +version = "0.6.0" dependencies = [ "bumpalo", "chrono", @@ -1164,7 +1164,7 @@ dependencies = [ [[package]] name = "kite_sql_serde_macros" -version = "0.2.4" +version = "0.2.5" dependencies = [ "darling", "proc-macro2", diff --git a/Cargo.toml b/Cargo.toml index 7cf94757..00406ada 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ [package] name = "kite_sql" -version = "0.5.0" +version = "0.6.0" edition = "2021" build = "build.rs" authors = ["Kould ", "Xwg "] @@ -56,7 +56,7 @@ required-features = ["pprof"] bumpalo = { version = "3", default-features = false, features = ["collections"] } ordered-float = { version = "4" } paste = { version = "1" } -kite_sql_serde_macros = { version = "0.2.4", path = "kite_sql_serde_macros" } +kite_sql_serde_macros = { version = "0.2.5", path = "kite_sql_serde_macros" } # Optional dependencies for features comfy-table = { version = "7", default-features = false, optional = true } diff --git a/Makefile b/Makefile index 6b763ded..9184ba38 100644 --- a/Makefile +++ b/Makefile @@ -3,6 +3,8 @@ CARGO ?= cargo GRCOV ?= grcov WASM_PACK ?= wasm-pack SQLLOGIC_PATH ?= tests/slt/**/*.slt +SQLITE_LOGIC_PATH ?= tests/kitesql-testing/data/sqlite/**/*.slt +SQLITE_LOGIC_ARGS ?= PYO3_PYTHON ?= /usr/bin/python3.12 TPCC_MEASURE_TIME ?= 15 TPCC_NUM_WARE ?= 1 @@ -33,9 +35,12 @@ FUZZ_PASS_EXIT_sql_exec = 70 # sql_gen: bounded generated queries, so crashes, timeouts and OOMs all fail. FUZZ_ARGS_sql_gen = -max_len=4096 FUZZ_PASS_EXIT_sql_gen = 0 +MIRI_TOOLCHAIN ?= nightly-2026-10-04 +MIRI_FEATURES ?= spill,decimal,lmdb +MIRI_TESTS ?= types::tuple:: execution::test:: COVERAGE_REPORT_ARGS ?= --llvm --ignore-not-existing --keep-only 'src/**' --ignore 'src/**/tests/**' --ignore 'tests/**' --ignore 'tpcc/**' --excl-start 'GRCOV_EXCL_START' --excl-stop 'GRCOV_EXCL_STOP' -.PHONY: test test-python test-wasm test-slt test-all codecov codecov-html wasm-build check tpcc tpcc-kitesql-rocksdb tpcc-kitesql-lmdb tpcc-lmdb-flamegraph tpcc-lmdb-heaptrack tpcc-sqlite tpcc-dual cargo-check build wasm-examples native-examples fmt clippy fuzz +.PHONY: test test-python test-wasm test-slt test-all codecov codecov-html wasm-build check tpcc tpcc-kitesql-rocksdb tpcc-kitesql-lmdb tpcc-lmdb-flamegraph tpcc-lmdb-heaptrack tpcc-sqlite tpcc-dual cargo-check build wasm-examples native-examples fmt clippy fuzz miri ## Run default Rust tests in the current environment (non-WASM). test: @@ -65,6 +70,12 @@ test-wasm: test-slt: $(CARGO) run -p sqllogictest-test -- --path '$(SQLLOGIC_PATH)' +.PHONY: test-slt-sqlite + +## Run the external SQLite corpus (select1ā€“5 excluded unless --include-select is passed). +test-slt-sqlite: + $(CARGO) run --release -p sqllogictest-test --bin sqlite -- --path '$(SQLITE_LOGIC_PATH)' $(SQLITE_LOGIC_ARGS) + ## Convenience target to run every suite in sequence. test-all: test test-wasm test-slt test-python @@ -103,6 +114,9 @@ fuzz: echo "fuzz: $(FUZZ_TARGET) failed (exit $$status), see fuzz/artifacts/$(FUZZ_TARGET)/"; exit $$status; \ fi +miri: + $(CARGO) +$(MIRI_TOOLCHAIN) miri test --lib --features '$(MIRI_FEATURES)' -- $(MIRI_TESTS) + ## Run formatting (check mode) across the workspace. fmt: $(CARGO) fmt --all -- --check diff --git a/README.md b/README.md index 12f731dc..ec954831 100755 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@

CI Codecov - Ask DeepWiki + Ask DeepWiki Discord @@ -131,14 +131,13 @@ print(list(db.run("select * from demo"))) ## TPC-C `make tpcc` runs the benchmark (`--backend rocksdb|lmdb`); `make tpcc-dual` cross-checks every statement against SQLite. -720-second single-thread run on an i9-13900HX (32 GB, KIOXIA EXCERIA PLUS G3), every backend pinned to the same P-core. Latencies are p90 in Āµs, including commit. +720-second single-thread run on an i9-13900HX (32 GB, KIOXIA EXCERIA PLUS G3), every backend pinned to the same P-core and run without fsync on commit. Latencies are p90 in Āµs, including commit. | Backend | TpmC | New-Order | Payment | Order-Status | Delivery | Stock-Level | | --- | ---: | ---: | ---: | ---: | ---: | ---: | -| KiteSQL LMDB | 159616 | 315 | 81 | 44 | 435 | 332 | -| KiteSQL RocksDB | 48035 | 459 | 258 | 200 | 11055 | 704 | -| SQLite balanced | 67102 | 283 | 67 | 48 | 338 | 473 | -| SQLite practical | 67153 | 336 | 66 | 41 | 427 | 339 | +| KiteSQL LMDB | 165956 | 303 | 75 | 43 | 436 | 318 | +| KiteSQL RocksDB | 50739 | 392 | 188 | 198 | 11743 | 812 | +| SQLite (WAL, `synchronous=OFF`) | 191484 | 274 | 41 | 33 | 343 | 268 | šŸ‘‰ [Details and how to reproduce](tpcc/README.md) diff --git a/examples/transaction.rs b/examples/transaction.rs index 00b4c868..8e8bdee8 100644 --- a/examples/transaction.rs +++ b/examples/transaction.rs @@ -16,7 +16,6 @@ mod app { use kite_sql::db::DataBaseBuilder; use kite_sql::errors::DatabaseError; - use kite_sql::types::tuple::Tuple; use kite_sql::types::value::DataValue; use std::fs; use std::io::ErrorKind; @@ -43,7 +42,7 @@ mod app { // Optimistic transactions are currently backed by RocksDB. let mut database = DataBaseBuilder::path(EXAMPLE_DB_PATH).build_optimistic()?; database.ddl("create table if not exists t1 (c1 int primary key, c2 int)")?; - let mut transaction = database.new_transaction()?; + let mut transaction = database.new_transaction(false)?; transaction .run("insert into t1 values(0, 0), (1, 1)")? @@ -57,29 +56,21 @@ mod app { let mut iter = database.run("select * from t1")?; iter.next_tuple(|_, tuple| { - assert_eq!( - tuple, - &Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]) - ); + assert_eq!(*tuple, vec![DataValue::Int32(0), DataValue::Int32(0)]); })? .unwrap(); iter.next_tuple(|_, tuple| { - assert_eq!( - tuple, - &Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(1)]) - ); + assert_eq!(*tuple, vec![DataValue::Int32(1), DataValue::Int32(1)]); })? .unwrap(); assert!(iter.next_tuple(|_, _| ())?.is_none()); iter.done()?; - let mut tx2 = database.new_transaction()?; + let mut tx2 = database.new_transaction(false)?; tx2.run("update t1 set c2 = 99 where c1 = 0")?.done()?; let mut c2_iter = database.run("select c2 from t1 where c1 = 0")?; assert_eq!( - c2_iter - .next_tuple(|_, tuple| tuple.values[0].i32())? - .unwrap(), + c2_iter.next_tuple(|_, tuple| tuple[0].i32())?.unwrap(), Some(0) ); c2_iter.done()?; diff --git a/fuzz/fuzz_targets/sql_exec.rs b/fuzz/fuzz_targets/sql_exec.rs index af0342fc..50ec7e31 100644 --- a/fuzz/fuzz_targets/sql_exec.rs +++ b/fuzz/fuzz_targets/sql_exec.rs @@ -72,7 +72,7 @@ fn reset(db: &mut Database) { let mut names = Vec::new(); if let Ok(mut iter) = db.run(show) { while let Ok(Some(())) = iter.next_tuple(|_, tuple| { - if let Some(DataValue::Utf8 { value, .. }) = tuple.values.first() { + if let Some(DataValue::Utf8 { value, .. }) = tuple.get(0) { names.push(value.clone()); } }) {} @@ -165,7 +165,7 @@ fn run_one(db: &mut Database, sql: &str) { // Drain and render every value, like the sqllogictest harness does. loop { match iter.next_tuple(|_, tuple| { - for value in tuple.values.iter() { + for value in tuple.iter() { let _ = value.to_string(); } }) { diff --git a/fuzz/fuzz_targets/sql_gen.rs b/fuzz/fuzz_targets/sql_gen.rs index c6d06d4a..608978f3 100644 --- a/fuzz/fuzz_targets/sql_gen.rs +++ b/fuzz/fuzz_targets/sql_gen.rs @@ -34,6 +34,7 @@ use kite_sql::db::{DataBaseBuilder, Database}; use kite_sql::errors::DatabaseError; use kite_sql::storage::lmdb::LmdbStorage; +use kite_sql::types::tuple::TupleImpl; use kite_sql::types::value::DataValue; use libfuzzer_sys::arbitrary::{self, Unstructured}; use libfuzzer_sys::fuzz_target; @@ -735,7 +736,11 @@ fn open_db() -> Database { fn check_tlp(db: &Database, select: &Select, predicate: &str) { let base_sql = select.render(false, false); let mut base = Vec::new(); - if for_each_row(db, &base_sql, |row| base.push(row.to_vec())).is_err() { + if for_each_row(db, &base_sql, |row| { + base.push(row.iter().cloned().collect()) + }) + .is_err() + { return; } let order_by = select.order_by.map(|(column, _)| column); @@ -778,8 +783,11 @@ fn rows_match( } let mut matched = true; let mut i = offset; + let mut key = Vec::new(); let rows = for_each_row(db, sql, |row| { - let in_base = match remaining.get_mut(row) { + key.clear(); + key.extend(row.iter().cloned()); + let in_base = match remaining.get_mut(key.as_slice()) { Some(count) if *count > 0 => { *count -= 1; true @@ -807,7 +815,7 @@ fn run(db: &Database, sql: &str) -> Result { fn for_each_row( db: &Database, sql: &str, - mut f: impl FnMut(&[DataValue]), + mut f: impl FnMut(&TupleImpl), ) -> Result { // Printed before running so a crashing statement is still visible. if debug() { @@ -815,7 +823,7 @@ fn for_each_row( } let mut iter = db.run(sql)?; let mut rows = 0; - while iter.next_tuple(|_, tuple| f(&tuple.values))?.is_some() { + while iter.next_tuple(|_, tuple| f(tuple))?.is_some() { rows += 1; } iter.done()?; diff --git a/kite_sql_serde_macros/Cargo.toml b/kite_sql_serde_macros/Cargo.toml index 3a3c6ba7..f04aa0c2 100644 --- a/kite_sql_serde_macros/Cargo.toml +++ b/kite_sql_serde_macros/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "kite_sql_serde_macros" -version = "0.2.4" +version = "0.2.5" edition = "2021" description = "Derive macros for KiteSQL" documentation = "https://docs.rs/kite_sql_serde_macros/latest/kite_sql_serde_macros/" diff --git a/kite_sql_serde_macros/src/orm.rs b/kite_sql_serde_macros/src/orm.rs index 24d53f26..1ebe4f3c 100644 --- a/kite_sql_serde_macros/src/orm.rs +++ b/kite_sql_serde_macros/src/orm.rs @@ -387,7 +387,7 @@ pub(crate) fn handle(ast: DeriveInput) -> Result { { fn from_query_row( schema: &::kite_sql::types::tuple::SchemaView<'_, '_>, - tuple: &mut ::kite_sql::types::tuple::Tuple, + tuple: &mut ::kite_sql::types::tuple::TupleImpl, ) -> ::std::result::Result { let mut __kite_orm_found_fields = 0usize; #(#field_index_declarations)* diff --git a/kite_sql_serde_macros/src/projection.rs b/kite_sql_serde_macros/src/projection.rs index fff67210..45fb0312 100644 --- a/kite_sql_serde_macros/src/projection.rs +++ b/kite_sql_serde_macros/src/projection.rs @@ -107,7 +107,7 @@ pub(crate) fn handle(ast: DeriveInput) -> Result { relation: &str, ) -> ::std::result::Result<::std::vec::Vec<::kite_sql::planner::ExprRef>, ::kite_sql::errors::DatabaseError> where - T: ::kite_sql::storage::Transaction, + T: ::kite_sql::storage::StorageTransaction, A: AsRef<[(usize, ::kite_sql::types::LogicalType)]>, { Ok(::std::vec![ @@ -121,7 +121,7 @@ pub(crate) fn handle(ast: DeriveInput) -> Result { { fn from_query_row( schema: &::kite_sql::types::tuple::SchemaView<'_, '_>, - tuple: &mut ::kite_sql::types::tuple::Tuple, + tuple: &mut ::kite_sql::types::tuple::TupleImpl, ) -> ::std::result::Result { let mut __kite_orm_found_fields = 0usize; #(#field_index_declarations)* diff --git a/kite_sql_serde_macros/src/reference_serialization.rs b/kite_sql_serde_macros/src/reference_serialization.rs index 087331d7..8cc7b38d 100644 --- a/kite_sql_serde_macros/src/reference_serialization.rs +++ b/kite_sql_serde_macros/src/reference_serialization.rs @@ -148,7 +148,7 @@ pub(crate) fn handle(ast: DeriveInput) -> Result { Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, Tx>>, reference_tables: &crate::serdes::ReferenceTables, @@ -229,7 +229,7 @@ pub(crate) fn handle(ast: DeriveInput) -> Result { Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, Tx>>, reference_tables: &crate::serdes::ReferenceTables, diff --git a/rfc/0001-greedy-join-reorder-owns-join-decisions.md b/rfc/0001-greedy-join-reorder-owns-join-decisions.md new file mode 100644 index 00000000..6a737735 --- /dev/null +++ b/rfc/0001-greedy-join-reorder-owns-join-decisions.md @@ -0,0 +1,427 @@ +# RFC 0001: GreedyJoinReorder Owns Join Decisions (Order, Implementation, and Index Choice) + +- Status: Draft +- Date: 2026-10-20 +- Modules: `optimizer/rule/normalization/join_reorder.rs` (`GreedyJoinReorder`, `OrientNestedLoopJoin`), + `optimizer/rule/normalization/parameterized_index.rs` (`ParameterizeInnerJoin`), + `planner/operator` (`PhysicalOption`), `optimizer/rule/implementation` (scan cost) + +## 1. Background + +### 1.1 Join decisions are spread across four places + +| Stage | Rule | Decides | Based on | +|---|---|---|---| +| implementation | `IndexScanImplementation`, etc. | Access path of each scan | Static range cost; unaware of join keys | +| after_batches 1 | `GreedyJoinReorder` | Join order and hash build side | Hash estimates only | +| after_batches 2 | `ParameterizeInnerJoin` | Whether to rewrite to an index probe, and which index | Rewrites whenever a probe is possible, using the first usable PK/Unique in `index_infos` | +| after_batches 5 | `OrientNestedLoopJoin` | Outer/inner side of a Cross NLJ | `estimated_work` | + +Greedy is passive: it receives a plan whose scans are already chosen, only moves joins around, and orders them +as if every join were a hash join. Afterwards, `ParameterizeInnerJoin` turns some joins into probes without any +cost comparison, and `OrientNestedLoopJoin` adjusts NLJ direction independently. The join order is chosen for +one implementation, while the one that finally executes may be another. + +### 1.2 Consequences + +- Join order ignores probes: a "small outer ā‹ˆ large unique inner" pair looks expensive as a hash join and is + postponed, although as a probe it should be merged first. +- Whether a probe is better is never compared: for a large outer ā‹ˆ small unique inner, the probe performs far + more seeks than a full hash read, yet it is still rewritten into a probe. +- Index choice is independent of the join: the probe uses the first usable index, and Greedy's + `index_statistics` only reads the index chosen during implementation, ignoring the uniqueness and NDV that + other indexes bound by the join keys would provide. +- Cost units are inconsistent, the per-probe cost is underestimated, and HashJoin / NLJ miss part of their cost + (see Section 4), so even if Greedy compared the options it could not get the right answer. + +## 2. Goals and Non-goals + +Goals: + +- `GreedyJoinReorder` becomes the only decision maker for joins inside a Greedy region: when merging each pair, + it decides at once the join order, the join implementation (Hash / IndexProbe / NestLoop) and direction, and + the index used by each leaf scan in the region. +- The downstream rule `ParameterizeInnerJoin` only executes and validates joins already decided by Greedy and no + longer makes its own choice; its original logic remains only as a fallback for joins Greedy did not handle. + `OrientNestedLoopJoin` is merged into Greedy and removed. +- Provide a unified cost model that supports these decisions. + +Non-goals: + +- Do not change the batch order of the optimizer pipeline (Greedy relies on the `row_count` produced by + implementation). +- Do not support inner join probes on non-unique indexes (Normal / Composite); see Section 8. +- Do not model I/O. + +## 3. Design: Greedy Owns the Decisions + +### 3.1 Responsibilities + +| Rule | Today | After | +|---|---|---| +| scan impl rule | Chooses the access path | Only provides a statically optimal default; leaves inside a Greedy region follow Greedy's decision | +| `GreedyJoinReorder` | Orders only, using hash estimates | Decides order + implementation + direction + probe index for each pair; rewrites probes directly and marks the rest `join_decided` | +| `ParameterizeInnerJoin` | Rewrites whenever it can probe | `join_decided`: no rewrite; otherwise the original logic is the fallback | +| `OrientNestedLoopJoin` | Orients every Cross NLJ | Removed; its responsibility moves into Greedy (3.10) | + +The access path of every leaf inside a Greedy region is decided by Greedy; the scan impl rule's result is only +one of Greedy's inputs: + +- Leaf as a probe inner side: the access path is the probe index chosen by Greedy, which replaces the scan's + original index and range during the rewrite. There is no case where "the probe is better but the scan uses + another index". +- Leaf as a hash / NLJ input: the access path does not use the join keys, so the candidates are exactly those + the scan impl rule already compared with static ranges, and its cheapest plan is also the best one here. + Greedy uses its `work` as the input cost of this shape; choosing Hash confirms it. + - Example: `a JOIN b ON a.id = b.a_id WHERE b.status = 1`, where b has `idx(a_id)` and `idx(status)`. Under + hash, b is read independently and its scan cannot see a's values, so `idx(a_id)` cannot restrict the range; + only the range for `status = 1` or a full scan is usable. Only a probe can pass `a.id` to b's scan and seek + on `idx(a_id)`. +- Scans outside a Greedy region, and scans under joins Greedy did not handle, are still decided by the scan impl + rule. + +Greedy runs after implementation (it relies on the `row_count` produced there), so each scan node already has a +statically optimal default and Greedy's decision is applied on top of it; for the same leaf, Greedy's decision has +the highest priority. + +### 3.2 Decision flow + +1. Collect the leaves and join edges in the region (existing logic). +2. For each pair, use the shared predicate function to find probe-able indexes on the leaf side (3.3). +3. The search is left-deep: start from the input with the fewest rows (ties broken by the lowest read cost). + Each round only evaluates "current result ā‹ˆ one remaining input", lists the feasible shapes, scores them, + and keeps the best shape within the pair (3.4). Remaining inputs are always leaves before they are merged, + so they keep the chance to be probed. +4. Among these pairs, pick the input to merge this round by `JoinCandidate` order; when the current result is + estimated at no more than 1 row, cross joins are not deprioritized. +5. Build the join with the chosen shape: decide left and right (the probe inner side and NLJ inner side go to the + right, the hash build side to the left); rewrite probes directly into MarkApply and mark the rest + `join_decided` (3.5). +6. The merged result becomes a new input for the next round; its `InputCost` comes from the cost function of the + chosen shape. + +### 3.3 Sources and probe candidates + +Greedy decides the usable probe indexes itself from the source's column information and the indexes on it, +without relying on the impl rule to keep candidates: + +```rust +struct ProbeIndex { + position: usize, // position in scan.index_infos + probe: Vec, // outer values and constants, in index column order + lookup_work: usize, // cost of one probe, see 4.3 +} +``` + +- Eligible leaves: `Filter*` + a `TableScan` without a limit. A merged join result cannot be a probe inner side, + so probes only appear in the "merged result ā‹ˆ leaf" shape, which matches how Greedy extends left-deep. +- Every index column either appears on some join edge (on this leaf's side) or is bound to a constant by an + equality in the leaf's filter; when both hold, it is treated as a constant to reduce the join-key requirement. +- Each pair calls the predicate function directly and takes the PK/Unique with the lowest `lookup_work` (ties + broken by index order); nothing is cached or truncated, and the number of candidates is bounded by the indexes + on the table. +- The output rows of a probe are tightened by uniqueness to `min(estimated rows, outer rows)`: the probed index is + bound by the join keys, so its uniqueness is known. + +### 3.4 Implementation shapes within a pair + +For a pair (L, R), the feasible shapes are: + +| Shape | Condition | +|---|---| +| Hash(build = L) / Hash(build = R) | The pair has join keys | +| Probe(inner = R) / Probe(inner = L) | That side is a leaf, and the predicate function finds an index using the pair's join keys | +| NestLoop(outer = L) / NestLoop(outer = R) | The pair has no join keys (cross) | + +Each shape is scored with the cost model in Section 4, and the lowest score within the pair wins. `JoinCandidate` +compares in this order (`derive(Ord)`): + +``` +cross_join ā†’ score ā†’ rows ā†’ left ā†’ right ā†’ swapped ā†’ probe +``` + +- When scores tie within a pair, the probe wins; without statistics, the default row count is used (see 3.9). + +### 3.5 Passing the decision down: `join_decided` + +```rust +// New in PhysicalOption; like startup_work, it is an in-pass annotation that does not take part in +// PartialEq / Hash and is not serialized +join_decided: bool, +``` + +When Greedy chooses a probe, it directly calls the rewrite function shared with `ParameterizeInnerJoin` to build +the MarkApply, so the decision and the execution use the same code and need no further validation. Hash / +NestLoop joins are marked `join_decided`, meaning "execute in this shape and do not decide again". + +### 3.6 Execution and fallback in downstream rules + +- `ParameterizeInnerJoin` + - `join_decided`: no rewrite. + - Otherwise: for joins Greedy did not handle (ON with a filter, `limit_pushed`, etc.), the original logic is the + fallback; it chooses the index with the same predicate function and computes `lookup_work` with the loader. +- NestLoop direction is decided by Greedy; no downstream rule adjusts direction any more (3.10). +- The replaced static range does not lose its filtering semantics: the Filter stays above the scan, and + `EliminateIndexFilter` runs after Parameterize. +- explain only shows the chosen index (existing behavior), not `join_decided` or `lookup_work`. + +### 3.7 Shared predicate function + +A pure predicate function (which does not modify the plan) is extracted from `parameterize`; both Greedy's +candidate generation and `ParameterizeInnerJoin`'s validation call it: + +- Every index column is in the scan's `columns`; +- The join value expression type exactly matches the index column type (`value.return_type == column.datatype()`); +- At least one join key is used (`uses_join_key`); filling the index with constants alone is not a probe. + +Greedy can only own the decision if its judgment matches execution. If each side kept its own copy, Greedy could +order joins for a probe while execution falls back to hash, which is worse than ignoring probes altogether. + +### 3.8 Index coverage: probes and scans use the same rule + +Today: coverage information (`covered_deserializers` / `cover_mapping`) is only computed in +`PushPredicateIntoScan`, and only for indexes that a static range can be extracted from the filter. Coverage for an +index used only by join keys is always `None`; when `parameterize` rewrites to a probe it only sets +`lookup = Probe` and `residual_predicate = None`, without recomputing coverage. The executor (`index_scan.rs`) +decides whether to look up the table by coverage, so a secondary Unique index that actually covers the columns the +scan needs still looks up the table during a probe, and its cost is counted with `seeks = 2`. As a result, probes are +systematically overestimated when compared with a covering index chosen by the scan. + +Fix: + +1. Extract the coverage computation from `PushPredicateIntoScan` into a pure function + `index_covering(scan, index) -> Option<(deserializers, mapping)>`, which only depends on the scan's `columns` + and the index columns, not on the range. +2. Greedy uses it to decide `seeks` when generating candidates (PK or covering Unique = 1, otherwise = 2); the + non-covering Ɨ2 rule in scan cost uses the same function. +3. The probe rewrite (shared by Greedy and the fallback path) also writes `covered_deserializers` / + `cover_mapping`, so the executor really skips the table lookup. + +This way, the Hash shape uses the `work` of the path the scan already chose (possibly a covering index), and the +Probe shape enumerates every candidate in `scan.index_infos` and prices it by its real coverage. The case "the scan +chose a covering index, but a probe is better" shows up directly in the in-pair comparison, and Greedy's choice +overrides the scan's. + +### 3.9 Default row count without statistics + +An input without statistics is priced at `DEFAULT_TABLE_ROWS` (1000) rows, and the N needed by a probe also falls +back to this value when missing; cost comparison proceeds as usual, with no heuristic branch. A static range on a +PK/Unique equality has an exact row count even without statistics, and the scan impl rule chooses it and reports +that row count. A join merged from default row counts is still recorded as having unknown rows, so the made-up +estimate does not leak to later rules. + +### 3.10 Merging `OrientNestedLoopJoin` + +`OrientNestedLoopJoin` only matches `Cross` + `on = None` + `limit_pushed = false` + `NestLoopJoin`, and decides the +direction by comparing `left_work + left_rows Ɨ right_work` with the opposite side, which is essentially the same +formula as `JoinImpl::NestLoop`. It handles three kinds of joins; all of them move into Greedy, and the rule and its +`Orient Nested Loop Join` batch are removed: + +| Source | Today | After | +|---|---|---| +| A pair without join keys inside a Greedy region | Greedy builds it with `swapped: connected && ...`, so cross joins are never swapped and Orient fixes the direction | Greedy compares `NestLoop(outer = L / R)` | +| A region rooted at `Cross` (e.g. `FROM a, b`, or under an outer join / an inner join with a filter) | Greedy requires an `Inner` root (`apply` entry) and skips it | Greedy's trigger is relaxed to `can_expand(plan)`; `can_expand` already accepts Cross | +| Joins under `FORCE_NEST_LOOP_JOIN` | Excluded by `can_expand`, so not in a region; Orient only orients the cross joins among them | They enter a region (only with joins carrying the same hint) and are reordered as usual, but each pair only compares NestLoop | + +The hint means "use NLJ, not hash"; it does not promise execution in the written order. + +Ordering cost: Orient used to run after `Parameterize Mark Apply`, `Eliminate Index Filter`, and +`Materialize Mark Apply`, so it saw the rewritten `estimated_work`. After moving into Greedy, if an NLJ inner subtree +contains a subquery-style MarkApply, the direction is decided by its pre-rewrite estimate, while materialization +actually lowers the cost of re-running the inner side. This deviation is accepted for now (cross joins are already +ordered last in Greedy, so such cases are rare) and is left as future work together with `rescan_work`. The actual +effect of the rewrite on the `estimated_work` of such MarkApply nodes has not been verified yet. + +## 4. Foundation: Cost Model + +Greedy's decisions only make sense when costs are comparable; this section is the prerequisite for moving the +decisions to Greedy. + +### 4.1 Today + +| Location | Unit | Accounts for row width | +|---|---|---| +| scan impl rule `cost` | Rows (non-covering index Ɨ2, discount for sort hints) | No | +| `calculate_work` / `calculate_startup_work` | CPU units: `CPU_ROW_WORK = 4`, `CPU_OPERATOR_WORK = 1` | No | +| `GreedyJoinReorder` | Number of values (rows Ɨ columns) | Yes | + +- Each MarkApply lookup only counts `startup_work(scan) + OP = 1`, cheaper than scanning one row sequentially (4). +- HashJoin takes the default branch and misses the full read of the probe side. +- Among NLJs only Cross is handled separately, and it does not count per-pair comparisons; non-Cross NLJs do not + count "re-run the inner side for each outer row" at all. + +### 4.2 Unified row cost and join cost functions + +```rust +// width is the number of columns; each row has a fixed overhead, so zero-column rows are not estimated as 0 +fn rows_work(rows: usize, width: usize) -> usize { + rows * (CPU_ROW_WORK + width * CPU_OPERATOR_WORK) +} + +enum JoinImpl { Hash, NestLoop { full: bool }, IndexProbe { lookup_work: usize } } +struct InputCost { rows: usize, width: usize, work: usize } +impl JoinImpl { + fn work(self, left: &InputCost, right: &InputCost, out_rows: usize, out_width: usize) -> usize; +} +``` + +``` +Hash = build.work // read the build side once + + probe.work // read the probe side once + + rows_work(build) + build.rows Ɨ HASH_BUILD_WORK // store whole build rows in the hash table; hashing, allocation, and freeing calibrated on LMDB + + probe.rows Ɨ OP // one hash table lookup per probe row + + rows_work(out) // emit matching rows +IndexProbe = outer.work // read the outer side once + + outer.rows Ɨ lookup_work // one inner index seek per outer row; the inner side is not fully read + + rows_work(out) // emit matching rows +NestLoop = outer.work // read the outer side once + + outer.rows Ɨ (inner.work // re-run the inner side for each outer row + + inner.rows Ɨ OP) // compare each outer row with every inner row + + rows_work(out) // emit matching rows + (+ inner.work, Full join only: scan the inner side once more for unmatched rows) +``` + +- The cost functions only take statistics and do not depend on `&LogicalPlan`, because Greedy evaluates joins that + have not been built yet. Greedy and `calculate_work` share them. +- `calculate_work` and `calculate_startup_work` switch to `rows_work`; the scan impl rule is unchanged (candidates of + the same scan have the same number of columns, so changing the unit does not change their order). +- `rescan_work` is not introduced yet and is treated as equal to `work` (exception in 6.2). + +### 4.3 Cost of one probe: `lookup_work` + +A Hash probe is an in-memory hash lookup (about `OP`), but requires reading both sides fully; a MarkApply probe is a +storage-level seek that rebuilds the inner executor every time, so it costs more per probe, but the inner side is not +read fully. + +``` +lookup_work = PROBE_RESET_WORK(= CPU_ROW_WORK) + + seeks Ɨ āŒˆlog2(N)āŒ‰ Ɨ OP // seeks: PK/covering Unique = 1, non-covering Unique = 2 + + rows_per_lookup Ɨ rows_work(1, inner_width) +``` + +- N is the `values_len` of the index histogram, or `DEFAULT_TABLE_ROWS` when absent. An inner Filter is only counted + per row. +- Example: with 1M inner rows of width 2, scanning one row = 6 and one lookup = 30; the probe wins when the outer + side is smaller than roughly 1/4 of the inner side. +- Stored in `PhysicalOption.lookup_work: Option`, computed by the predicate function with the loader and + written during the probe rewrite (shared by Greedy and the fallback path). Subquery-style MarkApply such as semi / + anti also switch to this formula. + +### 4.4 Pair score: work done in this step + +``` +score(Hash / NestLoop) = work(join) āˆ’ left.work āˆ’ right.work +score(Probe) = outer.rows Ɨ lookup_work + rows_work(out) āˆ’ inner_leaf.work +``` + +- The score only counts the intermediate work of this step. With all-hash joins each leaf is read once regardless of + order, so counting the reads would only postpone large tables. The existing Greedy formula + (`product Ɨ (left_width + right_width) + min(left_size, right_size)`) already follows this idea. +- The `āˆ’ inner_leaf.work` term of Probe is a correction against the baseline (every leaf is read once), not work + done; when the inner side is large and the outer side small, the score is negative and Greedy merges it first. + Scores use `i128`, and the code must comment on the meaning of this term. + +## 5. Pipeline and the Validity of `join_decided` + +Order of after_batches in `db.rs`: + +1. `Greedy Join Reorder` (once_topdown) +2. `Parameterize Mark Apply`: `ParameterizeMarkApply` + `ParameterizeInnerJoin` (fix_point_topdown(1)) +3. `Eliminate Index Filter` +4. `Materialize Mark Apply` +5. `Orient Nested Loop Join` (removed by this RFC) +6. ... + +There is no other batch between Greedy and Parameterize. In the same batch, `ParameterizeMarkApply` only matches +`MarkApply`, and its `parameterize_right_subtree` only handles `Project` and `TableScan` without touching Join nodes, +so the mark is not destroyed. The `estimated_work` of the rewritten MarkApply must still use the `lookup_work` +written by Greedy, so that `calculate_work` (including costs shown in explain) is consistent with Greedy's decision. + +## 6. Verified Executor Behavior + +### 6.1 MarkApply inner probe (`execution/dql/mark_apply.rs`) + +- `predicates` contains every `on` equality (constant equalities are `truncate`d before `parameterize` returns); + after `next_join_tuple` fetches an inner row it checks all conditions with `predicates_matched`, so extra + conditions are also checked when there are more join keys than `key_columns`. +- The original join's `filter` is wrapped into a `Filter` above the MarkApply; the existing Filter on the probed side + is kept. +- The inner side is reset for each outer row and at most 1 row is taken, so this RFC only supports PK/Unique probes. + +### 6.2 NestLoopJoin and `reset` + +NLJ does not materialize its inner side: for each outer row it calls `reset_right_input()` and scans the inner side +to the end; a Full join scans the inner side once more at the end. `reset` on `SeqScan` / `IndexScan` repositions in +the storage layer. Stateful operators clear their state on `reset` and reset their inputs recursively: + +| Operator | `reset` behavior | +|---|---| +| HashJoin | `state = Build`, `bump.reset()`, reset both sides | +| Sort / ExternalSort | Clear `rows`, reset input | +| TopK | `output = None`, `arena.reset()`, reset input | +| HashAgg / SimpleAgg | Clear output state, reset input | +| Window | `state.clear()`, `reset_functions()`, reset input | +| NestLoopJoin | Back to `PullLeft`, reset both sides | +| MarkApply | Only resets the left side; keeps the right-side cache when `materialize` is on and a cache exists | + +So `rescan_work = work` is accurate; the only exception is MarkApply with `materialize` enabled. + +### 6.3 Sort hints do not affect Greedy leaves + +The sort discount in scan cost comes from `sort_elimination_hint` / `stream_aggregate_hint`. We need to confirm that +when Greedy treats a leaf as a hash input or rewrites it into a probe, it neither loses an ordering that an upper +operator relies on nor uses a discounted `work`: + +- Hints are pushed top-down during implementation by `annotate_hints_and_physical_options` + (`heuristic/optimizer.rs`) and only pass through `Filter` / `Project` / `Limit` / `TopK` / `Sort`; any other + operator such as Join replaces them with an empty hint. So scans under a Join (i.e. Greedy leaves) have + `hint_sum = 0` and no sort discount. +- `EliminateRedundantSort` / `UseStreamAggregate`, which consume ordering, run in `apply_annotated_post_rules` of the + same stage, before Greedy; their `ensure_order` only descends through `Childrens::Only` and never crosses a Join. + A Join's own `PhysicalOption` is `SortOption::None`. +- If a Sort / Aggregate sits between a leaf and the Join, that subtree is not `Filter* + TableScan`, so it cannot be a + probe inner side, and its ordering is only consumed inside the subtree. + +Conclusion: changing a leaf's access path in Greedy does not break any upper ordering dependency, and the `work` used +by the Hash shape contains no sort discount. + +## 7. Implementation Plan + +The cost model is the prerequisite for Greedy's decisions, so costs are unified first and decisions are handed over +second, in two commits to make regressions easy to locate. + +1. Unify costs (behavior only): introduce `rows_work` and `JoinImpl` / `InputCost`, fix the HashJoin, NLJ, and + MarkApply estimates, and change the Greedy score to the work done in this step. Existing join orders and explain / + slt expectations will change and need to be recalibrated as a whole. +2. Hand over decisions: extract the shared predicate function and `index_covering`; Greedy compares shapes, rewrites + probes directly, and marks the rest `join_decided`; `ParameterizeInnerJoin` skips `join_decided` joins; Greedy's + trigger is relaxed to `can_expand` and it takes over NestLoop direction (`FORCE_NEST_LOOP_JOIN` regions only + compare NestLoop); remove `OrientNestedLoopJoin` and its batch. + +Tests (unit tests via `make test`, slt via `make test-slt`): + +- `rows_work` and each `JoinImpl::work`: Full join, zero-column rows, saturating arithmetic. +- Greedy: large outer Ɨ small unique inner chooses hash; small outer Ɨ large unique inner chooses probe and merges + it first; the break-even point; the cheapest of several usable unique indexes; a composite unique index plus + constant equalities; default row counts without statistics. +- Downstream rules: `join_decided` joins are not rewritten; unmarked joins take the fallback. +- NestLoop direction: cross pairs inside a Greedy region, regions rooted at Cross, and `FORCE_NEST_LOOP_JOIN` regions + all reach the same direction as the existing `OrientNestedLoopJoin` cases (migrate its unit tests), and only NLJ + appears under `FORCE_NEST_LOOP_JOIN`. +- Coverage: when the scan chose a covering index and another covering Unique can be probed, the probe is priced with + `seeks = 1` and chosen, the rewritten IndexScan carries `covered_deserializers`, and execution skips the table + lookup. +- Correctness: after a probe, all `on` conditions, the join filter, and static filters are kept (covered by slt). + +## 8. Future Work + +- `rescan_work` for MarkApply with `materialize` enabled, and the deviation where Greedy decides NLJ direction from + the pre-parameterization / pre-materialization estimate of a subquery-style MarkApply in the inner side (3.10). +- Probes on non-unique indexes: the executor must support emitting multiple rows per outer row, and resetting the + inner side per row is expensive, so evaluate case by case; `rows_per_lookup` would then be estimated from the + index NDV. +- Joint NDV and uniqueness bound for composite join keys: under the hash shape, estimate rows with any PK/Unique bound + by the join keys, not only the index chosen during implementation (row count is a property of the data, not of + which index the scan uses). Single-column keys already get NDV from column histograms, so the benefit is only for + composite keys; do it when a concrete case appears. +- I/O modeling. diff --git a/src/bin/shell.rs b/src/bin/shell.rs index 83945998..f0b153d3 100644 --- a/src/bin/shell.rs +++ b/src/bin/shell.rs @@ -15,7 +15,7 @@ #[cfg(not(target_arch = "wasm32"))] mod native { use comfy_table::{Cell, Table}; - use kite_sql::db::{DBTransaction, DataBaseBuilder, Database, ResultIter}; + use kite_sql::db::{DataBaseBuilder, Database, ResultIter, Transaction}; use kite_sql::errors::DatabaseError; use kite_sql::storage::rocksdb::RocksStorage; use rustyline::config::Configurer; @@ -46,7 +46,7 @@ Metacommands: .views Show views .schema NAME Describe a table or view -Transaction commands: +StorageTransaction commands: BEGIN COMMIT ROLLBACK @@ -211,7 +211,6 @@ Transaction commands: let mut row_count = 0usize; while let Some(row) = iter.next_tuple(|_, tuple| { tuple - .values .iter() .map(|value| Cell::new(format!("{value}"))) .collect::>() @@ -225,7 +224,7 @@ Transaction commands: fn run_sql<'a>( database: &Database, - tx: &mut Option>, + tx: &mut Option>, sql: &str, ) -> Result<(), DatabaseError> { let started = Instant::now(); @@ -241,7 +240,7 @@ Transaction commands: fn handle_meta_command<'a>( command: &str, database: &Database, - tx: &mut Option>, + tx: &mut Option>, ) -> Result { let trimmed = command.trim(); if trimmed.is_empty() { @@ -284,7 +283,7 @@ Transaction commands: fn handle_command<'a>( command: &str, database: &'a Database, - tx: &mut Option>, + tx: &mut Option>, ) -> Result { let trimmed = command.trim(); if trimmed.is_empty() { @@ -303,7 +302,7 @@ Transaction commands: if tx.is_some() { eprintln!("transaction already started"); } else { - *tx = Some(database.new_transaction()?); + *tx = Some(database.new_transaction(false)?); println!("OK"); } return Ok(Mode::Continue); diff --git a/src/binder/aggregate.rs b/src/binder/aggregate.rs index 41e0aa12..63538cdb 100644 --- a/src/binder/aggregate.rs +++ b/src/binder/aggregate.rs @@ -18,7 +18,7 @@ use crate::expression::visitor::{walk_expr, ExprVisitor}; use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; use crate::planner::MetaArena; use crate::planner::{ExprRef, LogicalPlan, PlanArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; use crate::{ expression::ScalarExpression, @@ -41,7 +41,7 @@ impl ExprVisitor for AggregateCallCollector<'_> { } } -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub fn bind_aggregate( &mut self, children: LogicalPlan, @@ -360,10 +360,10 @@ impl<'a> AggregateOutputBinder<'a> { } } - fn output_ref( + fn output_ref( &mut self, expr: ExprRef, - arena: &mut dyn MetaArena, + arena: &mut A, ) -> Result, DatabaseError> { let output_count = self.agg_calls.len() + self.group_by_exprs.len(); self.agg_calls @@ -395,12 +395,8 @@ impl<'a> AggregateOutputBinder<'a> { } } -impl ExprVisitorMut for AggregateOutputBinder<'_> { - fn visit( - &mut self, - expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { +impl ExprVisitorMut for AggregateOutputBinder<'_> { + fn visit(&mut self, expr: &mut ExprRef, arena: &mut A) -> Result<(), DatabaseError> { if let ScalarExpression::Alias { alias: crate::expression::AliasType::Name(_), .. @@ -510,10 +506,10 @@ mod tests { alias: AliasType::Name("g".to_string()), }); - let constant = arena.alloc_expression(ScalarExpression::Constant(1_i32.into())); + let identity = arena.alloc_expression(ScalarExpression::Constant(1_i32.into())); let mut target = arena.alloc_expression(ScalarExpression::Alias { - expr: constant, - alias: AliasType::Expr(group_expr), + expr: group_expr, + alias: AliasType::Expr(identity), }); { diff --git a/src/binder/alter_table.rs b/src/binder/alter_table.rs index 2cb1d303..c55c2d87 100644 --- a/src/binder/alter_table.rs +++ b/src/binder/alter_table.rs @@ -22,10 +22,10 @@ use crate::planner::operator::alter_table::change_column::{ use crate::planner::operator::alter_table::drop_column::DropColumnOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_add_column( &mut self, table_name: TableName, diff --git a/src/binder/analyze.rs b/src/binder/analyze.rs index 2021329d..19be98fc 100644 --- a/src/binder/analyze.rs +++ b/src/binder/analyze.rs @@ -19,10 +19,10 @@ use crate::planner::operator::analyze::AnalyzeOperator; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_analyze( &mut self, table_name: TableName, @@ -30,7 +30,7 @@ impl> Binder<'_, '_, T, A> { ) -> Result { let table = self .context - .source_and_bind(table_name.clone(), None, None, true)? + .source_and_bind(table_name.clone(), None, None, true, arena)? .and_then(|source| { if let Source::Table(table) = source { Some(table) @@ -47,6 +47,8 @@ impl> Binder<'_, '_, T, A> { table_name, index_metas, histogram_buckets: None, + hll_precision: None, + statistics_relative_error: None, }), Childrens::Only(Box::new(scan_op)), )) diff --git a/src/binder/copy.rs b/src/binder/copy.rs index 0e292ac2..5829c04a 100644 --- a/src/binder/copy.rs +++ b/src/binder/copy.rs @@ -65,7 +65,7 @@ impl FromStr for ExtSource { } } -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(super) fn bind_copy_to_file( &mut self, target: ExtSource, diff --git a/src/binder/create_index.rs b/src/binder/create_index.rs index 24702eda..b9f0d39b 100644 --- a/src/binder/create_index.rs +++ b/src/binder/create_index.rs @@ -19,11 +19,11 @@ use crate::planner::operator::create_index::CreateIndexOperator; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::index::IndexType; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_create_index_source( &mut self, table_name: TableName, @@ -31,7 +31,7 @@ impl> Binder<'_, '_, T, A> { ) -> Result { let source = self .context - .source_and_bind(table_name.clone(), None, None, false)? + .source_and_bind(table_name.clone(), None, None, false, arena)? .ok_or(DatabaseError::SourceNotFound)?; match source { Source::Table(table) => { diff --git a/src/binder/create_table.rs b/src/binder/create_table.rs index 31d22512..d84d7b3c 100644 --- a/src/binder/create_table.rs +++ b/src/binder/create_table.rs @@ -18,11 +18,11 @@ use crate::errors::DatabaseError; use crate::planner::operator::create_table::CreateTableOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; use std::collections::HashSet; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { // TODO: TableConstraint pub(crate) fn bind_create_table( &mut self, diff --git a/src/binder/create_view.rs b/src/binder/create_view.rs index 7f49753c..d25b7a9c 100644 --- a/src/binder/create_view.rs +++ b/src/binder/create_view.rs @@ -20,10 +20,10 @@ use crate::expression::{AliasType, ScalarExpression}; use crate::planner::operator::create_view::CreateViewOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, ExprRef, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_create_view( &mut self, view_name: TableName, diff --git a/src/binder/delete.rs b/src/binder/delete.rs index 4eb6617b..0419fb3c 100644 --- a/src/binder/delete.rs +++ b/src/binder/delete.rs @@ -18,10 +18,10 @@ use crate::errors::DatabaseError; use crate::planner::operator::delete::DeleteOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_delete( &mut self, table_name: TableName, diff --git a/src/binder/describe.rs b/src/binder/describe.rs index 3729331d..d6f58d7d 100644 --- a/src/binder/describe.rs +++ b/src/binder/describe.rs @@ -18,10 +18,10 @@ use crate::errors::DatabaseError; use crate::planner::operator::describe::DescribeOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_describe( &mut self, table_name: TableName, diff --git a/src/binder/distinct.rs b/src/binder/distinct.rs index 305eeb8a..ff8f4c2f 100644 --- a/src/binder/distinct.rs +++ b/src/binder/distinct.rs @@ -20,10 +20,10 @@ use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::operator::sort::SortField; use crate::planner::MetaArena; use crate::planner::{ExprRef, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub fn bind_distinct( &mut self, children: LogicalPlan, @@ -84,7 +84,11 @@ impl<'a> DistinctOutputBinder<'a> { Self { select_list } } - fn output_ref(&mut self, expr: ExprRef, arena: &mut dyn MetaArena) -> Option { + fn output_ref( + &mut self, + expr: ExprRef, + arena: &mut A, + ) -> Option { self.select_list .iter() .position(|candidate| { @@ -100,12 +104,8 @@ impl<'a> DistinctOutputBinder<'a> { } } -impl ExprVisitorMut for DistinctOutputBinder<'_> { - fn visit( - &mut self, - expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { +impl ExprVisitorMut for DistinctOutputBinder<'_> { + fn visit(&mut self, expr: &mut ExprRef, arena: &mut A) -> Result<(), DatabaseError> { let is_named_alias = matches!( arena.expression(*expr), ScalarExpression::Alias { @@ -197,10 +197,10 @@ mod tests { alias: AliasType::Name("v".to_string()), }); - let constant = arena.alloc_expression(ScalarExpression::Constant(1_i32.into())); + let identity = arena.alloc_expression(ScalarExpression::Constant(1_i32.into())); let mut target = arena.alloc_expression(ScalarExpression::Alias { - expr: constant, - alias: AliasType::Expr(expr), + expr, + alias: AliasType::Expr(identity), }); { diff --git a/src/binder/drop_index.rs b/src/binder/drop_index.rs index 149b61ea..20f76e24 100644 --- a/src/binder/drop_index.rs +++ b/src/binder/drop_index.rs @@ -18,10 +18,10 @@ use crate::errors::DatabaseError; use crate::planner::operator::drop_index::DropIndexOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_drop_index( &mut self, table_name: TableName, diff --git a/src/binder/drop_table.rs b/src/binder/drop_table.rs index 047c90fc..ea2e0950 100644 --- a/src/binder/drop_table.rs +++ b/src/binder/drop_table.rs @@ -18,10 +18,10 @@ use crate::errors::DatabaseError; use crate::planner::operator::drop_table::DropTableOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_drop_table( &mut self, table_name: TableName, diff --git a/src/binder/drop_view.rs b/src/binder/drop_view.rs index 936c6490..87eab8b5 100644 --- a/src/binder/drop_view.rs +++ b/src/binder/drop_view.rs @@ -18,10 +18,10 @@ use crate::errors::DatabaseError; use crate::planner::operator::drop_view::DropViewOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_drop_view( &mut self, view_name: TableName, diff --git a/src/binder/explain.rs b/src/binder/explain.rs index 0bd109af..811e3f87 100644 --- a/src/binder/explain.rs +++ b/src/binder/explain.rs @@ -16,10 +16,10 @@ use crate::binder::Binder; use crate::errors::DatabaseError; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_explain(&mut self, plan: LogicalPlan) -> Result { Ok(LogicalPlan::new( Operator::Explain, diff --git a/src/binder/expr.rs b/src/binder/expr.rs index 29df7e87..d1906d17 100644 --- a/src/binder/expr.rs +++ b/src/binder/expr.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::catalog::ColumnRef; use crate::errors::DatabaseError; use crate::expression; use crate::expression::agg::AggKind; @@ -26,7 +25,7 @@ use crate::expression::{AliasType, ScalarExpression, TypeCast}; use crate::planner::operator::mark_apply::MarkApplyQuantifier; use crate::planner::operator::scalar_subquery::ScalarSubqueryOperator; use crate::planner::{ExprRef, LogicalPlan, MetaArena, PlanArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::{DataValue, Utf8Type}; use crate::types::{CharLengthUnits, LogicalType}; @@ -38,51 +37,12 @@ macro_rules! try_default { }; } -impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> { - fn find_column_in_schema<'schema>( - schema_ref: impl IntoIterator, - arena: &PlanArena, - column_name: &str, - ) -> Option<(usize, ColumnRef)> { - schema_ref - .into_iter() - .enumerate() - .find(|(_, column)| arena.column(**column).name() == column_name) - .map(|(position, column)| (position, *column)) - } - - fn find_column_in_scope( - context: &BinderContext<'a, T>, - arena: &PlanArena, - column_name: &str, - ) -> Result, DatabaseError> { - let mut position_offset = 0; - let mut found = None; - - for bound_source in &context.bind_table { - let source = &bound_source.source; - - if let Some((position, column)) = - Self::find_column_in_schema(source.schema().iter(), arena, column_name) - { - if found.is_some() { - return Err(DatabaseError::AmbiguousColumn(column_name.to_string())); - } - found = Some(ScalarExpression::column_expr( - column, - position_offset + position, - )); - } - - position_offset += source.schema_len(); - } - - Ok(found) - } +impl<'a, T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> { pub(crate) fn bind_temp_table( &mut self, expr: ScalarExpression, sub_query: LogicalPlan, + correlated: bool, arena: &mut PlanArena, ) -> Result<(ScalarExpression, LogicalPlan), DatabaseError> { let (exprs, is_tuple) = match expr { @@ -94,16 +54,18 @@ impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> for (position, expr) in exprs.into_iter().enumerate() { let (alias_expr, alias_ref) = self.bind_temp_table_alias(expr, position, arena); - let predicate_alias = arena.alloc_expression(alias_expr); - // The projection evaluates against the subquery-local tuple, while the - // predicate evaluates against the combined left/right tuple. Clone the - // complete expression graph so position rewrites in either context do - // not mutate the other one through shared ExprRefs. - let projection_alias = predicate_alias.clone_expression(arena)?; + // The projection evaluates against the subquery-local tuple, while a + // correlated predicate evaluates `expr` against the combined tuple. + // Clone the complete expression graph so position rewrites in either + // context do not mutate the other one through shared ExprRefs. + let projection_alias = arena.alloc_expression(alias_expr).clone_expression(arena)?; if !is_tuple { let alias_plan = Self::build_project_plan(sub_query, vec![projection_alias]); - return Ok((arena.expression(predicate_alias).clone(), alias_plan)); + // An uncorrelated predicate reads the column produced by the projection. + let predicate = if correlated { expr } else { alias_ref }; + return Ok((arena.expression(predicate).clone(), alias_plan)); } + alias_exprs.push(projection_alias); alias_refs.push(alias_ref); } @@ -316,7 +278,7 @@ impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> )); } - let (alias_expr, sub_query) = self.bind_temp_table(column, sub_query, arena)?; + let (alias_expr, sub_query) = self.bind_temp_table(column, sub_query, correlated, arena)?; let alias_expr = arena.alloc_expression(alias_expr); let predicate = ScalarExpression::Binary { op: compare_op, @@ -366,7 +328,7 @@ impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> { let input = match self.context.using.get(column_name) { Some(column) => Some(column.visible_expr(arena)?), - None => Self::find_column_in_scope(&self.context, arena, column_name)?, + None => self.context.resolve_column(column_name, None, arena)?, }; if let Some(input) = input { return Ok(input); @@ -392,24 +354,20 @@ impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> } if let Some(table) = table_name.or(bind_table_name) { let mut is_outer = false; - let (source, position_offset) = - match Self::resolve_source_columns_in_scope(&self.context, table) { - Ok(source) => source, - Err(err) => { - if let Some(parent) = self.parent { - self.context.mark_join_outer_ref(); - is_outer = true; - Self::resolve_source_columns_in_scope(parent, table).map_err(|_| err)? - } else { - return Err(err); - } + let expr = match self.context.resolve_column(column_name, Some(table), arena) { + Ok(expr) => expr, + Err(err @ DatabaseError::InvalidTable { .. }) => { + if let Some(parent) = self.parent { + self.context.mark_join_outer_ref(); + is_outer = true; + parent.resolve_column(column_name, Some(table), arena)? + } else { + return Err(err); } - }; - let (position, column) = - Self::find_column_in_schema(source.schema().iter(), arena, column_name) - .ok_or_else(|| DatabaseError::column_not_found(column_name.to_string()))?; - - let expr = ScalarExpression::column_expr(column, position_offset + position); + } + Err(err) => return Err(err), + } + .ok_or_else(|| DatabaseError::column_not_found(column_name.to_string()))?; if is_outer && self.context.capture_scalar_outer { return self.capture_scalar_outer(expr, arena); } @@ -420,7 +378,7 @@ impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> |context: &BinderContext<'a, T>| -> Result, DatabaseError> { match context.using.get(column_name) { Some(using_column) => using_column.visible_expr(arena).map(Some), - None => Self::find_column_in_scope(context, arena, column_name), + None => context.resolve_column(column_name, None, arena), } }; let mut got_column = find_visible_column(&self.context)?; @@ -556,7 +514,13 @@ impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> if args.len() != 1 { return Err(DatabaseError::MisMatch("number of sum() parameters", "1")); } - args[0].return_type(arena).into_owned() + let ty = args[0].return_type(arena).into_owned(); + // Text is summed by its numeric prefix, as in SQLite. + if matches!(ty, LogicalType::Varchar(..) | LogicalType::Char(..)) { + LogicalType::Double + } else { + ty + } } AggKind::Min => { if args.len() != 1 { @@ -576,6 +540,21 @@ impl<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, '_, T, A> } LogicalType::Double } + AggKind::Total => { + if args.len() != 1 { + return Err(DatabaseError::MisMatch("number of total() parameters", "1")); + } + LogicalType::Double + } + AggKind::GroupConcat => { + if args.len() != 1 { + return Err(DatabaseError::MisMatch( + "number of group_concat() parameters", + "1", + )); + } + LogicalType::Varchar(None, CharLengthUnits::Characters) + } }; Ok(ScalarExpression::AggCall { distinct: is_distinct, diff --git a/src/binder/insert.rs b/src/binder/insert.rs index 124561d7..0778edb1 100644 --- a/src/binder/insert.rs +++ b/src/binder/insert.rs @@ -19,11 +19,11 @@ use crate::planner::operator::insert::InsertOperator; use crate::planner::operator::values::ValuesOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, ExprRef, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::tuple::Schema; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_insert_values( &mut self, table_name: TableName, diff --git a/src/binder/mod.rs b/src/binder/mod.rs index 18e13ef8..61189d7a 100644 --- a/src/binder/mod.rs +++ b/src/binder/mod.rs @@ -51,8 +51,9 @@ mod truncate; mod update; mod window; +pub use crate::planner::operator::CommandType; #[cfg(feature = "parser")] -pub use parser::{command_type, CommandType, Statement}; +pub use parser::{command_type, Statement}; #[cfg(feature = "orm")] pub use select::{BindPlanFrom, BindPlanSelectList}; #[cfg(feature = "orm")] @@ -67,7 +68,7 @@ use crate::expression::{ScalarExpression, TypeCast}; use crate::planner::operator::join::JoinType; use crate::planner::operator::mark_apply::MarkApplyQuantifier; use crate::planner::{ExprRef, LogicalPlan, PlanArena, PlanRef, ScalarQueryRef}; -use crate::storage::{TableCache, Transaction, ViewCache}; +use crate::storage::{StorageTransaction, TableCache, ViewCache}; use crate::types::tuple::Schema; use crate::types::LogicalType; @@ -128,6 +129,31 @@ pub struct BoundSource<'a> { pub(crate) alias: Option, pub(crate) join_type: Option, pub(crate) source: Source<'a>, + // The first source with this visible name owns qualified name resolution. + qualified_visible: bool, +} + +// Sorted by (name, visible source name, source index, position). A source entry +// has no column; names are borrowed from existing metadata instead of copied. +#[derive(Debug, Clone, Copy)] +struct NameBinding { + source: usize, + column: Option, + position: usize, +} + +impl NameBinding { + fn key<'n>( + &self, + sources: &'n [BoundSource<'_>], + arena: &'n PlanArena<'_>, + ) -> (&'n str, Option<&'n str>) { + let source_name = sources[self.source].visible_name().as_ref(); + match self.column { + Some(column) => (arena.column(column).name(), Some(source_name)), + None => (source_name, None), + } + } } #[derive(Debug, Clone)] @@ -144,13 +170,6 @@ pub(crate) struct CteCheckpoint { } impl BoundSource<'_> { - pub(crate) fn matches_name(&self, table_name: &str) -> bool { - match &self.alias { - Some(alias) => alias.as_ref() == table_name, - None => self.table_name.as_ref() == table_name, - } - } - pub(crate) fn same_binding( &self, table_name: &TableName, @@ -240,7 +259,7 @@ pub(crate) struct BoundScalarQuery { param_bindings: Vec<(ScalarQueryRef, ExprRef)>, } -pub struct BinderContext<'a, T: Transaction> { +pub struct BinderContext<'a, T: StorageTransaction> { pub(crate) scala_functions: &'a ScalaFunctions, pub(crate) table_functions: &'a TableFunctions, pub(crate) table_cache: &'a TableCache, @@ -249,6 +268,10 @@ pub struct BinderContext<'a, T: Transaction> { // Tips: retain binding order so wildcard expansion and position derivation // follow FROM/JOIN order directly. pub(crate) bind_table: Vec>, + name_bindings: Vec, + // Sources before this index were inherited from an enclosing query by `fork` + // and do not share the current FROM namespace. + source_scope_start: usize, ctes: Vec, pub(crate) cte_depth: usize, // alias @@ -261,7 +284,7 @@ pub struct BinderContext<'a, T: Transaction> { using: HashMap, has_join_outer_ref: bool, // subquery - sub_queries: HashMap>, + sub_queries: Vec, scalar_queries: Vec, capture_scalar_outer: bool, scalar_outer_bindings: Vec<(ScalarQueryRef, ExprRef)>, @@ -303,7 +326,7 @@ impl Source<'_> { } } -impl<'a, T: Transaction> BinderContext<'a, T> { +impl<'a, T: StorageTransaction> BinderContext<'a, T> { pub fn new( table_cache: &'a TableCache, view_cache: &'a ViewCache, @@ -318,6 +341,8 @@ impl<'a, T: Transaction> BinderContext<'a, T> { view_cache, transaction, bind_table: Default::default(), + name_bindings: Vec::new(), + source_scope_start: 0, ctes: Default::default(), cte_depth: 0, expr_aliases: Default::default(), @@ -347,6 +372,8 @@ impl<'a, T: Transaction> BinderContext<'a, T> { view_cache: self.view_cache, transaction: self.transaction, bind_table: self.bind_table.clone(), + name_bindings: self.name_bindings.clone(), + source_scope_start: self.bind_table.len(), ctes: self.ctes.clone(), cte_depth: self.cte_depth, expr_aliases: self.expr_aliases.clone(), @@ -430,14 +457,7 @@ impl<'a, T: Transaction> BinderContext<'a, T> { } pub fn sub_query(&mut self, sub_query: SubQueryType) { - self.sub_queries - .entry(self.bind_step) - .or_default() - .push(sub_query) - } - - pub fn sub_queries_at_now(&mut self) -> Option> { - self.sub_queries.remove(&self.bind_step) + self.sub_queries.push(sub_query); } pub fn mark_join_outer_ref(&mut self) { @@ -476,55 +496,40 @@ impl<'a, T: Transaction> BinderContext<'a, T> { } } - #[allow(unused_assignments)] pub fn source_and_bind( &mut self, table_name: TableName, alias: Option<&TableName>, join_type: Option, only_table: bool, + arena: &PlanArena, ) -> Result>, DatabaseError> { - let mut source = None; - - source = if let Some(real_name) = self.table_aliases.get(table_name.as_ref()) { - self.transaction.table(self.table_cache, real_name.clone()) - } else { - self.transaction.table(self.table_cache, table_name.clone()) - }? - .map(Source::Table); - - if source.is_none() && !only_table { - source = if let Some(real_name) = self.table_aliases.get(table_name.as_ref()) { - self.transaction.view( - self.table_cache, - self.view_cache, - self.scala_functions, - self.table_functions, - real_name.clone(), - ) - } else { - self.transaction.view( - self.table_cache, - self.view_cache, - self.scala_functions, - self.table_functions, - table_name.clone(), - ) - }? - .map(Source::View); + if let Some(source) = self.bound_source(&table_name, alias, join_type) { + if !only_table || matches!(source, Source::Table(_)) { + return Ok(Some(source.clone())); + } } + let source = if only_table { + let name = self.table_aliases.get(&table_name).unwrap_or(&table_name); + self.transaction + .table(self.table_cache, name.clone())? + .map(Source::Table) + } else { + self.source(&table_name)? + }; if let Some(source) = &source { self.add_bound_source( table_name.clone(), alias.cloned(), join_type, source.clone(), + arena, ); } Ok(source) } - pub fn source(&self, table_name: &TableName) -> Result>, DatabaseError> { + pub fn source(&self, table_name: &TableName) -> Result>, DatabaseError> { let mut source = if let Some(real_name) = self.table_aliases.get(table_name.as_ref()) { self.transaction.table(self.table_cache, real_name.clone()) } else { @@ -556,27 +561,164 @@ impl<'a, T: Transaction> BinderContext<'a, T> { Ok(source) } + /// Rejects a FROM source whose visible name is already bound in the current + /// FROM namespace. Two such sources would also expose identical column + /// identities, which later stages cannot tell apart. + pub(crate) fn ensure_unique_source_name(&self, name: &TableName) -> Result<(), DatabaseError> { + if self.bind_table[self.source_scope_start..] + .iter() + .any(|bound_source| bound_source.visible_name() == name) + { + return Err(DatabaseError::DuplicateTableName(name.to_string())); + } + Ok(()) + } + + fn bound_source( + &self, + table_name: &TableName, + alias: Option<&TableName>, + join_type: Option, + ) -> Option<&Source<'a>> { + self.bind_table + .iter() + .find(|bound| bound.same_binding(table_name, alias, join_type)) + .map(|bound| &bound.source) + } + pub fn add_bound_source( &mut self, table_name: TableName, alias: Option, join_type: Option, source: Source<'a>, + arena: &PlanArena, ) { - if let Some(bound_source) = self - .bind_table - .iter_mut() - .find(|bound_source| bound_source.same_binding(&table_name, alias.as_ref(), join_type)) + if self + .bound_source(&table_name, alias.as_ref(), join_type) + .is_some() { - bound_source.source = source; return; } + let offset = self + .bind_table + .iter() + .map(|bound| bound.source.schema_len()) + .sum::(); + let index = self.bind_table.len(); + let qualified_visible = self + .name_binding(alias.as_ref().unwrap_or(&table_name), None, arena) + .is_none(); self.bind_table.push(BoundSource { table_name, alias, join_type, source, + qualified_visible, }); + self.name_bindings + .reserve(self.bind_table[index].source.schema_len() + 1); + self.insert_name_binding( + NameBinding { + source: index, + column: None, + position: offset, + }, + arena, + ); + for position in 0..self.bind_table[index].source.schema_len() { + let column = self.bind_table[index].source.schema()[position]; + self.insert_name_binding( + NameBinding { + source: index, + column: Some(column), + position: offset + position, + }, + arena, + ); + } + } + + fn insert_name_binding(&mut self, binding: NameBinding, arena: &PlanArena) { + let key = binding.key(&self.bind_table, arena); + let index = self.name_bindings.partition_point(|existing| { + existing + .key(&self.bind_table, arena) + .cmp(&key) + .then(existing.source.cmp(&binding.source)) + .then(existing.position.cmp(&binding.position)) + .is_lt() + }); + self.name_bindings.insert(index, binding); + } + + fn name_binding( + &self, + name: &str, + source: Option<&str>, + arena: &PlanArena, + ) -> Option { + let key = (name, source); + let index = self + .name_bindings + .partition_point(|binding| binding.key(&self.bind_table, arena) < key); + self.name_bindings + .get(index) + .filter(|binding| { + binding.key(&self.bind_table, arena) == key + && (source.is_none() || self.bind_table[binding.source].qualified_visible) + }) + .copied() + } + + fn resolve_column( + &self, + name: &str, + source: Option<&str>, + arena: &PlanArena, + ) -> Result, DatabaseError> { + // Fast path for single source + if let [bound] = self.bind_table.as_slice() { + if let Some(source) = source { + if bound.visible_name().as_ref() != source { + return Err(DatabaseError::invalid_table(source)); + } + } + return Ok(bound + .source + .schema() + .iter() + .enumerate() + .find(|(_, column)| arena.column(**column).name() == name) + .map(|(position, column)| ScalarExpression::column_expr(*column, position))); + } + let binding = if let Some(source) = source { + let binding = self.name_binding(name, Some(source), arena); + if binding.is_none() && self.name_binding(source, None, arena).is_none() { + return Err(DatabaseError::invalid_table(source)); + } + binding + } else { + let start = self + .name_bindings + .partition_point(|binding| binding.key(&self.bind_table, arena).0 < name); + let mut columns = self.name_bindings[start..] + .iter() + .take_while(|binding| binding.key(&self.bind_table, arena).0 == name) + .filter(|binding| binding.column.is_some()); + let Some(first) = columns.next() else { + return Ok(None); + }; + if columns.any(|binding| binding.source != first.source) { + return Err(DatabaseError::AmbiguousColumn(name.to_string())); + } + Some(*first) + }; + Ok(binding.and_then(|binding| { + binding + .column + .map(|column| ScalarExpression::column_expr(column, binding.position)) + })) } // Tips: The order of this index is based on Aggregate being bound first. @@ -625,7 +767,7 @@ impl<'a, T: Transaction> BinderContext<'a, T> { } } -pub struct Binder<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> { +pub struct Binder<'a, 'parent, T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>> { pub(crate) context: BinderContext<'a, T>, pub(crate) args: &'a A, pub(crate) force_spill: bool, @@ -634,7 +776,9 @@ pub struct Binder<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> pub(crate) parent: Option<&'parent BinderContext<'a, T>>, } -impl<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'parent, T, A> { +impl<'a, 'parent, T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>> + Binder<'a, 'parent, T, A> +{ pub fn new( context: BinderContext<'a, T>, args: &'a A, @@ -665,13 +809,20 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, ' false } - pub(crate) fn extend(&mut self, context: BinderContext<'a, T>) { + pub(crate) fn extend( + &mut self, + context: BinderContext<'a, T>, + arena: &PlanArena, + ) -> Result<(), DatabaseError> { for bound_source in context.bind_table { + self.context + .ensure_unique_source_name(bound_source.visible_name())?; self.context.add_bound_source( bound_source.table_name, bound_source.alias, bound_source.join_type, bound_source.source, + arena, ); } for (key, expr) in context.expr_aliases { @@ -680,6 +831,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, ' for (key, table_name) in context.table_aliases { self.context.table_aliases.insert(key, table_name); } + Ok(()) } } @@ -697,7 +849,9 @@ pub mod test { use crate::errors::DatabaseError; use crate::planner::{LogicalPlan, PlanArena, TableArenaCell}; use crate::storage::rocksdb::RocksStorage; - use crate::storage::{table_codec::TableCodec, Storage, TableCache, Transaction, ViewCache}; + use crate::storage::{ + table_codec::TableCodec, Storage, StorageTransaction, TableCache, ViewCache, + }; use crate::types::ColumnId; use crate::types::LogicalType::Integer; use std::path::PathBuf; @@ -829,6 +983,83 @@ pub mod test { Ok(storage) } + #[test] + fn sorted_bindings_preserve_names_positions_and_scope_isolation() -> Result<(), DatabaseError> { + use crate::binder::Source; + use crate::expression::ScalarExpression; + + let state = build_t1_table()?; + let arena = PlanArena::new(&state.table_arena); + let transaction = state.storage.transaction()?; + let scala_functions = Default::default(); + let table_functions = Default::default(); + let mut context = BinderContext::new( + &state.table_cache, + &state.view_cache, + &transaction, + &scala_functions, + &table_functions, + ); + let c1 = state.table.get_column_by_name("c1").unwrap(); + let c2 = state.table.get_column_by_name("c2").unwrap(); + // A source name may equal a column name; duplicates within that source + // retain their first position and are not ambiguous. + context.add_bound_source( + "a".into(), + Some("c1".into()), + None, + Source::Schema(vec![c1, c1]), + &arena, + ); + assert!(matches!( + context.resolve_column("c1", None, &arena)?, + Some(ScalarExpression::ColumnRef { column, position: 0 }) if column == c1 + )); + context.add_bound_source("b".into(), None, None, Source::Schema(vec![c2]), &arena); + assert!(matches!( + context.resolve_column("c2", Some("b"), &arena)?, + Some(ScalarExpression::ColumnRef { column, position: 2 }) if column == c2 + )); + // Re-registering a binding preserves its original columns and positions. + context.add_bound_source("b".into(), None, None, Source::Schema(vec![c1]), &arena); + assert!(matches!( + context.resolve_column("c2", Some("b"), &arena)?, + Some(ScalarExpression::ColumnRef { column, position: 2 }) if column == c2 + )); + assert!(matches!( + context.resolve_column("c1", Some("c1"), &arena)?, + Some(ScalarExpression::ColumnRef { position: 0, .. }) + )); + context.add_bound_source( + "a".into(), + Some("c1".into()), + None, + Source::Schema(vec![c2]), + &arena, + ); + assert_eq!(context.name_binding("b", None, &arena).unwrap().position, 2); + assert!(matches!( + context.resolve_column("c1", None, &arena)?, + Some(ScalarExpression::ColumnRef { column, position: 0 }) if column == c1 + )); + assert!(matches!( + context.resolve_column("c1", Some("a"), &arena), + Err(DatabaseError::InvalidTable { .. }) + )); + let mut fork = context.fork(); + fork.add_bound_source("c".into(), None, None, Source::Schema(vec![c1]), &arena); + assert!(matches!( + fork.resolve_column("c1", None, &arena), + Err(DatabaseError::AmbiguousColumn(_)) + )); + assert!(context.resolve_column("c1", None, &arena)?.is_some()); + assert!(context + .fork_empty() + .resolve_column("c1", None, &arena)? + .is_none()); + Ok(()) + } + #[test] pub fn test_valid_identifier() { assert!(is_valid_identifier("valid_table")); diff --git a/src/binder/parser.rs b/src/binder/parser.rs index 71357ab9..b684c7b3 100644 --- a/src/binder/parser.rs +++ b/src/binder/parser.rs @@ -13,14 +13,15 @@ // limitations under the License. use super::select::{ - BindPlanAggregated, BindPlanComplete, BindPlanDistinct, BindPlanFiltered, BindPlanFrom, - BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput, TableAliasInput, + AliasPurpose, BindPlanAggregated, BindPlanComplete, BindPlanDistinct, BindPlanFiltered, + BindPlanFrom, BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput, + TableAliasInput, }; use super::{is_valid_identifier, Binder, CteBinding, QueryBindStep, SetOperatorKind}; #[cfg(feature = "copy")] use crate::binder::copy::{ExtSource, FileFormat}; use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef, TableName}; -use crate::db::{BindSource, DBTransaction, Database, DatabaseIter, State, TransactionIter}; +use crate::db::{BindSource, Database, DatabaseIter, State, Transaction, TransactionIter}; use crate::errors::{DatabaseError, SqlErrorSpan}; use crate::expression; use crate::expression::agg::AggKind; @@ -35,10 +36,10 @@ use crate::planner::operator::mark_apply::MarkApplyQuantifier; use crate::planner::operator::project::ProjectOperator; use crate::planner::operator::recursive_cte::{RecursiveCteOperator, RecursiveScanOperator}; use crate::planner::operator::sort::SortField; -use crate::planner::operator::Operator; +use crate::planner::operator::{CommandType, Operator}; use crate::planner::MetaArena; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; -use crate::storage::{Storage, Transaction}; +use crate::storage::{Storage, StorageTransaction}; use crate::types::value::{DataValue, Utf8Type}; use crate::types::{CharLengthUnits, ColumnId, LogicalType}; pub(super) use sqlparser::ast::{ @@ -62,14 +63,6 @@ use std::slice; /// Parsed SQL statement type used by KiteSQL SQL frontend APIs. pub type Statement = sqlparser::ast::Statement; -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub enum CommandType { - DQL, - DML, - DDL, - Analyze, -} - pub(crate) trait AttachSpanSource { fn sql_error_span(self) -> Option; } @@ -133,6 +126,8 @@ pub fn command_type(stmt: &Statement) -> Result { | Statement::Delete(_) | Statement::Insert(_) => Ok(CommandType::DML), #[cfg(feature = "copy")] + Statement::Copy { to: true, .. } => Ok(CommandType::DQL), + #[cfg(feature = "copy")] Statement::Copy { .. } => Ok(CommandType::DML), stmt => Err(DatabaseError::UnsupportedStmt(stmt.to_string())), } @@ -157,6 +152,15 @@ impl State { } } +/// Statements that only read storage can run in a read transaction. +fn statement_is_read_only(statement: &Statement) -> Result { + Ok(match statement { + // Locking reads (`FOR UPDATE`) are treated as writes. + Statement::Query(query) => query.locks.is_empty(), + _ => command_type(statement)?.is_read_only(), + }) +} + fn statement_mutates_catalog_or_statistics(statement: &Statement) -> Result { Ok(matches!( command_type(statement)?, @@ -264,9 +268,15 @@ impl Database { .with_sql_context(sql)); } + let read_only = statements + .iter() + .try_fold(true, |read_only, stmt| { + Ok::<_, DatabaseError>(read_only && statement_is_read_only(stmt)?) + }) + .map_err(|err| err.with_sql_context(sql))?; let transaction = Box::into_raw(Box::new( self.storage - .transaction_with_isolation(self.transaction_isolation)?, + .transaction_with_isolation(self.transaction_isolation, read_only)?, )); let mut statements = statements.into_iter().peekable(); @@ -309,7 +319,7 @@ impl Database { } } -impl<'txn, S: Storage> DBTransaction<'txn, S> { +impl<'txn, S: Storage> Transaction<'txn, S> { /// Runs SQL inside the current transaction and returns the final result iterator. pub fn run<'a, T: AsRef>( &'a mut self, @@ -333,7 +343,7 @@ impl<'txn, S: Storage> DBTransaction<'txn, S> { .ok_or_else(|| DatabaseError::EmptyStatement.with_sql_context(sql))?; for statement in statements { - BindSource::execute(&mut *self, |state, tx| { + BindSource::execute(&mut *self, false, |state, tx| { state.build_plan(&[], tx, |binder, arena| binder.bind(&statement, arena)) }) .map_err(|err| err.with_sql_context(sql))? @@ -341,7 +351,7 @@ impl<'txn, S: Storage> DBTransaction<'txn, S> { .map_err(|err| err.with_sql_context(sql))?; } - BindSource::execute(self, |state, tx| { + BindSource::execute(self, false, |state, tx| { state.build_plan(&[], tx, |binder, arena| binder.bind(&last_statement, arena)) }) .map_err(|err| err.with_sql_context(sql)) @@ -350,7 +360,7 @@ impl<'txn, S: Storage> DBTransaction<'txn, S> { struct BindStatementStart<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'s mut Binder<'a, 'b, T, A>, @@ -365,12 +375,12 @@ struct UpdateExprTargetRemapper<'a> { target_schema: &'a [ColumnRef], } -impl ExprVisitorMut for UpdateExprTargetRemapper<'_> { +impl ExprVisitorMut for UpdateExprTargetRemapper<'_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, position: &mut usize, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { let Some(target_position) = self .target_schema @@ -389,7 +399,7 @@ impl ExprVisitorMut for UpdateExprTargetRemapper<'_> { impl<'s, 'a: 'b, 'b, 'arena, T, A> BindStatementStart<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn statement(self, stmt: &Statement) -> Result { @@ -1062,7 +1072,7 @@ where let source = self .binder .context - .source_and_bind(table_name.clone(), None, None, false)? + .source_and_bind(table_name.clone(), None, None, false, self.arena)? .ok_or(DatabaseError::TableNotFound)?; let values_len = expr_rows[0].len(); @@ -1216,6 +1226,7 @@ where Binder::<'a, 'b, T, A>::resolve_source_columns_in_scope( &self.binder.context, &table_name, + self.arena, )?; let target_schema = target_source.schema().to_vec(); @@ -1318,7 +1329,7 @@ where }) .collect(); plan = LogicalPlan::new( - Operator::Project(ProjectOperator { exprs }), + Operator::Project(ProjectOperator::new(exprs)), Childrens::Only(Box::new(plan)), ); } @@ -1407,7 +1418,7 @@ impl BindStatementComplete { impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanStart<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { #[allow(clippy::wrong_self_convention)] @@ -1440,7 +1451,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanFrom<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn select_list_from_sql( @@ -1457,7 +1468,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanSelectList<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn where_sql( @@ -1479,7 +1490,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanFiltered<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn aggregate_sql( @@ -1586,7 +1597,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> super::select::BindPlanWindowed<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn distinct_sql( @@ -1599,7 +1610,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanProjected<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn select_into_sql( @@ -1716,7 +1727,7 @@ impl TryFrom<&Value> for DataValue { )); } let mut bytes = Vec::with_capacity(hex.len() / 2); - for pair in hex.as_bytes().chunks_exact(2) { + for pair in hex.as_bytes().as_chunks::<2>().0 { let Some(high) = (pair[0] as char).to_digit(16) else { return Err(DatabaseError::InvalidValue("invalid hex digit".to_string())); }; @@ -2061,7 +2072,9 @@ fn copy_file_format(options: Vec) -> Result> Binder<'a, 'parent, T, A> { +impl<'a, 'parent, T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>> + Binder<'a, 'parent, T, A> +{ fn bind_table_ref_sql( &mut self, from: &TableWithJoins, @@ -2091,10 +2104,15 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, ' sql_optional_table_alias(alias.clone()), arena, ), + TableFactor::Derived { lateral: true, .. } => Err(DatabaseError::UnsupportedStmt( + "LATERAL derived tables are not supported".to_string(), + )), TableFactor::Derived { subquery, alias, .. } => { - let mut binder = Binder::new(self.context.fork(), self.args, Some(&self.context)); + // A non-LATERAL derived table cannot see the sources bound before + // it in the same FROM clause. + let mut binder = Binder::new(self.context.fork_empty(), self.args, None); let plan = binder.bind_query(subquery, arena)?; self.bind_derived_source( plan, @@ -2157,7 +2175,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, ' let right = binder.bind_single_table_ref_sql(relation, Some(join_type), arena)?; (right, binder.context) }; - self.extend(context); + self.extend(context, arena)?; let constraint = match joint_condition { Some(constraint) => self.bind_join_constraint_sql(constraint, arena)?, @@ -3103,7 +3121,14 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, ' let mut binder = Binder::new(self.context.fork(), self.args, Some(&self.context)); let plan = binder.bind_query(&cte.query, arena)?; let source_name = arena.temp_table(); - binder.bind_alias(plan, &alias.columns, alias.name.clone(), source_name, arena) + binder.bind_alias( + plan, + &alias.columns, + alias.name.clone(), + source_name, + AliasPurpose::Output, + arena, + ) } fn add_cte_plan( @@ -3165,6 +3190,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, ' &alias.columns, alias.name.clone(), source_name, + AliasPurpose::Output, arena, )?; let schema = anchor.output_schema(arena).clone(); @@ -3484,6 +3510,16 @@ mod tests { command_type(&database.prepare_all("truncate table t")?.pop().unwrap())?, CommandType::DML ); + #[cfg(feature = "copy")] + for (sql, expected) in [ + ("copy t to 'out.csv'", CommandType::DQL), + ("copy t from 'in.csv'", CommandType::DML), + ] { + assert_eq!( + command_type(&database.prepare_all(sql)?.pop().unwrap())?, + expected + ); + } let err = command_type(&database.prepare_all("start transaction")?.pop().unwrap()).unwrap_err(); @@ -3507,7 +3543,7 @@ mod tests { ); assert_unsupported(expect_err(database.ddl("select 1")), "`Database::ddl`"); - let transaction = database.new_transaction()?; + let transaction = database.new_transaction(false)?; assert_unsupported( expect_err(transaction.prepare_sql("analyze table t", params)), "DDL and ANALYZE", diff --git a/src/binder/select.rs b/src/binder/select.rs index 8aa225d8..f7327bc5 100644 --- a/src/binder/select.rs +++ b/src/binder/select.rs @@ -44,28 +44,28 @@ use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::operator::union::UnionOperator; use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena, ScalarQueryRef}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::tuple::Schema; use crate::types::{ColumnId, LogicalType}; -struct RightSidePositionGlobalizer<'a> { - right_schema: &'a Schema, - left_len: usize, +struct InputPositionBinder<'a> { + schema: &'a Schema, + offset: usize, } -impl ExprVisitorMut for RightSidePositionGlobalizer<'_> { +impl ExprVisitorMut for InputPositionBinder<'_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, position: &mut usize, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { - if self - .right_schema + if let Some(index) = self + .schema .iter() - .any(|right| arena.same_column(*right, *column)) + .position(|input| arena.same_column(*input, *column)) { - *position += self.left_len; + *position = self.offset + index; } Ok(()) } @@ -77,25 +77,6 @@ struct AppendedRightOutput { output_position: usize, } -struct MarkerPositionGlobalizer<'a> { - output_column: &'a ColumnRef, - left_len: usize, -} - -impl ExprVisitorMut for MarkerPositionGlobalizer<'_> { - fn visit_column_ref( - &mut self, - column: &mut ColumnRef, - position: &mut usize, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { - if arena.same_column(*column, *self.output_column) { - *position = self.left_len; - } - Ok(()) - } -} - struct ProjectionOutputBinder<'a> { project_exprs: &'a [ExprRef], } @@ -105,7 +86,11 @@ impl<'a> ProjectionOutputBinder<'a> { Self { project_exprs } } - fn output_ref(&mut self, expr: ExprRef, arena: &mut dyn MetaArena) -> Option { + fn output_ref( + &mut self, + expr: ExprRef, + arena: &mut A, + ) -> Option { self.project_exprs .iter() .position(|candidate| { @@ -121,12 +106,8 @@ impl<'a> ProjectionOutputBinder<'a> { } } -impl ExprVisitorMut for ProjectionOutputBinder<'_> { - fn visit( - &mut self, - expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { +impl ExprVisitorMut for ProjectionOutputBinder<'_> { + fn visit(&mut self, expr: &mut ExprRef, arena: &mut A) -> Result<(), DatabaseError> { if let Some(output_ref) = self.output_ref(*expr, arena) { *expr = arena.alloc_expression(output_ref); return Ok(()); @@ -137,7 +118,7 @@ impl ExprVisitorMut for ProjectionOutputBinder<'_> { pub(crate) struct BindPlanStart<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) binder: &'s mut Binder<'a, 'b, T, A>, @@ -146,7 +127,7 @@ where pub struct BindPlanFrom<'s, 'a, 'b, 'arena, T, A, M = ()> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) binder: &'s mut Binder<'a, 'b, T, A>, @@ -157,7 +138,7 @@ where pub struct BindPlanSelectList<'s, 'a, 'b, 'arena, T, A, M = ()> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) binder: &'s mut Binder<'a, 'b, T, A>, @@ -169,7 +150,7 @@ where pub(crate) struct BindPlanFiltered<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(super) binder: &'s mut Binder<'a, 'b, T, A>, @@ -180,7 +161,7 @@ where pub(crate) struct BindPlanAggregated<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'s mut Binder<'a, 'b, T, A>, @@ -193,7 +174,7 @@ where pub(crate) struct BindPlanHaving<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'s mut Binder<'a, 'b, T, A>, @@ -205,7 +186,7 @@ where pub(crate) struct BindPlanWindowed<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'s mut Binder<'a, 'b, T, A>, @@ -217,7 +198,7 @@ where pub(crate) struct BindPlanDistinct<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'s mut Binder<'a, 'b, T, A>, @@ -229,7 +210,7 @@ where pub(crate) struct BindPlanSorted<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'s mut Binder<'a, 'b, T, A>, @@ -240,7 +221,7 @@ where pub(crate) struct BindPlanProjected<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { plan: LogicalPlan, @@ -256,6 +237,11 @@ pub(crate) struct TableAliasInput { pub(crate) columns: Vec, } +pub(crate) enum AliasPurpose { + Source, + Output, +} + pub(crate) enum JoinConstraintInput { On(ExprRef), Using(Vec), @@ -265,7 +251,7 @@ pub(crate) enum JoinConstraintInput { impl<'s, 'a: 'b, 'b, 'arena, T, A, M> BindPlanFrom<'s, 'a, 'b, 'arena, T, A, M> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { #[cfg(feature = "orm")] @@ -294,7 +280,7 @@ where join_type: JoinType, constraint: JoinConstraintInput, ) -> Result { - self.binder.extend(right_context); + self.binder.extend(right_context, self.arena)?; self.plan = self .binder .bind_join_plans(self.plan, right_plan, join_type, constraint, self.arena)?; @@ -317,7 +303,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A, M> BindPlanSelectList<'s, 'a, 'b, 'arena, T, A, M> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { #[cfg(feature = "orm")] @@ -466,7 +452,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanStart<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { #[allow(clippy::wrong_self_convention)] @@ -485,7 +471,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A, M> BindPlanSelectList<'s, 'a, 'b, 'arena, T, A, M> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn filter_expr( @@ -509,7 +495,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanFiltered<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn aggregate( @@ -622,7 +608,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanAggregated<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn having( @@ -644,7 +630,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanHaving<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn window( @@ -702,7 +688,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanWindowed<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn distinct( @@ -754,7 +740,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanDistinct<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn order_by( @@ -775,7 +761,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanSorted<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn project( @@ -796,7 +782,7 @@ where impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanProjected<'s, 'a, 'b, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub(crate) fn insert_into( @@ -824,7 +810,7 @@ impl BindPlanComplete { } } -impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b, T, A> { +impl<'a: 'b, 'b, T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b, T, A> { pub(crate) fn bind_scalar_queries( &mut self, mut plan: LogicalPlan, @@ -908,18 +894,12 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b pub(crate) fn resolve_source_columns_in_scope<'context>( context: &'context BinderContext<'a, T>, table_name: &str, + arena: &PlanArena, ) -> Result<(&'context Source<'a>, usize), DatabaseError> { - let mut position_offset = 0; - - for bound_source in &context.bind_table { - if bound_source.matches_name(table_name) { - return Ok((&bound_source.source, position_offset)); - } - - position_offset += bound_source.source.schema_len(); - } - - Err(DatabaseError::invalid_table(table_name)) + context + .name_binding(table_name, None, arena) + .map(|binding| (&context.bind_table[binding.source].source, binding.position)) + .ok_or_else(|| DatabaseError::invalid_table(table_name)) } fn localize_join_condition_from_join_scope( @@ -950,12 +930,12 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b appended_outputs: &'a [AppendedRightOutput], } - impl ExprVisitorMut for AppendedRightOutputBinder<'_> { + impl ExprVisitorMut for AppendedRightOutputBinder<'_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, position: &mut usize, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { if let Some(output) = self.appended_outputs.iter().find(|output| { *position == output.child_position && arena.same_column(*column, output.column) @@ -1006,14 +986,14 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b if !left_cast.is_empty() { left_plan = LogicalPlan::new( - Operator::Project(ProjectOperator { exprs: left_cast }), + Operator::Project(ProjectOperator::new(left_cast)), Childrens::Only(Box::new(left_plan)), ); } if !right_cast.is_empty() { right_plan = LogicalPlan::new( - Operator::Project(ProjectOperator { exprs: right_cast }), + Operator::Project(ProjectOperator::new(right_cast)), Childrens::Only(Box::new(right_plan)), ); } @@ -1135,6 +1115,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b alias_column: &[String], table_alias: TableName, table_name: TableName, + purpose: AliasPurpose, arena: &mut PlanArena, ) -> Result { let input_schema = plan.output_schema(arena); @@ -1142,6 +1123,14 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b if !alias_column.is_empty() && alias_column.len() != input_schema_len { return Err(DatabaseError::MisMatch("alias", "columns")); } + let reused_source = match purpose { + AliasPurpose::Source => self + .context + .name_binding(&table_alias, None, arena) + .filter(|binding| binding.source < self.context.source_scope_start) + .map(|binding| binding.source), + AliasPurpose::Output => None, + }; let mut alias_exprs = Vec::with_capacity(input_schema_len); for (position, column) in input_schema.iter().copied().enumerate() { @@ -1150,20 +1139,22 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b } else { alias_column[position].clone() }; - let (mut alias_column, column_id, is_temp) = { + let alias_column = if let Some(source) = reused_source { + self.context.bind_table[source].source.schema()[position] + } else { let source_column = arena.column(column); - ( - source_column.clone(), + let mut alias_column = source_column.clone(); + alias_column.set_name(alias.clone()); + alias_column.set_ref_table( + table_alias.clone(), source_column.id().unwrap_or_default(), matches!( &source_column.summary().relation, ColumnRelation::Table { is_temp: true, .. } ), - ) + ); + arena.alloc_column(alias_column) }; - alias_column.set_name(alias.clone()); - alias_column.set_ref_table(table_alias.clone(), column_id, is_temp); - let alias_column = arena.alloc_column(alias_column); let expr = arena.alloc_expression(ScalarExpression::column_expr(column, position)); let alias_expr = @@ -1223,7 +1214,8 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b arena: &mut PlanArena, ) -> Result { let table_alias = alias.as_ref().map(|alias| alias.name.clone()); - + self.context + .ensure_unique_source_name(table_alias.as_ref().unwrap_or(&table_name))?; if let Some(plan_ref) = self.context.cte(&table_name).map(|cte| cte.plan_ref) { let mut plan = arena.plan(plan_ref).clone(); OperatorExprVisitorMut::new(&mut ExprCloner, arena).visit_plan(&mut plan)?; @@ -1233,6 +1225,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b &alias.columns, alias.name.clone(), table_name.clone(), + AliasPurpose::Source, arena, )?; } @@ -1242,6 +1235,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b table_alias, join_type, Source::Schema(output_schema), + arena, ); return Ok(plan); } @@ -1249,9 +1243,9 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b let with_pk = self.is_scan_with_pk(&table_name); let source = self .context - .source_and_bind(table_name.clone(), table_alias.as_ref(), join_type, false)? + .source(&table_name)? .ok_or(DatabaseError::SourceNotFound)?; - let mut plan = match source { + let mut plan = match &source { Source::Table(table) => { TableScanOperator::build(table_name.clone(), table, with_pk, arena)? } @@ -1267,22 +1261,21 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b } }; - if let Some(alias) = alias { + let source = if let Some(alias) = alias { plan = self.bind_alias( plan, &alias.columns, - alias.name.clone(), + alias.name, table_name.clone(), + AliasPurpose::Source, arena, )?; - let output_schema = plan.output_schema(arena).clone(); - self.context.add_bound_source( - table_name, - Some(alias.name), - join_type, - Source::Schema(output_schema), - ); - } + Source::Schema(plan.output_schema(arena).clone()) + } else { + source + }; + self.context + .add_bound_source(table_name, table_alias, join_type, source, arena); Ok(plan) } @@ -1294,6 +1287,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b arena: &mut PlanArena, ) -> Result { if let Some(alias) = alias { + self.context.ensure_unique_source_name(&alias.name)?; let source_name = arena.temp_table(); plan = self.bind_alias( @@ -1301,6 +1295,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b &alias.columns, alias.name.clone(), source_name.clone(), + AliasPurpose::Source, arena, )?; let output_schema = plan.output_schema(arena).clone(); @@ -1309,6 +1304,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b Some(alias.name), joint_type, Source::Schema(output_schema), + arena, ); } else { let passthrough_source = { @@ -1323,6 +1319,9 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b None } }; + if let Some(source_name) = &passthrough_source { + self.context.ensure_unique_source_name(source_name)?; + } let needs_virtual_source = passthrough_source.is_none(); let source_name = passthrough_source.unwrap_or_else(|| arena.temp_table()); @@ -1335,6 +1334,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b None, joint_type, Source::Schema(output_schema), + arena, ); } @@ -1356,17 +1356,30 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b let mut table_alias = None; let table_name: TableName = function.summary().name.clone(); + self.context.ensure_unique_source_name( + alias + .as_ref() + .map(|alias| &alias.name) + .unwrap_or(&table_name), + )?; let mut plan = FunctionScanOperator::build(function); if let Some(alias) = alias { table_alias = Some(alias.name.clone()); - plan = self.bind_alias(plan, &alias.columns, alias.name, table_name.clone(), arena)?; + plan = self.bind_alias( + plan, + &alias.columns, + alias.name, + table_name.clone(), + AliasPurpose::Source, + arena, + )?; } let source = Source::Schema(plan.output_schema(arena).clone()); self.context - .add_bound_source(table_name, table_alias, joint_type, source); + .add_bound_source(table_name, table_alias, joint_type, source, arena); Ok(plan) } @@ -1385,7 +1398,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b is_qualified_wildcard: bool, ) -> Result<(), DatabaseError> { let (source, position_offset) = - Self::resolve_source_columns_in_scope(context, table_name.as_ref())?; + Self::resolve_source_columns_in_scope(context, table_name.as_ref(), arena)?; let fn_not_on_using = |column: &ColumnRef, arena: &PlanArena<'_>| { let column_catalog = arena.column(*column); @@ -1481,8 +1494,9 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b )); } - if self.context.sub_queries.contains_key(&QueryBindStep::Where) { - children = self.bind_mark_queries(children, std::iter::once(&mut predicate), arena)?; + let pending = self.context.sub_queries.len(); + children = self.bind_mark_queries(children, std::iter::once(&mut predicate), arena)?; + if self.context.sub_queries.len() < pending { let passthrough_exprs = children .output_schema(arena) .iter() @@ -1506,9 +1520,9 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b exprs: impl Iterator, arena: &mut PlanArena, ) -> Result { - let mut output_positions = Vec::new(); - if let Some(sub_queries) = self.context.sub_queries_at_now() { - for sub_query in sub_queries { + for expr in exprs { + while let Some(index) = self.next_mark_query(*expr, arena)? { + let sub_query = self.context.sub_queries.remove(index); match sub_query { SubQueryType::ExistsSubQuery { plan, @@ -1516,7 +1530,6 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b output_column, } => { let left_schema = children.output_schema(arena).clone(); - output_positions.push((output_column, left_schema.len())); let (plan, predicates) = Self::prepare_mark_apply( left_schema.as_ref(), plan, @@ -1537,23 +1550,28 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b plan, correlated, output_column, - predicate: mut quantified_predicate, + predicate: quantified_predicate, .. } => { - if correlated { - quantified_predicate = Self::rewrite_correlated_quantified_predicate( - quantified_predicate, - arena, - ); + if self.context.is_step(&QueryBindStep::Where) { + if let ScalarExpression::Binary { left_expr, .. } = + &quantified_predicate + { + if left_expr.has_agg_call(arena)? { + return Err(DatabaseError::AggMiss( + "aggregate functions are not allowed in WHERE".into(), + )); + } + } } + let predicate = arena.alloc_expression(quantified_predicate); let left_schema = children.output_schema(arena).clone(); - output_positions.push((output_column, left_schema.len())); let (plan, predicates) = Self::prepare_mark_apply( left_schema.as_ref(), plan, correlated, true, - vec![arena.alloc_expression(quantified_predicate)], + vec![predicate], arena, )?; children = MarkApplyOperator::build_quantified( @@ -1566,17 +1584,43 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b } } } + InputPositionBinder { + schema: children.output_schema(arena), + offset: 0, + } + .visit(expr, arena)?; } - for expr in exprs { - for (output_column, left_len) in &output_positions { - MarkerPositionGlobalizer { - output_column, - left_len: *left_len, + Ok(children) + } + + fn next_mark_query( + &self, + expr: ExprRef, + arena: &PlanArena, + ) -> Result, DatabaseError> { + let mut next = None; + expr.any_referenced_column(arena, |arena, column| { + next = self.context.sub_queries.iter().position(|query| { + let output = match query { + SubQueryType::ExistsSubQuery { output_column, .. } + | SubQueryType::QuantifiedSubQuery { output_column, .. } => *output_column, + }; + arena.same_column(output, *column) + }); + next.is_some() + })?; + if let Some(index) = next { + if let SubQueryType::QuantifiedSubQuery { + predicate: ScalarExpression::Binary { left_expr, .. }, + .. + } = &self.context.sub_queries[index] + { + if let Some(dependency) = self.next_mark_query(*left_expr, arena)? { + return Ok(Some(dependency)); } - .visit(expr, arena)?; } } - Ok(children) + Ok(next) } fn ensure_mark_apply_right_outputs( @@ -1662,9 +1706,14 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b } let right_schema = plan.output_schema(arena); for expr in &mut apply_predicates { - RightSidePositionGlobalizer { - right_schema, - left_len, + InputPositionBinder { + schema: left_schema, + offset: 0, + } + .visit(expr, arena)?; + InputPositionBinder { + schema: right_schema, + offset: left_len, } .visit(expr, arena)?; } @@ -1672,36 +1721,6 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b Ok((plan, apply_predicates)) } - fn rewrite_correlated_quantified_predicate( - predicate: ScalarExpression, - arena: &PlanArena<'_>, - ) -> ScalarExpression { - let strip_projection_alias = |expr| match arena.expression(expr) { - ScalarExpression::Alias { - expr, - alias: AliasType::Expr(_), - } => *expr, - _ => expr, - }; - - match predicate { - ScalarExpression::Binary { - op, - left_expr, - right_expr, - ty, - .. - } => ScalarExpression::Binary { - op, - left_expr: strip_projection_alias(left_expr), - right_expr: strip_projection_alias(right_expr), - evaluator: None, - ty, - }, - predicate => predicate, - } - } - fn plan_has_correlated_refs( plan: &LogicalPlan, left_schema: &Schema, @@ -1915,7 +1934,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b select_list: Vec, ) -> LogicalPlan { LogicalPlan::new( - Operator::Project(ProjectOperator { exprs: select_list }), + Operator::Project(ProjectOperator::new(select_list)), Childrens::Only(Box::new(children)), ) } @@ -2154,7 +2173,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(usize, LogicalType)]>> Binder<'a, 'b // GRCOV_EXCL_START #[cfg(all(test, not(target_arch = "wasm32")))] mod tests { - use super::{ProjectionOutputBinder, RightSidePositionGlobalizer}; + use super::{InputPositionBinder, ProjectionOutputBinder}; use crate::binder::test::build_t1_table; use crate::catalog::{ColumnCatalog, ColumnDesc}; use crate::errors::DatabaseError; @@ -2208,8 +2227,7 @@ mod tests { } #[test] - fn test_right_side_position_globalizer_only_shifts_right_columns() -> Result<(), DatabaseError> - { + fn test_input_position_binder_only_binds_available_columns() -> Result<(), DatabaseError> { let table_arena = TableArenaCell::default(); let mut arena = PlanArena::new(&table_arena); let left_column = arena.alloc_column(ColumnCatalog::new( @@ -2223,7 +2241,7 @@ mod tests { ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), )); let right_schema = vec![right_column]; - let left_expr = arena.alloc_expression(ScalarExpression::column_expr(left_column, 0)); + let left_expr = arena.alloc_expression(ScalarExpression::column_expr(left_column, 7)); let right_expr = arena.alloc_expression(ScalarExpression::column_expr(right_column, 0)); let mut expr = arena.alloc_expression(ScalarExpression::Binary { op: crate::expression::BinaryOperator::Eq, @@ -2233,9 +2251,23 @@ mod tests { ty: LogicalType::Boolean, }); - RightSidePositionGlobalizer { - right_schema: &right_schema, - left_len: 2, + InputPositionBinder { + schema: &right_schema, + offset: 2, + } + .visit(&mut expr, &mut arena)?; + InputPositionBinder { + schema: &right_schema, + offset: 2, + } + .visit(&mut expr, &mut arena)?; + assert!(matches!( + arena.expression(left_expr), + ScalarExpression::ColumnRef { position: 7, .. } + )); + InputPositionBinder { + schema: &vec![left_column], + offset: 0, } .visit(&mut expr, &mut arena)?; @@ -2364,6 +2396,54 @@ mod tests { Ok(()) } + #[test] + fn test_select_mark_alias_is_produced_before_where() -> Result<(), DatabaseError> { + let tables = build_t1_table()?; + let mut arena = PlanArena::new(&tables.table_arena); + let mut plan = tables.plan_with_arena( + "select c1,c2 in (select c3 from t2) as matched from t1 where matched", + &mut arena, + )?; + let marker = *find_mark_apply(&plan).unwrap().output_column(); + let Operator::Project(project) = &plan.operator else { + unreachable!() + }; + let selected = project.exprs[1].unpack_alias(&arena); + assert!( + matches!(arena.expression(selected), ScalarExpression::ColumnRef { column, position: 2 } + if arena.same_column(*column, marker)) + ); + let mut input = &mut plan; + let mut marks = 0; + let mut filters = 0; + loop { + match (&input.operator, input.childrens.as_mut()) { + (Operator::Filter(filter), Childrens::Only(child)) => { + let schema = child.output_schema(&mut arena); + assert!(arena.same_column(schema[2], marker)); + let referenced = filter.predicate.unpack_alias(&arena); + assert!( + matches!(arena.expression(referenced), ScalarExpression::ColumnRef { column, position: 2 } + if arena.same_column(*column, marker)) + ); + filters += 1; + input = child; + } + (Operator::MarkApply(op), Childrens::Twins { left, .. }) => { + assert!(arena.same_column(*op.output_column(), marker)); + marks += 1; + input = left; + } + (_, Childrens::Only(child)) => input = child, + (_, Childrens::None) => break, + _ => unreachable!(), + } + } + assert_eq!(marks, 1); + assert_eq!(filters, 1); + Ok(()) + } + #[test] fn test_any_subquery_in_where_binds_as_mark_apply() -> Result<(), DatabaseError> { let table_states = build_t1_table()?; diff --git a/src/binder/show_table.rs b/src/binder/show_table.rs index c6a0891c..f230dfa9 100644 --- a/src/binder/show_table.rs +++ b/src/binder/show_table.rs @@ -16,10 +16,10 @@ use crate::binder::Binder; use crate::errors::DatabaseError; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_show_tables(&mut self) -> Result { Ok(LogicalPlan::new(Operator::ShowTable, Childrens::None)) } diff --git a/src/binder/show_view.rs b/src/binder/show_view.rs index dc90320e..7bf26dba 100644 --- a/src/binder/show_view.rs +++ b/src/binder/show_view.rs @@ -16,10 +16,10 @@ use crate::binder::Binder; use crate::errors::DatabaseError; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_show_views(&mut self) -> Result { Ok(LogicalPlan::new(Operator::ShowView, Childrens::None)) } diff --git a/src/binder/truncate.rs b/src/binder/truncate.rs index f3e1611d..8ae366d0 100644 --- a/src/binder/truncate.rs +++ b/src/binder/truncate.rs @@ -18,10 +18,10 @@ use crate::errors::DatabaseError; use crate::planner::operator::truncate::TruncateOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_truncate( &mut self, table_name: TableName, diff --git a/src/binder/update.rs b/src/binder/update.rs index bab3520b..1e9b8b4a 100644 --- a/src/binder/update.rs +++ b/src/binder/update.rs @@ -18,10 +18,10 @@ use crate::errors::DatabaseError; use crate::planner::operator::update::UpdateOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, ExprRef, LogicalPlan}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_update( &mut self, table_name: TableName, diff --git a/src/binder/window.rs b/src/binder/window.rs index 0cb195b8..176fa1c5 100644 --- a/src/binder/window.rs +++ b/src/binder/window.rs @@ -24,19 +24,15 @@ use crate::planner::operator::window::WindowOperator; use crate::planner::operator::Operator; use crate::planner::MetaArena; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::LogicalType; struct WindowCollector { windows: Vec<(WindowCall, ColumnRef)>, } -impl ExprVisitorMut for WindowCollector { - fn visit( - &mut self, - expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { +impl ExprVisitorMut for WindowCollector { + fn visit(&mut self, expr: &mut ExprRef, arena: &mut A) -> Result<(), DatabaseError> { let ScalarExpression::WindowCall(window) = arena.expression(*expr) else { return walk_mut_expr(self, expr, arena); }; @@ -71,12 +67,12 @@ struct WindowOutputBinder<'a> { base_position: usize, } -impl ExprVisitorMut for WindowOutputBinder<'_> { +impl ExprVisitorMut for WindowOutputBinder<'_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, position: &mut usize, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result<(), DatabaseError> { if let Some(output_position) = self .groups @@ -97,7 +93,7 @@ struct WindowGroup { output_columns: Vec, } -impl> Binder<'_, '_, T, A> { +impl> Binder<'_, '_, T, A> { pub(crate) fn bind_window_function( &mut self, kind: WindowFunctionKind, diff --git a/src/db.rs b/src/db.rs index f1db4e06..87d9d0cc 100644 --- a/src/db.rs +++ b/src/db.rs @@ -35,6 +35,7 @@ use crate::function::lower::Lower; use crate::function::numbers::Numbers; use crate::function::octet_length::OctetLength; use crate::function::upper::Upper; +use crate::optimizer::core::hll; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::heuristic::batch::HepBatchStrategy; use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline; @@ -53,12 +54,13 @@ use crate::storage::memory::MemoryStorage; use crate::storage::rocksdb::{OptimisticRocksStorage, RocksStorage, StorageConfig}; use crate::storage::table_codec::TableCodec; use crate::storage::{ - CheckpointableStorage, StatisticsMetaCache, Storage, TableCache, Transaction, + CheckpointableStorage, StatisticsMetaCache, Storage, StorageTransaction, TableCache, TransactionIsolationLevel, ViewCache, }; -use crate::types::tuple::{Schema, SchemaView, Tuple}; +use crate::types::tuple::{Schema, SchemaView, TupleImpl}; use crate::types::value::DataValue; use crate::types::LogicalType; +use ordered_float::OrderedFloat; #[cfg(any(feature = "parser", feature = "orm"))] pub use prepared::PreparedPlan; use std::collections::{HashMap, HashSet}; @@ -68,6 +70,8 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +/// Output schema, execution-time arena, and executor of a built statement. +pub(crate) type Execution<'a, A, T> = (Schema, ExecMetaArena, Executor<'a, T>); pub(crate) type ScalaFunctions = HashMap>; pub(crate) type TableFunctions = HashMap; @@ -80,12 +84,15 @@ pub enum CatalogKind { pub(crate) trait BindSource<'a> { type Iter: ResultIter; - type Transaction: Transaction; + type Transaction: StorageTransaction; type Storage: Storage; + /// `read_only` lets a fresh transaction use the storage's read path; executors + /// that already hold a transaction ignore it. fn execute>, A: MetaArena + 'a, F>( self, + read_only: bool, build: F, ) -> Result where @@ -112,6 +119,8 @@ pub struct DataBaseBuilder { ))] path: PathBuf, histogram_buckets: Option, + hll_precision: Option, + statistics_relative_error: Option, parser_recursion_limit: Option, transaction_isolation: Option, #[cfg(all(not(target_arch = "wasm32"), feature = "rocksdb"))] @@ -144,6 +153,8 @@ impl DataBaseBuilder { ))] path, histogram_buckets: None, + hll_precision: None, + statistics_relative_error: None, parser_recursion_limit: None, transaction_isolation: None, #[cfg(all(not(target_arch = "wasm32"), feature = "rocksdb"))] @@ -159,6 +170,26 @@ impl DataBaseBuilder { self } + /// Sets the HyperLogLog precision `p` used by `ANALYZE` to estimate column NDV. + /// + /// Each sketch has `2^p` registers (one byte each) with a standard error of + /// about `1.04 / sqrt(2^p)`. Valid range is 4..=26; unset uses 14 + /// (16 KiB, ~0.81% error). + pub fn hll_precision(mut self, precision: u8) -> Self { + self.hll_precision = Some(precision); + self + } + + /// Sets the relative error `Īµ` of the quantile (KLL) and frequency + /// (Count-Min) sketches built by `ANALYZE`. + /// + /// Sketch size grows roughly as `1 / Īµ` (Count-Min width is + /// `next_power_of_two(2 / Īµ)`). Must be in `(0, 1)`; unset uses 0.001. + pub fn statistics_relative_error(mut self, relative_error: f64) -> Self { + self.statistics_relative_error = Some(relative_error); + self + } + /// Sets the SQL parser recursion limit. Unset uses the parser's default limit. /// Higher limits may require a larger thread stack for deeply nested SQL. #[cfg(feature = "parser")] @@ -250,6 +281,8 @@ impl DataBaseBuilder { Self::_build::( storage, self.histogram_buckets, + self.hll_precision, + self.statistics_relative_error, self.parser_recursion_limit, self.transaction_isolation, ) @@ -263,6 +296,8 @@ impl DataBaseBuilder { Self::_build::( storage, self.histogram_buckets, + self.hll_precision, + self.statistics_relative_error, self.parser_recursion_limit, self.transaction_isolation, ) @@ -276,6 +311,8 @@ impl DataBaseBuilder { Self::_build::( storage, self.histogram_buckets, + self.hll_precision, + self.statistics_relative_error, self.parser_recursion_limit, self.transaction_isolation, ) @@ -290,6 +327,8 @@ impl DataBaseBuilder { Self::_build::( storage, self.histogram_buckets, + self.hll_precision, + self.statistics_relative_error, self.parser_recursion_limit, self.transaction_isolation, ) @@ -303,6 +342,8 @@ impl DataBaseBuilder { Self::_build::( storage, self.histogram_buckets, + self.hll_precision, + self.statistics_relative_error, self.parser_recursion_limit, self.transaction_isolation, ) @@ -317,6 +358,8 @@ impl DataBaseBuilder { Self::_build::( storage, self.histogram_buckets, + self.hll_precision, + self.statistics_relative_error, self.parser_recursion_limit, self.transaction_isolation, ) @@ -325,6 +368,8 @@ impl DataBaseBuilder { fn _build( storage: T, histogram_buckets: Option, + hll_precision: Option, + statistics_relative_error: Option, parser_recursion_limit: Option, transaction_isolation: Option, ) -> Result, DatabaseError> { @@ -333,6 +378,23 @@ impl DataBaseBuilder { "histogram buckets must be >= 1".to_string(), )); } + // Reject out-of-range values when the database opens, not at the first ANALYZE. + if let Some(precision) = hll_precision { + if !(hll::MIN_PRECISION..=hll::MAX_PRECISION).contains(&precision) { + return Err(DatabaseError::InvalidValue(format!( + "hll precision must be between {} and {}, got {precision}", + hll::MIN_PRECISION, + hll::MAX_PRECISION + ))); + } + } + if let Some(relative_error) = statistics_relative_error { + if !(relative_error > 0.0 && relative_error < 1.0) { + return Err(DatabaseError::InvalidValue(format!( + "statistics relative error must be in (0, 1), got {relative_error}" + ))); + } + } if matches!(parser_recursion_limit, Some(0)) { return Err(DatabaseError::InvalidValue( "parser recursion limit must be >= 1".to_string(), @@ -341,7 +403,7 @@ impl DataBaseBuilder { let transaction_isolation = transaction_isolation.unwrap_or_else(|| storage.default_transaction_isolation()); storage.validate_transaction_isolation(transaction_isolation)?; - let meta_cache = HashMap::default(); + let meta_cache = StatisticsMetaCache::default(); let table_cache = HashMap::default(); let view_cache = HashMap::default(); let table_arena = TableArenaCell::default(); @@ -355,6 +417,8 @@ impl DataBaseBuilder { table_arena, optimizer_pipeline: optimizer_pipeline(), histogram_buckets, + hll_precision, + statistics_relative_error: statistics_relative_error.map(OrderedFloat), parser_recursion_limit, _p: Default::default(), }; @@ -384,62 +448,78 @@ impl DataBaseBuilder { fn optimizer_pipeline() -> HepOptimizerPipeline { HepOptimizerPipeline::builder() .before_batch( - "Column Pruning".to_string(), + "Column Pruning", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::ColumnPruning], + &[NormalizationRuleImpl::ColumnPruning], ) .before_batch( - "Simplify Filter".to_string(), + "Simplify Filter", HepBatchStrategy::fix_point_topdown(10), - vec![NormalizationRuleImpl::SimplifyFilter], + &[NormalizationRuleImpl::SimplifyFilter], ) .before_batch( - "Constant Calculation".to_string(), + "Constant Calculation", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::ConstantCalculation], + &[NormalizationRuleImpl::ConstantCalculation], ) .before_batch( - "Predicate Pushdown".to_string(), + "Predicate Pushdown", HepBatchStrategy::fix_point_topdown(30), - vec![ + &[ + NormalizationRuleImpl::PropagateFilterConstants, + NormalizationRuleImpl::PropagateJoinConstants, NormalizationRuleImpl::PushPredicateThroughJoin, NormalizationRuleImpl::PushJoinPredicateIntoScan, ], ) .before_batch( - "Limit Pushdown".to_string(), + "Limit Pushdown", HepBatchStrategy::fix_point_topdown(10), - vec![ + &[ NormalizationRuleImpl::LimitProjectTranspose, NormalizationRuleImpl::PushLimitThroughJoin, ], ) .before_batch( - "TopK".to_string(), - HepBatchStrategy::once_topdown(), - vec![ + "TopK", + HepBatchStrategy::fix_point_topdown(1), + &[ NormalizationRuleImpl::MinMaxToTopK, NormalizationRuleImpl::TopK, ], ) .before_batch( - "Combine Operators".to_string(), + "Collapse Project", + HepBatchStrategy::fix_point_bottomup(10), + &[NormalizationRuleImpl::EliminateProject], + ) + .before_batch( + "Combine Operators", HepBatchStrategy::fix_point_topdown(10), - vec![ - NormalizationRuleImpl::CollapseProject, + &[ NormalizationRuleImpl::CollapseGroupByAgg, NormalizationRuleImpl::CombineFilter, ], ) .before_batch( - "Predicate Into Scan".to_string(), + "Predicate Into Scan", HepBatchStrategy::fix_point_topdown(10), - vec![NormalizationRuleImpl::PushPredicateIntoScan], + &[NormalizationRuleImpl::PushPredicateIntoScan], + ) + .before_batch( + "Guard Mark Apply", + HepBatchStrategy::fix_point_topdown(1), + &[NormalizationRuleImpl::GuardMarkApply], ) .after_batch( - "Parameterize Mark Apply".to_string(), + "Greedy Join Reorder", HepBatchStrategy::once_topdown(), - vec![ + &[NormalizationRuleImpl::GreedyJoinReorder], + ) + .after_batch( + "Parameterize Mark Apply", + HepBatchStrategy::fix_point_topdown(1), + &[ NormalizationRuleImpl::ParameterizeMarkApply, NormalizationRuleImpl::ParameterizeInnerJoin, ], @@ -447,21 +527,31 @@ fn optimizer_pipeline() -> HepOptimizerPipeline { // Only discard predicates after parameterization has chosen the final // lookup. A Probe must not inherit residuals from a replaced static range. .after_batch( - "Eliminate Index Filter".to_string(), - HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::EliminateIndexFilter], + "Eliminate Index Filter", + HepBatchStrategy::fix_point_topdown(1), + &[NormalizationRuleImpl::EliminateIndexFilter], ) .after_batch( - "Expression Remapper".to_string(), + "Materialize Mark Apply", + HepBatchStrategy::fix_point_bottomup(1), + &[NormalizationRuleImpl::MaterializeMarkApply], + ) + .after_batch( + "Collapse Project", + HepBatchStrategy::fix_point_bottomup(10), + &[NormalizationRuleImpl::EliminateProject], + ) + .after_batch( + "Expression Remapper", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::EvaluatorBind], + &[NormalizationRuleImpl::EvaluatorBind], ) .after_batch( - "Limit Into Scan".to_string(), + "Limit Into Scan", HepBatchStrategy::fix_point_topdown(10), - vec![NormalizationRuleImpl::PushLimitIntoTableScan], + &[NormalizationRuleImpl::PushLimitIntoTableScan], ) - .implementations(vec![ + .implementations(&[ // DQL ImplementationRuleImpl::SimpleAggregate, ImplementationRuleImpl::GroupByAggregate, @@ -470,6 +560,7 @@ fn optimizer_pipeline() -> HepOptimizerPipeline { ImplementationRuleImpl::HashJoin, ImplementationRuleImpl::Limit, ImplementationRuleImpl::Projection, + ImplementationRuleImpl::MarkApply, ImplementationRuleImpl::ScalarSubquery, ImplementationRuleImpl::SeqScan, ImplementationRuleImpl::IndexScan, @@ -498,7 +589,7 @@ fn optimizer_pipeline() -> HepOptimizerPipeline { .build() } -fn statement_stamp( +fn statement_stamp( transaction: &mut T, plan: &LogicalPlan, ) -> Result { @@ -528,6 +619,8 @@ pub(crate) struct State { table_arena: TableArenaCell, optimizer_pipeline: HepOptimizerPipeline, histogram_buckets: Option, + hll_precision: Option, + statistics_relative_error: Option>, #[cfg_attr(not(feature = "parser"), allow(dead_code))] parser_recursion_limit: Option, _p: PhantomData, @@ -610,7 +703,7 @@ impl State { Ok(()) } - pub(crate) fn build_plan<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>, F>( + pub(crate) fn build_plan<'a, T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, F>( &'a self, params: A, transaction: &T, @@ -645,6 +738,12 @@ impl State { if op.histogram_buckets.is_none() { op.histogram_buckets = self.histogram_buckets; } + if op.hll_precision.is_none() { + op.hll_precision = self.hll_precision; + } + if op.statistics_relative_error.is_none() { + op.statistics_relative_error = self.statistics_relative_error; + } } plan.populate_output_schema_recursive(&mut arena); @@ -656,14 +755,7 @@ impl State { transaction: &'a mut S::TransactionType<'txn>, plan: PlanInput<'a>, mut plan_arena: A, - ) -> Result< - ( - Schema, - ExecMetaArena, - Executor<'a, S::TransactionType<'txn>>, - ), - DatabaseError, - > + ) -> Result>, DatabaseError> where S: 'txn, { @@ -691,14 +783,7 @@ impl State { transaction: &'a mut S::TransactionType<'txn>, params: A, build: F, - ) -> Result< - ( - Schema, - ExecMetaArena>, - Executor<'a, S::TransactionType<'txn>>, - ), - DatabaseError, - > + ) -> Result, S::TransactionType<'txn>>, DatabaseError> where S: 'txn, A: AsRef<[(usize, LogicalType)]>, @@ -717,6 +802,8 @@ impl State { table_arena, optimizer_pipeline, histogram_buckets, + hll_precision, + statistics_relative_error, .. } = self; let mut plan_arena = PlanArena::new(table_arena); @@ -741,6 +828,12 @@ impl State { if op.histogram_buckets.is_none() { op.histogram_buckets = *histogram_buckets; } + if op.hll_precision.is_none() { + op.hll_precision = *hll_precision; + } + if op.statistics_relative_error.is_none() { + op.statistics_relative_error = *statistics_relative_error; + } } plan.populate_output_schema_recursive(&mut plan_arena); @@ -787,17 +880,13 @@ impl DDLApply { let table = table.transplant_to_table_arena(plan_arena)?; state.table_cache.insert(name.clone(), table); if clear_statistics { - state - .meta_cache - .retain(|(cached_table_name, _), _| cached_table_name != &name); + state.meta_cache.remove_table(&name); } catalog_changed = true; } DDLApply::DropTable { name } => { state.table_cache.remove(&name); - state - .meta_cache - .retain(|(cached_table_name, _), _| cached_table_name != &name); + state.meta_cache.remove_table(&name); catalog_changed = true; } DDLApply::UpsertView { view } => { @@ -810,18 +899,31 @@ impl DDLApply { state.view_cache.remove(&name); catalog_changed = true; } - DDLApply::UpsertStatisticsMeta { + DDLApply::UpsertIndexStatistics { table_name, index_id, meta, } => { - state.meta_cache.insert((table_name, index_id), meta); + state + .meta_cache + .indexes + .insert((table_name, index_id), meta); } - DDLApply::RemoveStatisticsMeta { + DDLApply::RemoveIndexStatistics { table_name, index_id, } => { - state.meta_cache.remove(&(table_name, index_id)); + state.meta_cache.indexes.remove(&(table_name, index_id)); + } + DDLApply::UpsertColumnStatistics { + table_name, + column_id, + meta, + } => { + state + .meta_cache + .columns + .insert((table_name, column_id), meta); } } Ok(catalog_changed) @@ -864,7 +966,7 @@ impl Database { { let transaction = Box::into_raw(Box::new( self.storage - .transaction_with_isolation(self.transaction_isolation)?, + .transaction_with_isolation(self.transaction_isolation, false)?, )); let state = std::ptr::from_mut(&mut self.state); let (schema, plan_arena, executor) = @@ -923,7 +1025,10 @@ impl Database { Ok(()) } CatalogKind::Table(name) => { - let transaction = self.storage.transaction()?; + let transaction = self.storage.transaction_with_isolation( + self.storage.default_transaction_isolation(), + true, + )?; let mut table_codec = TableCodec::default(); let table = transaction .load_table( @@ -932,23 +1037,56 @@ impl Database { name.clone(), )? .ok_or(DatabaseError::TableNotFound)?; + let mut indexes = Vec::new(); + let mut columns = Vec::new(); for index in table.indexes() { let index_id = self.state.table_arena.borrow().index(*index).id; - if let Some(meta) = transaction.statistics_meta( + if let Some(meta) = transaction.index_statistics( &mut table_codec, name.as_ref(), index_id, self.state.table_arena.borrow_mut(), )? { - self.state.meta_cache.insert((name.clone(), index_id), meta); + indexes.push((index_id, meta)); + } + } + for column in table.columns() { + let column_id = { + let table_arena = self.state.table_arena.borrow(); + let column = table_arena.column(*column); + column + .id() + .ok_or_else(|| DatabaseError::invalid_column(column.name()))? + }; + if let Some(meta) = transaction.column_statistics( + &mut table_codec, + name.as_ref(), + column_id, + self.state.table_arena.borrow_mut(), + )? { + columns.push((column_id, meta)); } } + self.state.meta_cache.remove_table(&name); + self.state.meta_cache.indexes.extend( + indexes + .into_iter() + .map(|(id, meta)| ((name.clone(), id), meta)), + ); + self.state.meta_cache.columns.extend( + columns + .into_iter() + .map(|(id, meta)| ((name.clone(), id), meta)), + ); self.state.table_cache.insert(name, table); self.state.recycle_table_arena()?; Ok(()) } CatalogKind::View(name) => { - let transaction = self.storage.transaction()?; + let transaction = self.storage.transaction_with_isolation( + self.storage.default_transaction_isolation(), + true, + )?; let mut table_codec = TableCodec::default(); let view = transaction .load_view( @@ -970,15 +1108,18 @@ impl Database { /// Opens a new explicit transaction. /// /// Statements executed through the returned transaction share the same - /// transactional context until [`DBTransaction::commit`] is called. - pub fn new_transaction(&self) -> Result, DatabaseError> { + /// transactional context until [`Transaction::commit`] is called. + /// A `read_only` transaction rejects statements that write; on LMDB it runs + /// alongside the single writer instead of waiting for it. + pub fn new_transaction(&self, read_only: bool) -> Result, DatabaseError> { let transaction = self .storage - .transaction_with_isolation(self.transaction_isolation)?; + .transaction_with_isolation(self.transaction_isolation, read_only)?; - Ok(DBTransaction { + Ok(Transaction { inner: transaction, state: &self.state, + read_only, }) } @@ -1007,6 +1148,7 @@ impl<'a, S: Storage> BindSource<'a> for &'a Database { fn execute>, A: MetaArena + 'a, F>( self, + read_only: bool, build: F, ) -> Result where @@ -1014,7 +1156,7 @@ impl<'a, S: Storage> BindSource<'a> for &'a Database { { let transaction = Box::into_raw(Box::new( self.storage - .transaction_with_isolation(self.transaction_isolation)?, + .transaction_with_isolation(self.transaction_isolation, read_only)?, )); let result = (|| { let transaction = unsafe { &mut *transaction }; @@ -1049,7 +1191,7 @@ impl<'a, S: Storage> BindSource<'a> for &'a Database { { let mut transaction = self .storage - .transaction_with_isolation(self.transaction_isolation)?; + .transaction_with_isolation(self.transaction_isolation, true)?; transaction.begin_statement_scope()?; let (plan, mut arena) = self.state.build_plan(params, &transaction, build)?; Ok(plan.explain(&mut arena, 0)) @@ -1077,7 +1219,7 @@ pub trait ResultIter { /// Advances to the next row and runs `f` while the executor row slot is borrowed. fn next_tuple( &mut self, - f: impl FnOnce(&SchemaView<'_, '_>, &mut Tuple) -> R, + f: impl FnOnce(&SchemaView<'_, '_>, &mut TupleImpl) -> R, ) -> Result, DatabaseError>; /// Finishes consuming the iterator and flushes any remaining work. @@ -1171,7 +1313,7 @@ impl DatabaseIter<'_, S> { #[inline] pub fn next_tuple( &mut self, - f: impl FnOnce(&SchemaView<'_, '_>, &mut Tuple) -> R, + f: impl FnOnce(&SchemaView<'_, '_>, &mut TupleImpl) -> R, ) -> Result, DatabaseError> { unsafe { (*self.inner).next_tuple(f) } } @@ -1195,7 +1337,7 @@ impl ResultIter for DatabaseIter<'_, S> { fn next_tuple( &mut self, - f: impl FnOnce(&SchemaView<'_, '_>, &mut Tuple) -> R, + f: impl FnOnce(&SchemaView<'_, '_>, &mut TupleImpl) -> R, ) -> Result, DatabaseError> { DatabaseIter::next_tuple(self, f) } @@ -1206,12 +1348,13 @@ impl ResultIter for DatabaseIter<'_, S> { } /// Explicit transaction handle created by [`Database::new_transaction`]. -pub struct DBTransaction<'a, S: Storage + 'a> { +pub struct Transaction<'a, S: Storage + 'a> { pub(crate) inner: S::TransactionType<'a>, pub(crate) state: &'a State, + read_only: bool, } -impl<'txn, S: Storage> DBTransaction<'txn, S> { +impl<'txn, S: Storage> Transaction<'txn, S> { /// Commits the current transaction. pub fn commit(self) -> Result<(), DatabaseError> { self.inner.commit()?; @@ -1220,7 +1363,7 @@ impl<'txn, S: Storage> DBTransaction<'txn, S> { } } -impl<'a, 'txn, S: Storage> BindSource<'a> for &'a mut DBTransaction<'txn, S> { +impl<'a, 'txn, S: Storage> BindSource<'a> for &'a mut Transaction<'txn, S> { type Iter = TransactionIter<'a, S::TransactionType<'txn>>; type Transaction = S::TransactionType<'txn>; @@ -1228,6 +1371,7 @@ impl<'a, 'txn, S: Storage> BindSource<'a> for &'a mut DBTransaction<'txn, S> { fn execute>, A: MetaArena + 'a, F>( self, + _read_only: bool, build: F, ) -> Result where @@ -1235,10 +1379,16 @@ impl<'a, 'txn, S: Storage> BindSource<'a> for &'a mut DBTransaction<'txn, S> { { self.inner.begin_statement_scope()?; let (plan, arena) = build(self.state, &self.inner)?; + let plan = plan.into(); + if self.read_only && !plan.root().operator.command_type().is_read_only() { + return Err(DatabaseError::UnsupportedStmt( + "cannot write in a read-only transaction".into(), + )); + } let transaction = std::ptr::from_mut(&mut self.inner); let (schema, arena, executor) = self.state - .execute(unsafe { &mut *transaction }, plan.into(), arena)?; + .execute(unsafe { &mut *transaction }, plan, arena)?; Ok(TransactionIter::new( schema, arena.boxed(), @@ -1262,13 +1412,13 @@ impl<'a, 'txn, S: Storage> BindSource<'a> for &'a mut DBTransaction<'txn, S> { } } -struct TransactionGuard<'a, T: Transaction + 'a> { +struct TransactionGuard<'a, T: StorageTransaction + 'a> { executor: Option>, transaction: *mut T, statement_scope_active: bool, } -impl TransactionGuard<'_, T> { +impl TransactionGuard<'_, T> { #[inline] fn finish_statement_scope( &mut self, @@ -1286,14 +1436,18 @@ impl TransactionGuard<'_, T> { } } -impl Drop for TransactionGuard<'_, T> { +impl Drop for TransactionGuard<'_, T> { fn drop(&mut self) { let _ = self.finish_statement_scope(&mut Vec::new()); } } /// Raw result iterator returned by transaction execution APIs. -pub struct TransactionIter<'a, T: Transaction + 'a, A: MetaArena + 'a = Box> { +pub struct TransactionIter< + 'a, + T: StorageTransaction + 'a, + A: MetaArena + 'a = Box, +> { // Drop execution before the metadata arena it borrows from. guard: TransactionGuard<'a, T>, plan_arena: ExecMetaArena, @@ -1301,7 +1455,7 @@ pub struct TransactionIter<'a, T: Transaction + 'a, A: MetaArena + 'a = Box, } -impl<'a, T: Transaction + 'a, A: MetaArena + 'a> TransactionIter<'a, T, A> { +impl<'a, T: StorageTransaction + 'a, A: MetaArena + 'a> TransactionIter<'a, T, A> { pub(crate) fn new( schema: Schema, plan_arena: ExecMetaArena, @@ -1334,7 +1488,7 @@ impl<'a, T: Transaction + 'a, A: MetaArena + 'a> TransactionIter<'a, T, A> { #[inline] pub fn next_tuple( &mut self, - f: impl FnOnce(&SchemaView<'_, '_>, &mut Tuple) -> R, + f: impl FnOnce(&SchemaView<'_, '_>, &mut TupleImpl) -> R, ) -> Result, DatabaseError> { let Some(executor) = self.guard.executor.as_mut() else { return Ok(None); @@ -1364,7 +1518,7 @@ impl<'a, T: Transaction + 'a, A: MetaArena + 'a> TransactionIter<'a, T, A> { } } -impl<'a, T: Transaction + 'a> TransactionIter<'a, T, PlanArena<'a>> { +impl<'a, T: StorageTransaction + 'a> TransactionIter<'a, T, PlanArena<'a>> { fn done_with_ddl_apply(mut self) -> Result<(PlanArena<'a>, Vec), DatabaseError> { while self.next_tuple(|_, _| ())?.is_some() {} let Self { @@ -1378,14 +1532,14 @@ impl<'a, T: Transaction + 'a> TransactionIter<'a, T, PlanArena<'a>> { } } -impl ResultIter for TransactionIter<'_, T, A> { +impl ResultIter for TransactionIter<'_, T, A> { fn schema(&self, f: impl FnOnce(&SchemaView<'_, '_>) -> R) -> R { TransactionIter::schema(self, f) } fn next_tuple( &mut self, - f: impl FnOnce(&SchemaView<'_, '_>, &mut Tuple) -> R, + f: impl FnOnce(&SchemaView<'_, '_>, &mut TupleImpl) -> R, ) -> Result, DatabaseError> { TransactionIter::next_tuple(self, f) } @@ -1400,16 +1554,22 @@ impl ResultIter for TransactionIter<'_, T, A> { pub(crate) mod test { use crate::binder::{Binder, BinderContext}; use crate::catalog::{ColumnCatalog, ColumnDesc}; - use crate::db::{CatalogKind, DataBaseBuilder, DatabaseError, ResultIter}; - use crate::expression::ScalarExpression; - use crate::planner::operator::join::JoinCondition; + use crate::db::{optimizer_pipeline, CatalogKind, DataBaseBuilder, DatabaseError, ResultIter}; + use crate::execution::dql::projection::Projection; + use crate::execution::{empty_context, execute_input}; + use crate::expression::{AliasType, BinaryOperator, ScalarExpression}; + use crate::optimizer::core::statistics_meta::StatisticMetaLoader; + use crate::planner::operator::join::{JoinCondition, JoinOperator, JoinType}; + use crate::planner::operator::project::ProjectOperator; + use crate::planner::operator::values::ValuesOperator; use crate::planner::operator::Operator; - use crate::planner::PlanArena; + use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena, TableArenaCell}; use crate::storage::memory::MemoryStorage; use crate::storage::{ - table_codec::TableCodec, Storage, TableCache, Transaction, TransactionIsolationLevel, + table_codec::TableCodec, StatisticsMetaCache, Storage, StorageTransaction, TableCache, + TransactionIsolationLevel, ViewCache, }; - use crate::types::tuple::Tuple; + use crate::types::tuple::{Tuple, TupleImpl}; use crate::types::value::DataValue; use crate::types::LogicalType; use chrono::{Datelike, Local}; @@ -1426,6 +1586,166 @@ pub(crate) mod test { use std::time::Duration; use tempfile::TempDir; + fn literal_input( + name: &str, + values: &[i32], + position: usize, + arena: &mut PlanArena, + ) -> Result<(LogicalPlan, ExprRef), DatabaseError> { + let mut column = ColumnCatalog::new( + name.into(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + ); + column.set_ref_table(name.into(), 0, false); + let column = arena.alloc_column(column); + let key = arena.alloc_expression(ScalarExpression::column_expr(column, position)); + let rows = values + .iter() + .map(|value| { + arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(*value))) + }) + .collect(); + Ok(( + LogicalPlan::new( + Operator::Values(ValuesOperator { + rows, + row_count: values.len(), + schema_ref: vec![column], + }), + Childrens::None, + ), + key, + )) + } + + #[test] + fn late_project_collapse_preserves_oriented_columns_and_aliases() -> Result<(), DatabaseError> { + for alias in [false, true] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (left, large) = literal_input("large", &[1, 2, 3, 4], 0, &mut arena)?; + let (right, small) = literal_input("small", &[99], 1, &mut arena)?; + let join = + JoinOperator::build(left, right, JoinCondition::None, JoinType::Cross, false); + let output = if alias { + arena.alloc_expression(ScalarExpression::Alias { + expr: large, + alias: AliasType::Name("wanted".into()), + }) + } else { + large + }; + // Repeating the selected column must not lose an output slot. + let mut plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![output, output, small])), + Childrens::Only(Box::new(join)), + ); + let schema = plan.output_schema(&mut arena).clone(); + let cache = StatisticsMetaCache::default(); + let loader = StatisticMetaLoader::new(&cache); + let mut plan = optimizer_pipeline() + .instantiate(plan) + .find_best(Some(&loader), &mut arena)?; + plan.populate_output_schema_recursive(&mut arena); + assert_eq!(plan.output_schema(&mut arena), &schema); + let Childrens::Only(join) = plan.childrens.as_ref() else { + unreachable!() + }; + assert!( + matches!(join.operator, Operator::Join(_)), + "restore Project should be collapsed" + ); + let Childrens::Twins { left, right } = join.childrens.as_ref() else { + unreachable!() + }; + assert_eq!(left.estimated_rows(), Some(1)); + assert_eq!(right.estimated_rows(), Some(4)); + let Operator::Project(project) = &plan.operator else { + unreachable!() + }; + let Childrens::Only(input) = plan.childrens.as_ref() else { + unreachable!() + }; + let storage = MemoryStorage::new(); + let transaction = storage.transaction()?; + let table_cache = TableCache::default(); + let view_cache = ViewCache::default(); + let context = empty_context(&table_cache, &view_cache, &cache); + let mut executor = + execute_input::<_, Projection>((project, input), context, arena, &transaction); + for value in [1, 2, 3, 4] { + assert_eq!( + executor.next_tuple()?.cloned().unwrap(), + vec![ + DataValue::Int32(value), + DataValue::Int32(value), + DataValue::Int32(99) + ] + ); + } + assert!(executor.next_tuple()?.is_none()); + } + Ok(()) + } + + #[test] + fn late_project_collapse_inlines_computed_output() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (left, large) = literal_input("large", &[1, 2, 3, 4], 0, &mut arena)?; + let (right, small) = literal_input("small", &[99], 1, &mut arena)?; + let join = JoinOperator::build(left, right, JoinCondition::None, JoinType::Cross, false); + let ten = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(10))); + let computed = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Plus, + left_expr: large, + right_expr: ten, + evaluator: None, + ty: LogicalType::Integer, + }); + let mut plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![computed, small])), + Childrens::Only(Box::new(join)), + ); + let schema = plan.output_schema(&mut arena).clone(); + let cache = StatisticsMetaCache::default(); + let loader = StatisticMetaLoader::new(&cache); + let mut plan = optimizer_pipeline() + .instantiate(plan) + .find_best(Some(&loader), &mut arena)?; + plan.populate_output_schema_recursive(&mut arena); + assert_eq!(plan.output_schema(&mut arena), &schema); + let Childrens::Only(child) = plan.childrens.as_ref() else { + unreachable!() + }; + assert!( + matches!(child.operator, Operator::Join(_)), + "computed expressions should be substituted into the output Project" + ); + let Operator::Project(project) = &plan.operator else { + unreachable!() + }; + let Childrens::Only(input) = plan.childrens.as_ref() else { + unreachable!() + }; + let storage = MemoryStorage::new(); + let transaction = storage.transaction()?; + let table_cache = TableCache::default(); + let view_cache = ViewCache::default(); + let context = empty_context(&table_cache, &view_cache, &cache); + let mut executor = + execute_input::<_, Projection>((project, input), context, arena, &transaction); + for value in [11, 12, 13, 14] { + assert_eq!( + executor.next_tuple()?.cloned().unwrap(), + vec![DataValue::Int32(value), DataValue::Int32(99)] + ); + } + assert!(executor.next_tuple()?.is_none()); + Ok(()) + } + fn assert_send_sync() {} #[test] @@ -1455,13 +1775,13 @@ pub(crate) mod test { cross join scalar_b b cross join scalar_nested c"; let mut iter = db.run(sql)?; assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![ + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![ DataValue::Int32(11), DataValue::Int32(22), DataValue::Int32(33), DataValue::Int32(44), - ]) + ] ); assert!(iter.next_tuple(|_, _| ())?.is_none()); iter.done()?; @@ -1470,14 +1790,14 @@ pub(crate) mod test { db.ddl("create view scalar_new as select (select 55) as v")?; let mut iter = db.run("select a.v, b.v from scalar_a a cross join scalar_new b")?; assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int32(11), DataValue::Int32(55)]) + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int32(11), DataValue::Int32(55)] ); iter.done()?; Ok(()) } - pub(crate) fn build_table( + pub(crate) fn build_table( table_cache: &mut TableCache, transaction: &mut T, plan_arena: &mut PlanArena, @@ -1514,28 +1834,21 @@ pub(crate) mod test { Ok(()) } - fn tuple_owned(tuple: &Tuple) -> Tuple { - Tuple { - pk: tuple.pk.clone(), - values: tuple.values.clone(), - } + fn next_tuple_owned(iter: &mut I) -> Result, DatabaseError> { + iter.next_tuple(|_, tuple| tuple.clone()) } - fn next_tuple_owned(iter: &mut I) -> Result, DatabaseError> { - iter.next_tuple(|_, tuple| tuple_owned(tuple)) - } - - fn next_values(iter: &mut I) -> Result>, DatabaseError> { - iter.next_tuple(|_, tuple| tuple.values.clone()) + fn next_values(iter: &mut I) -> Result, DatabaseError> { + iter.next_tuple(|_, tuple| tuple.clone()) } fn read_single_i32(mut iter: I) -> Result where I: ResultIter, { - let value = match iter.next_tuple(|_, tuple| match tuple.values.as_slice() { - [DataValue::Int32(value)] => *value, - other => panic!("expected a single Int32 column, got {other:?}"), + let value = match iter.next_tuple(|_, tuple| match (tuple.len(), tuple.get(0)) { + (1, Some(DataValue::Int32(value))) => *value, + _ => panic!("expected a single Int32 column, got {tuple:?}"), })? { Some(value) => value, None => panic!("expected one result row"), @@ -1550,9 +1863,9 @@ pub(crate) mod test { sql: &str, ) -> Result { let mut iter = database.run(sql)?; - let value = match iter.next_tuple(|_, tuple| match tuple.values.as_slice() { - [DataValue::Int32(value)] => *value, - other => panic!("expected a single Int32 column, got {other:?}"), + let value = match iter.next_tuple(|_, tuple| match (tuple.len(), tuple.get(0)) { + (1, Some(DataValue::Int32(value))) => *value, + _ => panic!("expected a single Int32 column, got {tuple:?}"), })? { Some(value) => value, None => panic!("expected one result row for query: {sql}"), @@ -1774,7 +2087,7 @@ pub(crate) mod test { unreachable!("expected join filter"); }; let mut referenced_columns = Vec::new(); - filter.all_referenced_columns(&source_plan_arena, |_, column| { + filter.all_referenced_columns(&source_plan_arena, |_, column, _| { referenced_columns.push(*column); true })?; @@ -1867,7 +2180,10 @@ pub(crate) mod test { }; assert_eq!(*left_position, 0); assert_eq!(*right_position, 0); - let (_, right_child) = join_plan.childrens.pop_twins(); + let (_, mut right_child) = join_plan.childrens.pop_twins(); + if matches!(right_child.operator, Operator::Project(_)) { + right_child = right_child.childrens.pop_only(); + } let Operator::Filter(filter_op) = right_child.operator else { unreachable!("expected pushed-down filter on right child"); }; @@ -1918,7 +2234,7 @@ pub(crate) mod test { let mut iter = kite_sql.execute(&statement, &[(1, DataValue::Int32(0))])?; let row = next_tuple_owned(&mut iter)?.unwrap(); - let plan = row.values[0].utf8().unwrap(); + let plan = row[0].utf8().unwrap(); assert_eq!( plan, "Projection [t1.a, t1.b] [Project => (Sort Option: Follow)] Filter (t1.b > 0), Is Having: false [Filter => (Sort Option: Follow)] TableScan t1 -> [t1.a, t1.b] [SeqScan => (Sort Option: None)]" @@ -1945,7 +2261,7 @@ pub(crate) mod test { ], )?; let row = next_tuple_owned(&mut iter)?.unwrap(); - let plan = row.values[0].utf8().unwrap(); + let plan = row[0].utf8().unwrap(); assert_eq!( plan, "Projection [(t1.a + $1), Max((t1.b + $2))] [Project => (Sort Option: Follow)] Aggregate [Max((t1.b + 0))] -> Group By [(t1.a + 0)] [HashAggregate => (Sort Option: None)] Filter (t1.b > 1), Is Having: false [Filter => (Sort Option: Follow)] TableScan t1 -> [t1.a, t1.b] [SeqScan => (Sort Option: None)]" @@ -1964,7 +2280,7 @@ pub(crate) mod test { ], )?; let row = next_tuple_owned(&mut iter)?.unwrap(); - let plan = row.values[0].utf8().unwrap(); + let plan = row[0].utf8().unwrap(); assert_eq!( plan, "Projection [l.a, l.b, t1.a, t1.b, 9] [Project => (Sort Option: Follow)] LeftOuter Join Where (l.a > 0) [NestLoopJoin => (Sort Option: None)] Projection [(t1.a) as (l.a), (t1.b) as (l.b)] [Project => (Sort Option: Follow)] Filter (t1.b > 0), Is Having: false [Filter => (Sort Option: Follow)] TableScan t1 -> [t1.a, t1.b] [SeqScan => (Sort Option: None)] Projection [t1.a, t1.b] [Project => (Sort Option: Follow)] TableScan t1 -> [t1.a, t1.b] [IndexScan By pk_index => ($3, +inf) => (Sort Option: OrderBy: (t1.a Asc Nulls Last) ignore_prefix_len: 0)]" @@ -2085,8 +2401,8 @@ pub(crate) mod test { kite_sql.ddl("create table t1 (a int primary key, b int)")?; - let mut tx_1 = kite_sql.new_transaction()?; - let mut tx_2 = kite_sql.new_transaction()?; + let mut tx_1 = kite_sql.new_transaction(false)?; + let mut tx_2 = kite_sql.new_transaction(false)?; tx_1.run("insert into t1 values(0, 0)")?.done()?; tx_1.run("insert into t1 values(1, 1)")?.done()?; @@ -2116,7 +2432,7 @@ pub(crate) mod test { tx_1.commit()?; tx_2.commit()?; - let mut tx_3 = kite_sql.new_transaction()?; + let mut tx_3 = kite_sql.new_transaction(false)?; let res = tx_3.run("create table t2 (a int primary key, b int)"); assert!(res.is_err()); @@ -2130,7 +2446,7 @@ pub(crate) mod test { kite_sql.ddl("create table t_multi_tx (a int primary key, b int)")?; - let mut tx = kite_sql.new_transaction()?; + let mut tx = kite_sql.new_transaction(false)?; let mut iter = tx.run( "insert into t_multi_tx values(0, 0); insert into t_multi_tx values(1, 1); select * from t_multi_tx order by a", )?; @@ -2163,7 +2479,7 @@ pub(crate) mod test { kite_sql.ddl("create table t_iter_drop (a int primary key, b int)")?; - let mut tx = kite_sql.new_transaction()?; + let mut tx = kite_sql.new_transaction(false)?; tx.run("insert into t_iter_drop values (0, 0), (1, 1)")? .done()?; @@ -2212,8 +2528,8 @@ pub(crate) mod test { kite_sql.ddl("create table t_rc (a int primary key, b int)")?; kite_sql.run("insert into t_rc values (1, 10)")?.done()?; - let mut reader = kite_sql.new_transaction()?; - let mut writer = kite_sql.new_transaction()?; + let mut reader = kite_sql.new_transaction(false)?; + let mut writer = kite_sql.new_transaction(false)?; assert_eq!( read_single_i32(reader.run("select b from t_rc where a = 1")?)?, @@ -2255,8 +2571,8 @@ pub(crate) mod test { kite_sql.ddl("create table t_rr (a int primary key, b int)")?; kite_sql.run("insert into t_rr values (1, 10)")?.done()?; - let mut reader = kite_sql.new_transaction()?; - let mut writer = kite_sql.new_transaction()?; + let mut reader = kite_sql.new_transaction(false)?; + let mut writer = kite_sql.new_transaction(false)?; assert_eq!( read_single_i32(reader.run("select b from t_rr where a = 1")?)?, @@ -2282,8 +2598,8 @@ pub(crate) mod test { kite_sql.ddl("create table t1 (a int primary key, b int)")?; - let mut tx_1 = kite_sql.new_transaction()?; - let mut tx_2 = kite_sql.new_transaction()?; + let mut tx_1 = kite_sql.new_transaction(false)?; + let mut tx_2 = kite_sql.new_transaction(false)?; tx_1.run("insert into t1 values(0, 0)")?.done()?; tx_1.run("insert into t1 values(1, 1)")?.done()?; @@ -2318,7 +2634,7 @@ pub(crate) mod test { assert!(tx_2.commit().is_err()); - let mut tx_3 = kite_sql.new_transaction()?; + let mut tx_3 = kite_sql.new_transaction(false)?; let res = tx_3.run("create table t2 (a int primary key, b int)"); assert!(res.is_err()); @@ -2338,6 +2654,34 @@ pub(crate) mod test { )); } + #[test] + fn test_invalid_statistics_relative_error() { + for relative_error in [0.0, 1.0, -0.5, f64::NAN] { + let result = DataBaseBuilder::path(".") + .statistics_relative_error(relative_error) + .build_in_memory(); + assert!(matches!( + result, + Err(DatabaseError::InvalidValue(message)) + if message == format!("statistics relative error must be in (0, 1), got {relative_error}") + )); + } + } + + #[test] + fn test_invalid_hll_precision() { + for precision in [3, 27] { + let result = DataBaseBuilder::path(".") + .hll_precision(precision) + .build_in_memory(); + assert!(matches!( + result, + Err(DatabaseError::InvalidValue(message)) + if message == format!("hll precision must be between 4 and 26, got {precision}") + )); + } + } + #[test] fn parser_recursion_limit_applies_to_sql_frontends() -> Result<(), DatabaseError> { let expression = format!("{}7{}", "(".repeat(8), ")".repeat(8)); @@ -2355,30 +2699,30 @@ pub(crate) mod test { assert_eq!(read_single_i32(db.run("select * from deep_view")?)?, 7); db.run_mut(&sql, |iter| { assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int32(7)]) + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int32(7)] ); Ok(()) })?; let prepared = db.prepare_sql(&sql, &[])?; let mut iter = db.execute(&prepared, [])?; assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int32(7)]) + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int32(7)] ); iter.done()?; - let mut tx = db.new_transaction()?; + let mut tx = db.new_transaction(false)?; let mut iter = tx.run(&sql)?; assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int32(7)]) + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int32(7)] ); iter.done()?; let prepared = tx.prepare_sql(&sql, &[])?; let mut iter = tx.execute(&prepared, [])?; assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int32(7)]) + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int32(7)] ); iter.done()?; tx.commit()?; @@ -2481,12 +2825,12 @@ pub(crate) mod test { snapshot.load(CatalogKind::View("checkpoint_view".into()))?; let mut rows = snapshot.run("select id,v from checkpoint_rows order by id")?; assert_eq!( - rows.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int32(1), DataValue::Int32(10)]) + rows.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int32(1), DataValue::Int32(10)] ); assert_eq!( - rows.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int32(2), DataValue::Int32(20)]) + rows.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int32(2), DataValue::Int32(20)] ); assert!(rows.next_tuple(|_, _| ())?.is_none()); rows.done()?; diff --git a/src/db/prepared.rs b/src/db/prepared.rs index 405639bd..1929a12e 100644 --- a/src/db/prepared.rs +++ b/src/db/prepared.rs @@ -1,7 +1,7 @@ use super::*; // GRCOV_EXCL_START #[cfg(test)] -use crate::planner::operator::{PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{PhysicalOption, PlanImpl, RowCountOption, SortOption}; // GRCOV_EXCL_STOP use crate::planner::{ExprRef, ParamArena, PlanInput}; // GRCOV_EXCL_START @@ -16,6 +16,7 @@ pub struct PreparedPlan<'db> { pub(crate) plan: LogicalPlan, pub(crate) arena: PlanArena<'db>, parameter_expressions: Vec, + pub(crate) read_only: bool, } impl<'db> PreparedPlan<'db> { @@ -40,7 +41,7 @@ impl Database { )); } let arena = prepared.bind_parameters(params.as_ref())?; - BindSource::execute(self, |_, _| { + BindSource::execute(self, prepared.read_only, |_, _| { Ok((PlanInput::Borrowed(&prepared.plan), arena)) }) } @@ -56,12 +57,12 @@ impl Database { let statement = self.state.parse_statement(sql)?; let transaction = self .storage - .transaction_with_isolation(self.transaction_isolation)?; + .transaction_with_isolation(self.transaction_isolation, true)?; self.state.prepare_plan(&statement, params, &transaction) } } -impl<'db, S: Storage> DBTransaction<'db, S> { +impl<'db, S: Storage> Transaction<'db, S> { /// Executes a prepared SQL or ORM plan inside this transaction. pub fn execute<'a>( &'a mut self, @@ -74,7 +75,7 @@ impl<'db, S: Storage> DBTransaction<'db, S> { )); } let arena = prepared.bind_parameters(params.as_ref())?; - BindSource::execute(self, |_, _| { + BindSource::execute(self, prepared.read_only, |_, _| { Ok((PlanInput::Borrowed(&prepared.plan), arena)) }) } @@ -114,7 +115,12 @@ impl State { }) } - pub(crate) fn prepare_plan_with<'a, T: Transaction, A: AsRef<[(usize, LogicalType)]>, F>( + pub(crate) fn prepare_plan_with< + 'a, + T: StorageTransaction, + A: AsRef<[(usize, LogicalType)]>, + F, + >( &'a self, params: A, transaction: &T, @@ -128,10 +134,12 @@ impl State { { let (plan, mut arena) = self.build_plan(params, transaction, build)?; let parameter_expressions = arena.parameter_expressions(); + let read_only = plan.operator.command_type().is_read_only(); Ok(PreparedPlan { plan, arena, parameter_expressions, + read_only, }) } } @@ -141,6 +149,7 @@ impl State { mod tests { use super::*; use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::execution::dql::index_scan::IndexScan; use crate::expression::range_detacher::Range; use crate::expression::{BinaryOperator, ScalarExpression}; use crate::planner::operator::filter::FilterOperator; @@ -170,17 +179,17 @@ mod tests { ); }); assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int64(expected)]) + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int64(expected)] ); assert!(iter.next_tuple(|_, _| ())?.is_none()); iter.done()?; } - let mut tx = db.new_transaction()?; + let mut tx = db.new_transaction(false)?; let mut iter = tx.execute(&plan, [(1, DataValue::Int64(7))])?; assert_eq!( - iter.next_tuple(|_, row| row.values.clone())?, - Some(vec![DataValue::Int64(31)]) + iter.next_tuple(|_, row| row.clone())?.unwrap(), + vec![DataValue::Int64(31)] ); iter.done()?; tx.commit()?; @@ -296,6 +305,7 @@ mod tests { scan.physical_option = Some(PhysicalOption::new( PlanImpl::IndexScan(Box::new(index.clone())), sort.clone(), + RowCountOption::Unknown, )); let filter = FilterOperator { predicate, @@ -313,7 +323,7 @@ mod tests { let Childrens::Only(scan) = plan.childrens.as_ref() else { panic!("expected scan") }; - let Operator::TableScan(scan_op) = &scan.operator else { + let Operator::TableScan(_) = &scan.operator else { panic!("expected table scan") }; let Some(PhysicalOption { @@ -329,12 +339,12 @@ mod tests { let mut execution_arena = crate::execution::ExecArena::< ::TransactionType<'_>, >::with_capacity(0); - let executor = crate::execution::dql::index_scan::IndexScan::new( - scan_op, + let ranges = IndexScan::ranges( info, info.lookup.as_ref().expect("lookup"), - ); - let ranges = executor.ranges(&mut execution_arena, &mut arena)?; + &mut execution_arena, + &mut arena, + )?; let expected_range = match &expected_index.lookup { Some(IndexLookup::Static(range)) => range, _ => unreachable!(), @@ -594,7 +604,7 @@ mod tests { .or_else(|| plan.childrens.iter().find_map(find_scan)) } let scan = find_scan(&plan.plan).expect("table scan"); - let Operator::TableScan(scan_op) = &scan.operator else { + let Operator::TableScan(_) = &scan.operator else { unreachable!() }; let Some(PhysicalOption { @@ -607,12 +617,12 @@ mod tests { let mut execution_arena = crate::execution::ExecArena::< ::TransactionType<'_>, >::with_capacity(0); - let ranges = crate::execution::dql::index_scan::IndexScan::new( - scan_op, + let ranges = IndexScan::ranges( info, info.lookup.as_ref().expect("lookup"), - ) - .ranges(&mut execution_arena, &mut bound_arena)?; + &mut execution_arena, + &mut bound_arena, + )?; let mut ranges = ranges; assert_eq!( ranges.next(), diff --git a/src/errors.rs b/src/errors.rs index e1f4d04b..e74e7409 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -80,6 +80,7 @@ pub enum DatabaseError { DuplicateColumn(String), DuplicateSourceHash(String), DuplicateIndex(String), + DuplicateTableName(String), DuplicatePrimaryKey, DuplicateUniqueValue, FunctionNotFound { @@ -96,6 +97,7 @@ pub enum DatabaseError { span: Option, }, InvalidIndex, + InvalidPlan(String), InvalidTable { name: String, span: Option, @@ -178,6 +180,9 @@ impl fmt::Display for DatabaseError { write!(f, "table or view: `{value}` hash already exists") } Self::DuplicateIndex(value) => write!(f, "index: `{value}` already exists"), + Self::DuplicateTableName(value) => { + write!(f, "table name: `{value}` specified more than once") + } Self::DuplicatePrimaryKey => f.write_str("duplicate primary key"), Self::DuplicateUniqueValue => { f.write_str("the column has been declared unique and the value already exists") @@ -200,6 +205,7 @@ impl fmt::Display for DatabaseError { write!(f, "invalid column: `{name}`{}", format_sql_error_loc(span)) } Self::InvalidIndex => f.write_str("invalid index"), + Self::InvalidPlan(value) => write!(f, "invalid plan: {value}"), Self::InvalidTable { name, span } => { write!(f, "invalid table: `{name}`{}", format_sql_error_loc(span)) } @@ -671,6 +677,14 @@ mod tests { DatabaseError::DuplicateIndex("idx".into()), "index: `idx` already exists", ), + ( + DatabaseError::DuplicateTableName("t".into()), + "table name: `t` specified more than once", + ), + ( + DatabaseError::InvalidPlan("join key".into()), + "invalid plan: join key", + ), ( DatabaseError::Incomparable(LogicalType::Integer, LogicalType::Boolean), "can not compare two types: Integer and Boolean", diff --git a/src/execution/ddl/add_column.rs b/src/execution/ddl/add_column.rs index c0b78dcb..535aface 100644 --- a/src/execution/ddl/add_column.rs +++ b/src/execution/ddl/add_column.rs @@ -20,7 +20,7 @@ use crate::execution::{ use crate::iter_ext::Itertools; use crate::planner::operator::alter_table::add_column::AddColumnOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::index::{Index, IndexType}; use crate::types::value::DataValue; @@ -28,7 +28,7 @@ pub struct AddColumn<'a> { op: &'a AddColumnOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for AddColumn<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for AddColumn<'a> { type Input = &'a AddColumnOperator; fn into_executor( @@ -42,7 +42,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for AddColumn<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for AddColumn<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for AddColumn<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -118,18 +118,16 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for AddColumn<'a> { }, |tuple| { if let Some(value) = &default_value { - tuple.values.push(value.clone()); + tuple.push(value.clone()); } else { - tuple.values.push(DataValue::Null); + tuple.push(DataValue::Null); } Ok(()) }, |transaction, table_codec, tuple| { - if let (Some(unique_index_id), Some(value), Some(tuple_id)) = ( - unique_index_id.as_ref(), - default_value.as_ref(), - tuple.pk.as_ref(), - ) { + if let (Some(unique_index_id), Some(value), Some(tuple_id)) = + (unique_index_id.as_ref(), default_value.as_ref(), tuple.pk()) + { let index = Index::new( *unique_index_id, std::slice::from_ref(value), diff --git a/src/execution/ddl/change_column.rs b/src/execution/ddl/change_column.rs index ccddade6..ce12e245 100644 --- a/src/execution/ddl/change_column.rs +++ b/src/execution/ddl/change_column.rs @@ -20,14 +20,14 @@ use crate::execution::{ use crate::iter_ext::Itertools; use crate::planner::operator::alter_table::change_column::{ChangeColumnOperator, NotNullChange}; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; pub struct ChangeColumn<'a> { op: &'a ChangeColumnOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for ChangeColumn<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for ChangeColumn<'a> { type Input = &'a ChangeColumnOperator; fn into_executor( @@ -41,7 +41,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for ChangeColumn<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ChangeColumn<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for ChangeColumn<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -128,9 +128,9 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ChangeColumn<'a> { }) }, |tuple| { - let value = std::mem::replace(&mut tuple.values[column_index], DataValue::Null); - tuple.values[column_index] = value.cast(&target_data_type)?; - if needs_not_null_validation && tuple.values[column_index].is_null() { + let value = std::mem::replace(&mut tuple[column_index], DataValue::Null); + tuple[column_index] = value.cast(&target_data_type)?; + if needs_not_null_validation && tuple[column_index].is_null() { return Err(DatabaseError::not_null_column(target_column_name.clone())); } Ok(()) @@ -154,7 +154,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ChangeColumn<'a> { .map(|column| plan_arena.column(*column).datatype().serializable()) }, |tuple| { - if tuple.values[column_index].is_null() { + if tuple[column_index].is_null() { return Err(DatabaseError::not_null_column(target_column_name.clone())); } Ok(()) diff --git a/src/execution/ddl/create_index.rs b/src/execution/ddl/create_index.rs index 2dd43683..d8197a84 100644 --- a/src/execution/ddl/create_index.rs +++ b/src/execution/ddl/create_index.rs @@ -21,7 +21,7 @@ use crate::expression::ScalarExpression; use crate::planner::operator::create_index::CreateIndexOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::index::Index; use crate::types::tuple::Schema; use crate::types::value::DataValue; @@ -33,7 +33,7 @@ pub struct CreateIndex<'a> { input: ExecId, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CreateIndex<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for CreateIndex<'a> { type Input = (&'a CreateIndexOperator, &'a LogicalPlan); fn into_executor( @@ -53,7 +53,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CreateIndex<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CreateIndex<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for CreateIndex<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -117,15 +117,15 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CreateIndex<'a> { }; while arena.next_tuple(self.input, plan_arena)? { - if arena.result_tuple().pk.is_none() { + if arena.result_tuple().pk().is_none() { continue; } arena.rewrite(&column_exprs, plan_arena, None::<&[DataValue]>)?; { let mut state = arena.local_state(plan_arena); let (tuple, transaction, table_codec) = state.tuple_transaction_codec_mut(); - let tuple_pk = tuple.pk.as_ref().ok_or(DatabaseError::PrimaryKeyNotFound)?; - let index = Index::new(index_id, &tuple.values, *ty); + let tuple_pk = tuple.pk().ok_or(DatabaseError::PrimaryKeyNotFound)?; + let index = Index::new(index_id, tuple, *ty); transaction.add_index(table_codec, table_name.as_ref(), index, tuple_pk)?; } } diff --git a/src/execution/ddl/create_table.rs b/src/execution/ddl/create_table.rs index 9f398f8f..cc90dafe 100644 --- a/src/execution/ddl/create_table.rs +++ b/src/execution/ddl/create_table.rs @@ -18,13 +18,13 @@ use crate::execution::{ }; use crate::planner::operator::create_table::CreateTableOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct CreateTable<'a> { op: &'a CreateTableOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CreateTable<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for CreateTable<'a> { type Input = &'a CreateTableOperator; fn into_executor( @@ -38,7 +38,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CreateTable<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CreateTable<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for CreateTable<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/ddl/create_view.rs b/src/execution/ddl/create_view.rs index f7184891..85d60576 100644 --- a/src/execution/ddl/create_view.rs +++ b/src/execution/ddl/create_view.rs @@ -18,13 +18,13 @@ use crate::execution::{ }; use crate::planner::operator::create_view::CreateViewOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct CreateView<'a> { op: &'a CreateViewOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CreateView<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for CreateView<'a> { type Input = &'a CreateViewOperator; fn into_executor( @@ -38,7 +38,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CreateView<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CreateView<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for CreateView<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/ddl/drop_column.rs b/src/execution/ddl/drop_column.rs index db7c2c7e..501d0620 100644 --- a/src/execution/ddl/drop_column.rs +++ b/src/execution/ddl/drop_column.rs @@ -20,13 +20,13 @@ use crate::execution::{ use crate::iter_ext::Itertools; use crate::planner::operator::alter_table::drop_column::DropColumnOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct DropColumn<'a> { op: &'a DropColumnOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropColumn<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for DropColumn<'a> { type Input = &'a DropColumnOperator; fn into_executor( @@ -40,7 +40,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropColumn<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for DropColumn<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for DropColumn<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -100,7 +100,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for DropColumn<'a> { .map(|(_, column)| plan_arena.column(*column).datatype().serializable()) }, |tuple| { - let _ = tuple.values.remove(column_index); + let _ = tuple.remove(column_index); Ok(()) }, |_, _, _| Ok(()), diff --git a/src/execution/ddl/drop_index.rs b/src/execution/ddl/drop_index.rs index a4c53446..5a8f5e2a 100644 --- a/src/execution/ddl/drop_index.rs +++ b/src/execution/ddl/drop_index.rs @@ -18,13 +18,13 @@ use crate::execution::{ }; use crate::planner::operator::drop_index::DropIndexOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct DropIndex<'a> { op: &'a DropIndexOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropIndex<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for DropIndex<'a> { type Input = &'a DropIndexOperator; fn into_executor( @@ -38,7 +38,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropIndex<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for DropIndex<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for DropIndex<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -62,7 +62,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for DropIndex<'a> { }; if let Some((table, index_id)) = dropped { arena.push_ddl_apply(DDLApply::upsert_table(table, false)); - arena.push_ddl_apply(DDLApply::RemoveStatisticsMeta { + arena.push_ddl_apply(DDLApply::RemoveIndexStatistics { table_name: table_name.clone(), index_id, }); diff --git a/src/execution/ddl/drop_table.rs b/src/execution/ddl/drop_table.rs index 3df6c183..9e6bf44b 100644 --- a/src/execution/ddl/drop_table.rs +++ b/src/execution/ddl/drop_table.rs @@ -18,13 +18,13 @@ use crate::execution::{ }; use crate::planner::operator::drop_table::DropTableOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct DropTable<'a> { op: &'a DropTableOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropTable<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for DropTable<'a> { type Input = &'a DropTableOperator; fn into_executor( @@ -38,7 +38,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropTable<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for DropTable<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for DropTable<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/ddl/drop_view.rs b/src/execution/ddl/drop_view.rs index ce2b1277..4a8718a1 100644 --- a/src/execution/ddl/drop_view.rs +++ b/src/execution/ddl/drop_view.rs @@ -18,13 +18,13 @@ use crate::execution::{ }; use crate::planner::operator::drop_view::DropViewOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct DropView<'a> { op: &'a DropViewOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropView<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for DropView<'a> { type Input = &'a DropViewOperator; fn into_executor( @@ -38,7 +38,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropView<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for DropView<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for DropView<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/ddl/mod.rs b/src/execution/ddl/mod.rs index ca1e68d3..49bb6ce9 100644 --- a/src/execution/ddl/mod.rs +++ b/src/execution/ddl/mod.rs @@ -26,9 +26,9 @@ pub(crate) mod truncate; use crate::catalog::TableName; use crate::errors::DatabaseError; use crate::storage::table_codec::TableCodec; -use crate::storage::{InnerIter, Transaction}; +use crate::storage::{InnerIter, StorageTransaction}; use crate::types::serialize::TupleValueSerializableImpl; -use crate::types::tuple::{Tuple, TupleId}; +use crate::types::tuple::{Tuple, TupleId, TupleImpl}; use crate::types::LogicalType; use std::collections::Bound; @@ -43,11 +43,11 @@ fn read_tuple_batch( old_values_len: usize, old_deserializers: &mut D, start_after: Option<&TupleId>, - batch: &mut Vec, + batch: &mut Vec, batch_size: usize, ) -> Result where - T: Transaction, + T: StorageTransaction, D: FnMut() -> I, I: IntoIterator, { @@ -76,16 +76,16 @@ where let tuple = if len < batch.len() { &mut batch[len] } else { - batch.push(Tuple { - pk: None, - values: Vec::with_capacity(old_values_len), - }); + batch.push(TupleImpl::from(Tuple::new( + None, + Vec::with_capacity(old_values_len), + ))); batch .last_mut() .expect("batch contains the tuple that was just pushed") }; TableCodec::decode_tuple_into( - tuple, + tuple.output_mut(), old_deserializers(), Some(tuple_id), value, @@ -108,10 +108,10 @@ pub(crate) fn visit_table_in_batches( mut visit: F, ) -> Result<(), DatabaseError> where - T: Transaction, + T: StorageTransaction, D: FnMut() -> I, I: IntoIterator, - F: FnMut(&Tuple) -> Result<(), DatabaseError>, + F: FnMut(&TupleImpl) -> Result<(), DatabaseError>, { let mut last_pk = None; let mut batch = Vec::with_capacity(REWRITE_BATCH_SIZE); @@ -131,7 +131,9 @@ where if batch_len == 0 { break; } - last_pk = batch.get(batch_len - 1).and_then(|tuple| tuple.pk.clone()); + last_pk = batch + .get(batch_len - 1) + .and_then(|tuple| tuple.pk().cloned()); for tuple in batch.iter().take(batch_len) { visit(tuple)?; @@ -158,13 +160,13 @@ pub(crate) fn rewrite_table_in_batches( mut after_write: G, ) -> Result<(), DatabaseError> where - T: Transaction, + T: StorageTransaction, D: FnMut() -> I, I: IntoIterator, S: FnMut() -> J, J: IntoIterator, - F: FnMut(&mut Tuple) -> Result<(), DatabaseError>, - G: FnMut(&mut T, &mut TableCodec, &Tuple) -> Result<(), DatabaseError>, + F: FnMut(&mut TupleImpl) -> Result<(), DatabaseError>, + G: FnMut(&mut T, &mut TableCodec, &TupleImpl) -> Result<(), DatabaseError>, { let mut last_pk = None; let mut batch = Vec::with_capacity(REWRITE_BATCH_SIZE); @@ -184,7 +186,9 @@ where if batch_len == 0 { break; } - last_pk = batch.get(batch_len - 1).and_then(|tuple| tuple.pk.clone()); + last_pk = batch + .get(batch_len - 1) + .and_then(|tuple| tuple.pk().cloned()); for tuple in batch.iter_mut().take(batch_len) { rewrite(tuple)?; diff --git a/src/execution/ddl/truncate.rs b/src/execution/ddl/truncate.rs index ffc54ee5..65f2177e 100644 --- a/src/execution/ddl/truncate.rs +++ b/src/execution/ddl/truncate.rs @@ -18,13 +18,13 @@ use crate::execution::{ }; use crate::planner::operator::truncate::TruncateOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct Truncate<'a> { op: &'a TruncateOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Truncate<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for Truncate<'a> { type Input = &'a TruncateOperator; fn into_executor( @@ -38,7 +38,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Truncate<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Truncate<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Truncate<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/ddl_apply.rs b/src/execution/ddl_apply.rs index 7809a3e9..90240945 100644 --- a/src/execution/ddl_apply.rs +++ b/src/execution/ddl_apply.rs @@ -14,8 +14,10 @@ use crate::catalog::view::View; use crate::catalog::{TableCatalog, TableName}; -use crate::optimizer::core::statistics_meta::StatisticsMeta; +use crate::optimizer::core::histogram::HistogramMeta; +use crate::optimizer::core::statistics_meta::IndexStatistics; use crate::types::index::IndexId; +use crate::types::ColumnId; pub(crate) enum DDLApply { UpsertTable { @@ -31,15 +33,20 @@ pub(crate) enum DDLApply { DropView { name: TableName, }, - UpsertStatisticsMeta { + UpsertIndexStatistics { table_name: TableName, index_id: IndexId, - meta: StatisticsMeta, + meta: IndexStatistics, }, - RemoveStatisticsMeta { + RemoveIndexStatistics { table_name: TableName, index_id: IndexId, }, + UpsertColumnStatistics { + table_name: TableName, + column_id: ColumnId, + meta: HistogramMeta, + }, } impl DDLApply { diff --git a/src/execution/dml/analyze.rs b/src/execution/dml/analyze.rs index d3144016..e7a3e839 100644 --- a/src/execution/dml/analyze.rs +++ b/src/execution/dml/analyze.rs @@ -20,26 +20,33 @@ use crate::execution::{ }; use crate::expression::ScalarExpression; use crate::iter_ext::Itertools; -use crate::optimizer::core::histogram::{HistogramBuilder, ANALYZE_STATISTICS_RELATIVE_ERROR}; -use crate::optimizer::core::statistics_meta::StatisticsMeta; +use crate::optimizer::core::histogram::{ + HistogramBuilder, HistogramMetaBuilder, DEFAULT_HLL_PRECISION, + DEFAULT_STATISTICS_RELATIVE_ERROR, +}; +use crate::optimizer::core::statistics_meta::IndexStatistics; use crate::planner::operator::analyze::AnalyzeOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::{table_codec::TableCodec, Transaction}; +use crate::storage::{table_codec::TableCodec, StorageTransaction}; use crate::types::index::IndexId; +use crate::types::tuple::TupleLike; use crate::types::value::{DataValue, Utf8Type}; -use crate::types::CharLengthUnits; +use crate::types::{CharLengthUnits, ColumnId}; +use ordered_float::OrderedFloat; use std::fmt::{self, Formatter}; -const DEFAULT_NUM_OF_BUCKETS: usize = 100; +const DEFAULT_NUM_OF_BUCKETS: usize = 10; pub struct Analyze<'a> { table_name: &'a TableName, input: Option, histogram_buckets: Option, + hll_precision: Option, + statistics_relative_error: Option>, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Analyze<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for Analyze<'a> { type Input = (&'a AnalyzeOperator, &'a LogicalPlan); fn into_executor( @@ -47,6 +54,8 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Analyze<'a> { AnalyzeOperator { table_name, histogram_buckets, + hll_precision, + statistics_relative_error, .. }, input_plan, @@ -67,11 +76,13 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Analyze<'a> { table_name, input, histogram_buckets: *histogram_buckets, + hll_precision: *hll_precision, + statistics_relative_error: *statistics_relative_error, })) } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Analyze<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Analyze<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -82,17 +93,36 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Analyze<'a> { return Ok(()); }; - let mut builders = { + let hll_precision = self.hll_precision.unwrap_or(DEFAULT_HLL_PRECISION); + let relative_error = self + .statistics_relative_error + .map_or(DEFAULT_STATISTICS_RELATIVE_ERROR, |error| error.0); + let (mut columns, mut builders) = { let table = arena .transaction() .table(arena.table_cache(), self.table_name.clone())? .ok_or(DatabaseError::TableNotFound)?; - table + let columns = table + .columns() + .enumerate() + .map(|(position, column)| { + let column = plan_arena.column(*column); + let Some(column_id) = column.id() else { + return Err(DatabaseError::invalid_column(column.name())); + }; + Ok(( + column_id, + position, + HistogramMetaBuilder::new(hll_precision)?, + )) + }) + .collect::, DatabaseError>>()?; + let indexes = table .indexes() .map(|index| { let index = plan_arena.index(*index); let index_id = index.id; - let builder = HistogramBuilder::new(index, ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let builder = HistogramBuilder::new(index, relative_error, hll_precision)?; let exprs = index.column_exprs(table).collect::, _>>()?; Ok(State { index_id, @@ -101,17 +131,21 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Analyze<'a> { histogram_buckets: self.histogram_buckets, }) }) - .collect::, DatabaseError>>()? + .collect::, DatabaseError>>()?; + (columns, indexes) }; while arena.next_tuple(input, plan_arena)? { - let tuple = arena.materialize_tuple(); + let tuple = arena.take_tuple(); + for (_, position, builder) in &mut columns { + builder.append(tuple.value_at(*position)); + } for State { exprs, builder, .. } in builders.iter_mut() { let value = match exprs.as_slice() { [expr] => expr.eval(plan_arena, Some(&tuple))?.into_owned(), _ => { arena.rewrite(exprs, plan_arena, Some(&tuple))?; - DataValue::Tuple(arena.materialize_tuple().values) + DataValue::Tuple(arena.result_mut().take_values()) } }; builder.append(value)?; @@ -119,8 +153,9 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Analyze<'a> { } let mut state = arena.local_state(plan_arena); let (transaction, table_codec, ddl_apply) = state.write_transaction_codec_ddl_apply_mut(); - let values = Self::persist_statistics_meta( + let values = Self::persist_statistics( self.table_name, + columns, builders, ddl_apply, transaction, @@ -128,7 +163,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Analyze<'a> { plan_arena, )?; - arena.produce_tuple(crate::types::tuple::Tuple::new(None, values)); + arena.produce(crate::types::tuple::Tuple::new(None, values).into()); Ok(()) } } @@ -141,8 +176,9 @@ struct State { } impl<'a> Analyze<'a> { - fn persist_statistics_meta( + fn persist_statistics( table_name: &TableName, + columns: Vec<(ColumnId, usize, HistogramMetaBuilder)>, builders: Vec, applies: &mut Vec, transaction: &mut U, @@ -150,6 +186,21 @@ impl<'a> Analyze<'a> { plan_arena: &(dyn MetaArena + '_), ) -> Result, DatabaseError> { let mut values = Vec::with_capacity(builders.len()); + for (column_id, _, builder) in columns { + let meta = builder.build(); + transaction.save_column_statistics( + table_codec, + table_name, + column_id, + &meta, + plan_arena, + )?; + applies.push(DDLApply::UpsertColumnStatistics { + table_name: table_name.clone(), + column_id, + meta, + }); + } for State { index_id, @@ -160,10 +211,10 @@ impl<'a> Analyze<'a> { { let (histogram, sketch, top_n) = builder.build(histogram_buckets.unwrap_or(DEFAULT_NUM_OF_BUCKETS))?; - let meta = StatisticsMeta::new(histogram, sketch, top_n); + let meta = IndexStatistics::new(histogram, sketch, top_n); - transaction.save_statistics_meta(table_codec, table_name, meta.clone(), plan_arena)?; - applies.push(DDLApply::UpsertStatisticsMeta { + transaction.save_index_statistics(table_codec, table_name, meta.clone(), plan_arena)?; + applies.push(DDLApply::UpsertIndexStatistics { table_name: table_name.clone(), index_id, meta, @@ -196,20 +247,95 @@ mod test { use crate::errors::DatabaseError; use crate::execution::dml::analyze::DEFAULT_NUM_OF_BUCKETS; use crate::expression::range_detacher::Range; + use crate::iter_ext::Itertools; use crate::optimizer::core::cm_sketch::COUNT_MIN_SKETCH_STORAGE_PAGE_LEN; + use crate::optimizer::core::histogram::{HistogramMetaBuilder, DEFAULT_HLL_PRECISION}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; - use crate::storage::{table_codec::TableCodec, InnerIter, Storage, Transaction}; + use crate::storage::{table_codec::TableCodec, InnerIter, Storage, StorageTransaction}; use crate::types::value::DataValue; use std::ops::Bound; use tempfile::TempDir; #[test] fn test_analyze() -> Result<(), DatabaseError> { - test_statistics_meta_roundtrip()?; + test_column_and_index_statistics_roundtrip()?; test_meta_loader_uses_cache()?; test_meta_loader_negative_cache()?; test_clean_expired_index()?; + test_hll_precision_is_configurable()?; + test_statistics_relative_error_is_configurable()?; + + Ok(()) + } + fn test_hll_precision_is_configurable() -> Result<(), DatabaseError> { + let rows = 1000; + let ndv_with = |precision: u8| -> Result { + let mut builder = HistogramMetaBuilder::new(precision)?; + for i in 0..rows { + builder.append(&DataValue::Int32(i)); + } + Ok(builder.build().distinct_values_len()) + }; + let (coarse, default) = (ndv_with(4)?, ndv_with(DEFAULT_HLL_PRECISION)?); + // Precision 4 has 16 registers, so its estimate must differ from the default sketch. + assert_ne!(coarse, default); + + let values = (0..rows).map(|i| format!("({i})")).join(", "); + for (builder, expected) in [ + (DataBaseBuilder::path("."), default), + (DataBaseBuilder::path(".").hll_precision(4), coarse), + ] { + let mut kite_sql = builder.build_in_memory()?; + kite_sql.ddl("create table t1 (a int primary key)")?; + kite_sql + .run(format!("insert into t1 values {values}"))? + .done()?; + kite_sql.analyze("t1")?; + + let table_name = "t1".to_string().into(); + let a_id = *kite_sql + .state + .table_cache() + .get(&table_name) + .unwrap() + .get_column_id_by_name("a") + .unwrap(); + let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache()); + let column = loader.load_column(&table_name, a_id).unwrap(); + assert_eq!(column.distinct_values_len(), expected); + let index = loader.load(&table_name, 0)?.unwrap(); + assert_eq!(index.histogram().distinct_values_len(), expected); + } + Ok(()) + } + + fn test_statistics_relative_error_is_configurable() -> Result<(), DatabaseError> { + // Count-Min width is next_power_of_two(2 / error): 0.001 -> 2048, 0.01 -> 256, + // and error_bound(n) = ceil(2n / width). + for (builder, width) in [ + (DataBaseBuilder::path("."), 2048), + ( + DataBaseBuilder::path(".").statistics_relative_error(0.01), + 256, + ), + ] { + let mut kite_sql = builder.build_in_memory()?; + kite_sql.ddl("create table t1 (a int primary key)")?; + // ANALYZE needs at least DEFAULT_NUM_OF_BUCKETS rows. + let values = (0..DEFAULT_NUM_OF_BUCKETS) + .map(|i| format!("({i})")) + .join(", "); + kite_sql + .run(format!("insert into t1 values {values}"))? + .done()?; + kite_sql.analyze("t1")?; + + let table_name = "t1".to_string().into(); + let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache()); + let index = loader.load(&table_name, 0)?.unwrap(); + assert_eq!(index.sketch().error_bound(width), 2); + } Ok(()) } @@ -227,20 +353,24 @@ mod test { database.ddl(sql) } - fn test_statistics_meta_roundtrip() -> Result<(), DatabaseError> { + fn test_column_and_index_statistics_roundtrip() -> Result<(), DatabaseError> { let temp_dir = TempDir::new().expect("unable to create temporary working directory"); let buckets = 10; let mut kite_sql = DataBaseBuilder::path(temp_dir.path()) .histogram_buckets(buckets) .build_rocksdb()?; - create_table(&mut kite_sql, "create table t1 (a int primary key, b int)")?; + create_table( + &mut kite_sql, + "create table t1 (a int primary key, b int, c varchar)", + )?; create_index(&mut kite_sql, "create index b_index on t1 (b)")?; create_index(&mut kite_sql, "create index p_index on t1 (a, b)")?; - for i in 0..DEFAULT_NUM_OF_BUCKETS + 1 { + // Keep this fixture independent of the default histogram bucket count. + for i in 0..101 { kite_sql - .run(format!("insert into t1 values({i}, {})", i % 20))? + .run(format!("insert into t1 values({i}, {}, NULL)", i % 20))? .done()?; } kite_sql.analyze("t1")?; @@ -248,6 +378,10 @@ mod test { let table_name = "t1".to_string().into(); let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache()); + let table = kite_sql.state.table_cache().get(&table_name).unwrap(); + let a_id = *table.get_column_id_by_name("a").unwrap(); + let b_id = *table.get_column_id_by_name("b").unwrap(); + let c_id = *table.get_column_id_by_name("c").unwrap(); let statistics_meta_pk_index = loader.load(&table_name, 0)?.unwrap(); assert_eq!(statistics_meta_pk_index.index_id(), 0); assert_eq!(statistics_meta_pk_index.histogram().values_len(), 101); @@ -262,8 +396,22 @@ mod test { assert_eq!(statistics_meta_p_index.index_id(), 2); assert_eq!(statistics_meta_p_index.histogram().values_len(), 101); assert_eq!(statistics_meta_p_index.histogram().buckets_len(), buckets); - let expected_frequency = - loader.collect_count(&table_name, 1, &Range::Eq(DataValue::Int32(7)))?; + for (column, distinct) in [(a_id, 101), (b_id, 20)] { + let meta = loader.load_column(&table_name, column).unwrap(); + assert_eq!(meta.values_len(), 101); + assert_eq!(meta.distinct_values_len(), distinct); + assert_eq!(meta.buckets_len(), 0); + } + let unindexed = loader.load_column(&table_name, c_id).unwrap(); + assert_eq!(unindexed.values_len(), 0); + assert_eq!(unindexed.null_count(), 101); + assert_eq!(unindexed.distinct_values_len(), 0); + let expected_frequency = loader.collect_count( + &table_name, + 1, + &Range::Eq(DataValue::Int32(7)), + &mut Vec::new(), + )?; // Reload the stored parts instead of reusing ANALYZE's in-memory cache. drop(kite_sql); @@ -277,13 +425,86 @@ mod test { assert_eq!(meta.histogram().buckets_len(), buckets); } assert_eq!( - loader.collect_count(&table_name, 1, &Range::Eq(DataValue::Int32(7)))?, + loader.collect_count( + &table_name, + 1, + &Range::Eq(DataValue::Int32(7)), + &mut Vec::new() + )?, expected_frequency ); + assert_eq!( + loader + .load_column(&table_name, a_id) + .unwrap() + .distinct_values_len(), + 101 + ); + assert_eq!( + loader + .load_column(&table_name, b_id) + .unwrap() + .distinct_values_len(), + 20 + ); + assert_eq!( + loader.load_column(&table_name, c_id).unwrap().null_count(), + 101 + ); + reopened.ddl("drop index t1.b_index")?; + let loader = StatisticMetaLoader::new(reopened.state.meta_cache()); + assert!(loader.load(&table_name, 1)?.is_none()); + assert_eq!( + loader + .load_column(&table_name, b_id) + .unwrap() + .distinct_values_len(), + 20 + ); + let tx = reopened.storage.transaction()?; + assert!(tx + .index_statistics( + &mut TableCodec::default(), + "t1", + 1, + reopened.state.table_arena().borrow_mut() + )? + .is_none()); + assert_eq!( + tx.column_statistics( + &mut TableCodec::default(), + "t1", + b_id, + reopened.state.table_arena().borrow_mut() + )? + .unwrap() + .distinct_values_len(), + 20 + ); + drop(tx); + reopened.ddl("alter table t1 drop column c")?; + let tx = reopened.storage.transaction()?; + assert!(tx + .column_statistics( + &mut TableCodec::default(), + "t1", + c_id, + reopened.state.table_arena().borrow_mut() + )? + .is_none()); + assert!(tx + .column_statistics( + &mut TableCodec::default(), + "t1", + a_id, + reopened.state.table_arena().borrow_mut() + )? + .is_some()); + drop(tx); let mut iter = reopened.run("select a from t1 where b=7 order by a")?; let mut rows = Vec::new(); while iter - .next_tuple(|_, row| rows.push(row.values[0].clone()))? + .next_tuple(|_, row| rows.push(row[0].clone()))? .is_some() {} iter.done()?; @@ -292,9 +513,9 @@ mod test { // A missing root must not publish partially reconstructed statistics. let mut codec = TableCodec::default(); let mut tx = reopened.storage.transaction()?; - codec.with_statistics_meta( + codec.with_index_statistics( "t1", - 1, + 2, None, reopened.state.table_arena().borrow(), |key, _| tx.remove(key), @@ -304,6 +525,8 @@ mod test { let mut reopened = DataBaseBuilder::path(temp_dir.path()).build_rocksdb()?; assert!(matches!(reopened.load(CatalogKind::Table(table_name)), Err(DatabaseError::InvalidValue(message)) if message == "statistics meta is incomplete")); + assert!(reopened.state.meta_cache().indexes.is_empty()); + assert!(reopened.state.meta_cache().columns.is_empty()); Ok(()) } @@ -325,12 +548,19 @@ mod test { let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache()); assert!(loader.load(&table_name, 1)?.is_some()); assert_eq!( - loader.collect_count(&table_name, 1, &Range::Eq(DataValue::Int32(7)))?, + loader + .collect_count( + &table_name, + 1, + &Range::Eq(DataValue::Int32(7)), + &mut Vec::new() + )? + .map(|(rows, _)| rows), Some(1) ); let mut transaction = kite_sql.storage.transaction()?; let mut table_codec = TableCodec::default(); - let keys: Vec> = table_codec.with_statistics_index_bound("t1", 1, |min, max| { + let keys: Vec> = table_codec.with_index_statistics_bound("t1", 1, |min, max| { let mut iter = transaction.range(Bound::Included(min), Bound::Included(max))?; let mut keys = Vec::new(); while let Some((key, _)) = iter.try_next()? { @@ -345,7 +575,14 @@ mod test { let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache()); assert!(loader.load(&table_name, 1)?.is_some()); assert_eq!( - loader.collect_count(&table_name, 1, &Range::Eq(DataValue::Int32(7)))?, + loader + .collect_count( + &table_name, + 1, + &Range::Eq(DataValue::Int32(7)), + &mut Vec::new() + )? + .map(|(rows, _)| rows), Some(1) ); @@ -363,7 +600,11 @@ mod test { let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache()); assert!(loader.load(&table_name, 1)?.is_none()); - assert!(!kite_sql.state.meta_cache().contains_key(&(table_name, 1))); + assert!(!kite_sql + .state + .meta_cache() + .indexes + .contains_key(&(table_name, 1))); Ok(()) } @@ -412,16 +653,27 @@ mod test { })?; let table_name = "t1".to_string().into(); let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache()); - let statistics_meta = loader.load(&table_name, 0)?.unwrap(); - let expected_keys = 1 + let index_statistics = loader.load(&table_name, 0)?.unwrap(); + let expected_keys = 2 + 1 - + statistics_meta + + index_statistics .sketch() .storage_page_count(COUNT_MIN_SKETCH_STORAGE_PAGE_LEN) + 1 - + statistics_meta.histogram().buckets_len(); + + index_statistics.histogram().buckets_len(); assert_eq!(keys, expected_keys); - + drop(transaction); + kite_sql.ddl("drop table t1")?; + assert!(kite_sql.state.meta_cache().columns.is_empty()); + assert!(kite_sql.state.meta_cache().indexes.is_empty()); + let transaction = kite_sql.storage.transaction()?; + table_codec.with_statistics_bound("t1", |min, max| { + assert!(transaction + .range(Bound::Included(min), Bound::Included(max))? + .try_next()? + .is_none()); + Ok(()) + })?; Ok(()) } } diff --git a/src/execution/dml/copy_from_file.rs b/src/execution/dml/copy_from_file.rs index 1d842de8..72cd8f19 100644 --- a/src/execution/dml/copy_from_file.rs +++ b/src/execution/dml/copy_from_file.rs @@ -20,8 +20,8 @@ use crate::execution::{ use crate::iter_ext::Itertools; use crate::planner::operator::copy_from_file::CopyFromFileOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::{Tuple, TupleImpl}; use crate::types::tuple_builder::TupleBuilder; use std::fs::File; use std::io::BufReader; @@ -30,7 +30,7 @@ pub struct CopyFromFile<'a> { op: &'a CopyFromFileOperator, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CopyFromFile<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for CopyFromFile<'a> { type Input = &'a CopyFromFileOperator; fn into_executor( @@ -44,7 +44,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for CopyFromFile<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CopyFromFile<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for CopyFromFile<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -87,7 +87,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CopyFromFile<'a> { let tuple_builder = TupleBuilder::new(column_types, Some(table.primary_key_indices())); let mut record = csv::StringRecord::new(); - let mut chunk = Tuple::new(None, Vec::with_capacity(column_count)); + let mut chunk = TupleImpl::from(Tuple::new(None, Vec::with_capacity(column_count))); while reader.read_record(&mut record)? { if !(record.len() == column_count || record.len() == column_count + 1 && record.get(column_count) == Some("")) @@ -95,7 +95,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CopyFromFile<'a> { return Err(DatabaseError::MisMatch("columns", "values")); } - tuple_builder.build_with_row(&mut chunk, record.iter().take(column_count))?; + tuple_builder.build_with_row(chunk.output_mut(), record.iter().take(column_count))?; let mut state = arena.local_state(plan_arena); let (transaction, table_codec) = state.transaction_codec_mut(); transaction.append_tuple(table_codec, &table_name, &chunk, &serializers, false)?; @@ -189,7 +189,7 @@ mod tests { let mut iter = db.run("select count(*) from test_copy")?; let count = iter - .next_tuple(|_, tuple| tuple.values[0].to_string())? + .next_tuple(|_, tuple| tuple[0].to_string())? .expect("count row"); assert_eq!(count, "2"); iter.done()?; diff --git a/src/execution/dml/copy_to_file.rs b/src/execution/dml/copy_to_file.rs index 6fc3bf47..60a2fef5 100644 --- a/src/execution/dml/copy_to_file.rs +++ b/src/execution/dml/copy_to_file.rs @@ -21,7 +21,7 @@ use crate::iter_ext::Itertools; use crate::planner::operator::copy_to_file::CopyToFileOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::tuple_builder::TupleBuilder; pub struct CopyToFile<'a> { @@ -30,7 +30,7 @@ pub struct CopyToFile<'a> { input: Option, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for CopyToFile<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for CopyToFile<'a> { type Input = (&'a CopyToFileOperator, &'a LogicalPlan); fn into_executor( @@ -83,7 +83,7 @@ impl CopyToFile<'_> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CopyToFile<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for CopyToFile<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -96,7 +96,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CopyToFile<'a> { let mut writer = self.create_writer()?; while arena.next_tuple(input, plan_arena)? { - writer.write_record(arena.result_tuple().values.iter().map(|v| v.to_string()))?; + writer.write_record(arena.result_tuple().iter().map(|v| v.to_string()))?; } writer.flush().map_err(DatabaseError::from)?; @@ -105,7 +105,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for CopyToFile<'a> { } else { format!("{} [{}]", self.op, self.column_names.iter().join(", ")) }; - arena.produce_tuple(TupleBuilder::build_result(message)); + arena.produce(TupleBuilder::build_result(message).into()); arena.resume(); Ok(()) } diff --git a/src/execution/dml/delete.rs b/src/execution/dml/delete.rs index 301e6c98..75e83700 100644 --- a/src/execution/dml/delete.rs +++ b/src/execution/dml/delete.rs @@ -20,7 +20,7 @@ use crate::execution::{ use crate::planner::operator::delete::DeleteOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::index::{Index, IndexType}; use crate::types::tuple_builder::TupleBuilder; @@ -29,7 +29,7 @@ pub struct Delete<'a> { input: Option, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Delete<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for Delete<'a> { type Input = (&'a DeleteOperator, &'a LogicalPlan); fn into_executor( @@ -50,7 +50,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Delete<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Delete<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Delete<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -89,14 +89,14 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Delete<'a> { let mut deleted_count = 0; while arena.next_tuple(input, plan_arena)? { - let tuple = arena.materialize_tuple(); - let Some(tuple_id) = tuple.pk.as_ref() else { + let tuple = arena.take_tuple(); + let Some(tuple_id) = tuple.pk() else { continue; }; for (index_id, index_ty, exprs) in index_templates.iter() { arena.rewrite(exprs, plan_arena, Some(&tuple))?; let mut state = arena.local_state(plan_arena); - let (values, transaction, table_codec) = state.index_values_transaction_codec_mut(); + let (values, transaction, table_codec) = state.tuple_transaction_codec_mut(); let index = Index::new(*index_id, values, *index_ty); transaction.del_index(table_codec, self.table_name, &index, tuple_id)?; } @@ -107,7 +107,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Delete<'a> { deleted_count += 1; } - arena.produce_tuple(TupleBuilder::build_result(deleted_count.to_string())); + arena.produce(TupleBuilder::build_result(deleted_count.to_string()).into()); arena.resume(); Ok(()) } diff --git a/src/execution/dml/insert.rs b/src/execution/dml/insert.rs index a2838170..037a6247 100644 --- a/src/execution/dml/insert.rs +++ b/src/execution/dml/insert.rs @@ -21,9 +21,9 @@ use crate::iter_ext::Itertools; use crate::planner::operator::insert::InsertOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::index::Index; -use crate::types::tuple::{Schema, Tuple}; +use crate::types::tuple::{Schema, Tuple, TupleImpl}; use crate::types::tuple_builder::TupleBuilder; use crate::types::value::DataValue; use crate::types::ColumnId; @@ -37,7 +37,7 @@ pub struct Insert<'a> { is_mapping_by_name: bool, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Insert<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for Insert<'a> { type Input = (&'a InsertOperator, &'a LogicalPlan); fn into_executor( @@ -91,7 +91,7 @@ impl Insert<'_> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Insert<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Insert<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -121,16 +121,22 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Insert<'a> { .iter() .map(|column| plan_arena.column(*column).datatype().serializable()) .collect_vec(); - let mut tuple = Tuple::new(None, Vec::with_capacity(table_snapshot.columns_len)); + let mut tuple = TupleImpl::from(Tuple::new( + None, + Vec::with_capacity(table_snapshot.columns_len), + )); while arena.next_tuple(input, plan_arena)? { let mut tuple_map = HashMap::with_capacity(self.input_schema.len()); - for (i, value) in arena.result_tuple_mut().values.drain(..).enumerate() { + let mut i = 0; + arena.result_mut().drain_values(|value| { let column = plan_arena.column(self.input_schema[i]); tuple_map.insert(Self::column_key(column, self.is_mapping_by_name), value); - } + i += 1; + }); - tuple.values.clear(); + let row = tuple.output_mut(); + row.values.clear(); for column in table_snapshot.columns.iter() { let column = plan_arena.column(*column); let value = match tuple_map @@ -144,19 +150,18 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Insert<'a> { if value.is_null() && !column.nullable() { return Err(DatabaseError::not_null_column(column.name().to_string())); } - tuple.values.push(value) + row.values.push(value) } - tuple.pk = Some(Tuple::primary_projection( + row.pk = Some(Tuple::primary_projection( table_snapshot.primary_key_indices, - &tuple.values, + row.values.as_slice(), )); for (index_meta, exprs) in table_snapshot.index_metas.iter() { let index_meta = plan_arena.index(*index_meta); - let tuple_id = tuple.pk.as_ref().ok_or(DatabaseError::PrimaryKeyNotFound)?; + let tuple_id = tuple.pk().ok_or(DatabaseError::PrimaryKeyNotFound)?; arena.rewrite(exprs, plan_arena, Some(&tuple))?; let mut state = arena.local_state(plan_arena); - let (values, transaction, table_codec) = - state.index_values_transaction_codec_mut(); + let (values, transaction, table_codec) = state.tuple_transaction_codec_mut(); let index = Index::new(index_meta.id, values, index_meta.ty); transaction.add_index(table_codec, self.table_name, index, tuple_id)?; } @@ -172,7 +177,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Insert<'a> { inserted_count += 1; } } - arena.produce_tuple(TupleBuilder::build_result(inserted_count.to_string())); + arena.produce(TupleBuilder::build_result(inserted_count.to_string()).into()); arena.resume(); Ok(()) } diff --git a/src/execution/dml/update.rs b/src/execution/dml/update.rs index 2a722afa..55e1e215 100644 --- a/src/execution/dml/update.rs +++ b/src/execution/dml/update.rs @@ -22,7 +22,7 @@ use crate::planner::operator::update::UpdateOperator; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::index::{Index, IndexMeta, IndexType}; use crate::types::tuple::{Schema, Tuple}; use crate::types::tuple_builder::TupleBuilder; @@ -36,7 +36,7 @@ pub struct Update<'a> { input: Option, } -impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Update<'a> { +impl<'a, T: StorageTransaction + 'a> WriteExecutor<'a, T> for Update<'a> { type Input = (&'a UpdateOperator, &'a LogicalPlan); fn into_executor( @@ -87,7 +87,7 @@ impl Update<'_> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Update<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Update<'a> { fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -136,8 +136,8 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Update<'a> { while arena.next_tuple(input, plan_arena)? { let mut is_overwrite = true; - let mut tuple = arena.materialize_tuple(); - let Some(old_pk) = tuple.pk.clone() else { + let mut tuple = arena.take_tuple(); + let Some(old_pk) = tuple.pk().cloned() else { continue; }; @@ -153,8 +153,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Update<'a> { arena.rewrite(exprs, plan_arena, Some(&tuple))?; let mut state = arena.local_state(plan_arena); - let (values, transaction, table_codec) = - state.index_values_transaction_codec_mut(); + let (values, transaction, table_codec) = state.tuple_transaction_codec_mut(); let old_index = Index::new(index_meta.id, values, index_meta.ty); transaction.del_index(table_codec, self.table_name, &old_index, &old_pk)?; } @@ -166,12 +165,11 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Update<'a> { let value = plan_arena .expression(*expr) .eval(plan_arena, Some(&tuple))?; - tuple.values[i] = value.into_owned(); + tuple[i] = value.into_owned(); } } - let new_pk = - Tuple::primary_projection(table_snapshot.primary_key_indices, &tuple.values); + let new_pk = Tuple::primary_projection(table_snapshot.primary_key_indices, &tuple); let primary_key_changed = new_pk != old_pk; if primary_key_changed { let mut state = arena.local_state(plan_arena); @@ -191,13 +189,12 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Update<'a> { } arena.rewrite(exprs, plan_arena, Some(&tuple))?; let mut state = arena.local_state(plan_arena); - let (values, transaction, table_codec) = - state.index_values_transaction_codec_mut(); + let (values, transaction, table_codec) = state.tuple_transaction_codec_mut(); let new_index = Index::new(index_meta.id, values, index_meta.ty); transaction.add_index(table_codec, self.table_name, new_index, &new_pk)?; } - tuple.pk = Some(new_pk); + tuple.set_pk(Some(new_pk)); let mut state = arena.local_state(plan_arena); let (transaction, table_codec) = state.transaction_codec_mut(); let stamp = if is_overwrite { 0 } else { table_codec.stamp() }; @@ -213,7 +210,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Update<'a> { updated_count += 1; } } - arena.produce_tuple(TupleBuilder::build_result(updated_count.to_string())); + arena.produce(TupleBuilder::build_result(updated_count.to_string()).into()); arena.resume(); Ok(()) } diff --git a/src/execution/dql/aggregate/avg.rs b/src/execution/dql/aggregate/avg.rs index b9afbac2..dd068541 100644 --- a/src/execution/dql/aggregate/avg.rs +++ b/src/execution/dql/aggregate/avg.rs @@ -13,7 +13,7 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::execution::dql::aggregate::sum::SumAccumulator; +use crate::execution::dql::aggregate::sum::{text_as_number, SumAccumulator}; use crate::execution::dql::aggregate::Accumulator; use crate::expression::BinaryOperator; use crate::types::evaluator::binary_create; @@ -48,13 +48,14 @@ impl Accumulator for AvgAccumulator { return Ok(()); } } + let value = text_as_number(value); let acc = if let Some(ref mut inner) = self.inner { inner } else { self.inner .get_or_insert(SumAccumulator::new(Cow::Owned(value.logical_type()))?) }; - acc.update_value(value)?; + acc.update_value(&value)?; self.count += 1; self.result = DataValue::Null; } diff --git a/src/execution/dql/aggregate/group_concat.rs b/src/execution/dql/aggregate/group_concat.rs new file mode 100644 index 00000000..07264efe --- /dev/null +++ b/src/execution/dql/aggregate/group_concat.rs @@ -0,0 +1,74 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::execution::dql::aggregate::Accumulator; +use crate::types::value::{DataValue, Utf8Type}; +use crate::types::CharLengthUnits; +use std::collections::HashSet; +use std::fmt::Write; + +/// SQLite `group_concat()`: non-NULL values in input order joined with `,`; NULL when there are none. +pub struct GroupConcatAccumulator { + distinct_values: Option>, + result: DataValue, +} + +impl GroupConcatAccumulator { + pub fn new(distinct: bool) -> Self { + Self { + distinct_values: distinct.then(HashSet::new), + result: DataValue::Null, + } + } +} + +impl Accumulator for GroupConcatAccumulator { + fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError> { + if value.is_null() { + return Ok(()); + } + if let Some(values) = &mut self.distinct_values { + if !values.insert(value.clone()) { + return Ok(()); + } + } + match &mut self.result { + DataValue::Utf8 { value: text, .. } => { + text.push(','); + let _ = write!(text, "{value}"); + } + result => { + *result = DataValue::Utf8 { + value: value.to_string(), + ty: Utf8Type::Variable(None), + unit: CharLengthUnits::Characters, + } + } + } + Ok(()) + } + + fn evaluate(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn result(&self) -> &DataValue { + &self.result + } + + fn result_owned(self: Box) -> DataValue { + self.result + } +} diff --git a/src/execution/dql/aggregate/hash_agg.rs b/src/execution/dql/aggregate/hash_agg.rs index 927913dd..3c9a3693 100644 --- a/src/execution/dql/aggregate/hash_agg.rs +++ b/src/execution/dql/aggregate/hash_agg.rs @@ -23,7 +23,7 @@ use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; use std::collections::hash_map::IntoIter as HashMapIntoIter; use std::collections::HashMap; @@ -37,7 +37,7 @@ pub struct HashAggExecutor<'a> { output: Option, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for HashAggExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for HashAggExecutor<'a> { type Input = (&'a AggregateOperator, &'a LogicalPlan); fn into_executor( @@ -64,7 +64,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for HashAggExecutor<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashAggExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for HashAggExecutor<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.output = None; + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -104,7 +113,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashAggExecutor<'a> { return Ok(()); }; - write_aggregate_output(arena.result_tuple_mut(), accs, group_keys)?; + write_aggregate_output(arena.output_tuple_mut(), accs, group_keys)?; arena.resume(); Ok(()) } @@ -204,9 +213,9 @@ mod test { let pipeline = HepOptimizerPipeline::builder() .before_batch( - "Expression Remapper".to_string(), + "Expression Remapper", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::EvaluatorBind], + &[NormalizationRuleImpl::EvaluatorBind], ) .build(); let plan = pipeline @@ -225,7 +234,10 @@ mod test { assert_eq!(tuples.len(), 2); - let vec_values = tuples.into_iter().map(|tuple| tuple.values).collect_vec(); + let vec_values = tuples + .into_iter() + .map(|mut tuple| tuple.take_values()) + .collect_vec(); assert!(vec_values.contains(&build_integers(vec![Some(3), Some(0)]))); assert!(vec_values.contains(&build_integers(vec![Some(5), Some(1)]))); diff --git a/src/execution/dql/aggregate/mod.rs b/src/execution/dql/aggregate/mod.rs index 87c685b5..eb783dc2 100644 --- a/src/execution/dql/aggregate/mod.rs +++ b/src/execution/dql/aggregate/mod.rs @@ -14,18 +14,22 @@ mod avg; mod count; +mod group_concat; pub mod hash_agg; mod min_max; pub mod simple_agg; pub mod stream_agg; pub mod stream_distinct; mod sum; +mod total; use crate::errors::DatabaseError; use crate::execution::dql::aggregate::avg::AvgAccumulator; use crate::execution::dql::aggregate::count::{CountAccumulator, DistinctCountAccumulator}; +use crate::execution::dql::aggregate::group_concat::GroupConcatAccumulator; use crate::execution::dql::aggregate::min_max::MinMaxAccumulator; use crate::execution::dql::aggregate::sum::{DistinctSumAccumulator, SumAccumulator}; +use crate::execution::dql::aggregate::total::TotalAccumulator; use crate::expression::agg::AggKind; use crate::expression::ScalarExpression; use crate::iter_ext::Itertools; @@ -65,6 +69,8 @@ pub(crate) fn create_accumulator( (AggKind::Min, _) => Box::new(MinMaxAccumulator::new(false)), (AggKind::Max, _) => Box::new(MinMaxAccumulator::new(true)), (AggKind::Avg, distinct) => Box::new(AvgAccumulator::new(distinct)), + (AggKind::Total, distinct) => Box::new(TotalAccumulator::new(distinct)), + (AggKind::GroupConcat, distinct) => Box::new(GroupConcatAccumulator::new(distinct)), }) } diff --git a/src/execution/dql/aggregate/simple_agg.rs b/src/execution/dql/aggregate/simple_agg.rs index 64e1fe2b..09966f67 100644 --- a/src/execution/dql/aggregate/simple_agg.rs +++ b/src/execution/dql/aggregate/simple_agg.rs @@ -22,14 +22,14 @@ use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct SimpleAggExecutor<'a> { agg_calls: &'a [ExprRef], input: ExecId, returned: bool, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for SimpleAggExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for SimpleAggExecutor<'a> { type Input = (&'a AggregateOperator, &'a LogicalPlan); fn into_executor( @@ -48,7 +48,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for SimpleAggExecutor<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SimpleAggExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for SimpleAggExecutor<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.returned = false; + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -81,7 +90,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SimpleAggExecutor<'a> { } } - let output = arena.result_tuple_mut(); + let output = arena.output_tuple_mut(); output.pk = None; output.values.clear(); output.values.reserve(accs.len()); diff --git a/src/execution/dql/aggregate/stream_agg.rs b/src/execution/dql/aggregate/stream_agg.rs index 63b58304..0a896921 100644 --- a/src/execution/dql/aggregate/stream_agg.rs +++ b/src/execution/dql/aggregate/stream_agg.rs @@ -23,7 +23,7 @@ use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; use std::mem; @@ -36,7 +36,7 @@ pub struct StreamAggExecutor<'a> { input: ExecId, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for StreamAggExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for StreamAggExecutor<'a> { type Input = (&'a AggregateOperator, &'a LogicalPlan); fn into_executor( @@ -64,7 +64,17 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for StreamAggExecutor<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for StreamAggExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for StreamAggExecutor<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.group_keys = None; + self.accs.clear(); + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -77,7 +87,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for StreamAggExecutor<'a> { return Ok(()); }; write_aggregate_output( - arena.result_tuple_mut(), + arena.output_tuple_mut(), mem::take(&mut self.accs), group_keys, )?; @@ -110,7 +120,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for StreamAggExecutor<'a> { update_accumulators(&mut next_accs, self.agg_calls, tuple, plan_arena)?; mem::swap(current_keys, &mut group_keys); let current_accs = mem::replace(&mut self.accs, next_accs); - write_aggregate_output(arena.result_tuple_mut(), current_accs, group_keys)?; + write_aggregate_output(arena.output_tuple_mut(), current_accs, group_keys)?; arena.resume(); return Ok(()); } @@ -196,7 +206,9 @@ mod tests { ))?; assert_eq!( - rows.into_iter().map(|row| row.values).collect::>(), + rows.into_iter() + .map(|mut row| row.take_values()) + .collect::>(), vec![ vec![30.into(), 2.into(), 1.into()], vec![12.into(), 2.into(), 2.into()], diff --git a/src/execution/dql/aggregate/stream_distinct.rs b/src/execution/dql/aggregate/stream_distinct.rs index da1250ee..89e79938 100644 --- a/src/execution/dql/aggregate/stream_distinct.rs +++ b/src/execution/dql/aggregate/stream_distinct.rs @@ -20,17 +20,17 @@ use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::TupleImpl; use crate::types::value::DataValue; pub struct StreamDistinctExecutor<'a> { groupby_exprs: &'a [ExprRef], input: ExecId, - last_keys: Option, + last_keys: Option, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for StreamDistinctExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for StreamDistinctExecutor<'a> { type Input = (&'a AggregateOperator, &'a LogicalPlan); fn into_executor( @@ -49,7 +49,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for StreamDistinctExecutor<'a> } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for StreamDistinctExecutor<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for StreamDistinctExecutor<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.last_keys = None; + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -58,7 +67,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for StreamDistinctExecutor<'a> loop { if !arena.next_tuple(self.input, plan_arena)? { if let Some(last_keys) = self.last_keys.take() { - drop(std::mem::replace(arena.result_tuple_mut(), last_keys)); + drop(std::mem::replace(arena.result_mut(), last_keys)); arena.resume(); } else { arena.finish(); @@ -68,14 +77,14 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for StreamDistinctExecutor<'a> arena.rewrite(self.groupby_exprs, plan_arena, None::<&[DataValue]>)?; if let Some(last_keys) = &mut self.last_keys { - if last_keys.values == arena.result_tuple().values { + if last_keys.values_eq(arena.result_tuple()) { continue; } - std::mem::swap(last_keys, arena.result_tuple_mut()); + std::mem::swap(last_keys, arena.result_mut()); arena.resume(); return Ok(()); } - self.last_keys = Some(arena.materialize_tuple()); + self.last_keys = Some(arena.take_tuple()); } } } @@ -130,9 +139,9 @@ mod tests { ) -> Result { HepOptimizerPipeline::builder() .before_batch( - "Expression Remapper".to_string(), + "Expression Remapper", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::EvaluatorBind], + &[NormalizationRuleImpl::EvaluatorBind], ) .build() .instantiate(plan) @@ -186,7 +195,7 @@ mod tests { let actual = tuples .into_iter() - .flat_map(|tuple| tuple.values) + .flat_map(|mut tuple| tuple.take_values()) .flat_map(|value| value.i32()) .collect_vec(); assert_eq!(actual, vec![1, 2, 3]); @@ -242,7 +251,10 @@ mod tests { &transaction, ))?; - let actual = tuples.into_iter().map(|tuple| tuple.values).collect_vec(); + let actual = tuples + .into_iter() + .map(|mut tuple| tuple.take_values()) + .collect_vec(); assert_eq!( actual, vec![ diff --git a/src/execution/dql/aggregate/sum.rs b/src/execution/dql/aggregate/sum.rs index 1e220167..0cdc5ea0 100644 --- a/src/execution/dql/aggregate/sum.rs +++ b/src/execution/dql/aggregate/sum.rs @@ -21,6 +21,46 @@ use crate::types::LogicalType; use std::borrow::Cow; use std::collections::HashSet; +/// Text in a numeric aggregate is read as its longest numeric prefix (0 without one), as SQLite does. +pub(crate) fn text_as_number(value: &DataValue) -> Cow<'_, DataValue> { + match value { + DataValue::Utf8 { value: text, .. } => { + Cow::Owned(DataValue::Float64(numeric_prefix(text).into())) + } + _ => Cow::Borrowed(value), + } +} + +fn numeric_prefix(text: &str) -> f64 { + let text = text.trim_start(); + let bytes = text.as_bytes(); + let digits = |mut i: usize| { + while bytes.get(i).is_some_and(u8::is_ascii_digit) { + i += 1; + } + i + }; + let start = usize::from(matches!(bytes.first(), Some(b'+' | b'-'))); + let mut end = digits(start); + let mut has_digits = end > start; + if bytes.get(end) == Some(&b'.') { + let fraction = digits(end + 1); + has_digits |= fraction > end + 1; + end = fraction; + } + if !has_digits { + return 0.0; + } + if matches!(bytes.get(end), Some(b'e' | b'E')) { + let sign = end + 1 + usize::from(matches!(bytes.get(end + 1), Some(b'+' | b'-'))); + let exponent = digits(sign); + if exponent > sign { + end = exponent; + } + } + text[..end].parse().unwrap_or(0.0) +} + pub struct SumAccumulator { result: DataValue, evaluator: BinaryEvaluatorRef, @@ -41,12 +81,13 @@ impl SumAccumulator { impl Accumulator for SumAccumulator { fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError> { if !value.is_null() { + let value = text_as_number(value); if self.result.is_null() { - self.result = value.clone(); + self.result = value.into_owned(); } else { self.result = self .evaluator - .binary_eval(Cow::Borrowed(&self.result), value)?; + .binary_eval(Cow::Borrowed(&self.result), &value)?; } } @@ -100,6 +141,22 @@ impl Accumulator for DistinctSumAccumulator { mod tests { use super::*; + #[test] + fn text_reads_its_numeric_prefix() { + for (text, expected) in [ + ("10", 10.0), + (" -1.5e2x", -150.0), + ("+.5", 0.5), + ("3.", 3.0), + ("7e", 7.0), + ("true", 0.0), + (".", 0.0), + ("", 0.0), + ] { + assert_eq!(numeric_prefix(text), expected, "{text:?}"); + } + } + #[test] fn sum_results() -> Result<(), DatabaseError> { let mut accumulator = SumAccumulator::new(Cow::Borrowed(&LogicalType::Integer))?; diff --git a/src/execution/dql/aggregate/total.rs b/src/execution/dql/aggregate/total.rs new file mode 100644 index 00000000..ac8c27fa --- /dev/null +++ b/src/execution/dql/aggregate/total.rs @@ -0,0 +1,67 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::execution::dql::aggregate::sum::text_as_number; +use crate::execution::dql::aggregate::Accumulator; +use crate::types::value::DataValue; +use crate::types::LogicalType; +use std::collections::HashSet; + +/// SQLite `total()`: always a double, and 0.0 rather than NULL when nothing is summed. +pub struct TotalAccumulator { + distinct_values: Option>, + result: DataValue, +} + +impl TotalAccumulator { + pub fn new(distinct: bool) -> Self { + Self { + distinct_values: distinct.then(HashSet::new), + result: DataValue::Float64(0.0.into()), + } + } +} + +impl Accumulator for TotalAccumulator { + fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError> { + if value.is_null() { + return Ok(()); + } + if let Some(values) = &mut self.distinct_values { + if !values.insert(value.clone()) { + return Ok(()); + } + } + let DataValue::Float64(value) = text_as_number(value) + .into_owned() + .cast(&LogicalType::Double)? + else { + unreachable!("cast to Double yields Float64"); + }; + let DataValue::Float64(total) = &mut self.result else { + unreachable!("total is always Float64"); + }; + *total = (total.0 + value.0).into(); + Ok(()) + } + + fn result(&self) -> &DataValue { + &self.result + } + + fn result_owned(self: Box) -> DataValue { + self.result + } +} diff --git a/src/execution/dql/describe.rs b/src/execution/dql/describe.rs index cbd18eff..6368653f 100644 --- a/src/execution/dql/describe.rs +++ b/src/execution/dql/describe.rs @@ -17,7 +17,7 @@ use crate::errors::DatabaseError; use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor}; use crate::planner::operator::describe::DescribeOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::{DataValue, Utf8Type}; use crate::types::CharLengthUnits; use std::sync::LazyLock; @@ -56,7 +56,7 @@ impl From<&DescribeOperator> for Describe { } } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Describe { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Describe { type Input = &'a DescribeOperator; fn into_executor( @@ -71,7 +71,17 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Describe { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Describe { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Describe { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.columns = None; + self.cursor = 0; + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -99,7 +109,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Describe { let column = plan_arena.column(column_ref); let default = describe_default(column, plan_arena); - let output = arena.result_tuple_mut(); + let output = arena.output_tuple_mut(); output.pk = None; output.values.clear(); fill_describe_row(&mut output.values, column, default); diff --git a/src/execution/dql/dummy.rs b/src/execution/dql/dummy.rs index 019053dc..a39ce217 100644 --- a/src/execution/dql/dummy.rs +++ b/src/execution/dql/dummy.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor}; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::tuple::Tuple; pub struct Dummy { @@ -30,7 +30,7 @@ impl Default for Dummy { } } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Dummy { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Dummy { type Input = (); fn into_executor( @@ -45,7 +45,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Dummy { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Dummy { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Dummy { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.row = Some(Tuple::default()); + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -55,7 +64,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Dummy { arena.finish(); return Ok(()); }; - arena.produce_tuple(row); + arena.produce(row.into()); Ok(()) } } diff --git a/src/execution/dql/explain.rs b/src/execution/dql/explain.rs index dd999aba..c8be66a8 100644 --- a/src/execution/dql/explain.rs +++ b/src/execution/dql/explain.rs @@ -16,7 +16,7 @@ use crate::errors::DatabaseError; use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor}; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::{DataValue, Utf8Type}; use crate::types::CharLengthUnits; @@ -25,7 +25,7 @@ pub struct Explain<'a> { emitted: bool, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Explain<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Explain<'a> { type Input = &'a LogicalPlan; fn into_executor( @@ -42,7 +42,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Explain<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Explain<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Explain<'a> { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.emitted = false; + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -54,7 +63,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Explain<'a> { } let plan = self.plan.explain(plan_arena, 0); - let output = arena.result_tuple_mut(); + let output = arena.output_tuple_mut(); output.pk = None; output.values.clear(); output.values.push(DataValue::Utf8 { diff --git a/src/execution/dql/external_sort.rs b/src/execution/dql/external_sort.rs index c8a914c2..9e181505 100644 --- a/src/execution/dql/external_sort.rs +++ b/src/execution/dql/external_sort.rs @@ -21,7 +21,7 @@ use crate::execution::{ use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::fs::File; use std::io::BufReader; use std::mem; @@ -49,7 +49,7 @@ pub struct ExternalSort<'a> { input: ExecId, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ExternalSort<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for ExternalSort<'a> { type Input = (&'a SortOperator, &'a LogicalPlan); fn into_executor( @@ -68,7 +68,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ExternalSort<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ExternalSort<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for ExternalSort<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.rows = None; + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -81,7 +90,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ExternalSort<'a> { return Ok(()); }; plan_arena.restore_outer_values(row.outer_values); - arena.produce_tuple(row.tuple); + arena.produce(row.tuple); return Ok(()); } @@ -91,10 +100,10 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ExternalSort<'a> { // pass. By default R <= 1,024 and B is about 1 MiB; one oversized row makes this a // soft bound. let sort_fields = self.sort_fields; - let mut rows = SpillVec::new().on_flush(move |rows| sort_segment(sort_fields, rows)); + let mut rows = SpillVec::new().on_flush(move |_, rows| sort_segment(sort_fields, rows)); let mut runs = Vec::new(); while arena.next_tuple(self.input, plan_arena)? { - let tuple = arena.materialize_tuple(); + let tuple = arena.take_tuple(); let row = SortRow::new(sort_fields, tuple, plan_arena)?; if let Some(segment) = rows.push(row)? { runs.push(Run::new(segment, 1)); @@ -154,7 +163,7 @@ struct RunCursor<'source> { impl<'source> RunCursor<'source> { fn new(source: &'source SpillReader, run: &Run) -> Result { - let mut reader = source.open_segment_reader()?; + let mut reader = source.open_segment_reader(None)?; reader.reset(run.first_segment)?; let mut cursor = Self { remaining_segments: run.segment_count - 1, @@ -289,7 +298,7 @@ mod test { use crate::expression::ScalarExpression; use crate::planner::operator::sort::SortField; use crate::planner::{ExecMetaArena, MetaArena}; - use crate::types::tuple::Tuple; + use crate::types::tuple::{Tuple, TupleLike}; use crate::types::value::{DataValue, Utf8Type}; use crate::types::{CharLengthUnits, LogicalType}; use std::cmp::Ordering; @@ -315,17 +324,17 @@ mod test { let mut plan_arena = ExecMetaArena::new(plan_arena); let mut rows = SpillVec::new() .limit(4, usize::MAX) - .on_flush(|rows| sort_segment(&sort_fields, rows)); + .on_flush(|_, rows| sort_segment(&sort_fields, rows)); for value in [DataValue::Int32(2), DataValue::Null, DataValue::Int32(1)] { let _ = rows.push(SortRow::new( &sort_fields, - Tuple::new(None, vec![value]), + Tuple::new(None, vec![value]).into(), &mut plan_arena, )?)?; } let values = finish_sort(rows, Vec::new(), &sort_fields, 2)? - .map(|row| row.map(|row| row.tuple.values[0].clone())) + .map(|row| row.map(|row| row.tuple.value_at(0).clone())) .collect::, _>>()?; assert_eq!( values, @@ -366,7 +375,7 @@ mod test { let mut plan_arena = ExecMetaArena::new(plan_arena); let mut rows = SpillVec::new() .limit(2, usize::MAX) - .on_flush(|rows| sort_segment(&sort_fields, rows)); + .on_flush(|_, rows| sort_segment(&sort_fields, rows)); let mut runs = Vec::new(); for (sequence, value) in values.into_iter().enumerate() { let tuple = Tuple::new( @@ -382,7 +391,7 @@ mod test { unit: CharLengthUnits::Characters, }, ); - let row = SortRow::new(&sort_fields, tuple, &mut plan_arena)?; + let row = SortRow::new(&sort_fields, tuple.into(), &mut plan_arena)?; if let Some(segment) = rows.push(row)? { runs.push(Run::new(segment, 1)); } @@ -392,7 +401,7 @@ mod test { .map(|row| { row.map(|row| { assert_eq!(row.outer_values.len(), 1); - let sequence = row.tuple.values[1].i32().unwrap(); + let sequence = row.tuple.value_at(1).i32().unwrap(); assert_eq!(row.outer_values[0].0.pos, sequence as usize); assert_eq!( row.outer_values[0].1.utf8(), @@ -404,7 +413,7 @@ mod test { .collect::, _>>()?; let positions = tuples .iter() - .map(|tuple| tuple.values[1].clone()) + .map(|tuple| tuple.value_at(1).clone()) .collect::>(); assert_eq!( positions, @@ -420,8 +429,8 @@ mod test { DataValue::Int32(6), ] ); - assert_eq!(tuples[0].pk, Some(DataValue::Int32(1))); - assert_eq!(tuples[1].pk, Some(DataValue::Int32(5))); + assert_eq!(tuples[0].pk().cloned(), Some(DataValue::Int32(1))); + assert_eq!(tuples[1].pk().cloned(), Some(DataValue::Int32(5))); Ok(()) } @@ -461,7 +470,7 @@ mod test { ]; let mut plan_arena = ExecMetaArena::new(plan_arena); - let mut rows = SpillVec::new().on_flush(|rows| sort_segment(&sort_fields, rows)); + let mut rows = SpillVec::new().on_flush(|_, rows| sort_segment(&sort_fields, rows)); let mut runs = Vec::new(); for position in 0..ROW_COUNT { let sequence = (position * 7919) % ROW_COUNT; @@ -472,7 +481,9 @@ mod test { }; let sequence = DataValue::Int32(sequence as i32); let tuple = Tuple::new(Some(sequence.clone()), vec![key, sequence]); - if let Some(segment) = rows.push(SortRow::new(&sort_fields, tuple, &mut plan_arena)?)? { + if let Some(segment) = + rows.push(SortRow::new(&sort_fields, tuple.into(), &mut plan_arena)?)? + { runs.push(Run::new(segment, 1)); } } @@ -494,8 +505,8 @@ mod test { let actual = tuples .iter() .map(|tuple| { - assert_eq!(tuple.pk.as_ref(), tuple.values.get(1)); - tuple.values[1].clone() + assert_eq!(tuple.pk(), Some(tuple.value_at(1))); + tuple.value_at(1).clone() }) .collect::>(); assert_eq!( diff --git a/src/execution/dql/filter.rs b/src/execution/dql/filter.rs index 757ed6ec..766040ef 100644 --- a/src/execution/dql/filter.rs +++ b/src/execution/dql/filter.rs @@ -20,13 +20,13 @@ use crate::planner::operator::filter::FilterOperator; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct Filter { predicate: ExprRef, input: ExecId, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Filter { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Filter { type Input = (&'a FilterOperator, &'a LogicalPlan); fn into_executor( @@ -44,7 +44,15 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Filter { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Filter { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Filter { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/dql/function_scan.rs b/src/execution/dql/function_scan.rs index e7b206c1..cff849dd 100644 --- a/src/execution/dql/function_scan.rs +++ b/src/execution/dql/function_scan.rs @@ -17,7 +17,7 @@ use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNo use crate::expression::function::table::TableFunction; use crate::planner::operator::function_scan::FunctionScanOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::tuple::Tuple; pub struct FunctionScan<'a> { @@ -34,7 +34,7 @@ impl<'a> From<&'a FunctionScanOperator> for FunctionScan<'a> { } } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for FunctionScan<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for FunctionScan<'a> { type Input = &'a FunctionScanOperator; fn into_executor( @@ -49,7 +49,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for FunctionScan<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for FunctionScan<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for FunctionScan<'a> { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.iter = None; + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -65,7 +74,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for FunctionScan<'a> { arena.finish(); return Ok(()); }; - arena.produce_tuple(tuple); + arena.produce(tuple.into()); Ok(()) } } diff --git a/src/execution/dql/index_scan.rs b/src/execution/dql/index_scan.rs index 164fafa4..940b0b45 100644 --- a/src/execution/dql/index_scan.rs +++ b/src/execution/dql/index_scan.rs @@ -14,22 +14,22 @@ use crate::errors::DatabaseError; use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor}; -use crate::expression::range_detacher::{IndexRangeColumn, Range, RangeDetacher}; +use crate::expression::range_detacher::{IndexRangeColumn, RangeDetacher}; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::operator::SortOption; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::{IndexIter, IndexRanges, Iter, Transaction}; -use crate::types::index::{IndexInfo, IndexLookup, RuntimeIndexProbe}; +use crate::storage::{IndexIter, IndexRanges, Iter, StorageTransaction}; +use crate::types::index::{IndexInfo, IndexLookup}; use std::borrow::Cow; -pub(crate) struct IndexScan<'a, T: Transaction + 'a> { +pub(crate) struct IndexScan<'a, T: StorageTransaction + 'a> { op: &'a TableScanOperator, info: &'a IndexInfo, lookup: &'a IndexLookup, iter: Option>, } -impl<'a, T: Transaction + 'a> IndexScan<'a, T> { +impl<'a, T: StorageTransaction + 'a> IndexScan<'a, T> { pub(crate) fn new( op: &'a TableScanOperator, info: &'a IndexInfo, @@ -44,24 +44,23 @@ impl<'a, T: Transaction + 'a> IndexScan<'a, T> { } pub(crate) fn ranges( - &self, + info: &'a IndexInfo, + lookup: &'a IndexLookup, arena: &mut ExecArena<'a, T>, plan_arena: &mut (dyn MetaArena + 'a), ) -> Result, DatabaseError> { - let info = self.info; - let mut range = match self.lookup { + let mut range = match lookup { IndexLookup::Static(range) => Cow::Borrowed(range), - IndexLookup::Probe => Cow::Owned(match arena.pop_runtime_probe() { - RuntimeIndexProbe::Eq(value) => Range::Eq(value), - RuntimeIndexProbe::Scope { min, max } => Range::Scope { min, max }, - }), + IndexLookup::Probe => Cow::Owned(arena.pop_runtime_probe().into()), }; - if plan_arena.has_bound_params() && range.has_parameter() { + if range.has_parameter() { let plan_arena = &*plan_arena; range .to_mut() - .bind_parameters(&|id| plan_arena.bound_param(id))?; + .bind_parameters(&|id| plan_arena.bound_param(id), &|id| { + plan_arena.init_value(id) + })?; } if let ( Some(predicate), @@ -84,7 +83,7 @@ impl<'a, T: Transaction + 'a> IndexScan<'a, T> { } } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for IndexScan<'a, T> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for IndexScan<'a, T> { type Input = (&'a TableScanOperator, &'a IndexInfo, &'a IndexLookup); fn into_executor( @@ -98,7 +97,27 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for IndexScan<'a, T> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for IndexScan<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for IndexScan<'a, T> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + if let Some(iter) = self.iter.as_mut() { + match self.lookup { + IndexLookup::Static(_) => iter.rewind(self.op.limit), + IndexLookup::Probe if self.info.residual_predicate.is_none() => { + iter.reset_probe(arena.pop_runtime_probe(), self.op.limit) + } + IndexLookup::Probe => { + let ranges = Self::ranges(self.info, self.lookup, arena, plan_arena)?; + iter.reset(ranges, self.op.limit); + } + } + } + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -107,7 +126,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for IndexScan<'a, T> { let iter = match &mut self.iter { Some(iter) => iter, None => { - let ranges = self.ranges(arena, plan_arena)?; + let ranges = Self::ranges(self.info, self.lookup, arena, plan_arena)?; let TableScanOperator { table_name, columns, @@ -132,7 +151,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for IndexScan<'a, T> { }; let state = arena.local_state(plan_arena); - if iter.next_tuple_into(state.table_codec, &mut state.result.tuple)? { + if iter.next_tuple_into(state.table_codec, state.result.output_mut())? { arena.resume(); } else { arena.finish(); @@ -140,3 +159,119 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for IndexScan<'a, T> { Ok(()) } } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::db::test::build_table; + use crate::execution::{empty_context, ExecStatus}; + use crate::planner::{PlanArena, TableArenaCell}; + use crate::storage::memory::MemoryStorage; + use crate::storage::table_codec::TableCodec; + use crate::storage::{StatisticsMetaCache, Storage, TableCache, ViewCache}; + use crate::types::index::RuntimeIndexProbe; + use crate::types::tuple::Tuple; + use crate::types::value::DataValue; + use crate::types::LogicalType; + use std::ops::Bound; + + #[test] + fn reset_rebinds_probe_after_partial_read_eof_and_missing_key() -> Result<(), DatabaseError> { + let storage = MemoryStorage::new(); + let mut tx = storage.transaction()?; + let mut tables = TableCache::default(); + let cell = TableArenaCell::default(); + let mut meta = PlanArena::new(&cell); + build_table(&mut tables, &mut tx, &mut meta)?; + let meta = PlanArena::new(&cell); + let serializers = [ + LogicalType::Integer.serializable(), + LogicalType::Boolean.serializable(), + LogicalType::Integer.serializable(), + ]; + let mut codec = TableCodec::default(); + for id in 1..=4 { + tx.append_tuple( + &mut codec, + "t1", + &Tuple::new( + Some(DataValue::Int32(id)), + vec![ + DataValue::Int32(id), + DataValue::Boolean(true), + DataValue::Int32(id * 10), + ], + ) + .into(), + &serializers, + false, + )?; + } + let table = tables.get("t1").unwrap(); + let op = TableScanOperator { + table_name: "t1".into(), + columns: table.columns().copied().collect(), + limit: (None, None), + index_infos: vec![], + with_pk: true, + }; + let info = IndexInfo { + meta: *table.indexes().next().unwrap(), + sort_option: SortOption::None, + lookup: None, + residual_predicate: None, + covered_deserializers: None, + cover_mapping: None, + sort_elimination_hint: None, + stream_aggregate_hint: None, + }; + let views = ViewCache::default(); + let statistics = StatisticsMetaCache::default(); + let mut arena = ExecArena::with_capacity(0); + arena.init_context(empty_context(&tables, &views, &statistics), &tx); + let mut meta = ExecMetaArena::new(meta); + let mut scan = IndexScan::new(&op, &info, &IndexLookup::Probe); + arena.push_runtime_probe(RuntimeIndexProbe::Scope { + min: Bound::Included(DataValue::Int32(1)), + max: Bound::Included(DataValue::Int32(4)), + }); + scan.next_tuple(&mut arena, &mut meta)?; + assert_eq!(arena.result_tuple()[0], DataValue::Int32(1)); + let iterator = scan.iter.as_ref().unwrap() as *const _; + for (probe, expected) in [ + (RuntimeIndexProbe::Eq(DataValue::Int32(3)), vec![3]), + (RuntimeIndexProbe::Eq(DataValue::Int32(99)), vec![]), + ( + RuntimeIndexProbe::Scope { + min: Bound::Included(DataValue::Int32(1)), + max: Bound::Included(DataValue::Int32(4)), + }, + vec![1, 2, 3, 4], + ), + ( + RuntimeIndexProbe::Scope { + min: Bound::Excluded(DataValue::Int32(1)), + max: Bound::Excluded(DataValue::Int32(4)), + }, + vec![2, 3], + ), + (RuntimeIndexProbe::Eq(DataValue::Int32(1)), vec![1]), + ] { + arena.push_runtime_probe(probe); + scan.reset(&mut arena, &mut meta)?; + // An initialized scan is ready now, not at its next pull. + assert!(arena.runtime_probe_stack.is_empty()); + for id in expected { + scan.next_tuple(&mut arena, &mut meta)?; + assert_eq!(arena.result_tuple()[0], DataValue::Int32(id)); + assert_eq!(scan.iter.as_ref().unwrap() as *const _, iterator); + } + scan.next_tuple(&mut arena, &mut meta)?; + assert_eq!(arena.result.status, Some(ExecStatus::End)); + assert!(arena.runtime_probe_stack.is_empty()); + } + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/execution/dql/join/hash/full_join.rs b/src/execution/dql/join/hash/full_join.rs index ff9de82a..bf170efd 100644 --- a/src/execution/dql/join/hash/full_join.rs +++ b/src/execution/dql/join/hash/full_join.rs @@ -14,19 +14,19 @@ use crate::errors::DatabaseError; use crate::execution::dql::join::hash::{ - filter, JoinProbeState, LeftDropState, LeftDropTuples, ProbeState, + filter, join_rows, JoinProbeState, LeftDropState, LeftDropTuples, ProbeState, }; use crate::execution::dql::join::hash_join::BuildState; -use crate::execution::dql::join::RowBitmap; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::{SplitTupleRef, Tuple}; +use crate::types::bitset::BitSet; +use crate::types::tuple::{SplitTupleRef, TupleImpl}; use crate::types::value::DataValue; pub(crate) struct FullJoinState { pub(crate) left_schema_len: usize, pub(crate) right_schema_len: usize, - pub(crate) bits: RowBitmap, + pub(crate) bits: BitSet, } impl JoinProbeState for FullJoinState { @@ -36,7 +36,7 @@ impl JoinProbeState for FullJoinState { build_state: Option<&mut BuildState>, filter_expr: Option<&ExprRef>, plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { if probe_state.is_keys_has_null { if probe_state.emitted_unmatched { probe_state.finished = true; @@ -44,10 +44,10 @@ impl JoinProbeState for FullJoinState { } probe_state.emitted_unmatched = true; probe_state.finished = true; - return Ok(Some(Self::take_full_right_row( - self.left_schema_len, - &mut probe_state.probe_tuple, - ))); + probe_state + .probe_tuple + .prepend(std::iter::repeat_n(DataValue::Null, self.left_schema_len)); + return Ok(Some(std::mem::take(&mut probe_state.probe_tuple))); } let Some(build_state) = build_state else { @@ -57,44 +57,38 @@ impl JoinProbeState for FullJoinState { } probe_state.emitted_unmatched = true; probe_state.finished = true; - return Ok(Some(Self::take_full_right_row( - self.left_schema_len, - &mut probe_state.probe_tuple, - ))); + probe_state + .probe_tuple + .prepend(std::iter::repeat_n(DataValue::Null, self.left_schema_len)); + return Ok(Some(std::mem::take(&mut probe_state.probe_tuple))); }; while probe_state.index < build_state.tuples.len() { - let (i, Tuple { values, pk }) = &build_state.tuples[probe_state.index]; + let (i, build_tuple) = &build_state.tuples[probe_state.index]; probe_state.index += 1; if let Some(filter_expr) = filter_expr { - let full_values = - SplitTupleRef::from_slices(values, &probe_state.probe_tuple.values); + let full_values = SplitTupleRef::new(build_tuple, &probe_state.probe_tuple); if !filter(&full_values, filter_expr, plan_arena)? { continue; } } - let full_values = Vec::from_iter( - values - .iter() - .chain(probe_state.probe_tuple.values.iter()) - .cloned(), - ); self.bits.insert(*i); probe_state.produced = true; - return Ok(Some(Tuple::new( - pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(), - full_values, + return Ok(Some(join_rows( + build_tuple, + &mut probe_state.probe_tuple, + false, ))); } probe_state.finished = true; if !probe_state.produced && !probe_state.emitted_unmatched { probe_state.emitted_unmatched = true; - return Ok(Some(Self::take_full_right_row( - self.left_schema_len, - &mut probe_state.probe_tuple, - ))); + probe_state + .probe_tuple + .prepend(std::iter::repeat_n(DataValue::Null, self.left_schema_len)); + return Ok(Some(std::mem::take(&mut probe_state.probe_tuple))); } Ok(None) } @@ -104,7 +98,7 @@ impl JoinProbeState for FullJoinState { left_drop_state: &mut LeftDropState, _filter_expr: Option<&ExprRef>, _plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { let full_schema_len = self.right_schema_len + self.left_schema_len; loop { @@ -113,7 +107,7 @@ impl JoinProbeState for FullJoinState { if self.bits.contains(i) { continue; } - left_tuple.values.resize(full_schema_len, DataValue::Null); + left_tuple.resize(full_schema_len, DataValue::Null); return Ok(Some(left_tuple)); } left_drop_state.current = None; @@ -129,14 +123,3 @@ impl JoinProbeState for FullJoinState { } } } - -impl FullJoinState { - pub(crate) fn take_full_right_row(left_schema_len: usize, probe_tuple: &mut Tuple) -> Tuple { - let mut tuple = std::mem::take(probe_tuple); - tuple - .values - .resize(tuple.values.len() + left_schema_len, DataValue::Null); - tuple.values.rotate_right(left_schema_len); - tuple - } -} diff --git a/src/execution/dql/join/hash/inner_join.rs b/src/execution/dql/join/hash/inner_join.rs index 9df19866..10d168d1 100644 --- a/src/execution/dql/join/hash/inner_join.rs +++ b/src/execution/dql/join/hash/inner_join.rs @@ -13,11 +13,11 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::execution::dql::join::hash::{filter, JoinProbeState, ProbeState}; +use crate::execution::dql::join::hash::{filter, join_rows, JoinProbeState, ProbeState}; use crate::execution::dql::join::hash_join::BuildState; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::{SplitTupleRef, Tuple}; +use crate::types::tuple::{SplitTupleRef, TupleImpl}; pub(crate) struct InnerJoinState; @@ -28,7 +28,7 @@ impl JoinProbeState for InnerJoinState { build_state: Option<&mut BuildState>, filter_expr: Option<&ExprRef>, plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { if probe_state.is_keys_has_null { probe_state.finished = true; return Ok(None); @@ -40,25 +40,22 @@ impl JoinProbeState for InnerJoinState { }; while probe_state.index < build_state.tuples.len() { - let (_, Tuple { values, pk }) = &build_state.tuples[probe_state.index]; + let (_, build_tuple) = &build_state.tuples[probe_state.index]; probe_state.index += 1; - if let Some(filter_expr) = filter_expr { - let full_values = - SplitTupleRef::from_slices(values, &probe_state.probe_tuple.values); + let full_values = SplitTupleRef::new(build_tuple, &probe_state.probe_tuple); if !filter(&full_values, filter_expr, plan_arena)? { continue; } } - let full_values = Vec::from_iter( - values - .iter() - .chain(probe_state.probe_tuple.values.iter()) - .cloned(), - ); - return Ok(Some(Tuple::new( - pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(), - full_values, + let is_last = probe_state.index == build_state.tuples.len(); + if is_last { + probe_state.finished = true; + } + return Ok(Some(join_rows( + build_tuple, + &mut probe_state.probe_tuple, + is_last, ))); } diff --git a/src/execution/dql/join/hash/left_join.rs b/src/execution/dql/join/hash/left_join.rs index 5d4ebc82..5933732e 100644 --- a/src/execution/dql/join/hash/left_join.rs +++ b/src/execution/dql/join/hash/left_join.rs @@ -14,19 +14,19 @@ use crate::errors::DatabaseError; use crate::execution::dql::join::hash::{ - filter, JoinProbeState, LeftDropState, LeftDropTuples, ProbeState, + filter, join_rows, JoinProbeState, LeftDropState, LeftDropTuples, ProbeState, }; use crate::execution::dql::join::hash_join::BuildState; -use crate::execution::dql::join::RowBitmap; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::{SplitTupleRef, Tuple}; +use crate::types::bitset::BitSet; +use crate::types::tuple::{SplitTupleRef, TupleImpl}; use crate::types::value::DataValue; pub(crate) struct LeftJoinState { pub(crate) left_schema_len: usize, pub(crate) right_schema_len: usize, - pub(crate) bits: RowBitmap, + pub(crate) bits: BitSet, } impl JoinProbeState for LeftJoinState { @@ -36,7 +36,7 @@ impl JoinProbeState for LeftJoinState { build_state: Option<&mut BuildState>, filter_expr: Option<&ExprRef>, plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { if probe_state.is_keys_has_null { probe_state.finished = true; return Ok(None); @@ -48,26 +48,20 @@ impl JoinProbeState for LeftJoinState { }; while probe_state.index < build_state.tuples.len() { - let (i, Tuple { values, pk }) = &build_state.tuples[probe_state.index]; + let (i, build_tuple) = &build_state.tuples[probe_state.index]; probe_state.index += 1; if let Some(filter_expr) = filter_expr { - let full_values = - SplitTupleRef::from_slices(values, &probe_state.probe_tuple.values); + let full_values = SplitTupleRef::new(build_tuple, &probe_state.probe_tuple); if !filter(&full_values, filter_expr, plan_arena)? { continue; } } - let full_values = Vec::from_iter( - values - .iter() - .chain(probe_state.probe_tuple.values.iter()) - .cloned(), - ); self.bits.insert(*i); - return Ok(Some(Tuple::new( - pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(), - full_values, + return Ok(Some(join_rows( + build_tuple, + &mut probe_state.probe_tuple, + false, ))); } @@ -80,7 +74,7 @@ impl JoinProbeState for LeftJoinState { left_drop_state: &mut LeftDropState, _filter_expr: Option<&ExprRef>, _plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { let full_schema_len = self.right_schema_len + self.left_schema_len; loop { @@ -89,7 +83,7 @@ impl JoinProbeState for LeftJoinState { if self.bits.contains(i) { continue; } - left_tuple.values.resize(full_schema_len, DataValue::Null); + left_tuple.resize(full_schema_len, DataValue::Null); return Ok(Some(left_tuple)); } left_drop_state.current = None; diff --git a/src/execution/dql/join/hash/mod.rs b/src/execution/dql/join/hash/mod.rs index 6f22e876..a1f28364 100644 --- a/src/execution/dql/join/hash/mod.rs +++ b/src/execution/dql/join/hash/mod.rs @@ -26,13 +26,13 @@ use crate::execution::dql::join::hash_join::BuildState; use crate::execution::dql::sort::BumpVec; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::{Tuple, TupleLike}; +use crate::types::tuple::{Tuple, TupleImpl, TupleLike}; use crate::types::value::DataValue; use std::collections::hash_map::IntoIter as HashMapIntoIter; pub(crate) struct ProbeState { pub(crate) is_keys_has_null: bool, - pub(crate) probe_tuple: Tuple, + pub(crate) probe_tuple: TupleImpl, pub(crate) index: usize, pub(crate) produced: bool, pub(crate) finished: bool, @@ -45,7 +45,7 @@ pub(crate) struct LeftDropState { } pub(crate) struct LeftDropTuples { - pub(crate) tuples: std::vec::IntoIter<(usize, Tuple)>, + pub(crate) tuples: std::vec::IntoIter<(usize, TupleImpl)>, } pub(crate) trait JoinProbeState { @@ -55,14 +55,14 @@ pub(crate) trait JoinProbeState { build_state: Option<&mut BuildState>, filter_expr: Option<&ExprRef>, plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError>; + ) -> Result, DatabaseError>; fn left_drop_next( &mut self, _left_drop_state: &mut LeftDropState, _filter_expr: Option<&ExprRef>, _plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { Ok(None) } } @@ -81,7 +81,7 @@ impl JoinProbeState for JoinProbeStateImpl { build_state: Option<&mut BuildState>, filter_expr: Option<&ExprRef>, plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { match self { JoinProbeStateImpl::Inner(state) => { state.probe_next(probe_state, build_state, filter_expr, plan_arena) @@ -103,7 +103,7 @@ impl JoinProbeState for JoinProbeStateImpl { left_drop_state: &mut LeftDropState, filter_expr: Option<&ExprRef>, plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { match self { JoinProbeStateImpl::Inner(state) => { state.left_drop_next(left_drop_state, filter_expr, plan_arena) @@ -135,3 +135,21 @@ pub(crate) fn filter( _ => Err(DatabaseError::InvalidType), } } + +pub(crate) fn join_rows( + build: &TupleImpl, + probe: &mut TupleImpl, + consume_probe: bool, +) -> TupleImpl { + if consume_probe { + let pk = build.pk().cloned().or_else(|| probe.take_pk()); + let mut row = std::mem::take(probe); + row.prepend(build.iter().cloned()); + row.set_pk(pk); + return row; + } + let mut values = Vec::with_capacity(build.len() + probe.len()); + values.extend(build.iter().cloned()); + values.extend(probe.iter().cloned()); + Tuple::new(build.pk().or(probe.pk()).cloned(), values).into() +} diff --git a/src/execution/dql/join/hash/right_join.rs b/src/execution/dql/join/hash/right_join.rs index cee9034d..8f52a02e 100644 --- a/src/execution/dql/join/hash/right_join.rs +++ b/src/execution/dql/join/hash/right_join.rs @@ -13,12 +13,12 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::execution::dql::join::hash::full_join::FullJoinState; -use crate::execution::dql::join::hash::{filter, JoinProbeState, ProbeState}; +use crate::execution::dql::join::hash::{filter, join_rows, JoinProbeState, ProbeState}; use crate::execution::dql::join::hash_join::BuildState; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::{SplitTupleRef, Tuple}; +use crate::types::tuple::{SplitTupleRef, TupleImpl}; +use crate::types::value::DataValue; pub(crate) struct RightJoinState { pub(crate) left_schema_len: usize, @@ -31,7 +31,7 @@ impl JoinProbeState for RightJoinState { build_state: Option<&mut BuildState>, filter_expr: Option<&ExprRef>, plan_arena: &(dyn MetaArena + '_), - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { if probe_state.is_keys_has_null { if probe_state.emitted_unmatched { probe_state.finished = true; @@ -39,10 +39,10 @@ impl JoinProbeState for RightJoinState { } probe_state.emitted_unmatched = true; probe_state.finished = true; - return Ok(Some(FullJoinState::take_full_right_row( - self.left_schema_len, - &mut probe_state.probe_tuple, - ))); + probe_state + .probe_tuple + .prepend(std::iter::repeat_n(DataValue::Null, self.left_schema_len)); + return Ok(Some(std::mem::take(&mut probe_state.probe_tuple))); } let Some(build_state) = build_state else { @@ -52,43 +52,37 @@ impl JoinProbeState for RightJoinState { } probe_state.emitted_unmatched = true; probe_state.finished = true; - return Ok(Some(FullJoinState::take_full_right_row( - self.left_schema_len, - &mut probe_state.probe_tuple, - ))); + probe_state + .probe_tuple + .prepend(std::iter::repeat_n(DataValue::Null, self.left_schema_len)); + return Ok(Some(std::mem::take(&mut probe_state.probe_tuple))); }; while probe_state.index < build_state.tuples.len() { - let (_, Tuple { values, pk }) = &build_state.tuples[probe_state.index]; + let (_, build_tuple) = &build_state.tuples[probe_state.index]; probe_state.index += 1; if let Some(filter_expr) = filter_expr { - let full_values = - SplitTupleRef::from_slices(values, &probe_state.probe_tuple.values); + let full_values = SplitTupleRef::new(build_tuple, &probe_state.probe_tuple); if !filter(&full_values, filter_expr, plan_arena)? { continue; } } - let full_values = Vec::from_iter( - values - .iter() - .chain(probe_state.probe_tuple.values.iter()) - .cloned(), - ); probe_state.produced = true; - return Ok(Some(Tuple::new( - pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(), - full_values, + return Ok(Some(join_rows( + build_tuple, + &mut probe_state.probe_tuple, + false, ))); } if !probe_state.produced && !probe_state.emitted_unmatched { probe_state.emitted_unmatched = true; probe_state.finished = true; - return Ok(Some(FullJoinState::take_full_right_row( - self.left_schema_len, - &mut probe_state.probe_tuple, - ))); + probe_state + .probe_tuple + .prepend(std::iter::repeat_n(DataValue::Null, self.left_schema_len)); + return Ok(Some(std::mem::take(&mut probe_state.probe_tuple))); } probe_state.finished = true; diff --git a/src/execution/dql/join/hash_join.rs b/src/execution/dql/join/hash_join.rs index b34405e6..e50db76f 100644 --- a/src/execution/dql/join/hash_join.rs +++ b/src/execution/dql/join/hash_join.rs @@ -20,7 +20,6 @@ use crate::execution::dql::join::hash::right_join::RightJoinState; use crate::execution::dql::join::hash::{ JoinProbeState, JoinProbeStateImpl, LeftDropState, ProbeState, }; -use crate::execution::dql::join::RowBitmap; use crate::execution::dql::sort::BumpVec; use crate::execution::{ build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, @@ -29,8 +28,9 @@ use crate::planner::operator::join::{JoinCondition, JoinOperator, JoinType}; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::bitset::BitSet; +use crate::types::tuple::TupleImpl; use crate::types::value::DataValue; use bumpalo::Bump; use std::collections::HashMap; @@ -127,7 +127,7 @@ impl HashJoin { fn eval_keys( on_keys: &[ExprRef], - tuple: &Tuple, + tuple: &TupleImpl, build_buf: &mut BumpVec<'_, DataValue>, plan_arena: &(dyn MetaArena + '_), ) -> Result<(), DatabaseError> { @@ -143,7 +143,7 @@ impl HashJoin { Ok(()) } - fn initialize_build<'a, T: Transaction + 'a, A: MetaArena + 'a>( + fn initialize_build<'a, T: StorageTransaction + 'a, A: MetaArena + 'a>( &mut self, arena: &mut ExecArena<'a, T>, plan_arena: &mut ExecMetaArena, @@ -155,13 +155,15 @@ impl HashJoin { // build phase: // 1.construct hashtable, one hash key may contains multiple rows indices. // 2.merged all left tuples. + // TODO: std `RandomState` (SipHash) costs ~10% of an LMDB PK-sized hash join; a faster + // non-DoS-resistant hasher would do for build keys. #[allow(clippy::mutable_key_type)] let mut build_map = HashMap::new(); let mut build_buf = BumpVec::with_capacity_in(self.on_left_keys.len(), &self.bump); let mut build_count = 0usize; while arena.next_tuple(self.left_input, plan_arena)? { - let tuple = arena.materialize_tuple(); + let tuple = arena.take_tuple(); Self::eval_keys(&self.on_left_keys, &tuple, &mut build_buf, plan_arena)?; match build_map.get_mut(&build_buf) { @@ -207,25 +209,27 @@ impl HashJoin { JoinType::LeftOuter => JoinProbeStateImpl::Left(LeftJoinState { left_schema_len, right_schema_len, - bits: RowBitmap::with_capacity(build_count), + bits: BitSet::with_capacity(build_count), }), JoinType::RightOuter => JoinProbeStateImpl::Right(RightJoinState { left_schema_len }), JoinType::Full => JoinProbeStateImpl::Full(FullJoinState { left_schema_len, right_schema_len, - bits: RowBitmap::with_capacity(build_count), + bits: BitSet::with_capacity(build_count), }), JoinType::Cross => unreachable!(), } } } +// TODO: dropping the build map (one heap `Vec` and owned `Tuple`s per key) costs ~13.5% of an +// LMDB PK-sized hash join plus allocator consolidation; allocate build rows in `bump` instead. #[derive(Default, Debug)] pub(crate) struct BuildState { - pub(crate) tuples: Vec<(usize, Tuple)>, + pub(crate) tuples: Vec<(usize, TupleImpl)>, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for HashJoin { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for HashJoin { type Input = (&'a JoinOperator, &'a LogicalPlan, &'a LogicalPlan); fn into_executor( @@ -250,7 +254,18 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for HashJoin { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashJoin { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for HashJoin { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.state = HashJoinState::Build; + self.bump.reset(); + arena.reset(self.left_input, plan_arena)?; + arena.reset(self.right_input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -277,7 +292,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashJoin { if !arena.next_tuple(self.right_input, plan_arena)? { break true; } - let tuple = arena.materialize_tuple(); + let tuple = arena.take_tuple(); Self::eval_keys( &self.on_right_keys, &tuple, @@ -318,7 +333,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashJoin { probe_buf, probe_state, }; - arena.produce_tuple(tuple); + arena.produce(tuple); return Ok(()); } @@ -349,7 +364,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashJoin { join_impl, left_drop, }; - arena.produce_tuple(tuple); + arena.produce(tuple); return Ok(()); } state = HashJoinState::End; @@ -393,9 +408,9 @@ mod test { ) -> Result { HepOptimizerPipeline::builder() .before_batch( - "Expression Remapper".to_string(), + "Expression Remapper", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::EvaluatorBind], + &[NormalizationRuleImpl::EvaluatorBind], ) .build() .instantiate(plan) @@ -527,15 +542,15 @@ mod test { assert_eq!(tuples.len(), 3); assert_eq!( - tuples[0].values, + tuples[0], build_integers(vec![Some(0), Some(2), Some(4), Some(0), Some(2), Some(4)]) ); assert_eq!( - tuples[1].values, + tuples[1], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(3), Some(5)]) ); assert_eq!( - tuples[2].values, + tuples[2], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(1), Some(1)]) ); @@ -588,19 +603,19 @@ mod test { assert_eq!(tuples.len(), 4); assert_eq!( - tuples[0].values, + tuples[0], build_integers(vec![Some(0), Some(2), Some(4), Some(0), Some(2), Some(4)]) ); assert_eq!( - tuples[1].values, + tuples[1], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(3), Some(5)]) ); assert_eq!( - tuples[2].values, + tuples[2], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(1), Some(1)]) ); assert_eq!( - tuples[3].values, + tuples[3], build_integers(vec![Some(3), Some(5), Some(7), None, None, None]) ); } @@ -654,19 +669,19 @@ mod test { assert_eq!(tuples.len(), 4); assert_eq!( - tuples[0].values, + tuples[0], build_integers(vec![Some(0), Some(2), Some(4), Some(0), Some(2), Some(4)]) ); assert_eq!( - tuples[1].values, + tuples[1], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(3), Some(5)]) ); assert_eq!( - tuples[2].values, + tuples[2], build_integers(vec![None, None, None, Some(4), Some(6), Some(8)]) ); assert_eq!( - tuples[3].values, + tuples[3], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(1), Some(1)]) ); @@ -763,7 +778,7 @@ mod test { assert_eq!(tuples.len(), 1); assert_eq!( - tuples[0].values, + tuples[0], vec![ DataValue::Int32(2), DataValue::Int32(5), @@ -820,23 +835,23 @@ mod test { assert_eq!(tuples.len(), 5); assert_eq!( - tuples[0].values, + tuples[0], build_integers(vec![Some(0), Some(2), Some(4), Some(0), Some(2), Some(4)]) ); assert_eq!( - tuples[1].values, + tuples[1], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(3), Some(5)]) ); assert_eq!( - tuples[2].values, + tuples[2], build_integers(vec![None, None, None, Some(4), Some(6), Some(8)]) ); assert_eq!( - tuples[3].values, + tuples[3], build_integers(vec![Some(1), Some(3), Some(5), Some(1), Some(1), Some(1)]) ); assert_eq!( - tuples[4].values, + tuples[4], build_integers(vec![Some(3), Some(5), Some(7), None, None, None]) ); diff --git a/src/execution/dql/join/index_probe_join.rs b/src/execution/dql/join/index_probe_join.rs new file mode 100644 index 00000000..fd7862db --- /dev/null +++ b/src/execution/dql/join/index_probe_join.rs @@ -0,0 +1,284 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::execution::dql::join::predicates_matched; +use crate::execution::{ + build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, +}; +use crate::planner::operator::index_probe_join::IndexProbeJoinOperator; +use crate::planner::{ExecMetaArena, LogicalPlan, MetaArena}; +use crate::storage::StorageTransaction; +use crate::types::index::RuntimeIndexProbe; + +pub struct IndexProbeJoin<'a> { + op: &'a IndexProbeJoinOperator, + left_input: ExecId, + right_input: ExecId, +} + +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for IndexProbeJoin<'a> { + type Input = (&'a IndexProbeJoinOperator, &'a LogicalPlan, &'a LogicalPlan); + + fn into_executor( + (op, left_input, right_input): Self::Input, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut (dyn MetaArena + 'a), + cache: ExecutionContext<'_>, + transaction: &T, + ) -> ExecId { + let left_input = build_read(arena, plan_arena, left_input, cache, transaction); + let right_input = build_read(arena, plan_arena, right_input, cache, transaction); + arena.push(ExecNode::IndexProbeJoin(Self { + op, + left_input, + right_input, + })) + } +} + +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for IndexProbeJoin<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + arena.reset(self.left_input, plan_arena) + } + + fn next_tuple( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + while arena.next_tuple(self.left_input, plan_arena)? { + let mut left = arena.take_tuple(); + let key = plan_arena + .expression(self.op.probe) + .eval(plan_arena, Some(&left))? + .into_owned(); + if key.has_null() { + arena.release(left); + continue; + } + arena.push_runtime_probe(RuntimeIndexProbe::Eq(key)); + arena.reset(self.right_input, plan_arena)?; + if arena.next_tuple(self.right_input, plan_arena)? + && predicates_matched(&self.op.predicates, &left, arena.result_tuple(), plan_arena)? + { + let mut right = arena.take_tuple(); + let pk = left.take_pk().or_else(|| right.take_pk()); + let joined = arena.concat(left, right, pk); + arena.produce(joined); + return Ok(()); + } + arena.release(left); + } + arena.finish(); + Ok(()) + } +} + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; + use crate::execution::{execute_input, ExecArena}; + use crate::expression::{BinaryOperator, ScalarExpression}; + use crate::planner::operator::values::ValuesOperator; + use crate::planner::operator::Operator; + use crate::planner::test::PlanArenaTestExt; + use crate::planner::{Childrens, ExprRef, PlanArena, TableArenaCell}; + use crate::storage::rocksdb::RocksStorage; + use crate::storage::{StatisticsMetaCache, Storage, TableCache, ViewCache}; + use crate::types::evaluator::binary_create; + use crate::types::value::DataValue; + use crate::types::{CharLengthUnits, LogicalType}; + use std::borrow::Cow; + use tempfile::TempDir; + + fn build_values( + arena: &mut PlanArena, + columns: Vec<(&str, LogicalType)>, + rows: Vec>, + ) -> LogicalPlan { + let schema_ref = columns + .into_iter() + .map(|(name, ty)| { + arena.alloc_column(ColumnCatalog::new( + name.to_string(), + true, + ColumnDesc::new(ty, None, true, None).unwrap(), + )) + }) + .collect(); + LogicalPlan::new( + Operator::Values(ValuesOperator::new( + arena.alloc_expression_rows(&rows), + rows.len(), + schema_ref, + )), + Childrens::None, + ) + } + + fn build_equality( + arena: &mut PlanArena, + left: (ColumnRef, usize), + right: (ColumnRef, usize), + ) -> Result { + let left_expr = arena.alloc_expression(ScalarExpression::column_expr(left.0, left.1)); + let right_expr = arena.alloc_expression(ScalarExpression::column_expr(right.0, right.1)); + Ok(arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Eq, + left_expr, + right_expr, + evaluator: Some(binary_create( + Cow::Owned(LogicalType::Integer), + BinaryOperator::Eq, + )?), + ty: LogicalType::Boolean, + })) + } + + fn build_test_storage() -> Result< + ( + TableCache, + ViewCache, + StatisticsMetaCache, + TempDir, + RocksStorage, + ), + DatabaseError, + > { + let temp_dir = TempDir::new().expect("unable to create temporary working directory"); + let storage = RocksStorage::new(temp_dir.path())?; + Ok(( + TableCache::default(), + ViewCache::default(), + StatisticsMetaCache::default(), + temp_dir, + storage, + )) + } + + // Values ignores runtime probes, so it stands in for the unique lookup and the pushed + // probes stay on the stack for inspection. + #[test] + fn index_probe_join_moves_single_match_and_reuses_right_slots() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut plan_arena = PlanArena::new(&table_arena); + let mut left = build_values( + &mut plan_arena, + vec![ + ("left_key", LogicalType::Integer), + ( + "payload", + LogicalType::Varchar(None, CharLengthUnits::Characters), + ), + ], + vec![ + vec![DataValue::Int32(2), DataValue::from("first".to_string())], + vec![DataValue::Null, DataValue::from("null".to_string())], + vec![DataValue::Int32(99), DataValue::from("missing".to_string())], + vec![DataValue::Int32(2), DataValue::from("again".to_string())], + ], + ); + let mut right = build_values( + &mut plan_arena, + vec![ + ("right_key", LogicalType::Integer), + ("flag", LogicalType::Boolean), + ], + vec![vec![DataValue::Int32(2), DataValue::Boolean(true)]], + ); + let left_column = left.output_schema(&mut plan_arena)[0]; + let right_schema = right.output_schema(&mut plan_arena).clone(); + let equality = build_equality(&mut plan_arena, (left_column, 0), (right_schema[0], 2))?; + let residual = + plan_arena.alloc_expression(ScalarExpression::column_expr(right_schema[1], 3)); + let probe = plan_arena.alloc_expression(ScalarExpression::column_expr(left_column, 0)); + let op = IndexProbeJoinOperator::new(probe, vec![equality, residual]); + let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; + let transaction = storage.transaction()?; + let cache = crate::execution::empty_context(&table_cache, &view_cache, &meta_cache); + let mut arena = ExecArena::with_capacity(0); + arena.init_context(cache, &transaction); + let root = >::into_executor( + (&op, &left, &right), + &mut arena, + &mut plan_arena, + cache, + &transaction, + ); + let address = arena.nodes.items.as_ptr(); + let count = arena.nodes.items.len(); + assert_eq!(count, 3); + let mut plan_arena = ExecMetaArena::new(plan_arena); + for payload in ["first", "again"] { + assert!(arena.next_tuple(root, &mut plan_arena)?); + assert_eq!( + arena.take_tuple(), + vec![ + DataValue::Int32(2), + DataValue::from(payload.to_string()), + DataValue::Int32(2), + DataValue::Boolean(true) + ] + ); + assert_eq!(arena.nodes.items.as_ptr(), address); + assert_eq!(arena.nodes.items.len(), count); + } + assert!(!arena.next_tuple(root, &mut plan_arena)?); + // The NULL key is skipped without probing. + assert_eq!( + arena.runtime_probe_stack, + [2, 99, 2].map(|key| RuntimeIndexProbe::Eq(DataValue::Int32(key))) + ); + Ok(()) + } + + #[test] + fn index_probe_join_skips_a_rejected_unique_match() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut plan_arena = PlanArena::new(&table_arena); + let mut left = build_values( + &mut plan_arena, + vec![("left_key", LogicalType::Integer)], + vec![vec![DataValue::Int32(1)]], + ); + let mut right = build_values( + &mut plan_arena, + vec![("flag", LogicalType::Boolean)], + vec![vec![DataValue::Boolean(false)]], + ); + let left_column = left.output_schema(&mut plan_arena)[0]; + let right_column = right.output_schema(&mut plan_arena)[0]; + let predicate = plan_arena.alloc_expression(ScalarExpression::column_expr(right_column, 1)); + let probe = plan_arena.alloc_expression(ScalarExpression::column_expr(left_column, 0)); + let op = IndexProbeJoinOperator::new(probe, vec![predicate]); + let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; + let transaction = storage.transaction()?; + let mut executor = execute_input::<_, IndexProbeJoin>( + (&op, &left, &right), + crate::execution::empty_context(&table_cache, &view_cache, &meta_cache), + plan_arena, + &transaction, + ); + assert!(executor.next_tuple()?.is_none()); + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/execution/dql/join/mod.rs b/src/execution/dql/join/mod.rs index c72cae3f..4d61c41f 100644 --- a/src/execution/dql/join/mod.rs +++ b/src/execution/dql/join/mod.rs @@ -12,46 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::errors::DatabaseError; use crate::planner::operator::join::JoinType; +use crate::planner::{ExprRef, MetaArena}; +use crate::types::tuple::{SplitTupleRef, TupleImpl}; +use crate::types::value::DataValue; mod hash; pub(crate) mod hash_join; +pub(crate) mod index_probe_join; pub(crate) mod nested_loop_join; -pub(crate) struct RowBitmap { - blocks: Vec, - len: usize, -} - -impl RowBitmap { - pub(crate) fn with_capacity(len: usize) -> Self { - let bits_per_block = usize::BITS as usize; - let blocks = len.div_ceil(bits_per_block); - RowBitmap { - blocks: vec![0; blocks], - len, - } - } - - pub(crate) fn insert(&mut self, index: usize) { - assert!( - index < self.len, - "insert at index {index} exceeds bitmap size {}", - self.len - ); - let bits_per_block = usize::BITS as usize; - self.blocks[index / bits_per_block] |= 1usize << (index % bits_per_block); - } - - pub(crate) fn contains(&self, index: usize) -> bool { - if index >= self.len { - return false; - } - let bits_per_block = usize::BITS as usize; - (self.blocks[index / bits_per_block] & (1usize << (index % bits_per_block))) != 0 - } -} - pub fn joins_nullable(join_type: &JoinType) -> (bool, bool) { match join_type { JoinType::Inner => (false, false), @@ -62,60 +33,23 @@ pub fn joins_nullable(join_type: &JoinType) -> (bool, bool) { } } -// GRCOV_EXCL_START -#[cfg(test)] -mod tests { - use super::RowBitmap; - - #[test] - fn row_bitmap_contains_false_for_empty_bitmap() { - let bitmap = RowBitmap::with_capacity(0); - - assert!(!bitmap.contains(0)); - assert!(!bitmap.contains(usize::MAX)); - } - - #[test] - fn row_bitmap_contains_false_for_unset_bits_within_capacity() { - let bitmap = RowBitmap::with_capacity(usize::BITS as usize + 1); - - assert!(!bitmap.contains(0)); - assert!(!bitmap.contains(usize::BITS as usize)); - } - - #[test] - fn row_bitmap_insert_marks_bits_across_blocks() { - let boundary = usize::BITS as usize; - let mut bitmap = RowBitmap::with_capacity(boundary + 2); - - bitmap.insert(0); - bitmap.insert(boundary - 1); - bitmap.insert(boundary); - bitmap.insert(boundary + 1); - - assert!(bitmap.contains(0)); - assert!(bitmap.contains(boundary - 1)); - assert!(bitmap.contains(boundary)); - assert!(bitmap.contains(boundary + 1)); - assert!(!bitmap.contains(1)); - } - - #[test] - fn row_bitmap_contains_returns_false_past_capacity() { - let mut bitmap = RowBitmap::with_capacity(2); - - bitmap.insert(1); - - assert!(!bitmap.contains(2)); - assert!(!bitmap.contains(usize::MAX)); - } - - #[test] - #[should_panic(expected = "exceeds bitmap size")] - fn row_bitmap_insert_panics_past_capacity() { - let mut bitmap = RowBitmap::with_capacity(1); - - bitmap.insert(1); +/// Evaluates join predicates bound to the `left ++ right` layout; NULL counts as no match. +pub(crate) fn predicates_matched( + predicates: &[ExprRef], + left_tuple: &TupleImpl, + right_tuple: &TupleImpl, + plan_arena: &(dyn MetaArena + '_), +) -> Result { + let values = SplitTupleRef::new(left_tuple, right_tuple); + for predicate in predicates { + match *plan_arena + .expression(*predicate) + .eval(plan_arena, Some(&values))? + { + DataValue::Boolean(true) => {} + DataValue::Boolean(false) | DataValue::Null => return Ok(false), + _ => return Err(DatabaseError::InvalidType), + } } + Ok(true) } -// GRCOV_EXCL_STOP diff --git a/src/execution/dql/join/nested_loop_join.rs b/src/execution/dql/join/nested_loop_join.rs index e31a2d7e..7b3c5bea 100644 --- a/src/execution/dql/join/nested_loop_join.rs +++ b/src/execution/dql/join/nested_loop_join.rs @@ -19,14 +19,14 @@ use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; use crate::errors::DatabaseError; -use crate::execution::dql::join::RowBitmap; use crate::execution::{ build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, }; use crate::planner::operator::join::{JoinCondition, JoinOperator, JoinType}; use crate::planner::ExprRef; -use crate::storage::Transaction; -use crate::types::tuple::{SplitTupleRef, Tuple}; +use crate::storage::StorageTransaction; +use crate::types::bitset::BitSet; +use crate::types::tuple::{SplitTupleRef, TupleImpl, TupleLike}; use crate::types::value::DataValue; /// Equivalent condition @@ -41,10 +41,10 @@ impl EqualCondition { /// Compare left tuple and right tuple on equivalent condition /// `left_tuple` must be from the [`NestedLoopJoin::left_input`] /// `right_tuple` must be from the [`NestedLoopJoin::right_input`] - fn equals( + fn equals( &self, - left_tuple: &Tuple, - right_tuple: &Tuple, + left_tuple: &L, + right_tuple: &R, arena: &(dyn MetaArena + '_), ) -> Result { if self.on_left_keys.is_empty() { @@ -74,41 +74,40 @@ impl EqualCondition { /// | Right/RightSemi/RightAnti/Full | left | right | /// |--------------------------------|----------------|----------------| /// | Full | left | right | -pub struct NestedLoopJoin<'a> { - right_input_plan: &'a LogicalPlan, +pub struct NestedLoopJoin { ty: JoinType, filter: Option, eq_cond: EqualCondition, left_input: ExecId, - right_pos: ExecId, + right_input: ExecId, state: NestedLoopJoinState, } enum NestedLoopJoinState { PullLeft { - right_bitmap: Option, + right_bitmap: Option, }, ScanRight { active_left: ActiveLeftState, - right_bitmap: Option, + right_bitmap: Option, }, EmitRightUnmatched { right_input: ExecId, - right_bitmap: Option, + right_bitmap: Option, right_emit_index: usize, }, End, } struct ActiveLeftState { - left_tuple: Tuple, + left_tuple: TupleImpl, right_input: ExecId, right_index: usize, has_matched: bool, first_matches: Vec, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for NestedLoopJoin<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for NestedLoopJoin { type Input = (&'a JoinOperator, &'a LogicalPlan, &'a LogicalPlan); fn into_executor( @@ -139,21 +138,29 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for NestedLoopJoin<'a> { }; let left_input = build_read(arena, plan_arena, left_plan, cache, transaction); - let right_pos = arena.nodes.position(); - build_read(arena, plan_arena, right_plan, cache, transaction); + let right_input = build_read(arena, plan_arena, right_plan, cache, transaction); arena.push(ExecNode::NestedLoopJoin(NestedLoopJoin { - right_input_plan: right_plan, ty: join_type, filter, eq_cond, left_input, - right_pos, + right_input, state: NestedLoopJoinState::PullLeft { right_bitmap: None }, })) } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.state = NestedLoopJoinState::PullLeft { right_bitmap: None }; + arena.reset(self.left_input, plan_arena)?; + arena.reset(self.right_input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -167,7 +174,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { if !arena.next_tuple(self.left_input, plan_arena)? { if matches!(self.ty, JoinType::Full) { state = NestedLoopJoinState::EmitRightUnmatched { - right_input: self.build_right_input(arena, plan_arena), + right_input: self.reset_right_input(arena, plan_arena)?, right_bitmap, right_emit_index: 0, }; @@ -177,12 +184,12 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { arena.finish(); return Ok(()); } - let left_tuple = arena.materialize_tuple(); + let left_tuple = arena.take_tuple(); state = NestedLoopJoinState::ScanRight { active_left: ActiveLeftState { left_tuple, - right_input: self.build_right_input(arena, plan_arena), + right_input: self.reset_right_input(arena, plan_arena)?, right_index: 0, has_matched: false, first_matches: Vec::new(), @@ -223,11 +230,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { active_left.has_matched = true; // A matched pair is always one output row, even when // neither side projects a column (e.g. `count(*)`). - Self::emit_tuple( - &active_left.left_tuple, - arena.result_tuple_mut(), - self.ty, - ); + Self::emit_tuple(&active_left.left_tuple, arena.result_mut(), self.ty); if matches!(self.ty, JoinType::Full) { if let Some(bits) = right_bitmap.as_mut() { bits.insert(idx); @@ -250,7 +253,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { bits.insert(idx); } } else { - let mut bits = RowBitmap::with_capacity(active_left.right_index); + let mut bits = BitSet::with_capacity(active_left.right_index); for idx in active_left.first_matches { bits.insert(idx); } @@ -262,11 +265,14 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { JoinType::LeftOuter | JoinType::RightOuter | JoinType::Full if !active_left.has_matched => { - let right_tuple = arena.result_tuple_mut(); - right_tuple.pk = None; - right_tuple.values.clear(); - right_tuple.values.resize(right_schema_len, DataValue::Null); - Self::emit_tuple(&active_left.left_tuple, right_tuple, self.ty); + let mut row = std::mem::take(&mut active_left.left_tuple); + let nulls = std::iter::repeat_n(DataValue::Null, right_schema_len); + if matches!(self.ty, JoinType::RightOuter) { + row.prepend(nulls); + } else { + row.extend(nulls); + } + *arena.result_mut() = row; true } _ => false, @@ -294,9 +300,10 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { }; if is_unmatched { - let values = &mut arena.result_tuple_mut().values; - values.resize(values.len() + self.eq_cond.left_len, DataValue::Null); - values.rotate_right(self.eq_cond.left_len); + arena.result_mut().prepend(std::iter::repeat_n( + DataValue::Null, + self.eq_cond.left_len, + )); self.state = NestedLoopJoinState::EmitRightUnmatched { right_input, right_bitmap, @@ -321,37 +328,31 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for NestedLoopJoin<'a> { } } -impl<'a> NestedLoopJoin<'a> { - fn build_right_input( +impl NestedLoopJoin { + fn reset_right_input<'a, T: StorageTransaction + 'a, A: MetaArena + 'a>( &mut self, arena: &mut ExecArena<'a, T>, - plan_arena: &mut (dyn MetaArena + 'a), - ) -> ExecId { - let cache = arena.context(); - let transaction = arena.transaction(); - // The same right-hand plan rebuilds the same slots, including nested joins. - arena.nodes.seek(self.right_pos); - build_read(arena, plan_arena, self.right_input_plan, cache, transaction) + plan_arena: &mut ExecMetaArena, + ) -> Result { + arena.reset(self.right_input, plan_arena)?; + Ok(self.right_input) } /// Emit a tuple according to the join type. /// - /// `left_tuple`: retained outer tuple (logical right for RightOuter). - /// `right_tuple`: current inner tuple (logical left for RightOuter), rewritten in place. + /// `left`: retained outer tuple (logical right for RightOuter). + /// `right`: current inner tuple (logical left for RightOuter). /// `ty`: the type of join - fn emit_tuple(left_tuple: &Tuple, right_tuple: &mut Tuple, ty: JoinType) { - let right_len = right_tuple.values.len(); - right_tuple.values.extend(left_tuple.values.iter().cloned()); - if matches!(ty, JoinType::RightOuter) { - if right_tuple.pk.is_none() { - right_tuple.pk = left_tuple.pk.clone(); - } + fn emit_tuple(left: &TupleImpl, right: &mut TupleImpl, ty: JoinType) { + let values = left.iter().cloned(); + let pk = if matches!(ty, JoinType::RightOuter) { + right.extend(values); + right.take_pk().or_else(|| left.pk().cloned()) } else { - right_tuple.values.rotate_left(right_len); - if left_tuple.pk.is_some() { - right_tuple.pk = left_tuple.pk.clone(); - } - } + right.prepend(values); + left.pk().cloned().or_else(|| right.take_pk()) + }; + right.set_pk(pk); } } @@ -386,18 +387,17 @@ mod test { ) -> Result { HepOptimizerPipeline::builder() .before_batch( - "Expression Remapper".to_string(), + "Expression Remapper", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::EvaluatorBind], + &[NormalizationRuleImpl::EvaluatorBind], ) .build() .instantiate(plan) .find_best(None, arena) } - fn tuple_to_strings(tuple: &Tuple) -> Vec> { + fn tuple_to_strings(tuple: &TupleImpl) -> Vec> { tuple - .values .iter() .map(|value| match value { DataValue::Null => None, @@ -525,12 +525,11 @@ mod test { (on_keys, values_t1, values_t2, filter) } - fn valid_result(expected: &mut HashSet>, actual: &[Tuple]) { + fn valid_result(expected: &mut HashSet>, actual: &[TupleImpl]) { assert_eq!(actual.len(), expected.len()); for tuple in actual { let values = tuple - .values .iter() .map(|v| { if matches!(v, DataValue::Null) { @@ -583,7 +582,7 @@ mod test { let mut plan_arena = ExecMetaArena::new(plan_arena); while arena.next_tuple(root, &mut plan_arena)? { rows += 1; - assert_eq!(arena.result_tuple().values.len(), 9); + assert_eq!(arena.result_tuple().len(), 9); assert_eq!(arena.nodes.items.len(), count); assert_eq!(arena.nodes.items.as_ptr(), address); } @@ -691,7 +690,7 @@ mod test { let tuples = try_collect(executor)?; assert_eq!( - tuples[0].values, + tuples[0], build_integers(vec![Some(0), Some(2), Some(4), None, None, None]) ); @@ -1029,7 +1028,7 @@ mod test { let tuples = try_collect(executor)?; assert_eq!( - tuples[0].values, + tuples[0], build_integers(vec![Some(0), Some(2), Some(4), None, None, None]) ); @@ -1190,7 +1189,7 @@ mod test { assert_eq!(tuples.len(), 1); assert_eq!( - tuples[0].values, + tuples[0], vec![ DataValue::Int32(2), DataValue::Int32(5), diff --git a/src/execution/dql/limit.rs b/src/execution/dql/limit.rs index a49a0665..1324f73f 100644 --- a/src/execution/dql/limit.rs +++ b/src/execution/dql/limit.rs @@ -19,7 +19,7 @@ use crate::execution::{ use crate::planner::operator::limit::LimitOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct Limit { offset: Option, limit: Option, @@ -28,7 +28,7 @@ pub struct Limit { emitted: usize, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Limit { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Limit { type Input = (&'a LimitOperator, &'a LogicalPlan); fn into_executor( @@ -49,7 +49,17 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Limit { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Limit { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Limit { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.skipped = 0; + self.emitted = 0; + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/dql/mark_apply.rs b/src/execution/dql/mark_apply.rs index 761c1335..4b889bd5 100644 --- a/src/execution/dql/mark_apply.rs +++ b/src/execution/dql/mark_apply.rs @@ -13,35 +13,98 @@ // limitations under the License. use crate::errors::DatabaseError; +use crate::execution::dql::join::predicates_matched; +use crate::execution::runs::{Runs, RunsBuilder, RunsCursor}; +#[cfg(feature = "spill")] +use crate::execution::spill::runs::SpillRuns; use crate::execution::{ build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, }; +use crate::expression::{BinaryOperator, ScalarExpression, TypeCast}; use crate::planner::operator::mark_apply::{MarkApplyKind, MarkApplyOperator, MarkApplyQuantifier}; -use crate::planner::LogicalPlan; -use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::planner::{ExecMetaArena, ExprRef, LogicalPlan, MetaArena}; +use crate::storage::StorageTransaction; +use crate::types::evaluator::BinaryEvaluatorRef; use crate::types::index::RuntimeIndexProbe; -use crate::types::tuple::{SplitTupleRef, Tuple}; +use crate::types::tuple::{SplitTupleRef, TupleImpl}; use crate::types::value::DataValue; +use std::borrow::Cow; +use std::cmp::Ordering; + +#[cfg(feature = "spill")] +type KeysBuilder = SpillRuns; +#[cfg(not(feature = "spill"))] +type KeysBuilder = crate::execution::runs::SortedVec; + +// The binder coerces both comparison sides to one type, so non-NULL keys of +// one apply share a variant and are totally ordered. +fn compare_keys(left: &DataValue, right: &DataValue) -> Ordering { + left.partial_cmp(right).unwrap_or(Ordering::Equal) +} + +/// Sorted non-NULL RHS comparison keys; NULL keys only decide UNKNOWN results. +struct MaterializedKeys { + runs: ::Runs, + is_empty: bool, + has_null: bool, + // Spilled reads size their buffer to one key when the key type is fixed-width. + fixed_width: Option, +} -#[derive(PartialEq, Eq)] -enum QuantifiedPredicateOutcome { - True, - False, - Null, - Skip, +// TODO: range ANY/ALL only need the RHS extrema rather than every key. +enum RightInput { + Streaming { root: ExecId }, + Materialize { root: ExecId }, + Materialized(MaterializedKeys), +} + +impl RightInput { + fn reset<'a, T: StorageTransaction + 'a, A: MetaArena + 'a>( + &self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + probe: Option, + ) -> Result { + let (Self::Streaming { root } | Self::Materialize { root }) = self else { + unreachable!("materialized input looks up its keys") + }; + if let Some(probe) = probe { + arena.push_runtime_probe(probe); + } + arena.reset(*root, plan_arena)?; + Ok(*root) + } + + // Return the first terminal predicate result, or None after exhausting the RHS. + fn scan<'a, T: StorageTransaction + 'a, A: MetaArena + 'a>( + &self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + probe: Option, + mut evaluate: impl FnMut( + &TupleImpl, + &ExecMetaArena, + ) -> Result, DatabaseError>, + ) -> Result, DatabaseError> { + let root = self.reset(arena, plan_arena, probe)?; + while arena.next_tuple(root, plan_arena)? { + if let Some(value) = evaluate(arena.result_tuple(), plan_arena)? { + return Ok(Some(value)); + } + } + Ok(None) + } } pub struct MarkApply<'a> { op: &'a MarkApplyOperator, - right_input_plan: &'a LogicalPlan, left_input: ExecId, - right_pos: ExecId, - // Retain a streaming inner input across next_tuple calls, not its result rows. - join_input: Option<(ExecId, Tuple)>, + right_root: ExecId, + // Temporarily None while scan owns and transitions the right input. + right_input: Option, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for MarkApply<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for MarkApply<'a> { type Input = (&'a MarkApplyOperator, &'a LogicalPlan, &'a LogicalPlan); fn into_executor( @@ -52,77 +115,80 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for MarkApply<'a> { transaction: &T, ) -> ExecId { let left_input = build_read(arena, plan_arena, left_input, cache, transaction); - let right_pos = arena.nodes.position(); - build_read(arena, plan_arena, right_input, cache, transaction); + let right_root = build_read(arena, plan_arena, right_input, cache, transaction); + let right_input = if op.materialize { + RightInput::Materialize { root: right_root } + } else { + RightInput::Streaming { root: right_root } + }; arena.push(ExecNode::MarkApply(Self { op, - right_input_plan: right_input, left_input, - right_pos, - join_input: None, + right_root, + right_input: Some(right_input), })) } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for MarkApply<'a> { - fn next_tuple( +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for MarkApply<'a> { + fn reset( &mut self, arena: &mut ExecArena<'a, T>, plan_arena: &mut ExecMetaArena, ) -> Result<(), DatabaseError> { - if matches!(self.op.kind, MarkApplyKind::InnerJoin) { - return self.next_join_tuple(arena, plan_arena); - } - if !arena.next_tuple(self.left_input, plan_arena)? { - arena.finish(); - return Ok(()); + // Materialization already makes the RHS reusable across left rows. + // Keep both complete caches and paused prefixes; only reset the left input. + if !self.op.materialize || self.right_input.is_none() { + self.right_input = Some(if self.op.materialize { + RightInput::Materialize { + root: self.right_root, + } + } else { + RightInput::Streaming { + root: self.right_root, + } + }); } - - let mut left_tuple = arena.materialize_tuple(); - let marker = self.mark_value(arena, plan_arena, &left_tuple)?; - - left_tuple.values.push(marker); - arena.produce_tuple(left_tuple); - Ok(()) + arena.reset(self.left_input, plan_arena) } -} -impl<'a> MarkApply<'a> { - fn next_join_tuple( + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, plan_arena: &mut ExecMetaArena, ) -> Result<(), DatabaseError> { - loop { - if let Some((root, left)) = &mut self.join_input { - while arena.next_tuple(*root, plan_arena)? { - let right = arena.result_tuple(); - if Self::predicates_matched(self.op.predicates(), left, right, plan_arena)? { - let output = arena.result_tuple_mut(); - let right_len = output.values.len(); - output.values.extend(left.values.iter().cloned()); - output.values.rotate_left(right_len); - output.pk = left.pk.clone().or(output.pk.take()); - arena.resume(); - return Ok(()); - } - } - } - if !arena.next_tuple(self.left_input, plan_arena)? { - self.join_input = None; - arena.finish(); - return Ok(()); + let Some(input) = self.right_input.take() else { + arena.finish(); + return Ok(()); + }; + if !arena.next_tuple(self.left_input, plan_arena)? { + // Return the temporarily owned cache so reset can reuse it after EOF. + if self.op.materialize { + self.right_input = Some(input); } - let left: Tuple = arena.materialize_tuple(); - let value = self.parameterized_probe_value(&left, plan_arena)?; - let root = self.build_right_input(arena, plan_arena, value); - self.join_input = Some((root, left)); + arena.finish(); + return Ok(()); } + let mut left = arena.take_tuple(); + let (input, marker) = match self.op.guard { + Some(guard) + if *plan_arena.expression(guard).eval(plan_arena, Some(&left))? + == DataValue::Boolean(false) => + { + (input, DataValue::Null) + } + _ => self.mark_value(input, arena, plan_arena, &left)?, + }; + self.right_input = Some(input); + left.push(marker); + arena.produce(left); + Ok(()) } +} +impl<'a> MarkApply<'a> { fn runtime_probe_for(&self, param_value: Option) -> Option { self.op.parameterized_probe()?; - match param_value { Some(value) => Some(RuntimeIndexProbe::Eq(value)), None if matches!( @@ -139,28 +205,9 @@ impl<'a> MarkApply<'a> { } } - fn build_right_input( - &self, - arena: &mut ExecArena<'a, T>, - plan_arena: &mut (dyn MetaArena + 'a), - param_value: Option, - ) -> ExecId { - arena.nodes.seek(self.right_pos); - if let Some(probe) = self.runtime_probe_for(param_value) { - arena.push_runtime_probe(probe); - } - build_read( - arena, - plan_arena, - self.right_input_plan, - arena.context(), - arena.transaction(), - ) - } - fn parameterized_probe_value( &self, - left_tuple: &Tuple, + left_tuple: &TupleImpl, plan_arena: &(dyn MetaArena + '_), ) -> Result, DatabaseError> { self.op @@ -174,193 +221,228 @@ impl<'a> MarkApply<'a> { .transpose() } - fn mark_value( + fn mark_value( &self, + mut input: RightInput, arena: &mut ExecArena<'a, T>, plan_arena: &mut ExecMetaArena, - left_tuple: &Tuple, - ) -> Result { + left_tuple: &TupleImpl, + ) -> Result<(RightInput, DataValue), DatabaseError> { let probe = self.parameterized_probe_value(left_tuple, plan_arena)?; - match self.op.kind { - MarkApplyKind::InnerJoin => unreachable!("inner join streams tuples"), + let value = match self.op.kind { MarkApplyKind::Exists => { - let right_input = self.build_right_input(arena, plan_arena, probe); - while arena.next_tuple(right_input, plan_arena)? { - let right_tuple = arena.result_tuple(); - if Self::predicates_matched( - self.op.predicates(), - left_tuple, - right_tuple, - plan_arena, - )? { - return Ok(DataValue::Boolean(true)); - } - } - Ok(DataValue::Boolean(false)) + let predicates = self.op.predicates(); + let runtime_probe = self.runtime_probe_for(probe); + input + .scan(arena, plan_arena, runtime_probe, |right, meta| { + Ok(predicates_matched(predicates, left_tuple, right, meta)? + .then_some(DataValue::Boolean(true))) + })? + .unwrap_or(DataValue::Boolean(false)) } - MarkApplyKind::Quantified(MarkApplyQuantifier::Any) => { - if let Some(probe_value) = probe { - if !probe_value.is_null() { - let right_input = - self.build_right_input(arena, plan_arena, Some(probe_value)); - while arena.next_tuple(right_input, plan_arena)? { - let right_tuple = arena.result_tuple(); - if self.quantified_predicate_outcome( - left_tuple, - right_tuple, - plan_arena, - )? == QuantifiedPredicateOutcome::True - { - return Ok(DataValue::Boolean(true)); - } - } - - let right_input = - self.build_right_input(arena, plan_arena, Some(DataValue::Null)); - while arena.next_tuple(right_input, plan_arena)? { - let right_tuple = arena.result_tuple(); - if self.quantified_predicate_outcome( - left_tuple, - right_tuple, - plan_arena, - )? == QuantifiedPredicateOutcome::Null - { - return Ok(DataValue::Null); - } - } - - return Ok(DataValue::Boolean(false)); + MarkApplyKind::Quantified(quantifier) => { + if let RightInput::Materialize { root } = input { + input = RightInput::Materialized( + self.materialize_keys(root, arena, plan_arena, left_tuple)?, + ); + } + match &input { + RightInput::Materialized(keys) => { + self.lookup_keys(keys, quantifier, plan_arena, left_tuple)? } + _ => self.scan_quantified_right_input( + &input, arena, plan_arena, probe, quantifier, left_tuple, + )?, } + } + }; + Ok((input, value)) + } - let right_input = self.build_right_input(arena, plan_arena, None); - self.scan_quantified_right_input( - arena, - plan_arena, - right_input, - MarkApplyQuantifier::Any, - left_tuple, - ) + fn quantified_comparison<'p>( + predicates: &[ExprRef], + plan_arena: &'p (dyn MetaArena + '_), + ) -> Result<(BinaryOperator, ExprRef, ExprRef, &'p BinaryEvaluatorRef), DatabaseError> { + let [predicate] = predicates else { + return Err(DatabaseError::InvalidType); + }; + let ScalarExpression::Binary { + op, + left_expr, + right_expr, + evaluator, + .. + } = plan_arena.expression(*predicate) + else { + return Err(DatabaseError::InvalidType); + }; + let evaluator = evaluator.as_ref().ok_or(DatabaseError::EvaluatorNotFound)?; + Ok((*op, *left_expr, *right_expr, evaluator)) + } + + fn materialize_keys( + &self, + root: ExecId, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + left_tuple: &TupleImpl, + ) -> Result { + let (_, _, right_expr, _) = Self::quantified_comparison(self.op.predicates(), plan_arena)?; + arena.reset(root, plan_arena)?; + let mut builder = KeysBuilder::new(compare_keys); + let (mut is_empty, mut has_null) = (true, false); + while arena.next_tuple(root, plan_arena)? { + // RHS positions follow the outer row, so evaluate against the joined layout. + let values = SplitTupleRef::new(left_tuple, arena.result_tuple()); + let key = plan_arena + .expression(right_expr) + .eval(plan_arena, Some(&values))? + .into_owned(); + is_empty = false; + if key.is_null() { + has_null = true; + } else { + builder.push(key)?; + } + } + Ok(MaterializedKeys { + runs: builder.finish()?, + is_empty, + has_null, + fixed_width: plan_arena + .expression(right_expr) + .return_type(plan_arena) + .fixed_encoded_len(), + }) + } + + fn lookup_keys( + &self, + keys: &MaterializedKeys, + quantifier: MarkApplyQuantifier, + plan_arena: &(dyn MetaArena + '_), + left_tuple: &TupleImpl, + ) -> Result { + let any = matches!(quantifier, MarkApplyQuantifier::Any); + if keys.is_empty { + return Ok(DataValue::Boolean(!any)); + } + let (op, left_expr, _, evaluator) = + Self::quantified_comparison(self.op.predicates(), plan_arena)?; + let probe = plan_arena + .expression(left_expr) + .eval(plan_arena, Some(left_tuple))?; + let mut cursor = keys.runs.cursor(keys.fixed_width)?; + let mut saw_null = false; + // `= ANY` / `<> ALL` reduce to membership. Tuple probes compare NULL + // fields three-valued, which membership cannot express. + if matches!( + (op, quantifier), + (BinaryOperator::Eq, MarkApplyQuantifier::Any) + | (BinaryOperator::NotEq, MarkApplyQuantifier::All) + ) && !matches!(probe.as_ref(), DataValue::Tuple(_)) + { + if probe.is_null() { + return Ok(DataValue::Null); } - MarkApplyKind::Quantified(MarkApplyQuantifier::All) => { - let right_input = self.build_right_input(arena, plan_arena, None); - self.scan_quantified_right_input( - arena, - plan_arena, - right_input, - MarkApplyQuantifier::All, - left_tuple, - ) + if cursor.contains_by(|key| Ok(compare_keys(key, &probe)))? { + return Ok(DataValue::Boolean(any)); + } + } else if let Some(value) = cursor.visit_all(|key| { + match evaluator.binary_eval(Cow::Borrowed(probe.as_ref()), key)? { + DataValue::Boolean(value) if value == any => Ok(Some(value)), + DataValue::Boolean(_) => Ok(None), + DataValue::Null => { + saw_null = true; + Ok(None) + } + _ => Err(DatabaseError::InvalidType), } + })? { + return Ok(DataValue::Boolean(value)); } + // No decisive key: a NULL key or comparison makes the result UNKNOWN. + Ok(if keys.has_null || saw_null { + DataValue::Null + } else { + DataValue::Boolean(!any) + }) } - fn scan_quantified_right_input( + fn scan_quantified_right_input( &self, + input: &RightInput, arena: &mut ExecArena<'a, T>, plan_arena: &mut ExecMetaArena, - right_input: ExecId, + probe: Option, quantifier: MarkApplyQuantifier, - left_tuple: &Tuple, + left_tuple: &TupleImpl, ) -> Result { let mut saw_null = false; - - while arena.next_tuple(right_input, plan_arena)? { - let right_tuple = arena.result_tuple(); - match self.quantified_predicate_outcome(left_tuple, right_tuple, plan_arena)? { - QuantifiedPredicateOutcome::True => { - if matches!(quantifier, MarkApplyQuantifier::Any) { - return Ok(DataValue::Boolean(true)); + let predicates = self.op.predicates(); + // One fold across all probed ranges, so the mark is the quantifier over their union. + for runtime_probe in self.quantified_probes(probe) { + let value = input.scan(arena, plan_arena, runtime_probe, |right, meta| { + match Self::quantified_row_value(predicates, left_tuple, right, meta)? { + Some(DataValue::Boolean(true)) + if matches!(quantifier, MarkApplyQuantifier::Any) => + { + return Ok(Some(DataValue::Boolean(true))); } - } - QuantifiedPredicateOutcome::False => { - if matches!(quantifier, MarkApplyQuantifier::All) { - return Ok(DataValue::Boolean(false)); + Some(DataValue::Boolean(false)) + if matches!(quantifier, MarkApplyQuantifier::All) => + { + return Ok(Some(DataValue::Boolean(false))); } + Some(DataValue::Null) => saw_null = true, + _ => {} } - QuantifiedPredicateOutcome::Null => saw_null = true, - QuantifiedPredicateOutcome::Skip => {} + Ok(None) + })?; + if let Some(value) = value { + return Ok(value); } } - - if saw_null { - Ok(DataValue::Null) + Ok(if saw_null { + DataValue::Null } else { - Ok(DataValue::Boolean(matches!( - quantifier, - MarkApplyQuantifier::All - ))) - } - } - - fn predicates_matched( - predicates: &[crate::planner::ExprRef], - left_tuple: &Tuple, - right_tuple: &Tuple, - plan_arena: &(dyn MetaArena + '_), - ) -> Result { - let values = SplitTupleRef::new(left_tuple, right_tuple); - - for predicate in predicates { - match *plan_arena - .expression(*predicate) - .eval(plan_arena, Some(&values))? - { - DataValue::Boolean(true) => {} - DataValue::Boolean(false) | DataValue::Null => return Ok(false), - _ => return Err(DatabaseError::InvalidType), - } - } - - Ok(true) + DataValue::Boolean(matches!(quantifier, MarkApplyQuantifier::All)) + }) } - fn quantified_predicate_outcome( + fn quantified_probes( &self, - left_tuple: &Tuple, - right_tuple: &Tuple, - plan_arena: &(dyn MetaArena + '_), - ) -> Result { - match self.eval_predicates(left_tuple, right_tuple, plan_arena)? { - Some(DataValue::Boolean(true)) => Ok(QuantifiedPredicateOutcome::True), - Some(DataValue::Boolean(false)) => Ok(QuantifiedPredicateOutcome::False), - Some(DataValue::Null) => Ok(QuantifiedPredicateOutcome::Null), - None => Ok(QuantifiedPredicateOutcome::Skip), - Some(_) => Err(DatabaseError::InvalidType), - } + probe: Option, + ) -> impl Iterator> { + let (first, null_keys) = match probe.filter(|value| !value.is_null()) { + Some(key) => ( + Some(RuntimeIndexProbe::Eq(key)), + Some(Some(RuntimeIndexProbe::Eq(DataValue::Null))), + ), + None => (self.runtime_probe_for(None), None), + }; + std::iter::once(first).chain(null_keys) } - - fn eval_predicates( - &self, - left_tuple: &Tuple, - right_tuple: &Tuple, + fn quantified_row_value( + predicates: &[ExprRef], + left_tuple: &TupleImpl, + right_tuple: &TupleImpl, plan_arena: &(dyn MetaArena + '_), ) -> Result, DatabaseError> { + let (comparison, correlated_predicates) = + predicates.split_first().ok_or(DatabaseError::InvalidType)?; + if !predicates_matched(correlated_predicates, left_tuple, right_tuple, plan_arena)? { + return Ok(None); + } let values = SplitTupleRef::new(left_tuple, right_tuple); - // probe_predicate is in predicate, always first - let (probe_predicate, correlated_predicates) = self - .op - .predicates() - .split_first() - .ok_or(DatabaseError::InvalidType)?; - - for predicate in correlated_predicates { - match *plan_arena - .expression(*predicate) - .eval(plan_arena, Some(&values))? - { - DataValue::Boolean(true) => {} - DataValue::Boolean(false) | DataValue::Null => return Ok(None), - _ => return Err(DatabaseError::InvalidType), - } + match plan_arena + .expression(*comparison) + .eval(plan_arena, Some(&values))? + .into_owned() + { + value @ (DataValue::Boolean(_) | DataValue::Null) => Ok(Some(value)), + _ => Err(DatabaseError::InvalidType), } - - Ok(Some( - plan_arena - .expression(*probe_predicate) - .eval(plan_arena, Some(&values))? - .into_owned(), - )) } } @@ -377,7 +459,7 @@ mod tests { use crate::planner::{Childrens, ExprRef, LogicalPlan}; use crate::storage::rocksdb::RocksStorage; use crate::storage::{StatisticsMetaCache, Storage, TableCache, ViewCache}; - use crate::types::evaluator::binary_create; + use crate::types::evaluator::{binary_create, cast_create}; use crate::types::index::RuntimeIndexProbe; use crate::types::tuple::Tuple; use crate::types::LogicalType; @@ -469,113 +551,6 @@ mod tests { })) } - #[test] - fn inner_join_apply_emits_all_matches_and_reuses_right_slots() -> Result<(), DatabaseError> { - let table_arena = crate::planner::TableArenaCell::default(); - let mut plan_arena = crate::planner::PlanArena::new(&table_arena); - let mut left = build_values( - &mut plan_arena, - "left_key", - vec![ - vec![DataValue::Int32(2)], - vec![DataValue::Null], - vec![DataValue::Int32(99)], - vec![DataValue::Int32(2)], - ], - ); - let mut right = build_values_with_schema( - &mut plan_arena, - vec![ - ("right_key", LogicalType::Integer), - ("flag", LogicalType::Boolean), - ], - vec![ - vec![DataValue::Int32(2), DataValue::Boolean(true)], - vec![DataValue::Int32(2), DataValue::Boolean(false)], - vec![DataValue::Int32(2), DataValue::Null], - vec![DataValue::Int32(2), DataValue::Boolean(true)], - vec![DataValue::Null, DataValue::Boolean(true)], - ], - ); - let left_column = left.output_schema(&mut plan_arena)[0]; - let right_schema = right.output_schema(&mut plan_arena).clone(); - let equality = - build_equality_predicate(&mut plan_arena, left_column, 0, right_schema[0], 1)?; - let residual = - plan_arena.alloc_expression(ScalarExpression::column_expr(right_schema[1], 2)); - let probe = plan_arena.alloc_expression(ScalarExpression::column_expr(left_column, 0)); - let mut op = MarkApplyOperator::new_inner_join(vec![equality, residual], probe); - // Values has no parameterized IndexScan to consume a runtime probe. - op.set_parameterized_probe(None); - let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; - let transaction = storage.transaction()?; - let cache = crate::execution::empty_context(&table_cache, &view_cache, &meta_cache); - let mut arena = ExecArena::with_capacity(0); - arena.init_context(cache, &transaction); - let root = >::into_executor( - (&op, &left, &right), - &mut arena, - &mut plan_arena, - cache, - &transaction, - ); - let address = arena.nodes.items.as_ptr(); - let count = arena.nodes.items.len(); - assert_eq!(count, 3); - let mut plan_arena = ExecMetaArena::new(plan_arena); - for _ in 0..4 { - assert!(arena.next_tuple(root, &mut plan_arena)?); - assert_eq!( - arena.materialize_tuple().values, - vec![ - DataValue::Int32(2), - DataValue::Int32(2), - DataValue::Boolean(true) - ] - ); - assert_eq!(arena.nodes.items.as_ptr(), address); - assert_eq!(arena.nodes.items.len(), count); - } - assert!(!arena.next_tuple(root, &mut plan_arena)?); - Ok(()) - } - - #[test] - fn inner_join_apply_does_not_read_past_the_returned_match() -> Result<(), DatabaseError> { - let table_arena = crate::planner::TableArenaCell::default(); - let mut plan_arena = crate::planner::PlanArena::new(&table_arena); - let mut left = build_values(&mut plan_arena, "left_key", vec![vec![DataValue::Int32(1)]]); - let mut right = build_values_with_schema( - &mut plan_arena, - vec![("flag", LogicalType::Boolean)], - // A later row that cannot be cast must not fail the first fetch. - vec![ - vec![DataValue::Boolean(true)], - vec![DataValue::from("not-a-boolean".to_string())], - ], - ); - let left_column = left.output_schema(&mut plan_arena)[0]; - let right_column = right.output_schema(&mut plan_arena)[0]; - let predicate = plan_arena.alloc_expression(ScalarExpression::column_expr(right_column, 1)); - let probe = plan_arena.alloc_expression(ScalarExpression::column_expr(left_column, 0)); - let mut op = MarkApplyOperator::new_inner_join(vec![predicate], probe); - op.set_parameterized_probe(None); - let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; - let transaction = storage.transaction()?; - let mut executor = execute_input::<_, MarkApply>( - (&op, &left, &right), - crate::execution::empty_context(&table_cache, &view_cache, &meta_cache), - plan_arena, - &transaction, - ); - assert_eq!( - executor.next_tuple()?.unwrap().values, - vec![DataValue::Int32(1), DataValue::Boolean(true)] - ); - assert!(executor.next_tuple().is_err()); - Ok(()) - } - #[test] fn mark_exists_apply_appends_boolean_match_column() -> Result<(), DatabaseError> { let table_arena = crate::planner::TableArenaCell::default(); @@ -615,15 +590,16 @@ mod tests { let mut tuples = Vec::new(); let mut plan_arena = ExecMetaArena::new(plan_arena); while arena.next_tuple(root, &mut plan_arena)? { - tuples.push(arena.materialize_tuple()); + tuples.push(arena.take_tuple()); assert_eq!(arena.nodes.items.len(), count); assert_eq!(arena.nodes.items.as_ptr(), address); } assert_eq!( tuples - .into_iter() - .flat_map(|tuple| tuple.values) + .iter() + .flat_map(TupleImpl::iter) + .cloned() .collect::>(), vec![ DataValue::Int32(1), @@ -669,7 +645,7 @@ mod tests { assert_eq!( tuples .into_iter() - .flat_map(|tuple| tuple.values) + .flat_map(|mut tuple| tuple.take_values()) .collect::>(), vec![ DataValue::Int32(1), @@ -733,18 +709,22 @@ mod tests { &transaction, ); - let exec: MarkApply = MarkApply { + let cache = arena.context(); + let right_root = build_read(&mut arena, &mut plan_arena, &right, cache, &transaction); + let mut exec: MarkApply = MarkApply { + right_root, op: &op, - right_input_plan: &right, + right_input: Some(RightInput::Streaming { root: right_root }), left_input: 0, - right_pos: arena.nodes.position(), - join_input: None, }; - let left_tuple = Tuple::new(None, vec![DataValue::Int32(2), DataValue::Int32(1)]); + let left_tuple: TupleImpl = + Tuple::new(None, vec![DataValue::Int32(2), DataValue::Int32(1)]).into(); let mut plan_arena = ExecMetaArena::new(plan_arena); + let input = exec.right_input.take().unwrap(); assert_eq!( - exec.mark_value(&mut arena, &mut plan_arena, &left_tuple)?, + exec.mark_value(input, &mut arena, &mut plan_arena, &left_tuple)? + .1, DataValue::Boolean(true) ); assert_eq!( @@ -787,18 +767,21 @@ mod tests { &transaction, ); - let exec: MarkApply = MarkApply { + let cache = arena.context(); + let right_root = build_read(&mut arena, &mut plan_arena, &right, cache, &transaction); + let mut exec: MarkApply = MarkApply { + right_root, op: &op, - right_input_plan: &right, + right_input: Some(RightInput::Streaming { root: right_root }), left_input: 0, - right_pos: arena.nodes.position(), - join_input: None, }; - let left_tuple = Tuple::new(None, vec![DataValue::Int32(2)]); + let left_tuple: TupleImpl = Tuple::new(None, vec![DataValue::Int32(2)]).into(); let mut plan_arena = ExecMetaArena::new(plan_arena); + let input = exec.right_input.take().unwrap(); assert_eq!( - exec.mark_value(&mut arena, &mut plan_arena, &left_tuple)?, + exec.mark_value(input, &mut arena, &mut plan_arena, &left_tuple)? + .1, DataValue::Boolean(true) ); assert_eq!( @@ -841,18 +824,21 @@ mod tests { &transaction, ); - let exec: MarkApply = MarkApply { + let cache = arena.context(); + let right_root = build_read(&mut arena, &mut plan_arena, &right, cache, &transaction); + let mut exec: MarkApply = MarkApply { + right_root, op: &op, - right_input_plan: &right, + right_input: Some(RightInput::Streaming { root: right_root }), left_input: 0, - right_pos: arena.nodes.position(), - join_input: None, }; - let left_tuple = Tuple::new(None, vec![DataValue::Null]); + let left_tuple: TupleImpl = Tuple::new(None, vec![DataValue::Null]).into(); let mut plan_arena = ExecMetaArena::new(plan_arena); + let input = exec.right_input.take().unwrap(); assert_eq!( - exec.mark_value(&mut arena, &mut plan_arena, &left_tuple)?, + exec.mark_value(input, &mut arena, &mut plan_arena, &left_tuple)? + .1, DataValue::Null ); assert_eq!( @@ -901,7 +887,7 @@ mod tests { assert_eq!( tuples .into_iter() - .flat_map(|tuple| tuple.values) + .flat_map(|mut tuple| tuple.take_values()) .collect::>(), vec![ DataValue::Int32(1), @@ -914,6 +900,402 @@ mod tests { Ok(()) } + #[test] + fn guard_false_skips_failing_rhs_but_true_and_null_evaluate_it() -> Result<(), DatabaseError> { + let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; + let transaction = storage.transaction()?; + let cache = crate::execution::empty_context(&table_cache, &view_cache, &meta_cache); + for kind in [ + MarkApplyKind::Exists, + MarkApplyKind::Quantified(MarkApplyQuantifier::Any), + MarkApplyKind::Quantified(MarkApplyQuantifier::All), + ] { + for materialize in [false, true] { + for (flag, use_guard) in [ + (DataValue::Boolean(false), true), + (DataValue::Boolean(true), true), + (DataValue::Null, true), + (DataValue::Boolean(false), false), + ] { + let tables = crate::planner::TableArenaCell::default(); + let mut meta = crate::planner::PlanArena::new(&tables); + let mut left = build_values_with_schema( + &mut meta, + vec![ + ("key", LogicalType::Integer), + ("flag", LogicalType::Boolean), + ], + vec![vec![DataValue::Int32(1), flag.clone()]], + ); + let mut right = + build_values(&mut meta, "right_key", vec![vec![DataValue::Int32(1)]]); + let text = DataValue::from("not an integer".to_string()); + let evaluator = cast_create(&text.logical_type(), &LogicalType::Integer)?; + let expr = meta.alloc_expression(text.into()); + let failing = meta.alloc_expression(ScalarExpression::TypeCast { + expr, + ty: LogicalType::Integer, + evaluator: Some(evaluator), + }); + let Operator::Values(values) = &mut right.operator else { + unreachable!() + }; + values.rows[0] = failing; + let schema = left.output_schema(&mut meta); + let left_column = schema[0]; + let flag_column = schema[1]; + let right_column = right.output_schema(&mut meta)[0]; + let predicate = + build_equality_predicate(&mut meta, left_column, 0, right_column, 2)?; + let mut op = + MarkApplyOperator::new_in(build_marker_column(&mut meta), vec![predicate]); + op.kind = kind; + op.materialize = materialize && !matches!(kind, MarkApplyKind::Exists); + if use_guard { + op.guard = Some( + meta.alloc_expression(ScalarExpression::column_expr(flag_column, 1)), + ); + } + let mut arena = ExecArena::with_capacity(0); + arena.init_context(cache, &transaction); + let left_input = build_read(&mut arena, &mut meta, &left, cache, &transaction); + let right_root = build_read(&mut arena, &mut meta, &right, cache, &transaction); + let right_input = if op.materialize { + RightInput::Materialize { root: right_root } + } else { + RightInput::Streaming { root: right_root } + }; + let mut exec = MarkApply { + right_root, + op: &op, + left_input, + right_input: Some(right_input), + }; + let mut meta = ExecMetaArena::new(meta); + let result = exec.next_tuple(&mut arena, &mut meta); + if use_guard && flag == DataValue::Boolean(false) { + result?; + assert_eq!( + *arena.result_tuple(), + vec![DataValue::Int32(1), flag, DataValue::Null] + ); + assert!(matches!( + exec.right_input, + Some(RightInput::Streaming { .. } | RightInput::Materialize { .. }) + )); + exec.next_tuple(&mut arena, &mut meta)?; + assert_eq!(exec.right_input.is_none(), !op.materialize); + } else { + assert!(matches!(result, Err(DatabaseError::ParseInt(_)))); + } + } + } + } + Ok(()) + } + + #[test] + fn guard_skips_preserve_materialized_rhs() -> Result<(), DatabaseError> { + let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; + let transaction = storage.transaction()?; + let cache = crate::execution::empty_context(&table_cache, &view_cache, &meta_cache); + let cases = [ + (99, DataValue::Boolean(false), DataValue::Null), + (1, DataValue::Boolean(true), DataValue::Boolean(true)), + (99, DataValue::Boolean(false), DataValue::Null), + (2, DataValue::Boolean(true), DataValue::Boolean(true)), + (99, DataValue::Null, DataValue::Null), + (2, DataValue::Boolean(false), DataValue::Null), + (1, DataValue::Boolean(true), DataValue::Boolean(true)), + ]; + for materialize in [false, true] { + let tables = crate::planner::TableArenaCell::default(); + let mut meta = crate::planner::PlanArena::new(&tables); + let mut left = build_values_with_schema( + &mut meta, + vec![ + ("key", LogicalType::Integer), + ("flag", LogicalType::Boolean), + ], + cases + .iter() + .map(|(key, flag, _)| vec![DataValue::Int32(*key), flag.clone()]) + .collect(), + ); + let mut right = build_values( + &mut meta, + "right_key", + vec![ + vec![DataValue::Int32(1)], + vec![DataValue::Int32(2)], + vec![DataValue::Null], + ], + ); + let schema = left.output_schema(&mut meta); + let left_column = schema[0]; + let flag_column = schema[1]; + let right_column = right.output_schema(&mut meta)[0]; + let predicate = build_equality_predicate(&mut meta, left_column, 0, right_column, 2)?; + let mut op = MarkApplyOperator::new_in(build_marker_column(&mut meta), vec![predicate]); + op.materialize = materialize; + op.guard = Some(meta.alloc_expression(ScalarExpression::column_expr(flag_column, 1))); + let mut arena = ExecArena::with_capacity(0); + arena.init_context(cache, &transaction); + let left_input = build_read(&mut arena, &mut meta, &left, cache, &transaction); + let right_root = build_read(&mut arena, &mut meta, &right, cache, &transaction); + let right_input = if materialize { + RightInput::Materialize { root: right_root } + } else { + RightInput::Streaming { root: right_root } + }; + let mut exec = MarkApply { + right_root, + op: &op, + left_input, + right_input: Some(right_input), + }; + let mut meta = ExecMetaArena::new(meta); + for (index, (key, flag, marker)) in cases.iter().enumerate() { + exec.next_tuple(&mut arena, &mut meta)?; + assert_eq!( + *arena.result_tuple(), + vec![DataValue::Int32(*key), flag.clone(), marker.clone()] + ); + if materialize { + if index == 0 { + assert!(matches!( + exec.right_input, + Some(RightInput::Materialize { .. }) + )); + } else { + // The first unguarded row materializes every RHS key at once. + let Some(RightInput::Materialized(keys)) = &exec.right_input else { + panic!("guard must retain the materialized RHS") + }; + assert!(!keys.is_empty && keys.has_null); + } + } else { + assert!(matches!( + exec.right_input, + Some(RightInput::Streaming { .. }) + )); + } + } + exec.next_tuple(&mut arena, &mut meta)?; + assert_eq!(exec.right_input.is_none(), !materialize); + } + Ok(()) + } + + #[test] + fn quantified_apply_materializes_once_and_reuses_keys_after_short_circuit( + ) -> Result<(), DatabaseError> { + let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; + let transaction = storage.transaction()?; + let cache = crate::execution::empty_context(&table_cache, &view_cache, &meta_cache); + // Exercise both a single segment and the default 1024-row spill boundary. + for (quantifier, repeated_rows) in [ + (MarkApplyQuantifier::Any, 1), + (MarkApplyQuantifier::All, 1), + (MarkApplyQuantifier::Any, 1025), + (MarkApplyQuantifier::All, 1025), + ] { + let table_arena = crate::planner::TableArenaCell::default(); + let mut meta = crate::planner::PlanArena::new(&table_arena); + let mut left = build_values(&mut meta, "left_key", vec![vec![DataValue::Int32(2)]; 8]); + let mut right_rows = vec![vec![DataValue::Int32(2)]; repeated_rows]; + right_rows.push(vec![DataValue::Null]); + let mut right = build_values(&mut meta, "right_key", right_rows); + let left_column = left.output_schema(&mut meta)[0]; + let right_column = right.output_schema(&mut meta)[0]; + let predicate = build_equality_predicate(&mut meta, left_column, 0, right_column, 1)?; + let mut op = MarkApplyOperator::new_quantified( + quantifier, + build_marker_column(&mut meta), + vec![predicate], + ); + op.materialize = true; + let empty = build_values(&mut meta, "empty", vec![]); + let mut arena = ExecArena::with_capacity(0); + arena.init_context(cache, &transaction); + let left_input = build_read(&mut arena, &mut meta, &left, cache, &transaction); + let right_root = build_read(&mut arena, &mut meta, &right, cache, &transaction); + let mut exec = MarkApply { + right_root, + op: &op, + right_input: Some(RightInput::Materialize { root: right_root }), + left_input, + }; + let mut meta = ExecMetaArena::new(meta); + for value in [ + DataValue::Int32(2), + DataValue::Int32(1), + DataValue::Null, + DataValue::Int32(2), + ] { + let expected = match (quantifier, &value) { + (MarkApplyQuantifier::Any, DataValue::Int32(2)) => DataValue::Boolean(true), + (MarkApplyQuantifier::All, DataValue::Int32(1)) => DataValue::Boolean(false), + _ => DataValue::Null, + }; + let tuple: TupleImpl = Tuple::new(None, vec![value]).into(); + let input = exec.right_input.take().unwrap(); + let (input, value) = exec.mark_value(input, &mut arena, &mut meta, &tuple)?; + exec.right_input = Some(input); + assert_eq!(value, expected); + // A short-circuited match still caches every RHS key. + let Some(RightInput::Materialized(keys)) = &exec.right_input else { + panic!("the first lookup must materialize the RHS") + }; + assert!(!keys.is_empty && keys.has_null); + exec.reset(&mut arena, &mut meta)?; + assert!(matches!( + exec.right_input, + Some(RightInput::Materialized(_)) + )); + } + // Exhausting the left must not discard the RHS cache either. + while arena.next_tuple(exec.left_input, &mut meta)? {} + exec.next_tuple(&mut arena, &mut meta)?; + assert!(matches!( + exec.right_input, + Some(RightInput::Materialized(_)) + )); + exec.reset(&mut arena, &mut meta)?; + exec.next_tuple(&mut arena, &mut meta)?; + assert_eq!( + *arena.result_tuple(), + vec![ + DataValue::Int32(2), + if matches!(quantifier, MarkApplyQuantifier::Any) { + DataValue::Boolean(true) + } else { + DataValue::Null + }, + ] + ); + assert!(matches!( + exec.right_input, + Some(RightInput::Materialized(_)) + )); + + // An empty RHS completes immediately, including for a NULL left value. + let cache = arena.context(); + exec.right_root = build_read(&mut arena, &mut meta, &empty, cache, &transaction); + exec.right_input = Some(RightInput::Materialize { + root: exec.right_root, + }); + let input = exec.right_input.take().unwrap(); + let (input, value) = exec.mark_value( + input, + &mut arena, + &mut meta, + &TupleImpl::from(Tuple::new(None, vec![DataValue::Null])), + )?; + exec.right_input = Some(input); + assert_eq!( + value, + DataValue::Boolean(matches!(quantifier, MarkApplyQuantifier::All)) + ); + assert!(matches!( + exec.right_input, + Some(RightInput::Materialized(_)) + )); + exec.reset(&mut arena, &mut meta)?; + assert!(matches!( + exec.right_input, + Some(RightInput::Materialized(_)) + )); + exec.next_tuple(&mut arena, &mut meta)?; + assert_eq!( + arena.result_tuple().iter().last(), + Some(&DataValue::Boolean(matches!( + quantifier, + MarkApplyQuantifier::All + ))) + ); + } + Ok(()) + } + + #[test] + fn paused_rhs_survives_left_subquery_rebuilds() -> Result<(), DatabaseError> { + let tables = crate::planner::TableArenaCell::default(); + let mut meta = crate::planner::PlanArena::new(&tables); + let mut left = build_values( + &mut meta, + "left_key", + vec![ + vec![DataValue::Int32(1)], + vec![DataValue::Int32(2)], + vec![DataValue::Int32(99)], + vec![DataValue::Int32(1)], + ], + ); + let mut lookup = build_values( + &mut meta, + "lookup_key", + vec![vec![DataValue::Int32(1)], vec![DataValue::Int32(2)]], + ); + let left_column = left.output_schema(&mut meta)[0]; + let lookup_column = lookup.output_schema(&mut meta)[0]; + let predicate = build_equality_predicate(&mut meta, left_column, 0, lookup_column, 1)?; + let marker = build_marker_column(&mut meta); + let left = MarkApplyOperator::build_exists(left, lookup, marker, vec![predicate]); + let mut right = build_values( + &mut meta, + "right_key", + vec![ + vec![DataValue::Int32(1)], + vec![DataValue::Int32(2)], + vec![DataValue::Int32(3)], + ], + ); + let right_column = right.output_schema(&mut meta)[0]; + let predicate = build_equality_predicate(&mut meta, left_column, 0, right_column, 2)?; + let marker = build_marker_column(&mut meta); + let mut op = + MarkApplyOperator::new_quantified(MarkApplyQuantifier::Any, marker, vec![predicate]); + op.materialize = true; + let (table_cache, view_cache, meta_cache, _temp_dir, storage) = build_test_storage()?; + let transaction = storage.transaction()?; + let tuples = try_collect(execute_input::<_, MarkApply>( + (&op, &left, &right), + crate::execution::empty_context(&table_cache, &view_cache, &meta_cache), + meta, + &transaction, + ))?; + assert_eq!( + tuples + .into_iter() + .map(|mut row| row.take_values()) + .collect::>(), + vec![ + vec![ + DataValue::Int32(1), + DataValue::Boolean(true), + DataValue::Boolean(true) + ], + vec![ + DataValue::Int32(2), + DataValue::Boolean(true), + DataValue::Boolean(true) + ], + vec![ + DataValue::Int32(99), + DataValue::Boolean(false), + DataValue::Boolean(false) + ], + vec![ + DataValue::Int32(1), + DataValue::Boolean(true), + DataValue::Boolean(true) + ], + ] + ); + Ok(()) + } + #[test] fn mark_in_apply_treats_null_predicate_as_not_matched() -> Result<(), DatabaseError> { let table_arena = crate::planner::TableArenaCell::default(); @@ -949,7 +1331,7 @@ mod tests { assert_eq!( tuples .into_iter() - .flat_map(|tuple| tuple.values) + .flat_map(|mut tuple| tuple.take_values()) .collect::>(), vec![ DataValue::Int32(1), @@ -1020,7 +1402,7 @@ mod tests { assert_eq!( tuples .into_iter() - .flat_map(|tuple| tuple.values) + .flat_map(|mut tuple| tuple.take_values()) .collect::>(), vec![ DataValue::Int32(1), diff --git a/src/execution/dql/projection.rs b/src/execution/dql/projection.rs index 5c9f03a3..769d787b 100644 --- a/src/execution/dql/projection.rs +++ b/src/execution/dql/projection.rs @@ -20,7 +20,7 @@ use crate::planner::operator::project::ProjectOperator; use crate::planner::ExprRef; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; pub struct Projection<'a> { @@ -28,7 +28,7 @@ pub struct Projection<'a> { input: ExecId, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Projection<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Projection<'a> { type Input = (&'a ProjectOperator, &'a LogicalPlan); fn into_executor( @@ -43,7 +43,15 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Projection<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Projection<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Projection<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/dql/recursive_cte.rs b/src/execution/dql/recursive_cte.rs index eb8c06c2..634309fe 100644 --- a/src/execution/dql/recursive_cte.rs +++ b/src/execution/dql/recursive_cte.rs @@ -25,29 +25,59 @@ use crate::execution::{ use crate::planner::operator::recursive_cte::RecursiveScanOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::TupleImpl; use std::mem; pub(crate) enum RecursiveInput { - One(Option), + One { + tuple: TupleImpl, + emitted: bool, + }, #[cfg(feature = "spill")] - Many(SpillReader), + Many(SpillReader), #[cfg(not(feature = "spill"))] - Many(std::vec::IntoIter), + Many { + rows: Vec, + position: usize, + }, } impl Iterator for RecursiveInput { - type Item = Result; + type Item = Result; fn next(&mut self) -> Option { match self { - Self::One(tuple) => tuple.take().map(Ok), + Self::One { tuple, emitted } => { + if *emitted { + None + } else { + *emitted = true; + Some(Ok(tuple.clone())) + } + } #[cfg(feature = "spill")] Self::Many(rows) => rows.next(), #[cfg(not(feature = "spill"))] - Self::Many(rows) => rows.next().map(Ok), + Self::Many { rows, position } => { + let row = rows.get(*position).cloned(); + *position += usize::from(row.is_some()); + row.map(Ok) + } + } + } +} + +impl RecursiveInput { + fn reset(&mut self) -> Result<(), DatabaseError> { + match self { + Self::One { emitted, .. } => *emitted = false, + #[cfg(feature = "spill")] + Self::Many(rows) => rows.reset()?, + #[cfg(not(feature = "spill"))] + Self::Many { position, .. } => *position = 0, } + Ok(()) } } @@ -56,24 +86,24 @@ enum RecursiveRows { #[default] Empty, One { - tuple: Tuple, + tuple: TupleImpl, output_done: bool, }, #[cfg(feature = "spill")] - Writing(SpillVec<'static, Tuple>), + Writing(SpillVec<'static, TupleImpl>), #[cfg(feature = "spill")] - Reading(SpillReader), + Reading(SpillReader), #[cfg(not(feature = "spill"))] - Writing(Vec), + Writing(Vec), #[cfg(not(feature = "spill"))] Reading { - rows: Vec, + rows: Vec, output_index: usize, }, } impl RecursiveRows { - fn push(&mut self, tuple: Tuple) -> Result<(), DatabaseError> { + fn push(&mut self, tuple: TupleImpl) -> Result<(), DatabaseError> { match self { Self::Empty => { *self = Self::One { @@ -132,7 +162,7 @@ impl RecursiveRows { } } - fn next_output(&mut self) -> Result, DatabaseError> { + fn next_output(&mut self) -> Result, DatabaseError> { match self { Self::Empty => Ok(None), Self::One { tuple, output_done } => { @@ -157,14 +187,17 @@ impl RecursiveRows { fn into_input(self) -> Result, DatabaseError> { match self { Self::Empty => Ok(None), - Self::One { tuple, .. } => Ok(Some(RecursiveInput::One(Some(tuple)))), + Self::One { tuple, .. } => Ok(Some(RecursiveInput::One { + tuple, + emitted: false, + })), #[cfg(feature = "spill")] Self::Reading(mut reader) => { reader.reset()?; Ok(Some(RecursiveInput::Many(reader))) } #[cfg(not(feature = "spill"))] - Self::Reading { rows, .. } => Ok(Some(RecursiveInput::Many(rows.into_iter()))), + Self::Reading { rows, .. } => Ok(Some(RecursiveInput::Many { rows, position: 0 })), Self::Writing(_) => unreachable!("recursive generation must be finished first"), } } @@ -176,7 +209,7 @@ enum RecursivePhase { Recursive, } -pub struct RecursiveCte<'a, T: Transaction + 'a> { +pub struct RecursiveCte<'a, T: StorageTransaction + 'a> { recursive_plan: &'a LogicalPlan, anchor_input: ExecId, recursive_arena: ExecArena<'a, T>, @@ -186,7 +219,7 @@ pub struct RecursiveCte<'a, T: Transaction + 'a> { phase: RecursivePhase, } -impl<'a, T: Transaction + 'a> RecursiveCte<'a, T> { +impl<'a, T: StorageTransaction + 'a> RecursiveCte<'a, T> { fn new( anchor_input: ExecId, recursive_plan: &'a LogicalPlan, @@ -226,7 +259,7 @@ impl<'a, T: Transaction + 'a> RecursiveCte<'a, T> { } } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for RecursiveCte<'a, T> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for RecursiveCte<'a, T> { type Input = (&'a LogicalPlan, &'a LogicalPlan); fn into_executor( @@ -247,7 +280,19 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for RecursiveCte<'a, T> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for RecursiveCte<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for RecursiveCte<'a, T> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.working = RecursiveRows::default(); + self.next = RecursiveRows::default(); + self.phase = RecursivePhase::Anchor; + self.recursive_arena.reset_for_rebuild(); + arena.reset(self.anchor_input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -257,14 +302,14 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for RecursiveCte<'a, T> { match self.phase { RecursivePhase::Anchor => { while arena.next_tuple(self.anchor_input, plan_arena)? { - self.next.push(arena.materialize_tuple())?; + self.next.push(arena.take_tuple())?; } self.working = mem::take(&mut self.next).finish()?; self.phase = RecursivePhase::Output; } RecursivePhase::Output => { if let Some(tuple) = self.working.next_output()? { - arena.produce_tuple(tuple); + arena.produce(tuple); return Ok(()); } if !self.start_recursive(plan_arena)? { @@ -278,7 +323,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for RecursiveCte<'a, T> { .recursive_arena .next_tuple(self.recursive_root, plan_arena)? { - self.next.push(self.recursive_arena.materialize_tuple())?; + self.next.push(self.recursive_arena.take_tuple())?; } self.recursive_arena.reset_for_rebuild(); self.working = mem::take(&mut self.next).finish()?; @@ -293,7 +338,7 @@ pub struct RecursiveScan { input: RecursiveInput, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for RecursiveScan { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for RecursiveScan { type Input = &'a RecursiveScanOperator; fn into_executor( @@ -308,14 +353,22 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for RecursiveScan { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for RecursiveScan { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for RecursiveScan { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.input.reset() + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, _plan_arena: &mut ExecMetaArena, ) -> Result<(), DatabaseError> { match self.input.next().transpose()? { - Some(tuple) => arena.produce_tuple(tuple), + Some(tuple) => arena.produce(tuple), None => arena.finish(), } Ok(()) @@ -349,7 +402,12 @@ mod tests { #[test] fn spilled_generation_is_written_once_and_replayed_for_scan() -> Result<(), DatabaseError> { let expected = (0..1100) - .map(|value| Tuple::new(None, vec![DataValue::Int32(value)])) + .map(|value| { + TupleImpl::from(crate::types::tuple::Tuple::new( + None, + vec![DataValue::Int32(value)], + )) + }) .collect::>(); let mut rows = RecursiveRows::default(); for tuple in expected.iter().cloned() { @@ -373,7 +431,12 @@ mod tests { #[test] fn memory_generation_is_replayed_for_scan() -> Result<(), DatabaseError> { let expected = (0..3) - .map(|value| Tuple::new(None, vec![DataValue::Int32(value)])) + .map(|value| { + TupleImpl::from(crate::types::tuple::Tuple::new( + None, + vec![DataValue::Int32(value)], + )) + }) .collect::>(); let mut rows = RecursiveRows::default(); for tuple in expected.iter().cloned() { @@ -449,9 +512,7 @@ mod tests { ty: LogicalType::Integer, }); let recursive = LogicalPlan::new( - Operator::Project(ProjectOperator { - exprs: vec![project_expr], - }), + Operator::Project(ProjectOperator::new(vec![project_expr])), Childrens::Only(Box::new(filter)), ); @@ -471,7 +532,7 @@ mod tests { assert_eq!( tuples .into_iter() - .flat_map(|tuple| tuple.values) + .flat_map(|mut tuple| tuple.take_values()) .collect::>(), vec![ DataValue::Int32(1), diff --git a/src/execution/dql/scalar_query_init.rs b/src/execution/dql/scalar_query_init.rs index 449b870d..b5539442 100644 --- a/src/execution/dql/scalar_query_init.rs +++ b/src/execution/dql/scalar_query_init.rs @@ -20,8 +20,8 @@ use crate::expression::ScalarExpression; use crate::planner::operator::scalar_query_init::ScalarQueryInitOperator; use crate::planner::ExecMetaArena; use crate::planner::{ExprRef, LogicalPlan, MetaArena, ScalarQueryRef}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::TupleImpl; use crate::types::value::DataValue; enum ScalarQueryInitState { @@ -37,22 +37,20 @@ pub struct ScalarQueryInit<'a> { reference: ScalarQueryRef, state: ScalarQueryInitState, param_bindings: &'a [(ScalarQueryRef, ExprRef)], - scratch_tuple: Tuple, + scratch_tuple: TupleImpl, - init_plan: &'a LogicalPlan, - init_pos: ExecId, init: ExecId, } impl<'a> ScalarQueryInit<'a> { - fn evaluate_init( + fn evaluate_init( &self, arena: &mut ExecArena<'a, T>, plan_arena: &mut ExecMetaArena, ) -> Result<(), DatabaseError> { let mut value = DataValue::Null; if arena.next_tuple(self.init, plan_arena)? { - std::mem::swap(&mut value, &mut arena.result_tuple_mut().values[0]); + std::mem::swap(&mut value, &mut arena.result_mut()[0]); } if arena.next_tuple(self.init, plan_arena)? { return Err(DatabaseError::InvalidValue( @@ -63,7 +61,7 @@ impl<'a> ScalarQueryInit<'a> { Ok(()) } - pub(crate) fn build( + pub(crate) fn build( op: &'a ScalarQueryInitOperator, input: ExecId, init_plan: &'a LogicalPlan, @@ -72,7 +70,6 @@ impl<'a> ScalarQueryInit<'a> { cache: ExecutionContext<'_>, transaction: &T, ) -> ExecId { - let init_pos = arena.nodes.position(); let init = build_read(arena, plan_arena, init_plan, cache, transaction); arena.push(ExecNode::ScalarQueryInit(Self { input, @@ -87,14 +84,12 @@ impl<'a> ScalarQueryInit<'a> { ScalarQueryInitState::Initialize }, param_bindings: &op.param_bindings, - scratch_tuple: Tuple::default(), - init_plan, - init_pos, + scratch_tuple: TupleImpl::default(), })) } } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ScalarQueryInit<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for ScalarQueryInit<'a> { type Input = ( &'a ScalarQueryInitOperator, &'a LogicalPlan, @@ -113,7 +108,21 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ScalarQueryInit<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ScalarQueryInit<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for ScalarQueryInit<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.state = if self.reference.is_outer_value { + ScalarQueryInitState::ReadOuter + } else { + ScalarQueryInitState::Initialize + }; + self.scratch_tuple.output_mut().clear(); + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -123,6 +132,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ScalarQueryInit<'a> { match self.state { ScalarQueryInitState::Initialize => { if plan_arena.init_value(self.reference).is_none() { + arena.reset(self.init, plan_arena)?; self.evaluate_init(arena, plan_arena)?; } self.state = ScalarQueryInitState::ReadInput; @@ -139,7 +149,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ScalarQueryInit<'a> { return Ok(()); } // Preserve the outer row and lend the previous subquery buffer to the arena. - std::mem::swap(arena.result_tuple_mut(), &mut self.scratch_tuple); + std::mem::swap(arena.result_mut(), &mut self.scratch_tuple); for (reference, expr) in self.param_bindings { let value = plan_arena .expression(*expr) @@ -150,18 +160,9 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ScalarQueryInit<'a> { self.state = ScalarQueryInitState::EvaluateOuter; } ScalarQueryInitState::EvaluateOuter => { - let previous = arena.nodes.position(); - arena.nodes.seek(self.init_pos); - self.init = build_read( - arena, - plan_arena, - self.init_plan, - arena.context(), - arena.transaction(), - ); - arena.nodes.seek(previous); + arena.reset(self.init, plan_arena)?; self.evaluate_init(arena, plan_arena)?; - std::mem::swap(arena.result_tuple_mut(), &mut self.scratch_tuple); + std::mem::swap(arena.result_mut(), &mut self.scratch_tuple); self.state = ScalarQueryInitState::ReadOuter; arena.resume(); return Ok(()); diff --git a/src/execution/dql/scalar_subquery.rs b/src/execution/dql/scalar_subquery.rs index 98d060a4..80f3aa42 100644 --- a/src/execution/dql/scalar_subquery.rs +++ b/src/execution/dql/scalar_subquery.rs @@ -19,7 +19,7 @@ use crate::execution::{ use crate::planner::operator::scalar_subquery::ScalarSubqueryOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; pub struct ScalarSubquery { @@ -28,7 +28,7 @@ pub struct ScalarSubquery { returned: bool, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ScalarSubquery { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for ScalarSubquery { type Input = (&'a ScalarSubqueryOperator, &'a LogicalPlan); fn into_executor( @@ -48,7 +48,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ScalarSubquery { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ScalarSubquery { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for ScalarSubquery { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.returned = false; + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -67,7 +76,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ScalarSubquery { self.returned = true; if !has_next { - let output = arena.result_tuple_mut(); + let output = arena.output_tuple_mut(); output.pk = None; output.values.clear(); output diff --git a/src/execution/dql/seq_scan.rs b/src/execution/dql/seq_scan.rs index 2e237084..9c031ed7 100644 --- a/src/execution/dql/seq_scan.rs +++ b/src/execution/dql/seq_scan.rs @@ -16,14 +16,14 @@ use crate::errors::DatabaseError; use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor}; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::{Iter, Transaction, TupleIter}; +use crate::storage::{Iter, StorageTransaction, TupleIter}; -pub(crate) struct SeqScan<'a, T: Transaction + 'a> { +pub(crate) struct SeqScan<'a, T: StorageTransaction + 'a> { op: &'a TableScanOperator, iter: Option>, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for SeqScan<'a, T> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for SeqScan<'a, T> { type Input = &'a TableScanOperator; fn into_executor( @@ -37,7 +37,24 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for SeqScan<'a, T> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SeqScan<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for SeqScan<'a, T> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + if let Some(iter) = &mut self.iter { + let state = arena.local_state(plan_arena); + iter.reset( + state.transaction(), + state.table_codec, + &self.op.table_name, + self.op.limit, + )?; + } + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -66,7 +83,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SeqScan<'a, T> { } }; - if iter.next_tuple_into(state.table_codec, &mut state.result.tuple)? { + if iter.next_tuple_into(state.table_codec, state.result.output_mut())? { arena.resume(); } else { arena.finish(); @@ -74,3 +91,84 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SeqScan<'a, T> { Ok(()) } } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::db::test::build_table; + use crate::execution::empty_context; + use crate::planner::{PlanArena, TableArenaCell}; + use crate::storage::memory::MemoryStorage; + use crate::storage::table_codec::TableCodec; + use crate::storage::{StatisticsMetaCache, Storage, TableCache, ViewCache}; + use crate::types::tuple::Tuple; + use crate::types::value::DataValue; + use crate::types::LogicalType; + + #[test] + fn reset_keeps_seq_iterator_and_restores_offset_limit() -> Result<(), DatabaseError> { + let storage = MemoryStorage::new(); + let mut tx = storage.transaction()?; + let mut tables = TableCache::default(); + let cell = TableArenaCell::default(); + let mut meta = PlanArena::new(&cell); + build_table(&mut tables, &mut tx, &mut meta)?; + let meta = PlanArena::new(&cell); + let serializers = [ + LogicalType::Integer.serializable(), + LogicalType::Boolean.serializable(), + LogicalType::Integer.serializable(), + ]; + let mut codec = TableCodec::default(); + for id in 1..=4 { + tx.append_tuple( + &mut codec, + "t1", + &Tuple::new( + Some(DataValue::Int32(id)), + vec![ + DataValue::Int32(id), + DataValue::Boolean(true), + DataValue::Int32(id * 10), + ], + ) + .into(), + &serializers, + false, + )?; + } + let op = TableScanOperator { + table_name: "t1".into(), + columns: tables.get("t1").unwrap().columns().copied().collect(), + limit: (Some(1), Some(2)), + index_infos: vec![], + with_pk: true, + }; + let views = ViewCache::default(); + let statistics = StatisticsMetaCache::default(); + let mut arena = ExecArena::with_capacity(0); + arena.init_context(empty_context(&tables, &views, &statistics), &tx); + let mut meta = ExecMetaArena::new(meta); + let mut scan = SeqScan { + op: &op, + iter: None, + }; + scan.next_tuple(&mut arena, &mut meta)?; + assert_eq!(arena.result_tuple()[0], DataValue::Int32(2)); + let iterator = scan.iter.as_ref().unwrap() as *const _; + for _ in 0..3 { + scan.reset(&mut arena, &mut meta)?; + scan.reset(&mut arena, &mut meta)?; + for id in [2, 3] { + scan.next_tuple(&mut arena, &mut meta)?; + assert_eq!(arena.result_tuple()[0], DataValue::Int32(id)); + assert_eq!(scan.iter.as_ref().unwrap() as *const _, iterator); + } + scan.next_tuple(&mut arena, &mut meta)?; + assert_eq!(arena.result.status, Some(crate::execution::ExecStatus::End)); + } + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/execution/dql/set_membership.rs b/src/execution/dql/set_membership.rs index a83a8679..b1d65932 100644 --- a/src/execution/dql/set_membership.rs +++ b/src/execution/dql/set_membership.rs @@ -19,19 +19,19 @@ use crate::execution::{ use crate::planner::operator::set_membership::SetMembershipKind; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::TupleImpl; use std::collections::HashMap; pub struct SetMembership { kind: SetMembershipKind, left_input: ExecId, right_input: ExecId, - right_counts: HashMap, + right_counts: HashMap, built: bool, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for SetMembership { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for SetMembership { type Input = (SetMembershipKind, &'a LogicalPlan, &'a LogicalPlan); fn into_executor( @@ -53,7 +53,18 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for SetMembership { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SetMembership { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for SetMembership { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.built = false; + self.right_counts.clear(); + arena.reset(self.left_input, plan_arena)?; + arena.reset(self.right_input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -64,7 +75,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SetMembership { if let Some(count) = self.right_counts.get_mut(arena.result_tuple()) { *count += 1; } else { - self.right_counts.insert(arena.materialize_tuple(), 1); + self.right_counts.insert(arena.take_tuple(), 1); } } self.built = true; @@ -91,7 +102,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SetMembership { } impl SetMembership { - fn consume_right_match(&mut self, tuple: &Tuple) -> bool { + fn consume_right_match(&mut self, tuple: &TupleImpl) -> bool { if let Some(count) = self.right_counts.get_mut(tuple) { if *count > 0 { *count -= 1; diff --git a/src/execution/dql/show_table.rs b/src/execution/dql/show_table.rs index ebd617ad..71b8ffc7 100644 --- a/src/execution/dql/show_table.rs +++ b/src/execution/dql/show_table.rs @@ -15,15 +15,15 @@ use crate::errors::DatabaseError; use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor}; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::{TableIter, Transaction}; +use crate::storage::{StorageTransaction, TableIter}; use crate::types::value::{DataValue, Utf8Type}; use crate::types::CharLengthUnits; -pub struct ShowTables<'a, T: Transaction + 'a> { +pub struct ShowTables<'a, T: StorageTransaction + 'a> { pub(crate) metas: Option>, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ShowTables<'a, T> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for ShowTables<'a, T> { type Input = (); fn into_executor( @@ -37,7 +37,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ShowTables<'a, T> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ShowTables<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for ShowTables<'a, T> { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.metas = None; + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -59,7 +68,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ShowTables<'a, T> { return Ok(()); }; - let output = arena.result_tuple_mut(); + let output = arena.output_tuple_mut(); output.pk = None; output.values.clear(); output.values.push(DataValue::Utf8 { diff --git a/src/execution/dql/show_view.rs b/src/execution/dql/show_view.rs index dfeef23d..9f5bcfc3 100644 --- a/src/execution/dql/show_view.rs +++ b/src/execution/dql/show_view.rs @@ -15,15 +15,15 @@ use crate::errors::DatabaseError; use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor}; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::{Transaction, ViewIter}; +use crate::storage::{StorageTransaction, ViewIter}; use crate::types::value::{DataValue, Utf8Type}; use crate::types::CharLengthUnits; -pub struct ShowViews<'a, T: Transaction + 'a> { +pub struct ShowViews<'a, T: StorageTransaction + 'a> { pub(crate) metas: Option>, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ShowViews<'a, T> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for ShowViews<'a, T> { type Input = (); fn into_executor( @@ -37,40 +37,42 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ShowViews<'a, T> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ShowViews<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for ShowViews<'a, T> { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.metas = None; + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, plan_arena: &mut ExecMetaArena, ) -> Result<(), DatabaseError> { if self.metas.is_none() { - let context = arena.context(); let mut state = arena.local_state(plan_arena); let (transaction, table_codec) = state.transaction_codec(); - self.metas = Some(transaction.views( - table_codec, - context.table_cache(), - plan_arena.table_arena_cell(), - context.scala_functions(), - context.table_functions(), - )?); + self.metas = Some(transaction.views(table_codec)?); } - let Some(view) = self + let Some(name) = self .metas .as_mut() .expect("show views iterator initialized") - .try_next()? + .try_next(plan_arena)? else { arena.finish(); return Ok(()); }; - let output = arena.result_tuple_mut(); + let output = arena.output_tuple_mut(); output.pk = None; output.values.clear(); output.values.push(DataValue::Utf8 { - value: view.name.to_string(), + value: name.to_string(), ty: Utf8Type::Variable(None), unit: CharLengthUnits::Characters, }); diff --git a/src/execution/dql/sort.rs b/src/execution/dql/sort.rs index cec626f6..389a6cc6 100644 --- a/src/execution/dql/sort.rs +++ b/src/execution/dql/sort.rs @@ -20,8 +20,8 @@ use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::LogicalPlan; use crate::planner::ScalarQueryRef; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::TupleImpl; use crate::types::value::DataValue; use bumpalo::Bump; use std::cmp::Ordering; @@ -49,6 +49,10 @@ impl<'a, T> NullableVec<'a, T> { self.0.len() } + pub(crate) fn clear(&mut self) { + while self.pop().is_some() {} + } + #[inline] pub(crate) fn pop(&mut self) -> Option { self.0.pop().map(|item| unsafe { item.assume_init() }) @@ -57,7 +61,7 @@ impl<'a, T> NullableVec<'a, T> { impl Drop for NullableVec<'_, T> { fn drop(&mut self) { - while self.pop().is_some() {} + self.clear(); } } @@ -146,7 +150,7 @@ pub(crate) fn compare_sort_keys<'a>( } pub(crate) struct SortTuple { - tuple: Tuple, + tuple: TupleImpl, outer_values: Vec<(ScalarQueryRef, DataValue)>, } @@ -157,7 +161,7 @@ pub struct Sort<'a> { input: ExecId, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Sort<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Sort<'a> { type Input = (&'a SortOperator, &'a LogicalPlan); fn into_executor( @@ -183,7 +187,16 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Sort<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Sort<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Sort<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.rows.clear(); + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -192,7 +205,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Sort<'a> { loop { if let Some((_, row)) = self.rows.pop() { plan_arena.restore_outer_values(row.outer_values); - arena.produce_tuple(row.tuple); + arena.produce(row.tuple); return Ok(()); } while arena.next_tuple(self.input, plan_arena)? { @@ -200,7 +213,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Sort<'a> { self.rows.put(( offset, SortTuple { - tuple: arena.materialize_tuple(), + tuple: arena.take_tuple(), outer_values: plan_arena.take_outer_values(), }, )); @@ -224,7 +237,7 @@ mod test { use crate::expression::ScalarExpression; use crate::planner::operator::sort::SortField; use crate::planner::{ExecMetaArena, PlanArena}; - use crate::types::tuple::Tuple; + use crate::types::tuple::{Tuple, TupleImpl}; use crate::types::value::DataValue; use crate::types::LogicalType; use bumpalo::Bump; @@ -303,12 +316,14 @@ mod test { &transaction, ); let mut metadata = ExecMetaArena::new(metadata); + let nodes = arena.nodes.items.len(); + assert!(arena.next_tuple(root, &mut metadata)?); + assert_eq!(*arena.result_tuple(), vec![DataValue::Int32(1)]); + arena.reset(root, &mut metadata)?; + assert_eq!(arena.nodes.items.len(), nodes); for expected in [1, 2, 3] { assert!(arena.next_tuple(root, &mut metadata)?); - assert_eq!( - arena.result_tuple().values, - vec![DataValue::Int32(expected)] - ); + assert_eq!(*arena.result_tuple(), vec![DataValue::Int32(expected)]); assert_eq!( metadata.init_value(result), Some(&DataValue::Int32(expected)) @@ -322,11 +337,22 @@ mod test { ); } assert!(!arena.next_tuple(root, &mut metadata)?); + arena.reset(root, &mut metadata)?; + for expected in [1, 2, 3] { + assert!(arena.next_tuple(root, &mut metadata)?); + assert_eq!(*arena.result_tuple(), vec![DataValue::Int32(expected)]); + assert_eq!( + metadata.init_value(result), + Some(&DataValue::Int32(expected)) + ); + } + assert!(!arena.next_tuple(root, &mut metadata)?); + assert_eq!(arena.nodes.items.len(), nodes); Ok(()) } #[test] - fn nullable_vec_drops_values() { + fn nullable_vec_clear_and_drop_release_values_without_reallocating() { struct DropValue<'a>(&'a Cell); impl Drop for DropValue<'_> { @@ -341,15 +367,24 @@ mod test { let mut values = NullableVec::new(&arena); values.put(DropValue(&dropped)); values.put(DropValue(&dropped)); + let allocation = (values.0.as_ptr(), values.0.capacity()); + values.clear(); + assert_eq!(dropped.get(), 2); + assert_eq!(values.len(), 0); + assert_eq!((values.0.as_ptr(), values.0.capacity()), allocation); + values.clear(); + assert_eq!(dropped.get(), 2); + values.put(DropValue(&dropped)); + assert_eq!((values.0.as_ptr(), values.0.capacity()), allocation); } - assert_eq!(dropped.get(), 2); + assert_eq!(dropped.get(), 3); } fn sorted_rows<'a>( sort_fields: &[SortField], mut tuples: NullableVec<'a, (usize, SortTuple)>, plan_arena: &mut ExecMetaArena>, - ) -> Result + 'a, DatabaseError> { + ) -> Result + 'a, DatabaseError> { sort_tuples(sort_fields, &mut tuples, plan_arena)?; let mut rows = Vec::with_capacity(tuples.len()); while let Some((_, row)) = tuples.pop() { @@ -392,91 +427,91 @@ mod test { vec.put(( 0_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Null]), + tuple: Tuple::new(None, vec![DataValue::Null]).into(), outer_values: Vec::new(), }, )); vec.put(( 1_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Int32(0)]), + tuple: Tuple::new(None, vec![DataValue::Int32(0)]).into(), outer_values: Vec::new(), }, )); vec.put(( 2_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Int32(1)]), + tuple: Tuple::new(None, vec![DataValue::Int32(1)]).into(), outer_values: Vec::new(), }, )); vec }; - let fn_asc_and_nulls_last_eq = |mut iter: Box>| { + let fn_asc_and_nulls_last_eq = |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1)]) + assert_eq!(tuple, vec![DataValue::Int32(1)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null]) } else { unreachable!() } }; - let fn_desc_and_nulls_last_eq = |mut iter: Box>| { + let fn_desc_and_nulls_last_eq = |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1)]) + assert_eq!(tuple, vec![DataValue::Int32(1)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null]) } else { unreachable!() } }; - let fn_asc_and_nulls_first_eq = |mut iter: Box>| { + let fn_asc_and_nulls_first_eq = |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1)]) + assert_eq!(tuple, vec![DataValue::Int32(1)]) } else { unreachable!() } }; - let fn_desc_and_nulls_first_eq = |mut iter: Box>| { + let fn_desc_and_nulls_first_eq = |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1)]) + assert_eq!(tuple, vec![DataValue::Int32(1)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0)]) } else { unreachable!() } @@ -563,175 +598,175 @@ mod test { vec.put(( 0_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Null, DataValue::Null]), + tuple: Tuple::new(None, vec![DataValue::Null, DataValue::Null]).into(), outer_values: Vec::new(), }, )); vec.put(( 1_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]), + tuple: Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]).into(), outer_values: Vec::new(), }, )); vec.put(( 2_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]), + tuple: Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]).into(), outer_values: Vec::new(), }, )); vec.put(( 3_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]), + tuple: Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]).into(), outer_values: Vec::new(), }, )); vec.put(( 4_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]), + tuple: Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]).into(), outer_values: Vec::new(), }, )); vec.put(( 5_usize, SortTuple { - tuple: Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]), + tuple: Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]).into(), outer_values: Vec::new(), }, )); vec }; let fn_asc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq = - |mut iter: Box>| { + |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Int32(0)]) } else { unreachable!() } }; let fn_asc_1_and_nulls_last_1_and_asc_2_and_nulls_first_2_eq = - |mut iter: Box>| { + |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Int32(0)]) } else { unreachable!() } }; let fn_desc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq = - |mut iter: Box>| { + |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Int32(0)]) } else { unreachable!() } }; let fn_desc_1_and_nulls_last_1_and_asc_2_and_nulls_first_2_eq = - |mut iter: Box>| { + |mut iter: Box>| { if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(1), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(1), DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Int32(0), DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Int32(0), DataValue::Int32(0)]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Null]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Null]) } else { unreachable!() } if let Some(tuple) = iter.next() { - assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Int32(0)]) + assert_eq!(tuple, vec![DataValue::Null, DataValue::Int32(0)]) } else { unreachable!() } diff --git a/src/execution/dql/top_k.rs b/src/execution/dql/top_k.rs index ac925004..5490d0d8 100644 --- a/src/execution/dql/top_k.rs +++ b/src/execution/dql/top_k.rs @@ -22,8 +22,8 @@ use crate::planner::operator::top_k::TopKOperator; use crate::planner::{ExecMetaArena, MetaArena}; use crate::planner::{LogicalPlan, ScalarQueryRef}; use crate::storage::table_codec::BumpBytes; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::TupleImpl; use crate::types::value::DataValue; use bumpalo::Bump; use std::cmp::Ordering; @@ -34,7 +34,7 @@ use std::mem::transmute; struct CmpItem<'a> { key: BumpVec<'a, u8>, sequence: usize, - tuple: Tuple, + tuple: TupleImpl, outer_values: Vec<(ScalarQueryRef, DataValue)>, } @@ -65,7 +65,7 @@ fn top_sort<'a>( full_key: &mut BumpBytes<'a>, sort_fields: &[SortField], heap: &mut BTreeSet>, - tuple: &mut Tuple, + tuple: &mut TupleImpl, keep_count: usize, sequence: usize, plan_arena: &mut ExecMetaArena, @@ -117,7 +117,7 @@ pub struct TopK<'a> { input: ExecId, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for TopK<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for TopK<'a> { type Input = (&'a TopKOperator, &'a LogicalPlan); fn into_executor( @@ -146,7 +146,17 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for TopK<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for TopK<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for TopK<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.output = None; + self.arena.reset(); + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -164,7 +174,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for TopK<'a> { &mut key_scratch, self.sort_fields, &mut set, - arena.result_tuple_mut(), + arena.result_mut(), keep_count, sequence, plan_arena, @@ -186,7 +196,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for TopK<'a> { if let Some(item) = self.output.as_mut().and_then(std::iter::Iterator::next) { plan_arena.restore_outer_values(item.outer_values); - arena.produce_tuple(item.tuple); + arena.produce(item.tuple); } else { arena.finish(); } @@ -204,7 +214,7 @@ mod test { use crate::expression::ScalarExpression; use crate::planner::operator::sort::SortField; use crate::planner::ExecMetaArena; - use crate::types::tuple::Tuple; + use crate::types::tuple::{Tuple, TupleLike}; use crate::types::value::DataValue; use crate::types::LogicalType; use bumpalo::Bump; @@ -219,7 +229,7 @@ mod test { CmpItem { key, sequence, - tuple: Tuple::new(None, vec![DataValue::Int32(value)]), + tuple: Tuple::new(None, vec![DataValue::Int32(value)]).into(), outer_values: Vec::new(), } }; @@ -266,48 +276,48 @@ mod test { let fn_asc_and_nulls_last_eq = |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Int32(0)]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0)]) } else { unreachable!() } if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Int32(1)]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1)]) } else { unreachable!() } }; let fn_desc_and_nulls_last_eq = |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Int32(1)]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1)]) } else { unreachable!() } if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Int32(0)]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0)]) } else { unreachable!() } }; let fn_asc_and_nulls_first_eq = |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Null]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Null]) } else { unreachable!() } if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Int32(0)]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0)]) } else { unreachable!() } }; let fn_desc_and_nulls_first_eq = |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Null]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Null]) } else { unreachable!() } if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Int32(1)]) + assert_eq!(reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1)]) } else { unreachable!() } @@ -319,7 +329,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null]).into(), 2, 0, &mut plan_arena, @@ -328,7 +338,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0)]).into(), 2, 1, &mut plan_arena, @@ -337,7 +347,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1)]), + &mut Tuple::new(None, vec![DataValue::Int32(1)]).into(), 2, 2, &mut plan_arena, @@ -350,7 +360,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null]).into(), 2, 3, &mut plan_arena, @@ -359,7 +369,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0)]).into(), 2, 4, &mut plan_arena, @@ -368,7 +378,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1)]), + &mut Tuple::new(None, vec![DataValue::Int32(1)]).into(), 2, 5, &mut plan_arena, @@ -381,7 +391,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null]).into(), 2, 6, &mut plan_arena, @@ -390,7 +400,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0)]).into(), 2, 7, &mut plan_arena, @@ -399,7 +409,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1)]), + &mut Tuple::new(None, vec![DataValue::Int32(1)]).into(), 2, 8, &mut plan_arena, @@ -412,7 +422,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null]).into(), 2, 9, &mut plan_arena, @@ -421,7 +431,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0)]).into(), 2, 10, &mut plan_arena, @@ -430,7 +440,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1)]), + &mut Tuple::new(None, vec![DataValue::Int32(1)]).into(), 2, 11, &mut plan_arena, @@ -484,13 +494,16 @@ mod test { let fn_asc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq = |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Null, DataValue::Null]) + assert_eq!( + reverse.tuple.as_slice().unwrap(), + vec![DataValue::Null, DataValue::Null] + ) } else { unreachable!() } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Null, DataValue::Int32(0)] ) } else { @@ -498,7 +511,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0), DataValue::Null] ) } else { @@ -506,7 +519,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0), DataValue::Int32(0)] ) } else { @@ -517,7 +530,7 @@ mod test { |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0), DataValue::Null] ) } else { @@ -525,7 +538,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0), DataValue::Int32(0)] ) } else { @@ -533,7 +546,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1), DataValue::Null] ) } else { @@ -541,7 +554,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1), DataValue::Int32(0)] ) } else { @@ -551,13 +564,16 @@ mod test { let fn_desc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq = |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { - assert_eq!(reverse.tuple.values, vec![DataValue::Null, DataValue::Null]) + assert_eq!( + reverse.tuple.as_slice().unwrap(), + vec![DataValue::Null, DataValue::Null] + ) } else { unreachable!() } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Null, DataValue::Int32(0)] ) } else { @@ -565,7 +581,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1), DataValue::Null] ) } else { @@ -573,7 +589,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1), DataValue::Int32(0)] ) } else { @@ -584,7 +600,7 @@ mod test { |mut heap: BTreeSet>| { if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1), DataValue::Null] ) } else { @@ -592,7 +608,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(1), DataValue::Int32(0)] ) } else { @@ -600,7 +616,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0), DataValue::Null] ) } else { @@ -608,7 +624,7 @@ mod test { } if let Some(reverse) = heap.pop_first() { assert_eq!( - reverse.tuple.values, + reverse.tuple.as_slice().unwrap(), vec![DataValue::Int32(0), DataValue::Int32(0)] ) } else { @@ -622,7 +638,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]).into(), 4, 12, &mut plan_arena, @@ -631,7 +647,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]).into(), 4, 13, &mut plan_arena, @@ -640,7 +656,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]).into(), 4, 14, &mut plan_arena, @@ -649,7 +665,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]).into(), 4, 15, &mut plan_arena, @@ -658,7 +674,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]).into(), 4, 16, &mut plan_arena, @@ -667,7 +683,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]).into(), 4, 17, &mut plan_arena, @@ -680,7 +696,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]).into(), 4, 18, &mut plan_arena, @@ -689,7 +705,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]).into(), 4, 19, &mut plan_arena, @@ -698,7 +714,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]).into(), 4, 20, &mut plan_arena, @@ -707,7 +723,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]).into(), 4, 21, &mut plan_arena, @@ -716,7 +732,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]).into(), 4, 22, &mut plan_arena, @@ -725,7 +741,7 @@ mod test { &mut key_scratch, &fn_sort_fields(true, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]).into(), 4, 23, &mut plan_arena, @@ -738,7 +754,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]).into(), 4, 24, &mut plan_arena, @@ -747,7 +763,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]).into(), 4, 25, &mut plan_arena, @@ -756,7 +772,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]).into(), 4, 26, &mut plan_arena, @@ -765,7 +781,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]).into(), 4, 27, &mut plan_arena, @@ -774,7 +790,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]).into(), 4, 28, &mut plan_arena, @@ -783,7 +799,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, true, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]).into(), 4, 29, &mut plan_arena, @@ -796,7 +812,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Null]).into(), 4, 30, &mut plan_arena, @@ -805,7 +821,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Null]).into(), 4, 31, &mut plan_arena, @@ -814,7 +830,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Null]).into(), 4, 32, &mut plan_arena, @@ -823,7 +839,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Null, DataValue::Int32(0)]).into(), 4, 33, &mut plan_arena, @@ -832,7 +848,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(0), DataValue::Int32(0)]).into(), 4, 34, &mut plan_arena, @@ -841,7 +857,7 @@ mod test { &mut key_scratch, &fn_sort_fields(false, false, true, true), &mut indices, - &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]), + &mut Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(0)]).into(), 4, 35, &mut plan_arena, diff --git a/src/execution/dql/union.rs b/src/execution/dql/union.rs index 1d409344..753c992c 100644 --- a/src/execution/dql/union.rs +++ b/src/execution/dql/union.rs @@ -18,14 +18,14 @@ use crate::execution::{ }; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; pub struct Union { left_input: ExecId, right_input: ExecId, reading_left: bool, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Union { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Union { type Input = (&'a LogicalPlan, &'a LogicalPlan); fn into_executor( @@ -45,7 +45,17 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Union { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Union { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Union { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.reading_left = true; + arena.reset(self.left_input, plan_arena)?; + arena.reset(self.right_input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, diff --git a/src/execution/dql/values.rs b/src/execution/dql/values.rs index 9d75aa43..e18ff9fa 100644 --- a/src/execution/dql/values.rs +++ b/src/execution/dql/values.rs @@ -17,33 +17,26 @@ use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNo use crate::planner::operator::values::ValuesOperator; use crate::planner::ExprRef; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Schema; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; pub struct Values<'a> { + op: &'a ValuesOperator, rows: std::slice::Iter<'a, ExprRef>, remaining_rows: usize, - schema_ref: &'a Schema, } impl<'a> From<&'a ValuesOperator> for Values<'a> { - fn from( - ValuesOperator { - rows, - row_count, - schema_ref, - }: &'a ValuesOperator, - ) -> Self { + fn from(op: &'a ValuesOperator) -> Self { Values { - rows: rows.iter(), - remaining_rows: *row_count, - schema_ref, + op, + rows: op.rows.iter(), + remaining_rows: op.row_count, } } } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Values<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Values<'a> { type Input = &'a ValuesOperator; fn into_executor( @@ -58,7 +51,17 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Values<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Values<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Values<'a> { + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.rows = self.op.rows.iter(); + self.remaining_rows = self.op.row_count; + Ok(()) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -69,13 +72,13 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Values<'a> { return Ok(()); } self.remaining_rows -= 1; - let width = self.schema_ref.len(); + let width = self.op.schema_ref.len(); - let output = arena.result_tuple_mut(); + let output = arena.output_tuple_mut(); output.pk = None; output.values.clear(); for (i, expr) in self.rows.by_ref().take(width).enumerate() { - let ty = plan_arena.column(self.schema_ref[i]).datatype(); + let ty = plan_arena.column(self.op.schema_ref[i]).datatype(); output.values.push( plan_arena .expression(*expr) diff --git a/src/execution/dql/window.rs b/src/execution/dql/window.rs index 32cc9aa1..e9e8900b 100644 --- a/src/execution/dql/window.rs +++ b/src/execution/dql/window.rs @@ -21,8 +21,8 @@ use crate::planner::operator::sort::SortField; use crate::planner::operator::window::WindowOperator; use crate::planner::LogicalPlan; use crate::planner::{ExecMetaArena, MetaArena}; -use crate::storage::Transaction; -use crate::types::tuple::Tuple; +use crate::storage::StorageTransaction; +use crate::types::tuple::TupleImpl; use crate::types::value::DataValue; mod function; @@ -44,8 +44,8 @@ enum Boundary { #[derive(Default)] struct WindowState { - buffered: Vec<(usize, Tuple)>, - pending: Option<(Tuple, Boundary)>, + buffered: Vec<(usize, TupleImpl)>, + pending: Option<(TupleImpl, Boundary)>, sort_values: Vec, started: bool, partition_rows: usize, @@ -54,6 +54,14 @@ struct WindowState { } impl WindowState { + fn clear(&mut self) { + self.buffered.clear(); + self.pending = None; + self.sort_values.clear(); + self.started = false; + self.begin_partition(); + } + fn begin_partition(&mut self) { self.partition_rows = 0; self.peer_start = 0; @@ -76,7 +84,7 @@ pub struct Window<'a> { input: ExecId, } -impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Window<'a> { +impl<'a, T: StorageTransaction + 'a> ReadExecutor<'a, T> for Window<'a> { type Input = (&'a WindowOperator, &'a LogicalPlan); fn into_executor( @@ -124,7 +132,7 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Window<'a> { impl<'a> Window<'a> { fn update_keys( &mut self, - tuple: &Tuple, + tuple: &TupleImpl, plan_arena: &(dyn MetaArena + '_), ) -> Result, DatabaseError> { let mut boundary = (!self.state.started).then_some(Boundary::Partition); @@ -158,10 +166,9 @@ impl<'a> Window<'a> { if self.state.buffered.is_empty() { return Ok(()); } - let output_offset = self.state.buffered[0].1.values.len(); + let output_offset = self.state.buffered[0].1.len(); for (_, row) in &mut self.state.buffered { - row.values - .resize(output_offset + self.functions.len(), DataValue::Null); + row.resize(output_offset + self.functions.len(), DataValue::Null); } let len = self.state.buffered.len(); for (slot, function) in self.functions.iter_mut().enumerate() { @@ -180,7 +187,7 @@ impl<'a> Window<'a> { fn eval( &mut self, - tuple: Tuple, + tuple: TupleImpl, boundary: Option, plan_arena: &(dyn MetaArena + '_), ) -> Result { @@ -218,7 +225,18 @@ impl<'a> Window<'a> { } } -impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window<'a> { +impl<'a, T: StorageTransaction + 'a> ExecutorNode<'a, T> for Window<'a> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + self.state.clear(); + self.input_exhausted = false; + self.reset_functions()?; + arena.reset(self.input, plan_arena) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -228,7 +246,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window<'a> { loop { if output_ready { if let Some((_, tuple)) = self.state.buffered.pop() { - arena.produce_tuple(tuple); + arena.produce(tuple); return Ok(()); } } @@ -240,7 +258,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window<'a> { let (tuple, boundary) = if let Some((tuple, boundary)) = self.state.pending.take() { (tuple, Some(boundary)) } else if arena.next_tuple(self.input, plan_arena)? { - (arena.materialize_tuple(), None) + (arena.take_tuple(), None) } else { self.eval_functions(plan_arena)?; self.input_exhausted = true; @@ -287,8 +305,55 @@ mod tests { } } - fn tuple(values: &[i32]) -> Tuple { - Tuple::new(None, values.iter().copied().map(DataValue::from).collect()) + fn tuple(values: &[i32]) -> TupleImpl { + crate::types::tuple::Tuple::new(None, values.iter().copied().map(DataValue::from).collect()) + .into() + } + + #[test] + fn clear_restores_window_state_and_keeps_buffer_capacity() { + let mut state = WindowState { + buffered: Vec::with_capacity(8), + pending: Some((tuple(&[2, 20]), Boundary::Peer)), + sort_values: Vec::with_capacity(4), + started: true, + partition_rows: 5, + peer_start: 3, + peer_index: 2, + }; + state.buffered.push((0, tuple(&[1, 10]))); + state.sort_values.push(DataValue::Int32(10)); + let buffered = (state.buffered.as_ptr(), state.buffered.capacity()); + let keys = (state.sort_values.as_ptr(), state.sort_values.capacity()); + for _ in 0..2 { + state.clear(); + assert!(state.buffered.is_empty()); + assert!(state.pending.is_none()); + assert!(state.sort_values.is_empty()); + assert!(!state.started); + assert_eq!( + (state.partition_rows, state.peer_start, state.peer_index), + (0, 0, 0) + ); + assert_eq!( + (state.buffered.as_ptr(), state.buffered.capacity()), + buffered + ); + assert_eq!( + (state.sort_values.as_ptr(), state.sort_values.capacity()), + keys + ); + } + state.buffered.push((0, tuple(&[3, 30]))); + state.sort_values.push(DataValue::Int32(30)); + assert_eq!( + (state.buffered.as_ptr(), state.buffered.capacity()), + buffered + ); + assert_eq!( + (state.sort_values.as_ptr(), state.sort_values.capacity()), + keys + ); } #[test] @@ -317,7 +382,7 @@ mod tests { for (value, expected) in [(10, [1_i64, 1]), (10, [2, 1]), (20, [3, 3])] { window.eval(tuple(&[1, value]), None, &plan_arena)?; let row = window.state.buffered.pop().unwrap().1; - assert_eq!(row.values[2..], expected.map(DataValue::from)); + assert!(row.iter().skip(2).eq(&expected.map(DataValue::from))); } Ok(()) } diff --git a/src/execution/dql/window/function.rs b/src/execution/dql/window/function.rs index 53808f3d..1558bf28 100644 --- a/src/execution/dql/window/function.rs +++ b/src/execution/dql/window/function.rs @@ -18,7 +18,7 @@ use crate::expression::agg::AggKind; use crate::expression::window::WindowFunctionKind; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::Tuple; +use crate::types::tuple::TupleImpl; use crate::types::value::DataValue; use crate::types::LogicalType; use std::ops::Range; @@ -30,7 +30,7 @@ pub(super) trait WindowFunction { fn evaluate( &mut self, - rows: &mut [(usize, Tuple)], + rows: &mut [(usize, TupleImpl)], peer: Range, peer_start: usize, peer_index: usize, @@ -44,7 +44,7 @@ struct RowNumber; impl WindowFunction for RowNumber { fn evaluate( &mut self, - rows: &mut [(usize, Tuple)], + rows: &mut [(usize, TupleImpl)], peer: Range, _peer_start: usize, _peer_index: usize, @@ -52,7 +52,7 @@ impl WindowFunction for RowNumber { _arena: &(dyn MetaArena + '_), ) -> Result<(), DatabaseError> { for (row_index, row) in &mut rows[peer] { - row.values[output_position] = DataValue::Int64((*row_index + 1) as i64); + row[output_position] = DataValue::Int64((*row_index + 1) as i64); } Ok(()) } @@ -65,7 +65,7 @@ struct Rank { impl WindowFunction for Rank { fn evaluate( &mut self, - rows: &mut [(usize, Tuple)], + rows: &mut [(usize, TupleImpl)], peer: Range, peer_start: usize, peer_index: usize, @@ -78,7 +78,7 @@ impl WindowFunction for Rank { peer_start + 1 }; for (_, row) in &mut rows[peer] { - row.values[output_position] = DataValue::Int64(rank as i64); + row[output_position] = DataValue::Int64(rank as i64); } Ok(()) } @@ -99,7 +99,7 @@ impl WindowFunction for Aggregate { fn evaluate( &mut self, - rows: &mut [(usize, Tuple)], + rows: &mut [(usize, TupleImpl)], peer: Range, _peer_start: usize, _peer_index: usize, @@ -116,7 +116,7 @@ impl WindowFunction for Aggregate { accumulator.evaluate()?; let result = accumulator.result(); for (_, row) in &mut rows[peer] { - row.values[output_position] = result.clone(); + row[output_position] = result.clone(); } Ok(()) } diff --git a/src/execution/mod.rs b/src/execution/mod.rs index 64c29a74..810dc761 100644 --- a/src/execution/mod.rs +++ b/src/execution/mod.rs @@ -20,6 +20,7 @@ pub(crate) mod ddl; mod ddl_apply; pub(crate) mod dml; pub(crate) mod dql; +pub(crate) mod runs; #[cfg(feature = "spill")] pub(crate) mod spill; @@ -61,6 +62,7 @@ use crate::execution::dql::filter::Filter; use crate::execution::dql::function_scan::FunctionScan; use crate::execution::dql::index_scan::IndexScan; use crate::execution::dql::join::hash_join::HashJoin; +use crate::execution::dql::join::index_probe_join::IndexProbeJoin; use crate::execution::dql::limit::Limit; use crate::execution::dql::projection::Projection; use crate::execution::dql::recursive_cte::{RecursiveCte, RecursiveInput, RecursiveScan}; @@ -80,10 +82,9 @@ use crate::planner::operator::{Operator, PhysicalOption, PlanImpl}; use crate::planner::MetaArena; use crate::planner::{ExecMetaArena, ExprRef, LogicalPlan, PlanKeeper}; use crate::storage::table_codec::TableCodec; -use crate::storage::{StatisticsMetaCache, TableCache, Transaction, ViewCache}; +use crate::storage::{StatisticsMetaCache, StorageTransaction, TableCache, ViewCache}; use crate::types::index::RuntimeIndexProbe; -use crate::types::tuple::{Tuple, TupleLike}; -use crate::types::value::DataValue; +use crate::types::tuple::{ConcatPool, Input, Tuple, TupleId, TupleImpl}; #[derive(Clone, Copy)] pub(crate) struct ExecutionContext<'a> { @@ -115,14 +116,6 @@ impl<'a> ExecutionContext<'a> { self.table_cache } - pub(crate) fn scala_functions(self) -> &'a ScalaFunctions { - self.scala_functions - } - - pub(crate) fn table_functions(self) -> &'a TableFunctions { - self.table_functions - } - fn is_same_context(&self, other: ExecutionContext<'_>) -> bool { std::ptr::eq(self.table_cache, other.table_cache) && std::ptr::eq(self.view_cache, other.view_cache) @@ -142,8 +135,16 @@ pub(crate) enum ExecStatus { #[derive(Debug, Default)] pub(crate) struct ExecResult { - pub(crate) tuple: Tuple, + pub(crate) tuple: TupleImpl, pub(crate) status: Option, + concat_pool: ConcatPool, +} + +impl ExecResult { + #[inline] + pub(crate) fn output_mut(&mut self) -> &mut Tuple { + self.concat_pool.retain_leftmost(&mut self.tuple) + } } /// Resolves either an arena-backed expression or a direct scalar expression. @@ -163,7 +164,7 @@ impl RewriteExpression for ScalarExpression { } } -pub struct Executor<'a, T: Transaction + 'a> { +pub struct Executor<'a, T: StorageTransaction + 'a> { arena: ExecArena<'a, T>, root: ExecId, // Never read: it only keeps the plan that `arena` borrows from alive, and must stay @@ -172,7 +173,7 @@ pub struct Executor<'a, T: Transaction + 'a> { keeper: PlanKeeper<'a>, } -impl<'a, T: Transaction + 'a> Executor<'a, T> { +impl<'a, T: StorageTransaction + 'a> Executor<'a, T> { pub(crate) fn new(arena: ExecArena<'a, T>, root: ExecId, keeper: PlanKeeper<'a>) -> Self { Self { arena, @@ -184,11 +185,11 @@ impl<'a, T: Transaction + 'a> Executor<'a, T> { pub(crate) fn next_tuple( &mut self, plan_arena: &mut ExecMetaArena, - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { if !self.arena.next_tuple(self.root, plan_arena)? { return Ok(None); } - Ok(Some(self.arena.result_tuple_mut())) + Ok(Some(self.arena.result_mut())) } pub(crate) fn take_ddl_apply(&mut self) -> Vec { @@ -197,7 +198,7 @@ impl<'a, T: Transaction + 'a> Executor<'a, T> { } #[allow(clippy::large_enum_variant)] -pub(crate) enum ExecNode<'a, T: Transaction + 'a> { +pub(crate) enum ExecNode<'a, T: StorageTransaction + 'a> { AddColumn(AddColumn<'a>), Analyze(Analyze<'a>), ChangeColumn(ChangeColumn<'a>), @@ -226,7 +227,8 @@ pub(crate) enum ExecNode<'a, T: Transaction + 'a> { Insert(Insert<'a>), Limit(Limit), MarkApply(MarkApply<'a>), - NestedLoopJoin(NestedLoopJoin<'a>), + IndexProbeJoin(IndexProbeJoin<'a>), + NestedLoopJoin(NestedLoopJoin), Projection(Projection<'a>), RecursiveCte(RecursiveCte<'a, T>), RecursiveScan(RecursiveScan), @@ -248,7 +250,18 @@ pub(crate) enum ExecNode<'a, T: Transaction + 'a> { Window(Window<'a>), } -pub(crate) trait ExecutorNode<'a, T: Transaction + 'a>: Sized { +pub(crate) trait ExecutorNode<'a, T: StorageTransaction + 'a>: Sized { + // Reads restore their local state and reset inputs. Writes must not be replayed. + fn reset( + &mut self, + _arena: &mut ExecArena<'a, T>, + _plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + Err(DatabaseError::UnsupportedStmt( + "resetting a write executor".into(), + )) + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -256,7 +269,151 @@ pub(crate) trait ExecutorNode<'a, T: Transaction + 'a>: Sized { ) -> Result<(), DatabaseError>; } -impl<'a, T: Transaction + 'a> ExecNode<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecNode<'a, T> { + fn reset( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + match self { + ExecNode::AddColumn(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::Analyze(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::ChangeColumn(exec) => { + >::reset(exec, arena, plan_arena) + } + #[cfg(feature = "copy")] + ExecNode::CopyFromFile(exec) => { + >::reset(exec, arena, plan_arena) + } + #[cfg(feature = "copy")] + ExecNode::CopyToFile(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::CreateIndex(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::CreateTable(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::CreateView(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::Delete(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::Describe(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::DropColumn(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::DropIndex(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::DropTable(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::DropView(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::Dummy(exec) => >::reset(exec, arena, plan_arena), + ExecNode::Explain(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + #[cfg(feature = "spill")] + ExecNode::ExternalSort(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::Filter(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::FunctionScan(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::HashAgg(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::HashJoin(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::IndexScan(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::Insert(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::Limit(exec) => >::reset(exec, arena, plan_arena), + ExecNode::MarkApply(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::IndexProbeJoin(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::NestedLoopJoin(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::Projection(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::RecursiveCte(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::RecursiveScan(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::ScalarQueryInit(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::ScalarSubquery(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::SetMembership(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::SeqScan(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::ShowTables(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::ShowViews(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::SimpleAgg(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::Sort(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::StreamAgg(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::StreamDistinct(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::TopK(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::Truncate(exec) => { + >::reset(exec, arena, plan_arena) + } + ExecNode::Union(exec) => >::reset(exec, arena, plan_arena), + ExecNode::Update(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::Values(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + ExecNode::Window(exec) => { + as ExecutorNode<'a, T>>::reset(exec, arena, plan_arena) + } + } + } + fn next_tuple( &mut self, arena: &mut ExecArena<'a, T>, @@ -341,8 +498,11 @@ impl<'a, T: Transaction + 'a> ExecNode<'a, T> { ExecNode::MarkApply(exec) => { as ExecutorNode<'a, T>>::next_tuple(exec, arena, plan_arena) } + ExecNode::IndexProbeJoin(exec) => { + as ExecutorNode<'a, T>>::next_tuple(exec, arena, plan_arena) + } ExecNode::NestedLoopJoin(exec) => { - as ExecutorNode<'a, T>>::next_tuple(exec, arena, plan_arena) + >::next_tuple(exec, arena, plan_arena) } ExecNode::Projection(exec) => { as ExecutorNode<'a, T>>::next_tuple(exec, arena, plan_arena) @@ -407,17 +567,14 @@ impl<'a, T: Transaction + 'a> ExecNode<'a, T> { } } -struct ExecNodes<'a, T: Transaction + 'a> { +struct ExecNodes<'a, T: StorageTransaction + 'a> { items: Vec>>, pos: ExecId, executing: usize, } -impl<'a, T: Transaction + 'a> ExecNodes<'a, T> { - fn position(&self) -> ExecId { - self.pos - } - +impl<'a, T: StorageTransaction + 'a> ExecNodes<'a, T> { + #[cfg(test)] fn seek(&mut self, pos: ExecId) { assert!(pos <= self.items.len(), "node position out of bounds"); self.pos = pos; @@ -442,7 +599,7 @@ impl<'a, T: Transaction + 'a> ExecNodes<'a, T> { } } -pub(crate) struct ExecArena<'a, T: Transaction + 'a> { +pub(crate) struct ExecArena<'a, T: StorageTransaction + 'a> { nodes: ExecNodes<'a, T>, result: ExecResult, table_codec: TableCodec, @@ -453,7 +610,7 @@ pub(crate) struct ExecArena<'a, T: Transaction + 'a> { recursive_input: Option, } -pub(crate) struct ExecArenaLocalState<'b, 'a, T: Transaction + 'a> { +pub(crate) struct ExecArenaLocalState<'b, 'a, T: StorageTransaction + 'a> { transaction: *mut T, pub(crate) table_codec: &'b mut TableCodec, pub(crate) context: ExecutionContext<'a>, @@ -462,7 +619,7 @@ pub(crate) struct ExecArenaLocalState<'b, 'a, T: Transaction + 'a> { ddl_apply: &'b mut Vec, } -impl<'b, 'a, T: Transaction + 'a> ExecArenaLocalState<'b, 'a, T> { +impl<'b, 'a, T: StorageTransaction + 'a> ExecArenaLocalState<'b, 'a, T> { pub(crate) fn transaction(&self) -> &'a T { unsafe { &*self.transaction } } @@ -471,7 +628,7 @@ impl<'b, 'a, T: Transaction + 'a> ExecArenaLocalState<'b, 'a, T> { unsafe { (&mut *self.transaction, &mut *self.table_codec) } } - pub(crate) fn tuple_transaction_codec_mut(&mut self) -> (&Tuple, &mut T, &mut TableCodec) { + pub(crate) fn tuple_transaction_codec_mut(&mut self) -> (&TupleImpl, &mut T, &mut TableCodec) { unsafe { ( &self.result.tuple, @@ -481,18 +638,6 @@ impl<'b, 'a, T: Transaction + 'a> ExecArenaLocalState<'b, 'a, T> { } } - pub(crate) fn index_values_transaction_codec_mut( - &mut self, - ) -> (&[DataValue], &mut T, &mut TableCodec) { - unsafe { - ( - &self.result.tuple.values, - &mut *self.transaction, - &mut *self.table_codec, - ) - } - } - pub(crate) fn transaction_codec(&mut self) -> (&'a T, &mut TableCodec) { unsafe { (&*self.transaction, &mut *self.table_codec) } } @@ -510,7 +655,7 @@ impl<'b, 'a, T: Transaction + 'a> ExecArenaLocalState<'b, 'a, T> { } } -impl<'a, T: Transaction + 'a> ExecArena<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecArena<'a, T> { pub(crate) fn set_statement_stamp(&mut self, stamp: u64) { self.table_codec.set_stamp(stamp); } @@ -533,7 +678,7 @@ impl<'a, T: Transaction + 'a> ExecArena<'a, T> { } } -impl<'a, T: Transaction + 'a> ExecArena<'a, T> { +impl<'a, T: StorageTransaction + 'a> ExecArena<'a, T> { pub(crate) fn init_context(&mut self, context: ExecutionContext<'a>, transaction: &'a T) { if let Some(current) = &self.context { debug_assert!(current.is_same_context(context)); @@ -621,56 +766,82 @@ impl<'a, T: Transaction + 'a> ExecArena<'a, T> { debug_assert!(self.runtime_probe_stack.is_empty()); debug_assert!(self.ddl_apply.is_empty()); self.nodes.clear(); - self.result.tuple = Tuple::default(); + self.result.tuple = TupleImpl::default(); self.result.status = None; self.recursive_input = None; } #[inline] - pub(crate) fn result_tuple(&self) -> &Tuple { + pub(crate) fn result_tuple(&self) -> &TupleImpl { &self.result.tuple } #[inline] - pub(crate) fn result_tuple_mut(&mut self) -> &mut Tuple { + pub(crate) fn result_mut(&mut self) -> &mut TupleImpl { &mut self.result.tuple } - pub(crate) fn materialize_tuple(&mut self) -> Tuple { + #[inline] + pub(crate) fn output_tuple_mut(&mut self) -> &mut Tuple { + self.result.output_mut() + } + + #[inline] + pub(crate) fn concat( + &mut self, + left: TupleImpl, + right: TupleImpl, + pk: Option, + ) -> TupleImpl { + self.result.concat_pool.concat(left, right, pk) + } + + #[inline] + pub(crate) fn release(&mut self, tuple: TupleImpl) { + self.result.concat_pool.release(tuple); + } + + pub(crate) fn take_tuple(&mut self) -> TupleImpl { std::mem::take(&mut self.result.tuple) } - pub(crate) fn rewrite( + pub(crate) fn rewrite( &mut self, exprs: &[E], arena: &dyn MetaArena, input: Option<&TInput>, ) -> Result<(), DatabaseError> { - let values = &mut self.result.tuple.values; - let base = values.len(); - values.reserve(exprs.len()); - + let base = { + let values = &mut self.result.tuple.rightmost_mut().values; + values.reserve(exprs.len()); + values.len() + }; for expr in exprs { let value = { - let input_values = &values[..base]; let expression = expr.expression(arena); match input { Some(input) => expression.eval(arena, Some(input)), - None => expression.eval(arena, Some(&input_values)), + None => expression.eval(arena, Some(&self.result.tuple)), } .map(|value| value.into_owned()) }; match value { - Ok(value) => values.push(value), + Ok(value) => self.result.tuple.push(value), Err(error) => { - values.truncate(base); + self.result.tuple.rightmost_mut().values.truncate(base); return Err(error); } } } - values.rotate_left(base); - values.truncate(exprs.len()); + let pk = self.result.tuple.take_pk(); + let tuple = self + .result + .concat_pool + .retain_rightmost(&mut self.result.tuple); + tuple.values.rotate_left(base); + tuple.values.truncate(exprs.len()); + tuple.pk = pk; Ok(()) } @@ -685,11 +856,25 @@ impl<'a, T: Transaction + 'a> ExecArena<'a, T> { } #[inline] - pub(crate) fn produce_tuple(&mut self, tuple: Tuple) { + pub(crate) fn produce(&mut self, tuple: TupleImpl) { self.result.tuple = tuple; self.resume(); } + pub(crate) fn reset( + &mut self, + id: ExecId, + plan_arena: &mut ExecMetaArena, + ) -> Result<(), DatabaseError> { + let slot = &self.nodes.items[id] as *const std::cell::RefCell>; + self.nodes.executing += 1; + // SAFETY: the same slot/RefCell invariant as next_tuple. Reset may recurse + // into inputs, but cannot relocate slots or overwrite an active node. + let result = unsafe { (&*slot).borrow_mut().reset(self, plan_arena) }; + self.nodes.executing -= 1; + result + } + pub(crate) fn next_tuple( &mut self, id: ExecId, @@ -715,7 +900,7 @@ impl<'a, T: Transaction + 'a> ExecArena<'a, T> { } } -pub(crate) trait ReadExecutor<'a, T: Transaction + 'a>: Sized { +pub(crate) trait ReadExecutor<'a, T: StorageTransaction + 'a>: Sized { type Input; fn into_executor( @@ -727,7 +912,7 @@ pub(crate) trait ReadExecutor<'a, T: Transaction + 'a>: Sized { ) -> ExecId; } -pub(crate) trait WriteExecutor<'a, T: Transaction + 'a>: Sized { +pub(crate) trait WriteExecutor<'a, T: StorageTransaction + 'a>: Sized { type Input; fn into_executor( @@ -747,7 +932,7 @@ pub(crate) fn build_read<'a, T>( transaction: &T, ) -> ExecId where - T: Transaction + 'a, + T: StorageTransaction + 'a, { macro_rules! read { ($executor:ty, $input:expr) => { @@ -801,6 +986,10 @@ where let (left, right) = plan.childrens.twins(); read!(MarkApply<'a>, (op, left, right)) } + Operator::IndexProbeJoin(op) => { + let (left, right) = plan.childrens.twins(); + read!(IndexProbeJoin<'a>, (op, left, right)) + } Operator::Join(op) => { let use_hash_join = matches!( &op.on, @@ -817,7 +1006,7 @@ where if use_hash_join { read!(HashJoin, (op, left, right)) } else { - read!(NestedLoopJoin<'a>, (op, left, right)) + read!(NestedLoopJoin, (op, left, right)) } } Operator::Project(op) => read!(Projection<'a>, (op, plan.childrens.only())), @@ -876,7 +1065,7 @@ pub(crate) fn build_write<'a, T>( transaction: &'a T, ) -> ExecId where - T: Transaction + 'a, + T: StorageTransaction + 'a, { arena.init_context(cache, transaction); let transaction_ref: &T = transaction; @@ -950,13 +1139,13 @@ mod test_utils { ) } - pub(crate) struct TestExecutor<'a, T: Transaction + 'a> { + pub(crate) struct TestExecutor<'a, T: StorageTransaction + 'a> { executor: Executor<'a, T>, plan_arena: ExecMetaArena>, } - impl TestExecutor<'_, T> { - pub(crate) fn next_tuple(&mut self) -> Result, DatabaseError> { + impl TestExecutor<'_, T> { + pub(crate) fn next_tuple(&mut self) -> Result, DatabaseError> { self.executor.next_tuple(&mut self.plan_arena) } } @@ -968,7 +1157,7 @@ mod test_utils { transaction: &'a T, ) -> TestExecutor<'a, T> where - T: Transaction + 'a, + T: StorageTransaction + 'a, E: ReadExecutor<'a, T>, { let mut arena = ExecArena::with_capacity(0); @@ -994,7 +1183,7 @@ mod test_utils { transaction: &'a T, ) -> TestExecutor<'a, T> where - T: Transaction + 'a, + T: StorageTransaction + 'a, E: WriteExecutor<'a, T>, { let mut arena = ExecArena::with_capacity(0); @@ -1012,14 +1201,14 @@ mod test_utils { } } - pub fn try_collect( + pub fn try_collect( executor: TestExecutor<'_, T>, - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { let mut executor = executor; let mut tuples = Vec::new(); while let Some(tuple) = executor.next_tuple()? { - tuples.push(tuple.clone()); + tuples.push(std::mem::take(tuple)); } Ok(tuples) } @@ -1065,5 +1254,41 @@ mod test { assert_eq!(arena.nodes.items.len(), 1); assert_eq!(arena.nodes.items.as_ptr(), slot); } + + #[test] + fn rewrite_projects_joined_row() -> Result<(), DatabaseError> { + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::types::value::DataValue; + use crate::types::LogicalType; + + let table_arena = TableArenaCell::default(); + let mut plan_arena = PlanArena::new(&table_arena); + let column = plan_arena.alloc_column(ColumnCatalog::new( + "c".to_string(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + )); + let exprs = [3, 0, 1].map(|position| { + plan_arena.alloc_expression(ScalarExpression::column_expr(column, position)) + }); + let ints = |values: [i32; 2]| Tuple::new(None, values.map(DataValue::Int32).to_vec()); + + let mut arena = ExecArena::<'_, MemoryTransaction>::with_capacity(0); + for base in [0, 10] { + arena.result.tuple = TupleImpl::concat( + ints([base, base + 1]).into(), + ints([base + 2, base + 3]).into(), + Some(DataValue::Int32(base)), + ); + arena.rewrite(&exprs, &plan_arena, None::<&[DataValue]>)?; + assert!(matches!(arena.result_tuple(), TupleImpl::Tuple(_))); + assert_eq!(arena.result_tuple().pk(), Some(&DataValue::Int32(base))); + assert_eq!( + *arena.result_tuple(), + [base + 3, base, base + 1].map(DataValue::Int32) + ); + } + Ok(()) + } } // GRCOV_EXCL_STOP diff --git a/src/execution/runs.rs b/src/execution/runs.rs new file mode 100644 index 00000000..4c1a8680 --- /dev/null +++ b/src/execution/runs.rs @@ -0,0 +1,213 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Sorted runs of DataValues, shared by the in-memory `Vec` and spilled implementations. + +use crate::errors::DatabaseError; +use crate::types::value::DataValue; +use std::cmp::Ordering; + +pub(crate) type RunsCompare = fn(&DataValue, &DataValue) -> Ordering; + +/// Accepts unsorted values; `finish` yields runs each sorted by `compare`. +pub(crate) trait RunsBuilder: Sized { + type Runs: Runs; + + fn new(compare: RunsCompare) -> Self; + + fn push(&mut self, value: DataValue) -> Result<(), DatabaseError>; + + fn finish(self) -> Result; +} + +pub(crate) trait Runs { + type Cursor<'a>: RunsCursor + where + Self: 'a; + + /// `fixed_width` is the encoded width shared by every value, sizing spilled reads. + fn cursor(&self, fixed_width: Option) -> Result, DatabaseError>; +} + +pub(crate) trait RunsCursor { + fn run_count(&self) -> usize; + + /// First index in `run` whose value does not compare `Less` against the target. + fn lower_bound_by( + &mut self, + run: usize, + compare: impl FnMut(&DataValue) -> Result, + ) -> Result; + + /// Visits `run` from `start`, stopping at the first `Some` result. + fn visit_run( + &mut self, + run: usize, + start: usize, + visit: impl FnMut(&DataValue) -> Result, DatabaseError>, + ) -> Result, DatabaseError>; + + fn contains_by( + &mut self, + mut compare: impl FnMut(&DataValue) -> Result, + ) -> Result { + for run in 0..self.run_count() { + let index = self.lower_bound_by(run, &mut compare)?; + if self + .visit_run(run, index, |value| { + Ok(Some(compare(value)? == Ordering::Equal)) + })? + .unwrap_or(false) + { + return Ok(true); + } + } + Ok(false) + } + + /// Visits every value of every run, stopping at the first `Some` result. + fn visit_all( + &mut self, + mut visit: impl FnMut(&DataValue) -> Result, DatabaseError>, + ) -> Result, DatabaseError> { + for run in 0..self.run_count() { + if let Some(result) = self.visit_run(run, 0, &mut visit)? { + return Ok(Some(result)); + } + } + Ok(None) + } +} + +pub(crate) fn lower_bound( + rows: usize, + mut compare_at: impl FnMut(usize) -> Result, +) -> Result { + let (mut low, mut high) = (0, rows); + while low < high { + let mid = low + (high - low) / 2; + if compare_at(mid)? == Ordering::Less { + low = mid + 1; + } else { + high = mid; + } + } + Ok(low) +} + +/// Builds a single in-memory run, sorted once by `finish`. +/// Spill builds use `SpillRuns` instead. +#[cfg(any(test, not(feature = "spill")))] +pub(crate) struct SortedVec { + values: Vec, + compare: RunsCompare, +} + +#[cfg(any(test, not(feature = "spill")))] +impl RunsBuilder for SortedVec { + type Runs = Vec; + + fn new(compare: RunsCompare) -> Self { + Self { + values: Vec::new(), + compare, + } + } + + fn push(&mut self, value: DataValue) -> Result<(), DatabaseError> { + self.values.push(value); + Ok(()) + } + + fn finish(mut self) -> Result, DatabaseError> { + self.values.sort_by(self.compare); + Ok(self.values) + } +} + +impl Runs for Vec { + type Cursor<'a> = &'a [DataValue]; + + fn cursor(&self, _: Option) -> Result, DatabaseError> { + Ok(self.as_slice()) + } +} + +/// A sorted slice is one run; an empty slice has none. +impl RunsCursor for &[DataValue] { + fn run_count(&self) -> usize { + usize::from(!self.is_empty()) + } + + fn lower_bound_by( + &mut self, + _: usize, + mut compare: impl FnMut(&DataValue) -> Result, + ) -> Result { + lower_bound(self.len(), |index| compare(&self[index])) + } + + fn visit_run( + &mut self, + _: usize, + start: usize, + mut visit: impl FnMut(&DataValue) -> Result, DatabaseError>, + ) -> Result, DatabaseError> { + for value in self.iter().skip(start) { + if let Some(result) = visit(value)? { + return Ok(Some(result)); + } + } + Ok(None) + } +} + +// GRCOV_EXCL_START +#[cfg(test)] +mod tests { + use super::*; + + fn compare(left: &DataValue, right: &DataValue) -> Ordering { + left.partial_cmp(right).unwrap() + } + + #[test] + fn vec_runs_sort_once_and_search() -> Result<(), DatabaseError> { + let mut builder = SortedVec::new(compare); + for value in [5, 1, 3, 1] { + builder.push(DataValue::Int32(value))?; + } + let runs = builder.finish()?; + let mut cursor = runs.cursor(None)?; + assert_eq!(cursor.run_count(), 1); + for target in 0..7 { + assert_eq!( + cursor.contains_by(|value| Ok(compare(value, &DataValue::Int32(target))))?, + [1, 3, 5].contains(&target) + ); + } + let mut visited = Vec::new(); + cursor.visit_all::<()>(|value| { + visited.push(value.clone()); + Ok(None) + })?; + assert_eq!(visited, [1, 1, 3, 5].map(DataValue::Int32).to_vec()); + let empty = SortedVec::new(compare).finish()?; + let mut cursor = empty.cursor(None)?; + assert_eq!(cursor.run_count(), 0); + assert!(!cursor.contains_by(|_| Ok(Ordering::Equal))?); + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/execution/spill/codec.rs b/src/execution/spill/codec.rs index 37ec0e17..dce65171 100644 --- a/src/execution/spill/codec.rs +++ b/src/execution/spill/codec.rs @@ -16,21 +16,21 @@ use super::SpillCodec; use crate::errors::DatabaseError; use crate::planner::operator::sort::SortField; use crate::planner::{ExecMetaArena, MetaArena, ScalarQueryRef}; -use crate::types::tuple::Tuple; +use crate::types::tuple::{Tuple, TupleImpl}; use crate::types::value::DataValue; use std::io::{Read, Write}; use std::mem::size_of; pub(crate) struct SortRow { pub(crate) sort_values: Vec, - pub(crate) tuple: Tuple, + pub(crate) tuple: TupleImpl, pub(crate) outer_values: Vec<(ScalarQueryRef, DataValue)>, } impl SortRow { pub(crate) fn new( sort_fields: &[SortField], - tuple: Tuple, + tuple: TupleImpl, arena: &mut ExecMetaArena, ) -> Result { let sort_values = sort_fields @@ -60,11 +60,17 @@ impl SpillCodec for SortRow { fn decode(reader: &mut R) -> Result { Ok(Self { sort_values: Vec::::decode(reader)?, - tuple: Tuple::decode(reader)?, + tuple: TupleImpl::decode(reader)?, outer_values: Vec::decode(reader)?, }) } + fn encoded_len(&self) -> Result { + Ok(self.sort_values.encoded_len()? + + self.tuple.encoded_len()? + + self.outer_values.encoded_len()?) + } + fn estimated_size(&self) -> usize { size_of::() .saturating_add( @@ -80,7 +86,7 @@ impl SpillCodec for SortRow { .saturating_add( self.tuple .estimated_size() - .saturating_sub(size_of::()), + .saturating_sub(size_of::()), ) } } @@ -95,6 +101,10 @@ impl SpillCodec for (A, B) { Ok((A::decode(reader)?, B::decode(reader)?)) } + fn encoded_len(&self) -> Result { + Ok(self.0.encoded_len()? + self.1.encoded_len()?) + } + fn estimated_size(&self) -> usize { size_of::() .saturating_add(self.0.estimated_size().saturating_sub(size_of::())) @@ -133,6 +143,10 @@ impl SpillCodec for ScalarQueryRef { }) } + fn encoded_len(&self) -> Result { + Ok(17) + } + fn estimated_size(&self) -> usize { size_of::() } @@ -147,6 +161,48 @@ impl SpillCodec for DataValue { Self::decode_reference_value(reader) } + fn encoded_len(&self) -> Result { + use crate::types::value::Utf8Type; + let payload = match self { + DataValue::Parameter { .. } | DataValue::InitValue { .. } => { + return Err(DatabaseError::InvalidValue( + "unbound parameter cannot be serialized".into(), + )) + } + DataValue::Null => 0, + DataValue::Boolean(_) | DataValue::Int8(_) | DataValue::UInt8(_) => 1, + DataValue::Int16(_) | DataValue::UInt16(_) => 2, + DataValue::Float32(_) + | DataValue::Int32(_) + | DataValue::UInt32(_) + | DataValue::Date32(_) => 4, + DataValue::Float64(_) + | DataValue::Int64(_) + | DataValue::UInt64(_) + | DataValue::Date64(_) => 8, + DataValue::Time32(..) => 12, + DataValue::Time64(..) => 17, + #[cfg(feature = "decimal")] + DataValue::Decimal(_) => 16, + DataValue::Binary { value, len } => { + let _: u32 = value.len().try_into()?; + 4 + value.len() + 1 + if len.is_some() { 4 } else { 0 } + } + DataValue::Utf8 { value, ty, .. } => { + let _: u32 = value.len().try_into()?; + 4 + value.len() + + 1 + + match ty { + Utf8Type::Variable(None) => 2, + Utf8Type::Variable(Some(_)) => 6, + Utf8Type::Fixed(_) => 5, + } + } + DataValue::Tuple(values) => values.encoded_len()?, + }; + Ok(1 + payload) // Every DataValue starts with a one-byte variant tag. + } + fn estimated_size(&self) -> usize { size_of::().saturating_add(estimated_dynamic_value_size(self)) } @@ -154,12 +210,7 @@ impl SpillCodec for DataValue { impl SpillCodec for Vec { fn encode(&self, writer: &mut W) -> Result<(), DatabaseError> { - let len: u32 = self.len().try_into()?; - writer.write_all(&len.to_le_bytes())?; - for value in self { - value.encode(writer)?; - } - Ok(()) + encode_seq(self.len(), self.iter(), writer) } fn decode(reader: &mut R) -> Result { @@ -173,6 +224,10 @@ impl SpillCodec for Vec { Ok(values) } + fn encoded_len(&self) -> Result { + seq_encoded_len(self.len(), self.iter()) + } + fn estimated_size(&self) -> usize { size_of::() .saturating_add(self.capacity().saturating_mul(size_of::())) @@ -191,10 +246,7 @@ impl SpillCodec for Option { writer.write_all(&[1])?; value.encode(writer) } - None => { - writer.write_all(&[0])?; - Ok(()) - } + None => Ok(writer.write_all(&[0])?), } } @@ -210,6 +262,13 @@ impl SpillCodec for Option { } } + fn encoded_len(&self) -> Result { + Ok(1 + match self { + Some(value) => value.encoded_len()?, + None => 0, + }) + } + fn estimated_size(&self) -> usize { size_of::().saturating_add( self.as_ref() @@ -232,6 +291,10 @@ impl SpillCodec for Tuple { )) } + fn encoded_len(&self) -> Result { + Ok(self.pk.encoded_len()? + self.values.encoded_len()?) + } + fn estimated_size(&self) -> usize { size_of::() .saturating_add( @@ -254,6 +317,61 @@ impl SpillCodec for Tuple { } } +impl SpillCodec for TupleImpl { + // Same layout as `Tuple`, so a spilled Concat decodes back as a flat tuple. + fn encode(&self, writer: &mut W) -> Result<(), DatabaseError> { + self.pk_slot().encode(writer)?; + encode_seq(self.len(), self.iter(), writer) + } + + fn decode(reader: &mut R) -> Result { + Tuple::decode(reader).map(TupleImpl::Tuple) + } + + fn encoded_len(&self) -> Result { + Ok(self.pk_slot().encoded_len()? + seq_encoded_len(self.len(), self.iter())?) + } + + fn estimated_size(&self) -> usize { + match self { + TupleImpl::Tuple(tuple) => tuple.estimated_size(), + TupleImpl::Concat(_) => size_of::() + .saturating_add(self.len().saturating_mul(size_of::())) + .saturating_add( + self.iter() + .chain(self.pk()) + .map(estimated_dynamic_value_size) + .fold(0usize, usize::saturating_add), + ), + } + } +} + +#[inline] +fn encode_seq<'v, T: SpillCodec + 'v, W: Write>( + len: usize, + values: impl IntoIterator, + writer: &mut W, +) -> Result<(), DatabaseError> { + let len: u32 = len.try_into()?; + writer.write_all(&len.to_le_bytes())?; + values + .into_iter() + .try_for_each(|value| value.encode(writer)) +} + +#[inline] +fn seq_encoded_len<'v, T: SpillCodec + 'v>( + len: usize, + values: impl IntoIterator, +) -> Result { + let _: u32 = len.try_into()?; + values.into_iter().try_fold( + size_of::(), + |len, value| Ok(len + value.encoded_len()?), + ) +} + fn estimated_dynamic_value_size(value: &DataValue) -> usize { match value { DataValue::Utf8 { value, .. } => value.capacity(), @@ -282,6 +400,7 @@ mod tests { fn round_trip(value: &T) -> Result<(), DatabaseError> { let mut bytes = Vec::new(); value.encode(&mut bytes)?; + assert_eq!(value.encoded_len()?, bytes.len()); let mut reader = bytes.as_slice(); assert_eq!(&T::decode(&mut reader)?, value); assert!(reader.is_empty()); @@ -354,6 +473,12 @@ mod tests { .encode(&mut Vec::new()), Err(DatabaseError::InvalidValue(_)) )); + assert!(DataValue::Tuple(vec![DataValue::Parameter { + id: 1, + ty: LogicalType::Integer + }]) + .encoded_len() + .is_err()); assert!(matches!( DataValue::decode(&mut &[255][..]), Err(DatabaseError::InvalidValue(_)) @@ -381,11 +506,23 @@ mod tests { round_trip(&(reference, DataValue::Null))?; } for pk in [None, Some(DataValue::new_utf8("primary".into()))] { - let tuple = Tuple::new(pk, values.clone()); + let tuple = Tuple::new(pk.clone(), values.clone()); round_trip(&tuple)?; + let (left, right) = values.split_at(values.len() / 2); + let joined = TupleImpl::concat( + Tuple::new(None, left.to_vec()).into(), + Tuple::new(None, right.to_vec()).into(), + pk, + ); + round_trip(&joined)?; + let mut flat_bytes = Vec::new(); + tuple.encode(&mut flat_bytes)?; + let mut joined_bytes = Vec::new(); + joined.encode(&mut joined_bytes)?; + assert_eq!(joined_bytes, flat_bytes); let row = SortRow { sort_values: vec![DataValue::Int32(2), DataValue::Null], - tuple, + tuple: joined, outer_values: vec![( ScalarQueryRef { arena_id: 0, @@ -397,6 +534,7 @@ mod tests { }; let mut bytes = Vec::new(); row.encode(&mut bytes)?; + assert_eq!(row.encoded_len()?, bytes.len()); let mut reader = bytes.as_slice(); let decoded = SortRow::decode(&mut reader)?; assert!(reader.is_empty()); diff --git a/src/execution/spill/mod.rs b/src/execution/spill/mod.rs index 1a293485..b602244c 100644 --- a/src/execution/spill/mod.rs +++ b/src/execution/spill/mod.rs @@ -21,11 +21,13 @@ use std::path::PathBuf; use std::sync::atomic::{AtomicU64, Ordering}; mod codec; +pub(crate) mod runs; pub(crate) use codec::SortRow; const DEFAULT_MAX_ROWS: usize = 1024; const DEFAULT_MAX_BYTES: usize = 1024 * 1024; +const DEFAULT_READ_BUFFER_SIZE: usize = 8 * 1024; static NEXT_SPILL_FILE_ID: AtomicU64 = AtomicU64::new(0); pub(crate) type SegmentOffset = u64; @@ -35,6 +37,9 @@ pub(crate) trait SpillCodec: Sized { fn decode(reader: &mut R) -> Result; + /// Exact byte length of encode(), without allocating or encoding the value. + fn encoded_len(&self) -> Result; + fn estimated_size(&self) -> usize; } @@ -55,7 +60,17 @@ struct WriteState<'on_flush, T: SpillCodec> { max_bytes: usize, } -type OnFlush<'on_flush, T> = Box) -> Result<(), DatabaseError> + 'on_flush>; +/// Tells the flush callback whether it prepares a memory tail or a disk segment. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum FlushMode { + /// Preparing an in-memory reader, including the tail of a spilled reader. + Memory, + /// Writing a segment to disk, whether explicitly or due to a size limit. + Disk, +} + +type OnFlush<'on_flush, T> = + Box) -> Result<(), DatabaseError> + 'on_flush>; enum ReadState { Memory(std::vec::IntoIter), @@ -96,7 +111,7 @@ impl<'on_flush, T: SpillCodec> SpillVec<'on_flush, T> { pub(crate) fn on_flush(mut self, on_flush: F) -> Self where - F: FnMut(&mut Vec) -> Result<(), DatabaseError> + 'on_flush, + F: FnMut(FlushMode, &mut Vec) -> Result<(), DatabaseError> + 'on_flush, { if let Ok(state) = &mut self.writer { state.on_flush = Some(Box::new(on_flush)); @@ -186,6 +201,14 @@ impl Iterator for SpillReader { } impl SpillReader { + /// The unconsumed memory suffix. Borrowing it prevents advancing the reader. + pub(crate) fn memory_tail(&self) -> &[T] { + match &self.state { + ReadState::Memory(rows) | ReadState::Spilled { tail: rows, .. } => rows.as_slice(), + ReadState::Exhausted(_) => &[], + } + } + /// Rewinds a fully flushed spill file. The reader must not contain an in-memory tail. pub(crate) fn reset(&mut self) -> Result<(), DatabaseError> { let ReadState::Spilled { reader, tail, .. } = &mut self.state else { @@ -201,17 +224,21 @@ impl SpillReader { reader.reset(0) } - pub(crate) fn open_segment_reader<'source>( - &'source self, - ) -> Result, T>, DatabaseError> { + pub(crate) fn open_segment_reader( + &self, + capacity: Option, + ) -> Result, T>, DatabaseError> { let ReadState::Spilled { _file_guard, .. } = &self.state else { return Err(DatabaseError::InvalidValue( "cannot open a segment reader for an in-memory SpillVec".to_string(), )); }; - Ok(SegmentReader::new(BufReader::new(File::open( - &_file_guard.path, - )?))) + Ok(SegmentReader::new(BufReader::with_capacity( + capacity + .unwrap_or(DEFAULT_READ_BUFFER_SIZE) + .clamp(1, DEFAULT_READ_BUFFER_SIZE), + File::open(&_file_guard.path)?, + ))) } } @@ -256,6 +283,20 @@ impl SegmentReader<'_, R, T> { } impl SegmentReader<'_, R, T> { + /// Reads a nested value without parsing a segment header. Call reset before + /// resuming segment iteration. Relative seeking retains BufReader's buffer. + pub(crate) fn read_at(&mut self, offset: u64) -> Result { + self.remaining_rows = 0; + let current = self.reader.stream_position()?; + match i64::try_from(i128::from(offset) - i128::from(current)) { + Ok(delta) if delta != i64::MIN => self.reader.seek_relative(delta)?, + _ => { + self.reader.seek(SeekFrom::Start(offset))?; + } + } + V::decode(&mut self.reader) + } + pub(crate) fn reset(&mut self, offset: SegmentOffset) -> Result<(), DatabaseError> { self.reader.seek(SeekFrom::Start(offset))?; self.remaining_rows = 0; @@ -313,7 +354,7 @@ impl WriteState<'_, T> { return Ok(None); } if let Some(on_flush) = &mut self.on_flush { - on_flush(&mut self.buffer)?; + on_flush(FlushMode::Disk, &mut self.buffer)?; } let Some(file) = self.file.as_mut() else { return Err(DatabaseError::InvalidValue( @@ -328,7 +369,7 @@ impl WriteState<'_, T> { fn into_read(mut self) -> Result, DatabaseError> { if let Some(on_flush) = &mut self.on_flush { - on_flush(&mut self.buffer)?; + on_flush(FlushMode::Memory, &mut self.buffer)?; } let Some(file) = self.file.take() else { return Ok(ReadState::Memory(self.buffer.into_iter())); @@ -496,9 +537,9 @@ mod tests { })?; let source = values.into_iter(); - let mut first_reader = source.open_segment_reader()?; + let mut first_reader = source.open_segment_reader(None)?; first_reader.reset(first)?; - let mut second_reader = source.open_segment_reader()?; + let mut second_reader = source.open_segment_reader(None)?; second_reader.reset(second)?; assert_eq!(first_reader.next().transpose()?, Some(row(1))); @@ -539,7 +580,7 @@ mod tests { let values = SpillVec::from(vec![row(1)]).into_iter(); assert!(matches!( - values.open_segment_reader(), + values.open_segment_reader(None), Err(DatabaseError::InvalidValue(message)) if message == "cannot open a segment reader for an in-memory SpillVec" )); @@ -608,6 +649,10 @@ mod tests { Err(DatabaseError::InvalidValue("encode failed".to_string())) } + fn encoded_len(&self) -> Result { + Err(DatabaseError::InvalidValue("encode failed".to_string())) + } + fn decode(_: &mut R) -> Result { Ok(Self) } @@ -642,6 +687,10 @@ mod tests { Err(DatabaseError::InvalidValue("decode failed".to_string())) } + fn encoded_len(&self) -> Result { + Ok(0) + } + fn estimated_size(&self) -> usize { 0 } @@ -697,20 +746,20 @@ mod tests { } std::fs::remove_file(path.expect("expected spilled reader"))?; - assert!(source.open_segment_reader().is_err()); + assert!(source.open_segment_reader(None).is_err()); Ok(()) } #[test] fn on_flush_runs_before_memory_read_and_segment_flush() -> Result<(), DatabaseError> { - fn reverse(rows: &mut [Vec]) -> Result<(), DatabaseError> { + fn reverse(_: FlushMode, rows: &mut [Vec]) -> Result<(), DatabaseError> { rows.reverse(); Ok(()) } let mut memory = SpillVec::new() .limit(3, usize::MAX) - .on_flush(|rows| reverse(rows)); + .on_flush(|mode, rows| reverse(mode, rows)); let _ = memory.push(row(1))?; let _ = memory.push(row(2))?; assert_eq!( @@ -720,7 +769,7 @@ mod tests { let mut spilled = SpillVec::new() .limit(2, usize::MAX) - .on_flush(|rows| reverse(rows)); + .on_flush(|mode, rows| reverse(mode, rows)); for value in 1..=4 { let _ = spilled.push(row(value))?; } @@ -731,6 +780,54 @@ mod tests { Ok(()) } + #[test] + fn on_flush_receives_memory_or_disk_mode() -> Result<(), DatabaseError> { + for (count, limit, explicit_flush, expected) in [ + (3, 4, false, vec![FlushMode::Memory]), + ( + 8, + 3, + false, + vec![FlushMode::Disk, FlushMode::Disk, FlushMode::Memory], + ), + (4, 5, true, vec![FlushMode::Disk, FlushMode::Memory]), + ] { + let mut modes = Vec::new(); + let mut values = SpillVec::new() + .limit(limit, usize::MAX) + .on_flush(|mode, _| { + modes.push(mode); + Ok(()) + }); + for value in 1..=count { + values.push(row(value))?; + if explicit_flush && value == 2 { + values.flush()?; + } + } + assert_eq!( + values.into_iter().collect::, _>>()?, + (1..=count).map(row).collect::>() + ); + assert_eq!(modes, expected); + } + Ok(()) + } + + #[test] + fn on_flush_errors_propagate_from_memory_and_disk() -> Result<(), DatabaseError> { + let callback = |_: FlushMode, _: &mut Vec>| { + Err(DatabaseError::InvalidValue("callback invoked".into())) + }; + let mut memory = SpillVec::new().limit(2, usize::MAX).on_flush(callback); + memory.push(row(1))?; + assert!(memory.into_iter().next().unwrap().is_err()); + + let mut disk = SpillVec::new().limit(1, usize::MAX).on_flush(callback); + assert!(disk.push(row(1)).is_err()); + Ok(()) + } + #[test] fn tuple_codec_preserves_primary_key_and_values() -> Result<(), DatabaseError> { let tuple = Tuple::new( diff --git a/src/execution/spill/runs.rs b/src/execution/spill/runs.rs new file mode 100644 index 00000000..2e6b8617 --- /dev/null +++ b/src/execution/spill/runs.rs @@ -0,0 +1,501 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{FlushMode, SegmentOffset, SegmentReader, SpillCodec, SpillReader, SpillVec}; +use crate::errors::DatabaseError; +use crate::execution::runs::{lower_bound, Runs, RunsBuilder, RunsCompare, RunsCursor}; +use crate::types::value::DataValue; +use std::cmp::Ordering; +use std::fs::File; +use std::io::BufReader; +use std::mem::size_of; + +pub(crate) struct SpillRuns { + runs: Vec, + spill: SpillVec<'static, DataValue>, +} + +/// Disk layout, entirely encoded as ordinary DataValues: +/// UInt32(N), N UInt32 offsets relative to the records area, N original values. +/// Extend and rotate the existing allocation; no scratch encoding or heap Vec. +fn prepare_disk_run(values: &mut Vec) -> Result<(), DatabaseError> { + let rows = values.len(); + let count = u32::try_from(rows)?; + let slots = rows + .checked_add(1) + .ok_or_else(|| DatabaseError::InvalidValue("run length overflow".into()))?; + let len = rows + .checked_add(slots) + .ok_or_else(|| DatabaseError::InvalidValue("run length overflow".into()))?; + values.resize(len, DataValue::Null); + values.rotate_right(slots); + values[0] = DataValue::UInt32(count); + let mut offset = 0u32; + for index in 0..rows { + let bytes = u32::try_from(values[slots + index].encoded_len()?)?; + values[index + 1] = DataValue::UInt32(offset); + offset = offset.checked_add(bytes).ok_or_else(|| { + DatabaseError::InvalidValue("run payload exceeds UInt32 offsets".into()) + })?; + } + Ok(()) +} + +impl SpillRuns { + /// Ends the current run early; later values start a new one. + #[cfg(test)] + pub(crate) fn flush(&mut self) -> Result<(), DatabaseError> { + if let Some(offset) = self.spill.flush()? { + self.runs.push(offset); + } + Ok(()) + } +} + +/// Every SpillVec flush sorts its bounded segment into one run. +impl RunsBuilder for SpillRuns { + type Runs = SpillRunsReader; + + fn new(compare: RunsCompare) -> Self { + Self { + runs: Vec::new(), + spill: SpillVec::new().on_flush(move |mode, values| { + values.sort_by(compare); + match mode { + FlushMode::Disk => prepare_disk_run(values), + FlushMode::Memory => Ok(()), + } + }), + } + } + + fn push(&mut self, value: DataValue) -> Result<(), DatabaseError> { + if let Some(offset) = self.spill.push(value)? { + self.runs.push(offset); + } + Ok(()) + } + + fn finish(self) -> Result { + Ok(SpillRunsReader { + runs: self.runs, + spill: self.spill.into_iter(), + }) + } +} + +pub(crate) struct SpillRunsReader { + runs: Vec, + spill: SpillReader, +} + +impl Runs for SpillRunsReader { + type Cursor<'a> = SpillRunsCursor<'a>; + + fn cursor(&self, fixed_width: Option) -> Result, DatabaseError> { + let memory = self.spill.memory_tail(); + if self.runs.is_empty() { + return Ok(SpillRunsCursor::Memory(memory)); + } + let directory_width = DataValue::UInt32(0).encoded_len()?; + let reader = self + .spill + .open_segment_reader(fixed_width.map(|width| width.max(directory_width)))?; + Ok(SpillRunsCursor::Spilled { + runs: &self.runs, + memory, + reader, + }) + } +} + +pub(crate) enum SpillRunsCursor<'a> { + Memory(&'a [DataValue]), + Spilled { + runs: &'a [SegmentOffset], + memory: &'a [DataValue], + reader: SegmentReader<'a, BufReader, DataValue>, + }, +} + +fn disk_value( + reader: &mut SegmentReader<'_, BufReader, DataValue>, + index: usize, + rows: usize, + directory: u64, +) -> Result { + let width = DataValue::UInt32(0).encoded_len()? as u64; + let DataValue::UInt32(offset) = reader.read_at(directory + index as u64 * width)? else { + return Err(DatabaseError::InvalidValue("invalid run offset".into())); + }; + // TODO: accept single-BufReader directory/record jumps for now. Profile + // before caching the directory or introducing a second reader. + reader.read_at(directory + rows as u64 * width + u64::from(offset)) +} + +fn disk_layout( + reader: &mut SegmentReader<'_, BufReader, DataValue>, + offset: SegmentOffset, +) -> Result<(usize, u64), DatabaseError> { + let offset = offset + size_of::() as u64; + let DataValue::UInt32(rows) = reader.read_at(offset)? else { + return Err(DatabaseError::InvalidValue("invalid run row count".into())); + }; + Ok(( + rows as usize, + offset + DataValue::UInt32(0).encoded_len()? as u64, + )) +} + +impl RunsCursor for SpillRunsCursor<'_> { + fn run_count(&self) -> usize { + match self { + Self::Memory(memory) => memory.run_count(), + Self::Spilled { runs, memory, .. } => runs.len() + memory.run_count(), + } + } + + fn lower_bound_by( + &mut self, + run: usize, + mut compare: impl FnMut(&DataValue) -> Result, + ) -> Result { + match self { + Self::Memory(memory) => memory.lower_bound_by(run, compare), + Self::Spilled { runs, memory, .. } if run == runs.len() => { + memory.lower_bound_by(0, compare) + } + Self::Spilled { runs, reader, .. } => { + let (rows, directory) = disk_layout(reader, runs[run])?; + lower_bound(rows, |index| { + compare(&disk_value(reader, index, rows, directory)?) + }) + } + } + } + + fn visit_run( + &mut self, + run: usize, + start: usize, + mut visit: impl FnMut(&DataValue) -> Result, DatabaseError>, + ) -> Result, DatabaseError> { + match self { + Self::Memory(memory) => memory.visit_run(run, start, visit), + Self::Spilled { runs, memory, .. } if run == runs.len() => { + memory.visit_run(0, start, visit) + } + Self::Spilled { runs, reader, .. } => { + let (rows, directory) = disk_layout(reader, runs[run])?; + for index in start..rows { + if let Some(result) = visit(&disk_value(reader, index, rows, directory)?)? { + return Ok(Some(result)); + } + } + Ok(None) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn compare(left: &DataValue, right: &DataValue) -> Ordering { + match (left, right) { + (DataValue::Null, DataValue::Null) => Ordering::Equal, + (DataValue::Null, _) => Ordering::Less, + (_, DataValue::Null) => Ordering::Greater, + _ => left.partial_cmp(right).unwrap(), + } + } + + fn descending(left: &DataValue, right: &DataValue) -> Ordering { + compare(right, left) + } + + fn runs_of( + compare: RunsCompare, + batches: impl IntoIterator>, + ) -> Result { + let mut runs = SpillRuns::new(compare); + for batch in batches { + runs.flush()?; + for value in batch { + runs.push(value)?; + } + } + Ok(runs) + } + + #[test] + fn original_memory_and_disk_values_match() -> Result<(), DatabaseError> { + for values in [ + vec![ + DataValue::Int64(3), + DataValue::Null, + DataValue::Int64(-2), + DataValue::Int64(3), + ], + vec![ + DataValue::new_utf8("äø­ę–‡".repeat(100)), + DataValue::Null, + DataValue::new_utf8("".into()), + DataValue::new_utf8("abc".into()), + ], + vec![ + DataValue::Tuple(vec![DataValue::Int64(2), DataValue::Null]), + DataValue::Null, + DataValue::Tuple(vec![DataValue::Int64(1), DataValue::Null]), + ], + vec![DataValue::Null; 9], + ] { + let mut expected = values.clone(); + expected.sort_by(compare); + for flushed in [false, true] { + let mut runs = runs_of(compare, [values.clone()])?; + if flushed { + runs.flush()?; + } + let source = runs.finish()?; + if !flushed { + assert_eq!(source.spill.memory_tail(), expected); + } + for _ in 0..2 { + let mut cursor = source.cursor(None)?; + let mut actual = Vec::new(); + cursor.visit_run::<()>(0, 0, |value| { + actual.push(value.clone()); + Ok(None) + })?; + assert_eq!(actual, expected); + for target in &expected { + assert!(cursor.contains_by(|value| Ok(compare(value, target)))?); + } + assert_eq!( + cursor.lower_bound_by(0, |_| Ok(Ordering::Less))?, + expected.len() + ); + } + } + } + Ok(()) + } + + #[test] + fn disk_runs_and_original_memory_tail() -> Result<(), DatabaseError> { + let source = runs_of( + compare, + [ + vec![DataValue::Int64(4), DataValue::Int64(2)], + vec![DataValue::Int64(7), DataValue::Int64(1)], + vec![DataValue::Int64(9)], + ], + )? + .finish()?; + let mut cursor = source.cursor(Some(9))?; + assert_eq!(cursor.run_count(), 3); + assert!( + matches!(&cursor, SpillRunsCursor::Spilled { runs, memory, .. } if runs.len() == 2 && memory.len() == 1) + ); + assert_eq!(source.spill.memory_tail(), &[DataValue::Int64(9)]); + for target in 0..11 { + assert_eq!( + cursor.contains_by(|value| Ok(compare(value, &DataValue::Int64(target))))?, + [1, 2, 4, 7, 9].contains(&target) + ); + } + let mut other = source.cursor(Some(9))?; + assert_eq!( + other.lower_bound_by(1, |value| Ok(compare(value, &DataValue::Int64(6))))?, + 1 + ); + assert!(cursor.contains_by(|value| Ok(compare(value, &DataValue::Int64(2))))?); + Ok(()) + } + + #[test] + fn automatic_spill_and_928_memory_keys() -> Result<(), DatabaseError> { + let source = runs_of( + compare, + [(0..super::super::DEFAULT_MAX_ROWS + 1) + .rev() + .map(|value| DataValue::Int64(value as i64)) + .collect()], + )? + .finish()?; + let mut cursor = source.cursor(Some(9))?; + assert_eq!(cursor.run_count(), 2); + for target in [0, 1, 1023, 1024] { + assert!(cursor.contains_by(|value| Ok(compare(value, &DataValue::Int64(target))))?); + } + let runs = runs_of(compare, [(0..928).rev().map(DataValue::UInt32).collect()])?; + assert!(!runs.spill.is_spilled()); + let source = runs.finish()?; + let mut cursor = source.cursor(Some(5))?; + assert!(matches!(&cursor, SpillRunsCursor::Memory(_))); + let mut comparisons = 0; + assert!(cursor.contains_by(|value| { + comparisons += 1; + Ok(compare(value, &DataValue::UInt32(900))) + })?); + assert!(comparisons <= 12); + assert_eq!( + SpillRuns::new(compare).finish()?.cursor(None)?.run_count(), + 0 + ); + Ok(()) + } + + #[test] + fn cursor_uses_one_fixed_record_capped_at_default() -> Result<(), DatabaseError> { + let default = super::super::DEFAULT_READ_BUFFER_SIZE; + for (value, expected) in [ + (DataValue::Boolean(true), 5), + (DataValue::Int64(42), 9), + (DataValue::Tuple(vec![DataValue::Int32(1); 5]), 30), + (DataValue::Tuple(vec![DataValue::Int64(1); 5]), 50), + ] { + let width = value.encoded_len()?; + let mut runs = runs_of( + compare, + [vec![value.clone(), DataValue::Null, value.clone()]], + )?; + runs.flush()?; + let source = runs.finish()?; + let mut cursor = source.cursor(Some(width))?; + let SpillRunsCursor::Spilled { reader, .. } = &cursor else { + panic!("expected disk cursor"); + }; + assert_eq!(reader.reader.capacity(), expected); + assert!(cursor.contains_by(|key| Ok(compare(key, &value)))?); + assert!(cursor.contains_by(|key| Ok(compare(key, &DataValue::Null)))?); + for (width, expected) in [(None, default), (Some(0), 5), (Some(default + 1), default)] { + let cursor = source.cursor(width)?; + let SpillRunsCursor::Spilled { reader, .. } = &cursor else { + panic!("expected disk cursor"); + }; + assert_eq!(reader.reader.capacity(), expected); + } + } + Ok(()) + } + + #[test] + fn directory_offsets_match_original_encoding() -> Result<(), DatabaseError> { + let original = vec![ + DataValue::Null, + DataValue::Int8(1), + DataValue::Int64(9), + DataValue::new_utf8("äø­ę–‡".into()), + DataValue::Tuple(vec![DataValue::Null, DataValue::UInt32(4)]), + ]; + let mut values = original.clone(); + prepare_disk_run(&mut values)?; + assert_eq!(values.len(), original.len() * 2 + 1); + assert_eq!(values[0], DataValue::UInt32(original.len() as u32)); + assert_eq!(&values[original.len() + 1..], original); + let mut offset = 0; + for (index, value) in original.iter().enumerate() { + assert_eq!(values[index + 1], DataValue::UInt32(offset)); + let mut bytes = Vec::new(); + value.encode(&mut bytes)?; + offset += bytes.len() as u32; + } + let mut empty = Vec::new(); + prepare_disk_run(&mut empty)?; + assert_eq!(empty, [DataValue::UInt32(0)]); + Ok(()) + } + + #[test] + fn null_last_and_descending_runs() -> Result<(), DatabaseError> { + for flushed in [false, true] { + let mut runs = runs_of( + descending, + [vec![ + DataValue::Int64(1), + DataValue::Null, + DataValue::Int64(5), + ]], + )?; + if flushed { + runs.flush()?; + } + let source = runs.finish()?; + let mut cursor = source.cursor(Some(9))?; + let mut actual = Vec::new(); + cursor.visit_run::<()>(0, 0, |value| { + actual.push(value.clone()); + Ok(None) + })?; + assert_eq!( + actual, + [DataValue::Int64(5), DataValue::Int64(1), DataValue::Null] + ); + for target in [DataValue::Int64(1), DataValue::Null, DataValue::Int64(5)] { + assert!(cursor.contains_by(|value| Ok(compare(&target, value)))?); + } + assert!(!cursor.contains_by(|value| Ok(compare(&DataValue::Int64(3), value)))?); + } + Ok(()) + } + + #[test] + fn directory_entries_stay_outside_search_range() -> Result<(), DatabaseError> { + let mut runs = runs_of( + compare, + [vec![ + DataValue::new_utf8("a".into()), + DataValue::new_utf8("c".into()), + ]], + )?; + runs.flush()?; + let source = runs.finish()?; + let mut cursor = source.cursor(None)?; + assert_eq!(cursor.lower_bound_by(0, |_| Ok(Ordering::Less))?, 2); + assert!(!cursor.contains_by(|value| Ok(compare(value, &DataValue::new_utf8("z".into()))))?); + Ok(()) + } + + #[test] + fn spill_limits_cut_sorted_runs() -> Result<(), DatabaseError> { + let max_rows = super::super::DEFAULT_MAX_ROWS; + let rows = max_rows * 2 + 3; + let mut runs = SpillRuns::new(compare); + for value in (0..rows as i64).rev() { + runs.push(DataValue::Int64(value))?; + } + let source = runs.finish()?; + let mut cursor = source.cursor(Some(9))?; + // Two full disk runs plus the sorted memory tail. + assert_eq!(cursor.run_count(), 3); + for run in 0..cursor.run_count() { + let mut previous: Option = None; + cursor.visit_run::<()>(run, 0, |value| { + if let Some(previous) = &previous { + assert_eq!(compare(previous, value), Ordering::Less); + } + previous = Some(value.clone()); + Ok(None) + })?; + } + for target in [0, 2, max_rows as i64, rows as i64 - 1] { + assert!(cursor.contains_by(|value| Ok(compare(value, &DataValue::Int64(target))))?); + } + assert!(!cursor.contains_by(|value| Ok(compare(value, &DataValue::Int64(-1))))?); + Ok(()) + } +} diff --git a/src/expression/agg.rs b/src/expression/agg.rs index 13198aef..8f8b4e7a 100644 --- a/src/expression/agg.rs +++ b/src/expression/agg.rs @@ -21,6 +21,10 @@ pub enum AggKind { Min, Sum, Count, + /// SQLite `total()`: a float sum that is 0.0, not NULL, over no rows. + Total, + /// SQLite `group_concat()`: non-NULL values joined with `,`. + GroupConcat, } impl AggKind { @@ -31,6 +35,8 @@ impl AggKind { "min" => Some(Self::Min), "sum" => Some(Self::Sum), "count" => Some(Self::Count), + "total" => Some(Self::Total), + "group_concat" => Some(Self::GroupConcat), _ => None, } } @@ -42,6 +48,8 @@ impl AggKind { Self::Min => "min", Self::Sum => "sum", Self::Count => "count", + Self::Total => "total", + Self::GroupConcat => "group_concat", } } @@ -52,6 +60,8 @@ impl AggKind { AggKind::Min => false, AggKind::Sum => true, AggKind::Count => true, + AggKind::Total => true, + AggKind::GroupConcat => true, } } } diff --git a/src/expression/eq_col.rs b/src/expression/eq_col.rs index 6df7123b..14e3b379 100644 --- a/src/expression/eq_col.rs +++ b/src/expression/eq_col.rs @@ -30,22 +30,22 @@ use crate::types::evaluator::{BinaryEvaluatorRef, CastEvaluatorRef, UnaryEvaluat use crate::types::value::DataValue; use crate::types::LogicalType; -pub(super) fn eq_ignore_colref_pos( +pub(super) fn eq_ignore_colref_pos( lhs: ExprRef, rhs: ExprRef, - arena: &(dyn MetaArena + '_), + arena: &A, ) -> bool { EqIgnoreColRefPosVisitor::equals(lhs, rhs, arena) } -struct EqIgnoreColRefPosVisitor<'a, 'arena> { +struct EqIgnoreColRefPosVisitor<'a, A: ?Sized> { rhs: ExprRef, - arena: &'a (dyn MetaArena + 'arena), + arena: &'a A, equal: bool, } -impl<'a, 'arena> EqIgnoreColRefPosVisitor<'a, 'arena> { - fn equals(lhs: ExprRef, rhs: ExprRef, arena: &'a (dyn MetaArena + 'arena)) -> bool { +impl<'a, A: MetaArena + ?Sized> EqIgnoreColRefPosVisitor<'a, A> { + fn equals(lhs: ExprRef, rhs: ExprRef, arena: &'a A) -> bool { let mut visitor = Self { rhs, arena, @@ -75,8 +75,8 @@ impl<'a, 'arena> EqIgnoreColRefPosVisitor<'a, 'arena> { } } -impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { - fn visit(&mut self, lhs: ExprRef, arena: &(dyn MetaArena + '_)) -> Result<(), DatabaseError> { +impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, A> { + fn visit(&mut self, lhs: ExprRef, arena: &A) -> Result<(), DatabaseError> { let lhs = lhs.unpack_alias(arena); self.rhs = self.rhs.unpack_alias(arena); if lhs == self.rhs { @@ -89,7 +89,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { &mut self, id: ScalarQueryRef, ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = matches!(self.rhs(), ScalarExpression::OuterParam { id: rhs, ty: rhs_ty } if id == *rhs && ty == rhs_ty); Ok(()) @@ -99,7 +99,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { &mut self, id: ScalarQueryRef, ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = matches!(self.rhs(), ScalarExpression::OuterValue { id: rhs, ty: rhs_ty } if id == *rhs && ty == rhs_ty); Ok(()) @@ -109,7 +109,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { &mut self, id: ScalarQueryRef, ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::InitValue { @@ -126,7 +126,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { Ok(()) } - fn visit_column_ref(&mut self, lhs: &ColumnRef) -> Result<(), DatabaseError> { + fn visit_column_ref(&mut self, lhs: &ColumnRef, _: usize) -> Result<(), DatabaseError> { self.equal = matches!( self.rhs(), ScalarExpression::ColumnRef { column: rhs, .. } @@ -140,7 +140,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_expr: ExprRef, lhs_ty: &LogicalType, lhs_evaluator: Option<&CastEvaluatorRef>, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::TypeCast { @@ -161,7 +161,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { &mut self, lhs_negated: bool, lhs_expr: ExprRef, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::IsNull { @@ -179,7 +179,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_expr: ExprRef, lhs_evaluator: Option<&UnaryEvaluatorRef>, lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::Unary { @@ -205,7 +205,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_right: ExprRef, lhs_evaluator: Option<&BinaryEvaluatorRef>, lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::Binary { @@ -232,7 +232,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_kind: &AggKind, lhs_args: &[ExprRef], lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::AggCall { @@ -251,11 +251,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { Ok(()) } - fn visit_window( - &mut self, - lhs: &WindowCall, - _arena: &(dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { + fn visit_window(&mut self, lhs: &WindowCall, _arena: &A) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::WindowCall(rhs) => { lhs.function.kind == rhs.function.kind @@ -281,7 +277,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_negated: bool, lhs_expr: ExprRef, lhs_args: &[ExprRef], - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::In { @@ -305,7 +301,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_expr: ExprRef, lhs_left: ExprRef, lhs_right: ExprRef, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::Between { @@ -330,7 +326,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_expr: ExprRef, lhs_for: Option, lhs_from: Option, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::SubString { @@ -351,7 +347,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { &mut self, lhs_expr: ExprRef, lhs_in: ExprRef, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::Position { @@ -371,7 +367,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_expr: ExprRef, lhs_what: Option, lhs_where: Option<&TrimWhereField>, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::Trim { @@ -393,11 +389,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { Ok(()) } - fn visit_tuple( - &mut self, - lhs: &[ExprRef], - _arena: &(dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { + fn visit_tuple(&mut self, lhs: &[ExprRef], _arena: &A) -> Result<(), DatabaseError> { self.equal = matches!(self.rhs(), ScalarExpression::Tuple(rhs) if self.refs_equal(lhs, rhs)); Ok(()) @@ -406,7 +398,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { fn visit_scala_function( &mut self, lhs: &ScalarFunction, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = matches!( self.rhs(), @@ -419,7 +411,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { fn visit_table_function( &mut self, lhs: &TableFunction, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = matches!( self.rhs(), @@ -435,7 +427,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_left: ExprRef, lhs_right: ExprRef, lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::If { @@ -459,7 +451,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_left: ExprRef, lhs_right: ExprRef, lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::IfNull { @@ -481,7 +473,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_left: ExprRef, lhs_right: ExprRef, lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::NullIf { @@ -502,7 +494,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { &mut self, lhs_exprs: &[ExprRef], lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::Coalesce { @@ -520,7 +512,7 @@ impl ExprVisitor for EqIgnoreColRefPosVisitor<'_, '_> { lhs_pairs: &[(ExprRef, ExprRef)], lhs_else: Option, lhs_ty: &LogicalType, - _arena: &(dyn MetaArena + '_), + _arena: &A, ) -> Result<(), DatabaseError> { self.equal = match self.rhs() { ScalarExpression::CaseWhen { diff --git a/src/expression/evaluator.rs b/src/expression/evaluator.rs index 51212489..7911e33d 100644 --- a/src/expression/evaluator.rs +++ b/src/expression/evaluator.rs @@ -18,10 +18,10 @@ use crate::planner::PlanArena; use crate::errors::DatabaseError; use crate::expression::function::scala::ScalarFunction; -use crate::expression::{AliasType, BinaryOperator, ScalarExpression, TrimWhereField}; +use crate::expression::{BinaryOperator, ScalarExpression, TrimWhereField}; use crate::planner::MetaArena; use crate::types::evaluator::binary_create; -use crate::types::tuple::TupleLike; +use crate::types::tuple::{Input, InputRef}; use crate::types::value::{DataValue, Utf8Type}; use crate::types::{CharLengthUnits, LogicalType}; use std::borrow::Cow; @@ -43,7 +43,7 @@ macro_rules! eval_to_num { } impl ScalarExpression { - pub fn eval<'a, T: TupleLike + ?Sized>( + pub fn eval<'a, T: Input + ?Sized>( &'a self, arena: &'a (dyn MetaArena + '_), tuple: Option<&'a T>, @@ -67,17 +67,14 @@ impl ScalarExpression { let Some(tuple) = tuple else { return Ok(Cow::Owned(DataValue::Null)); }; - Ok(Cow::Borrowed(tuple.value_at(*position))) + Ok(tuple.column(*position)) } - ScalarExpression::Alias { expr, alias } => { + // `AliasType::Expr` only carries the output identity. + ScalarExpression::Alias { expr, .. } => { let Some(tuple) = tuple else { return Ok(Cow::Owned(DataValue::Null)); }; - if let AliasType::Expr(inner_expr) = alias { - arena.expression(*inner_expr).eval(arena, Some(tuple)) - } else { - arena.expression(*expr).eval(arena, Some(tuple)) - } + arena.expression(*expr).eval(arena, Some(tuple)) } ScalarExpression::TypeCast { expr, evaluator, .. @@ -280,7 +277,7 @@ impl ScalarExpression { } ScalarExpression::ScalaFunction(ScalarFunction { inner, args, .. }) => { let value = match tuple { - Some(tuple) => inner.eval(args, arena, Some(&tuple))?, + Some(tuple) => inner.eval(args, arena, Some(&InputRef(tuple)))?, None => inner.eval(args, arena, None)?, }; value.cast(inner.return_type()).map(Cow::Owned) @@ -419,6 +416,7 @@ fn trim_string(value: &str, trim_what: &str, trim_where: Option) #[cfg(test)] mod tests { use super::*; + use crate::expression::AliasType; use crate::planner::test::PlanArenaTestExt; use crate::planner::{ExecMetaArena, ExprRef, TableArenaCell}; diff --git a/src/expression/function/scala.rs b/src/expression/function/scala.rs index c2e755da..9160fb15 100644 --- a/src/expression/function/scala.rs +++ b/src/expression/function/scala.rs @@ -16,7 +16,7 @@ use crate::errors::DatabaseError; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::LogicalType; use kite_sql_serde_macros::ReferenceSerialization; @@ -66,7 +66,7 @@ pub trait ScalarFunctionImpl: Debug + Send + Sync { &self, args: &[ExprRef], arena: &(dyn MetaArena + '_), - tuple: Option<&dyn TupleLike>, + tuple: Option<&dyn Input>, ) -> Result; // TODO: Exploiting monotonicity when optimizing `ScalarFunctionImpl::monotonicity()` diff --git a/src/expression/mod.rs b/src/expression/mod.rs index af3629c9..a58f2464 100644 --- a/src/expression/mod.rs +++ b/src/expression/mod.rs @@ -292,13 +292,13 @@ mod chrono_scalar_expression { pub struct BindEvaluator; -impl ExprVisitorMut for BindEvaluator { +impl ExprVisitorMut for BindEvaluator { fn visit_type_cast( &mut self, expr: &mut ExprRef, ty: &mut LogicalType, evaluator: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; let from = expr.return_type(arena); @@ -317,7 +317,7 @@ impl ExprVisitorMut for BindEvaluator { expr: &mut ExprRef, evaluator: &mut Option, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; @@ -344,7 +344,7 @@ impl ExprVisitorMut for BindEvaluator { right_expr: &mut ExprRef, evaluator: &mut Option, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(left_expr, arena)?; self.visit(right_expr, arena)?; @@ -366,7 +366,7 @@ impl ExprVisitorMut for BindEvaluator { expr: &mut ExprRef, args: &mut [ExprRef], evaluator: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; let mut ty = expr.return_type(arena).into_owned(); @@ -391,7 +391,7 @@ impl ExprVisitorMut for BindEvaluator { right_expr: &mut ExprRef, left_evaluator: &mut Option, right_evaluator: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; self.visit(left_expr, arena)?; @@ -442,19 +442,19 @@ impl ExprVisitor for HasCountStar { } pub trait TypeCast: Sized { - fn return_type<'a>(&'a self, arena: &'a (dyn MetaArena + '_)) -> Cow<'a, LogicalType>; + fn return_type<'a, A: MetaArena + ?Sized>(&'a self, arena: &'a A) -> Cow<'a, LogicalType>; - fn into_expr( + fn into_expr( self, ty: LogicalType, evaluator: CastEvaluatorRef, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Self; - fn type_cast( + fn type_cast( self, ty: Cow<'_, LogicalType>, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result { let from = self.return_type(arena); if from.as_ref() == ty.as_ref() { @@ -466,7 +466,7 @@ pub trait TypeCast: Sized { } impl TypeCast for ScalarExpression { - fn return_type<'a>(&'a self, arena: &'a (dyn MetaArena + '_)) -> Cow<'a, LogicalType> { + fn return_type<'a, A: MetaArena + ?Sized>(&'a self, arena: &'a A) -> Cow<'a, LogicalType> { match self { ScalarExpression::Constant(value) => Cow::Owned(value.logical_type()), ScalarExpression::ColumnRef { column, .. } => { @@ -509,11 +509,11 @@ impl TypeCast for ScalarExpression { } } - fn into_expr( + fn into_expr( self, ty: LogicalType, evaluator: CastEvaluatorRef, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Self { ScalarExpression::TypeCast { expr: arena.alloc_expression(self), @@ -524,15 +524,15 @@ impl TypeCast for ScalarExpression { } impl TypeCast for ExprRef { - fn return_type<'a>(&'a self, arena: &'a (dyn MetaArena + '_)) -> Cow<'a, LogicalType> { + fn return_type<'a, A: MetaArena + ?Sized>(&'a self, arena: &'a A) -> Cow<'a, LogicalType> { arena.expression(*self).return_type(arena) } - fn into_expr( + fn into_expr( self, ty: LogicalType, evaluator: CastEvaluatorRef, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Self { arena.alloc_expression(ScalarExpression::TypeCast { expr: self, @@ -549,10 +549,10 @@ impl ScalarExpression { } impl Explain for ExprRef { - fn fmt(&self, arena: &dyn MetaArena, f: &mut fmt::Formatter<'_>) -> fmt::Result { - fn write_exprs( + fn fmt(&self, arena: &A, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn write_exprs( exprs: &[ExprRef], - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut fmt::Formatter<'_>, ) -> fmt::Result { for (index, expr) in exprs.iter().enumerate() { @@ -804,13 +804,17 @@ impl ExprRef { SortField::from(self).nulls_last() } - pub(crate) fn eq_ignore_colref_pos(self, other: ExprRef, arena: &(dyn MetaArena + '_)) -> bool { + pub(crate) fn eq_ignore_colref_pos( + self, + other: ExprRef, + arena: &A, + ) -> bool { eq_col::eq_ignore_colref_pos(self, other, arena) } - pub(crate) fn clone_expression( + pub(crate) fn clone_expression( self, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result { let mut cloned = self; crate::expression::visitor_mut::ExprCloner.visit(&mut cloned, arena)?; @@ -818,13 +822,7 @@ impl ExprRef { } pub fn unpack_alias(self, arena: &(impl MetaArena + ?Sized)) -> ExprRef { - if let ScalarExpression::Alias { - alias: AliasType::Expr(expr), - .. - } = arena.expression(self) - { - expr.unpack_alias(arena) - } else if let ScalarExpression::Alias { expr, .. } = arena.expression(self) { + if let ScalarExpression::Alias { expr, .. } = arena.expression(self) { expr.unpack_alias(arena) } else { self @@ -860,7 +858,7 @@ impl ExprRef { Ok(()) } - fn visit_column_ref(&mut self, col: &ColumnRef) -> Result<(), DatabaseError> { + fn visit_column_ref(&mut self, col: &ColumnRef, _: usize) -> Result<(), DatabaseError> { self.any = (self.f)(self.arena, col); Ok(()) } @@ -878,7 +876,7 @@ impl ExprRef { pub fn all_referenced_columns( self, arena: &(dyn MetaArena + '_), - mut predicate: impl FnMut(&(dyn MetaArena + '_), &ColumnRef) -> bool, + mut predicate: impl FnMut(&(dyn MetaArena + '_), &ColumnRef, usize) -> bool, ) -> Result { struct ColumnRefVisitor<'a, 'arena, F> { f: &'a mut F, @@ -886,8 +884,8 @@ impl ExprRef { arena: &'a (dyn MetaArena + 'arena), } - impl bool> ExprVisitor - for ColumnRefVisitor<'_, '_, F> + impl bool> + ExprVisitor for ColumnRefVisitor<'_, '_, F> { fn visit( &mut self, @@ -900,8 +898,12 @@ impl ExprRef { Ok(()) } - fn visit_column_ref(&mut self, col: &ColumnRef) -> Result<(), DatabaseError> { - self.all = (self.f)(self.arena, col); + fn visit_column_ref( + &mut self, + col: &ColumnRef, + position: usize, + ) -> Result<(), DatabaseError> { + self.all = (self.f)(self.arena, col, position); Ok(()) } } @@ -980,11 +982,11 @@ impl ExprRef { Ok(checker.0) } - pub fn output_name(self, arena: &(dyn MetaArena + '_)) -> String { + pub fn output_name(self, arena: &A) -> String { self.explain(arena).to_string() } - pub fn output_column_ref(self, arena: &mut (dyn MetaArena + '_)) -> ColumnRef { + pub fn output_column_ref(self, arena: &mut A) -> ColumnRef { match arena.expression(self) { ScalarExpression::ColumnRef { column, .. } => *column, ScalarExpression::Alias { @@ -1108,7 +1110,7 @@ mod test { use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; use crate::storage::rocksdb::RocksStorage; use crate::storage::rocksdb::RocksTransaction; - use crate::storage::{Storage, Transaction}; + use crate::storage::{Storage, StorageTransaction}; use crate::types::evaluator::{binary_create, cast_create, unary_create}; use crate::types::value::{DataValue, Utf8Type}; use crate::types::CharLengthUnits; diff --git a/src/expression/range_detacher.rs b/src/expression/range_detacher.rs index 211e1481..1eff00a0 100644 --- a/src/expression/range_detacher.rs +++ b/src/expression/range_detacher.rs @@ -18,6 +18,7 @@ use crate::expression::{BinaryOperator, ScalarExpression}; use crate::iter_ext::Itertools; use crate::planner::ExprRef; use crate::planner::MetaArena; +use crate::planner::ScalarQueryRef; use crate::types::index::IndexMetaRef; use crate::types::value::{DataValue, DataValueRef}; use crate::types::{ColumnId, LogicalType}; @@ -198,25 +199,36 @@ fn build_tree<'a, T: Borrow + 'a, I: EqRangeInput + ?Sized>( } impl Range { - pub(crate) fn bind_parameters<'p, F>(&mut self, get: &F) -> Result<(), DatabaseError> + pub(crate) fn bind_parameters<'p, F, G>( + &mut self, + params: &F, + init_values: &G, + ) -> Result<(), DatabaseError> where F: Fn(usize) -> Option<&'p DataValue>, + G: Fn(ScalarQueryRef) -> Option<&'p DataValue>, { - match self { + let has_null = match self { Self::Scope { min, max } => { + let mut has_null = false; for bound in [min, max] { if let Bound::Included(value) | Bound::Excluded(value) = bound { - value.bind_parameters(get)?; + has_null |= value.bind_parameters(params, init_values)?; } } + has_null } - Self::Eq(value) => value.bind_parameters(get)?, + Self::Eq(value) => value.bind_parameters(params, init_values)?, Self::SortedRanges(ranges) => { for range in ranges { - range.bind_parameters(get)?; + range.bind_parameters(params, init_values)?; } + false } - Self::Dummy => {} + Self::Dummy => false, + }; + if has_null { + *self = Self::Dummy; } Ok(()) } @@ -427,13 +439,13 @@ impl<'a, M: RangeColumnMatcher, A: MetaArena + ?Sized> RangeDetacher<'a, M, A> { left_expr .unpack_bound_col(self.arena, false) .map(|(column, _)| column), - right_expr.unpack_val(self.arena), + self.range_val(right_expr, op), ) { return self .new_range(op, col, val, false) .map(|range| range.map(DetachedPredicate::consumed)); } else if let (Some(val), Some(col)) = ( - left_expr.unpack_val(self.arena), + self.range_val(left_expr, op), right_expr .unpack_bound_col(self.arena, false) .map(|(column, _)| column), @@ -505,6 +517,50 @@ impl<'a, M: RangeColumnMatcher, A: MetaArena + ?Sized> RangeDetacher<'a, M, A> { _ => None, } } + ScalarExpression::In { + negated, + expr, + args, + .. + } => { + let Some((column, ty)) = expr + .in_point_column(self.arena) + .filter(|_| !args.is_empty()) + else { + return Ok(None); + }; + if !self.matches_column(column) { + return Ok(None); + } + let mut values = Vec::with_capacity(args.len()); + for arg in args { + let Some(value) = arg.in_point(ty, self.arena) else { + return Ok(None); + }; + if value.is_null() { + // NOT IN with NULL cannot be true in a Filter; positive + // IN ignores NULL points but must not match NULL rows. + if negated { + return Ok(Some(DetachedPredicate::consumed(Range::Dummy))); + } + } else { + values.push(value.clone()); + } + } + // Keep NOT IN residual, as with the old NotEq expansion. + if negated { + return Ok(None); + } + // `in_point` keeps every value in the column's type, so they are comparable. + values.sort_unstable_by(|left, right| left.partial_cmp(right).unwrap()); + values.dedup(); + let range = match values.len() { + 0 => Range::Dummy, + 1 => Range::Eq(values.pop().unwrap()), + _ => Range::SortedRanges(values.into_iter().map(Range::Eq).collect()), + }; + Some(DetachedPredicate::consumed(range)) + } ScalarExpression::Alias { expr, .. } | ScalarExpression::TypeCast { expr, .. } => { self.detach(expr)? } @@ -557,7 +613,6 @@ impl<'a, M: RangeColumnMatcher, A: MetaArena + ?Sized> RangeDetacher<'a, M, A> { | ScalarExpression::ColumnRef { .. } => None, // FIXME: support [RangeDetacher::_detach] ScalarExpression::Unary { .. } - | ScalarExpression::In { .. } | ScalarExpression::Between { .. } | ScalarExpression::SubString { .. } | ScalarExpression::Position { .. } @@ -575,6 +630,19 @@ impl<'a, M: RangeColumnMatcher, A: MetaArena + ?Sized> RangeDetacher<'a, M, A> { }) } + fn range_val(&self, expr: ExprRef, op: BinaryOperator) -> Option { + expr.unpack_val(self.arena) + .or_else(|| match self.arena.expression(expr) { + ScalarExpression::InitValue { id, ty } if op != BinaryOperator::Spaceship => { + Some(DataValue::InitValue { + id: *id, + ty: ty.clone(), + }) + } + _ => None, + }) + } + #[allow(clippy::result_large_err)] fn merge_binary( op: BinaryOperator, @@ -1302,7 +1370,7 @@ mod test { use crate::optimizer::rule::normalization::NormalizationRuleImpl; use crate::planner::operator::filter::FilterOperator; use crate::planner::operator::Operator; - use crate::planner::{ExprRef, LogicalPlan, MetaArena, PlanArena}; + use crate::planner::{ExprRef, LogicalPlan, MetaArena, PlanArena, TableArenaCell}; use crate::types::value::DataValue; use crate::types::LogicalType; use std::ops::Bound; @@ -1313,9 +1381,9 @@ mod test { ) -> Result, DatabaseError> { let pipeline = HepOptimizerPipeline::builder() .before_batch( - "test_simplify_filter".to_string(), + "test_simplify_filter", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::SimplifyFilter], + &[NormalizationRuleImpl::SimplifyFilter], ) .build(); let best_plan = pipeline.instantiate(plan).find_best(None, arena)?; @@ -1370,6 +1438,165 @@ mod test { }) } + fn point_list( + arena: &mut PlanArena, + column: ColumnRef, + values: &[DataValue], + negated: bool, + ) -> ExprRef { + let expr = arena.alloc_expression(ScalarExpression::column_expr(column, 0)); + let args = values + .iter() + .map(|value| arena.alloc_expression(ScalarExpression::Constant(value.clone()))) + .collect(); + arena.alloc_expression(ScalarExpression::In { + negated, + expr, + args, + evaluator: None, + }) + } + + #[test] + fn detach_in_deduplicates_points_and_never_selects_null() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let name: TableName = "points".into(); + let column = test_column(&mut arena, &name, 0, "c1", true)?; + for (values, expected) in [ + ( + vec![ + DataValue::Int32(3), + DataValue::Null, + DataValue::Int32(1), + DataValue::Int32(3), + ], + Range::SortedRanges(vec![ + Range::Eq(DataValue::Int32(1)), + Range::Eq(DataValue::Int32(3)), + ]), + ), + ( + vec![DataValue::Null, DataValue::Int32(1)], + Range::Eq(DataValue::Int32(1)), + ), + (vec![DataValue::Null], Range::Dummy), + ] { + let predicate = point_list(&mut arena, column, &values, false); + let detached = RangeDetacher::new(name.as_ref(), &0, &mut arena) + .detach(predicate)? + .unwrap(); + assert_eq!(detached.range, expected); + assert_eq!(detached.residual, None); + } + Ok(()) + } + + #[test] + fn detach_not_in_keeps_residual_unless_null_proves_no_match() -> Result<(), DatabaseError> { + for nullable in [false, true] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let name: TableName = "points".into(); + let column = test_column(&mut arena, &name, 0, "c1", nullable)?; + let predicate = point_list( + &mut arena, + column, + &[ + DataValue::Int32(3), + DataValue::Int32(1), + DataValue::Int32(3), + ], + true, + ); + assert_eq!( + RangeDetacher::new(name.as_ref(), &0, &mut arena).detach(predicate)?, + None + ); + let predicate = point_list( + &mut arena, + column, + &[DataValue::Int32(1), DataValue::Null], + true, + ); + assert_eq!( + RangeDetacher::new(name.as_ref(), &0, &mut arena) + .detach(predicate)? + .unwrap() + .range, + Range::Dummy + ); + } + Ok(()) + } + + #[test] + fn in_detachment_keeps_other_column_residual_and_intersects_bounds() -> Result<(), DatabaseError> + { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let name: TableName = "points".into(); + let column = test_column(&mut arena, &name, 0, "c1", true)?; + let other = test_column(&mut arena, &name, 1, "c2", true)?; + let points = point_list( + &mut arena, + column, + &[ + DataValue::Int32(1), + DataValue::Int32(3), + DataValue::Int32(5), + ], + false, + ); + let lower = cmp_predicate(&mut arena, column, BinaryOperator::Gt, 2); + let limited = and_predicate(&mut arena, points, lower); + let residual = cmp_predicate(&mut arena, other, BinaryOperator::Eq, 7); + let combined = and_predicate(&mut arena, limited, residual); + let detached = RangeDetacher::new(name.as_ref(), &0, &mut arena) + .detach(combined)? + .unwrap(); + assert_eq!( + detached.range, + Range::SortedRanges(vec![ + Range::Eq(DataValue::Int32(3)), + Range::Eq(DataValue::Int32(5)) + ]) + ); + assert_eq!(detached.residual, Some(residual)); + assert_eq!( + RangeDetacher::new(name.as_ref(), &1, &mut arena).detach(points)?, + None + ); + Ok(()) + } + + #[test] + fn unsafe_point_lists_are_not_detached() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let name: TableName = "points".into(); + let column = test_column(&mut arena, &name, 0, "c1", true)?; + for values in [ + vec![DataValue::Int64(2147483648)], + vec![DataValue::Parameter { + id: 0, + ty: LogicalType::Integer, + }], + ] { + let predicate = point_list(&mut arena, column, &values, false); + assert_eq!( + RangeDetacher::new(name.as_ref(), &0, &mut arena).detach(predicate)?, + None + ); + } + let predicate = point_list(&mut arena, column, &[], false); + assert_eq!( + RangeDetacher::new(name.as_ref(), &0, &mut arena).detach(predicate)?, + None + ); + Ok(()) + } + #[test] fn test_detach_consumes_and_predicates_with_residual() -> Result<(), DatabaseError> { let table_state = build_t1_table()?; @@ -2749,6 +2976,59 @@ mod test { Ok(()) } + #[test] + fn init_value_ranges_bind_and_drop_null_points() -> Result<(), DatabaseError> { + use crate::planner::{MetaArena, ScalarQueryRef, TableArena}; + + let mut arena = TableArena::default(); + let first = arena.alloc_scalar_query_ref(false); + let second = arena.alloc_scalar_query_ref(false); + let init = |id| DataValue::InitValue { + id, + ty: LogicalType::Bigint, + }; + let seven = DataValue::Int32(7); + let null = DataValue::Null; + let get = |id: ScalarQueryRef| match id.pos { + 0 => Some(&seven), + 1 => Some(&null), + _ => None, + }; + + let mut point = Range::Eq(DataValue::Tuple(vec![DataValue::Int32(1), init(first)])); + assert!(point.has_parameter()); + point.bind_parameters(&|_| None, &get)?; + assert!(!point.has_parameter()); + assert_eq!( + point, + Range::Eq(DataValue::Tuple(vec![ + DataValue::Int32(1), + DataValue::Int64(7) + ])) + ); + + let mut null_point = Range::Eq(DataValue::Tuple(vec![DataValue::Null, init(second)])); + null_point.bind_parameters(&|_| None, &get)?; + assert_eq!(null_point, Range::Dummy); + + let mut ranges = Range::SortedRanges(vec![ + Range::Eq(init(first)), + Range::Scope { + min: Bound::Excluded(init(second)), + max: Bound::Excluded(DataValue::Null), + }, + ]); + ranges.bind_parameters(&|_| None, &get)?; + assert_eq!( + ranges, + Range::SortedRanges(vec![Range::Eq(DataValue::Int64(7)), Range::Dummy]) + ); + + let mut unset = Range::Eq(init(first)); + assert!(unset.bind_parameters(&|_| None, &|_| None).is_err()); + Ok(()) + } + #[test] fn parameterized_range_union_keeps_residual_until_values_are_bound() -> Result<(), DatabaseError> { @@ -2775,11 +3055,14 @@ mod test { let mut ranges = Range::SortedRanges(vec![left, right]); let low = DataValue::Int32(8); let high = DataValue::Int32(3); - ranges.bind_parameters(&|id| match id { - 1 => Some(&low), - 2 => Some(&high), - _ => None, - })?; + ranges.bind_parameters( + &|id| match id { + 1 => Some(&low), + 2 => Some(&high), + _ => None, + }, + &|_| None, + )?; assert!(!ranges.has_parameter()); assert_eq!( ranges, diff --git a/src/expression/simplify.rs b/src/expression/simplify.rs index 59aa9d9f..18041394 100644 --- a/src/expression/simplify.rs +++ b/src/expression/simplify.rs @@ -26,16 +26,16 @@ use std::borrow::Cow; pub struct ConstantCalculator; impl ConstantCalculator { - pub fn new(_arena: &(dyn MetaArena + '_)) -> Self { + pub fn new(_arena: &A) -> Self { Self } } -impl ExprVisitorMut for ConstantCalculator { +impl ExprVisitorMut for ConstantCalculator { fn visit_expression( &mut self, expr: &mut ScalarExpression, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result { match expr { ScalarExpression::Unary { @@ -107,14 +107,16 @@ impl ExprVisitorMut for ConstantCalculator { } } +const IN_EXPANSION_LIMIT: usize = 5; + #[derive(Debug, Default)] pub struct Simplify; -impl ExprVisitorMut for Simplify { +impl ExprVisitorMut for Simplify { fn visit_expression( &mut self, expr: &mut ScalarExpression, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result { match expr { ScalarExpression::Unary { @@ -198,6 +200,15 @@ impl ExprVisitorMut for Simplify { if args.is_empty() { return Ok(false); } + // Keep point lists compact without changing pairwise coercion or + // repeated evaluation of non-column operands in the old rewrite. + if args.len() > IN_EXPANSION_LIMIT + && arg_expr.in_point_column(arena).is_some_and(|(_, ty)| { + args.iter().all(|arg| arg.in_point(ty, arena).is_some()) + }) + { + return Ok(true); + } let (op_1, op_2) = if *negated { (BinaryOperator::NotEq, BinaryOperator::And) @@ -302,9 +313,9 @@ impl Simplify { } } - fn take_range_comparison( + fn take_range_comparison( expr: ExprRef, - arena: &(dyn MetaArena + '_), + arena: &A, ) -> Option { match arena.expression(expr) { expression @ ScalarExpression::Binary { op, .. } @@ -316,9 +327,9 @@ impl Simplify { } } - fn take_negated_range_comparison( + fn take_negated_range_comparison( expr: ExprRef, - arena: &(dyn MetaArena + '_), + arena: &A, ) -> Option { let mut expression = arena.expression(expr).clone(); match &mut expression { @@ -330,17 +341,17 @@ impl Simplify { } } - fn boolean_constant(expr: ExprRef, arena: &(dyn MetaArena + '_)) -> Option { + fn boolean_constant(expr: ExprRef, arena: &A) -> Option { match arena.expression(expr) { ScalarExpression::Constant(DataValue::Boolean(value)) => Some(*value), _ => None, } } - fn take_range_comparison_with_polarity( + fn take_range_comparison_with_polarity( expr: ExprRef, positive: bool, - arena: &(dyn MetaArena + '_), + arena: &A, ) -> Option { if positive { Self::take_range_comparison(expr, arena) @@ -349,11 +360,11 @@ impl Simplify { } } - fn take_bool_normalized_range_comparison( + fn take_bool_normalized_range_comparison( op: BinaryOperator, left_expr: ExprRef, right_expr: ExprRef, - arena: &(dyn MetaArena + '_), + arena: &A, ) -> Option { let is_eq = matches!(op, BinaryOperator::Eq); if !matches!(op, BinaryOperator::Eq | BinaryOperator::NotEq) { @@ -412,11 +423,11 @@ impl Simplify { /// (float and decimal arithmetic round), never `*`/`/` (sign-dependent and /// not invertible for integers), and any overflow while folding aborts. /// Returns `None` when nothing could be peeled. - fn isolate_column( + fn isolate_column( mut expr: ExprRef, mut op: BinaryOperator, value_expr: ExprRef, - arena: &(dyn MetaArena + '_), + arena: &A, ) -> Option<(ExprRef, BinaryOperator, DataValue)> { let mut value = value_expr.unpack_val(arena)?; if value.has_parameter() { @@ -502,6 +513,39 @@ impl Simplify { } impl ExprRef { + pub(crate) fn in_point_column( + mut self, + arena: &A, + ) -> Option<(ColumnRef, &LogicalType)> { + while let ScalarExpression::Alias { expr, .. } = arena.expression(self) { + self = *expr; + } + let ScalarExpression::ColumnRef { column, .. } = arena.expression(self) else { + return None; + }; + let ty = arena.column(*column).datatype(); + (!matches!( + ty, + LogicalType::Tuple(_) | LogicalType::Float | LogicalType::Double + )) + .then_some((*column, ty)) + } + + pub(crate) fn in_point<'v, A: MetaArena + ?Sized>( + self, + ty: &LogicalType, + arena: &'v A, + ) -> Option<&'v DataValue> { + match arena.expression(self) { + ScalarExpression::Constant(value) + if !value.has_parameter() && (value.is_null() || value.logical_type() == *ty) => + { + Some(value) + } + _ => None, + } + } + pub(crate) fn unpack_val(self, arena: &A) -> Option { match arena.expression(self) { ScalarExpression::Constant(val) => Some(val.clone()), diff --git a/src/expression/visitor.rs b/src/expression/visitor.rs index e65c02ac..a1a81f9f 100644 --- a/src/expression/visitor.rs +++ b/src/expression/visitor.rs @@ -79,18 +79,20 @@ pub trait ExprVisitor: Sized { fn visit_constant(&mut self, _value: &DataValue) -> Result<(), DatabaseError> { Ok(()) } - fn visit_column_ref(&mut self, _column: &ColumnRef) -> Result<(), DatabaseError> { + fn visit_column_ref( + &mut self, + _column: &ColumnRef, + _position: usize, + ) -> Result<(), DatabaseError> { Ok(()) } fn visit_alias( &mut self, expr: ExprRef, - alias: &AliasType, + _alias: &AliasType, arena: &A, ) -> Result<(), DatabaseError> { - if let AliasType::Expr(alias_expr) = alias { - self.visit(*alias_expr, arena)?; - } + // The alias target only carries the output identity. self.visit(expr, arena) } fn visit_type_cast( @@ -324,7 +326,9 @@ pub fn walk_expr>( ScalarExpression::OuterParam { id, ty } => visitor.visit_outer_param(*id, ty, arena), ScalarExpression::OuterValue { id, ty } => visitor.visit_outer_value(*id, ty, arena), ScalarExpression::InitValue { id, ty } => visitor.visit_init_value(*id, ty, arena), - ScalarExpression::ColumnRef { column, .. } => visitor.visit_column_ref(column), + ScalarExpression::ColumnRef { column, position } => { + visitor.visit_column_ref(column, *position) + } ScalarExpression::Alias { expr, alias } => visitor.visit_alias(*expr, alias, arena), ScalarExpression::TypeCast { expr, diff --git a/src/expression/visitor_mut.rs b/src/expression/visitor_mut.rs index 41c0cb66..7f492292 100644 --- a/src/expression/visitor_mut.rs +++ b/src/expression/visitor_mut.rs @@ -29,14 +29,23 @@ use crate::types::LogicalType; pub(crate) struct ExprCloner; -impl ExprVisitorMut for ExprCloner { - fn visit( +impl ExprVisitorMut for ExprCloner { + fn visit(&mut self, expr: &mut ExprRef, arena: &mut A) -> Result<(), DatabaseError> { + *expr = arena.alloc_expression(arena.expression(*expr).clone()); + walk_mut_expr(self, expr, arena) + } + + // A deep clone must not share the alias target with the original. + fn visit_alias( &mut self, expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), + alias: &mut AliasType, + arena: &mut A, ) -> Result<(), DatabaseError> { - *expr = arena.alloc_expression(arena.expression(*expr).clone()); - walk_mut_expr(self, expr, arena) + if let AliasType::Expr(alias_expr) = alias { + self.visit(alias_expr, arena)?; + } + self.visit(expr, arena) } } @@ -44,12 +53,12 @@ pub(crate) struct PositionShift { pub(crate) delta: isize, } -impl ExprVisitorMut for PositionShift { +impl ExprVisitorMut for PositionShift { fn visit_column_ref( &mut self, _column: &mut ColumnRef, position: &mut usize, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result<(), DatabaseError> { if self.delta.is_negative() { *position = position.saturating_sub(self.delta.unsigned_abs()); @@ -60,12 +69,12 @@ impl ExprVisitorMut for PositionShift { } } -pub trait ExprVisitorMut: Sized { +pub trait ExprVisitorMut: Sized { fn visit_outer_param( &mut self, _id: &mut ScalarQueryRef, _ty: &mut LogicalType, - _arena: &mut dyn MetaArena, + _arena: &mut A, ) -> Result<(), DatabaseError> { Ok(()) } @@ -74,7 +83,7 @@ pub trait ExprVisitorMut: Sized { &mut self, _id: &mut ScalarQueryRef, _ty: &mut LogicalType, - _arena: &mut dyn MetaArena, + _arena: &mut A, ) -> Result<(), DatabaseError> { Ok(()) } @@ -83,16 +92,12 @@ pub trait ExprVisitorMut: Sized { &mut self, _id: &mut ScalarQueryRef, _ty: &mut LogicalType, - _arena: &mut dyn MetaArena, + _arena: &mut A, ) -> Result<(), DatabaseError> { Ok(()) } - fn visit( - &mut self, - expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { + fn visit(&mut self, expr: &mut ExprRef, arena: &mut A) -> Result<(), DatabaseError> { if !self.visit_expression_ref(expr, arena)? { return Ok(()); } @@ -110,7 +115,7 @@ pub trait ExprVisitorMut: Sized { fn visit_expression_ref( &mut self, _expr: &mut ExprRef, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result { Ok(true) } @@ -118,7 +123,7 @@ pub trait ExprVisitorMut: Sized { fn visit_expression( &mut self, _expr: &mut ScalarExpression, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result { Ok(true) } @@ -126,7 +131,7 @@ pub trait ExprVisitorMut: Sized { fn visit_constant( &mut self, _value: &mut DataValue, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result<(), DatabaseError> { Ok(()) } @@ -135,7 +140,7 @@ pub trait ExprVisitorMut: Sized { &mut self, _column: &mut ColumnRef, _position: &mut usize, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result<(), DatabaseError> { Ok(()) } @@ -143,12 +148,10 @@ pub trait ExprVisitorMut: Sized { fn visit_alias( &mut self, expr: &mut ExprRef, - alias: &mut AliasType, - arena: &mut (dyn MetaArena + '_), + _alias: &mut AliasType, + arena: &mut A, ) -> Result<(), DatabaseError> { - if let AliasType::Expr(alias_expr) = alias { - self.visit(alias_expr, arena)?; - } + // The alias target only carries the output identity. self.visit(expr, arena) } @@ -157,7 +160,7 @@ pub trait ExprVisitorMut: Sized { expr: &mut ExprRef, _ty: &mut LogicalType, _evaluator: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena) } @@ -166,7 +169,7 @@ pub trait ExprVisitorMut: Sized { &mut self, _negated: bool, expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena) } @@ -177,7 +180,7 @@ pub trait ExprVisitorMut: Sized { expr: &mut ExprRef, _evaluator: &mut Option, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena) } @@ -189,7 +192,7 @@ pub trait ExprVisitorMut: Sized { right_expr: &mut ExprRef, _evaluator: &mut Option, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(left_expr, arena)?; self.visit(right_expr, arena) @@ -201,7 +204,7 @@ pub trait ExprVisitorMut: Sized { _kind: &mut AggKind, args: &mut [ExprRef], _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { for arg in args { self.visit(arg, arena)?; @@ -212,7 +215,7 @@ pub trait ExprVisitorMut: Sized { fn visit_window( &mut self, window: &mut WindowCall, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { for expr in window .function @@ -232,7 +235,7 @@ pub trait ExprVisitorMut: Sized { expr: &mut ExprRef, args: &mut [ExprRef], _evaluator: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; for arg in args { @@ -250,7 +253,7 @@ pub trait ExprVisitorMut: Sized { right_expr: &mut ExprRef, _left_evaluator: &mut Option, _right_evaluator: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; self.visit(left_expr, arena)?; @@ -262,7 +265,7 @@ pub trait ExprVisitorMut: Sized { expr: &mut ExprRef, for_expr: &mut Option, from_expr: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; if let Some(for_expr) = for_expr { @@ -278,7 +281,7 @@ pub trait ExprVisitorMut: Sized { &mut self, expr: &mut ExprRef, in_expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; self.visit(in_expr, arena) @@ -289,7 +292,7 @@ pub trait ExprVisitorMut: Sized { expr: &mut ExprRef, trim_what_expr: &mut Option, _trim_where: &mut Option, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena)?; if let Some(trim_what_expr) = trim_what_expr { @@ -306,16 +309,12 @@ pub trait ExprVisitorMut: Sized { &mut self, expr: &mut ExprRef, _pos: usize, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(expr, arena) } - fn visit_tuple( - &mut self, - exprs: &mut [ExprRef], - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { + fn visit_tuple(&mut self, exprs: &mut [ExprRef], arena: &mut A) -> Result<(), DatabaseError> { for expr in exprs { self.visit(expr, arena)?; } @@ -325,7 +324,7 @@ pub trait ExprVisitorMut: Sized { fn visit_scala_function( &mut self, function: &mut ScalarFunction, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { for arg in &mut function.args { self.visit(arg, arena)?; @@ -336,7 +335,7 @@ pub trait ExprVisitorMut: Sized { fn visit_table_function( &mut self, function: &mut TableFunction, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { for arg in &mut function.args { self.visit(arg, arena)?; @@ -350,7 +349,7 @@ pub trait ExprVisitorMut: Sized { left_expr: &mut ExprRef, right_expr: &mut ExprRef, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(condition, arena)?; self.visit(left_expr, arena)?; @@ -362,7 +361,7 @@ pub trait ExprVisitorMut: Sized { left_expr: &mut ExprRef, right_expr: &mut ExprRef, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(left_expr, arena)?; self.visit(right_expr, arena) @@ -373,7 +372,7 @@ pub trait ExprVisitorMut: Sized { left_expr: &mut ExprRef, right_expr: &mut ExprRef, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { self.visit(left_expr, arena)?; self.visit(right_expr, arena) @@ -383,7 +382,7 @@ pub trait ExprVisitorMut: Sized { &mut self, exprs: &mut [ExprRef], _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { for expr in exprs { self.visit(expr, arena)?; @@ -397,7 +396,7 @@ pub trait ExprVisitorMut: Sized { expr_pairs: &mut [(ExprRef, ExprRef)], else_expr: &mut Option, _ty: &mut LogicalType, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { if let Some(expr) = operand_expr { self.visit(expr, arena)?; @@ -413,10 +412,10 @@ pub trait ExprVisitorMut: Sized { } } -pub fn walk_mut_expr( +pub fn walk_mut_expr, A: MetaArena + ?Sized>( visitor: &mut V, expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result<(), DatabaseError> { let mut expression = std::mem::replace(&mut *arena.expression_mut(*expr), ScalarExpression::Empty); diff --git a/src/function/abs.rs b/src/function/abs.rs index 9f5cd34c..e5dbbd84 100644 --- a/src/function/abs.rs +++ b/src/function/abs.rs @@ -18,7 +18,7 @@ use crate::expression::function::scala::ScalarFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::LogicalType; use ordered_float::OrderedFloat; @@ -62,7 +62,7 @@ impl ScalarFunctionImpl for Abs { &self, exprs: &[ExprRef], arena: &(dyn MetaArena + '_), - tuples: Option<&dyn TupleLike>, + tuples: Option<&dyn Input>, ) -> Result { let value = arena.expression(exprs[0]).eval(arena, tuples)?; let overflow = || DatabaseError::OverFlow; diff --git a/src/function/char_length.rs b/src/function/char_length.rs index e2a49bbc..a999586d 100644 --- a/src/function/char_length.rs +++ b/src/function/char_length.rs @@ -18,7 +18,7 @@ use crate::expression::function::scala::ScalarFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::CharLengthUnits; use crate::types::LogicalType; @@ -46,7 +46,7 @@ impl ScalarFunctionImpl for CharLength { &self, exprs: &[ExprRef], arena: &(dyn MetaArena + '_), - tuples: Option<&dyn TupleLike>, + tuples: Option<&dyn Input>, ) -> Result { let mut value = arena.expression(exprs[0]).eval(arena, tuples)?; // SQL: the length of NULL is NULL, not 0. diff --git a/src/function/current_date.rs b/src/function/current_date.rs index 1deb47bf..7d4353bf 100644 --- a/src/function/current_date.rs +++ b/src/function/current_date.rs @@ -18,7 +18,7 @@ use crate::expression::function::scala::ScalarFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::LogicalType; use chrono::{Datelike, Local}; @@ -48,7 +48,7 @@ impl ScalarFunctionImpl for CurrentDate { &self, _: &[ExprRef], _: &(dyn MetaArena + '_), - _: Option<&dyn TupleLike>, + _: Option<&dyn Input>, ) -> Result { Ok(DataValue::Date32(Local::now().num_days_from_ce())) } diff --git a/src/function/current_timestamp.rs b/src/function/current_timestamp.rs index b8a00f4d..be3f72a4 100644 --- a/src/function/current_timestamp.rs +++ b/src/function/current_timestamp.rs @@ -18,7 +18,7 @@ use crate::expression::function::scala::ScalarFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::LogicalType; use chrono::Utc; @@ -48,7 +48,7 @@ impl ScalarFunctionImpl for CurrentTimeStamp { &self, _: &[ExprRef], _: &(dyn MetaArena + '_), - _: Option<&dyn TupleLike>, + _: Option<&dyn Input>, ) -> Result { Ok(DataValue::Time64(Utc::now().timestamp(), 0, false)) } diff --git a/src/function/lower.rs b/src/function/lower.rs index 8e4e8b53..a4fc59b8 100644 --- a/src/function/lower.rs +++ b/src/function/lower.rs @@ -18,7 +18,7 @@ use crate::expression::function::scala::ScalarFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::CharLengthUnits; use crate::types::LogicalType; @@ -48,7 +48,7 @@ impl ScalarFunctionImpl for Lower { &self, exprs: &[ExprRef], arena: &(dyn MetaArena + '_), - tuples: Option<&dyn TupleLike>, + tuples: Option<&dyn Input>, ) -> Result { let mut value = arena.expression(exprs[0]).eval(arena, tuples)?; if !matches!(value.logical_type(), LogicalType::Varchar(_, _)) { diff --git a/src/function/octet_length.rs b/src/function/octet_length.rs index 0a0ad13b..80474722 100644 --- a/src/function/octet_length.rs +++ b/src/function/octet_length.rs @@ -18,7 +18,7 @@ use crate::expression::function::scala::ScalarFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::CharLengthUnits; use crate::types::LogicalType; @@ -47,7 +47,7 @@ impl ScalarFunctionImpl for OctetLength { &self, exprs: &[ExprRef], arena: &(dyn MetaArena + '_), - tuples: Option<&dyn TupleLike>, + tuples: Option<&dyn Input>, ) -> Result { let mut value = arena.expression(exprs[0]).eval(arena, tuples)?; // SQL: the length of NULL is NULL, not 0. diff --git a/src/function/upper.rs b/src/function/upper.rs index c6da1e46..b121d83b 100644 --- a/src/function/upper.rs +++ b/src/function/upper.rs @@ -18,7 +18,7 @@ use crate::expression::function::scala::ScalarFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::ExprRef; use crate::planner::MetaArena; -use crate::types::tuple::TupleLike; +use crate::types::tuple::Input; use crate::types::value::DataValue; use crate::types::CharLengthUnits; use crate::types::LogicalType; @@ -48,7 +48,7 @@ impl ScalarFunctionImpl for Upper { &self, exprs: &[ExprRef], arena: &(dyn MetaArena + '_), - tuples: Option<&dyn TupleLike>, + tuples: Option<&dyn Input>, ) -> Result { let mut value = arena.expression(exprs[0]).eval(arena, tuples)?; if !matches!(value.logical_type(), LogicalType::Varchar(_, _)) { diff --git a/src/macros/mod.rs b/src/macros/mod.rs index 18a87a1d..71b7ddbc 100644 --- a/src/macros/mod.rs +++ b/src/macros/mod.rs @@ -41,13 +41,13 @@ macro_rules! from_tuple { impl ::kite_sql::orm::FromQueryRow for $struct_name { fn from_query_row( schema: &::kite_sql::types::tuple::SchemaView<'_, '_>, - tuple: &mut ::kite_sql::types::tuple::Tuple, + tuple: &mut ::kite_sql::types::tuple::TupleImpl, ) -> ::std::result::Result { - fn try_get(tuple: &mut ::kite_sql::types::tuple::Tuple, schema: &::kite_sql::types::tuple::SchemaView<'_, '_>, field_name: &str) -> Option<::kite_sql::types::value::DataValue> { + fn try_get(tuple: &mut ::kite_sql::types::tuple::TupleImpl, schema: &::kite_sql::types::tuple::SchemaView<'_, '_>, field_name: &str) -> Option<::kite_sql::types::value::DataValue> { let ty = ::kite_sql::types::LogicalType::type_trans::()?; let idx = schema.position(field_name)?; - std::mem::replace(&mut tuple.values[idx], ::kite_sql::types::value::DataValue::Null).cast(&ty).ok() + std::mem::replace(&mut tuple[idx], ::kite_sql::types::value::DataValue::Null).cast(&ty).ok() } let mut struct_instance = $struct_name::default(); @@ -102,7 +102,7 @@ macro_rules! scala_function { } } impl ::kite_sql::expression::function::scala::ScalarFunctionImpl for $struct_name { #[allow(unused_variables, clippy::redundant_closure_call)] - fn eval(&self, args: &[::kite_sql::planner::ExprRef], arena: &dyn ::kite_sql::planner::MetaArena, tuple: Option<&dyn ::kite_sql::types::tuple::TupleLike>) -> Result<::kite_sql::types::value::DataValue, ::kite_sql::errors::DatabaseError> { + fn eval(&self, args: &[::kite_sql::planner::ExprRef], arena: &dyn ::kite_sql::planner::MetaArena, tuple: Option<&dyn ::kite_sql::types::tuple::Input>) -> Result<::kite_sql::types::value::DataValue, ::kite_sql::errors::DatabaseError> { let mut _index = 0; $closure($({ diff --git a/src/optimizer/core/cm_sketch.rs b/src/optimizer/core/cm_sketch.rs index 76ee3a7d..32f4f530 100644 --- a/src/optimizer/core/cm_sketch.rs +++ b/src/optimizer/core/cm_sketch.rs @@ -17,7 +17,7 @@ use crate::expression::range_detacher::Range; use crate::planner::MetaArena; use crate::serdes::stable_hash::{StableHasher, CM_SKETCH_HASH_KEYS}; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; use kite_sql_serde_macros::ReferenceSerialization; use std::borrow::Borrow; @@ -373,7 +373,7 @@ impl ReferenceSerialization for CountMinSketch { Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/optimizer/core/histogram.rs b/src/optimizer/core/histogram.rs index a2f28dcb..2eb9e5f7 100644 --- a/src/optimizer/core/histogram.rs +++ b/src/optimizer/core/histogram.rs @@ -28,19 +28,27 @@ use ordered_float::OrderedFloat; use std::borrow::Cow; use std::cmp; use std::collections::Bound; -use std::mem; use std::sync::OnceLock; const ANALYZE_STATISTICS_CONFIDENCE: f64 = 0.95; -pub(crate) const ANALYZE_STATISTICS_RELATIVE_ERROR: f64 = 0.001; +// Relative error of the KLL quantile and Count-Min sketches when +// `DataBaseBuilder::statistics_relative_error` is unset. +pub(crate) const DEFAULT_STATISTICS_RELATIVE_ERROR: f64 = 0.001; +// HyperLogLog precision used for NDV when `DataBaseBuilder::hll_precision` is +// unset: 2^14 registers (16 KiB per sketch, ~0.81% standard error). +pub(crate) const DEFAULT_HLL_PRECISION: u8 = 14; + +pub struct HistogramMetaBuilder { + null_count: usize, + values_len: usize, + hll: HyperLogLog, +} pub struct HistogramBuilder { index_id: IndexId, - null_count: usize, - values_len: usize, + meta: HistogramMetaBuilder, quantile: KllSketchBuilder, sketch: CountMinSketch, - hll: HyperLogLog, top_n: ColumnTopN, } @@ -54,7 +62,6 @@ struct BoundComparator { #[derive(Debug, Clone, PartialEq, ReferenceSerialization)] pub struct HistogramMeta { - index_id: IndexId, number_of_distinct_value: usize, null_count: usize, values_len: usize, @@ -68,6 +75,7 @@ pub struct HistogramMeta { // Equal depth histogram #[derive(Debug)] pub struct Histogram { + index_id: IndexId, meta: HistogramMeta, buckets: Vec, comparator: OnceLock, @@ -76,6 +84,7 @@ pub struct Histogram { impl Clone for Histogram { fn clone(&self) -> Self { Self { + index_id: self.index_id, meta: self.meta.clone(), buckets: self.buckets.clone(), comparator: OnceLock::new(), @@ -85,7 +94,7 @@ impl Clone for Histogram { impl PartialEq for Histogram { fn eq(&self, other: &Self) -> bool { - self.meta == other.meta && self.buckets == other.buckets + self.index_id == other.index_id && self.meta == other.meta && self.buckets == other.buckets } } @@ -97,37 +106,78 @@ pub struct Bucket { // repeat: u64, } -impl HistogramBuilder { - pub fn new(index_meta: &IndexMeta, relative_error: f64) -> Result { +impl HistogramMetaBuilder { + pub fn new(hll_precision: u8) -> Result { Ok(Self { - index_id: index_meta.id, null_count: 0, values_len: 0, + hll: HyperLogLog::with_precision(hll_precision)?, + }) + } + + pub fn append(&mut self, value: &DataValue) -> bool { + if value.is_null() { + self.null_count += 1; + false + } else { + self.values_len += 1; + self.hll.add(value); + true + } + } + + pub fn build(self) -> HistogramMeta { + HistogramMeta { + number_of_distinct_value: if self.values_len == 0 { + 0 + } else { + self.hll.estimate().clamp(1, self.values_len) + }, + null_count: self.null_count, + values_len: self.values_len, + buckets_len: 0, + correlation: 0.0, + } + } +} + +impl HistogramBuilder { + pub fn new( + index_meta: &IndexMeta, + relative_error: f64, + hll_precision: u8, + ) -> Result { + Ok(Self { + index_id: index_meta.id, + meta: HistogramMetaBuilder::new(hll_precision)?, + // TODO: derive the KLL rank error from the histogram bucket count + // (e.g. 1 / (2 * buckets)) instead of sharing `relative_error` with the + // Count-Min sketch; 100 buckets only need ~0.5% rank error, not 0.1%. quantile: KllSketchBuilder::with_relative_error(relative_error)?, sketch: CountMinSketch::with_relative_error( ANALYZE_STATISTICS_CONFIDENCE, relative_error, )?, - hll: HyperLogLog::with_relative_error(relative_error)?, top_n: ColumnTopN::default(), }) } pub fn append(&mut self, value: DataValue) -> Result<(), DatabaseError> { - if value.is_null() { - self.null_count += 1; - } else { + if self.meta.append(&value) { self.sketch.increment(&value); - self.hll.add(&value); self.top_n .add_with_size(ANALYZE_STATISTICS_TOP_N_SIZE, value.clone(), 1); self.quantile.insert(value)?; - self.values_len += 1; } Ok(()) } + // TODO: shrink the Count-Min sketch to the observed row count before returning it. + // Its width is fixed by `relative_error` (2048 x 5 counters, ~80 KiB at 0.001) + // and persisted per index even for tiny tables; since the width is a power of + // two, rows can be folded in half (summing both halves) down to about + // `next_power_of_two(values_len)` without losing accuracy on large tables. pub fn build( self, number_of_buckets: usize, @@ -138,19 +188,17 @@ impl HistogramBuilder { )); } - let values_len = self.values_len; + let values_len = self.meta.values_len; if number_of_buckets > values_len { return Err(DatabaseError::TooManyBuckets(number_of_buckets, values_len)); } let HistogramBuilder { index_id, - null_count, + meta, quantile, mut sketch, - hll, top_n, - .. } = self; let mut buckets = Vec::with_capacity(number_of_buckets); let quantile = quantile.build()?; @@ -178,20 +226,16 @@ impl HistogramBuilder { count: (upper_rank - lower_rank + 1) as u64, }); } - sketch.add(&DataValue::Null, self.null_count); - let number_of_distinct_value = hll.estimate().clamp(1, values_len); + let mut meta = meta.build(); + sketch.add(&DataValue::Null, meta.null_count); + meta.buckets_len = buckets.len(); + meta.correlation = correlation; let top_n = top_n.finish_with_size(ANALYZE_STATISTICS_TOP_N_SIZE); Ok(( Histogram { - meta: HistogramMeta { - index_id, - number_of_distinct_value, - null_count, - values_len, - buckets_len: buckets.len(), - correlation, - }, + index_id, + meta, buckets, comparator: OnceLock::new(), }, @@ -282,7 +326,11 @@ fn is_above( } impl Histogram { - pub fn from_parts(meta: HistogramMeta, buckets: Vec) -> Result { + pub fn from_parts( + index_id: IndexId, + meta: HistogramMeta, + buckets: Vec, + ) -> Result { if meta.buckets_len != buckets.len() { return Err(DatabaseError::InvalidValue(format!( "histogram bucket count mismatch: meta={}, actual={}", @@ -292,14 +340,15 @@ impl Histogram { } Ok(Self { + index_id, meta, buckets, comparator: OnceLock::new(), }) } - pub fn into_parts(self) -> (HistogramMeta, Vec) { - (self.meta, self.buckets) + pub fn into_parts(self) -> (IndexId, HistogramMeta, Vec) { + (self.index_id, self.meta, self.buckets) } fn comparator(&self) -> Result<&BoundComparator, DatabaseError> { @@ -331,7 +380,7 @@ impl Histogram { } pub fn index_id(&self) -> IndexId { - self.meta.index_id + self.index_id } pub fn values_len(&self) -> usize { @@ -422,6 +471,7 @@ impl Histogram { ranges: &[Range], sketch: &CountMinSketch, top_n: &ColumnTopN, + buf: &mut Vec, ) -> Result { if self.buckets.is_empty() || ranges.is_empty() { return Ok(0); @@ -444,31 +494,24 @@ impl Histogram { let mut count = 0; let mut binary_i = 0; let mut bucket_i = 0; - let mut bucket_idxs = Vec::new(); - let mut buf = Vec::new(); while bucket_i < self.buckets.len() && binary_i < ranges.len() { let is_dummy = self._collect_count( ranges, &mut binary_i, &mut bucket_i, - &mut bucket_idxs, &mut count, sketch, top_n, comparator, - &mut buf, + buf, )?; if is_dummy { return Ok(0); } } - Ok(bucket_idxs - .iter() - .map(|idx| self.buckets[*idx].count as usize) - .sum::() - + count) + Ok(count) } #[allow(clippy::too_many_arguments)] @@ -477,7 +520,6 @@ impl Histogram { ranges: &[Range], binary_i: &mut usize, bucket_i: &mut usize, - bucket_idxs: &mut Vec, count: &mut usize, sketch: &CountMinSketch, top_n: &ColumnTopN, @@ -505,7 +547,9 @@ impl Histogram { && (is_under(comparator, &bucket.upper, max, false)? || is_eq(&bucket.upper, max)) { - bucket_idxs.push(mem::replace(bucket_i, *bucket_i + 1)); + // A fully covered bucket counts its rows, without the NULLs added above. + *count += bucket.count as usize; + *bucket_i += 1; } else if is_above(comparator, &bucket.lower, max, false)? { *binary_i += 1; } else if is_under(comparator, &bucket.upper, min, true)? { @@ -670,8 +714,8 @@ fn scope_lower_includes_null(min: &Bound) -> bool { } impl HistogramMeta { - pub fn index_id(&self) -> IndexId { - self.index_id + pub fn null_count(&self) -> usize { + self.null_count } pub fn values_len(&self) -> usize { @@ -694,7 +738,8 @@ mod tests { use crate::expression::range_detacher::Range; use crate::optimizer::core::cm_sketch::CountMinSketch; use crate::optimizer::core::histogram::{ - Bucket, HistogramBuilder, ANALYZE_STATISTICS_RELATIVE_ERROR, + Bucket, HistogramBuilder, HistogramMetaBuilder, DEFAULT_HLL_PRECISION, + DEFAULT_STATISTICS_RELATIVE_ERROR, }; use crate::optimizer::core::top_n::ColumnTopN; use crate::types::index::{IndexMeta, IndexType}; @@ -702,6 +747,36 @@ mod tests { use crate::types::LogicalType; use std::ops::Bound; + #[test] + fn column_meta_counts_values_without_histogram_buckets() -> Result<(), DatabaseError> { + let empty = HistogramMetaBuilder::new(DEFAULT_HLL_PRECISION)?.build(); + assert_eq!(empty.values_len(), 0); + assert_eq!(empty.null_count(), 0); + assert_eq!(empty.distinct_values_len(), 0); + assert_eq!(empty.buckets_len(), 0); + let mut builder = HistogramMetaBuilder::new(DEFAULT_HLL_PRECISION)?; + for value in [ + DataValue::Null, + DataValue::Int32(1), + DataValue::Int32(1), + DataValue::Int32(2), + ] { + builder.append(&value); + } + let meta = builder.build(); + assert_eq!(meta.values_len(), 3); + assert_eq!(meta.null_count(), 1); + assert_eq!(meta.distinct_values_len(), 2); + assert_eq!(meta.buckets_len(), 0); + let mut builder = HistogramMetaBuilder::new(DEFAULT_HLL_PRECISION)?; + builder.append(&DataValue::Null); + let meta = builder.build(); + assert_eq!(meta.values_len(), 0); + assert_eq!(meta.null_count(), 1); + assert_eq!(meta.distinct_values_len(), 0); + Ok(()) + } + fn index_meta() -> IndexMeta { IndexMeta { id: 0, @@ -716,7 +791,11 @@ mod tests { #[test] fn test_sort_tuples_on_histogram() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; builder.append(DataValue::Int32(0))?; builder.append(DataValue::Int32(1))?; @@ -782,7 +861,11 @@ mod tests { #[test] fn test_rev_sort_tuples_on_histogram() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; builder.append(DataValue::Int32(14))?; builder.append(DataValue::Int32(13))?; @@ -846,7 +929,11 @@ mod tests { #[test] fn test_non_average_on_histogram() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; builder.append(DataValue::Int32(14))?; builder.append(DataValue::Int32(13))?; @@ -905,7 +992,11 @@ mod tests { #[test] fn parameterized_tuple_ranges_accumulate_prefix_selectivity() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for value in 0..100 { builder.append(DataValue::Tuple(vec![DataValue::Int32(value); 4]))?; } @@ -923,12 +1014,22 @@ mod tests { max: Bound::Excluded(DataValue::Tuple(upper)), }; assert_eq!( - histogram.collect_count(&[scope(prefix.clone(), prefix.clone())], &sketch, &top_n)?, + histogram.collect_count( + &[scope(prefix.clone(), prefix.clone())], + &sketch, + &top_n, + &mut Vec::new() + )?, 100 ); let longer = vec![parameter(1), parameter(2), parameter(3)]; assert_eq!( - histogram.collect_count(&[scope(longer.clone(), longer)], &sketch, &top_n)?, + histogram.collect_count( + &[scope(longer.clone(), longer)], + &sketch, + &top_n, + &mut Vec::new() + )?, 10 ); let mut lower = prefix.clone(); @@ -936,16 +1037,31 @@ mod tests { let mut upper = prefix.clone(); upper.push(parameter(4)); assert_eq!( - histogram.collect_count(&[scope(lower.clone(), prefix)], &sketch, &top_n)?, + histogram.collect_count( + &[scope(lower.clone(), prefix)], + &sketch, + &top_n, + &mut Vec::new() + )?, 25 ); assert_eq!( - histogram.collect_count(&[scope(lower.clone(), upper.clone())], &sketch, &top_n)?, + histogram.collect_count( + &[scope(lower.clone(), upper.clone())], + &sketch, + &top_n, + &mut Vec::new() + )?, 2 ); let full = DataValue::Tuple(vec![parameter(1), parameter(2), parameter(3), parameter(4)]); assert_eq!( - histogram.collect_count(&[Range::Eq(full.clone())], &sketch, &top_n)?, + histogram.collect_count( + &[Range::Eq(full.clone())], + &sketch, + &top_n, + &mut Vec::new() + )?, 1 ); assert_eq!( @@ -955,7 +1071,8 @@ mod tests { max: Bound::Included(full), }], &sketch, - &top_n + &top_n, + &mut Vec::new() )?, 0 ); @@ -963,7 +1080,7 @@ mod tests { lower.push(parameter(5)); upper.push(parameter(5)); assert_eq!( - histogram.collect_count(&[scope(lower, upper)], &sketch, &top_n)?, + histogram.collect_count(&[scope(lower, upper)], &sketch, &top_n, &mut Vec::new())?, 2 ); histogram.buckets[0].upper = DataValue::Tuple(vec![ @@ -972,7 +1089,12 @@ mod tests { ]); let nested = vec![DataValue::Tuple(vec![parameter(1), parameter(2)])]; assert_eq!( - histogram.collect_count(&[scope(nested.clone(), nested)], &sketch, &top_n)?, + histogram.collect_count( + &[scope(nested.clone(), nested)], + &sketch, + &top_n, + &mut Vec::new() + )?, 100 ); Ok(()) @@ -1170,7 +1292,11 @@ mod tests { #[test] fn parameterized_range_count_uses_shape_and_ndv() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for value in 0..1_000 { builder.append(DataValue::Int32(value % 100))?; } @@ -1204,7 +1330,7 @@ mod tests { ), ] { assert_eq!( - histogram.collect_count(&[range], &sketch, &top_n)?, + histogram.collect_count(&[range], &sketch, &top_n, &mut Vec::new())?, expected ); } @@ -1213,7 +1339,11 @@ mod tests { #[test] fn test_collect_count() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; builder.append(DataValue::Int32(14))?; builder.append(DataValue::Int32(13))?; @@ -1248,6 +1378,7 @@ mod tests { ], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_1, 9); @@ -1259,6 +1390,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_2, 11); @@ -1270,6 +1402,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_3, 7); @@ -1281,6 +1414,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_4, 12); @@ -1292,6 +1426,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_5, 8); @@ -1303,6 +1438,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_6, 12); @@ -1314,6 +1450,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_7, 13); @@ -1325,6 +1462,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_8, 13); @@ -1336,6 +1474,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_9, 13); @@ -1347,6 +1486,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_10, 2); @@ -1358,6 +1498,7 @@ mod tests { }], &sketch, &top_n, + &mut Vec::new(), )?; assert_eq!(count_11, 2); @@ -1367,7 +1508,11 @@ mod tests { #[test] fn test_builder_uses_hll_for_distinct_values() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for _ in 0..50 { builder.append(DataValue::Int32(1))?; @@ -1386,7 +1531,11 @@ mod tests { #[test] fn test_eq_count_uses_cm_sketch_when_error_is_small() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for _ in 0..70 { builder.append(DataValue::Int32(1))?; @@ -1399,7 +1548,12 @@ mod tests { let top_n = ColumnTopN::default(); assert_eq!( - histogram.collect_count(&[Range::Eq(DataValue::Int32(1))], &sketch, &top_n)?, + histogram.collect_count( + &[Range::Eq(DataValue::Int32(1))], + &sketch, + &top_n, + &mut Vec::new() + )?, 70 ); @@ -1408,7 +1562,11 @@ mod tests { #[test] fn test_eq_count_uses_top_n_when_value_is_cached() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for _ in 0..70 { builder.append(DataValue::Int32(1))?; @@ -1421,7 +1579,12 @@ mod tests { sketch.add(&DataValue::Int32(1), 1000); assert_eq!( - histogram.collect_count(&[Range::Eq(DataValue::Int32(1))], &sketch, &top_n)?, + histogram.collect_count( + &[Range::Eq(DataValue::Int32(1))], + &sketch, + &top_n, + &mut Vec::new() + )?, 70 ); @@ -1430,7 +1593,11 @@ mod tests { #[test] fn test_eq_count_falls_back_to_average_when_cm_error_is_large() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for value in 0..10_000 { builder.append(DataValue::Int32(value))?; @@ -1445,7 +1612,12 @@ mod tests { assert!(estimated_count.saturating_sub(error_bound) <= average_count); assert_eq!( - histogram.collect_count(&[Range::Eq(DataValue::Int32(7))], &sketch, &top_n)?, + histogram.collect_count( + &[Range::Eq(DataValue::Int32(7))], + &sketch, + &top_n, + &mut Vec::new() + )?, average_count ); @@ -1454,7 +1626,11 @@ mod tests { #[test] fn test_eq_count_uses_cm_sketch_for_clear_heavy_hitter() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for _ in 0..1_000 { builder.append(DataValue::Int32(7))?; @@ -1472,7 +1648,12 @@ mod tests { assert!(estimated_count.saturating_sub(error_bound) > average_count); assert_eq!( - histogram.collect_count(&[Range::Eq(DataValue::Int32(7))], &sketch, &top_n)?, + histogram.collect_count( + &[Range::Eq(DataValue::Int32(7))], + &sketch, + &top_n, + &mut Vec::new() + )?, estimated_count ); @@ -1481,7 +1662,11 @@ mod tests { #[test] fn test_eq_count_falls_back_when_top_n_error_is_large() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for value in 0..10_000 { builder.append(DataValue::Int32(value))?; @@ -1502,7 +1687,12 @@ mod tests { assert!(expected < entry.count()); assert_eq!( - histogram.collect_count(&[Range::Eq(DataValue::Int32(7))], &sketch, &top_n)?, + histogram.collect_count( + &[Range::Eq(DataValue::Int32(7))], + &sketch, + &top_n, + &mut Vec::new() + )?, expected ); @@ -1511,7 +1701,11 @@ mod tests { #[test] fn test_collect_count_ignores_tuple_prefix_endpoint_count() -> Result<(), DatabaseError> { - let mut builder = HistogramBuilder::new(&index_meta(), ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta(), + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for value in 0..15 { builder.append(DataValue::Tuple(vec![ @@ -1525,13 +1719,13 @@ mod tests { min: Bound::Excluded(DataValue::Tuple(vec![DataValue::Int32(0)])), max: Bound::Excluded(DataValue::Tuple(vec![DataValue::Int32(8)])), }]; - let clean_count = histogram.collect_count(&ranges, &sketch, &top_n)?; + let clean_count = histogram.collect_count(&ranges, &sketch, &top_n, &mut Vec::new())?; sketch.increment(&DataValue::Tuple(vec![DataValue::Int32(0)])); sketch.increment(&DataValue::Tuple(vec![DataValue::Int32(8)])); assert_eq!( - histogram.collect_count(&ranges, &sketch, &top_n)?, + histogram.collect_count(&ranges, &sketch, &top_n, &mut Vec::new())?, clean_count ); @@ -1547,7 +1741,7 @@ mod tests { }; let mut sketch = CountMinSketch::with_relative_error( super::ANALYZE_STATISTICS_CONFIDENCE, - ANALYZE_STATISTICS_RELATIVE_ERROR, + DEFAULT_STATISTICS_RELATIVE_ERROR, )?; let real_key = DataValue::Tuple(vec![DataValue::Int32(8), DataValue::Int32(8)]); diff --git a/src/optimizer/core/hll.rs b/src/optimizer/core/hll.rs index 8dae20ae..f9807604 100644 --- a/src/optimizer/core/hll.rs +++ b/src/optimizer/core/hll.rs @@ -18,8 +18,8 @@ use std::cmp; use std::hash::{Hash, Hasher}; use std::marker::PhantomData; -const MIN_PRECISION: u8 = 4; -const MAX_PRECISION: u8 = 26; +pub(crate) const MIN_PRECISION: u8 = 4; +pub(crate) const MAX_PRECISION: u8 = 26; #[derive(Debug, Clone)] pub struct HyperLogLog { @@ -29,17 +29,16 @@ pub struct HyperLogLog { } impl HyperLogLog { - pub fn with_relative_error(relative_error: f64) -> Result { - if relative_error <= 0.0 || relative_error.is_nan() { + /// Creates a sketch with `2^precision` registers, whose standard error is + /// about `1.04 / sqrt(2^precision)`. + pub fn with_precision(precision: u8) -> Result { + if !(MIN_PRECISION..=MAX_PRECISION).contains(&precision) { return Err(DatabaseError::InvalidValue(format!( - "hyperloglog relative error must be greater than zero, got {relative_error}" + "hyperloglog precision must be between {MIN_PRECISION} and {MAX_PRECISION}, got {precision}" ))); } - - let register_count = Self::optimal_register_count(relative_error)?; - let precision = register_count.trailing_zeros() as u8; Ok(Self { - registers: vec![0; register_count], + registers: vec![0; 1 << precision], precision, phantom_k: PhantomData, }) @@ -80,32 +79,6 @@ impl HyperLogLog { estimate.round() as usize } - fn optimal_register_count(relative_error: f64) -> Result { - let count = (1.04 / relative_error).powi(2).ceil(); - if !count.is_finite() || count > usize::MAX as f64 { - return Err(DatabaseError::InvalidValue(format!( - "hyperloglog relative error is too small: {relative_error}" - ))); - } - - let min_count = 1usize << MIN_PRECISION; - let count = cmp::max(min_count, count as usize) - .checked_next_power_of_two() - .ok_or_else(|| { - DatabaseError::InvalidValue(format!( - "hyperloglog register count overflow for relative error: {relative_error}" - )) - })?; - let precision = count.trailing_zeros() as u8; - if precision > MAX_PRECISION { - return Err(DatabaseError::InvalidValue(format!( - "hyperloglog relative error is too small: {relative_error}" - ))); - } - - Ok(count) - } - fn hash(value: &Q) -> u64 { let (key0, key1) = HLL_HASH_KEYS; let mut hasher = StableHasher::new_with_keys(key0, key1); @@ -130,7 +103,7 @@ mod tests { #[test] fn estimate_ignores_duplicates() { - let mut hll = HyperLogLog::::with_relative_error(0.01).unwrap(); + let mut hll = HyperLogLog::::with_precision(14).unwrap(); for _ in 0..1000 { hll.add(&7); } @@ -138,9 +111,22 @@ mod tests { assert_eq!(hll.estimate(), 1); } + #[test] + fn precision_sets_register_count() { + assert_eq!( + HyperLogLog::::with_precision(14) + .unwrap() + .registers + .len(), + 1 << 14 + ); + assert!(HyperLogLog::::with_precision(3).is_err()); + assert!(HyperLogLog::::with_precision(27).is_err()); + } + #[test] fn estimate_distinct_values_within_error_budget() { - let mut hll = HyperLogLog::::with_relative_error(0.02).unwrap(); + let mut hll = HyperLogLog::::with_precision(12).unwrap(); for value in 0..10_000 { hll.add(&value); } diff --git a/src/optimizer/core/rule.rs b/src/optimizer/core/rule.rs index ffc62673..386cdbff 100644 --- a/src/optimizer/core/rule.rs +++ b/src/optimizer/core/rule.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption}; +use crate::planner::operator::PhysicalOption; use crate::planner::{LogicalPlan, PlanArena}; use std::cmp::Ordering; @@ -28,7 +28,14 @@ pub trait MatchPattern { pub trait NormalizationRule { /// Returns true when the plan tree is modified. - fn apply(&self, plan: &mut LogicalPlan, arena: &mut PlanArena) -> Result; + /// Local rewrites must invalidate changed descendants before reading their schemas. + /// The optimizer invalidates the rewrite root and its ancestors, not untouched subtrees. + fn apply( + &self, + plan: &mut LogicalPlan, + arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, + ) -> Result; } fn compare_costs(candidate_cost: Option, best_cost: Option) -> Ordering { @@ -55,12 +62,18 @@ pub fn keep_best_physical_option( } } +#[derive(Default)] +pub struct HintState { + pub(crate) histogram_buf: Vec, +} + pub trait ImplementationRule: MatchPattern { fn update_best_option( &self, - op: &Operator, + plan: &LogicalPlan, arena: &PlanArena, loader: &StatisticMetaLoader<'_>, + state: &mut HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError>; } diff --git a/src/optimizer/core/statistics_meta.rs b/src/optimizer/core/statistics_meta.rs index d0a3260f..27042787 100644 --- a/src/optimizer/core/statistics_meta.rs +++ b/src/optimizer/core/statistics_meta.rs @@ -21,6 +21,7 @@ use crate::optimizer::core::top_n::ColumnTopN; use crate::storage::StatisticsMetaCache; use crate::types::index::IndexId; use crate::types::value::DataValue; +use crate::types::ColumnId; use kite_sql_serde_macros::ReferenceSerialization; use std::slice; @@ -37,9 +38,17 @@ impl<'a> StatisticMetaLoader<'a> { &self, table_name: &TableName, index_id: IndexId, - ) -> Result, DatabaseError> { + ) -> Result, DatabaseError> { let key = (table_name.clone(), index_id); - Ok(self.cache.get(&key)) + Ok(self.cache.indexes.get(&key)) + } + + pub fn load_column( + &self, + table_name: &TableName, + column_id: ColumnId, + ) -> Option<&HistogramMeta> { + self.cache.columns.get(&(table_name.clone(), column_id)) } pub fn collect_count( @@ -47,7 +56,8 @@ impl<'a> StatisticMetaLoader<'a> { table_name: &TableName, index_id: IndexId, range: &Range, - ) -> Result, DatabaseError> { + buf: &mut Vec, + ) -> Result, DatabaseError> { let Some(entry) = self.load(table_name, index_id)? else { return Ok(None); }; @@ -59,21 +69,21 @@ impl<'a> StatisticMetaLoader<'a> { entry .histogram() - .collect_count(ranges, entry.sketch(), entry.top_n()) - .map(Some) + .collect_count(ranges, entry.sketch(), entry.top_n(), buf) + .map(|rows| Some((rows, entry.histogram().distinct_values_len()))) } } #[derive(Debug, Clone, ReferenceSerialization)] -pub struct StatisticsMetaRoot { +pub struct IndexStatisticsRoot { index_id: IndexId, histogram_meta: HistogramMeta, } -impl StatisticsMetaRoot { - pub fn new(histogram_meta: HistogramMeta) -> Self { +impl IndexStatisticsRoot { + pub fn new(index_id: IndexId, histogram_meta: HistogramMeta) -> Self { Self { - index_id: histogram_meta.index_id(), + index_id, histogram_meta, } } @@ -92,16 +102,16 @@ impl StatisticsMetaRoot { } #[derive(Debug, Clone)] -pub struct StatisticsMeta { +pub struct IndexStatistics { index_id: IndexId, histogram: Histogram, sketch: CountMinSketch, top_n: ColumnTopN, } -impl StatisticsMeta { +impl IndexStatistics { pub fn new(histogram: Histogram, sketch: CountMinSketch, top_n: ColumnTopN) -> Self { - StatisticsMeta { + IndexStatistics { index_id: histogram.index_id(), histogram, sketch, @@ -110,12 +120,13 @@ impl StatisticsMeta { } pub fn from_parts( - root: StatisticsMetaRoot, + root: IndexStatisticsRoot, buckets: Vec, sketch: CountMinSketch, top_n: Option, ) -> Result { - let histogram = Histogram::from_parts(root.into_histogram_meta(), buckets)?; + let histogram = + Histogram::from_parts(root.index_id(), root.into_histogram_meta(), buckets)?; Ok(Self::new(histogram, sketch, top_n.unwrap_or_default())) } @@ -123,14 +134,14 @@ impl StatisticsMeta { pub fn into_parts( self, ) -> ( - StatisticsMetaRoot, + IndexStatisticsRoot, Vec, CountMinSketch, ColumnTopN, ) { - let (histogram_meta, buckets) = self.histogram.into_parts(); + let (index_id, histogram_meta, buckets) = self.histogram.into_parts(); ( - StatisticsMetaRoot::new(histogram_meta), + IndexStatisticsRoot::new(index_id, histogram_meta), buckets, self.sketch, self.top_n, @@ -158,8 +169,12 @@ impl StatisticsMeta { #[cfg(all(test, not(target_arch = "wasm32")))] mod tests { use crate::errors::DatabaseError; - use crate::optimizer::core::histogram::{HistogramBuilder, ANALYZE_STATISTICS_RELATIVE_ERROR}; - use crate::optimizer::core::statistics_meta::StatisticsMeta; + use crate::expression::range_detacher::Range; + use crate::optimizer::core::histogram::{ + HistogramBuilder, DEFAULT_HLL_PRECISION, DEFAULT_STATISTICS_RELATIVE_ERROR, + }; + use crate::optimizer::core::statistics_meta::{IndexStatistics, StatisticMetaLoader}; + use crate::storage::StatisticsMetaCache; use crate::types::index::{IndexMeta, IndexType}; use crate::types::value::DataValue; use crate::types::LogicalType; @@ -176,7 +191,11 @@ mod tests { ty: IndexType::PrimaryKey { is_multiple: false }, }; - let mut builder = HistogramBuilder::new(&index, ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index, + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; builder.append(DataValue::Int32(14))?; builder.append(DataValue::Int32(13))?; @@ -199,18 +218,37 @@ mod tests { let (histogram, sketch, top_n) = builder.build(4)?; let expected_estimate = sketch.estimate(&DataValue::Int32(7)); let expected_top_n_count = top_n.get(&DataValue::Int32(7)); - let meta = StatisticsMeta::new(histogram.clone(), sketch, top_n); + let meta = IndexStatistics::new(histogram.clone(), sketch, top_n); let (root, buckets, sketch, top_n) = meta.into_parts(); - let statistics_meta = StatisticsMeta::from_parts(root, buckets, sketch, Some(top_n))?; + let index_statistics = IndexStatistics::from_parts(root, buckets, sketch, Some(top_n))?; - assert_eq!(histogram, statistics_meta.histogram); + let mut cache = StatisticsMetaCache::default(); + cache.indexes.insert( + (index.table_name.clone(), index.id), + index_statistics.clone(), + ); + let loader = StatisticMetaLoader::new(&cache); + assert_eq!( + loader.collect_count( + &index.table_name, + index.id, + &Range::Eq(DataValue::Int32(7)), + &mut Vec::new() + )?, + Some((1, histogram.distinct_values_len())) + ); + assert_eq!( + loader.collect_count(&index.table_name, 1, &Range::Dummy, &mut Vec::new())?, + None + ); + assert_eq!(histogram, index_statistics.histogram); assert_eq!( expected_estimate, - statistics_meta.sketch().estimate(&DataValue::Int32(7)) + index_statistics.sketch().estimate(&DataValue::Int32(7)) ); assert_eq!( expected_top_n_count, - statistics_meta.top_n().get(&DataValue::Int32(7)) + index_statistics.top_n().get(&DataValue::Int32(7)) ); Ok(()) diff --git a/src/optimizer/heuristic/batch.rs b/src/optimizer/heuristic/batch.rs index 25efa273..fb351496 100644 --- a/src/optimizer/heuristic/batch.rs +++ b/src/optimizer/heuristic/batch.rs @@ -22,7 +22,7 @@ use std::array; #[derive(Clone)] pub struct HepBatch { #[allow(dead_code)] - pub name: String, + pub name: &'static str, pub strategy: HepBatchStrategy, pub steps: Vec, } @@ -90,13 +90,13 @@ impl HepLocalRewriteBatch { impl HepBatch { pub fn new( - name: String, + name: &'static str, strategy: HepBatchStrategy, - rules: Vec, + rules: &[NormalizationRuleImpl], ) -> Self { let mut steps = Vec::new(); - for rule in rules { + for &rule in rules { match rule.pass_kind() { NormalizationPassKind::WholeTreePass(kind) => match steps.last_mut() { Some(HepBatchStep::WholeTree(pass)) if pass.kind == kind => { @@ -124,28 +124,51 @@ impl HepBatch { } } +#[derive(Clone, Copy)] +pub enum HepMatchOrder { + TopDown, + BottomUp, +} + #[derive(Clone, Copy)] pub enum HepBatchStrategy { - /// An execution_ap strategy for rules that indicates the maximum number of executions. If the - /// execution_ap reaches fix point (i.e. converge) before maxIterations, it will stop. - /// - /// Fix Point means that plan tree not changed after applying all rules. - MaxTimes(usize), + /// Each local rule may succeed once across the tree. + Once(HepMatchOrder), + /// Stop at a fix point or after the configured number of full-tree iterations. + MaxTimes { + max_iteration: usize, + order: HepMatchOrder, + }, #[allow(dead_code)] - LoopIfApplied, + LoopIfApplied(HepMatchOrder), } impl HepBatchStrategy { pub fn once_topdown() -> Self { - HepBatchStrategy::MaxTimes(1) + Self::Once(HepMatchOrder::TopDown) + } + + #[cfg(test)] + pub fn once_bottomup() -> Self { + Self::Once(HepMatchOrder::BottomUp) } pub fn fix_point_topdown(max_iteration: usize) -> Self { - HepBatchStrategy::MaxTimes(max_iteration) + Self::MaxTimes { + max_iteration, + order: HepMatchOrder::TopDown, + } + } + + pub fn fix_point_bottomup(max_iteration: usize) -> Self { + Self::MaxTimes { + max_iteration, + order: HepMatchOrder::BottomUp, + } } #[allow(dead_code)] pub fn loop_if_applied() -> Self { - HepBatchStrategy::LoopIfApplied + Self::LoopIfApplied(HepMatchOrder::TopDown) } } diff --git a/src/optimizer/heuristic/optimizer.rs b/src/optimizer/heuristic/optimizer.rs index 8b853407..2fdd21c4 100644 --- a/src/optimizer/heuristic/optimizer.rs +++ b/src/optimizer/heuristic/optimizer.rs @@ -14,11 +14,11 @@ use crate::errors::DatabaseError; use crate::optimizer::core::rule::{ - BestPhysicalOption, ImplementationRule, MatchPattern, NormalizationRule, + BestPhysicalOption, HintState, ImplementationRule, MatchPattern, NormalizationRule, }; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::heuristic::batch::{ - HepBatch, HepBatchStep, HepBatchStrategy, HepLocalRewriteBatch, HepWholeTreePass, + HepBatch, HepBatchStep, HepBatchStrategy, HepLocalRewriteBatch, HepMatchOrder, HepWholeTreePass, }; use crate::optimizer::rule::implementation::{ImplementationRuleImpl, ImplementationRuleRootTag}; use crate::optimizer::rule::normalization::{ @@ -26,7 +26,7 @@ use crate::optimizer::rule::normalization::{ evaluator_bind_current, NormalizationRuleImpl, OrderHintKind, ScanOrderHint, WholeTreePassKind, }; use crate::planner::operator::table_scan::TableScanOperator; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::planner::{Childrens, LogicalPlan, PlanArena}; use std::array; use std::ops::Not; @@ -38,7 +38,6 @@ pub struct HepOptimizer<'a> { before_batches: &'a [HepBatch], after_batches: &'a [HepBatch], implementation_index: &'a ImplementationRuleIndex, - max_local_rules_len: usize, plan: LogicalPlan, } @@ -48,13 +47,11 @@ impl<'a> HepOptimizer<'a> { before_batches: &'a [HepBatch], after_batches: &'a [HepBatch], implementation_index: &'a ImplementationRuleIndex, - max_local_rules_len: usize, ) -> Self { Self { before_batches, after_batches, implementation_index, - max_local_rules_len, plan, } } @@ -64,12 +61,13 @@ impl<'a> HepOptimizer<'a> { loader: Option<&StatisticMetaLoader<'_>>, arena: &mut PlanArena, ) -> Result { - let mut applied_rules = Vec::with_capacity(self.max_local_rules_len); + let mut applied_rules = Vec::new(); Self::apply_batches( &mut self.plan, self.before_batches, &mut applied_rules, arena, + loader, )?; if let Some(loader) = loader { @@ -84,6 +82,7 @@ impl<'a> HepOptimizer<'a> { self.implementation_index, &apply_no_sort_hints, &apply_no_stream_aggregate_hints, + &mut HintState::default(), arena, )?; } @@ -93,6 +92,7 @@ impl<'a> HepOptimizer<'a> { self.after_batches, &mut applied_rules, arena, + loader, )?; Ok(self.plan) @@ -104,18 +104,25 @@ impl<'a> HepOptimizer<'a> { batches: &[HepBatch], applied_rules: &mut Vec, arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, ) -> Result<(), DatabaseError> { for batch in batches { match batch.strategy { - HepBatchStrategy::MaxTimes(max_iteration) => { + HepBatchStrategy::Once(order) => { + Self::apply_batch(plan, batch, order, applied_rules, arena, loader)?; + } + HepBatchStrategy::MaxTimes { + max_iteration, + order, + } => { for _ in 0..max_iteration { - if !Self::apply_batch(plan, batch, applied_rules, arena)? { + if !Self::apply_batch(plan, batch, order, applied_rules, arena, loader)? { break; } } } - HepBatchStrategy::LoopIfApplied => { - while Self::apply_batch(plan, batch, applied_rules, arena)? {} + HepBatchStrategy::LoopIfApplied(order) => { + while Self::apply_batch(plan, batch, order, applied_rules, arena, loader)? {} } } } @@ -126,19 +133,25 @@ impl<'a> HepOptimizer<'a> { fn apply_batch( plan: &mut LogicalPlan, batch: &HepBatch, + order: HepMatchOrder, applied_rules: &mut Vec, arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let mut applied = false; for step in &batch.steps { match step { HepBatchStep::WholeTree(pass) => { - if Self::apply_whole_tree_pass(plan, pass, arena)? { + if Self::apply_whole_tree_pass(plan, pass, arena, loader)? { applied = true; } } HepBatchStep::LocalRewrite(rules) => { - if Self::apply_local_rules(plan, rules, applied_rules, arena)? { + applied_rules.clear(); + if matches!(batch.strategy, HepBatchStrategy::Once(_)) { + applied_rules.resize(rules.len(), false); + } + if Self::apply_local_rules(plan, rules, order, applied_rules, arena, loader)? { applied = true; } } @@ -151,12 +164,13 @@ impl<'a> HepOptimizer<'a> { plan: &mut LogicalPlan, pass: &HepWholeTreePass, arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { match pass.kind { WholeTreePassKind::ColumnPruning => { let mut applied = false; for rule in &pass.rules { - applied |= rule.apply(plan, arena)?; + applied |= rule.apply(plan, arena, loader)?; } if applied { plan.reset_output_schema_cache_recursive(); @@ -247,6 +261,7 @@ impl<'a> HepOptimizer<'a> { implementation_index: &ImplementationRuleIndex, inherited_sort_hints: &'plan ScanHintApplier<'plan>, inherited_stream_aggregate_hints: &'plan ScanHintApplier<'plan>, + state: &mut HintState, arena: &mut PlanArena, ) -> Result<(), DatabaseError> { if let Operator::TableScan(scan_op) = &mut plan.operator { @@ -254,25 +269,6 @@ impl<'a> HepOptimizer<'a> { inherited_stream_aggregate_hints(scan_op, arena)?; } - { - let LogicalPlan { - operator, - physical_option, - .. - } = plan; - if let Some(option) = implementation_index.direct_physical_option(operator) { - *physical_option = Some(option); - } else { - let mut best_physical_option: BestPhysicalOption = None; - for rule in implementation_index.for_matching_operator(operator) { - rule.update_best_option(operator, arena, loader, &mut best_physical_option)?; - } - if let Some((option, _)) = best_physical_option { - *physical_option = Some(option); - } - } - } - { let LogicalPlan { operator, @@ -290,6 +286,7 @@ impl<'a> HepOptimizer<'a> { implementation_index, child_sort_hints, child_stream_aggregate_hints, + state, arena, ), Childrens::Twins { left, right } => { @@ -299,6 +296,7 @@ impl<'a> HepOptimizer<'a> { implementation_index, child_sort_hints, child_stream_aggregate_hints, + state, arena, )?; Self::annotate_hints_and_physical_options( @@ -307,6 +305,7 @@ impl<'a> HepOptimizer<'a> { implementation_index, child_sort_hints, child_stream_aggregate_hints, + state, arena, ) } @@ -316,6 +315,18 @@ impl<'a> HepOptimizer<'a> { })?; } + if let Some(option) = implementation_index.direct_physical_option(&plan.operator) { + plan.physical_option = Some(option); + } else { + let mut best_physical_option: BestPhysicalOption = None; + for rule in implementation_index.for_matching_operator(&plan.operator) { + rule.update_best_option(plan, arena, loader, state, &mut best_physical_option)?; + } + if let Some((option, _)) = best_physical_option { + plan.physical_option = Some(option); + } + } + apply_annotated_post_rules(plan, arena)?; Ok(()) @@ -395,59 +406,66 @@ impl<'a> HepOptimizer<'a> { fn apply_local_rules( plan: &mut LogicalPlan, rules: &HepLocalRewriteBatch, - applied_rules: &mut Vec, + order: HepMatchOrder, + applied_rules: &mut [bool], arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { - applied_rules.clear(); - applied_rules.resize(rules.len(), false); - Self::apply_local_rules_inner(plan, rules, applied_rules, arena) + if !applied_rules.is_empty() && applied_rules.iter().all(|done| *done) { + return Ok(false); + } + let mut applied = false; + if matches!(order, HepMatchOrder::TopDown) { + applied |= Self::apply_local_rules_current(plan, rules, applied_rules, arena, loader)?; + } + + let child_applied = match plan.childrens.as_mut() { + Childrens::Only(child) => { + Self::apply_local_rules(child, rules, order, applied_rules, arena, loader)? + } + Childrens::Twins { left, right } => { + let left_applied = + Self::apply_local_rules(left, rules, order, applied_rules, arena, loader)?; + let right_applied = + Self::apply_local_rules(right, rules, order, applied_rules, arena, loader)?; + left_applied || right_applied + } + Childrens::None => false, + }; + if child_applied { + plan.reset_output_schema_cache(); + } + applied |= child_applied; + + if matches!(order, HepMatchOrder::BottomUp) { + applied |= Self::apply_local_rules_current(plan, rules, applied_rules, arena, loader)?; + } + Ok(applied) } - fn apply_local_rules_inner( + fn apply_local_rules_current( plan: &mut LogicalPlan, rules: &HepLocalRewriteBatch, applied_rules: &mut [bool], arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let mut applied = false; let mut next_rule_idx = 0; - while let Some(idx) = rules.next_matching_rule_index_from(&plan.operator, next_rule_idx) { - let rule = rules.rules[idx]; next_rule_idx = idx + 1; - if applied_rules[idx] { + if applied_rules.get(idx).copied().unwrap_or(false) { continue; } - let applied_rule = rule.apply(plan, arena)?; - if applied_rule { - applied_rules[idx] = true; - applied = true; - plan.reset_output_schema_cache_recursive(); - } - } - - match plan.childrens.as_mut() { - Childrens::Only(child) => { - let child_applied = - Self::apply_local_rules_inner(child, rules, applied_rules, arena)?; - applied |= child_applied; - if child_applied { - plan.reset_output_schema_cache(); - } - } - Childrens::Twins { left, right } => { - let left_applied = - Self::apply_local_rules_inner(left, rules, applied_rules, arena)?; - let right_applied = - Self::apply_local_rules_inner(right, rules, applied_rules, arena)?; - applied |= left_applied || right_applied; - if left_applied || right_applied { - plan.reset_output_schema_cache(); + if rules.rules[idx].apply(plan, arena, loader)? { + if let Some(done) = applied_rules.get_mut(idx) { + *done = true; } + applied = true; + // Rules invalidate changed descendants; untouched subtrees retain their caches. + plan.reset_output_schema_cache(); } - Childrens::None => {} } - Ok(applied) } } @@ -458,9 +476,9 @@ struct ImplementationRuleIndex { } impl ImplementationRuleIndex { - fn new(implementations: Vec) -> Self { + fn new(implementations: &[ImplementationRuleImpl]) -> Self { let mut groups = array::from_fn(|_| Vec::new()); - for implementation in implementations { + for &implementation in implementations { groups[implementation.root_tag() as usize].push(implementation); } Self { groups } @@ -483,6 +501,7 @@ impl ImplementationRuleIndex { Some(PhysicalOption::new( PlanImpl::HashAggregate, SortOption::None, + RowCountOption::Unknown, )) } Operator::Aggregate(op) @@ -492,25 +511,35 @@ impl ImplementationRuleIndex { Some(PhysicalOption::new( PlanImpl::SimpleAggregate, SortOption::None, + RowCountOption::Estimated { rows: 1, ndv: None }, )) } - Operator::Dummy if self.contains(ImplementationRuleImpl::Dummy) => { - Some(PhysicalOption::new(PlanImpl::Dummy, SortOption::None)) - } - Operator::Filter(_) if self.contains(ImplementationRuleImpl::Filter) => { - Some(PhysicalOption::new(PlanImpl::Filter, SortOption::Follow)) - } + Operator::Dummy if self.contains(ImplementationRuleImpl::Dummy) => Some( + PhysicalOption::new(PlanImpl::Dummy, SortOption::None, RowCountOption::Unknown), + ), Operator::Join(join_op) if self.contains(ImplementationRuleImpl::HashJoin) => { - Some(PhysicalOption::new(join_op.plan_impl(), SortOption::None)) - } - Operator::Limit(_) if self.contains(ImplementationRuleImpl::Limit) => { - Some(PhysicalOption::new(PlanImpl::Limit, SortOption::Follow)) + Some(PhysicalOption::new( + join_op.plan_impl(), + SortOption::None, + RowCountOption::Unknown, + )) } + Operator::Limit(_) if self.contains(ImplementationRuleImpl::Limit) => Some( + PhysicalOption::new(PlanImpl::Limit, SortOption::Follow, RowCountOption::Unknown), + ), Operator::MarkApply(_) if self.contains(ImplementationRuleImpl::MarkApply) => { - Some(PhysicalOption::new(PlanImpl::MarkApply, SortOption::Follow)) + Some(PhysicalOption::new( + PlanImpl::MarkApply, + SortOption::Follow, + RowCountOption::Follow, + )) } Operator::Project(_) if self.contains(ImplementationRuleImpl::Projection) => { - Some(PhysicalOption::new(PlanImpl::Project, SortOption::Follow)) + Some(PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + )) } Operator::ScalarQueryInit(_) if self.contains(ImplementationRuleImpl::ScalarQueryInit) => @@ -518,6 +547,7 @@ impl ImplementationRuleIndex { Some(PhysicalOption::new( PlanImpl::ScalarQueryInit, SortOption::Follow, + RowCountOption::Follow, )) } Operator::ScalarSubquery(_) @@ -526,12 +556,14 @@ impl ImplementationRuleIndex { Some(PhysicalOption::new( PlanImpl::ScalarSubquery, SortOption::Follow, + RowCountOption::Follow, )) } Operator::FunctionScan(_) if self.contains(ImplementationRuleImpl::FunctionScan) => { Some(PhysicalOption::new( PlanImpl::FunctionScan, SortOption::None, + RowCountOption::Unknown, )) } Operator::Sort(op) if self.contains(ImplementationRuleImpl::Sort) => { @@ -541,6 +573,7 @@ impl ImplementationRuleIndex { fields: op.sort_fields.clone(), ignore_prefix_len: 0, }, + RowCountOption::Follow, )) } Operator::TopK(op) if self.contains(ImplementationRuleImpl::TopK) => { @@ -550,57 +583,91 @@ impl ImplementationRuleIndex { fields: op.sort_fields.clone(), ignore_prefix_len: 0, }, + RowCountOption::Unknown, )) } - Operator::Values(_) if self.contains(ImplementationRuleImpl::Values) => { - Some(PhysicalOption::new(PlanImpl::Values, SortOption::None)) - } - Operator::Window(op) if self.contains(ImplementationRuleImpl::Window) => { - Some(PhysicalOption::new(PlanImpl::Window, op.sort_option())) - } - Operator::Analyze(_) if self.contains(ImplementationRuleImpl::Analyze) => { - Some(PhysicalOption::new(PlanImpl::Analyze, SortOption::None)) + Operator::Values(op) if self.contains(ImplementationRuleImpl::Values) => { + Some(PhysicalOption::new( + PlanImpl::Values, + SortOption::None, + RowCountOption::Estimated { + rows: op.rows.len() / op.schema_ref.len(), + ndv: None, + }, + )) } + Operator::Window(op) if self.contains(ImplementationRuleImpl::Window) => Some( + PhysicalOption::new(PlanImpl::Window, op.sort_option(), RowCountOption::Follow), + ), + Operator::Analyze(_) if self.contains(ImplementationRuleImpl::Analyze) => Some( + PhysicalOption::new(PlanImpl::Analyze, SortOption::None, RowCountOption::Unknown), + ), #[cfg(feature = "copy")] Operator::CopyFromFile(_) if self.contains(ImplementationRuleImpl::CopyFromFile) => { Some(PhysicalOption::new( PlanImpl::CopyFromFile, SortOption::None, + RowCountOption::Unknown, )) } #[cfg(feature = "copy")] Operator::CopyToFile(_) if self.contains(ImplementationRuleImpl::CopyToFile) => { - Some(PhysicalOption::new(PlanImpl::CopyToFile, SortOption::None)) - } - Operator::Delete(_) if self.contains(ImplementationRuleImpl::Delete) => { - Some(PhysicalOption::new(PlanImpl::Delete, SortOption::None)) - } - Operator::Insert(_) if self.contains(ImplementationRuleImpl::Insert) => { - Some(PhysicalOption::new(PlanImpl::Insert, SortOption::None)) - } - Operator::Update(_) if self.contains(ImplementationRuleImpl::Update) => { - Some(PhysicalOption::new(PlanImpl::Update, SortOption::None)) + Some(PhysicalOption::new( + PlanImpl::CopyToFile, + SortOption::None, + RowCountOption::Unknown, + )) } + Operator::Delete(_) if self.contains(ImplementationRuleImpl::Delete) => Some( + PhysicalOption::new(PlanImpl::Delete, SortOption::None, RowCountOption::Unknown), + ), + Operator::Insert(_) if self.contains(ImplementationRuleImpl::Insert) => Some( + PhysicalOption::new(PlanImpl::Insert, SortOption::None, RowCountOption::Unknown), + ), + Operator::Update(_) if self.contains(ImplementationRuleImpl::Update) => Some( + PhysicalOption::new(PlanImpl::Update, SortOption::None, RowCountOption::Unknown), + ), Operator::AddColumn(_) if self.contains(ImplementationRuleImpl::AddColumn) => { - Some(PhysicalOption::new(PlanImpl::AddColumn, SortOption::None)) + Some(PhysicalOption::new( + PlanImpl::AddColumn, + SortOption::None, + RowCountOption::Unknown, + )) } Operator::ChangeColumn(_) if self.contains(ImplementationRuleImpl::ChangeColumn) => { Some(PhysicalOption::new( PlanImpl::ChangeColumn, SortOption::None, + RowCountOption::Unknown, )) } Operator::CreateTable(_) if self.contains(ImplementationRuleImpl::CreateTable) => { - Some(PhysicalOption::new(PlanImpl::CreateTable, SortOption::None)) + Some(PhysicalOption::new( + PlanImpl::CreateTable, + SortOption::None, + RowCountOption::Unknown, + )) } Operator::DropColumn(_) if self.contains(ImplementationRuleImpl::DropColumn) => { - Some(PhysicalOption::new(PlanImpl::DropColumn, SortOption::None)) + Some(PhysicalOption::new( + PlanImpl::DropColumn, + SortOption::None, + RowCountOption::Unknown, + )) } Operator::DropTable(_) if self.contains(ImplementationRuleImpl::DropTable) => { - Some(PhysicalOption::new(PlanImpl::DropTable, SortOption::None)) + Some(PhysicalOption::new( + PlanImpl::DropTable, + SortOption::None, + RowCountOption::Unknown, + )) } Operator::Truncate(_) if self.contains(ImplementationRuleImpl::Truncate) => { - Some(PhysicalOption::new(PlanImpl::Truncate, SortOption::None)) + Some(PhysicalOption::new( + PlanImpl::Truncate, + SortOption::None, + RowCountOption::Unknown, + )) } _ => None, } @@ -626,18 +693,136 @@ mod tests { use crate::expression::range_detacher::Range; use crate::expression::ScalarExpression; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; - use crate::optimizer::heuristic::batch::HepBatchStrategy; - use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline; + use crate::optimizer::heuristic::batch::{HepBatch, HepBatchStrategy, HepMatchOrder}; + use crate::optimizer::heuristic::optimizer::{HepOptimizer, HepOptimizerPipeline}; use crate::optimizer::rule::implementation::ImplementationRuleImpl; use crate::optimizer::rule::normalization::NormalizationRuleImpl; + use crate::planner::operator::filter::FilterOperator; use crate::planner::operator::sort::SortField; use crate::planner::operator::{Operator, PlanImpl, SortOption}; - use crate::storage::{Storage, Transaction}; + use crate::planner::{Childrens, LogicalPlan, PlanArena, TableArenaCell}; + use crate::storage::{Storage, StorageTransaction}; use crate::types::index::IndexLookup; use crate::types::value::DataValue; use std::ops::Bound; use tempfile::TempDir; + fn filter_chain(length: usize, arena: &mut PlanArena) -> LogicalPlan { + let mut plan = LogicalPlan::new(Operator::Dummy, Childrens::None); + for value in 0..length { + let predicate = + arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(value as i32))); + plan = FilterOperator::build(predicate, plan, false); + } + plan + } + + #[test] + fn local_rule_order_and_sibling_application() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let batch = HepBatch::new( + "Combine Filters", + HepBatchStrategy::fix_point_bottomup(1), + &[NormalizationRuleImpl::CombineFilter], + ); + let mut applied_rules = Vec::new(); + let mut topdown = filter_chain(3, &mut arena); + assert!(HepOptimizer::apply_batch( + &mut topdown, + &batch, + HepMatchOrder::TopDown, + &mut applied_rules, + &mut arena, + None, + )?); + assert!(matches!( + topdown.childrens.only().operator, + Operator::Filter(_) + )); + let mut bottomup = filter_chain(3, &mut arena); + assert!(HepOptimizer::apply_batch( + &mut bottomup, + &batch, + HepMatchOrder::BottomUp, + &mut applied_rules, + &mut arena, + None, + )?); + assert!(matches!( + bottomup.childrens.only().operator, + Operator::Dummy + )); + assert!(!HepOptimizer::apply_batch( + &mut bottomup, + &batch, + HepMatchOrder::BottomUp, + &mut applied_rules, + &mut arena, + None, + )?); + + for strategy in [ + HepBatchStrategy::fix_point_topdown(1), + HepBatchStrategy::fix_point_bottomup(1), + ] { + let plan = LogicalPlan::new( + Operator::Dummy, + Childrens::Twins { + left: Box::new(filter_chain(2, &mut arena)), + right: Box::new(filter_chain(2, &mut arena)), + }, + ); + let pipeline = HepOptimizerPipeline::builder() + .before_batch( + "Combine Filters", + strategy, + &[NormalizationRuleImpl::CombineFilter], + ) + .build(); + let result = pipeline.instantiate(plan).find_best(None, &mut arena)?; + let (left, right) = result.childrens.twins(); + assert!(matches!(left.operator, Operator::Filter(_))); + assert!(matches!(right.operator, Operator::Filter(_))); + assert!(matches!(left.childrens.only().operator, Operator::Dummy)); + assert!(matches!(right.childrens.only().operator, Operator::Dummy)); + } + Ok(()) + } + + #[test] + fn once_stops_each_rule_after_its_first_success() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + for strategy in [ + HepBatchStrategy::once_topdown(), + HepBatchStrategy::once_bottomup(), + ] { + let plan = LogicalPlan::new( + Operator::Dummy, + Childrens::Twins { + left: Box::new(filter_chain(2, &mut arena)), + right: Box::new(filter_chain(2, &mut arena)), + }, + ); + let pipeline = HepOptimizerPipeline::builder() + .before_batch( + "Combine Filters", + strategy, + &[NormalizationRuleImpl::CombineFilter], + ) + .build(); + let result = pipeline.instantiate(plan).find_best(None, &mut arena)?; + let (left, right) = result.childrens.twins(); + assert!(matches!(left.childrens.only().operator, Operator::Dummy)); + assert!(matches!( + right.childrens.only().operator, + Operator::Filter(_) + )); + } + Ok(()) + } + #[test] fn test_find_best_selects_cheapest_scan() -> Result<(), DatabaseError> { let temp_dir = TempDir::new().expect("unable to create temporary working directory"); @@ -684,25 +869,25 @@ mod tests { let plan = binder.bind(&stmt, &mut plan_arena)?; let pipeline = HepOptimizerPipeline::builder() .before_batch( - "Simplify Filter".to_string(), + "Simplify Filter", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::SimplifyFilter], + &[NormalizationRuleImpl::SimplifyFilter], ) .before_batch( - "Predicate Pushdown".to_string(), + "Predicate Pushdown", HepBatchStrategy::fix_point_topdown(10), - vec![ + &[ NormalizationRuleImpl::PushPredicateThroughJoin, NormalizationRuleImpl::PushJoinPredicateIntoScan, NormalizationRuleImpl::PushPredicateIntoScan, ], ) .after_batch( - "Eliminate Index Filter".to_string(), + "Eliminate Index Filter", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::EliminateIndexFilter], + &[NormalizationRuleImpl::EliminateIndexFilter], ) - .implementations(vec![ + .implementations(&[ ImplementationRuleImpl::Projection, ImplementationRuleImpl::Filter, ImplementationRuleImpl::HashJoin, @@ -763,7 +948,6 @@ pub struct HepOptimizerPipeline { before_batches: Vec, after_batches: Vec, implementation_index: ImplementationRuleIndex, - max_local_rules_len: usize, } impl HepOptimizerPipeline { @@ -771,31 +955,7 @@ impl HepOptimizerPipeline { HepOptimizerPipelineBuilder { before_batches: vec![], after_batches: vec![], - implementations: vec![], - } - } - - pub fn new( - before_batches: Vec, - after_batches: Vec, - implementations: Vec, - ) -> Self { - let max_local_rules_len = before_batches - .iter() - .chain(after_batches.iter()) - .flat_map(|batch| batch.steps.iter()) - .filter_map(|step| match step { - HepBatchStep::LocalRewrite(rules) => Some(rules.len()), - HepBatchStep::WholeTree(_) => None, - }) - .max() - .unwrap_or(0); - - Self { - before_batches, - after_batches, - implementation_index: ImplementationRuleIndex::new(implementations), - max_local_rules_len, + implementation_index: ImplementationRuleIndex::new(&[]), } } @@ -805,7 +965,6 @@ impl HepOptimizerPipeline { &self.before_batches, &self.after_batches, &self.implementation_index, - self.max_local_rules_len, ) } } @@ -813,15 +972,15 @@ impl HepOptimizerPipeline { pub struct HepOptimizerPipelineBuilder { before_batches: Vec, after_batches: Vec, - implementations: Vec, + implementation_index: ImplementationRuleIndex, } impl HepOptimizerPipelineBuilder { pub fn before_batch( mut self, - name: String, + name: &'static str, strategy: HepBatchStrategy, - rules: Vec, + rules: &[NormalizationRuleImpl], ) -> Self { self.before_batches .push(HepBatch::new(name, strategy, rules)); @@ -830,25 +989,25 @@ impl HepOptimizerPipelineBuilder { pub fn after_batch( mut self, - name: String, + name: &'static str, strategy: HepBatchStrategy, - rules: Vec, + rules: &[NormalizationRuleImpl], ) -> Self { self.after_batches .push(HepBatch::new(name, strategy, rules)); self } - pub fn implementations(mut self, implementations: Vec) -> Self { - self.implementations = implementations; + pub fn implementations(mut self, implementations: &[ImplementationRuleImpl]) -> Self { + self.implementation_index = ImplementationRuleIndex::new(implementations); self } pub fn build(self) -> HepOptimizerPipeline { - HepOptimizerPipeline::new( - self.before_batches, - self.after_batches, - self.implementations, - ) + HepOptimizerPipeline { + before_batches: self.before_batches, + after_batches: self.after_batches, + implementation_index: self.implementation_index, + } } } diff --git a/src/optimizer/plan_utils.rs b/src/optimizer/plan_utils.rs index 94c9f6c4..19977509 100644 --- a/src/optimizer/plan_utils.rs +++ b/src/optimizer/plan_utils.rs @@ -12,8 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::planner::operator::Operator; -use crate::planner::{Childrens, LogicalPlan}; +use crate::catalog::{ColumnRef, TableName}; +use crate::errors::DatabaseError; +use crate::expression::visitor::ExprVisitor; +use crate::expression::{AliasType, ScalarExpression}; +use crate::planner::operator::project::ProjectOperator; +use crate::planner::operator::table_scan::TableScanOperator; +use crate::planner::operator::visitor::{ChildrenVisit, OperatorVisitor}; +use crate::planner::operator::{Operator, PhysicalOption, RowCountOption}; +use crate::planner::{Childrens, LogicalPlan, MetaArena, PlanArena}; +use crate::types::ColumnId; use std::mem; #[allow(dead_code)] @@ -85,6 +93,31 @@ pub fn children(plan: &LogicalPlan) -> Vec<&LogicalPlan> { } } +pub(crate) fn predicate_scan_mut<'plan>( + plan: &'plan mut LogicalPlan, + arena: &PlanArena, +) -> Option<( + &'plan mut TableScanOperator, + &'plan mut Option, +)> { + match (&mut plan.operator, plan.childrens.as_mut()) { + (Operator::TableScan(scan), _) if scan.limit == (None, None) => { + Some((scan, &mut plan.physical_option)) + } + (Operator::Sort(_), Childrens::Only(input)) => predicate_scan_mut(input, arena), + (Operator::ScalarQueryInit(init), Childrens::Twins { left, .. }) + if init.param_bindings.is_empty() + && matches!( + arena.expression(init.value), + ScalarExpression::InitValue { .. } + ) => + { + predicate_scan_mut(left, arena) + } + _ => None, + } +} + pub fn replace_with_only_child(plan: &mut LogicalPlan) -> bool { if let Childrens::Only(child) = take_childrens(plan) { *plan = *child; @@ -94,18 +127,11 @@ pub fn replace_with_only_child(plan: &mut LogicalPlan) -> bool { } } -#[allow(dead_code)] -pub fn replace_child_with_only_child(plan: &mut LogicalPlan, child_idx: usize) -> bool { - if let Some(child_plan) = child_mut(plan, child_idx) { - return replace_with_only_child(child_plan); - } - false -} - pub fn wrap_child_with(plan: &mut LogicalPlan, child_idx: usize, operator: Operator) -> bool { if let Some(slot) = child_mut(plan, child_idx) { let previous = mem::replace(slot, LogicalPlan::new(operator, Childrens::None)); *slot.childrens = Childrens::Only(Box::new(previous)); + plan.reset_output_schema_cache(); true } else { false @@ -116,6 +142,136 @@ fn take_childrens(plan: &mut LogicalPlan) -> Childrens { mem::replace(&mut *plan.childrens, Childrens::None) } +struct ScanColumnVisitor<'plan, 'arena, A> { + column: Option, + source: Option<(&'plan TableName, ColumnId)>, + arena: &'arena A, +} + +impl<'plan, A: MetaArena> OperatorVisitor<'plan> for ScanColumnVisitor<'plan, '_, A> { + fn visit_plan(&mut self, plan: &'plan LogicalPlan) -> Result<(), DatabaseError> { + if self.column.is_none() { + return Ok(()); + } + self.visit_operator( + &plan.operator, + plan.physical_option.as_ref(), + &mut ChildrenVisit::empty(), + )?; + let follow = match &plan.operator { + Operator::Project(_) | Operator::Filter(_) => true, + Operator::Join(_) | Operator::Aggregate(_) | Operator::Limit(_) | Operator::TopK(_) => { + false + } + _ => plan + .physical_option + .as_ref() + .is_some_and(|option| matches!(option.row_count, RowCountOption::Follow)), + }; + if follow && self.column.is_some() { + self.visit_children(&plan.childrens)?; + } + Ok(()) + } + + fn visit_children(&mut self, childrens: &'plan Childrens) -> Result<(), DatabaseError> { + match childrens { + Childrens::Only(input) | Childrens::Twins { left: input, .. } => self.visit_plan(input), + Childrens::None => Ok(()), + } + } + + fn visit_table_scan(&mut self, scan: &'plan TableScanOperator) -> Result<(), DatabaseError> { + if scan.limit == (None, None) { + if let Some(source) = self.column.and_then(|column| { + scan.columns + .iter() + .find(|source| self.arena.same_column(**source, column)) + }) { + self.source = self + .arena + .column(*source) + .id() + .map(|id| (&scan.table_name, id)); + } + } + self.column = None; + Ok(()) + } + + fn visit_project( + &mut self, + project: &'plan ProjectOperator, + _children: &mut ChildrenVisit<'plan>, + ) -> Result<(), DatabaseError> { + let Some(column) = self.column.take() else { + return Ok(()); + }; + let arena = self.arena; + for &candidate in &project.exprs { + let source = match arena.expression(candidate) { + ScalarExpression::Alias { + expr, + alias: AliasType::Name(_), + } if arena.column(column).id().is_none() + && arena.column(column).table_name().is_none() + && candidate.output_name(arena) == arena.column(column).name() => + { + *expr + } + ScalarExpression::Alias { + expr, + alias: AliasType::Expr(alias), + } if matches!(arena.expression(*alias), ScalarExpression::ColumnRef { column: target, .. } + if arena.same_column(*target, column)) => + { + *expr + } + ScalarExpression::ColumnRef { column: target, .. } + if arena.same_column(*target, column) => + { + candidate + } + _ => continue, + }; + return ExprVisitor::visit(self, source, arena); + } + Ok(()) + } +} + +impl ExprVisitor for ScanColumnVisitor<'_, '_, A> { + fn visit_expression( + &mut self, + expr: &ScalarExpression, + _arena: &A, + ) -> Result { + Ok(matches!( + expr, + ScalarExpression::Alias { .. } | ScalarExpression::ColumnRef { .. } + )) + } + + fn visit_column_ref(&mut self, column: &ColumnRef, _: usize) -> Result<(), DatabaseError> { + self.column = Some(*column); + Ok(()) + } +} + +pub(crate) fn scan_column<'plan>( + plan: &'plan LogicalPlan, + column: ColumnRef, + arena: &PlanArena, +) -> Result, DatabaseError> { + let mut visitor = ScanColumnVisitor { + column: Some(column), + source: None, + arena, + }; + visitor.visit_plan(plan)?; + Ok(visitor.source) +} + // GRCOV_EXCL_START #[cfg(test)] mod tests { @@ -192,20 +348,10 @@ mod tests { let mut parent = LogicalPlan::new( Operator::Dummy, Childrens::Twins { - left: Box::new(LogicalPlan::new( - Operator::ShowView, - Childrens::Only(Box::new(leaf(Operator::ShowTable))), - )), + left: Box::new(leaf(Operator::ShowTable)), right: Box::new(leaf(Operator::Dummy)), }, ); - assert!(replace_child_with_only_child(&mut parent, 0)); - assert!(matches!( - child(&parent, 0).unwrap().operator, - Operator::ShowTable - )); - assert!(!replace_child_with_only_child(&mut parent, 1)); - assert!(!replace_child_with_only_child(&mut parent, 2)); assert!(wrap_child_with(&mut parent, 1, Operator::ShowView)); let wrapped = child(&parent, 1).unwrap(); diff --git a/src/optimizer/rule/implementation/ddl/add_column.rs b/src/optimizer/rule/implementation/ddl/add_column.rs index d2f3bf17..0a82b516 100644 --- a/src/optimizer/rule/implementation/ddl/add_column.rs +++ b/src/optimizer/rule/implementation/ddl/add_column.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct AddColumnImplementation; single_mapping!( AddColumnImplementation, ADD_COLUMN_PATTERN, - PhysicalOption::new(PlanImpl::AddColumn, SortOption::None) + PhysicalOption::new( + PlanImpl::AddColumn, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/ddl/change_column.rs b/src/optimizer/rule/implementation/ddl/change_column.rs index e9613793..7bc31c6e 100644 --- a/src/optimizer/rule/implementation/ddl/change_column.rs +++ b/src/optimizer/rule/implementation/ddl/change_column.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct ChangeColumnImplementation; single_mapping!( ChangeColumnImplementation, CHANGE_COLUMN_PATTERN, - PhysicalOption::new(PlanImpl::ChangeColumn, SortOption::None) + PhysicalOption::new( + PlanImpl::ChangeColumn, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/ddl/create_table.rs b/src/optimizer/rule/implementation/ddl/create_table.rs index c50fc810..f3b0ec41 100644 --- a/src/optimizer/rule/implementation/ddl/create_table.rs +++ b/src/optimizer/rule/implementation/ddl/create_table.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct CreateTableImplementation; single_mapping!( CreateTableImplementation, CREATE_TABLE_PATTERN, - PhysicalOption::new(PlanImpl::CreateTable, SortOption::None) + PhysicalOption::new( + PlanImpl::CreateTable, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/ddl/drop_column.rs b/src/optimizer/rule/implementation/ddl/drop_column.rs index a8b86ab7..d4684982 100644 --- a/src/optimizer/rule/implementation/ddl/drop_column.rs +++ b/src/optimizer/rule/implementation/ddl/drop_column.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct DropColumnImplementation; single_mapping!( DropColumnImplementation, DROP_COLUMN_PATTERN, - PhysicalOption::new(PlanImpl::DropColumn, SortOption::None) + PhysicalOption::new( + PlanImpl::DropColumn, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/ddl/drop_table.rs b/src/optimizer/rule/implementation/ddl/drop_table.rs index 351663d4..26e09322 100644 --- a/src/optimizer/rule/implementation/ddl/drop_table.rs +++ b/src/optimizer/rule/implementation/ddl/drop_table.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct DropTableImplementation; single_mapping!( DropTableImplementation, DROP_TABLE_PATTERN, - PhysicalOption::new(PlanImpl::DropTable, SortOption::None) + PhysicalOption::new( + PlanImpl::DropTable, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/ddl/truncate.rs b/src/optimizer/rule/implementation/ddl/truncate.rs index c9cb9248..6e3ff66c 100644 --- a/src/optimizer/rule/implementation/ddl/truncate.rs +++ b/src/optimizer/rule/implementation/ddl/truncate.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct TruncateImplementation; single_mapping!( TruncateImplementation, TRUNCATE_PATTERN, - PhysicalOption::new(PlanImpl::Truncate, SortOption::None) + PhysicalOption::new( + PlanImpl::Truncate, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/dml/analyze.rs b/src/optimizer/rule/implementation/dml/analyze.rs index 0f571861..6bcc6dc8 100644 --- a/src/optimizer/rule/implementation/dml/analyze.rs +++ b/src/optimizer/rule/implementation/dml/analyze.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,5 @@ pub struct AnalyzeImplementation; single_mapping!( AnalyzeImplementation, ANALYZE_PATTERN, - PhysicalOption::new(PlanImpl::Analyze, SortOption::None) + PhysicalOption::new(PlanImpl::Analyze, SortOption::None, RowCountOption::Unknown) ); diff --git a/src/optimizer/rule/implementation/dml/copy_from_file.rs b/src/optimizer/rule/implementation/dml/copy_from_file.rs index 593fc843..b1cf3cd0 100644 --- a/src/optimizer/rule/implementation/dml/copy_from_file.rs +++ b/src/optimizer/rule/implementation/dml/copy_from_file.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct CopyFromFileImplementation; single_mapping!( CopyFromFileImplementation, COPY_FROM_FILE_PATTERN, - PhysicalOption::new(PlanImpl::CopyFromFile, SortOption::None) + PhysicalOption::new( + PlanImpl::CopyFromFile, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/dml/copy_to_file.rs b/src/optimizer/rule/implementation/dml/copy_to_file.rs index d1f9f234..f668c230 100644 --- a/src/optimizer/rule/implementation/dml/copy_to_file.rs +++ b/src/optimizer/rule/implementation/dml/copy_to_file.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct CopyToFileImplementation; single_mapping!( CopyToFileImplementation, COPY_TO_FILE_PATTERN, - PhysicalOption::new(PlanImpl::CopyToFile, SortOption::None) + PhysicalOption::new( + PlanImpl::CopyToFile, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/dml/delete.rs b/src/optimizer/rule/implementation/dml/delete.rs index 80f3cb18..436b8d49 100644 --- a/src/optimizer/rule/implementation/dml/delete.rs +++ b/src/optimizer/rule/implementation/dml/delete.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,5 @@ pub struct DeleteImplementation; single_mapping!( DeleteImplementation, DELETE_PATTERN, - PhysicalOption::new(PlanImpl::Delete, SortOption::None) + PhysicalOption::new(PlanImpl::Delete, SortOption::None, RowCountOption::Unknown) ); diff --git a/src/optimizer/rule/implementation/dml/insert.rs b/src/optimizer/rule/implementation/dml/insert.rs index 5f71fa9a..0300c55a 100644 --- a/src/optimizer/rule/implementation/dml/insert.rs +++ b/src/optimizer/rule/implementation/dml/insert.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,5 @@ pub struct InsertImplementation; single_mapping!( InsertImplementation, INSERT_PATTERN, - PhysicalOption::new(PlanImpl::Insert, SortOption::None) + PhysicalOption::new(PlanImpl::Insert, SortOption::None, RowCountOption::Unknown) ); diff --git a/src/optimizer/rule/implementation/dml/update.rs b/src/optimizer/rule/implementation/dml/update.rs index df120c2b..075a4cdc 100644 --- a/src/optimizer/rule/implementation/dml/update.rs +++ b/src/optimizer/rule/implementation/dml/update.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,5 @@ pub struct UpdateImplementation; single_mapping!( UpdateImplementation, UPDATE_PATTERN, - PhysicalOption::new(PlanImpl::Update, SortOption::None) + PhysicalOption::new(PlanImpl::Update, SortOption::None, RowCountOption::Unknown) ); diff --git a/src/optimizer/rule/implementation/dql/aggregate.rs b/src/optimizer/rule/implementation/dql/aggregate.rs index 2466076f..598a8abd 100644 --- a/src/optimizer/rule/implementation/dql/aggregate.rs +++ b/src/optimizer/rule/implementation/dql/aggregate.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -45,7 +45,11 @@ pub struct GroupByAggregateImplementation; single_mapping!( GroupByAggregateImplementation, GROUP_BY_AGGREGATE_PATTERN, - PhysicalOption::new(PlanImpl::HashAggregate, SortOption::None) + PhysicalOption::new( + PlanImpl::HashAggregate, + SortOption::None, + RowCountOption::Unknown + ) ); pub struct SimpleAggregateImplementation; @@ -53,5 +57,9 @@ pub struct SimpleAggregateImplementation; single_mapping!( SimpleAggregateImplementation, SIMPLE_AGGREGATE_PATTERN, - PhysicalOption::new(PlanImpl::SimpleAggregate, SortOption::None) + PhysicalOption::new( + PlanImpl::SimpleAggregate, + SortOption::None, + RowCountOption::Estimated { rows: 1, ndv: None } + ) ); diff --git a/src/optimizer/rule/implementation/dql/dummy.rs b/src/optimizer/rule/implementation/dql/dummy.rs index 9fd3cdd7..5e83d9c8 100644 --- a/src/optimizer/rule/implementation/dql/dummy.rs +++ b/src/optimizer/rule/implementation/dql/dummy.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,5 @@ pub struct DummyImplementation; single_mapping!( DummyImplementation, DUMMY_PATTERN, - PhysicalOption::new(PlanImpl::Dummy, SortOption::None) + PhysicalOption::new(PlanImpl::Dummy, SortOption::None, RowCountOption::Unknown) ); diff --git a/src/optimizer/rule/implementation/dql/filter.rs b/src/optimizer/rule/implementation/dql/filter.rs index cd254bc2..65390320 100644 --- a/src/optimizer/rule/implementation/dql/filter.rs +++ b/src/optimizer/rule/implementation/dql/filter.rs @@ -12,13 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::catalog::ColumnRef; use crate::errors::DatabaseError; +use crate::expression::visitor::ExprVisitor; +use crate::expression::{BinaryOperator, ScalarExpression}; use crate::optimizer::core::pattern::Pattern; -use crate::optimizer::core::rule::BestPhysicalOption; -use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; +use crate::optimizer::core::rule::{ + keep_best_physical_option, BestPhysicalOption, HintState, ImplementationRule, MatchPattern, +}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; -use crate::single_mapping; +use crate::optimizer::plan_utils::scan_column; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; +use crate::planner::{Childrens, ExprRef, LogicalPlan, MetaArena, PlanArena}; +use crate::types::value::DataValue; +use std::collections::HashSet; use std::sync::LazyLock; static FILTER_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -28,8 +35,390 @@ static FILTER_PATTERN: LazyLock = LazyLock::new(|| Pattern { #[derive(Clone)] pub struct FilterImplementation; -single_mapping!( - FilterImplementation, - FILTER_PATTERN, - PhysicalOption::new(PlanImpl::Filter, SortOption::Follow) -); +impl MatchPattern for FilterImplementation { + fn pattern(&self) -> &Pattern { + &FILTER_PATTERN + } +} + +impl ImplementationRule for FilterImplementation { + fn update_best_option( + &self, + plan: &LogicalPlan, + arena: &PlanArena, + loader: &StatisticMetaLoader<'_>, + _: &mut HintState, + best_physical_option: &mut BestPhysicalOption, + ) -> Result<(), DatabaseError> { + let rows = estimate_filter_rows(plan, arena, loader)?; + let row_count = rows.map_or(RowCountOption::Follow, |rows| RowCountOption::Estimated { + rows, + ndv: None, + }); + keep_best_physical_option( + best_physical_option, + PhysicalOption::new(PlanImpl::Filter, SortOption::Follow, row_count), + None, + ); + Ok(()) + } +} + +struct PointFilterVisitor<'a, A> { + arena: &'a A, + column: Option, + values: HashSet<&'a DataValue>, + supported: bool, +} + +impl<'a, A: MetaArena> PointFilterVisitor<'a, A> { + fn bind_column(&mut self, column: ColumnRef) { + if self + .column + .is_some_and(|previous| !self.arena.same_column(previous, column)) + { + self.supported = false; + } else { + self.column = Some(column); + } + } + + fn add_value(&mut self, column: ColumnRef, expr: ExprRef) { + let arena = self.arena; + let ScalarExpression::Constant(value) = arena.expression(expr) else { + self.supported = false; + return; + }; + if value.is_null() { + return; + } + if value.logical_type() != *arena.column(column).datatype() { + self.supported = false; + return; + } + self.values.insert(value); + } +} + +impl ExprVisitor for PointFilterVisitor<'_, A> { + fn visit_expression_ref(&mut self, expr: ExprRef, _arena: &A) -> Result { + if !self.supported { + return Ok(false); + } + let arena = self.arena; + match arena.expression(expr) { + ScalarExpression::Binary { + op: BinaryOperator::Or, + .. + } => return Ok(true), + ScalarExpression::Binary { + op: BinaryOperator::Eq, + left_expr, + right_expr, + .. + } => { + let (column, value) = + match (arena.expression(*left_expr), arena.expression(*right_expr)) { + ( + ScalarExpression::ColumnRef { column, .. }, + ScalarExpression::Constant(_), + ) => (*column, *right_expr), + ( + ScalarExpression::Constant(_), + ScalarExpression::ColumnRef { column, .. }, + ) => (*column, *left_expr), + _ => { + self.supported = false; + return Ok(false); + } + }; + self.bind_column(column); + self.add_value(column, value); + } + ScalarExpression::In { + negated: false, + expr, + args, + .. + } => { + let ScalarExpression::ColumnRef { column, .. } = arena.expression(*expr) else { + self.supported = false; + return Ok(false); + }; + self.bind_column(*column); + for &arg in args { + self.add_value(*column, arg); + } + } + _ => self.supported = false, + } + Ok(false) + } +} + +fn estimate_filter_rows( + plan: &LogicalPlan, + arena: &PlanArena, + loader: &StatisticMetaLoader<'_>, +) -> Result, DatabaseError> { + let (Operator::Filter(filter), Childrens::Only(input)) = + (&plan.operator, plan.childrens.as_ref()) + else { + return Ok(None); + }; + let Some(rows) = input.estimated_rows() else { + return Ok(None); + }; + // Do not apply the same point selectivity again to an already restricted + // index or an earlier estimated filter. Start with unrestricted scan inputs. + if !matches!( + input.row_count_option().map(|option| &option.plan), + Some(PlanImpl::SeqScan) + ) { + return Ok(None); + } + let mut visitor = PointFilterVisitor { + arena, + column: None, + values: HashSet::new(), + supported: true, + }; + visitor.visit(filter.predicate, arena)?; + let Some(column) = visitor.column.filter(|_| visitor.supported) else { + return Ok(None); + }; + let Some((table, column)) = scan_column(input, column, arena)? else { + return Ok(None); + }; + let Some(meta) = loader.load_column(table, column) else { + return Ok(None); + }; + let ndv = meta.distinct_values_len(); + if ndv == 0 || visitor.values.is_empty() || meta.values_len() == 0 { + return Ok(Some(0)); + } + // Uniform single-column estimate only: no products of correlated selectivities. + let non_null_fraction = + meta.values_len() as f64 / meta.values_len().saturating_add(meta.null_count()) as f64; + let selectivity = visitor.values.len().min(ndv) as f64 / ndv as f64; + Ok(Some( + ((rows as f64 * non_null_fraction * selectivity).ceil() as usize).min(rows), + )) +} + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::expression::AliasType; + use crate::optimizer::core::histogram::{HistogramMetaBuilder, DEFAULT_HLL_PRECISION}; + use crate::planner::operator::filter::FilterOperator; + use crate::planner::operator::project::ProjectOperator; + use crate::planner::operator::table_scan::TableScanOperator; + use crate::planner::TableArenaCell; + use crate::storage::StatisticsMetaCache; + use crate::types::LogicalType; + + fn input(arena: &mut PlanArena, rows: usize) -> (LogicalPlan, [ExprRef; 2]) { + let mut columns = Vec::new(); + let mut keys = Vec::new(); + for id in 0..2 { + let mut column = ColumnCatalog::new( + format!("c{id}"), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), + ); + column.set_ref_table("t".into(), id, false); + let column = arena.alloc_column(column); + columns.push(column); + keys.push(arena.alloc_expression(ScalarExpression::column_expr(column, id as usize))); + } + let mut plan = LogicalPlan::new( + Operator::TableScan(TableScanOperator { + table_name: "t".into(), + columns, + limit: (None, None), + index_infos: vec![], + with_pk: false, + }), + Childrens::None, + ); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { rows, ndv: None }, + )); + (plan, [keys[0], keys[1]]) + } + + fn statistics(non_null: usize, nulls: usize, ndv: usize) -> StatisticsMetaCache { + let mut builder = HistogramMetaBuilder::new(DEFAULT_HLL_PRECISION).unwrap(); + for i in 0..non_null { + builder.append(&DataValue::Int32((i % ndv) as i32)); + } + for _ in 0..nulls { + builder.append(&DataValue::Null); + } + let mut cache = StatisticsMetaCache::default(); + cache.columns.insert(("t".into(), 0), builder.build()); + cache + } + + fn binary(arena: &mut PlanArena, op: BinaryOperator, left: ExprRef, right: ExprRef) -> ExprRef { + arena.alloc_expression(ScalarExpression::Binary { + op, + left_expr: left, + right_expr: right, + evaluator: None, + ty: LogicalType::Boolean, + }) + } + + fn equal(arena: &mut PlanArena, column: ExprRef, value: DataValue) -> ExprRef { + let value = arena.alloc_expression(ScalarExpression::Constant(value)); + binary(arena, BinaryOperator::Eq, column, value) + } + + fn in_list( + arena: &mut PlanArena, + column: ExprRef, + values: &[DataValue], + negated: bool, + ) -> ExprRef { + let args = values + .iter() + .map(|value| arena.alloc_expression(ScalarExpression::Constant(value.clone()))) + .collect(); + arena.alloc_expression(ScalarExpression::In { + negated, + expr: column, + args, + evaluator: None, + }) + } + + fn estimate( + arena: &PlanArena, + cache: &StatisticsMetaCache, + predicate: ExprRef, + input: &LogicalPlan, + ) -> Result, DatabaseError> { + let plan = FilterOperator::build(predicate, input.clone(), false); + estimate_filter_rows(&plan, arena, &StatisticMetaLoader::new(cache)) + } + + #[test] + fn point_predicates_use_column_stats() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (scan, keys) = input(&mut arena, 100); + let cache = statistics(100, 0, 10); + let eq = equal(&mut arena, keys[0], DataValue::Int32(1)); + let value = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(1))); + let reversed = binary(&mut arena, BinaryOperator::Eq, value, keys[0]); + let values = [1, 1, 2].map(DataValue::Int32); + let list = in_list( + &mut arena, + keys[0], + &[&values[..], &[DataValue::Null]].concat(), + false, + ); + let union = binary(&mut arena, BinaryOperator::Or, eq, list); + let many = in_list( + &mut arena, + keys[0], + &(0..20).map(DataValue::Int32).collect::>(), + false, + ); + let null = equal(&mut arena, keys[0], DataValue::Null); + // Duplicates and NULLs do not count as points; selectivity is capped at 1. + for (predicate, rows) in [ + (eq, 10), + (reversed, 10), + (list, 20), + (union, 20), + (many, 100), + (null, 0), + ] { + assert_eq!(estimate(&arena, &cache, predicate, &scan)?, Some(rows)); + } + // NULLs in the column scale the estimate by the non-NULL fraction. + assert_eq!( + estimate(&arena, &statistics(75, 25, 5), eq, &scan)?, + Some(15) + ); + let (empty, _) = input(&mut arena, 0); + assert_eq!(estimate(&arena, &cache, eq, &empty)?, Some(0)); + + // A column renamed by a projection still resolves to the scan's statistics. + let alias = arena.alloc_expression(ScalarExpression::Alias { + expr: keys[0], + alias: AliasType::Name("alias".into()), + }); + let column = alias.output_column_ref(&mut arena); + let key = arena.alloc_expression(ScalarExpression::column_expr(column, 0)); + let mut project = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![alias])), + Childrens::Only(Box::new(scan)), + ); + project.physical_option = Some(PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + )); + let aliased = equal(&mut arena, key, DataValue::Int32(1)); + assert_eq!(estimate(&arena, &cache, aliased, &project)?, Some(10)); + Ok(()) + } + + #[test] + fn unsupported_predicates_and_restricted_inputs_keep_follow() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (scan, keys) = input(&mut arena, 100); + let cache = statistics(100, 0, 10); + let eq = equal(&mut arena, keys[0], DataValue::Int32(1)); + let other = equal(&mut arena, keys[1], DataValue::Int32(2)); + let multi_column = binary(&mut arena, BinaryOperator::Or, eq, other); + let column_eq = binary(&mut arena, BinaryOperator::Eq, keys[0], keys[1]); + let not_in = in_list(&mut arena, keys[0], &[DataValue::Int32(1)], true); + let wrong_type = equal(&mut arena, keys[0], DataValue::Int64(1)); + for predicate in [multi_column, column_eq, not_in, wrong_type] { + assert_eq!(estimate(&arena, &cache, predicate, &scan)?, None); + } + + // Inputs already restricted by a limit or an earlier estimate are not reduced again. + let mut limited_scan = scan.clone(); + let Operator::TableScan(op) = &mut limited_scan.operator else { + unreachable!() + }; + op.limit = (Some(1), Some(5)); + assert_eq!(estimate(&arena, &cache, eq, &limited_scan)?, None); + let mut filtered = FilterOperator::build(eq, scan.clone(), false); + filtered.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Estimated { + rows: 10, + ndv: None, + }, + )); + assert_eq!(estimate(&arena, &cache, eq, &filtered)?, None); + + // Without column statistics the rule falls back to the input row count. + let plan = FilterOperator::build(eq, scan, false); + let mut best = None; + let empty_cache = StatisticsMetaCache::default(); + FilterImplementation.update_best_option( + &plan, + &arena, + &StatisticMetaLoader::new(&empty_cache), + &mut HintState::default(), + &mut best, + )?; + assert_eq!(best.unwrap().0.row_count, RowCountOption::Follow); + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/optimizer/rule/implementation/dql/function_scan.rs b/src/optimizer/rule/implementation/dql/function_scan.rs index 1b7c7f92..bf59d42a 100644 --- a/src/optimizer/rule/implementation/dql/function_scan.rs +++ b/src/optimizer/rule/implementation/dql/function_scan.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct FunctionScanImplementation; single_mapping!( FunctionScanImplementation, FUNCTION_SCAN_PATTERN, - PhysicalOption::new(PlanImpl::FunctionScan, SortOption::None) + PhysicalOption::new( + PlanImpl::FunctionScan, + SortOption::None, + RowCountOption::Unknown + ) ); diff --git a/src/optimizer/rule/implementation/dql/join.rs b/src/optimizer/rule/implementation/dql/join.rs index 6e703ae7..e63d8641 100644 --- a/src/optimizer/rule/implementation/dql/join.rs +++ b/src/optimizer/rule/implementation/dql/join.rs @@ -16,7 +16,8 @@ use crate::errors::DatabaseError; use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, RowCountOption, SortOption}; +use crate::planner::LogicalPlan; use std::sync::LazyLock; static JOIN_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -35,15 +36,20 @@ impl MatchPattern for JoinImplementation { impl ImplementationRule for JoinImplementation { fn update_best_option( &self, - op: &Operator, + plan: &LogicalPlan, _: &crate::planner::PlanArena, _: &StatisticMetaLoader<'_>, + _: &mut crate::optimizer::core::rule::HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { - let Operator::Join(operator) = op else { + let Operator::Join(operator) = &plan.operator else { unreachable!() }; - let physical_option = PhysicalOption::new(operator.plan_impl(), SortOption::None); + let physical_option = PhysicalOption::new( + operator.plan_impl(), + SortOption::None, + RowCountOption::Unknown, + ); crate::optimizer::core::rule::keep_best_physical_option( best_physical_option, physical_option, diff --git a/src/optimizer/rule/implementation/dql/limit.rs b/src/optimizer/rule/implementation/dql/limit.rs index 10e17b6b..6654468b 100644 --- a/src/optimizer/rule/implementation/dql/limit.rs +++ b/src/optimizer/rule/implementation/dql/limit.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,5 @@ pub struct LimitImplementation; single_mapping!( LimitImplementation, LIMIT_PATTERN, - PhysicalOption::new(PlanImpl::Limit, SortOption::Follow) + PhysicalOption::new(PlanImpl::Limit, SortOption::Follow, RowCountOption::Unknown) ); diff --git a/src/optimizer/rule/implementation/dql/mark_apply.rs b/src/optimizer/rule/implementation/dql/mark_apply.rs index c79d3dc4..ec778e06 100644 --- a/src/optimizer/rule/implementation/dql/mark_apply.rs +++ b/src/optimizer/rule/implementation/dql/mark_apply.rs @@ -14,11 +14,11 @@ use crate::errors::DatabaseError; use crate::optimizer::core::pattern::Pattern; -use crate::optimizer::core::rule::BestPhysicalOption; +use crate::optimizer::core::rule::{keep_best_physical_option, BestPhysicalOption}; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; -use crate::single_mapping; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; +use crate::planner::{LogicalPlan, PlanArena}; use std::sync::LazyLock; static MARK_APPLY_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -28,8 +28,34 @@ static MARK_APPLY_PATTERN: LazyLock = LazyLock::new(|| Pattern { #[derive(Clone)] pub struct MarkApplyImplementation; -single_mapping!( - MarkApplyImplementation, - MARK_APPLY_PATTERN, - PhysicalOption::new(PlanImpl::MarkApply, SortOption::Follow) -); +impl MatchPattern for MarkApplyImplementation { + fn pattern(&self) -> &Pattern { + &MARK_APPLY_PATTERN + } +} + +impl ImplementationRule for MarkApplyImplementation { + fn update_best_option( + &self, + plan: &LogicalPlan, + _: &PlanArena, + _: &StatisticMetaLoader<'_>, + _: &mut crate::optimizer::core::rule::HintState, + best_physical_option: &mut BestPhysicalOption, + ) -> Result<(), DatabaseError> { + if !matches!(plan.operator, Operator::MarkApply(_)) { + return Ok(()); + } + // A mark adds one value to each left row. + keep_best_physical_option( + best_physical_option, + PhysicalOption::new( + PlanImpl::MarkApply, + SortOption::Follow, + RowCountOption::Follow, + ), + None, + ); + Ok(()) + } +} diff --git a/src/optimizer/rule/implementation/dql/projection.rs b/src/optimizer/rule/implementation/dql/projection.rs index 02806e9a..cb0e7d32 100644 --- a/src/optimizer/rule/implementation/dql/projection.rs +++ b/src/optimizer/rule/implementation/dql/projection.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct ProjectionImplementation; single_mapping!( ProjectionImplementation, PROJECTION_PATTERN, - PhysicalOption::new(PlanImpl::Project, SortOption::Follow) + PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow + ) ); diff --git a/src/optimizer/rule/implementation/dql/scalar_query_init.rs b/src/optimizer/rule/implementation/dql/scalar_query_init.rs index 9be962ba..a612ff44 100644 --- a/src/optimizer/rule/implementation/dql/scalar_query_init.rs +++ b/src/optimizer/rule/implementation/dql/scalar_query_init.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct ScalarQueryInitImplementation; single_mapping!( ScalarQueryInitImplementation, SCALAR_APPLY_PATTERN, - PhysicalOption::new(PlanImpl::ScalarQueryInit, SortOption::Follow) + PhysicalOption::new( + PlanImpl::ScalarQueryInit, + SortOption::Follow, + RowCountOption::Follow + ) ); diff --git a/src/optimizer/rule/implementation/dql/scalar_subquery.rs b/src/optimizer/rule/implementation/dql/scalar_subquery.rs index dcaaf88a..ba61d91f 100644 --- a/src/optimizer/rule/implementation/dql/scalar_subquery.rs +++ b/src/optimizer/rule/implementation/dql/scalar_subquery.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,9 @@ pub struct ScalarSubqueryImplementation; single_mapping!( ScalarSubqueryImplementation, SCALAR_SUBQUERY_PATTERN, - PhysicalOption::new(PlanImpl::ScalarSubquery, SortOption::Follow) + PhysicalOption::new( + PlanImpl::ScalarSubquery, + SortOption::Follow, + RowCountOption::Follow + ) ); diff --git a/src/optimizer/rule/implementation/dql/sort.rs b/src/optimizer/rule/implementation/dql/sort.rs index 6ea317e2..c92c1900 100644 --- a/src/optimizer/rule/implementation/dql/sort.rs +++ b/src/optimizer/rule/implementation/dql/sort.rs @@ -16,7 +16,8 @@ use crate::errors::DatabaseError; use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; +use crate::planner::LogicalPlan; use std::sync::LazyLock; static SORT_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -35,12 +36,13 @@ impl MatchPattern for SortImplementation { impl ImplementationRule for SortImplementation { fn update_best_option( &self, - op: &Operator, + plan: &LogicalPlan, _: &crate::planner::PlanArena, _: &StatisticMetaLoader<'_>, + _: &mut crate::optimizer::core::rule::HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { - if let Operator::Sort(op) = op { + if let Operator::Sort(op) = &plan.operator { crate::optimizer::core::rule::keep_best_physical_option( best_physical_option, PhysicalOption::new( @@ -49,6 +51,7 @@ impl ImplementationRule for SortImplementation { fields: op.sort_fields.clone(), ignore_prefix_len: 0, }, + RowCountOption::Follow, ), None, ); diff --git a/src/optimizer/rule/implementation/dql/table_scan.rs b/src/optimizer/rule/implementation/dql/table_scan.rs index e40663af..3ba4b987 100644 --- a/src/optimizer/rule/implementation/dql/table_scan.rs +++ b/src/optimizer/rule/implementation/dql/table_scan.rs @@ -15,9 +15,12 @@ use crate::errors::DatabaseError; use crate::expression::range_detacher::Range; use crate::optimizer::core::pattern::Pattern; -use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; +use crate::optimizer::core::rule::{ + keep_best_physical_option, BestPhysicalOption, ImplementationRule, MatchPattern, +}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; +use crate::planner::{LogicalPlan, PlanArena}; use crate::types::index::{IndexLookup, IndexType}; use std::sync::LazyLock; @@ -37,12 +40,13 @@ impl MatchPattern for SeqScanImplementation { impl ImplementationRule for SeqScanImplementation { fn update_best_option( &self, - op: &Operator, - arena: &crate::planner::PlanArena, + plan: &LogicalPlan, + arena: &PlanArena, loader: &StatisticMetaLoader<'_>, + _: &mut crate::optimizer::core::rule::HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { - if let Operator::TableScan(scan_op) = op { + if let Operator::TableScan(scan_op) = &plan.operator { let cost = scan_op .index_infos .iter() @@ -55,13 +59,15 @@ impl ImplementationRule for SeqScanImplementation { .map(|index_info| loader.load(&scan_op.table_name, arena.index(index_info.meta).id)) .transpose()? .flatten() - .map(|statistics_meta| statistics_meta.histogram().values_len()); + .map(|index_statistics| index_statistics.histogram().values_len()); - crate::optimizer::core::rule::keep_best_physical_option( - best_physical_option, - PhysicalOption::new(PlanImpl::SeqScan, SortOption::None), - cost, + let option = PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + cost.map(|rows| RowCountOption::Estimated { rows, ndv: None }) + .unwrap_or(RowCountOption::Unknown), ); + keep_best_physical_option(best_physical_option, option, cost); Ok(()) } else { unreachable!("invalid operator!") @@ -80,23 +86,35 @@ impl MatchPattern for IndexScanImplementation { impl ImplementationRule for IndexScanImplementation { fn update_best_option( &self, - op: &Operator, - arena: &crate::planner::PlanArena, + plan: &LogicalPlan, + arena: &PlanArena, loader: &StatisticMetaLoader<'_>, + state: &mut crate::optimizer::core::rule::HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { - if let Operator::TableScan(scan_op) = op { + if let Operator::TableScan(scan_op) = &plan.operator { for index_info in scan_op.index_infos.iter() { let Some(IndexLookup::Static(range)) = &index_info.lookup else { continue; }; let mut cost = None; + let mut estimated_rows = None; + let mut ndv = None; let index_meta = arena.index(index_info.meta); - if let Some(mut row_count) = - loader.collect_count(&scan_op.table_name, index_meta.id, range)? - { + // A PK/Unique equality range has an exact row count even without statistics. + let counted = loader + .collect_count( + &scan_op.table_name, + index_meta.id, + range, + &mut state.histogram_buf, + )? + .or_else(|| unique_eq_row_count(index_meta.ty, range).map(|rows| (rows, rows))); + if let Some((mut row_count, distinct_count)) = counted { row_count = adjust_index_row_count(index_meta.ty, range, row_count); + estimated_rows = Some(row_count); + ndv = Some(distinct_count.min(row_count)); if index_info.covered_deserializers.is_none() && !matches!(index_meta.ty, IndexType::PrimaryKey { .. }) { @@ -142,14 +160,14 @@ impl ImplementationRule for IndexScanImplementation { } } - crate::optimizer::core::rule::keep_best_physical_option( - best_physical_option, - PhysicalOption::new( - PlanImpl::IndexScan(Box::new(index_info.clone())), - index_info.sort_option.clone(), - ), - cost, + let option = PhysicalOption::new( + PlanImpl::IndexScan(Box::new(index_info.clone())), + index_info.sort_option.clone(), + estimated_rows + .map(|rows| RowCountOption::Estimated { rows, ndv }) + .unwrap_or(RowCountOption::Unknown), ); + keep_best_physical_option(best_physical_option, option, cost); } Ok(()) diff --git a/src/optimizer/rule/implementation/dql/top_k.rs b/src/optimizer/rule/implementation/dql/top_k.rs index 5c024221..310d0a2c 100644 --- a/src/optimizer/rule/implementation/dql/top_k.rs +++ b/src/optimizer/rule/implementation/dql/top_k.rs @@ -16,7 +16,8 @@ use crate::errors::DatabaseError; use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; +use crate::planner::LogicalPlan; use std::sync::LazyLock; static TOPK_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -35,12 +36,13 @@ impl MatchPattern for TopKImplementation { impl ImplementationRule for TopKImplementation { fn update_best_option( &self, - op: &Operator, + plan: &LogicalPlan, _: &crate::planner::PlanArena, _: &StatisticMetaLoader<'_>, + _: &mut crate::optimizer::core::rule::HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { - if let Operator::TopK(op) = op { + if let Operator::TopK(op) = &plan.operator { crate::optimizer::core::rule::keep_best_physical_option( best_physical_option, PhysicalOption::new( @@ -49,6 +51,7 @@ impl ImplementationRule for TopKImplementation { fields: op.sort_fields.clone(), ignore_prefix_len: 0, }, + RowCountOption::Unknown, ), None, ); diff --git a/src/optimizer/rule/implementation/dql/values.rs b/src/optimizer/rule/implementation/dql/values.rs index 99ec57f0..badfcd2f 100644 --- a/src/optimizer/rule/implementation/dql/values.rs +++ b/src/optimizer/rule/implementation/dql/values.rs @@ -17,7 +17,7 @@ use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::BestPhysicalOption; use crate::optimizer::core::rule::{ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; use crate::single_mapping; use std::sync::LazyLock; @@ -31,5 +31,5 @@ pub struct ValuesImplementation; single_mapping!( ValuesImplementation, VALUES_PATTERN, - PhysicalOption::new(PlanImpl::Values, SortOption::None) + PhysicalOption::new(PlanImpl::Values, SortOption::None, RowCountOption::Unknown) ); diff --git a/src/optimizer/rule/implementation/dql/window.rs b/src/optimizer/rule/implementation/dql/window.rs index ab234f8d..1e3186b1 100644 --- a/src/optimizer/rule/implementation/dql/window.rs +++ b/src/optimizer/rule/implementation/dql/window.rs @@ -16,7 +16,8 @@ use crate::errors::DatabaseError; use crate::optimizer::core::pattern::Pattern; use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption}; +use crate::planner::LogicalPlan; use std::sync::LazyLock; static WINDOW_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -35,15 +36,16 @@ impl MatchPattern for WindowImplementation { impl ImplementationRule for WindowImplementation { fn update_best_option( &self, - op: &Operator, + plan: &LogicalPlan, _: &crate::planner::PlanArena, _: &StatisticMetaLoader<'_>, + _: &mut crate::optimizer::core::rule::HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { - if let Operator::Window(op) = op { + if let Operator::Window(op) = &plan.operator { crate::optimizer::core::rule::keep_best_physical_option( best_physical_option, - PhysicalOption::new(PlanImpl::Window, op.sort_option()), + PhysicalOption::new(PlanImpl::Window, op.sort_option(), RowCountOption::Follow), None, ); } diff --git a/src/optimizer/rule/implementation/macros.rs b/src/optimizer/rule/implementation/macros.rs index 15014862..e0835533 100644 --- a/src/optimizer/rule/implementation/macros.rs +++ b/src/optimizer/rule/implementation/macros.rs @@ -24,9 +24,10 @@ macro_rules! single_mapping { impl ImplementationRule for $ty { fn update_best_option( &self, - _: &Operator, + _: &$crate::planner::LogicalPlan, _: &$crate::planner::PlanArena, _: &StatisticMetaLoader<'_>, + _: &mut $crate::optimizer::core::rule::HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { //TODO: CostModel diff --git a/src/optimizer/rule/implementation/mod.rs b/src/optimizer/rule/implementation/mod.rs index c1a5879f..6e4e5501 100644 --- a/src/optimizer/rule/implementation/mod.rs +++ b/src/optimizer/rule/implementation/mod.rs @@ -19,7 +19,9 @@ pub(crate) mod macros; use crate::errors::DatabaseError; use crate::optimizer::core::pattern::Pattern; -use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; +use crate::optimizer::core::rule::{ + BestPhysicalOption, HintState, ImplementationRule, MatchPattern, +}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::rule::implementation::ddl::add_column::AddColumnImplementation; use crate::optimizer::rule::implementation::ddl::change_column::ChangeColumnImplementation; @@ -55,6 +57,7 @@ use crate::optimizer::rule::implementation::dql::top_k::TopKImplementation; use crate::optimizer::rule::implementation::dql::values::ValuesImplementation; use crate::optimizer::rule::implementation::dql::window::WindowImplementation; use crate::planner::operator::Operator; +use crate::planner::{LogicalPlan, PlanArena}; #[repr(usize)] #[derive(Debug, Copy, Clone, Eq, PartialEq)] @@ -101,6 +104,8 @@ impl ImplementationRuleRootTag { Operator::Join(_) => Some(Self::Join), Operator::Limit(_) => Some(Self::Limit), Operator::MarkApply(_) => Some(Self::MarkApply), + // The join planner sets its physical option when building it. + Operator::IndexProbeJoin(_) => None, Operator::Project(_) => Some(Self::Project), Operator::ScalarQueryInit(_) => Some(Self::ScalarQueryInit), Operator::ScalarSubquery(_) => Some(Self::ScalarSubquery), @@ -258,14 +263,21 @@ impl ImplementationRuleImpl { impl ImplementationRule for ImplementationRuleImpl { fn update_best_option( &self, - operator: &Operator, - arena: &crate::planner::PlanArena, + plan: &LogicalPlan, + arena: &PlanArena, loader: &StatisticMetaLoader<'_>, + state: &mut HintState, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { macro_rules! update { ($implementation:expr) => { - $implementation.update_best_option(operator, arena, loader, best_physical_option)? + $implementation.update_best_option( + plan, + arena, + loader, + state, + best_physical_option, + )? }; } @@ -324,7 +336,9 @@ mod tests { }; use crate::planner::operator::function_scan::FunctionScanOperator; use crate::planner::operator::top_k::TopKOperator; - use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; + use crate::planner::operator::{ + Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption, + }; use crate::planner::{Childrens, LogicalPlan, PlanArena, TableArenaCell}; use crate::storage::StatisticsMetaCache; use crate::types::value::DataValue; @@ -377,7 +391,14 @@ mod tests { let statistics_cache = StatisticsMetaCache::default(); let loader = StatisticMetaLoader::new(&statistics_cache); let mut best = None; - rule.update_best_option(operator, arena, &loader, &mut best)?; + let plan = LogicalPlan::new(operator.clone(), Childrens::None); + rule.update_best_option( + &plan, + arena, + &loader, + &mut crate::optimizer::core::rule::HintState::default(), + &mut best, + )?; Ok(best .expect("implementation rule should produce an option") @@ -418,11 +439,15 @@ mod tests { |op| matches!(op, Operator::Dummy), PlanImpl::Dummy, )?; - assert_sql_rule( + with_operator_from_sql( "select * from t1 where c1 = 1", - ImplementationRuleImpl::Filter, |op| matches!(op, Operator::Filter(_)), - PlanImpl::Filter, + |filter, arena| { + let option = best_option(ImplementationRuleImpl::Filter, &filter, arena)?; + assert_eq!(option.plan, PlanImpl::Filter); + assert_eq!(option.row_count, RowCountOption::Follow); + Ok(()) + }, )?; assert_sql_rule( "select * from t1 limit 1", @@ -430,11 +455,15 @@ mod tests { |op| matches!(op, Operator::Limit(_)), PlanImpl::Limit, )?; - assert_sql_rule( + with_operator_from_sql( "select * from t1 where exists(select * from t2)", - ImplementationRuleImpl::MarkApply, |op| matches!(op, Operator::MarkApply(_)), - PlanImpl::MarkApply, + |operator, arena| { + let option = best_option(ImplementationRuleImpl::MarkApply, &operator, arena)?; + assert_eq!(option.plan, PlanImpl::MarkApply); + assert_eq!(option.row_count, RowCountOption::Follow); + Ok(()) + }, )?; assert_sql_rule( "select c1 from t1", diff --git a/src/optimizer/rule/normalization/column_pruning.rs b/src/optimizer/rule/normalization/column_pruning.rs index 85fb9fa8..a932b571 100644 --- a/src/optimizer/rule/normalization/column_pruning.rs +++ b/src/optimizer/rule/normalization/column_pruning.rs @@ -16,14 +16,17 @@ use crate::catalog::ColumnRef; use crate::errors::DatabaseError; use crate::expression::agg::AggKind; use crate::expression::visitor::ExprVisitor; -use crate::expression::{AliasType, HasCountStar, ScalarExpression}; +use crate::expression::{HasCountStar, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::rule::normalization::{ remap_expr_positions, remap_exprs_positions, PositionRemapper, }; use crate::planner::operator::join::JoinCondition; -use crate::planner::operator::visitor::{OperatorExprVisitor, OperatorVisitor}; -use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; +use crate::planner::operator::visitor::{ChildrenVisit, OperatorExprVisitor, OperatorVisitor}; +use crate::planner::operator::visitor_mut::{ + ChildrenVisitMut, OperatorExprVisitorMut, OperatorVisitorMut, +}; use crate::planner::operator::Operator; use crate::planner::MetaArena; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; @@ -95,19 +98,11 @@ impl ExprVisitor for ReferencedColumnCollector<'_, '_> { fn visit_column_ref( &mut self, column: &crate::catalog::ColumnRef, + _: usize, ) -> Result<(), DatabaseError> { self.referenced_columns.insert(*column, self.arena); Ok(()) } - - fn visit_alias( - &mut self, - expr: ExprRef, - _ty: &AliasType, - arena: &(dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { - self.visit(expr, arena) - } } impl ApplyOutcome { @@ -130,8 +125,11 @@ impl ColumnPruning { referenced_columns, arena, }; - OperatorExprVisitor::new(&mut collector, arena as &dyn MetaArena) - .visit_operator(operator)?; + OperatorExprVisitor::new(&mut collector, arena as &dyn MetaArena).visit_operator( + operator, + None, + &mut ChildrenVisit::empty(), + )?; struct ReferencedOperatorColumnCollector<'a, 'p> { referenced_columns: &'a mut ReferencedColumns, @@ -142,14 +140,10 @@ impl ColumnPruning { fn visit_mark_apply( &mut self, op: &'a crate::planner::operator::mark_apply::MarkApplyOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { - if !matches!( - op.kind, - crate::planner::operator::mark_apply::MarkApplyKind::InnerJoin - ) { - self.referenced_columns - .insert(*op.output_column(), self.arena); - } + self.referenced_columns + .insert(*op.output_column(), self.arena); Ok(()) } @@ -174,6 +168,7 @@ impl ColumnPruning { fn visit_union( &mut self, op: &'a crate::planner::operator::union::UnionOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { self.referenced_columns.extend( op.left_schema_ref @@ -188,6 +183,7 @@ impl ColumnPruning { fn visit_set_membership( &mut self, op: &'a crate::planner::operator::set_membership::SetMembershipOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { self.referenced_columns.extend( op.left_schema_ref @@ -202,6 +198,7 @@ impl ColumnPruning { fn visit_delete( &mut self, op: &'a crate::planner::operator::delete::DeleteOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { self.referenced_columns .extend(op.primary_keys.iter().copied(), self.arena); @@ -213,7 +210,7 @@ impl ColumnPruning { referenced_columns, arena, } - .visit_operator(operator) + .visit_operator(operator, None, &mut ChildrenVisit::empty()) } fn extend_expr_referenced_columns<'a>( @@ -270,19 +267,7 @@ impl ColumnPruning { &mut PositionRemapper::new(removed_positions, remapped_exprs), arena, ) - .visit_operator(operator, None) - } - - fn remap_exprs_after_child_change<'a>( - exprs: impl IntoIterator, - removed_positions: &[usize], - remapped_exprs: &mut HashSet, - arena: &mut PlanArena<'_>, - ) -> Result<(), DatabaseError> { - if removed_positions.is_empty() { - return Ok(()); - } - remap_exprs_positions(exprs, removed_positions, remapped_exprs, arena) + .visit_operator(operator, None, &mut ChildrenVisitMut::empty()) } fn apply_only_child( @@ -520,6 +505,7 @@ impl ColumnPruning { | Operator::Limit(_) | Operator::ScalarQueryInit(_) | Operator::MarkApply(_) + | Operator::IndexProbeJoin(_) | Operator::ScalarSubquery(_) | Operator::Join(_) | Operator::Filter(_) @@ -557,7 +543,11 @@ impl ColumnPruning { )?; changed |= outcome.changed; outcome.removed_positions.truncate(main_removed_end); - } else if matches!(operator, Operator::MarkApply(_)) { + } else if matches!( + operator, + Operator::MarkApply(_) | Operator::IndexProbeJoin(_) + ) { + // Both inputs stay intact, so probe/predicate positions need no remapping. Self::extend_operator_referenced_columns( operator, &mut required_columns, @@ -663,21 +653,6 @@ impl ColumnPruning { ); } } - } else if let Operator::MarkApply(op) = operator { - Self::offset_removed_positions( - &mut outcome.removed_positions - [right_removed_start..right_removed_end], - old_left_outputs_len, - ); - let removed_positions = - &outcome.removed_positions[left_removed_start..right_removed_end]; - Self::remap_exprs_after_child_change( - op.predicates_mut().iter_mut(), - removed_positions, - &mut outcome.remapped_exprs, - arena, - )?; - outcome.removed_positions.truncate(right_removed_start); } else { Self::offset_removed_positions( &mut outcome.removed_positions @@ -835,7 +810,8 @@ impl NormalizationRule for ColumnPruning { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let outcome = Self::_apply( ReferencedColumns::with_arena_capacity(arena), @@ -868,9 +844,9 @@ mod tests { HepOptimizerPipeline::builder() .before_batch( - format!("column_pruning::{sql}"), + "column_pruning", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::ColumnPruning], + &[NormalizationRuleImpl::ColumnPruning], ) .build() .instantiate(plan) @@ -1020,9 +996,9 @@ mod tests { let pipeline = HepOptimizerPipeline::builder() .before_batch( - "test_column_pruning".to_string(), + "test_column_pruning", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::ColumnPruning], + &[NormalizationRuleImpl::ColumnPruning], ) .build(); let best_plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; diff --git a/src/optimizer/rule/normalization/combine_operators.rs b/src/optimizer/rule/normalization/combine_operators.rs index b45b8903..12a91bd3 100644 --- a/src/optimizer/rule/normalization/combine_operators.rs +++ b/src/optimizer/rule/normalization/combine_operators.rs @@ -13,209 +13,16 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; use crate::expression::{BinaryOperator, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::plan_utils::{only_child_mut, replace_with_only_child}; -use crate::optimizer::rule::normalization::strip_alias; use crate::planner::operator::filter::FilterOperator; -use crate::planner::operator::project::ProjectOperator; use crate::planner::operator::Operator; -use crate::planner::MetaArena; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; use crate::types::LogicalType; use std::mem; -fn is_passthrough_project(op: &ProjectOperator, arena: &PlanArena<'_>) -> bool { - op.exprs.iter().all(|expr| { - matches!( - arena.expression(strip_alias(*expr, arena)), - ScalarExpression::ColumnRef { .. } - ) - }) -} - -fn passthrough_source_position(expr: ExprRef, arena: &PlanArena<'_>) -> Option { - match arena.expression(strip_alias(expr, arena)) { - ScalarExpression::ColumnRef { position, .. } => Some(*position), - _ => None, - } -} - -fn collapse_match_expr(expr: ExprRef, arena: &PlanArena<'_>) -> ExprRef { - match arena.expression(expr) { - ScalarExpression::Alias { expr, .. } => *expr, - _ => expr, - } -} - -fn rewrite_column_position( - mut expr: ExprRef, - new_position: usize, - arena: &mut PlanArena<'_>, -) -> Result<(), DatabaseError> { - struct Rewriter(usize); - impl ExprVisitorMut for Rewriter { - fn visit_column_ref( - &mut self, - _column: &mut crate::catalog::ColumnRef, - position: &mut usize, - _arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { - *position = self.0; - Ok(()) - } - } - Rewriter(new_position).visit(&mut expr, arena) -} - -fn remap_passthrough_project_exprs( - parent_exprs: &mut [ExprRef], - child_exprs: &[ExprRef], - remapped_positions: &mut Vec, - arena: &mut PlanArena<'_>, -) -> Result { - remapped_positions.clear(); - - for parent_expr in parent_exprs.iter() { - let parent_match_expr = collapse_match_expr(*parent_expr, arena); - let Some(position) = child_exprs - .iter() - .find(|child_expr| parent_match_expr.eq_ignore_colref_pos(**child_expr, arena)) - .and_then(|expr| passthrough_source_position(*expr, arena)) - else { - return Ok(false); - }; - remapped_positions.push(position); - } - - for (parent_expr, position) in parent_exprs - .iter_mut() - .zip(remapped_positions.iter().copied()) - { - rewrite_column_position(*parent_expr, position, arena)?; - } - - Ok(true) -} - -fn groupby_exprs_match( - parent_exprs: &[ExprRef], - child_exprs: &[ExprRef], - arena: &crate::planner::PlanArena, -) -> bool { - parent_exprs.len() == child_exprs.len() - && parent_exprs - .iter() - .zip(child_exprs.iter()) - .all(|(parent_expr, child_expr)| parent_expr.eq_ignore_colref_pos(*child_expr, arena)) -} - -/// Combine two adjacent project operators into one. -pub struct CollapseProject; - -impl NormalizationRule for CollapseProject { - fn apply( - &self, - plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, - ) -> Result { - let Operator::Project(parent_op) = &mut plan.operator else { - return Ok(false); - }; - - let mut removed = false; - let mut remapped_positions = Vec::with_capacity(parent_op.exprs.len()); - while let Childrens::Only(child) = plan.childrens.as_mut() { - match &child.operator { - Operator::Project(child_op) - if is_passthrough_project(child_op, arena) - && remap_passthrough_project_exprs( - &mut parent_op.exprs, - &child_op.exprs, - &mut remapped_positions, - arena, - )? => - { - removed |= replace_with_only_child(child.as_mut()); - } - _ => break, - } - } - - Ok(removed) - } -} - -// Removing a project must preserve every input slot; renaming is allowed. -fn can_remove_project( - op: &ProjectOperator, - childrens: &mut Childrens, - arena: &mut PlanArena<'_>, -) -> bool { - let Childrens::Only(input) = childrens else { - return false; - }; - let schema = input.output_schema(arena); - if op.exprs.len() != schema.len() { - return false; - } - for (i, expr) in op.exprs.iter().enumerate() { - let mut source_expr = *expr; - while let ScalarExpression::Alias { expr, .. } = arena.expression(source_expr) { - source_expr = *expr; - } - let ScalarExpression::ColumnRef { - column: source_column, - .. - } = arena.expression(source_expr) - else { - return false; - }; - let ScalarExpression::ColumnRef { - position: execution_position, - .. - } = arena.expression(strip_alias(*expr, arena)) - else { - return false; - }; - if *source_column != schema[i] || *execution_position != i { - return false; - } - } - true -} - -struct ProjectionSubstitution<'a>(&'a [ExprRef]); - -impl ExprVisitorMut for ProjectionSubstitution<'_> { - fn visit( - &mut self, - expr: &mut ExprRef, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { - if let ScalarExpression::ColumnRef { position, .. } = arena.expression(*expr) { - let Some(source) = self.0.get(*position) else { - return Err(DatabaseError::InvalidValue( - "invalid projection slot".into(), - )); - }; - *expr = source.clone_expression(arena)?; - return Ok(()); - } - walk_mut_expr(self, expr, arena) - } - - fn visit_alias( - &mut self, - expr: &mut ExprRef, - _: &mut crate::expression::AliasType, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { - self.visit(expr, arena) - } -} - /// Combine two adjacent filter operators into one. pub struct CombineFilter; @@ -223,7 +30,8 @@ impl NormalizationRule for CombineFilter { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let parent_filter = match mem::replace(&mut plan.operator, Operator::Dummy) { Operator::Filter(op) => op, @@ -232,8 +40,6 @@ impl NormalizationRule for CombineFilter { return Ok(false); } }; - let mut parent_filter = parent_filter; - let cursor = match only_child_mut(plan) { Some(child) => child, None => { @@ -242,60 +48,54 @@ impl NormalizationRule for CombineFilter { } }; - let mut changed = false; - loop { - match &mut cursor.operator { - Operator::Filter(child_op) => { - let FilterOperator { - predicate, - having, - is_optimized: _, - } = parent_filter; - let child_predicate = child_op.predicate; - child_op.predicate = arena.alloc_expression(ScalarExpression::Binary { - op: BinaryOperator::And, - left_expr: predicate, - right_expr: child_predicate, - evaluator: None, - ty: LogicalType::Boolean, - }); - child_op.having = having || child_op.having; - child_op.is_optimized = false; - - return Ok(replace_with_only_child(plan)); - } - Operator::Project(project_op) => { - if !can_remove_project(project_op, &mut cursor.childrens, arena) { - plan.operator = Operator::Filter(parent_filter); - return Ok(changed); - } - let mut predicate = parent_filter.predicate.clone_expression(arena)?; - ProjectionSubstitution(&project_op.exprs).visit(&mut predicate, arena)?; - parent_filter.predicate = predicate; - parent_filter.is_optimized = false; - if replace_with_only_child(cursor) { - changed = true; - continue; - } - plan.operator = Operator::Filter(parent_filter); - return Ok(false); - } - _ => { - plan.operator = Operator::Filter(parent_filter); - return Ok(changed); - } + match &mut cursor.operator { + Operator::Filter(child_op) => { + let FilterOperator { + predicate, + having, + is_optimized: _, + } = parent_filter; + let child_predicate = child_op.predicate; + child_op.predicate = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::And, + left_expr: predicate, + right_expr: child_predicate, + evaluator: None, + ty: LogicalType::Boolean, + }); + child_op.having = having || child_op.having; + child_op.is_optimized = false; + + Ok(replace_with_only_child(plan)) + } + _ => { + plan.operator = Operator::Filter(parent_filter); + Ok(false) } } } } +fn groupby_exprs_match( + parent_exprs: &[ExprRef], + child_exprs: &[ExprRef], + arena: &PlanArena<'_>, +) -> bool { + parent_exprs.len() == child_exprs.len() + && parent_exprs + .iter() + .zip(child_exprs) + .all(|(parent, child)| parent.eq_ignore_colref_pos(*child, arena)) +} + pub struct CollapseGroupByAgg; impl NormalizationRule for CollapseGroupByAgg { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let can_collapse = { let LogicalPlan { @@ -337,12 +137,9 @@ mod tests { use crate::optimizer::core::rule::NormalizationRule; use crate::optimizer::heuristic::batch::HepBatchStrategy; use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline; - use crate::optimizer::rule::normalization::combine_operators::{ - CollapseGroupByAgg, CollapseProject, - }; + use crate::optimizer::rule::normalization::combine_operators::CollapseGroupByAgg; use crate::optimizer::rule::normalization::NormalizationRuleImpl; use crate::planner::operator::aggregate::AggregateOperator; - use crate::planner::operator::project::ProjectOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; @@ -351,113 +148,6 @@ mod tests { arena.alloc_expression(ScalarExpression::column_expr(column, position)) } - #[test] - fn test_collapse_project() -> Result<(), DatabaseError> { - let table_state = build_t1_table()?; - let mut arena = PlanArena::new(&table_state.table_arena); - let plan = - table_state.plan_with_arena("select c1 from (select c1, c2 from t1) t", &mut arena)?; - - let pipeline = HepOptimizerPipeline::builder() - .before_batch( - "test_collapse_project".to_string(), - HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::CollapseProject], - ) - .build(); - let best_plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; - - if let Operator::Project(op) = &best_plan.operator { - assert_eq!(op.exprs.len(), 1); - } else { - unreachable!("Should be a project operator") - } - - let alias_project = best_plan.childrens.pop_only(); - assert!( - matches!(alias_project.operator, Operator::Project(_)), - "Derived-table alias projection should be preserved" - ); - let scan_op = alias_project.childrens.pop_only(); - assert!(matches!(scan_op.operator, Operator::TableScan(_))); - assert!(matches!(scan_op.childrens.as_ref(), Childrens::None)); - - Ok(()) - } - - #[test] - fn test_collapse_project_with_alias() -> Result<(), DatabaseError> { - let table_state = build_t1_table()?; - let mut arena = PlanArena::new(&table_state.table_arena); - let plan = table_state - .plan_with_arena("select t.x from (select c1 as x from t1) t", &mut arena)?; - let original = plan.clone(); - let original_child = original.childrens.pop_only(); - assert!(matches!(original_child.operator, Operator::Project(_))); - let original_grandchild = original_child.childrens.pop_only(); - assert!(matches!(original_grandchild.operator, Operator::Project(_))); - - let pipeline = HepOptimizerPipeline::builder() - .before_batch( - "test_collapse_project_with_alias".to_string(), - HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::CollapseProject], - ) - .build(); - let best_plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; - if let Operator::Project(op) = &best_plan.operator { - assert_eq!(op.exprs.len(), 1); - } else { - unreachable!("Should be a project operator"); - } - - let scan_op = best_plan.childrens.pop_only(); - assert!(matches!(scan_op.operator, Operator::Project(_))); - let scan_child = scan_op.childrens.pop_only(); - assert!( - !matches!(scan_child.operator, Operator::Project(_)), - "Child project should be collapsed" - ); - - Ok(()) - } - - #[test] - fn test_collapse_project_remaps_reordered_passthrough_positions() -> Result<(), DatabaseError> { - let table_arena = crate::planner::TableArenaCell::default(); - let mut arena = PlanArena::new(&table_arena); - let child = LogicalPlan::new( - Operator::Project(ProjectOperator { - exprs: vec![ - column_expr(&mut arena, "c2", 1), - column_expr(&mut arena, "c1", 0), - ], - }), - Childrens::Only(Box::new(LogicalPlan::new(Operator::Dummy, Childrens::None))), - ); - let mut plan = LogicalPlan::new( - Operator::Project(ProjectOperator { - exprs: vec![column_expr(&mut arena, "c2", 0)], - }), - Childrens::Only(Box::new(child)), - ); - - assert!(CollapseProject.apply(&mut plan, &mut arena)?); - - let Operator::Project(op) = &plan.operator else { - unreachable!("expected project"); - }; - let ScalarExpression::ColumnRef { position, .. } = arena.expression(op.exprs[0]) else { - unreachable!("expected column ref"); - }; - assert_eq!(*position, 1); - assert!(matches!( - plan.childrens.pop_only().operator, - Operator::Dummy - )); - Ok(()) - } - #[test] fn test_combine_filter_keeps_reordered_projection() -> Result<(), DatabaseError> { let table_state = build_t1_table()?; @@ -469,7 +159,7 @@ mod tests { project.exprs.reverse(); } } - assert!(!super::CombineFilter.apply(&mut filter, &mut arena)?); + assert!(!super::CombineFilter.apply(&mut filter, &mut arena, None)?); assert!(matches!( filter.childrens.pop_only().operator, Operator::Project(_) @@ -486,11 +176,17 @@ mod tests { &mut arena, )?; + // Projects are eliminated before filters are combined. let pipeline = HepOptimizerPipeline::builder() .before_batch( - "test_combine_filter".to_string(), + "Collapse Project", + HepBatchStrategy::fix_point_topdown(10), + &[NormalizationRuleImpl::EliminateProject], + ) + .before_batch( + "test_combine_filter", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::CombineFilter], + &[NormalizationRuleImpl::CombineFilter], ) .build(); let best_plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; @@ -518,9 +214,9 @@ mod tests { let pipeline = HepOptimizerPipeline::builder() .before_batch( - "test_collapse_group_by_agg".to_string(), + "test_collapse_group_by_agg", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::CollapseGroupByAgg], + &[NormalizationRuleImpl::CollapseGroupByAgg], ) .build(); @@ -552,7 +248,7 @@ mod tests { let expr = column_expr(&mut arena, "c2", 0); let mut plan = AggregateOperator::build(child, vec![], vec![expr], true, false); - assert!(CollapseGroupByAgg.apply(&mut plan, &mut arena)?); + assert!(CollapseGroupByAgg.apply(&mut plan, &mut arena, None)?); let Operator::Aggregate(op) = &plan.operator else { unreachable!("expected aggregate"); }; diff --git a/src/optimizer/rule/normalization/compilation_in_advance.rs b/src/optimizer/rule/normalization/compilation_in_advance.rs index e26a2111..4cdcda2f 100644 --- a/src/optimizer/rule/normalization/compilation_in_advance.rs +++ b/src/optimizer/rule/normalization/compilation_in_advance.rs @@ -15,7 +15,10 @@ use crate::errors::DatabaseError; use crate::expression::BindEvaluator; use crate::optimizer::core::rule::NormalizationRule; -use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::planner::operator::visitor_mut::{ + ChildrenVisitMut, OperatorExprVisitorMut, OperatorVisitorMut, +}; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan, PlanArena}; @@ -27,7 +30,11 @@ pub(crate) fn evaluator_bind_current( arena: &mut PlanArena, ) -> Result<(), DatabaseError> { let mut evaluator = BindEvaluator; - OperatorExprVisitorMut::new(&mut evaluator, arena).visit_operator(&mut plan.operator, None) + OperatorExprVisitorMut::new(&mut evaluator, arena).visit_operator( + &mut plan.operator, + None, + &mut ChildrenVisitMut::empty(), + ) } impl EvaluatorBind { @@ -40,6 +47,7 @@ impl EvaluatorBind { plan.operator, Operator::ScalarQueryInit(_) | Operator::MarkApply(_) + | Operator::IndexProbeJoin(_) | Operator::Join(_) | Operator::Union(_) | Operator::SetMembership(_) @@ -60,7 +68,8 @@ impl NormalizationRule for EvaluatorBind { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { Self::_apply(plan, arena)?; Ok(true) @@ -222,7 +231,7 @@ mod tests { right: Box::new(right_filter), }, ); - assert!(EvaluatorBind.apply(&mut join, &mut arena)?); + assert!(EvaluatorBind.apply(&mut join, &mut arena, None)?); let Operator::Join(op) = &join.operator else { unreachable!() }; @@ -243,7 +252,7 @@ mod tests { let union_left = filter(&mut arena, column); let union_right = filter(&mut arena, column); let mut union = UnionOperator::build(vec![column], vec![column], union_left, union_right); - EvaluatorBind.apply(&mut union, &mut arena)?; + EvaluatorBind.apply(&mut union, &mut arena, None)?; assert!(union.childrens.iter().all(|child| { matches!(&child.operator, Operator::Filter(op) if is_bound(op.predicate, &arena)) })); diff --git a/src/optimizer/rule/normalization/constant_propagation.rs b/src/optimizer/rule/normalization/constant_propagation.rs new file mode 100644 index 00000000..444afdb0 --- /dev/null +++ b/src/optimizer/rule/normalization/constant_propagation.rs @@ -0,0 +1,622 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Constant propagation over equality predicates. +//! +//! `a = b AND a = 1` implies `b = 1`. The derived predicate is added (never replaces anything), +//! so later rules can turn it into a static index range on the other column. +//! +//! - [`PropagateFilterConstants`] derives constants inside one `Filter`'s top-level conjuncts. +//! - [`PropagateJoinConstants`] derives constants across a join's equi-keys and wraps the +//! target input with a new `Filter`. Only the sides whose rows may be filtered without +//! changing the join result are targeted: both sides of an inner join, the right side of a +//! left join and the left side of a right join. Constants come from filters below the join +//! and, for outer joins, also from the preserved side's ON conjuncts that stay on the join. +//! +//! Only plain table columns of identical types participate, and NULL constants are ignored: +//! `col = NULL` is never true, so it carries no value to propagate. + +use crate::catalog::{ColumnRef, ColumnRelation}; +use crate::errors::DatabaseError; +use crate::expression::{BinaryOperator, ScalarExpression}; +use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::optimizer::plan_utils::wrap_child_with; +use crate::planner::operator::filter::FilterOperator; +use crate::planner::operator::join::{JoinCondition, JoinType}; +use crate::planner::operator::Operator; +use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; +use crate::types::LogicalType; + +#[derive(Clone, Copy)] +struct Fact { + column: ColumnRef, + position: usize, + constant: ExprRef, + left: bool, +} + +fn lookup(facts: &[Fact], arena: &PlanArena<'_>, column: ColumnRef, left: bool) -> Option { + facts + .iter() + .find(|fact| fact.left == left && arena.same_column(fact.column, column)) + .map(|fact| fact.constant) +} + +fn for_each_conjunct(expr: ExprRef, arena: &PlanArena<'_>, f: &mut impl FnMut(ExprRef)) { + match arena.expression(expr) { + ScalarExpression::Binary { + op: BinaryOperator::And, + left_expr, + right_expr, + .. + } => { + let (left_expr, right_expr) = (*left_expr, *right_expr); + for_each_conjunct(left_expr, arena, f); + for_each_conjunct(right_expr, arena, f); + } + _ => f(expr), + } +} + +fn table_column(expr: ExprRef, arena: &PlanArena<'_>) -> Option<(ColumnRef, usize)> { + match expr.unpack_alias_ref(arena) { + ScalarExpression::ColumnRef { column, position } + if matches!( + arena.column(*column).summary().relation, + ColumnRelation::Table { .. } + ) => + { + Some((*column, *position)) + } + _ => None, + } +} + +fn eq_operands(expr: ExprRef, arena: &PlanArena<'_>) -> Option<(ExprRef, ExprRef)> { + match expr.unpack_alias_ref(arena) { + ScalarExpression::Binary { + op: BinaryOperator::Eq, + left_expr, + right_expr, + .. + } => Some((*left_expr, *right_expr)), + _ => None, + } +} + +fn as_fact( + expr: ExprRef, + arena: &PlanArena<'_>, + left: bool, + has_col_connection: &mut bool, +) -> Option { + let (lhs, rhs) = eq_operands(expr, arena)?; + let ((column, position), constant) = match (table_column(lhs, arena), table_column(rhs, arena)) + { + (Some(column), None) => (column, rhs.unpack_alias(arena)), + (None, Some(column)) => (column, lhs.unpack_alias(arena)), + (Some(_), Some(_)) => { + *has_col_connection = true; + return None; + } + (None, None) => return None, + }; + matches!(arena.expression(constant), ScalarExpression::Constant(value) if !value.is_null()) + .then_some(Fact { + column, + position, + constant, + left, + }) +} + +type ColumnSlot = (ColumnRef, usize); + +fn as_column_pair( + lhs: ExprRef, + rhs: ExprRef, + arena: &PlanArena<'_>, +) -> Option<(ColumnSlot, ColumnSlot)> { + let lhs = table_column(lhs, arena)?; + let rhs = table_column(rhs, arena)?; + (arena.column(lhs.0).datatype() == arena.column(rhs.0).datatype() + && !arena.same_column(lhs.0, rhs.0)) + .then_some((lhs, rhs)) +} + +fn and_fact( + acc: Option, + fact: Fact, + arena: &mut PlanArena<'_>, +) -> Result { + let left_expr = arena.alloc_expression(ScalarExpression::ColumnRef { + column: fact.column, + position: fact.position, + }); + let right_expr = fact.constant.clone_expression(arena)?; + let eq = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Eq, + left_expr, + right_expr, + evaluator: None, + ty: LogicalType::Boolean, + }); + Ok(match acc { + Some(left_expr) => arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::And, + left_expr, + right_expr: eq, + evaluator: None, + ty: LogicalType::Boolean, + }), + None => eq, + }) +} + +pub struct PropagateFilterConstants; + +impl NormalizationRule for PropagateFilterConstants { + fn apply( + &self, + plan: &mut LogicalPlan, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, + ) -> Result { + let Operator::Filter(filter) = &plan.operator else { + return Ok(false); + }; + let predicate = filter.predicate; + + let mut facts = Vec::new(); + let mut has_col_connection = false; + for_each_conjunct(predicate, arena, &mut |conjunct| { + facts.extend(as_fact(conjunct, arena, false, &mut has_col_connection)); + }); + // Without both a constant and a column pair nothing can be derived. + if facts.is_empty() || !has_col_connection { + return Ok(false); + } + + let base = facts.len(); + loop { + let before = facts.len(); + for_each_conjunct(predicate, arena, &mut |conjunct| { + let Some((lhs, rhs)) = eq_operands(conjunct, arena) else { + return; + }; + let Some((lhs, rhs)) = as_column_pair(lhs, rhs, arena) else { + return; + }; + for (from, to) in [(lhs, rhs), (rhs, lhs)] { + if lookup(&facts, arena, to.0, false).is_some() { + continue; + } + if let Some(constant) = lookup(&facts, arena, from.0, false) { + facts.push(Fact { + column: to.0, + position: to.1, + constant, + left: false, + }); + } + } + }); + if facts.len() == before { + break; + } + } + if facts.len() == base { + return Ok(false); + } + + let mut acc = predicate; + for &fact in &facts[base..] { + acc = and_fact(Some(acc), fact, arena)?; + } + if let Operator::Filter(filter) = &mut plan.operator { + filter.predicate = acc; + } + Ok(true) + } +} + +fn collect_facts(plan: &LogicalPlan, arena: &PlanArena<'_>, left: bool, facts: &mut Vec) { + match (&plan.operator, plan.childrens.as_ref()) { + (Operator::Filter(filter), Childrens::Only(child)) => { + for_each_conjunct(filter.predicate, arena, &mut |conjunct| { + facts.extend(as_fact(conjunct, arena, left, &mut false)); + }); + collect_facts(child, arena, left, facts); + } + (Operator::Join(join), Childrens::Twins { left: l, right: r }) => match join.join_type { + JoinType::Inner | JoinType::Cross => { + collect_facts(l, arena, left, facts); + collect_facts(r, arena, left, facts); + } + JoinType::LeftOuter => collect_facts(l, arena, left, facts), + JoinType::RightOuter => collect_facts(r, arena, left, facts), + JoinType::Full => {} + }, + _ => {} + } +} + +pub struct PropagateJoinConstants; + +impl NormalizationRule for PropagateJoinConstants { + fn apply( + &self, + plan: &mut LogicalPlan, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, + ) -> Result { + let Operator::Join(join) = &plan.operator else { + return Ok(false); + }; + let (to_left, to_right) = match join.join_type { + JoinType::Inner => (true, true), + JoinType::LeftOuter => (false, true), + JoinType::RightOuter => (true, false), + JoinType::Full | JoinType::Cross => return Ok(false), + }; + let JoinCondition::On { on, filter } = &join.on else { + return Ok(false); + }; + if on.is_empty() { + return Ok(false); + } + let join_filter = *filter; + + let mut facts = Vec::new(); + if let (Some(join_filter), true) = (join_filter, to_left != to_right) { + let preserved_left = to_right; + let Childrens::Twins { left, right } = plan.childrens.as_mut() else { + return Ok(false); + }; + let preserved = if preserved_left { left } else { right }; + let preserved_columns = preserved.output_schema(arena); + for_each_conjunct(join_filter, arena, &mut |conjunct| { + facts.extend( + as_fact(conjunct, arena, preserved_left, &mut false) + .filter(|fact| preserved_columns.contains(&fact.column)), + ); + }); + } + let Childrens::Twins { left, right } = plan.childrens.as_ref() else { + return Ok(false); + }; + collect_facts(left, arena, true, &mut facts); + collect_facts(right, arena, false, &mut facts); + if facts.is_empty() { + return Ok(false); + } + + let base = facts.len(); + for &(left_key, right_key) in on { + let Some((l, r)) = as_column_pair(left_key, right_key, arena) else { + continue; + }; + for (allowed, from, to, to_left_side) in + [(to_right, l, r, false), (to_left, r, l, true)] + { + if !allowed || lookup(&facts, arena, to.0, to_left_side).is_some() { + continue; + } + if let Some(constant) = lookup(&facts, arena, from.0, !to_left_side) { + facts.push(Fact { + column: to.0, + position: to.1, + constant, + left: to_left_side, + }); + } + } + } + if facts.len() == base { + return Ok(false); + } + + let (mut left_acc, mut right_acc) = (None, None); + for &fact in &facts[base..] { + let acc = if fact.left { + &mut left_acc + } else { + &mut right_acc + }; + *acc = Some(and_fact(acc.take(), fact, arena)?); + } + let mut applied = false; + for (child_idx, acc) in [(0, left_acc), (1, right_acc)] { + if let Some(predicate) = acc { + applied |= wrap_child_with( + plan, + child_idx, + Operator::Filter(FilterOperator { + predicate, + is_optimized: false, + having: false, + }), + ); + } + } + Ok(applied) + } +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::{eq_operands, for_each_conjunct, table_column}; + use super::{PropagateFilterConstants, PropagateJoinConstants}; + use crate::binder::test::build_t1_table; + use crate::errors::DatabaseError; + use crate::expression::ScalarExpression; + use crate::optimizer::core::rule::NormalizationRule; + use crate::optimizer::heuristic::batch::HepBatchStrategy; + use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline; + use crate::optimizer::plan_utils::wrap_child_with; + use crate::optimizer::rule::normalization::NormalizationRuleImpl; + use crate::planner::operator::filter::FilterOperator; + use crate::planner::operator::Operator; + use crate::planner::{Childrens, LogicalPlan, PlanArena}; + + /// `column = constant` conjuncts of the filters stacked directly on `plan`, as + /// sorted `"column = value"` strings. + fn facts(plan: &LogicalPlan, arena: &PlanArena<'_>) -> Vec { + let mut out = Vec::new(); + let mut node = plan; + while let (Operator::Filter(filter), Childrens::Only(child)) = + (&node.operator, node.childrens.as_ref()) + { + for_each_conjunct(filter.predicate, arena, &mut |conjunct| { + let Some((lhs, rhs)) = eq_operands(conjunct, arena) else { + return; + }; + let (column, constant) = match (table_column(lhs, arena), table_column(rhs, arena)) + { + (Some((column, _)), None) => (column, rhs), + (None, Some((column, _))) => (column, lhs), + _ => return, + }; + if let ScalarExpression::Constant(value) = constant.unpack_alias_ref(arena) { + out.push(format!("{} = {}", arena.column(column).name(), value)); + } + }); + node = child; + } + out.sort(); + out + } + + fn find_filter(plan: &LogicalPlan) -> &LogicalPlan { + match (&plan.operator, plan.childrens.as_ref()) { + (Operator::Filter(_), _) => plan, + (_, Childrens::Only(child)) => find_filter(child), + _ => unreachable!("expected a filter"), + } + } + + fn find_join(plan: &LogicalPlan) -> (&LogicalPlan, &LogicalPlan) { + match (&plan.operator, plan.childrens.as_ref()) { + (Operator::Join(_), Childrens::Twins { left, right }) => (left, right), + (_, Childrens::Only(child)) => find_join(child), + _ => unreachable!("expected a join"), + } + } + + fn find_join_mut(plan: &mut LogicalPlan) -> &mut LogicalPlan { + if matches!(plan.operator, Operator::Join(_)) { + return plan; + } + match plan.childrens.as_mut() { + Childrens::Only(child) => find_join_mut(child), + _ => unreachable!("expected a join"), + } + } + + /// Runs the rules of the default "Predicate Pushdown" batch (same order as `db.rs`). + fn propagate_join(sql: &str) -> Result<(Vec, Vec), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let plan = table_state.plan_with_arena(sql, &mut arena)?; + let mut plan = HepOptimizerPipeline::builder() + .before_batch( + "test_propagate_join_constants", + HepBatchStrategy::fix_point_topdown(30), + &[ + NormalizationRuleImpl::PropagateFilterConstants, + NormalizationRuleImpl::PropagateJoinConstants, + NormalizationRuleImpl::PushPredicateThroughJoin, + NormalizationRuleImpl::PushJoinPredicateIntoScan, + ], + ) + .build() + .instantiate(plan) + .find_best(None, &mut arena)?; + + // Fix point reached: a second application derives nothing new. + assert!(!PropagateJoinConstants.apply(find_join_mut(&mut plan), &mut arena, None)?); + + let (left, right) = find_join(&plan); + Ok((facts(left, &arena), facts(right, &arena))) + } + + #[test] + fn filter_derives_constant_through_column_equality() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let mut plan = + table_state.plan_with_arena("select * from t1 where c1 = c2 and c2 = 1", &mut arena)?; + let filter = find_filter_mut(&mut plan); + + assert!(PropagateFilterConstants.apply(filter, &mut arena, None)?); + assert_eq!(facts(filter, &arena), ["c1 = 1", "c2 = 1"]); + // Idempotent: the derived fact is now known, so nothing is added again. + assert!(!PropagateFilterConstants.apply(filter, &mut arena, None)?); + Ok(()) + } + + #[test] + fn filter_closes_equality_chains() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let mut plan = table_state.plan_with_arena( + "select * from t1, t2 where c1 = c3 and c3 = c2 and c2 = 1", + &mut arena, + )?; + let filter = find_filter_mut(&mut plan); + + assert!(PropagateFilterConstants.apply(filter, &mut arena, None)?); + assert_eq!(facts(filter, &arena), ["c1 = 1", "c2 = 1", "c3 = 1"]); + Ok(()) + } + + #[test] + fn filter_ignores_null_and_non_equality() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + for sql in [ + "select * from t1 where c1 = c2 and c2 = null", + "select * from t1 where c1 > c2 and c2 = 1", + "select * from t1 where c1 = c2 + 1 and c2 = 1", + ] { + let mut arena = PlanArena::new(&table_state.table_arena); + let mut plan = table_state.plan_with_arena(sql, &mut arena)?; + let filter = find_filter_mut(&mut plan); + assert!( + !PropagateFilterConstants.apply(filter, &mut arena, None)?, + "{sql}" + ); + } + Ok(()) + } + + #[test] + fn inner_join_propagates_both_ways() -> Result<(), DatabaseError> { + assert_eq!( + propagate_join("select * from t1 join t2 on c1 = c3 where c1 = 1")?, + (vec!["c1 = 1".into()], vec!["c3 = 1".into()]) + ); + assert_eq!( + propagate_join("select * from t1 join t2 on c1 = c3 where c3 = 1")?, + (vec!["c1 = 1".into()], vec!["c3 = 1".into()]) + ); + Ok(()) + } + + #[test] + fn left_join_only_constrains_right_side() -> Result<(), DatabaseError> { + assert_eq!( + propagate_join("select * from t1 left join t2 on c1 = c3 where c1 = 1")?, + (vec!["c1 = 1".into()], vec!["c3 = 1".into()]) + ); + // The non-preserved side's ON constant is pushed down, never derived onto the left. + assert_eq!( + propagate_join("select * from t1 left join t2 on c1 = c3 and c3 = 1")?, + (vec![], vec!["c3 = 1".into()]) + ); + // The preserved side's ON constant stays on the join but still derives the right side. + assert_eq!( + propagate_join("select * from t1 left join t2 on c1 = c3 and c1 = 1")?, + (vec![], vec!["c3 = 1".into()]) + ); + Ok(()) + } + + #[test] + fn right_join_only_constrains_left_side() -> Result<(), DatabaseError> { + assert_eq!( + propagate_join("select * from t1 right join t2 on c1 = c3 where c3 = 1")?, + (vec!["c1 = 1".into()], vec!["c3 = 1".into()]) + ); + assert_eq!( + propagate_join("select * from t1 right join t2 on c1 = c3 and c1 = 1")?, + (vec!["c1 = 1".into()], vec![]) + ); + assert_eq!( + propagate_join("select * from t1 right join t2 on c1 = c3 and c3 = 1")?, + (vec!["c1 = 1".into()], vec![]) + ); + Ok(()) + } + + fn propagate_from_side( + join: &str, + fact_on_left: bool, + ) -> Result<(bool, Vec, Vec), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let mut plan = table_state + .plan_with_arena(format!("select * from t1 {join} t2 on c1 = c3"), &mut arena)?; + let (fact_sql, child_idx) = if fact_on_left { + ("select * from t1 where c1 = 1", 0) + } else { + ("select * from t2 where c3 = 1", 1) + }; + // Same table arena, so the fact's column has the input's summary and position. + let fact_plan = table_state.plan_with_arena(fact_sql, &mut arena)?; + let Operator::Filter(filter) = &find_filter(&fact_plan).operator else { + unreachable!() + }; + let predicate = filter.predicate; + + let join_plan = find_join_mut(&mut plan); + assert!(wrap_child_with( + join_plan, + child_idx, + Operator::Filter(FilterOperator { + predicate, + is_optimized: false, + having: false, + }), + )); + let applied = PropagateJoinConstants.apply(join_plan, &mut arena, None)?; + let (left, right) = find_join(&plan); + Ok((applied, facts(left, &arena), facts(right, &arena))) + } + + #[test] + fn join_type_limits_propagation_direction() -> Result<(), DatabaseError> { + let one = |column: &str| vec![format!("{column} = 1")]; + // (join, fact on left, expected (applied, left facts, right facts)) + let cases = [ + ("join", true, (true, one("c1"), one("c3"))), + ("join", false, (true, one("c1"), one("c3"))), + ("left join", true, (true, one("c1"), one("c3"))), + // Filtering the preserved left side would drop unmatched rows. + ("left join", false, (false, vec![], one("c3"))), + ("right join", true, (false, one("c1"), vec![])), + ("right join", false, (true, one("c1"), one("c3"))), + ("full join", true, (false, one("c1"), vec![])), + ("full join", false, (false, vec![], one("c3"))), + ]; + for (join, fact_on_left, expected) in cases { + assert_eq!( + propagate_from_side(join, fact_on_left)?, + expected, + "{join}, fact on left: {fact_on_left}" + ); + } + Ok(()) + } + + fn find_filter_mut(plan: &mut LogicalPlan) -> &mut LogicalPlan { + if matches!(plan.operator, Operator::Filter(_)) { + return plan; + } + match plan.childrens.as_mut() { + Childrens::Only(child) => find_filter_mut(child), + _ => unreachable!("expected a filter"), + } + } +} diff --git a/src/optimizer/rule/normalization/elimination.rs b/src/optimizer/rule/normalization/elimination.rs index 2e483454..901126df 100644 --- a/src/optimizer/rule/normalization/elimination.rs +++ b/src/optimizer/rule/normalization/elimination.rs @@ -13,23 +13,276 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::expression::ScalarExpression; +use crate::expression::visitor::ExprVisitor; +use crate::expression::visitor_mut::ExprVisitorMut; +use crate::expression::{AliasType, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; -use crate::optimizer::plan_utils::{only_child_mut, replace_with_only_child, wrap_child_with}; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::optimizer::plan_utils::{ + only_child_mut, predicate_scan_mut, replace_with_only_child, wrap_child_with, +}; +use crate::planner::operator::join::{JoinCondition, JoinType}; use crate::planner::operator::limit::LimitOperator; use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::operator::table_scan::TableScanOperator; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; -use crate::planner::{Childrens, ExprRef, LogicalPlan}; +use crate::planner::operator::visitor::{ChildrenVisit, OperatorExprVisitor, OperatorVisitor}; +use crate::planner::operator::visitor_mut::{ + ChildrenVisitMut, OperatorExprVisitorMut, OperatorVisitorMut, +}; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption}; +use crate::planner::{Childrens, ExprRef, LogicalPlan, MetaArena, PlanArena}; use crate::types::index::{IndexLookup, IndexOrderHint}; +pub struct EliminateProject; + +struct Substitute<'m> { + exprs: Option<&'m [ExprRef]>, + offset: usize, +} + +struct OpaqueCheck { + found: bool, +} + +impl ExprVisitor for OpaqueCheck { + fn visit_expression_ref(&mut self, expr: ExprRef, arena: &A) -> Result { + if self.found { + return Ok(false); + } + self.found = matches!( + arena.expression(expr), + ScalarExpression::ScalaFunction(_) | ScalarExpression::TableFunction(_) + ); + Ok(!self.found) + } +} + +fn contains_opaque( + exprs: impl Iterator, + arena: &A, +) -> Result { + let mut check = OpaqueCheck { found: false }; + for expr in exprs { + ExprVisitor::visit(&mut check, expr, arena)?; + if check.found { + break; + } + } + Ok(check.found) +} + +impl ExprVisitorMut for Substitute<'_> { + fn visit_expression_ref( + &mut self, + expr: &mut ExprRef, + arena: &mut A, + ) -> Result { + match arena.expression(*expr) { + ScalarExpression::ColumnRef { position, column } => { + let (position, column) = (*position, *column); + if let Some(replacement) = self.exprs.and_then(|exprs| exprs.get(position)) { + // Do not revisit or mutate the replacement's input expression. + *expr = *replacement; + } else if self.offset != 0 { + *expr = arena.alloc_expression(ScalarExpression::column_expr( + column, + position + self.offset, + )); + } + Ok(false) + } + ScalarExpression::Constant(_) + | ScalarExpression::OuterParam { .. } + | ScalarExpression::OuterValue { .. } + | ScalarExpression::InitValue { .. } + | ScalarExpression::Empty => Ok(false), + _ => { + // Reached only for a matched Project rewrite or a nonzero Cross offset. + // Copy before changing child refs: a retained operator may share this node + // but still require its original input positions. + let expression = arena.expression(*expr).clone(); + *expr = arena.alloc_expression(expression); + Ok(true) + } + } + } +} + +fn movable_project(plan: &LogicalPlan, arena: &PlanArena) -> Result { + let (Operator::Project(op), Childrens::Only(_)) = (&plan.operator, plan.childrens.as_ref()) + else { + return Ok(false); + }; + Ok(!contains_opaque(op.exprs.iter().copied(), arena)?) +} + +fn lift_cross_project( + plan: &mut LogicalPlan, + arena: &mut PlanArena, +) -> Result { + let Childrens::Twins { left, right } = plan.childrens.as_mut() else { + return Ok(false); + }; + let (input, other, from_right) = if movable_project(left, arena)? { + (left, right, false) + } else if movable_project(right, arena)? { + (right, left, true) + } else { + return Ok(false); + }; + let (Operator::Project(op), Childrens::Only(base)) = + (&mut input.operator, input.childrens.as_mut()) + else { + return Ok(false); + }; + let offset = if from_right { + other.output_schema(arena).len() + } else { + base.output_schema(arena).len() + }; + if from_right && offset != 0 { + let mut shift = Substitute { + exprs: None, + offset, + }; + for expr in &mut op.exprs { + ExprVisitorMut::visit(&mut shift, expr, arena)?; + } + } + let identity_start = if from_right { 0 } else { offset }; + let columns = other.output_schema(arena); + op.exprs + .extend(columns.iter().enumerate().map(|(index, column)| { + arena.alloc_expression(ScalarExpression::column_expr( + *column, + identity_start + index, + )) + })); + if from_right { + op.exprs.rotate_right(offset); + } + + let mut project = input.take(); + **input = std::mem::replace(&mut *project.childrens, Childrens::None).pop_only(); + let mut join = plan.take(); + + join.reset_output_schema_cache(); + + *project.childrens = Childrens::Only(Box::new(join)); + project.physical_option = None; + project.reset_output_schema_cache(); + + *plan = project; + Ok(true) +} + +impl NormalizationRule for EliminateProject { + fn apply( + &self, + plan: &mut LogicalPlan, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, + ) -> Result { + if matches!(&plan.operator, Operator::Join(op) + if op.join_type == JoinType::Cross && op.on == JoinCondition::None) + { + return lift_cross_project(plan, arena); + } + if !matches!( + plan.operator, + Operator::Project(_) + | Operator::Filter(_) + | Operator::Sort(_) + | Operator::Limit(_) + | Operator::TopK(_) + ) { + return Ok(false); + } + let Childrens::Only(child) = plan.childrens.as_ref() else { + return Ok(false); + }; + if !movable_project(child, arena)? { + return Ok(false); + } + let Operator::Project(project) = &child.operator else { + return Ok(false); + }; + let mut check = OpaqueCheck { found: false }; + OperatorExprVisitor::new(&mut check, arena).visit_operator( + &plan.operator, + None, + &mut ChildrenVisit::empty(), + )?; + if check.found { + return Ok(false); + } + let mut substitute = Substitute { + exprs: Some(&project.exprs), + offset: 0, + }; + let merge = matches!(plan.operator, Operator::Project(_)); + let schema = if let Operator::Project(op) = &plan.operator { + Some( + op.exprs + .iter() + .map(|expr| expr.output_column_ref(arena)) + .collect::>(), + ) + } else { + None + }; + OperatorExprVisitorMut::new(&mut substitute, arena).visit_operator( + &mut plan.operator, + None, + &mut ChildrenVisitMut::empty(), + )?; + match &mut plan.operator { + Operator::Project(op) => { + if let Some(schema) = schema { + for (expr, column) in op.exprs.iter_mut().zip(schema) { + let current = expr.output_column_ref(arena); + if arena.column(current).name() != arena.column(column).name() { + let alias = AliasType::Name(arena.column(column).name().to_string()); + *expr = arena + .alloc_expression(ScalarExpression::Alias { expr: *expr, alias }); + } + } + } + } + Operator::Filter(op) => op.is_optimized = false, + _ => (), + } + let Childrens::Only(child) = plan.childrens.as_mut() else { + return Ok(false); + }; + if merge { + replace_with_only_child(child); + } else { + // Lift the Project: its expressions keep the parent's output layout. + let mut project = child.take(); + let mut parent = plan.take(); + + std::mem::swap(&mut parent.childrens, &mut project.childrens); + + parent.reset_output_schema_cache(); + + *project.childrens = Childrens::Only(Box::new(parent)); + project.physical_option = None; + *plan = project; + } + plan.reset_output_schema_cache(); + Ok(true) + } +} + pub struct EliminateRedundantSort; impl NormalizationRule for EliminateRedundantSort { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let (sort_fields, topk_limit) = match &plan.operator { Operator::Sort(sort_op) => (sort_op.sort_fields.clone(), None), @@ -56,7 +309,11 @@ impl NormalizationRule for EliminateRedundantSort { offset, limit: Some(limit), }); - plan.physical_option = Some(PhysicalOption::new(PlanImpl::Limit, SortOption::Follow)); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::Limit, + SortOption::Follow, + RowCountOption::Unknown, + )); return Ok(true); } @@ -70,20 +327,23 @@ impl NormalizationRule for EliminateIndexFilter { fn apply( &self, plan: &mut LogicalPlan, - _: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { if !matches!(plan.operator, Operator::Filter(_)) { return Ok(false); } let residual = { - let Some(child) = only_child_mut(plan) else { + let Some((_, physical_option)) = + only_child_mut(plan).and_then(|child| predicate_scan_mut(child, arena)) + else { return Ok(false); }; let Some(PhysicalOption { plan: PlanImpl::IndexScan(index_info), .. - }) = child.physical_option.as_ref() + }) = physical_option.as_ref() else { return Ok(false); }; @@ -97,7 +357,11 @@ impl NormalizationRule for EliminateIndexFilter { let Operator::Filter(filter_op) = &mut plan.operator else { unreachable!("filter operator checked before residual rewrite"); }; + if filter_op.predicate == residual { + return Ok(false); + } filter_op.predicate = residual; + filter_op.is_optimized = false; return Ok(true); } @@ -175,7 +439,7 @@ pub(crate) fn apply_scan_order_hint( ScanOrderHint::SortFields(fields) => fields[index].expr, ScanOrderHint::GroupBy(groupby_exprs) => groupby_exprs[index], }; - if !expr.all_referenced_columns(arena, |arena, column| { + if !expr.all_referenced_columns(arena, |arena, column, _| { scan_op .columns .iter() @@ -264,7 +528,8 @@ impl NormalizationRule for UseStreamAggregate { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let Operator::Aggregate(op) = &plan.operator else { return Ok(false); @@ -296,7 +561,11 @@ impl NormalizationRule for UseStreamAggregate { return Ok(false); } - plan.physical_option = Some(PhysicalOption::new(implementation, SortOption::Follow)); + plan.physical_option = Some(PhysicalOption::new( + implementation, + SortOption::Follow, + RowCountOption::Unknown, + )); Ok(true) } } @@ -307,7 +576,8 @@ impl NormalizationRule for ForceSpillAggregate { fn apply( &self, plan: &mut LogicalPlan, - _: &mut crate::planner::PlanArena, + _: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let (implementation, sort_fields) = match (&plan.operator, &plan.physical_option) { ( @@ -341,8 +611,16 @@ impl NormalizationRule for ForceSpillAggregate { return Ok(false); } let sort = only_child_mut(plan).expect("aggregate child was wrapped with sort"); - sort.physical_option = Some(PhysicalOption::new(PlanImpl::Sort, sort_option)); - plan.physical_option = Some(PhysicalOption::new(implementation, SortOption::Follow)); + sort.physical_option = Some(PhysicalOption::new( + PlanImpl::Sort, + sort_option, + RowCountOption::Follow, + )); + plan.physical_option = Some(PhysicalOption::new( + implementation, + SortOption::Follow, + RowCountOption::Unknown, + )); Ok(true) } } @@ -353,14 +631,14 @@ pub(crate) fn apply_annotated_post_rules( ) -> Result { let mut changed = false; - if EliminateRedundantSort.apply(plan, arena)? { + if EliminateRedundantSort.apply(plan, arena, None)? { changed = true; } - if UseStreamAggregate.apply(plan, arena)? { + if UseStreamAggregate.apply(plan, arena, None)? { changed = true; } // Run last so an existing ordered child is reused before a forced spill adds a Sort. - if ForceSpillAggregate.apply(plan, arena)? { + if ForceSpillAggregate.apply(plan, arena, None)? { changed = true; } @@ -376,6 +654,7 @@ fn ensure_order( if let Some(PhysicalOption { plan: PlanImpl::IndexScan(index_info), sort_option, + .. }) = plan.physical_option.as_mut() { if covers(required, &index_info.sort_option, |required, provided| { @@ -494,26 +773,488 @@ pub(crate) fn covers( #[cfg(all(test, not(target_arch = "wasm32")))] mod tests { use super::{ - EliminateIndexFilter, EliminateRedundantSort, ForceSpillAggregate, UseStreamAggregate, + EliminateIndexFilter, EliminateProject, EliminateRedundantSort, ForceSpillAggregate, + UseStreamAggregate, }; + use crate::binder::test::build_t1_table; use crate::catalog::{ColumnCatalog, TableName}; use crate::errors::DatabaseError; + use crate::expression::agg::AggKind; use crate::expression::range_detacher::Range; - use crate::expression::ScalarExpression; + use crate::expression::{BinaryOperator, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; + use crate::optimizer::heuristic::batch::HepBatchStrategy; + use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline; + use crate::optimizer::rule::normalization::NormalizationRuleImpl; use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::operator::filter::FilterOperator; + use crate::planner::operator::join::{JoinCondition, JoinOperator, JoinType}; use crate::planner::operator::limit::LimitOperator; + use crate::planner::operator::project::ProjectOperator; + use crate::planner::operator::scalar_query_init::ScalarQueryInitOperator; use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::operator::top_k::TopKOperator; - use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; - use crate::planner::{Childrens, ExprRef, LogicalPlan}; + use crate::planner::operator::values::ValuesOperator; + use crate::planner::operator::{ + Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption, + }; + use crate::planner::{Childrens, ExprRef, LogicalPlan, MetaArena, PlanArena}; use crate::types::index::{IndexInfo, IndexLookup, IndexMeta, IndexType}; use crate::types::value::DataValue; use crate::types::ColumnId; use crate::types::LogicalType; use std::ops::Bound; + fn column(arena: &mut PlanArena, source: ExprRef, position: usize) -> ExprRef { + let column = source.output_column_ref(arena); + arena.alloc_expression(ScalarExpression::column_expr(column, position)) + } + + fn column_expr(arena: &mut PlanArena, name: &str, position: usize) -> ExprRef { + let column = arena.alloc_column(ColumnCatalog::new_dummy(name.to_string())); + arena.alloc_expression(ScalarExpression::column_expr(column, position)) + } + + fn binary(arena: &mut PlanArena, left: ExprRef, op: BinaryOperator, value: i32) -> ExprRef { + let right = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(value))); + arena.alloc_expression(ScalarExpression::Binary { + op, + left_expr: left, + right_expr: right, + evaluator: None, + ty: if matches!(op, BinaryOperator::Gt) { + LogicalType::Boolean + } else { + LogicalType::Integer + }, + }) + } + + fn source(arena: &mut PlanArena, mut plan: LogicalPlan) -> (LogicalPlan, ExprRef) { + let Operator::Project(op) = plan.operator else { + unreachable!() + }; + let expr = binary(arena, op.exprs[1], BinaryOperator::Plus, 1); + let input = plan.childrens.pop_only(); + plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![expr, op.exprs[0]])), + Childrens::Only(Box::new(input)), + ); + (plan, expr) + } + + fn eliminate_projects( + plan: &mut LogicalPlan, + arena: &mut PlanArena, + ) -> Result { + let original = plan.clone(); + let pipeline = HepOptimizerPipeline::builder() + .before_batch( + "eliminate_projects", + HepBatchStrategy::fix_point_bottomup(10), + &[NormalizationRuleImpl::EliminateProject], + ) + .build(); + *plan = pipeline.instantiate(plan.take()).find_best(None, arena)?; + Ok(*plan != original) + } + + #[test] + fn substitutes_through_filter_and_sort_without_cloning_sources() -> Result<(), DatabaseError> { + let table = build_t1_table()?; + let mut arena = PlanArena::new(&table.table_arena); + let scan = table.plan_with_arena("select c1, c2 from t1", &mut arena)?; + let (input, replacement) = source(&mut arena, scan); + let predicate_column = column(&mut arena, replacement, 0); + let predicate = binary(&mut arena, predicate_column, BinaryOperator::Gt, 10); + let filter = LogicalPlan::new( + Operator::Filter(FilterOperator { + predicate, + having: false, + is_optimized: false, + }), + Childrens::Only(Box::new(input)), + ); + let sort_column = column(&mut arena, replacement, 0); + let sort = LogicalPlan::new( + Operator::Sort(SortOperator { + sort_fields: vec![SortField { + expr: sort_column, + asc: true, + nulls_first: false, + }], + }), + Childrens::Only(Box::new(filter)), + ); + let result_column = column(&mut arena, replacement, 0); + let result = binary(&mut arena, result_column, BinaryOperator::Multiply, 2); + let mut plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![result])), + Childrens::Only(Box::new(sort)), + ); + + assert!(eliminate_projects(&mut plan, &mut arena)?); + let Operator::Project(op) = &plan.operator else { + unreachable!() + }; + let ScalarExpression::Binary { left_expr, .. } = op.exprs[0].unpack_alias_ref(&arena) + else { + unreachable!() + }; + assert_eq!(*left_expr, replacement); + let Childrens::Only(sort) = plan.childrens.as_ref() else { + unreachable!() + }; + let Operator::Sort(op) = &sort.operator else { + unreachable!() + }; + assert_eq!(op.sort_fields[0].expr, replacement); + let Childrens::Only(filter) = sort.childrens.as_ref() else { + unreachable!() + }; + let Operator::Filter(op) = &filter.operator else { + unreachable!() + }; + let ScalarExpression::Binary { left_expr, .. } = arena.expression(op.predicate) else { + unreachable!() + }; + assert_eq!(*left_expr, replacement); + let Childrens::Only(scan) = filter.childrens.as_ref() else { + unreachable!() + }; + assert!(matches!(scan.operator, Operator::TableScan(_))); + assert!(!eliminate_projects(&mut plan, &mut arena)?); + Ok(()) + } + + #[test] + fn aggregate_preserves_input_project_and_separate_output_slots() -> Result<(), DatabaseError> { + let table = build_t1_table()?; + let mut arena = PlanArena::new(&table.table_arena); + let scan = table.plan_with_arena("select c1, c2 from t1", &mut arena)?; + let (input, replacement) = source(&mut arena, scan); + let arg = column(&mut arena, replacement, 0); + let sum = arena.alloc_expression(ScalarExpression::AggCall { + distinct: false, + kind: AggKind::Sum, + args: vec![arg], + ty: LogicalType::Integer, + }); + let aggregate = LogicalPlan::new( + Operator::Aggregate(AggregateOperator { + agg_calls: vec![sum], + groupby_exprs: vec![], + is_distinct: false, + force_spill: false, + }), + Childrens::Only(Box::new(input)), + ); + let output = column(&mut arena, sum, 0); + let mut plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![output])), + Childrens::Only(Box::new(aggregate)), + ); + + let original = plan.clone(); + assert!(!eliminate_projects(&mut plan, &mut arena)?); + assert_eq!(plan, original); + let Operator::Project(op) = &plan.operator else { + unreachable!() + }; + assert_eq!(op.exprs[0], output); + let Childrens::Only(aggregate) = plan.childrens.as_ref() else { + unreachable!() + }; + let Operator::Aggregate(op) = &aggregate.operator else { + unreachable!() + }; + let ScalarExpression::AggCall { args, .. } = arena.expression(op.agg_calls[0]) else { + unreachable!() + }; + assert_eq!(args, &[arg]); + let Childrens::Only(input) = aggregate.childrens.as_ref() else { + unreachable!() + }; + let Operator::Project(op) = &input.operator else { + unreachable!() + }; + assert_eq!(op.exprs[0], replacement); + assert!(matches!( + input.childrens.only().operator, + Operator::TableScan(_) + )); + Ok(()) + } + + #[test] + fn mapping_transfer_preserves_shared_source_expressions_across_plan_nodes( + ) -> Result<(), DatabaseError> { + let table = build_t1_table()?; + let mut arena = PlanArena::new(&table.table_arena); + let scan = table.plan_with_arena("select c1, c2 from t1", &mut arena)?; + let (input, replacement) = source(&mut arena, scan); + let reference = column(&mut arena, replacement, 0); + let shared = binary(&mut arena, reference, BinaryOperator::Gt, 10); + let filter = LogicalPlan::new( + Operator::Filter(FilterOperator { + predicate: shared, + having: false, + is_optimized: false, + }), + Childrens::Only(Box::new(input)), + ); + let mut plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![shared])), + Childrens::Only(Box::new(filter)), + ); + + assert!(eliminate_projects(&mut plan, &mut arena)?); + let Operator::Project(project) = &plan.operator else { + unreachable!() + }; + let Operator::Filter(filter) = &plan.childrens.only().operator else { + unreachable!() + }; + for rewritten in [project.exprs[0], filter.predicate] { + let ScalarExpression::Binary { left_expr, .. } = rewritten.unpack_alias_ref(&arena) + else { + unreachable!() + }; + assert_eq!(*left_expr, replacement); + assert_ne!(rewritten, shared); + } + let ScalarExpression::Binary { left_expr, .. } = arena.expression(shared) else { + unreachable!() + }; + assert_eq!(*left_expr, reference); + let ScalarExpression::Binary { left_expr, .. } = arena.expression(replacement) else { + unreachable!() + }; + let ScalarExpression::ColumnRef { position, .. } = left_expr.unpack_alias_ref(&arena) + else { + unreachable!() + }; + assert_eq!(*position, 1); + assert!(!eliminate_projects(&mut plan, &mut arena)?); + Ok(()) + } + + fn values(arena: &mut PlanArena, names: &[&str]) -> LogicalPlan { + let schema = names + .iter() + .map(|name| arena.alloc_column(ColumnCatalog::new_dummy(name.to_string()))) + .collect(); + LogicalPlan::new( + Operator::Values(ValuesOperator::new(vec![], 0, schema)), + Childrens::None, + ) + } + + fn permutation(arena: &mut PlanArena, mut input: LogicalPlan, slots: &[usize]) -> LogicalPlan { + let schema = input.output_schema(arena).clone(); + let exprs = slots + .iter() + .map(|position| { + arena.alloc_expression(ScalarExpression::column_expr(schema[*position], *position)) + }) + .collect(); + LogicalPlan::new( + Operator::Project(ProjectOperator::new(exprs)), + Childrens::Only(Box::new(input)), + ) + } + + #[test] + fn cross_join_composes_p4_and_p9_into_final_computations() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let outer_left = values(&mut arena, &["d2", "e2", "c5"]); + let t3 = values(&mut arena, &["a3", "e3", "x3"]); + let t4 = values(&mut arena, &["b4", "d4"]); + let inner = JoinOperator::build(t3, t4, JoinCondition::None, JoinType::Inner, false); + let t9 = values(&mut arena, &["b9", "d9"]); + let t1 = values(&mut arena, &["a1", "e1"]); + let cross4 = JoinOperator::build(t9, t1, JoinCondition::None, JoinType::Cross, false); + let p4 = permutation(&mut arena, cross4, &[2, 3, 0, 1]); + let cross9 = JoinOperator::build(inner, p4, JoinCondition::None, JoinType::Cross, false); + let p9 = permutation(&mut arena, cross9, &[5, 6, 7, 8, 0, 1, 2, 3, 4]); + let outer = + JoinOperator::build(outer_left, p9, JoinCondition::None, JoinType::Cross, false); + let mut plan = permutation(&mut arena, outer, &(0..12).collect::>()); + let Operator::Project(op) = &mut plan.operator else { + unreachable!() + }; + for expr in &mut op.exprs { + *expr = binary(&mut arena, *expr, BinaryOperator::Multiply, 2); + } + let schema = plan.output_schema(&mut arena).clone(); + + assert!(eliminate_projects(&mut plan, &mut arena)?); + assert_eq!(plan.output_schema(&mut arena), &schema); + assert_eq!(plan.explain(&mut arena, 0).matches("Projection").count(), 1); + let Operator::Project(op) = &plan.operator else { + unreachable!() + }; + for (expr, expected) in op.exprs.iter().zip([0, 1, 2, 10, 11, 8, 9, 3, 4, 5, 6, 7]) { + let ScalarExpression::Binary { left_expr, .. } = expr.unpack_alias_ref(&arena) else { + unreachable!() + }; + let ScalarExpression::ColumnRef { position, .. } = left_expr.unpack_alias_ref(&arena) + else { + unreachable!() + }; + assert_eq!(*position, expected); + } + let Childrens::Only(input) = plan.childrens.as_mut() else { + unreachable!() + }; + let names: Vec<_> = input + .output_schema(&mut arena) + .iter() + .map(|column| arena.column(*column).name().to_string()) + .collect(); + assert_eq!( + names, + ["d2", "e2", "c5", "a3", "e3", "x3", "b4", "d4", "b9", "d9", "a1", "e1"] + ); + assert!(!eliminate_projects(&mut plan, &mut arena)?); + Ok(()) + } + + #[test] + fn cross_join_remaps_by_new_left_width() -> Result<(), DatabaseError> { + // (left slots, right slots, outer slots, expected positions); an empty left + // mapping must not be treated as identity. + type Case<'a> = (&'a [usize], Option<&'a [usize]>, &'a [usize], &'a [usize]); + let cases: [Case; 3] = [ + (&[2], Some(&[1, 0]), &[0, 1, 2], &[2, 4, 3]), + (&[2], None, &[0, 1, 2], &[2, 3, 4]), + (&[], None, &[0, 1], &[3, 4]), + ]; + for (left_slots, right_slots, outer_slots, expected) in cases { + let table_arena = crate::planner::TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let left = values(&mut arena, &["l0", "l1", "l2"]); + let left = permutation(&mut arena, left, left_slots); + let right = values(&mut arena, &["r0", "r1"]); + let right = match right_slots { + Some(slots) => permutation(&mut arena, right, slots), + None => right, + }; + let cross = + JoinOperator::build(left, right, JoinCondition::None, JoinType::Cross, false); + let mut plan = permutation(&mut arena, cross, outer_slots); + + assert!(eliminate_projects(&mut plan, &mut arena)?); + assert_eq!(plan.explain(&mut arena, 0).matches("Projection").count(), 1); + let Operator::Project(op) = &plan.operator else { + unreachable!() + }; + let positions: Vec<_> = op + .exprs + .iter() + .map(|expr| match expr.unpack_alias_ref(&arena) { + ScalarExpression::ColumnRef { position, .. } => *position, + _ => unreachable!(), + }) + .collect(); + assert_eq!(positions, expected); + assert!(!eliminate_projects(&mut plan, &mut arena)?); + } + Ok(()) + } + + #[test] + fn cross_join_inlines_computed_right_output_without_changing_source( + ) -> Result<(), DatabaseError> { + let table = build_t1_table()?; + let mut arena = PlanArena::new(&table.table_arena); + let scan = table.plan_with_arena("select c1, c2 from t1", &mut arena)?; + let (right, original) = source(&mut arena, scan); + let left = values(&mut arena, &["l0"]); + let cross = JoinOperator::build(left, right, JoinCondition::None, JoinType::Cross, false); + let mut plan = permutation(&mut arena, cross, &[0, 1, 2]); + + assert!(eliminate_projects(&mut plan, &mut arena)?); + assert_eq!(plan.explain(&mut arena, 0).matches("Projection").count(), 1); + let Operator::Project(op) = &plan.operator else { + unreachable!() + }; + for (expr, expected) in [(op.exprs[1], 2), (original, 1)] { + let ScalarExpression::Binary { left_expr, .. } = expr.unpack_alias_ref(&arena) else { + unreachable!() + }; + let ScalarExpression::ColumnRef { position, .. } = left_expr.unpack_alias_ref(&arena) + else { + unreachable!() + }; + assert_eq!(*position, expected); + } + Ok(()) + } + + #[test] + fn test_collapse_project() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let plan = + table_state.plan_with_arena("select c1 from (select c1, c2 from t1) t", &mut arena)?; + + let pipeline = HepOptimizerPipeline::builder() + .before_batch( + "test_collapse_project", + HepBatchStrategy::fix_point_bottomup(10), + &[NormalizationRuleImpl::EliminateProject], + ) + .build(); + let best_plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; + + if let Operator::Project(op) = &best_plan.operator { + assert_eq!(op.exprs.len(), 1); + } else { + unreachable!("Should be a project operator") + } + + let scan_op = best_plan.childrens.pop_only(); + assert!(matches!(scan_op.operator, Operator::TableScan(_))); + assert!(matches!(scan_op.childrens.as_ref(), Childrens::None)); + + Ok(()) + } + + #[test] + fn test_collapse_project_remaps_reordered_passthrough_positions() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let child = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![ + column_expr(&mut arena, "c2", 1), + column_expr(&mut arena, "c1", 0), + ])), + Childrens::Only(Box::new(LogicalPlan::new(Operator::Dummy, Childrens::None))), + ); + let mut plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![column_expr(&mut arena, "c2", 0)])), + Childrens::Only(Box::new(child)), + ); + + assert!(EliminateProject.apply(&mut plan, &mut arena, None)?); + + let Operator::Project(op) = &plan.operator else { + unreachable!("expected project"); + }; + let ScalarExpression::ColumnRef { position, .. } = op.exprs[0].unpack_alias_ref(&arena) + else { + unreachable!("expected column ref"); + }; + assert_eq!(*position, 1); + assert!(matches!( + plan.childrens.pop_only().operator, + Operator::Dummy + )); + Ok(()) + } + fn make_sort_field(arena: &mut crate::planner::PlanArena, name: &str) -> SortField { make_sort_field_with_position(arena, name, 0) } @@ -544,6 +1285,7 @@ mod tests { leaf.physical_option = Some(PhysicalOption::new( PlanImpl::IndexScan(Box::new(index_info)), index_sort_option, + RowCountOption::Unknown, )); let predicate = @@ -556,7 +1298,11 @@ mod tests { }), Childrens::Only(Box::new(leaf)), ); - filter.physical_option = Some(PhysicalOption::new(PlanImpl::Filter, SortOption::Follow)); + filter.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Unknown, + )); LogicalPlan::new( Operator::Sort(SortOperator { @@ -632,6 +1378,7 @@ mod tests { scan.physical_option = Some(PhysicalOption::new( PlanImpl::IndexScan(Box::new(index_info)), sort_option, + RowCountOption::Unknown, )); let mut filter = LogicalPlan::new( @@ -642,7 +1389,11 @@ mod tests { }), Childrens::Only(Box::new(scan)), ); - filter.physical_option = Some(PhysicalOption::new(PlanImpl::Filter, SortOption::Follow)); + filter.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Unknown, + )); filter } @@ -704,35 +1455,45 @@ mod tests { } #[test] - fn exact_index_filter_is_removed_after_physical_selection() -> Result<(), DatabaseError> { - let table_arena = crate::planner::TableArenaCell::default(); - let mut arena = crate::planner::PlanArena::new(&table_arena); - let predicate = - arena.alloc_expression(ScalarExpression::Constant(DataValue::Boolean(true))); - let mut plan = build_filter_with_selected_index(&mut arena, predicate, None); - - let rule = EliminateIndexFilter; - assert!(rule.apply(&mut plan, &mut arena)?); - assert!(matches!(plan.operator, Operator::TableScan(_))); - Ok(()) - } - - #[test] - fn partial_index_filter_keeps_residual_after_physical_selection() -> Result<(), DatabaseError> { + fn index_filter_elimination_preserves_residual_and_initializer() -> Result<(), DatabaseError> { let table_arena = crate::planner::TableArenaCell::default(); let mut arena = crate::planner::PlanArena::new(&table_arena); let predicate = arena.alloc_expression(ScalarExpression::Constant(DataValue::Boolean(true))); let residual = arena.alloc_expression(ScalarExpression::Constant(DataValue::Boolean(false))); - let mut plan = build_filter_with_selected_index(&mut arena, predicate, Some(residual)); - + let id = arena.alloc_scalar_query_ref(false); + let value = arena.alloc_expression(ScalarExpression::InitValue { + id, + ty: LogicalType::Integer, + }); let rule = EliminateIndexFilter; - assert!(rule.apply(&mut plan, &mut arena)?); - let Operator::Filter(filter_op) = &plan.operator else { - unreachable!("residual predicate should keep filter"); - }; - assert_eq!(filter_op.predicate, residual); + for residual in [None, Some(residual)] { + for through_init in [false, true] { + let mut plan = build_filter_with_selected_index(&mut arena, predicate, residual); + if through_init { + let scan = plan.childrens.pop_only(); + *plan.childrens = Childrens::Only(Box::new(ScalarQueryInitOperator::build( + scan, + LogicalPlan::new(Operator::Dummy, Childrens::None), + value, + vec![], + ))); + } + let input = plan.childrens.only().clone(); + assert!(rule.apply(&mut plan, &mut arena, None)?); + if let Some(residual) = residual { + let Operator::Filter(filter) = &plan.operator else { + panic!("residual predicate should keep filter") + }; + assert_eq!(filter.predicate, residual); + assert_eq!(plan.childrens.only(), &input); + } else { + assert_eq!(plan, input); + } + assert!(!rule.apply(&mut plan, &mut arena, None)?); + } + } Ok(()) } @@ -769,7 +1530,7 @@ mod tests { index_info.lookup = Some(IndexLookup::Probe); let rule = EliminateIndexFilter; - assert!(!rule.apply(&mut plan, &mut arena)?); + assert!(!rule.apply(&mut plan, &mut arena, None)?); let Operator::Filter(filter) = &plan.operator else { panic!("probe cannot discharge the static filter") }; @@ -785,7 +1546,7 @@ mod tests { let mut plan = build_plan(&mut arena, vec![sort_field.clone()], vec![sort_field], 0); let rule = EliminateRedundantSort; - assert!(rule.apply(&mut plan, &mut arena)?); + assert!(rule.apply(&mut plan, &mut arena, None)?); assert!(matches!(plan.operator, Operator::Filter(_))); Ok(()) } @@ -808,7 +1569,7 @@ mod tests { }); let rule = EliminateRedundantSort; - assert!(rule.apply(&mut plan, &mut arena)?); + assert!(rule.apply(&mut plan, &mut arena, None)?); match plan.operator { Operator::Limit(limit_op) => { assert_eq!(limit_op.limit, Some(10)); @@ -864,7 +1625,7 @@ mod tests { let mut arena = crate::planner::PlanArena::new(&table_arena); let (mut plan, _) = desc_topk_plan(&mut arena, false); - assert!(EliminateRedundantSort.apply(&mut plan, &mut arena)?); + assert!(EliminateRedundantSort.apply(&mut plan, &mut arena, None)?); assert!(matches!(plan.operator, Operator::Limit(_))); assert!(scan_is_reverse(&plan)); Ok(()) @@ -876,7 +1637,7 @@ mod tests { let mut arena = crate::planner::PlanArena::new(&table_arena); let (mut plan, _) = desc_topk_plan(&mut arena, true); - assert!(!EliminateRedundantSort.apply(&mut plan, &mut arena)?); + assert!(!EliminateRedundantSort.apply(&mut plan, &mut arena, None)?); assert!(matches!(plan.operator, Operator::TopK(_))); assert!(!scan_is_reverse(&plan)); Ok(()) @@ -894,9 +1655,13 @@ mod tests { offset: None, limit: Some(3), }); - filter.physical_option = Some(PhysicalOption::new(PlanImpl::Limit, SortOption::Follow)); + filter.physical_option = Some(PhysicalOption::new( + PlanImpl::Limit, + SortOption::Follow, + RowCountOption::Unknown, + )); - assert!(!EliminateRedundantSort.apply(&mut plan, &mut arena)?); + assert!(!EliminateRedundantSort.apply(&mut plan, &mut arena, None)?); assert!(!scan_is_reverse(&plan)); Ok(()) } @@ -934,7 +1699,7 @@ mod tests { info.lookup = lookup; } super::mark_sort_preserving_indexes(&mut plan, std::slice::from_ref(&c2), &arena)?; - assert!(EliminateRedundantSort.apply(&mut plan, &mut arena)?); + assert!(EliminateRedundantSort.apply(&mut plan, &mut arena, None)?); } Ok(()) } @@ -962,7 +1727,7 @@ mod tests { }); let rule = EliminateRedundantSort; - assert!(rule.apply(&mut plan, &mut arena)?); + assert!(rule.apply(&mut plan, &mut arena, None)?); assert!(matches!(plan.operator, Operator::Limit(_))); Ok(()) } @@ -1063,16 +1828,18 @@ mod tests { child.physical_option = Some(PhysicalOption::new( PlanImpl::IndexScan(Box::new(index_info)), sort_option, + RowCountOption::Unknown, )); } } plan.physical_option = Some(PhysicalOption::new( PlanImpl::HashAggregate, SortOption::None, + RowCountOption::Unknown, )); let rule = UseStreamAggregate; - assert!(rule.apply(&mut plan, &mut arena)?); + assert!(rule.apply(&mut plan, &mut arena, None)?); assert!(matches!( plan.physical_option, Some(PhysicalOption { @@ -1102,14 +1869,16 @@ mod tests { child.physical_option = Some(PhysicalOption::new( PlanImpl::IndexScan(Box::new(scan_op.index_infos[0].clone())), sort_option, + RowCountOption::Unknown, )); plan.physical_option = Some(PhysicalOption::new( PlanImpl::HashAggregate, SortOption::None, + RowCountOption::Unknown, )); let rule = UseStreamAggregate; - assert!(rule.apply(&mut plan, &mut arena)?); + assert!(rule.apply(&mut plan, &mut arena, None)?); assert!(matches!( plan.physical_option, Some(PhysicalOption { @@ -1118,7 +1887,7 @@ mod tests { }) )); let force_spill = ForceSpillAggregate; - assert!(!force_spill.apply(&mut plan, &mut arena)?); + assert!(!force_spill.apply(&mut plan, &mut arena, None)?); assert!(!matches!( plan.childrens.as_ref(), Childrens::Only(child) if matches!(child.operator, Operator::Sort(_)) @@ -1132,8 +1901,9 @@ mod tests { unordered.physical_option = Some(PhysicalOption::new( PlanImpl::HashAggregate, SortOption::None, + RowCountOption::Unknown, )); - assert!(!rule.apply(&mut unordered, &mut arena)?); + assert!(!rule.apply(&mut unordered, &mut arena, None)?); assert!(matches!( unordered.physical_option, Some(PhysicalOption { @@ -1158,10 +1928,11 @@ mod tests { plan.physical_option = Some(PhysicalOption::new( PlanImpl::HashAggregate, SortOption::None, + RowCountOption::Unknown, )); let rule = ForceSpillAggregate; - assert!(rule.apply(&mut plan, &mut arena)?); + assert!(rule.apply(&mut plan, &mut arena, None)?); assert!(matches!( plan.physical_option, Some(PhysicalOption { @@ -1200,7 +1971,7 @@ mod tests { super::mark_sort_preserving_indexes(&mut plan, &[c2], &arena)?; let rule = EliminateRedundantSort; - assert!(!rule.apply(&mut plan, &mut arena)?); + assert!(!rule.apply(&mut plan, &mut arena, None)?); assert!(matches!(plan.operator, Operator::Sort(_))); Ok(()) } @@ -1238,6 +2009,7 @@ mod tests { scan_plan.physical_option = Some(PhysicalOption::new( PlanImpl::IndexScan(Box::new(index_info.clone())), index_info.sort_option, + RowCountOption::Unknown, )); } @@ -1250,7 +2022,11 @@ mod tests { }), Childrens::Only(Box::new(scan_plan)), ); - filter.physical_option = Some(PhysicalOption::new(PlanImpl::Filter, SortOption::Follow)); + filter.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Unknown, + )); let mut plan = LogicalPlan::new( Operator::Sort(SortOperator { @@ -1265,7 +2041,7 @@ mod tests { }; super::mark_sort_preserving_indexes(&mut plan, &sort_fields, &arena)?; let rule = EliminateRedundantSort; - assert!(rule.apply(&mut plan, &mut arena)?); + assert!(rule.apply(&mut plan, &mut arena, None)?); assert!(matches!(plan.operator, Operator::Filter(_))); let table_plan = plan.childrens.pop_only(); diff --git a/src/optimizer/rule/normalization/guard_mark_apply.rs b/src/optimizer/rule/normalization/guard_mark_apply.rs new file mode 100644 index 00000000..1251bc18 --- /dev/null +++ b/src/optimizer/rule/normalization/guard_mark_apply.rs @@ -0,0 +1,980 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::catalog::ColumnRef; +use crate::errors::DatabaseError; +use crate::expression::visitor::ExprVisitor; +use crate::expression::{BinaryOperator, ScalarExpression, TypeCast, UnaryOperator}; +use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::optimizer::plan_utils::{only_child, only_child_mut}; +use crate::planner::operator::Operator; +use crate::planner::{Childrens, ExprRef, LogicalPlan, MetaArena, PlanArena}; +use crate::types::evaluator::BinaryEvaluatorRef; +use crate::types::value::DataValue; +use crate::types::LogicalType; +use std::iter::successors; + +pub struct GuardMarkApply; + +impl NormalizationRule for GuardMarkApply { + fn apply( + &self, + plan: &mut LogicalPlan, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, + ) -> Result { + if !matches!(plan.operator, Operator::Project(_)) { + return Ok(false); + } + let Some(filter) = skip_row_order(plan) else { + return Ok(false); + }; + let (Operator::Filter(filter_op), Childrens::Only(mark)) = + (&filter.operator, filter.childrens.as_ref()) + else { + return Ok(false); + }; + let predicate = filter_op.predicate; + let mut markers = Vec::new(); + let mut input = mark.as_ref(); + while let (Operator::MarkApply(op), Childrens::Twins { left, .. }) = + (&input.operator, input.childrens.as_ref()) + { + markers.push(*op.output_column()); + input = left; + } + if markers.is_empty() { + return Ok(false); + } + // No marker in this segment may be read above the filter. + let segment = successors(Some(&*plan), |node| only_child(node)) + .take_while(|node| !matches!(node.operator, Operator::Filter(_))); + for node in segment { + if !node + .operator + .visit_referenced_columns(arena, &mut |arena, column| { + !markers + .iter() + .any(|marker| arena.same_column(*column, *marker)) + })? + { + return Ok(false); + } + } + // Do not skip values consumed by another apply, even if Filter does not need them. + let mut consumed = Vec::new(); + let mut input = mark.as_ref(); + for index in 0..markers.len() { + let (Operator::MarkApply(op), Childrens::Twins { left, right }) = + (&input.operator, input.childrens.as_ref()) + else { + unreachable!() + }; + let later = &markers[index + 1..]; + let mut consume = |arena: &dyn MetaArena, column: &ColumnRef| { + if let Some(marker) = later.iter().find(|m| arena.same_column(*column, **m)) { + consumed.push(*marker); + } + true + }; + for expr in op.predicates.iter().chain(op.parameterized_probe.iter()) { + expr.all_referenced_columns(arena, |arena, column, _| consume(arena, column))?; + } + right.visit_referenced_columns(arena, &mut |arena, column| consume(arena, column))?; + input = left; + } + let mut changed = false; + let mut guards = Vec::new(); + let Some(mut input) = mark_below_filter(plan) else { + return Ok(false); + }; + for (index, marker) in markers.iter().enumerate() { + let (Operator::MarkApply(op), Childrens::Twins { left, .. }) = + (&mut input.operator, input.childrens.as_mut()) + else { + unreachable!() + }; + if op.guard.is_none() + && !consumed + .iter() + .any(|column| arena.same_column(*column, *marker)) + { + // Top-down plan order is the reverse of execution order. + let unavailable = &markers[..=index]; + guards.clear(); + let mut visitor = GuardVisitor { + marker: *marker, + guards: &mut guards, + }; + visitor.visit(predicate, arena)?; + let mut guard = None; + for &(sibling, when_true) in &guards { + let skip = decisive_condition(sibling, when_true, unavailable, arena)?; + let needed = skip.map(|expr| negate(expr, arena)); + guard = combine_conditions(BinaryOperator::And, guard, needed, true, arena); + } + if let Some(guard) = guard { + let schema = left.output_schema(arena); + if guard.all_referenced_columns(arena, |arena, column, _| { + !unavailable + .iter() + .any(|marker| arena.same_column(*marker, *column)) + && schema + .iter() + .any(|candidate| arena.same_column(*candidate, *column)) + })? { + op.guard = Some(guard.clone_expression(arena)?); + changed = true; + } + } + } + input = left; + } + Ok(changed) + } +} + +// Sort/Limit/TopK between the projection and the filter only reorder or cut rows. +fn skip_row_order(plan: &LogicalPlan) -> Option<&LogicalPlan> { + let mut child = only_child(plan)?; + while matches!( + child.operator, + Operator::Sort(_) | Operator::Limit(_) | Operator::TopK(_) + ) { + child = only_child(child)?; + } + Some(child) +} + +fn mark_below_filter(plan: &mut LogicalPlan) -> Option<&mut LogicalPlan> { + let mut node = plan; + while !matches!(node.operator, Operator::Filter(_)) { + node = only_child_mut(node)?; + } + only_child_mut(node) +} + +struct GuardVisitor<'a> { + marker: ColumnRef, + guards: &'a mut Vec<(ExprRef, bool)>, +} + +impl ExprVisitor> for GuardVisitor<'_> { + fn visit_expression_ref(&mut self, _: ExprRef, _: &PlanArena) -> Result { + Ok(true) + } + + fn visit_expression( + &mut self, + expr: &ScalarExpression, + _: &PlanArena, + ) -> Result { + Ok(matches!( + expr, + ScalarExpression::Binary { + op: BinaryOperator::And | BinaryOperator::Or, + .. + } | ScalarExpression::Unary { + op: UnaryOperator::Not, + .. + } | ScalarExpression::TypeCast { .. } + | ScalarExpression::Alias { .. } + )) + } + + fn visit_binary( + &mut self, + op: &BinaryOperator, + left: ExprRef, + right: ExprRef, + _: Option<&BinaryEvaluatorRef>, + _: &LogicalType, + arena: &PlanArena, + ) -> Result<(), DatabaseError> { + let left_marker = left.any_referenced_column(arena, |arena, column| { + arena.same_column(*column, self.marker) + })?; + let right_marker = right.any_referenced_column(arena, |arena, column| { + arena.same_column(*column, self.marker) + })?; + let (branch, sibling) = match (left_marker, right_marker) { + (true, false) => (left, right), + (false, true) => (right, left), + // Repeated occurrences can only use conditions collected above this node. + _ => return Ok(()), + }; + self.guards + .push((sibling, matches!(op, BinaryOperator::Or))); + self.visit(branch, arena) + } +} + +// Return only sufficient TRUE/FALSE conditions; unavailable marks and NULL are not evidence. +fn decisive_condition( + expr: ExprRef, + when_true: bool, + unavailable: &[ColumnRef], + arena: &mut PlanArena, +) -> Result, DatabaseError> { + if !expr.any_referenced_column(arena, |arena, column| { + unavailable + .iter() + .any(|marker| arena.same_column(*column, *marker)) + })? { + let mut visitor = SimpleGuardVisitor { + valid: true, + context: GuardContext::Predicate, + }; + visitor.visit(expr, arena)?; + if visitor.valid { + return Ok(Some(if when_true { expr } else { negate(expr, arena) })); + } + } + match arena.expression(expr.unpack_alias(arena)).clone() { + ScalarExpression::Binary { + op: op @ (BinaryOperator::And | BinaryOperator::Or), + left_expr, + right_expr, + .. + } => { + let left = decisive_condition(left_expr, when_true, unavailable, arena)?; + let right = decisive_condition(right_expr, when_true, unavailable, arena)?; + let either = matches!(op, BinaryOperator::Or) == when_true; + Ok(combine_conditions( + if either { + BinaryOperator::Or + } else { + BinaryOperator::And + }, + left, + right, + either, + arena, + )) + } + ScalarExpression::Unary { + op: UnaryOperator::Not, + expr, + .. + } => decisive_condition(expr, !when_true, unavailable, arena), + ScalarExpression::TypeCast { + expr, + ty: LogicalType::Boolean, + .. + } if *expr.return_type(arena) == LogicalType::Boolean => { + decisive_condition(expr, when_true, unavailable, arena) + } + _ => Ok(None), + } +} + +fn combine_conditions( + op: BinaryOperator, + left: Option, + right: Option, + allow_one: bool, + arena: &mut PlanArena, +) -> Option { + match (left, right) { + (Some(left_expr), Some(right_expr)) => { + Some(arena.alloc_expression(ScalarExpression::Binary { + op, + left_expr, + right_expr, + evaluator: None, + ty: LogicalType::Boolean, + })) + } + (Some(expr), None) | (None, Some(expr)) if allow_one => Some(expr), + _ => None, + } +} + +fn negate(expr: ExprRef, arena: &mut PlanArena) -> ExprRef { + if let ScalarExpression::Unary { + op: UnaryOperator::Not, + expr, + .. + } = arena.expression(expr) + { + return *expr; + } + arena.alloc_expression(ScalarExpression::Unary { + op: UnaryOperator::Not, + expr, + evaluator: None, + ty: LogicalType::Boolean, + }) +} + +#[derive(Clone, Copy)] +enum GuardContext { + Predicate, + NullOperand, + NumericOperand, +} + +struct SimpleGuardVisitor { + valid: bool, + context: GuardContext, +} + +impl ExprVisitor> for SimpleGuardVisitor { + fn visit_expression_ref( + &mut self, + expr: ExprRef, + arena: &PlanArena, + ) -> Result { + if !self.valid { + return Ok(false); + } + let unaliased = expr.unpack_alias(arena); + if unaliased != expr { + self.visit(unaliased, arena)?; + return Ok(false); + } + Ok(true) + } + + fn visit_expression( + &mut self, + expr: &ScalarExpression, + arena: &PlanArena, + ) -> Result { + self.valid = match (self.context, expr) { + ( + GuardContext::NullOperand, + ScalarExpression::Constant(_) | ScalarExpression::ColumnRef { .. }, + ) => true, + (GuardContext::Predicate, ScalarExpression::Constant(DataValue::Boolean(_))) => true, + (GuardContext::Predicate, ScalarExpression::ColumnRef { column, .. }) => { + arena.column(*column).datatype() == &LogicalType::Boolean + } + ( + GuardContext::Predicate, + ScalarExpression::IsNull { .. } + | ScalarExpression::Unary { + op: UnaryOperator::Not, + .. + } + | ScalarExpression::Binary { + op: + BinaryOperator::And + | BinaryOperator::Or + | BinaryOperator::Eq + | BinaryOperator::NotEq + | BinaryOperator::Gt + | BinaryOperator::GtEq + | BinaryOperator::Lt + | BinaryOperator::LtEq, + .. + }, + ) => true, + (GuardContext::NumericOperand, ScalarExpression::Constant(value)) => { + value.is_null() || value.logical_type().is_numeric() + } + (GuardContext::NumericOperand, ScalarExpression::ColumnRef { column, .. }) => { + arena.column(*column).datatype().is_numeric() + } + (GuardContext::NumericOperand, ScalarExpression::TypeCast { expr, ty, .. }) => { + let from = expr.return_type(arena); + // Implicit numeric widening cannot fail; bounded Decimal and narrowing casts can. + (from.is_numeric() || *from == LogicalType::SqlNull) + && ty.is_numeric() + && !matches!(ty, LogicalType::Decimal(..)) + && LogicalType::can_implicit_cast(&from, ty) + } + // Do not move arbitrary functions or potentially failing calculations earlier. + _ => false, + }; + Ok(self.valid) + } + + fn visit_is_null( + &mut self, + _: bool, + expr: ExprRef, + arena: &PlanArena, + ) -> Result<(), DatabaseError> { + let previous = self.context; + self.context = GuardContext::NullOperand; + self.visit(expr, arena)?; + self.context = previous; + Ok(()) + } + + fn visit_binary( + &mut self, + op: &BinaryOperator, + left: ExprRef, + right: ExprRef, + _: Option<&BinaryEvaluatorRef>, + _: &LogicalType, + arena: &PlanArena, + ) -> Result<(), DatabaseError> { + let previous = self.context; + self.context = if matches!(op, BinaryOperator::And | BinaryOperator::Or) { + GuardContext::Predicate + } else { + GuardContext::NumericOperand + }; + self.visit(left, arena)?; + self.visit(right, arena)?; + self.context = previous; + Ok(()) + } +} + +// GRCOV_EXCL_START +#[cfg(test)] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::expression::AliasType; + use crate::optimizer::rule::normalization::compilation_in_advance::EvaluatorBind; + use crate::planner::operator::filter::FilterOperator; + use crate::planner::operator::mark_apply::MarkApplyOperator; + use crate::planner::operator::project::ProjectOperator; + use crate::planner::operator::values::ValuesOperator; + use crate::planner::TableArenaCell; + use crate::types::evaluator::cast_create; + use crate::types::tuple::Tuple; + + struct Fixture { + plan: LogicalPlan, + flag: ExprRef, + marker: ExprRef, + } + + impl Fixture { + fn new( + arena: &mut PlanArena, + op: BinaryOperator, + marker_first: bool, + ) -> Result { + let flag_column = arena.alloc_column(ColumnCatalog::new( + "flag".to_string(), + true, + ColumnDesc::new(LogicalType::Boolean, None, false, None)?, + )); + let marker_column = arena.alloc_column(ColumnCatalog::new( + "marker".to_string(), + true, + ColumnDesc::new(LogicalType::Boolean, None, false, None)?, + )); + let value = arena.alloc_expression(DataValue::Boolean(true).into()); + let left = LogicalPlan::new( + Operator::Values(ValuesOperator::new(vec![value], 1, vec![flag_column])), + Childrens::None, + ); + let right = LogicalPlan::new( + Operator::Values(ValuesOperator::new(vec![value], 1, vec![flag_column])), + Childrens::None, + ); + let mark = MarkApplyOperator::build_exists(left, right, marker_column, vec![]); + let flag = arena.alloc_expression(ScalarExpression::column_expr(flag_column, 0)); + let marker = arena.alloc_expression(ScalarExpression::column_expr(marker_column, 1)); + let (left_expr, right_expr) = if marker_first { + (marker, flag) + } else { + (flag, marker) + }; + let predicate = arena.alloc_expression(ScalarExpression::Binary { + op, + left_expr, + right_expr, + evaluator: None, + ty: LogicalType::Boolean, + }); + let filter = FilterOperator::build(predicate, mark, false); + let plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![flag])), + Childrens::Only(Box::new(filter)), + ); + Ok(Self { plan, flag, marker }) + } + + fn add_mark( + &mut self, + arena: &mut PlanArena, + op: BinaryOperator, + ) -> Result { + let column = arena.alloc_column(ColumnCatalog::new( + "upper_marker".into(), + true, + ColumnDesc::new(LogicalType::Boolean, None, false, None)?, + )); + let upper = arena.alloc_expression(ScalarExpression::column_expr(column, 2)); + let value = arena.alloc_expression(DataValue::Boolean(true).into()); + let right = LogicalPlan::new( + Operator::Values(ValuesOperator::new(vec![value], 1, vec![column])), + Childrens::None, + ); + let Childrens::Only(filter) = self.plan.childrens.as_mut() else { + unreachable!() + }; + let Childrens::Only(mark) = filter.childrens.as_mut() else { + unreachable!() + }; + let lower = mark.take(); + **mark = MarkApplyOperator::build_exists(lower, right, column, vec![]); + let markers = arena.alloc_expression(ScalarExpression::Binary { + op: if matches!(op, BinaryOperator::And) { + BinaryOperator::Or + } else { + BinaryOperator::And + }, + left_expr: self.marker, + right_expr: upper, + evaluator: None, + ty: LogicalType::Boolean, + }); + let predicate = arena.alloc_expression(ScalarExpression::Binary { + op, + left_expr: self.flag, + right_expr: markers, + evaluator: None, + ty: LogicalType::Boolean, + }); + let Operator::Filter(filter) = &mut filter.operator else { + unreachable!() + }; + filter.predicate = predicate; + Ok(upper) + } + + fn mark(&mut self) -> &mut MarkApplyOperator { + let Childrens::Only(filter) = self.plan.childrens.as_mut() else { + unreachable!() + }; + let Childrens::Only(mark) = filter.childrens.as_mut() else { + unreachable!() + }; + let Operator::MarkApply(op) = &mut mark.operator else { + unreachable!() + }; + op + } + + fn predicate(&mut self) -> ExprRef { + let filter = self.plan.childrens.only(); + let Operator::Filter(op) = &filter.operator else { + unreachable!() + }; + op.predicate + } + } + + #[test] + fn independent_guards_preserve_all_marker_combinations() -> Result<(), DatabaseError> { + for op in [BinaryOperator::And, BinaryOperator::Or] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut fixture = Fixture::new(&mut arena, op, false)?; + fixture.add_mark(&mut arena, op)?; + assert!(GuardMarkApply.apply(&mut fixture.plan, &mut arena, None)?); + let upper = fixture.mark().guard.unwrap(); + let Childrens::Only(filter) = fixture.plan.childrens.as_mut() else { + unreachable!() + }; + let Childrens::Only(mark) = filter.childrens.as_mut() else { + unreachable!() + }; + let Childrens::Twins { left, .. } = mark.childrens.as_mut() else { + unreachable!() + }; + let Operator::MarkApply(lower) = &left.operator else { + unreachable!() + }; + let lower = lower.guard.unwrap(); + assert_ne!(upper, lower); + let predicate = fixture.predicate(); + EvaluatorBind.apply(&mut fixture.plan, &mut arena, None)?; + for flag in [ + DataValue::Boolean(false), + DataValue::Boolean(true), + DataValue::Null, + ] { + let input = Tuple::new(None, vec![flag.clone()]); + let skip_lower = *arena.expression(lower).eval(&arena, Some(&input))? + == DataValue::Boolean(false); + for a in [ + DataValue::Boolean(false), + DataValue::Boolean(true), + DataValue::Null, + ] { + for b in [ + DataValue::Boolean(false), + DataValue::Boolean(true), + DataValue::Null, + ] { + let original = Tuple::new(None, vec![flag.clone(), a.clone(), b.clone()]); + let a_out = if skip_lower { + DataValue::Null + } else { + a.clone() + }; + let input = Tuple::new(None, vec![flag.clone(), a_out.clone()]); + let skip_upper = *arena.expression(upper).eval(&arena, Some(&input))? + == DataValue::Boolean(false); + let guarded = Tuple::new( + None, + vec![ + flag.clone(), + a_out, + if skip_upper { + DataValue::Null + } else { + b.clone() + }, + ], + ); + assert_eq!( + arena + .expression(predicate) + .eval(&arena, Some(&original))? + .into_owned(), + arena + .expression(predicate) + .eval(&arena, Some(&guarded))? + .into_owned() + ); + } + } + } + assert!(!GuardMarkApply.apply(&mut fixture.plan, &mut arena, None)?); + } + Ok(()) + } + + #[test] + fn independent_protection_keeps_projected_markers_unguarded() -> Result<(), DatabaseError> { + for escaped in [0, 1, 2] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut fixture = Fixture::new(&mut arena, BinaryOperator::And, false)?; + let upper = fixture.add_mark(&mut arena, BinaryOperator::And)?; + if escaped == 0 { + // (P AND lower) OR upper cannot skip both marks when P=false. + let protected = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::And, + left_expr: fixture.flag, + right_expr: fixture.marker, + evaluator: None, + ty: LogicalType::Boolean, + }); + let predicate = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Or, + left_expr: protected, + right_expr: upper, + evaluator: None, + ty: LogicalType::Boolean, + }); + let Childrens::Only(filter) = fixture.plan.childrens.as_mut() else { + unreachable!() + }; + let Operator::Filter(filter) = &mut filter.operator else { + unreachable!() + }; + filter.predicate = predicate; + } else { + let Operator::Project(project) = &mut fixture.plan.operator else { + unreachable!() + }; + project + .exprs + .push(if escaped == 1 { fixture.marker } else { upper }); + } + assert_eq!( + GuardMarkApply.apply(&mut fixture.plan, &mut arena, None)?, + escaped == 0 + ); + assert_eq!(fixture.mark().guard.is_some(), escaped == 0); + let Childrens::Only(filter) = fixture.plan.childrens.as_mut() else { + unreachable!() + }; + let Childrens::Only(mark) = filter.childrens.as_mut() else { + unreachable!() + }; + let Childrens::Twins { left, .. } = mark.childrens.as_mut() else { + unreachable!() + }; + let Operator::MarkApply(lower) = &left.operator else { + unreachable!() + }; + assert_eq!(lower.guard.is_some(), escaped == 0); + } + Ok(()) + } + + #[test] + fn marker_in_another_branch_is_not_guarded() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut fixture = Fixture::new(&mut arena, BinaryOperator::Or, false)?; + let guarded_branch = fixture.predicate(); + let predicate = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Or, + left_expr: guarded_branch, + right_expr: fixture.marker, + evaluator: None, + ty: LogicalType::Boolean, + }); + let Childrens::Only(filter) = fixture.plan.childrens.as_mut() else { + unreachable!() + }; + let Operator::Filter(op) = &mut filter.operator else { + unreachable!() + }; + op.predicate = predicate; + assert!(!GuardMarkApply.apply(&mut fixture.plan, &mut arena, None)?); + assert!(fixture.mark().guard.is_none()); + Ok(()) + } + + #[test] + fn three_mark_guards_preserve_all_729_truth_combinations() -> Result<(), DatabaseError> { + use crate::expression::visitor_mut::ExprVisitorMut; + use crate::expression::BindEvaluator; + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut columns = Vec::new(); + let mut exprs = Vec::new(); + for (position, name) in ["q2", "p2", "q3", "m1", "m2", "m3"].iter().enumerate() { + let column = arena.alloc_column(ColumnCatalog::new( + (*name).into(), + true, + ColumnDesc::new(LogicalType::Boolean, None, false, None)?, + )); + columns.push(column); + exprs.push(arena.alloc_expression(ScalarExpression::column_expr(column, position))); + } + let binary = |op, left_expr, right_expr, arena: &mut PlanArena| { + arena.alloc_expression(ScalarExpression::Binary { + op, + left_expr, + right_expr, + evaluator: None, + ty: LogicalType::Boolean, + }) + }; + let p2_or_m2 = binary(BinaryOperator::Or, exprs[1], exprs[4], &mut arena); + let b2 = binary(BinaryOperator::And, exprs[0], p2_or_m2, &mut arena); + let b3 = binary(BinaryOperator::And, exprs[2], exprs[5], &mut arena); + let first = binary(BinaryOperator::Or, exprs[3], b2, &mut arena); + let mut predicate = binary(BinaryOperator::Or, first, b3, &mut arena); + let mut guards = Vec::new(); + let mut conditions = Vec::new(); + for index in 3..6 { + conditions.clear(); + let mut visitor = GuardVisitor { + marker: columns[index], + guards: &mut conditions, + }; + visitor.visit(predicate, &arena)?; + let mut guard = None; + for &(sibling, when_true) in &conditions { + let skip = decisive_condition(sibling, when_true, &columns[index..], &mut arena)?; + let needed = skip.map(|expr| negate(expr, &mut arena)); + guard = combine_conditions(BinaryOperator::And, guard, needed, true, &mut arena); + } + let mut guard = guard.unwrap(); + BindEvaluator.visit(&mut guard, &mut arena)?; + guards.push(guard); + } + BindEvaluator.visit(&mut predicate, &mut arena)?; + let values = [ + DataValue::Boolean(false), + DataValue::Boolean(true), + DataValue::Null, + ]; + let mut skipped = [0; 3]; + for combination in 0..729 { + let mut digits = combination; + let original = Tuple::new( + None, + (0..6) + .map(|_| { + let value = values[digits % 3].clone(); + digits /= 3; + value + }) + .collect(), + ); + let mut guarded = Tuple::new(None, original.values[..3].to_vec()); + for index in 0..3 { + let skip = *arena + .expression(guards[index]) + .eval(&arena, Some(&guarded))? + == DataValue::Boolean(false); + skipped[index] += usize::from(skip); + guarded.values.push(if skip { + DataValue::Null + } else { + original.values[index + 3].clone() + }); + } + assert_eq!( + arena.expression(predicate).eval(&arena, Some(&original))?, + arena.expression(predicate).eval(&arena, Some(&guarded))?, + "combination {combination}" + ); + } + assert!(skipped[1] > 0 && skipped[2] > 0); + Ok(()) + } + + #[test] + fn unsafe_sibling_is_not_evaluated_before_a_matching_marker() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut fixture = Fixture::new(&mut arena, BinaryOperator::Or, true)?; + let value = DataValue::from("not a boolean".to_string()); + let evaluator = cast_create(&value.logical_type(), &LogicalType::Boolean)?; + let expr = arena.alloc_expression(value.into()); + let failing = arena.alloc_expression(ScalarExpression::TypeCast { + expr, + ty: LogicalType::Boolean, + evaluator: Some(evaluator), + }); + let predicate = fixture.predicate(); + let ScalarExpression::Binary { right_expr, .. } = &mut *arena.expression_mut(predicate) + else { + unreachable!() + }; + *right_expr = failing; + assert!(!GuardMarkApply.apply(&mut fixture.plan, &mut arena, None)?); + assert!(fixture.mark().guard.is_none()); + EvaluatorBind.apply(&mut fixture.plan, &mut arena, None)?; + let matched = Tuple::new( + None, + vec![DataValue::Boolean(false), DataValue::Boolean(true)], + ); + assert_eq!( + *arena.expression(predicate).eval(&arena, Some(&matched))?, + DataValue::Boolean(true) + ); + let unmatched = Tuple::new( + None, + vec![DataValue::Boolean(false), DataValue::Boolean(false)], + ); + assert!(matches!( + arena.expression(predicate).eval(&arena, Some(&unmatched)), + Err(DatabaseError::ParseBool(_)) + )); + Ok(()) + } + + #[test] + fn simple_guard_whitelist() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut column = |name: &str, ty, position| -> Result { + let column = arena.alloc_column(ColumnCatalog::new( + name.into(), + true, + ColumnDesc::new(ty, None, false, None)?, + )); + Ok(arena.alloc_expression(ScalarExpression::column_expr(column, position))) + }; + let flag = column("flag", LogicalType::Boolean, 0)?; + let number = column("number", LogicalType::Bigint, 1)?; + let zero = arena.alloc_expression(DataValue::Int64(0).into()); + let mut alloc = |expr| arena.alloc_expression(expr); + let cast = |expr, ty| ScalarExpression::TypeCast { + expr, + ty, + evaluator: None, + }; + let not = |expr, ty| ScalarExpression::Unary { + op: UnaryOperator::Not, + expr, + evaluator: None, + ty, + }; + let widening = alloc(cast(number, LogicalType::Double)); + let narrowing = alloc(cast(number, LogicalType::Integer)); + let divide = alloc(ScalarExpression::Binary { + op: BinaryOperator::Divide, + left_expr: number, + right_expr: zero, + evaluator: None, + ty: LogicalType::Bigint, + }); + let not_number = alloc(not(number, LogicalType::Bigint)); + let mut cases = vec![ + (flag, true), + (number, false), + (alloc(not(flag, LogicalType::Boolean)), true), + ( + alloc(ScalarExpression::Alias { + expr: flag, + alias: AliasType::Expr(number), + }), + true, + ), + ( + alloc(ScalarExpression::IsNull { + negated: true, + expr: number, + }), + true, + ), + // IS NULL only accepts a plain column or constant operand. + ( + alloc(ScalarExpression::IsNull { + negated: false, + expr: not_number, + }), + false, + ), + ]; + // Comparisons accept numeric operands and widening casts, never failing calculations. + for (operand, valid) in [ + (number, true), + (widening, true), + (narrowing, false), + (divide, false), + ] { + for op in [ + BinaryOperator::Eq, + BinaryOperator::NotEq, + BinaryOperator::Gt, + BinaryOperator::GtEq, + BinaryOperator::Lt, + BinaryOperator::LtEq, + ] { + let comparison = alloc(ScalarExpression::Binary { + op, + left_expr: operand, + right_expr: zero, + evaluator: None, + ty: LogicalType::Boolean, + }); + cases.push((comparison, valid)); + } + } + for (expr, valid) in cases { + let mut visitor = SimpleGuardVisitor { + valid: true, + context: GuardContext::Predicate, + }; + visitor.visit(expr, &arena)?; + assert_eq!(visitor.valid, valid); + } + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/optimizer/rule/normalization/join_reorder.rs b/src/optimizer/rule/normalization/join_reorder.rs new file mode 100644 index 00000000..4a1fd94f --- /dev/null +++ b/src/optimizer/rule/normalization/join_reorder.rs @@ -0,0 +1,1417 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::catalog::ColumnRef; +use crate::errors::DatabaseError; +use crate::expression::visitor_mut::ExprVisitorMut; +use crate::expression::ScalarExpression; +use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::optimizer::plan_utils::scan_column; +use crate::optimizer::rule::normalization::parameterized_index::{ + probe_index, probe_join, ProbeIndex, +}; +use crate::planner::operator::join::{JoinCondition, JoinOperator, JoinType}; +use crate::planner::operator::project::ProjectOperator; +use crate::planner::operator::{ + rows_work, InputCost, JoinImpl, Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption, + DEFAULT_TABLE_ROWS, +}; +use crate::planner::{Childrens, ExprRef, LogicalPlan, MetaArena, PlanArena}; +use crate::types::index::{IndexMeta, IndexMetaRef, IndexType}; +use crate::types::tuple::Schema; +use std::borrow::Borrow; +use std::mem::swap; +use std::ops::Range; + +pub(crate) struct GreedyJoinReorder; + +/// Whether a cross nested loop should drive from its right input. +fn nest_loop_swap(left: &InputCost, right: &InputCost) -> bool { + if left.rows == 0 { + false + } else if right.rows == 0 { + true + } else { + // Both directions emit the same output, so it is left out of the comparison. + let nest_loop = JoinImpl::NestLoop { full: false }; + nest_loop.work(right, left, 0, 0) < nest_loop.work(left, right, 0, 0) + } +} + +fn project_join_output( + join: LogicalPlan, + original_schema: Schema, + positions: impl Iterator + Clone, + arena: &mut PlanArena, +) -> LogicalPlan { + if positions.clone().eq(0..original_schema.len()) { + return join; + } + let exprs = original_schema + .into_iter() + .zip(positions) + .map(|(column, position)| { + arena.alloc_expression(ScalarExpression::column_expr(column, position)) + }) + .collect(); + let mut plan = LogicalPlan::new( + Operator::Project(ProjectOperator::new(exprs)), + Childrens::Only(Box::new(join)), + ); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + )); + plan +} + +// Prefer equijoins, then less work in this step, then the original input order. +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +struct JoinCandidate { + cross_join: bool, + work: i128, + rows: usize, + leaf: usize, + swapped: bool, + probe: bool, +} + +fn column_ndv>( + left: &LogicalPlan, + right: &LogicalPlan, + keys: impl IntoIterator, + loader: Option<&StatisticMetaLoader<'_>>, + arena: &PlanArena, +) -> Result, DatabaseError> { + // scan_column cannot trace columns through a join. + if matches!(&left.operator, Operator::Join(_)) || matches!(&right.operator, Operator::Join(_)) { + return Ok(None); + } + let Some(loader) = loader else { + return Ok(None); + }; + let (Some(left_rows), Some(right_rows)) = (left.estimated_rows(), right.estimated_rows()) + else { + return Ok(None); + }; + let mut result = None; + let mut divisor = 0; + for key in keys { + let &(left_key, right_key) = key.borrow(); + let ( + ScalarExpression::ColumnRef { column: l, .. }, + ScalarExpression::ColumnRef { column: r, .. }, + ) = (arena.expression(left_key), arena.expression(right_key)) + else { + continue; + }; + for (l, r) in [(*l, *r), (*r, *l)] { + let Some((left_table, left_column)) = scan_column(left, l, arena)? else { + continue; + }; + let Some((right_table, right_column)) = scan_column(right, r, arena)? else { + continue; + }; + let Some(left_meta) = loader.load_column(left_table, left_column) else { + continue; + }; + let Some(right_meta) = loader.load_column(right_table, right_column) else { + continue; + }; + let left_ndv = left_meta.distinct_values_len().min(left_rows); + let right_ndv = right_meta.distinct_values_len().min(right_rows); + if left_ndv == 0 || right_ndv == 0 { + return Ok(Some((left_ndv, right_ndv))); + } + // Do not multiply selection rates for potentially correlated keys. + let candidate = left_ndv.max(right_ndv); + if candidate > divisor { + divisor = candidate; + result = Some((left_ndv, right_ndv)); + } + } + } + Ok(result) +} + +fn index_statistics(plan: &LogicalPlan) -> Option<(IndexMetaRef, Option)> { + let option = plan.row_count_option()?; + let (PlanImpl::IndexScan(index), RowCountOption::Estimated { ndv, .. }) = + (&option.plan, &option.row_count) + else { + return None; + }; + Some((index.meta, *ndv)) +} + +struct JoinInputStatistics<'a> { + rows: usize, + schema: &'a [ColumnRef], + index: Option<(&'a IndexMeta, Option)>, +} + +fn index_column(column: ColumnRef, index: &IndexMeta, arena: &PlanArena) -> bool { + let column = arena.column(column); + column.table_name() == Some(&index.table_name) + && column.id().is_some_and(|id| index.column_ids.contains(&id)) +} + +fn covers_index>( + index: &IndexMeta, + other: &JoinInputStatistics, + keys: impl Fn() -> I, + require_other_index: bool, + arena: &PlanArena, +) -> bool { + index.column_ids.iter().all(|id| { + keys().any(|(left, right)| { + let ( + ScalarExpression::ColumnRef { column: left, .. }, + ScalarExpression::ColumnRef { column: right, .. }, + ) = (arena.expression(left), arena.expression(right)) + else { + return false; + }; + [(*left, *right), (*right, *left)] + .iter() + .any(|&(own, opposite)| { + index_column(own, index, arena) + && arena.column(own).id() == Some(*id) + && other + .schema + .iter() + .any(|column| arena.same_column(*column, opposite)) + && (!require_other_index + || other + .index + .is_some_and(|(index, _)| index_column(opposite, index, arena))) + }) + }) + }) +} + +#[derive(Clone, Copy)] +struct JoinInputEstimate { + rows: usize, + ndv: Option, + unique: bool, +} + +fn estimate_join_rows(left: JoinInputEstimate, right: JoinInputEstimate) -> usize { + // J = min(ceil(L * R / D), left_unique ? R : inf, right_unique ? L : inf). + // L/R are rows, N_L/N_R are NDVs; D = max(min(N_L, L), min(N_R, R), 1). + // Missing NDV means D = 1; either known NDV of zero means J = 0. + let mut rows = left.rows.saturating_mul(right.rows); + if left.ndv == Some(0) || right.ndv == Some(0) { + return 0; + } + if let (Some(left_ndv), Some(right_ndv)) = (left.ndv, right.ndv) { + let ndv = left_ndv + .min(left.rows) + .max(right_ndv.min(right.rows)) + .max(1); + rows = ((left.rows as u128 * right.rows as u128).div_ceil(ndv as u128)) + .min(usize::MAX as u128) as usize; + } + if left.unique { + rows = rows.min(right.rows); + } + if right.unique { + rows = rows.min(left.rows); + } + rows +} + +fn join_rows>( + left: &JoinInputStatistics, + right: &JoinInputStatistics, + keys: impl Fn() -> I, + ndvs: Option<(usize, usize)>, + arena: &PlanArena, +) -> usize { + let mut left_estimate = JoinInputEstimate { + rows: left.rows, + ndv: None, + unique: false, + }; + let mut right_estimate = JoinInputEstimate { + rows: right.rows, + ndv: None, + unique: false, + }; + if let (Some((left_index, Some(left_ndv))), Some((right_index, Some(right_ndv)))) = + (left.index, right.index) + { + if covers_index(left_index, right, &keys, true, arena) + && covers_index(right_index, left, &keys, true, arena) + { + left_estimate.ndv = Some(left_ndv); + right_estimate.ndv = Some(right_ndv); + } + } + if left_estimate.ndv.is_none() || right_estimate.ndv.is_none() { + if let Some((left_ndv, right_ndv)) = ndvs { + left_estimate.ndv = Some(left_ndv); + right_estimate.ndv = Some(right_ndv); + } + } + for (own, other, estimate) in [ + (left, right, &mut left_estimate), + (right, left, &mut right_estimate), + ] { + if let Some((index, _)) = own.index { + if matches!(index.ty, IndexType::PrimaryKey { .. } | IndexType::Unique) + && covers_index(index, other, &keys, false, arena) + { + estimate.unique = true; + } + } + } + // TODO: Estimate expression keys and intermediate Join key distributions; + // do not inherit uniqueness through joins that can duplicate an input row. + estimate_join_rows(left_estimate, right_estimate) +} + +// A region only spans joins sharing the `FORCE_NEST_LOOP_JOIN` setting of its root. +fn can_expand(plan: &LogicalPlan, forced: bool) -> bool { + let Operator::Join(JoinOperator { + join_type, + on, + force_nested_loop, + limit_pushed: false, + }) = &plan.operator + else { + return false; + }; + *force_nested_loop == forced + && (matches!((join_type, on), (JoinType::Cross, JoinCondition::None)) + || matches!((join_type, on), + (JoinType::Inner, JoinCondition::On { on, filter: None }) if !on.is_empty())) +} + +// Counts a region's inputs. +fn region_inputs(plan: &LogicalPlan, forced: bool) -> usize { + if !can_expand(plan, forced) { + return 1; + } + let (left, right) = plan.childrens.twins(); + region_inputs(left, forced) + region_inputs(right, forced) +} + +// Collects input schemas and join keys; returns whether every input has a row estimate. +// Joins inside the region are rebuilt by the reorder, so their own rows are not estimated here. +fn inspect<'plan>( + plan: &'plan mut LogicalPlan, + schemas: &mut Vec<&'plan Schema>, + keys: &mut JoinKeys, + arena: &mut PlanArena, + forced: bool, +) -> Result { + if !can_expand(plan, forced) { + let estimated = plan.estimated_rows().is_some(); + schemas.push(plan.output_schema(arena)); + return Ok(estimated); + } + let (Operator::Join(join), Childrens::Twins { left, right }) = + (&plan.operator, plan.childrens.as_mut()) + else { + unreachable!() + }; + let on = match &join.on { + JoinCondition::On { on, .. } => on.as_slice(), + JoinCondition::None => &[], + }; + let start = schemas.len(); + let left_estimated = inspect(left, schemas, keys, arena, forced)?; + let middle = schemas.len(); + let right_estimated = inspect(right, schemas, keys, arena, forced)?; + + for &(left, right) in on { + let mut endpoints = [None, None]; + for (endpoint, (expr, range)) in endpoints + .iter_mut() + .zip([(left, start..middle), (right, middle..schemas.len())]) + { + for index in range { + if key_in_schema(expr, schemas[index], arena)? { + *endpoint = Some(index); + break; + } + } + if endpoint.is_none() { + return Err(DatabaseError::InvalidPlan(format!( + "join key `{}` is not produced by any join input", + expr.output_name(arena) + ))); + } + } + if endpoints[0] == endpoints[1] { + return Err(DatabaseError::InvalidPlan(format!( + "join key `{} = {}` compares columns of the same join input", + left.output_name(arena), + right.output_name(arena) + ))); + } + // SAFETY: The loop above returns early if either endpoint is None. + let [l, r] = endpoints.map(Option::unwrap); + keys.insert(1u64 << l, 1u64 << r, (left, right)); + } + Ok(left_estimated && right_estimated) +} + +/// An input's cost, read once at its (default) rows when its estimate is incomplete. +fn input_cost_or_default( + input: &mut LogicalPlan, + rows: usize, + arena: &mut PlanArena, +) -> Result { + Ok(match input.price(arena).and_then(|cost| cost.input()) { + Some(cost) => cost, + None => { + let width = input.output_schema(arena).len(); + InputCost { + rows, + width, + work: rows_work(rows, width), + } + } + }) +} + +/// Collects a region's inputs with their costs; a leaf's cost never changes. `start` tracks +/// the smallest input (cheapest to read on ties, earliest on full ties). +fn flatten( + mut plan: LogicalPlan, + leaves: &mut Vec<(LogicalPlan, InputCost, u64)>, + start: &mut usize, + forced: bool, + arena: &mut PlanArena, +) -> Result<(), DatabaseError> { + if can_expand(&plan, forced) { + let (left, right) = plan.childrens.pop_twins(); + flatten(left, leaves, start, forced, arena)?; + flatten(right, leaves, start, forced, arena)?; + } else { + // Without statistics an input is assumed to have the default rows (RFC 0001 Ā§3.9). + let rows = plan.estimated_rows().unwrap_or(DEFAULT_TABLE_ROWS); + let cost = input_cost_or_default(&mut plan, rows, arena)?; + if leaves + .get(*start) + .is_some_and(|(_, best, _)| (cost.rows, cost.work) < (best.rows, best.work)) + { + *start = leaves.len(); + } + let mask = 1u64 << leaves.len(); + leaves.push((plan, cost, mask)); + } + Ok(()) +} + +fn key_in_schema(expr: ExprRef, schema: &Schema, arena: &PlanArena) -> Result { + expr.all_referenced_columns(arena, |arena, column, _| { + schema + .iter() + .any(|candidate| arena.same_column(*candidate, *column)) + }) +} + +/// Join keys between original inputs, each edge stored once per endpoint as +/// `(endpoint bit, other bit, endpoint key, other key)`. Sorted by endpoint bit, so an input's +/// edges are contiguous; equal endpoints keep insertion order. +#[derive(Default)] +struct JoinKeys(Vec<(u64, u64, ExprRef, ExprRef)>); + +impl JoinKeys { + fn insert(&mut self, left: u64, right: u64, (l, r): (ExprRef, ExprRef)) { + for entry in [(left, right, l, r), (right, left, r, l)] { + let at = self.0.partition_point(|e| e.0 <= entry.0); + self.0.insert(at, entry); + } + } + + fn range(&self, leaf: u64) -> Range { + let start = self.0.partition_point(|e| e.0 < leaf); + let len = self.0[start..].iter().take_while(|e| e.0 == leaf).count(); + start..start + len + } + + fn get(&self, current: u64, leaf: u64) -> impl Iterator + '_ { + self.0[self.range(leaf)] + .iter() + .filter(move |e| e.1 & current != 0) + .map(|&(_, _, leaf_key, other_key)| (other_key, leaf_key)) + } + + fn take(&mut self, current: u64, leaf: u64) -> Vec<(ExprRef, ExprRef)> { + let range = self.range(leaf); + self.0 + .drain(range) + .filter(|e| e.1 & current != 0) + .map(|(_, _, leaf_key, other_key)| (other_key, leaf_key)) + .collect() + } +} + +struct BindPosition<'a>(&'a Schema); + +impl ExprVisitorMut for BindPosition<'_> { + fn visit_column_ref( + &mut self, + column: &mut ColumnRef, + position: &mut usize, + _: &mut A, + ) -> Result<(), DatabaseError> { + // Join keys keep the refs their inputs output. + *position = self + .0 + .iter() + .position(|candidate| candidate == column) + .ok_or(DatabaseError::SourceNotFound)?; + Ok(()) + } +} + +impl NormalizationRule for GreedyJoinReorder { + fn apply( + &self, + plan: &mut LogicalPlan, + arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, + ) -> Result { + if plan + .physical_option + .as_ref() + .is_some_and(|option| option.join_decided) + { + return Ok(false); + } + // `FORCE_NEST_LOOP_JOIN` regions are reordered too, but only with nested loops. + let forced = matches!( + &plan.operator, + Operator::Join(JoinOperator { + force_nested_loop: true, + .. + }) + ); + if !can_expand(plan, forced) { + return Ok(false); + } + let input_count = region_inputs(plan, forced); + if input_count > u64::BITS as usize { + return Ok(false); + } + let original_schema = plan.output_schema(arena).clone(); + let mut schemas = Vec::new(); + let mut keys = JoinKeys::default(); + // Inputs without statistics are priced with default rows, but joins built from them keep + // an unknown row count for later rules. + let estimated = inspect(plan, &mut schemas, &mut keys, arena, forced)?; + // Each original column keeps its input bit and position within its current join. + let mut positions = Vec::with_capacity(original_schema.len()); + for (input, schema) in schemas.iter().enumerate() { + positions.extend((0..schema.len()).map(|position| (1u64 << input, position))); + } + let mut leaves = Vec::with_capacity(schemas.len()); + // Left-deep: start from the smallest input, then join one remaining leaf per step, so + // every leaf can still be probed (RFC 0001 Ā§3.2). + let mut start = 0; + flatten(plan.take(), &mut leaves, &mut start, forced, arena)?; + let mut probe_keys = Vec::new(); + let (mut current, mut current_cost, mut current_mask) = leaves.remove(start); + while !leaves.is_empty() { + let mut best: Option<(JoinCandidate, Option)> = None; + // A result of at most one row may take a Cartesian product before an equijoin. + let single = current_cost.rows <= 1; + for (k, (leaf, leaf_cost, leaf_mask)) in leaves.iter_mut().enumerate() { + let leaf_cost = *leaf_cost; + // Keys run from `current` to `leaf`, matching the unswapped join sides. + let candidate_keys = || keys.get(current_mask, *leaf_mask); + let product = { + let ndvs = column_ndv(¤t, leaf, candidate_keys(), loader, arena)?; + let current_stats = index_statistics(¤t); + let leaf_stats = index_statistics(leaf); + let current_schema = current.output_schema(arena); + let leaf_schema = leaf.output_schema(arena); + join_rows( + &JoinInputStatistics { + rows: current_cost.rows, + schema: current_schema, + index: current_stats.map(|(index, ndv)| (arena.index(index), ndv)), + }, + &JoinInputStatistics { + rows: leaf_cost.rows, + schema: leaf_schema, + index: leaf_stats.map(|(index, ndv)| (arena.index(index), ndv)), + }, + candidate_keys, + ndvs, + arena, + ) + }; + let connected = candidate_keys().next().is_some(); + let hash = connected && !forced; + // Every input is read once in any order, so only this step's work is compared. + // A probe never reads its inner input, so its step can be negative. + let step = + |implementation: JoinImpl, outer: &InputCost, inner: &InputCost, rows| { + implementation.work(outer, inner, rows, outer.width + inner.width) as i128 + - outer.work as i128 + - inner.work as i128 + }; + // (work, rows, swapped, probe index). Unswapped, `current` is the hash build side + // and the nested loop / probe outer side. + let (work, rows, swapped, probe) = if hash { + let keep = step(JoinImpl::Hash, ¤t_cost, &leaf_cost, product); + let swap = step(JoinImpl::Hash, &leaf_cost, ¤t_cost, product); + let mut shape = (keep.min(swap), product, swap < keep, None); + // Only a leaf can be probed: `leaf` always, `current` before its first join. + for swapped in [false, true] { + let (inner, outer_cost, inner_cost) = if swapped { + (¤t, &leaf_cost, ¤t_cost) + } else { + (&*leaf, ¤t_cost, &leaf_cost) + }; + // Probe keys run from the outer value to the probed column. + probe_keys.clear(); + probe_keys.extend(candidate_keys().map(|(c, l)| { + if swapped { + (l, c) + } else { + (c, l) + } + })); + let key_count = probe_keys.len(); + let Some(index) = + probe_index(inner, &mut probe_keys, key_count, arena, loader)? + else { + continue; + }; + // A unique probe emits at most one row per outer row. + let rows = product.min(outer_cost.rows); + let lookup_work = index.lookup_work; + let work = step( + JoinImpl::IndexProbe { lookup_work }, + outer_cost, + inner_cost, + rows, + ); + if work <= shape.0 { + shape = (work, rows, swapped, Some(index)); + } + } + shape + } else { + let swapped = nest_loop_swap(¤t_cost, &leaf_cost); + let (outer, inner) = if swapped { + (&leaf_cost, ¤t_cost) + } else { + (¤t_cost, &leaf_cost) + }; + let work = step(JoinImpl::NestLoop { full: false }, outer, inner, product); + (work, product, swapped, None) + }; + let candidate = JoinCandidate { + cross_join: !connected && !single, + work, + rows, + leaf: k, + swapped, + probe: probe.is_some(), + }; + if best + .as_ref() + .is_none_or(|(previous, _)| candidate < *previous) + { + best = Some((candidate, probe)); + } + } + // SAFETY: `leaves` is not empty, so the loop above chose a candidate. + let ( + JoinCandidate { + rows: product, + leaf: k, + swapped, + .. + }, + probe, + ) = best.unwrap(); + let (leaf, _, leaf_mask) = leaves.remove(k); + let mut on = keys.take(current_mask, leaf_mask); + let (mut left, mut right, right_mask) = if swapped { + for (c, l) in &mut on { + swap(c, l); + } + (leaf, current, current_mask) + } else { + (current, leaf, leaf_mask) + }; + let left_schema = left.output_schema(arena); + let right_schema = right.output_schema(arena); + for (input, position) in &mut positions { + if *input & right_mask != 0 { + *position += left_schema.len(); + } + } + for (l, r) in &mut on { + *l = l.clone_expression(arena)?; + *r = r.clone_expression(arena)?; + BindPosition(left_schema).visit(l, arena)?; + BindPosition(right_schema).visit(r, arena)?; + } + let row_count = if estimated { + RowCountOption::Estimated { + rows: product, + ndv: None, + } + } else { + RowCountOption::Unknown + }; + current = if let Some(mut index) = probe { + // Probe values were chosen before binding; they come from the outer (left) input. + for value in &mut index.probe { + *value = value.clone_expression(arena)?; + BindPosition(left_schema).visit(value, arena)?; + } + probe_join(left, right, on, index, row_count, arena)? + } else { + let (condition, join_type) = if on.is_empty() { + (JoinCondition::None, JoinType::Cross) + } else { + (JoinCondition::On { on, filter: None }, JoinType::Inner) + }; + let mut join = JoinOperator::build(left, right, condition, join_type, forced); + let Operator::Join(op) = &join.operator else { + unreachable!() + }; + let implementation = op.plan_impl(); + let mut option = PhysicalOption::new(implementation, SortOption::None, row_count); + option.join_decided = true; + join.physical_option = Some(option); + join + }; + current_cost = input_cost_or_default(&mut current, product, arena)?; + current_mask |= leaf_mask; + } + *plan = project_join_output( + current, + original_schema, + positions.iter().map(|(_, position)| *position), + arena, + ); + Ok(true) + } +} + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::expression::BinaryOperator; + use crate::optimizer::core::histogram::{HistogramMetaBuilder, DEFAULT_HLL_PRECISION}; + use crate::planner::operator::sort::{SortField, SortOperator}; + use crate::planner::operator::table_scan::TableScanOperator; + use crate::planner::TableArenaCell; + use crate::storage::StatisticsMetaCache; + use crate::types::index::IndexInfo; + use crate::types::value::DataValue; + use crate::types::LogicalType; + + fn input(name: &str, rows: usize, arena: &mut PlanArena) -> (LogicalPlan, ExprRef) { + let mut column = ColumnCatalog::new( + name.into(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), + ); + column.set_ref_table(name.into(), 0, false); + let column = arena.alloc_column(column); + let key = arena.alloc_expression(ScalarExpression::column_expr(column, 0)); + let mut plan = LogicalPlan::new( + Operator::TableScan(TableScanOperator { + table_name: name.into(), + columns: vec![column], + limit: (None, None), + index_infos: Vec::new(), + with_pk: false, + }), + Childrens::None, + ); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { rows, ndv: None }, + )); + (plan, key) + } + + fn indexed_input( + name: &str, + rows: usize, + ndv: Option, + arena: &mut PlanArena, + ) -> (LogicalPlan, ExprRef, IndexMetaRef) { + let (mut plan, key) = input(name, rows, arena); + let meta = arena.alloc_index(IndexMeta { + id: 1, + table_name: name.into(), + name: format!("idx_{name}"), + column_ids: vec![0], + ty: IndexType::Normal, + pk_ty: LogicalType::Integer, + value_ty: LogicalType::Integer, + }); + let info = IndexInfo { + meta, + sort_option: SortOption::None, + lookup: None, + residual_predicate: None, + covered_deserializers: None, + cover_mapping: None, + sort_elimination_hint: None, + stream_aggregate_hint: None, + }; + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::IndexScan(Box::new(info)), + SortOption::None, + RowCountOption::Estimated { rows, ndv }, + )); + (plan, key, meta) + } + + #[test] + fn malformed_regions_are_planning_errors() { + // (outer join key, expected message). Each side of a key is only searched in its own + // input range, so a left column used on the right is reported as not produced. + type OuterKey = fn(ExprRef, ExprRef, &mut PlanArena) -> (ExprRef, ExprRef); + let cases: [(OuterKey, &str); 2] = [ + (|a, _, _| (a, a), "is not produced by any join input"), + ( + |_, c, arena| (input("z", 1, arena).1, c), + "is not produced by any join input", + ), + ]; + for (outer_key, message) in cases { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (a, ak) = input("a", 1, &mut arena); + let (b, bk) = input("b", 1, &mut arena); + let (c, ck) = input("c", 1, &mut arena); + let ab = JoinOperator::build( + a, + b, + JoinCondition::On { + on: vec![(ak, bk)], + filter: None, + }, + JoinType::Inner, + false, + ); + let mut plan = JoinOperator::build( + ab, + c, + JoinCondition::On { + on: vec![outer_key(ak, ck, &mut arena)], + filter: None, + }, + JoinType::Inner, + false, + ); + let Err(DatabaseError::InvalidPlan(error)) = + GreedyJoinReorder.apply(&mut plan, &mut arena, None) + else { + panic!("expected invalid plan: {message}"); + }; + assert!(error.contains(message), "{error}"); + } + } + + #[test] + fn input_masks_support_64_inputs_and_skip_larger_regions() -> Result<(), DatabaseError> { + for count in [64, 65] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (mut plan, mut previous) = input("t0", 1, &mut arena); + for index in 1..count { + let (next, key) = input(&format!("t{index}"), 1, &mut arena); + plan = JoinOperator::build( + plan, + next, + JoinCondition::On { + on: vec![(previous, key)], + filter: None, + }, + JoinType::Inner, + false, + ); + previous = key; + } + let schema = plan.output_schema(&mut arena).clone(); + let original = plan.clone(); + assert_eq!( + GreedyJoinReorder.apply(&mut plan, &mut arena, None)?, + count == 64 + ); + assert_eq!(plan.output_schema(&mut arena), &schema); + if count == 65 { + assert_eq!(plan, original); + } + } + Ok(()) + } + + #[test] + fn column_ndv_is_read_from_scan_without_an_index() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (mut left, l) = input("l", 1000, &mut arena); + let (right, r) = input("r", 200, &mut arena); + let mut cache = StatisticsMetaCache::default(); + for (table, rows) in [("l", 1000), ("r", 200)] { + let mut builder = HistogramMetaBuilder::new(DEFAULT_HLL_PRECISION)?; + for value in 0..rows { + builder.append(&DataValue::Int32(value % 100)); + } + cache.columns.insert((table.into(), 0), builder.build()); + } + let loader = StatisticMetaLoader::new(&cache); + assert_eq!( + column_ndv(&left, &right, [(l, r)], Some(&loader), &arena)?, + Some((100, 100)) + ); + assert_eq!(column_ndv(&left, &right, [(l, r)], None, &arena)?, None); + assert_eq!( + column_ndv( + &left, + &right, + [] as [(ExprRef, ExprRef); 0], + Some(&loader), + &arena + )?, + None + ); + left.physical_option.as_mut().unwrap().row_count = + RowCountOption::Estimated { rows: 8, ndv: None }; + assert_eq!( + column_ndv(&left, &right, [(l, r)], Some(&loader), &arena)?, + Some((8, 100)) + ); + left.physical_option.as_mut().unwrap().row_count = RowCountOption::Estimated { + rows: 1000, + ndv: None, + }; + let mut nulls = HistogramMetaBuilder::new(DEFAULT_HLL_PRECISION)?; + nulls.append(&DataValue::Null); + cache.columns.insert(("r".into(), 0), nulls.build()); + assert_eq!( + column_ndv( + &left, + &right, + [(l, r)], + Some(&StatisticMetaLoader::new(&cache)), + &arena + )?, + Some((100, 0)) + ); + cache.columns.remove(&("r".into(), 0)); + assert_eq!( + column_ndv( + &left, + &right, + [(l, r)], + Some(&StatisticMetaLoader::new(&cache)), + &arena + )?, + None + ); + Ok(()) + } + + #[test] + fn complete_index_keys_reduce_rows_and_bound_uniqueness() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (mut left, l, li) = indexed_input("l", 1000, Some(100), &mut arena); + let (mut right, r, ri) = indexed_input("r", 200, Some(100), &mut arena); + let one = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(1))); + let expr_key = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Plus, + left_expr: l, + right_expr: one, + evaluator: None, + ty: LogicalType::Integer, + }); + let left_schema = left.output_schema(&mut arena); + let right_schema = right.output_schema(&mut arena); + let l_input = JoinInputStatistics { + rows: 1000, + schema: left_schema, + index: Some((arena.index(li), Some(100))), + }; + let r_input = JoinInputStatistics { + rows: 200, + schema: right_schema, + index: Some((arena.index(ri), Some(100))), + }; + assert_eq!( + join_rows(&l_input, &r_input, || std::iter::once((l, r)), None, &arena), + 2000 + ); + assert_eq!( + join_rows(&l_input, &r_input, std::iter::empty, None, &arena), + 200000 + ); + let mut unique_right = arena.index(ri).clone(); + unique_right.ty = IndexType::Unique; + let r_input = JoinInputStatistics { + rows: 200, + schema: right_schema, + index: Some((&unique_right, None)), + }; + assert_eq!( + join_rows(&l_input, &r_input, || std::iter::once((l, r)), None, &arena), + 1000 + ); + let mut unique_left = arena.index(li).clone(); + unique_left.ty = IndexType::PrimaryKey { is_multiple: false }; + let l_input = JoinInputStatistics { + rows: 1000, + schema: left_schema, + index: Some((&unique_left, None)), + }; + assert_eq!( + join_rows(&l_input, &r_input, || std::iter::once((l, r)), None, &arena), + 200 + ); + unique_right.column_ids.push(1); + let r_input = JoinInputStatistics { + rows: 200, + schema: right_schema, + index: Some((&unique_right, Some(100))), + }; + let l_input = JoinInputStatistics { + rows: 1000, + schema: left_schema, + index: Some((arena.index(li), Some(100))), + }; + assert_eq!( + join_rows(&l_input, &r_input, || std::iter::once((l, r)), None, &arena), + 200000 + ); + let r_input = JoinInputStatistics { + rows: 200, + schema: right_schema, + index: Some((arena.index(ri), Some(100))), + }; + assert_eq!( + join_rows( + &l_input, + &r_input, + || std::iter::once((expr_key, r)), + None, + &arena + ), + 200000 + ); + Ok(()) + } + + #[test] + fn composite_ndv_requires_the_whole_key() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (_left, l, li) = indexed_input("l", 1000, Some(100), &mut arena); + let (_right, r, ri) = indexed_input("r", 200, Some(100), &mut arena); + let mut lc = ColumnCatalog::new( + "l_second".into(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + ); + lc.set_ref_table("l".into(), 1, false); + let lc = arena.alloc_column(lc); + let mut rc = ColumnCatalog::new( + "r_second".into(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + ); + rc.set_ref_table("r".into(), 1, false); + let rc = arena.alloc_column(rc); + let l2 = arena.alloc_expression(ScalarExpression::column_expr(lc, 1)); + let r2 = arena.alloc_expression(ScalarExpression::column_expr(rc, 1)); + let mut l_index = arena.index(li).clone(); + let mut r_index = arena.index(ri).clone(); + for index in [&mut l_index, &mut r_index] { + index.column_ids.push(1); + index.ty = IndexType::Composite; + index.value_ty = LogicalType::Tuple(vec![LogicalType::Integer, LogicalType::Integer]); + } + let ScalarExpression::ColumnRef { column: l_col, .. } = arena.expression(l) else { + unreachable!() + }; + let ScalarExpression::ColumnRef { column: r_col, .. } = arena.expression(r) else { + unreachable!() + }; + let l_schema = [*l_col, lc]; + let r_schema = [*r_col, rc]; + let l_input = JoinInputStatistics { + rows: 1000, + schema: &l_schema, + index: Some((&l_index, Some(100))), + }; + let r_input = JoinInputStatistics { + rows: 200, + schema: &r_schema, + index: Some((&r_index, Some(100))), + }; + assert_eq!( + join_rows(&l_input, &r_input, || std::iter::once((l, r)), None, &arena), + 200000 + ); + assert_eq!( + join_rows( + &l_input, + &r_input, + || [(l, r), (l2, r2)].into_iter(), + None, + &arena + ), + 2000 + ); + Ok(()) + } + + fn materialized_right(arena: &mut PlanArena) -> LogicalPlan { + let (b, b_key) = input("b", 10, arena); + let (c, c_key) = input("c", 20, arena); + let mut join = JoinOperator::build( + b, + c, + JoinCondition::On { + on: vec![(b_key, c_key)], + filter: None, + }, + JoinType::Inner, + false, + ); + join.physical_option = Some(PhysicalOption::new( + PlanImpl::HashJoin, + SortOption::None, + RowCountOption::Estimated { + rows: 200, + ndv: None, + }, + )); + let schema = join.output_schema(arena); + let exprs = schema + .iter() + .enumerate() + .map(|(position, column)| { + arena.alloc_expression(ScalarExpression::column_expr(*column, position)) + }) + .collect(); + let mut project = LogicalPlan::new( + Operator::Project(ProjectOperator::new(exprs)), + Childrens::Only(Box::new(join)), + ); + project.physical_option = Some(PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + )); + project + } + + #[test] + fn unknown_startup_does_not_become_zero() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let (mut scan, key) = input("source", 16, &mut arena); + scan.physical_option.as_mut().unwrap().row_count = RowCountOption::Unknown; + let mut sort = LogicalPlan::new( + Operator::Sort(SortOperator { + sort_fields: vec![SortField { + expr: key, + asc: true, + nulls_first: false, + }], + }), + Childrens::Only(Box::new(scan)), + ); + sort.physical_option = Some(PhysicalOption::new( + PlanImpl::Sort, + SortOption::None, + RowCountOption::Follow, + )); + assert_eq!(sort.price(&mut arena).and_then(|cost| cost.startup), None); + assert_eq!( + sort.price(&mut arena) + .and_then(|cost| cost.input()) + .map(|cost| cost.work), + None + ); + + Ok(()) + } + + #[test] + fn nested_loop_moves_materialization_left_and_restores_positions() -> Result<(), DatabaseError> + { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let (mut left, _) = input("a", 100, &mut arena); + let mut right = materialized_right(&mut arena); + assert_eq!( + left.price(&mut arena).and_then(|cost| cost.startup), + Some(0) + ); + // The 10-row build side is read (50), stored whole (50) and hashed. + assert_eq!( + right.price(&mut arena).and_then(|cost| cost.startup), + Some(100 + 10 * crate::planner::operator::HASH_BUILD_WORK) + ); + let mut plan = + JoinOperator::build(left, right, JoinCondition::None, JoinType::Cross, false); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::NestLoopJoin, + SortOption::None, + RowCountOption::Estimated { + rows: 20000, + ndv: None, + }, + )); + let schema = plan.output_schema(&mut arena).clone(); + assert!(GreedyJoinReorder.apply(&mut plan, &mut arena, None)?); + assert_eq!(plan.output_schema(&mut arena), &schema); + assert_eq!(plan.estimated_rows(), Some(20000)); + let Operator::Project(project) = &plan.operator else { + unreachable!() + }; + for (expr, position) in project.exprs.iter().zip([2, 0, 1]) { + assert!(matches!( + arena.expression(*expr), + ScalarExpression::ColumnRef { position: actual, .. } if *actual == position + )); + } + let join = plan.childrens.only(); + let (left, right) = join.childrens.twins(); + assert_eq!(left.referenced_table(), vec!["b".into(), "c".into()]); + assert_eq!(right.referenced_table(), vec!["a".into()]); + let mut join = join.clone(); + assert!(!GreedyJoinReorder.apply(&mut join, &mut arena, None)?); + Ok(()) + } + + fn three_inputs(arena: &mut PlanArena) -> LogicalPlan { + let (a, a_key) = input("a", 1000, arena); + let (b, b_key) = input("b", 10, arena); + let (c, c_key) = input("c", 20, arena); + let ab = JoinOperator::build( + a, + b, + JoinCondition::On { + on: vec![(a_key, b_key)], + filter: None, + }, + JoinType::Inner, + false, + ); + let b_column = b_key.output_column_ref(arena); + let b_key = arena.alloc_expression(ScalarExpression::column_expr(b_column, 1)); + JoinOperator::build( + ab, + c, + JoinCondition::On { + on: vec![(b_key, c_key)], + filter: None, + }, + JoinType::Inner, + false, + ) + } + + #[test] + fn smallest_connected_product_preserves_schema_and_key_positions() -> Result<(), DatabaseError> + { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let mut plan = three_inputs(&mut arena); + let schema = plan.output_schema(&mut arena).clone(); + assert!(GreedyJoinReorder.apply(&mut plan, &mut arena, None)?); + assert_eq!(plan.output_schema(&mut arena), &schema); + assert_eq!(plan.estimated_rows(), Some(200000)); + // BC (200 rows) is merged first and builds against A (1000 rows). + let (bc, a) = plan.childrens.only().childrens.twins(); + assert_eq!(a.referenced_table(), vec!["a".into()]); + assert_eq!(bc.referenced_table(), vec!["b".into(), "c".into()]); + assert_eq!(bc.estimated_rows(), Some(200)); + let Operator::Join(op) = &bc.operator else { + unreachable!() + }; + let JoinCondition::On { on, .. } = &op.on else { + unreachable!() + }; + for expr in [on[0].0, on[0].1] { + assert!(matches!( + arena.expression(expr), + ScalarExpression::ColumnRef { position: 0, .. } + )); + } + Ok(()) + } + + #[test] + fn nonadjacent_pair_restores_output_order_and_keeps_all_keys() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let (a, a_key) = input("a", 10, &mut arena); + let (b, b_key) = input("b", 1000, &mut arena); + let (c, c_key) = input("c", 20, &mut arena); + let ab = JoinOperator::build( + a, + b, + JoinCondition::On { + on: vec![(a_key, b_key)], + filter: None, + }, + JoinType::Inner, + false, + ); + let b_column = b_key.output_column_ref(&mut arena); + let b_key = arena.alloc_expression(ScalarExpression::column_expr(b_column, 1)); + let mut plan = JoinOperator::build( + ab, + c, + JoinCondition::On { + on: vec![(a_key, c_key), (b_key, c_key)], + filter: None, + }, + JoinType::Inner, + false, + ); + let schema = plan.output_schema(&mut arena).clone(); + assert!(GreedyJoinReorder.apply(&mut plan, &mut arena, None)?); + assert_eq!(plan.output_schema(&mut arena), &schema); + let Operator::Project(project) = &plan.operator else { + unreachable!() + }; + for (expr, expected) in project.exprs.iter().zip([0, 2, 1]) { + assert!( + matches!(arena.expression(*expr), ScalarExpression::ColumnRef { position, .. } if *position == expected) + ); + } + let join = plan.childrens.only(); + let (left, right) = join.childrens.twins(); + assert_eq!(left.referenced_table(), vec!["a".into(), "c".into()]); + assert_eq!(right.referenced_table(), vec!["b".into()]); + let Operator::Join(op) = &join.operator else { + unreachable!() + }; + let JoinCondition::On { on, .. } = &op.on else { + unreachable!() + }; + assert_eq!(on.len(), 2); + let Operator::Join(op) = &left.operator else { + unreachable!() + }; + let JoinCondition::On { on, .. } = &op.on else { + unreachable!() + }; + assert_eq!(on.len(), 1); + Ok(()) + } + + #[test] + fn unestimated_region_defers_cartesian_products_by_input_order() -> Result<(), DatabaseError> { + // FROM a, c, d, b WHERE a = b AND c = b: FROM order crosses a, c and d + // before any join key applies. Without estimates the region is priced with + // default rows, and only d stays a cross join. + for known_b in [false, true] { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let unknown = |name: &str, arena: &mut PlanArena| { + let (mut plan, key) = input(name, 10, arena); + if !(known_b && name == "b") { + plan.physical_option.as_mut().unwrap().row_count = RowCountOption::Unknown; + } + (plan, key) + }; + let (a, a_key) = unknown("a", &mut arena); + let (c, c_key) = unknown("c", &mut arena); + let (d, _) = unknown("d", &mut arena); + let (b, b_key) = unknown("b", &mut arena); + let ac = JoinOperator::build(a, c, JoinCondition::None, JoinType::Cross, false); + let acd = JoinOperator::build(ac, d, JoinCondition::None, JoinType::Cross, false); + let c_column = c_key.output_column_ref(&mut arena); + let c_key = arena.alloc_expression(ScalarExpression::column_expr(c_column, 1)); + let mut plan = JoinOperator::build( + acd, + b, + JoinCondition::On { + on: vec![(a_key, b_key), (c_key, b_key)], + filter: None, + }, + JoinType::Inner, + false, + ); + let schema = plan.output_schema(&mut arena).clone(); + assert!(GreedyJoinReorder.apply(&mut plan, &mut arena, None)?); + assert_eq!(plan.output_schema(&mut arena), &schema); + + let cross = plan.childrens.only(); + assert!(matches!( + &cross.operator, + Operator::Join(join) if join.join_type == JoinType::Cross + )); + let (abc, d) = cross.childrens.twins(); + assert_eq!(d.referenced_table(), vec!["d".into()]); + // Without estimates, the narrower C becomes the build side of AB. + let (c, ab) = abc.childrens.twins(); + assert_eq!(c.referenced_table(), vec!["c".into()]); + let mut tables = ab.referenced_table(); + tables.sort(); + assert_eq!(tables, vec!["a".into(), "b".into()]); + for join in [cross, abc, ab] { + assert_eq!(join.row_count_option(), None); + } + for join in [abc, ab] { + let Operator::Join(JoinOperator { + join_type: JoinType::Inner, + on: JoinCondition::On { on, .. }, + .. + }) = &join.operator + else { + unreachable!() + }; + assert_eq!(on.len(), 1); + } + } + Ok(()) + } + + #[test] + fn outer_joins_keep_original_order() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let mut plan = three_inputs(&mut arena); + if let Operator::Join(op) = &mut plan.operator { + op.join_type = JoinType::LeftOuter; + } + let original = plan.clone(); + assert!(!GreedyJoinReorder.apply(&mut plan, &mut arena, None)?); + assert_eq!(plan, original); + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/optimizer/rule/normalization/materialize_mark_apply.rs b/src/optimizer/rule/normalization/materialize_mark_apply.rs new file mode 100644 index 00000000..a49936e3 --- /dev/null +++ b/src/optimizer/rule/normalization/materialize_mark_apply.rs @@ -0,0 +1,198 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::NormalizationRule; +use crate::errors::DatabaseError; +use crate::expression::ScalarExpression; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::planner::operator::mark_apply::{MarkApplyKind, MarkApplyOperator}; +use crate::planner::operator::Operator; +use crate::planner::{Childrens, LogicalPlan, PlanArena}; + +pub(crate) struct MaterializeMarkApply; + +impl MaterializeMarkApply { + fn can_materialize(op: &MarkApplyOperator, arena: &PlanArena) -> bool { + // The first predicate compares values; additional predicates bind outer rows. + // The binder builds that comparison as `outer_operand op subquery_output`, + // so the RHS keys can be cached by evaluating its right side alone. + matches!(op.kind, MarkApplyKind::Quantified(_)) + && op.parameterized_probe().is_none() + && matches!( + op.predicates(), + [predicate] if matches!(arena.expression(*predicate), ScalarExpression::Binary { .. }) + ) + } +} + +impl NormalizationRule for MaterializeMarkApply { + fn apply( + &self, + plan: &mut LogicalPlan, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, + ) -> Result { + let (Operator::MarkApply(op), Childrens::Twins { left, right }) = + (&mut plan.operator, plan.childrens.as_ref()) + else { + return Ok(false); + }; + if op.materialize || !Self::can_materialize(op, arena) { + return Ok(false); + } + // Streaming rebuilds the RHS once per left row, which nested subqueries multiply, + // so materialize by default. Stream only when estimates show the cache would not + // pay off: no repeated left input, or fewer rebuilds than retained RHS rows. + let stream = match (left.estimated_rows(), right.estimated_rows()) { + (Some(left_rows), _) if left_rows <= 1 => true, + (Some(left_rows), Some(right_rows)) => left_rows < right_rows, + _ => false, + }; + if stream { + return Ok(false); + } + op.materialize = true; + Ok(true) + } +} + +// GRCOV_EXCL_START +#[cfg(test)] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::expression::BinaryOperator; + use crate::planner::operator::mark_apply::MarkApplyQuantifier; + use crate::planner::operator::values::ValuesOperator; + use crate::planner::operator::{PhysicalOption, PlanImpl, RowCountOption, SortOption}; + use crate::planner::TableArenaCell; + use crate::types::value::DataValue; + use crate::types::LogicalType; + + #[test] + fn selects_materialization_without_changing_probe_or_correlated_paths( + ) -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let key = arena.alloc_column(ColumnCatalog::new( + "key".to_string(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + )); + let marker = arena.alloc_column(ColumnCatalog::new( + "marker".to_string(), + true, + ColumnDesc::new(LogicalType::Boolean, None, false, None)?, + )); + let right_key = arena.alloc_column(ColumnCatalog::new( + "right_key".to_string(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + )); + let value = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(2))); + let left = LogicalPlan::new( + Operator::Values(ValuesOperator::new(vec![value; 8], 8, vec![key])), + Childrens::None, + ); + let right = LogicalPlan::new( + Operator::Values(ValuesOperator::new(vec![value; 2], 2, vec![right_key])), + Childrens::None, + ); + let left_ref = arena.alloc_expression(ScalarExpression::column_expr(key, 0)); + let right_ref = arena.alloc_expression(ScalarExpression::column_expr(right_key, 1)); + let compare = |arena: &mut PlanArena, left_expr, right_expr| { + arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Eq, + left_expr, + right_expr, + evaluator: None, + ty: LogicalType::Boolean, + }) + }; + let predicate = compare(&mut arena, left_ref, right_ref); + let constant = arena.alloc_expression(ScalarExpression::Constant(DataValue::Boolean(true))); + let op = + MarkApplyOperator::new_quantified(MarkApplyQuantifier::Any, marker, vec![predicate]); + let mut plan = LogicalPlan::new( + Operator::MarkApply(op), + Childrens::Twins { + left: Box::new(left), + right: Box::new(right), + }, + ); + for (left_rows, right_rows, materialize) in [ + (Some(0), Some(2), false), + (Some(1), Some(0), false), + (Some(1), None, false), + (Some(2), Some(8), false), + (Some(2), Some(2), true), + (Some(8), Some(2), true), + (None, None, true), + (Some(8), None, true), + ] { + let Operator::MarkApply(op) = &mut plan.operator else { + unreachable!() + }; + op.materialize = false; + let Childrens::Twins { left, right } = plan.childrens.as_mut() else { + unreachable!() + }; + let estimated = |rows: Option| { + rows.map(|rows| { + PhysicalOption::new( + PlanImpl::Values, + SortOption::None, + RowCountOption::Estimated { rows, ndv: None }, + ) + }) + }; + left.physical_option = estimated(left_rows); + right.physical_option = estimated(right_rows); + assert_eq!( + MaterializeMarkApply.apply(&mut plan, &mut arena, None)?, + materialize + ); + } + assert!(!MaterializeMarkApply.apply(&mut plan, &mut arena, None)?); + let Operator::MarkApply(op) = &mut plan.operator else { + unreachable!() + }; + assert!(op.materialize); + op.materialize = false; + op.set_parameterized_probe(Some(predicate)); + assert!(!MaterializeMarkApply.apply(&mut plan, &mut arena, None)?); + let Operator::MarkApply(op) = &mut plan.operator else { + unreachable!() + }; + op.set_parameterized_probe(None); + op.predicates.push(predicate); + assert!(!MaterializeMarkApply.apply(&mut plan, &mut arena, None)?); + let Operator::MarkApply(op) = &mut plan.operator else { + unreachable!() + }; + op.predicates.pop(); + // Only a binary comparison exposes the RHS side whose keys are cached. + let Operator::MarkApply(op) = &mut plan.operator else { + unreachable!() + }; + op.predicates[0] = constant; + assert!(!MaterializeMarkApply.apply(&mut plan, &mut arena, None)?); + let Operator::MarkApply(op) = &mut plan.operator else { + unreachable!() + }; + op.predicates[0] = predicate; + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/optimizer/rule/normalization/min_max_top_k.rs b/src/optimizer/rule/normalization/min_max_top_k.rs index 28bfae1c..bbf062e1 100644 --- a/src/optimizer/rule/normalization/min_max_top_k.rs +++ b/src/optimizer/rule/normalization/min_max_top_k.rs @@ -16,11 +16,12 @@ use crate::errors::DatabaseError; use crate::expression::agg::AggKind; use crate::expression::ScalarExpression; use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::plan_utils::{only_child, wrap_child_with}; use crate::planner::operator::sort::SortField; use crate::planner::operator::top_k::TopKOperator; use crate::planner::operator::Operator; -use crate::planner::LogicalPlan; +use crate::planner::{LogicalPlan, PlanArena}; pub struct MinMaxToTopK; @@ -28,7 +29,8 @@ impl NormalizationRule for MinMaxToTopK { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let Operator::Aggregate(op) = &plan.operator else { return Ok(false); @@ -119,7 +121,7 @@ mod tests { let mut plan = table_state.plan_with_arena("select min(c1) from t1", &mut arena)?; let agg_plan = find_aggregate_mut(&mut plan); - assert!(MinMaxToTopK.apply(agg_plan, &mut arena)?); + assert!(MinMaxToTopK.apply(agg_plan, &mut arena, None)?); let agg_plan = find_aggregate(&plan); let Operator::Aggregate(op) = &agg_plan.operator else { @@ -157,7 +159,7 @@ mod tests { let mut plan = table_state.plan_with_arena("select max(c2) from t1", &mut arena)?; let agg_plan = find_aggregate_mut(&mut plan); - assert!(MinMaxToTopK.apply(agg_plan, &mut arena)?); + assert!(MinMaxToTopK.apply(agg_plan, &mut arena, None)?); let agg_plan = find_aggregate(&plan); let child = match agg_plan.childrens.as_ref() { diff --git a/src/optimizer/rule/normalization/mod.rs b/src/optimizer/rule/normalization/mod.rs index 87eea173..6e80de89 100644 --- a/src/optimizer/rule/normalization/mod.rs +++ b/src/optimizer/rule/normalization/mod.rs @@ -14,13 +14,14 @@ use crate::errors::DatabaseError; use crate::expression::visitor_mut::ExprVisitorMut; -use crate::expression::{AliasType, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::rule::normalization::column_pruning::ColumnPruning; -use crate::optimizer::rule::normalization::combine_operators::{ - CollapseGroupByAgg, CollapseProject, CombineFilter, -}; +use crate::optimizer::rule::normalization::combine_operators::{CollapseGroupByAgg, CombineFilter}; use crate::optimizer::rule::normalization::compilation_in_advance::EvaluatorBind; +use crate::optimizer::rule::normalization::constant_propagation::{ + PropagateFilterConstants, PropagateJoinConstants, +}; use crate::planner::operator::Operator; use crate::planner::MetaArena; @@ -34,13 +35,20 @@ use crate::optimizer::rule::normalization::pushdown_predicates::{ use crate::optimizer::rule::normalization::simplification::ConstantCalculation; use crate::optimizer::rule::normalization::simplification::SimplifyFilter; use crate::optimizer::rule::normalization::top_k::TopK; -use crate::planner::{ExprRef, LogicalPlan}; +use crate::planner::{ExprRef, LogicalPlan, PlanArena}; use std::collections::HashSet; mod column_pruning; mod combine_operators; mod compilation_in_advance; +mod constant_propagation; mod elimination; +mod guard_mark_apply; +mod join_reorder; +mod materialize_mark_apply; mod min_max_top_k; +use guard_mark_apply::GuardMarkApply; +use join_reorder::GreedyJoinReorder; +use materialize_mark_apply::MaterializeMarkApply; mod parameterized_index; mod pushdown_limit; mod pushdown_predicates; @@ -48,8 +56,8 @@ mod simplification; mod top_k; pub(crate) use compilation_in_advance::evaluator_bind_current; pub(crate) use elimination::{ - apply_annotated_post_rules, apply_scan_order_hint, EliminateIndexFilter, OrderHintKind, - ScanOrderHint, + apply_annotated_post_rules, apply_scan_order_hint, EliminateIndexFilter, EliminateProject, + OrderHintKind, ScanOrderHint, }; pub(crate) use parameterized_index::{ParameterizeInnerJoin, ParameterizeMarkApply}; pub(crate) use simplification::constant_calculation_current; @@ -58,7 +66,7 @@ pub(crate) use simplification::constant_calculation_current; pub enum NormalizationRuleImpl { ColumnPruning, // Combine operators - CollapseProject, + EliminateProject, CollapseGroupByAgg, CombineFilter, // PushDown limit @@ -66,6 +74,8 @@ pub enum NormalizationRuleImpl { PushLimitThroughJoin, PushLimitIntoTableScan, // PushDown predicates + PropagateFilterConstants, + PropagateJoinConstants, PushPredicateThroughJoin, PushJoinPredicateIntoScan, // Tips: need to be used with `SimplifyFilter` @@ -80,6 +90,9 @@ pub enum NormalizationRuleImpl { ParameterizeMarkApply, ParameterizeInnerJoin, EliminateIndexFilter, + GreedyJoinReorder, + MaterializeMarkApply, + GuardMarkApply, } #[derive(Debug, Copy, Clone, Eq, PartialEq)] @@ -108,6 +121,8 @@ impl NormalizationRuleRootTag { match operator { Operator::Aggregate(_) => Some(Self::Aggregate), Operator::MarkApply(_) => Some(Self::MarkApply), + // Built after the rules that could rewrite it; nothing normalizes it further. + Operator::IndexProbeJoin(_) => None, Operator::ScalarQueryInit(_) => Some(Self::Any), Operator::Filter(_) => Some(Self::Filter), Operator::Join(_) => Some(Self::Join), @@ -169,23 +184,28 @@ impl NormalizationRuleImpl { pub fn root_tag(&self) -> NormalizationRuleRootTag { match self { NormalizationRuleImpl::ColumnPruning => NormalizationRuleRootTag::Any, - NormalizationRuleImpl::CollapseProject => NormalizationRuleRootTag::Project, + NormalizationRuleImpl::EliminateProject => NormalizationRuleRootTag::Any, NormalizationRuleImpl::CollapseGroupByAgg => NormalizationRuleRootTag::Aggregate, + NormalizationRuleImpl::GuardMarkApply => NormalizationRuleRootTag::Project, NormalizationRuleImpl::CombineFilter => NormalizationRuleRootTag::Filter, NormalizationRuleImpl::LimitProjectTranspose | NormalizationRuleImpl::PushLimitThroughJoin | NormalizationRuleImpl::PushLimitIntoTableScan | NormalizationRuleImpl::TopK => NormalizationRuleRootTag::Limit, - NormalizationRuleImpl::PushPredicateThroughJoin + NormalizationRuleImpl::PropagateFilterConstants + | NormalizationRuleImpl::PushPredicateThroughJoin | NormalizationRuleImpl::PushPredicateIntoScan | NormalizationRuleImpl::SimplifyFilter => NormalizationRuleRootTag::Filter, - NormalizationRuleImpl::PushJoinPredicateIntoScan => NormalizationRuleRootTag::Join, + NormalizationRuleImpl::PropagateJoinConstants + | NormalizationRuleImpl::PushJoinPredicateIntoScan => NormalizationRuleRootTag::Join, NormalizationRuleImpl::ConstantCalculation => NormalizationRuleRootTag::Any, NormalizationRuleImpl::EvaluatorBind => NormalizationRuleRootTag::Any, NormalizationRuleImpl::MinMaxToTopK => NormalizationRuleRootTag::Aggregate, NormalizationRuleImpl::EliminateIndexFilter => NormalizationRuleRootTag::Filter, - NormalizationRuleImpl::ParameterizeInnerJoin => NormalizationRuleRootTag::Join, - NormalizationRuleImpl::ParameterizeMarkApply => NormalizationRuleRootTag::MarkApply, + NormalizationRuleImpl::ParameterizeInnerJoin + | NormalizationRuleImpl::GreedyJoinReorder => NormalizationRuleRootTag::Join, + NormalizationRuleImpl::ParameterizeMarkApply + | NormalizationRuleImpl::MaterializeMarkApply => NormalizationRuleRootTag::MarkApply, } } } @@ -194,57 +214,67 @@ impl NormalizationRule for NormalizationRuleImpl { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { match self { - NormalizationRuleImpl::ColumnPruning => ColumnPruning.apply(plan, arena), - NormalizationRuleImpl::CollapseProject => CollapseProject.apply(plan, arena), - NormalizationRuleImpl::CollapseGroupByAgg => CollapseGroupByAgg.apply(plan, arena), - NormalizationRuleImpl::CombineFilter => CombineFilter.apply(plan, arena), + NormalizationRuleImpl::ColumnPruning => ColumnPruning.apply(plan, arena, loader), + NormalizationRuleImpl::GreedyJoinReorder => { + GreedyJoinReorder.apply(plan, arena, loader) + } + NormalizationRuleImpl::GuardMarkApply => GuardMarkApply.apply(plan, arena, loader), + NormalizationRuleImpl::EliminateProject => EliminateProject.apply(plan, arena, loader), + NormalizationRuleImpl::CollapseGroupByAgg => { + CollapseGroupByAgg.apply(plan, arena, loader) + } + NormalizationRuleImpl::CombineFilter => CombineFilter.apply(plan, arena, loader), NormalizationRuleImpl::LimitProjectTranspose => { - LimitProjectTranspose.apply(plan, arena) + LimitProjectTranspose.apply(plan, arena, loader) + } + NormalizationRuleImpl::PushLimitThroughJoin => { + PushLimitThroughJoin.apply(plan, arena, loader) + } + NormalizationRuleImpl::PushLimitIntoTableScan => { + PushLimitIntoScan.apply(plan, arena, loader) + } + NormalizationRuleImpl::PropagateFilterConstants => { + PropagateFilterConstants.apply(plan, arena, loader) + } + NormalizationRuleImpl::PropagateJoinConstants => { + PropagateJoinConstants.apply(plan, arena, loader) } - NormalizationRuleImpl::PushLimitThroughJoin => PushLimitThroughJoin.apply(plan, arena), - NormalizationRuleImpl::PushLimitIntoTableScan => PushLimitIntoScan.apply(plan, arena), NormalizationRuleImpl::PushPredicateThroughJoin => { - PushPredicateThroughJoin.apply(plan, arena) + PushPredicateThroughJoin.apply(plan, arena, loader) } NormalizationRuleImpl::PushJoinPredicateIntoScan => { - PushJoinPredicateIntoScan.apply(plan, arena) + PushJoinPredicateIntoScan.apply(plan, arena, loader) } - NormalizationRuleImpl::SimplifyFilter => SimplifyFilter.apply(plan, arena), + NormalizationRuleImpl::SimplifyFilter => SimplifyFilter.apply(plan, arena, loader), NormalizationRuleImpl::PushPredicateIntoScan => { - PushPredicateIntoScan.apply(plan, arena) + PushPredicateIntoScan.apply(plan, arena, loader) + } + NormalizationRuleImpl::ConstantCalculation => { + ConstantCalculation.apply(plan, arena, loader) + } + NormalizationRuleImpl::EvaluatorBind => EvaluatorBind.apply(plan, arena, loader), + NormalizationRuleImpl::MinMaxToTopK => MinMaxToTopK.apply(plan, arena, loader), + NormalizationRuleImpl::TopK => TopK.apply(plan, arena, loader), + NormalizationRuleImpl::EliminateIndexFilter => { + EliminateIndexFilter.apply(plan, arena, loader) } - NormalizationRuleImpl::ConstantCalculation => ConstantCalculation.apply(plan, arena), - NormalizationRuleImpl::EvaluatorBind => EvaluatorBind.apply(plan, arena), - NormalizationRuleImpl::MinMaxToTopK => MinMaxToTopK.apply(plan, arena), - NormalizationRuleImpl::TopK => TopK.apply(plan, arena), - NormalizationRuleImpl::EliminateIndexFilter => EliminateIndexFilter.apply(plan, arena), NormalizationRuleImpl::ParameterizeInnerJoin => { - ParameterizeInnerJoin.apply(plan, arena) + ParameterizeInnerJoin.apply(plan, arena, loader) } NormalizationRuleImpl::ParameterizeMarkApply => { - ParameterizeMarkApply.apply(plan, arena) + ParameterizeMarkApply.apply(plan, arena, loader) + } + NormalizationRuleImpl::MaterializeMarkApply => { + MaterializeMarkApply.apply(plan, arena, loader) } } } } -pub(crate) fn strip_alias(expr: ExprRef, arena: &crate::planner::PlanArena<'_>) -> ExprRef { - match arena.expression(expr) { - ScalarExpression::Alias { - expr, - alias: AliasType::Name(_), - } => strip_alias(*expr, arena), - ScalarExpression::Alias { - alias: AliasType::Expr(alias_expr), - .. - } => strip_alias(*alias_expr, arena), - _ => expr, - } -} - pub(crate) fn remap_position(position: &mut usize, removed_positions: &[usize]) { match removed_positions.binary_search(position) { Ok(_) => { @@ -277,11 +307,11 @@ impl<'positions, 'visited> PositionRemapper<'positions, 'visited> { } } -impl ExprVisitorMut for PositionRemapper<'_, '_> { +impl ExprVisitorMut for PositionRemapper<'_, '_> { fn visit_expression_ref( &mut self, expr: &mut ExprRef, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result { Ok(self.visited.insert(*expr)) } @@ -290,39 +320,27 @@ impl ExprVisitorMut for PositionRemapper<'_, '_> { &mut self, _column: &mut crate::catalog::ColumnRef, position: &mut usize, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result<(), DatabaseError> { remap_position(position, self.removed_positions); Ok(()) } - - fn visit_alias( - &mut self, - expr: &mut ExprRef, - alias: &mut AliasType, - arena: &mut (dyn MetaArena + '_), - ) -> Result<(), DatabaseError> { - match alias { - AliasType::Expr(alias_expr) => self.visit(alias_expr, arena), - AliasType::Name(_) => self.visit(expr, arena), - } - } } -pub(crate) fn remap_expr_positions( +pub(crate) fn remap_expr_positions( mut expr: ExprRef, removed_positions: &[usize], visited: &mut HashSet, - arena: &mut crate::planner::PlanArena<'_>, + arena: &mut A, ) -> Result<(), DatabaseError> { PositionRemapper::new(removed_positions, visited).visit(&mut expr, arena) } -pub(crate) fn remap_exprs_positions<'a>( +pub(crate) fn remap_exprs_positions<'a, A: MetaArena + ?Sized>( exprs: impl IntoIterator, removed_positions: &[usize], visited: &mut HashSet, - arena: &mut crate::planner::PlanArena<'_>, + arena: &mut A, ) -> Result<(), DatabaseError> { let mut remapper = PositionRemapper::new(removed_positions, visited); for expr in exprs { diff --git a/src/optimizer/rule/normalization/parameterized_index.rs b/src/optimizer/rule/normalization/parameterized_index.rs index 8ee20085..52cb67eb 100644 --- a/src/optimizer/rule/normalization/parameterized_index.rs +++ b/src/optimizer/rule/normalization/parameterized_index.rs @@ -17,17 +17,25 @@ use crate::errors::DatabaseError; use crate::expression::visitor_mut::{ExprVisitorMut, PositionShift}; use crate::expression::{BinaryOperator, ScalarExpression, TypeCast}; use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::optimizer::rule::normalization::pushdown_predicates::{index_covering, is_index_covers}; use crate::planner::operator::filter::FilterOperator; +use crate::planner::operator::index_probe_join::IndexProbeJoinOperator; use crate::planner::operator::join::{JoinCondition, JoinType}; -use crate::planner::operator::mark_apply::{MarkApplyKind, MarkApplyOperator, MarkApplyQuantifier}; +use crate::planner::operator::mark_apply::{MarkApplyKind, MarkApplyQuantifier}; use crate::planner::operator::project::ProjectOperator; use crate::planner::operator::table_scan::TableScanOperator; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{ + rows_work, Operator, PhysicalOption, PlanImpl, RowCountOption, SortOption, CPU_OPERATOR_WORK, + DEFAULT_TABLE_ROWS, PROBE_RESET_WORK, +}; use crate::planner::MetaArena; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; -use crate::types::index::{IndexLookup, IndexType}; +use crate::types::index::{IndexLookup, IndexMetaRef, IndexType}; use crate::types::tuple::Schema; use crate::types::LogicalType; +use std::borrow::Cow; +use std::mem; pub(crate) struct ParameterizeMarkApply; @@ -35,12 +43,11 @@ impl NormalizationRule for ParameterizeMarkApply { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let (op, new_probe) = match (&mut plan.operator, plan.childrens.as_mut()) { - (Operator::MarkApply(op), Childrens::Twins { left, right }) - if !matches!(op.kind, MarkApplyKind::InnerJoin) => - { + (Operator::MarkApply(op), Childrens::Twins { left, right }) => { let probe = find_parameterized_probe( op.kind, op.predicates(), @@ -48,9 +55,11 @@ impl NormalizationRule for ParameterizeMarkApply { right.output_schema(arena), arena, )?; - let new_probe = probe.and_then(|(right_column, left_expr)| { - parameterize_right_subtree(right, &right_column, arena).then_some(left_expr) - }); + let new_probe = if let Some((right_column, left_expr)) = probe { + parameterize_right_subtree(right, &right_column, left_expr, arena)? + } else { + None + }; (op, new_probe) } _ => return Ok(false), @@ -87,7 +96,7 @@ fn find_parameterized_probe( Ok(None) } } - MarkApplyKind::InnerJoin | MarkApplyKind::Quantified(MarkApplyQuantifier::All) => Ok(None), + MarkApplyKind::Quantified(MarkApplyQuantifier::All) => Ok(None), } } @@ -142,7 +151,7 @@ fn extract_parameterized_probe_side( if !schema_contains_column(right_schema, &right_column, arena) { return Ok(None); } - if !left_expr.all_referenced_columns(arena, |arena, candidate| { + if !left_expr.all_referenced_columns(arena, |arena, candidate, _| { schema_contains_column(left_schema, candidate, arena) })? { return Ok(None); @@ -159,45 +168,58 @@ fn extract_parameterized_probe_side( fn parameterize_right_subtree( plan: &mut LogicalPlan, right_column: &ColumnRef, - arena: &dyn MetaArena, -) -> bool { - if matches!(plan.operator, Operator::TableScan(_)) { - let index_info = { - let Operator::TableScan(scan_op) = &mut plan.operator else { - unreachable!(); + mut left_expr: ExprRef, + arena: &mut PlanArena<'_>, +) -> Result, DatabaseError> { + match (&mut plan.operator, plan.childrens.as_mut()) { + (Operator::Project(project), Childrens::Only(child)) => { + let Some(mut expr) = project.exprs.iter().copied().find(|expr| { + let column = expr.output_column_ref(arena); + arena.same_column(column, *right_column) + }) else { + return Ok(None); }; - let Some(target_index) = - pick_parameterized_index_position(scan_op, right_column, arena) + // Follow the source, not the projection's output relation/name. + // Computed and cast keys are not index probes. + while let ScalarExpression::Alias { expr: source, .. } = arena.expression(expr) { + expr = *source; + } + let ScalarExpression::ColumnRef { column, .. } = *arena.expression(expr) else { + return Ok(None); + }; + parameterize_right_subtree(child, &column, left_expr, arena) + } + (Operator::TableScan(scan), _) => { + if scan.limit != (None, None) { + return Ok(None); + } + let Some(position) = pick_parameterized_index_position(scan, right_column, arena) else { - return false; + return Ok(None); }; - scan_op.index_infos[target_index].lookup = Some(IndexLookup::Probe); - scan_op.index_infos[target_index].clone() - }; - let sort_option = index_info.sort_option.clone(); - plan.physical_option = Some(PhysicalOption::new( - PlanImpl::IndexScan(Box::new(index_info)), - sort_option, - )); - return true; - } - - let passthrough = matches!( - plan.operator, - Operator::Filter(_) - | Operator::Project(_) - | Operator::Limit(_) - | Operator::Sort(_) - | Operator::TopK(_) - ); - - if !passthrough { - return false; - } - - match plan.childrens.as_mut() { - Childrens::Only(child) => parameterize_right_subtree(child, right_column, arena), - _ => false, + let key_type = &arena.index(scan.index_infos[position].meta).value_ty; + let left_type = left_expr.return_type(arena); + if left_type.as_ref() != key_type { + let common = LogicalType::max_logical_type(&left_type, key_type)?; + if common.as_ref() != key_type { + return Ok(None); + } + left_expr = left_expr.type_cast(Cow::Owned(common.into_owned()), arena)?; + } + scan.index_infos[position].lookup = Some(IndexLookup::Probe); + let index_info = scan.index_infos[position].clone(); + let sort_option = index_info.sort_option.clone(); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::IndexScan(Box::new(index_info)), + sort_option, + RowCountOption::Unknown, + )); + Ok(Some(left_expr)) + } + (Operator::Filter(_) | Operator::Sort(_), Childrens::Only(child)) => { + parameterize_right_subtree(child, right_column, left_expr, arena) + } + _ => Ok(None), } } @@ -221,6 +243,7 @@ fn pick_parameterized_index_position( .filter(|(_, index_info)| { let index_meta = arena.index(index_info.meta); index_meta.table_name == *table_name + && index_meta.column_ids.len() == 1 && index_meta.column_ids.first().copied() == Some(column_id) }) .min_by_key(|(_, index_info)| index_priority(arena.index(index_info.meta).ty)) @@ -271,20 +294,53 @@ fn constant_keys(expr: ExprRef, keys: &mut Vec<(ExprRef, ExprRef)>, arena: &Plan } } -fn parameterize( - plan: &mut LogicalPlan, - mut keys: Vec<(ExprRef, ExprRef)>, +/// A PK/Unique index of a `Filter* -> TableScan` input that join keys can probe. +pub(crate) struct ProbeIndex { + position: usize, + pub(crate) probe: Vec, + probed_keys: u64, + pub(crate) lookup_work: usize, +} + +#[inline] +fn mark_probed(probed_keys: &mut u64, position: usize) { + if position < u64::BITS as usize { + *probed_keys |= 1 << position; + } +} + +#[inline] +fn is_probed(probed_keys: u64, position: usize) -> bool { + position < u64::BITS as usize && probed_keys & (1 << position) != 0 +} + +pub(crate) fn probe_index( + plan: &LogicalPlan, + keys: &mut Vec<(ExprRef, ExprRef)>, + join_key_count: usize, arena: &PlanArena<'_>, -) -> Option> { - match (&mut plan.operator, plan.childrens.as_mut()) { + loader: Option<&StatisticMetaLoader<'_>>, +) -> Result, DatabaseError> { + match (&plan.operator, plan.childrens.as_ref()) { (Operator::Filter(filter), Childrens::Only(child)) => { - constant_keys(filter.predicate, &mut keys, arena); - parameterize(child, keys, arena) + let key_count = keys.len(); + constant_keys(filter.predicate, keys, arena); + let probe = probe_index(child, keys, join_key_count, arena, loader); + // Keep filter keys local to this traversal, including failed attempts. + keys.truncate(key_count); + probe } (Operator::TableScan(scan), _) if scan.limit == (None, None) => { - 'indexes: for info in &mut scan.index_infos { + let mut best: Option = None; + let mut probe = Vec::new(); + 'indexes: for (position, info) in scan.index_infos.iter().enumerate() { let meta = arena.index(info.meta); - let mut probe = Vec::with_capacity(meta.column_ids.len()); + if !matches!(meta.ty, IndexType::PrimaryKey { .. } | IndexType::Unique) { + continue; + } + probe.clear(); + let mut uses_join_key = false; + let mut probed_keys = 0u64; 'columns: for id in &meta.column_ids { let Some(candidate) = scan .columns @@ -293,7 +349,7 @@ fn parameterize( else { continue 'indexes; }; - for &(value, column) in &keys { + for (position, &(value, column)) in keys.iter().enumerate() { let ScalarExpression::ColumnRef { column, .. } = arena.expression(column.unpack_alias(arena)) else { @@ -304,33 +360,154 @@ fn parameterize( == arena.column(*candidate).datatype() { probe.push(value); + if position < join_key_count { + uses_join_key = true; + mark_probed(&mut probed_keys, position); + } continue 'columns; } } continue 'indexes; } - info.lookup = Some(IndexLookup::Probe); - info.residual_predicate = None; - plan.physical_option = Some(PhysicalOption::new( - PlanImpl::IndexScan(Box::new(info.clone())), - info.sort_option.clone(), - )); - return Some(probe); + // Constants appended by filters may complete a composite key, + // but must not trigger a per-outer-row probe on their own. + if !uses_join_key { + continue; + } + let lookup_work = lookup_work(scan, info.meta, arena, loader)?; + // Equal work keeps the index order. + if best + .as_ref() + .is_none_or(|best| lookup_work < best.lookup_work) + { + best = Some(ProbeIndex { + position, + probe: mem::take(&mut probe), + probed_keys, + lookup_work, + }); + } } - None + Ok(best) } - // LIMIT/aggregation/projection are not row-local filters. Moving a probe - // below them could change which rows the original inner input produces. - _ => None, + _ => Ok(None), } } +/// RFC 0001 Ā§4.3: reset the probed input, seek the index (twice when it must read the +/// table row too), and read one row. Without statistics the index has `DEFAULT_TABLE_ROWS`. +fn lookup_work( + scan: &TableScanOperator, + index: IndexMetaRef, + arena: &PlanArena<'_>, + loader: Option<&StatisticMetaLoader<'_>>, +) -> Result { + let meta = arena.index(index); + let rows = match loader { + Some(loader) => loader + .load(&scan.table_name, meta.id)? + .map(|statistics| statistics.histogram().values_len()), + None => None, + } + .unwrap_or(DEFAULT_TABLE_ROWS); + let seeks = if matches!(meta.ty, IndexType::PrimaryKey { .. }) + || is_index_covers(&scan.columns, index, arena) + { + 1 + } else { + 2 + }; + let depth = (usize::BITS - rows.saturating_sub(1).leading_zeros()) as usize; + Ok(PROBE_RESET_WORK + .saturating_add(seeks * depth * CPU_OPERATOR_WORK) + .saturating_add(rows_work(1, scan.columns.len()))) +} + +fn use_probe_index(plan: &mut LogicalPlan, position: usize, arena: &PlanArena<'_>) { + match (&mut plan.operator, plan.childrens.as_mut()) { + (Operator::Filter(_), Childrens::Only(child)) => use_probe_index(child, position, arena), + (Operator::TableScan(scan), _) => { + let info = &mut scan.index_infos[position]; + info.lookup = Some(IndexLookup::Probe); + info.residual_predicate = None; + (info.covered_deserializers, info.cover_mapping) = + match index_covering(&scan.columns, info.meta, arena) { + Some((deserializers, mapping)) => (Some(deserializers), mapping), + None => (None, None), + }; + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::IndexScan(Box::new(info.clone())), + info.sort_option.clone(), + RowCountOption::Unknown, + )); + } + // SAFETY: `probe_index` only accepts filters over a table scan. + _ => unreachable!(), + } +} + +/// Builds the `IndexProbeJoin` probing `right` with `index` for every `left` row. `on` holds +/// the join keys bound to each side; keys not used by the probe are checked after it. +pub(crate) fn probe_join( + mut left: LogicalPlan, + mut right: LogicalPlan, + on: Vec<(ExprRef, ExprRef)>, + index: ProbeIndex, + row_count: RowCountOption, + arena: &mut PlanArena<'_>, +) -> Result { + use_probe_index(&mut right, index.position, arena); + let left_len = left.output_schema(arena).len(); + let probe = if index.probe.len() == 1 { + index.probe[0] + } else { + arena.alloc_expression(ScalarExpression::Tuple(index.probe)) + }; + let mut predicates = Vec::with_capacity(on.len()); + for (i, (left_expr, mut right_expr)) in on.into_iter().enumerate() { + if is_probed(index.probed_keys, i) { + continue; + } + PositionShift { + delta: left_len as isize, + } + .visit(&mut right_expr, arena)?; + predicates.push(arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Eq, + left_expr, + right_expr, + evaluator: None, + ty: LogicalType::Boolean, + })); + } + let mut apply = LogicalPlan::new( + Operator::IndexProbeJoin(IndexProbeJoinOperator::new(probe, predicates)), + Childrens::Twins { + left: Box::new(left), + right: Box::new(right), + }, + ); + let mut option = PhysicalOption::new(PlanImpl::IndexProbeJoin, SortOption::Follow, row_count); + option.lookup_work = Some(index.lookup_work); + apply.physical_option = Some(option); + Ok(apply) +} + impl NormalizationRule for ParameterizeInnerJoin { fn apply( &self, plan: &mut LogicalPlan, arena: &mut PlanArena<'_>, + loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { + if plan + .physical_option + .as_ref() + .is_some_and(|option| option.join_decided) + { + return Ok(false); + } + let rows = plan.estimated_rows(); let (Operator::Join(join), Childrens::Twins { left, right }) = (&mut plan.operator, plan.childrens.as_mut()) else { @@ -346,13 +523,16 @@ impl NormalizationRule for ParameterizeInnerJoin { return Ok(false); } let mut keys = on.clone(); - let (project, probe_keys) = if let Some(probe) = parameterize(right, keys.clone(), arena) { - (None, probe) + // Filter traversal only appends constants after these outer-dependent keys. + let join_key_count = keys.len(); + let right_probe = probe_index(right, &mut keys, join_key_count, arena, loader)?; + let (project, index) = if let Some(index) = right_probe { + (None, index) } else { for (l, r) in &mut keys { std::mem::swap(l, r); } - let Some(probe) = parameterize(left, keys.clone(), arena) else { + let Some(index) = probe_index(left, &mut keys, join_key_count, arena, loader)? else { return Ok(false); }; let left_schema = left.output_schema(arena); @@ -376,43 +556,25 @@ impl NormalizationRule for ParameterizeInnerJoin { .collect(); std::mem::swap(left, right); let mut project = LogicalPlan::new( - Operator::Project(ProjectOperator { exprs }), + Operator::Project(ProjectOperator::new(exprs)), Childrens::None, ); - project.physical_option = - Some(PhysicalOption::new(PlanImpl::Project, SortOption::Follow)); - (Some(project), probe) + project.physical_option = Some(PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + )); + (Some(project), index) }; let filter = filter.take(); - let (mut left, right) = plan.take().childrens.pop_twins(); - let left_len = left.output_schema(arena).len(); - let probe = if probe_keys.len() == 1 { - probe_keys[0] - } else { - arena.alloc_expression(ScalarExpression::Tuple(probe_keys)) - }; - let mut predicates = Vec::with_capacity(keys.len()); - for (left_expr, mut right_expr) in keys { - PositionShift { - delta: left_len as isize, - } - .visit(&mut right_expr, arena)?; - predicates.push(arena.alloc_expression(ScalarExpression::Binary { - op: BinaryOperator::Eq, - left_expr, - right_expr, - evaluator: None, - ty: LogicalType::Boolean, - })); - } - let apply = LogicalPlan::new( - Operator::MarkApply(MarkApplyOperator::new_inner_join(predicates, probe)), - Childrens::Twins { - left: Box::new(left), - right: Box::new(right), - }, - ); + let (left, right) = plan.take().childrens.pop_twins(); + // This is the same join output, not the per-key cardinality of the probe. + let row_count = rows.map_or(RowCountOption::Unknown, |rows| RowCountOption::Estimated { + rows, + ndv: None, + }); + let apply = probe_join(left, right, keys, index, row_count, arena)?; *plan = if let Some(mut project) = project { project.childrens = Box::new(Childrens::Only(Box::new(apply))); project @@ -422,7 +584,11 @@ impl NormalizationRule for ParameterizeInnerJoin { if let Some(filter) = filter { // The original join filter uses the original left/right output slots. *plan = FilterOperator::build(filter, plan.take(), false); - plan.physical_option = Some(PhysicalOption::new(PlanImpl::Filter, SortOption::Follow)); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Follow, + )); } Ok(true) } @@ -545,10 +711,11 @@ mod tests { let table_arena = TableArenaCell::default(); let mut arena = PlanArena::new(&table_arena); let column = column(&mut arena, "value"); + let probe = arena.alloc_expression(ScalarExpression::from(2)); let mut plan = LogicalPlan::new(Operator::Dummy, Childrens::None); - assert!(!ParameterizeMarkApply.apply(&mut plan, &mut arena)?); - assert!(!parameterize_right_subtree(&mut plan, &column, &arena)); + assert!(!ParameterizeMarkApply.apply(&mut plan, &mut arena, None)?); + assert!(parameterize_right_subtree(&mut plan, &column, probe, &mut arena)?.is_none()); let mut filter = LogicalPlan::new( Operator::Filter(FilterOperator { @@ -558,7 +725,7 @@ mod tests { }), Childrens::None, ); - assert!(!parameterize_right_subtree(&mut filter, &column, &arena)); + assert!(parameterize_right_subtree(&mut filter, &column, probe, &mut arena)?.is_none()); assert_eq!( index_priority(IndexType::PrimaryKey { is_multiple: false }), @@ -579,11 +746,13 @@ mod inner_join_tests { use super::*; use crate::catalog::{ColumnCatalog, ColumnDesc}; use crate::expression::range_detacher::Range; + use crate::expression::{AliasType, UnaryOperator}; use crate::planner::operator::join::JoinOperator; - use crate::planner::operator::mark_apply::MarkApplyKind; + use crate::planner::operator::limit::LimitOperator; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::TableArenaCell; use crate::types::index::{IndexInfo, IndexMeta, IndexType}; + use crate::types::tuple::Tuple; use crate::types::value::DataValue; fn scan( @@ -647,6 +816,24 @@ mod inner_join_tests { ) } + fn unique_scan( + arena: &mut PlanArena<'_>, + name: &str, + types: &[LogicalType], + indexes: &[&[usize]], + ) -> LogicalPlan { + let mut plan = scan(arena, name, types, indexes); + let Operator::TableScan(scan) = &mut plan.operator else { + unreachable!() + }; + for info in &mut scan.index_infos { + let mut meta = arena.index(info.meta).clone(); + meta.ty = IndexType::Unique; + info.meta = arena.alloc_index(meta); + } + plan + } + fn column(plan: &LogicalPlan, position: usize, arena: &mut PlanArena<'_>) -> ExprRef { let Operator::TableScan(scan) = &plan.operator else { panic!("expected scan") @@ -694,6 +881,294 @@ mod inner_join_tests { ) } + #[test] + fn scalar_probe_rejects_composite_unique_keys() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut right = scan( + &mut arena, + "right", + &[LogicalType::Integer, LogicalType::Integer], + &[&[0, 1], &[0]], + ); + let (column, single) = { + let Operator::TableScan(scan) = &mut right.operator else { + unreachable!() + }; + let mut composite = arena.index(scan.index_infos[0].meta).clone(); + composite.ty = IndexType::Unique; + scan.index_infos[0].meta = arena.alloc_index(composite); + (scan.columns[0], scan.index_infos.pop().unwrap()) + }; + let probe = arena.alloc_expression(ScalarExpression::from(2)); + assert!(parameterize_right_subtree(&mut right, &column, probe, &mut arena)?.is_none()); + assert!(right.physical_option.is_none()); + let Operator::TableScan(scan) = &mut right.operator else { + unreachable!() + }; + assert!(scan.index_infos[0].lookup.is_none()); + scan.index_infos.push(single); + // Prefer a usable scalar index even though UNIQUE normally ranks higher. + assert_eq!( + pick_parameterized_index_position(scan, &column, &arena), + Some(1) + ); + assert_eq!( + parameterize_right_subtree(&mut right, &column, probe, &mut arena)?, + Some(probe) + ); + let PlanImpl::IndexScan(index) = &right.physical_option.as_ref().unwrap().plan else { + unreachable!() + }; + assert_eq!(arena.index(index.meta).column_ids, vec![0]); + assert!(matches!(index.lookup, Some(IndexLookup::Probe))); + Ok(()) + } + + #[test] + fn scalar_probe_normalizes_only_safe_key_types() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut right = scan(&mut arena, "right", &[LogicalType::Bigint], &[&[0]]); + let Operator::TableScan(op) = &right.operator else { + unreachable!() + }; + let column = op.columns[0]; + let original = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(2))); + let probe = parameterize_right_subtree(&mut right, &column, original, &mut arena)?.unwrap(); + assert_eq!(probe.return_type(&arena).as_ref(), &LogicalType::Bigint); + assert_eq!( + *arena.expression(probe).eval(&arena, None::<&Tuple>)?, + DataValue::Int64(2) + ); + assert_eq!( + *arena.expression(original).eval(&arena, None::<&Tuple>)?, + DataValue::Int32(2) + ); + + for value in [ + DataValue::Float64(2.5.into()), + DataValue::from("hello".to_string()), + ] { + let mut right = scan(&mut arena, "right", &[LogicalType::Bigint], &[&[0]]); + let Operator::TableScan(op) = &right.operator else { + unreachable!() + }; + let column = op.columns[0]; + let original = arena.alloc_expression(ScalarExpression::Constant(value)); + assert!( + parameterize_right_subtree(&mut right, &column, original, &mut arena)?.is_none() + ); + assert!(right.physical_option.is_none()); + let Operator::TableScan(op) = &right.operator else { + unreachable!() + }; + assert!(op.index_infos[0].lookup.is_none()); + } + + let mut right = scan(&mut arena, "narrow", &[LogicalType::Integer], &[&[0]]); + let Operator::TableScan(op) = &right.operator else { + unreachable!() + }; + let column = op.columns[0]; + let large = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int64(1 << 40))); + assert!(parameterize_right_subtree(&mut right, &column, large, &mut arena)?.is_none()); + assert!(right.physical_option.is_none()); + Ok(()) + } + + #[test] + fn projected_probe_resolves_output_alias_but_rejects_computed_keys_and_limits( + ) -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let base = scan(&mut arena, "right", &[LogicalType::Integer], &[&[0]]); + let source = column(&base, 0, &mut arena); + let output = arena.alloc_column(ColumnCatalog::new( + "renamed".to_string(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + )); + let alias_column = arena.alloc_expression(ScalarExpression::column_expr(output, 0)); + let alias = arena.alloc_expression(ScalarExpression::Alias { + expr: source, + alias: AliasType::Expr(alias_column), + }); + let project = |input| { + LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![alias])), + Childrens::Only(Box::new(input)), + ) + }; + let probe = arena.alloc_expression(ScalarExpression::from(2)); + let mut projected = project(base.clone()); + assert_eq!( + parameterize_right_subtree(&mut projected, &output, probe, &mut arena)?, + Some(probe) + ); + let input = projected.childrens.only(); + let Operator::TableScan(input) = &input.operator else { + unreachable!() + }; + assert_eq!(input.index_infos[0].lookup, Some(IndexLookup::Probe)); + assert!(matches!( + projected + .childrens + .only() + .physical_option + .as_ref() + .unwrap() + .plan, + PlanImpl::IndexScan(_) + )); + + let mut limited = project(LimitOperator::build(Some(1), Some(1), base.clone())); + assert!(parameterize_right_subtree(&mut limited, &output, probe, &mut arena)?.is_none()); + let mut pushed = base.clone(); + let Operator::TableScan(op) = &mut pushed.operator else { + unreachable!() + }; + op.limit = (None, Some(1)); + let mut pushed = project(pushed); + assert!(parameterize_right_subtree(&mut pushed, &output, probe, &mut arena)?.is_none()); + + let computed = arena.alloc_expression(ScalarExpression::Unary { + op: UnaryOperator::Minus, + expr: source, + evaluator: None, + ty: LogicalType::Integer, + }); + let computed = arena.alloc_expression(ScalarExpression::Alias { + expr: computed, + alias: AliasType::Expr(alias_column), + }); + let mut projected = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![computed])), + Childrens::Only(Box::new(base)), + ); + assert!(parameterize_right_subtree(&mut projected, &output, probe, &mut arena)?.is_none()); + Ok(()) + } + + #[test] + fn parameterized_join_preserves_rows_without_guessing_probe_cardinality( + ) -> Result<(), DatabaseError> { + for rows in [None, Some(0), Some(120)] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let mut left = scan(&mut arena, "outer", &[LogicalType::Integer], &[]); + let mut right = unique_scan(&mut arena, "inner", &[LogicalType::Integer], &[&[0]]); + let l = column(&left, 0, &mut arena); + let r = column(&right, 0, &mut arena); + left.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 129, + ndv: None, + }, + )); + right.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 120, + ndv: None, + }, + )); + let mut plan = join(left, right, vec![(l, r)], None); + if let Some(rows) = rows { + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::HashJoin, + SortOption::None, + RowCountOption::Estimated { rows, ndv: None }, + )); + } + assert!(ParameterizeInnerJoin.apply(&mut plan, &mut arena, None)?); + let option = plan.physical_option.as_ref().unwrap(); + assert_eq!(option.plan, PlanImpl::IndexProbeJoin); + assert_eq!(plan.estimated_rows(), rows); + let Childrens::Twins { right, .. } = plan.childrens.as_ref() else { + unreachable!() + }; + assert_eq!(right.estimated_rows(), None); + assert_eq!( + right.physical_option.as_ref().unwrap().row_count, + RowCountOption::Unknown + ); + assert_eq!( + plan.price(&mut arena).and_then(|cost| cost.startup), + Some(0) + ); + } + Ok(()) + } + + #[test] + fn constant_only_index_does_not_block_outer_key_probe() -> Result<(), DatabaseError> { + for swap in [false, true] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let left_indexes: &[&[usize]] = if swap { &[&[0]] } else { &[] }; + let right_indexes: &[&[usize]] = if swap { &[&[1]] } else { &[&[1], &[0]] }; + let left = unique_scan(&mut arena, "outer", &[LogicalType::Integer], left_indexes); + let right = unique_scan( + &mut arena, + "inner", + &[LogicalType::Integer, LogicalType::Integer], + right_indexes, + ); + let l = column(&left, 0, &mut arena); + let r = column(&right, 0, &mut arena); + let static_key = column(&right, 1, &mut arena); + let constant = + arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(782))); + let predicate = binary(&mut arena, BinaryOperator::Eq, constant, static_key); + let right = FilterOperator::build(predicate, right, false); + let original_right = right.clone(); + let mut plan = join(left, right, vec![(l, r)], None); + let schema = plan.output_schema(&mut arena).clone(); + + assert!(ParameterizeInnerJoin.apply(&mut plan, &mut arena, None)?); + assert_eq!(plan.output_schema(&mut arena), &schema); + let apply_plan = if swap { + let Childrens::Only(child) = plan.childrens.as_ref() else { + panic!("expected output-restoring projection") + }; + child.as_ref() + } else { + &plan + }; + let Operator::IndexProbeJoin(apply) = &apply_plan.operator else { + panic!("expected index probe join") + }; + assert!(apply.predicates.is_empty()); + assert_eq!(apply.probe, if swap { r } else { l }); + let Childrens::Twins { left, right } = apply_plan.childrens.as_ref() else { + unreachable!() + }; + let scan = if swap { + assert_eq!(left.as_ref(), &original_right); + right.as_ref() + } else { + let Operator::Filter(filter) = &right.operator else { + panic!("static filter must survive") + }; + assert_eq!(filter.predicate, predicate); + let Childrens::Only(child) = right.childrens.as_ref() else { + unreachable!() + }; + child.as_ref() + }; + let PlanImpl::IndexScan(info) = &scan.physical_option.as_ref().unwrap().plan else { + panic!("expected index scan") + }; + assert_eq!(arena.index(info.meta).column_ids, vec![0]); + assert_eq!(info.lookup, Some(IndexLookup::Probe)); + } + Ok(()) + } + #[test] fn composite_probe_uses_index_order_and_preserves_static_filter() -> Result<(), DatabaseError> { let tables = TableArenaCell::default(); @@ -704,12 +1179,12 @@ mod inner_join_tests { &[const { LogicalType::Integer }; 2], &[], ); - // First index cannot be probed. Second requires constant + two join keys. - let mut right = scan( + // Skip the constant-only and unmatched indexes; the last uses both key sources. + let mut right = unique_scan( &mut arena, "inner", &[const { LogicalType::Integer }; 4], - &[&[3], &[0, 2, 1]], + &[&[0], &[3], &[0, 2, 1]], ); let l0 = column(&left, 0, &mut arena); let l1 = column(&left, 1, &mut arena); @@ -724,28 +1199,31 @@ mod inner_join_tests { let Operator::TableScan(scan) = &mut right.operator else { unreachable!() }; - let info = &mut scan.index_infos[1]; + let info = &mut scan.index_infos[2]; + let mut meta = arena.index(info.meta).clone(); + meta.ty = IndexType::PrimaryKey { is_multiple: true }; + info.meta = arena.alloc_index(meta); info.lookup = Some(IndexLookup::Static(Range::Eq(DataValue::Int32(7)))); info.residual_predicate = Some(residual); let index = info.meta; right.physical_option = Some(PhysicalOption::new( PlanImpl::IndexScan(Box::new(info.clone())), SortOption::None, + RowCountOption::Unknown, )); let right = FilterOperator::build(predicate, right, false); let mut plan = join(left, right, vec![(l0, r1), (l1, r2)], None); - assert!(ParameterizeInnerJoin.apply(&mut plan, &mut arena)?); - let Operator::MarkApply(apply) = &plan.operator else { - panic!("expected apply without projection") + assert!(ParameterizeInnerJoin.apply(&mut plan, &mut arena, None)?); + let Operator::IndexProbeJoin(apply) = &plan.operator else { + panic!("expected index probe join without projection") }; - assert_eq!(apply.kind, MarkApplyKind::InnerJoin); + assert!(apply.predicates.is_empty()); assert_eq!( - arena.expression(*apply.parameterized_probe().unwrap()), + arena.expression(apply.probe), &ScalarExpression::Tuple(vec![constant, l1, l0]) ); assert_eq!((position(l0, &arena), position(l1, &arena)), (0, 1)); - assert_eq!((position(r1, &arena), position(r2, &arena)), (3, 4)); let Childrens::Twins { right, .. } = plan.childrens.as_ref() else { unreachable!() }; @@ -771,7 +1249,7 @@ mod inner_join_tests { fn swapped_join_restores_unequal_widths_and_filter_positions() -> Result<(), DatabaseError> { let tables = TableArenaCell::default(); let mut arena = PlanArena::new(&tables); - let mut left = scan( + let mut left = unique_scan( &mut arena, "left", &[const { LogicalType::Integer }; 3], @@ -785,9 +1263,42 @@ mod inner_join_tests { let filter_left = arena.alloc_expression(ScalarExpression::column_expr(schema[2], 2)); let filter_right = arena.alloc_expression(ScalarExpression::column_expr(schema[3], 3)); let predicate = binary(&mut arena, BinaryOperator::Gt, filter_left, filter_right); + left.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 100, + ndv: None, + }, + )); + right.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 10, + ndv: None, + }, + )); let mut plan = join(left, right, vec![(l, r)], Some(predicate)); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::HashJoin, + SortOption::None, + RowCountOption::Estimated { + rows: 30, + ndv: None, + }, + )); - assert!(ParameterizeInnerJoin.apply(&mut plan, &mut arena)?); + assert!(ParameterizeInnerJoin.apply(&mut plan, &mut arena, None)?); + assert_eq!(plan.estimated_rows(), Some(30)); + assert_eq!( + plan.physical_option.as_ref().unwrap().row_count, + RowCountOption::Follow + ); + assert_eq!( + plan.price(&mut arena).and_then(|cost| cost.startup), + Some(0) + ); assert_eq!(plan.output_schema(&mut arena), &schema); let Operator::Filter(filter) = &plan.operator else { panic!("join filter must remain above projection") @@ -806,6 +1317,10 @@ mod inner_join_tests { let Operator::Project(project_op) = &project.operator else { panic!("expected reorder projection") }; + assert_eq!( + project.physical_option.as_ref().unwrap().row_count, + RowCountOption::Follow + ); assert_eq!( project_op .exprs @@ -817,20 +1332,17 @@ mod inner_join_tests { let Childrens::Only(apply_plan) = project.childrens.as_ref() else { unreachable!() }; - let Operator::MarkApply(apply) = &apply_plan.operator else { - panic!("expected apply") - }; - assert_eq!(apply.parameterized_probe(), Some(&r)); - assert_eq!((position(r, &arena), position(l, &arena)), (0, 2)); - let ScalarExpression::Binary { - left_expr, - right_expr, - .. - } = arena.expression(apply.predicates()[0]) - else { - unreachable!() + let Operator::IndexProbeJoin(apply) = &apply_plan.operator else { + panic!("expected index probe join") }; - assert_eq!((*left_expr, *right_expr), (r, l)); + assert_eq!(apply_plan.estimated_rows(), Some(30)); + assert_eq!( + apply_plan.physical_option.as_ref().unwrap().plan, + PlanImpl::IndexProbeJoin + ); + assert_eq!(apply.probe, r); + assert_eq!(position(r, &arena), 0); + assert!(apply.predicates.is_empty()); let Childrens::Twins { left, right } = apply_plan.childrens.as_ref() else { unreachable!() }; @@ -848,7 +1360,15 @@ mod inner_join_tests { #[test] fn rejected_probe_leaves_join_and_expression_positions_unchanged() -> Result<(), DatabaseError> { - for case in ["missing_key", "type_mismatch", "inner_limit", "outer_join"] { + for case in [ + "missing_key", + "type_mismatch", + "inner_limit", + "outer_join", + "normal_index", + "composite_index", + "join_decided", + ] { let tables = TableArenaCell::default(); let mut arena = PlanArena::new(&tables); let left = scan(&mut arena, "left", &[LogicalType::Integer], &[]); @@ -857,8 +1377,16 @@ mod inner_join_tests { } else { LogicalType::Integer }; - let index: &[usize] = if case == "missing_key" { &[0, 1] } else { &[0] }; - let mut right = scan(&mut arena, "right", &[ty, LogicalType::Integer], &[index]); + let index: &[usize] = if matches!(case, "missing_key" | "composite_index") { + &[0, 1] + } else { + &[0] + }; + let mut right = if matches!(case, "normal_index" | "composite_index") { + scan(&mut arena, "right", &[ty, LogicalType::Integer], &[index]) + } else { + unique_scan(&mut arena, "right", &[ty, LogicalType::Integer], &[index]) + }; if case == "inner_limit" { let Operator::TableScan(scan) = &mut right.operator else { unreachable!() @@ -867,16 +1395,29 @@ mod inner_join_tests { } let l = column(&left, 0, &mut arena); let r = column(&right, 0, &mut arena); - let mut plan = join(left, right, vec![(l, r)], None); + let mut keys = vec![(l, r)]; + if case == "composite_index" { + keys.push((l, column(&right, 1, &mut arena))); + } + let mut plan = join(left, right, keys, None); if case == "outer_join" { let Operator::Join(join) = &mut plan.operator else { unreachable!() }; join.join_type = JoinType::LeftOuter; } + if case == "join_decided" { + let mut option = PhysicalOption::new( + PlanImpl::HashJoin, + SortOption::None, + RowCountOption::Unknown, + ); + option.join_decided = true; + plan.physical_option = Some(option); + } let before = plan.clone(); assert!( - !ParameterizeInnerJoin.apply(&mut plan, &mut arena)?, + !ParameterizeInnerJoin.apply(&mut plan, &mut arena, None)?, "{case}" ); assert_eq!(plan, before, "{case}"); diff --git a/src/optimizer/rule/normalization/pushdown_limit.rs b/src/optimizer/rule/normalization/pushdown_limit.rs index e33dc91c..84f1a46c 100644 --- a/src/optimizer/rule/normalization/pushdown_limit.rs +++ b/src/optimizer/rule/normalization/pushdown_limit.rs @@ -14,11 +14,12 @@ use crate::errors::DatabaseError; use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::plan_utils::{only_child_mut, replace_with_only_child, wrap_child_with}; use crate::planner::operator::join::JoinType; use crate::planner::operator::limit::LimitOperator; use crate::planner::operator::Operator; -use crate::planner::LogicalPlan; +use crate::planner::{LogicalPlan, PlanArena}; pub struct LimitProjectTranspose; @@ -26,7 +27,8 @@ impl NormalizationRule for LimitProjectTranspose { fn apply( &self, plan: &mut LogicalPlan, - _: &mut crate::planner::PlanArena, + _: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let operator = std::mem::replace(&mut plan.operator, Operator::Dummy); @@ -46,6 +48,7 @@ impl NormalizationRule for LimitProjectTranspose { &mut child.operator, Operator::Limit(limit_op.clone()), )); + child.reset_output_schema_cache(); } } @@ -71,7 +74,8 @@ impl NormalizationRule for PushLimitThroughJoin { fn apply( &self, plan: &mut LogicalPlan, - _: &mut crate::planner::PlanArena, + _: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let limit_op = match &plan.operator { Operator::Limit(LimitOperator { @@ -110,7 +114,8 @@ impl NormalizationRule for PushLimitIntoScan { fn apply( &self, plan: &mut LogicalPlan, - _: &mut crate::planner::PlanArena, + _: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let (offset, limit) = match &plan.operator { Operator::Limit(limit_op) => (limit_op.offset, limit_op.limit), @@ -148,9 +153,9 @@ mod tests { let pipeline = HepOptimizerPipeline::builder() .before_batch( - "test_limit_project_transpose".to_string(), + "test_limit_project_transpose", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::LimitProjectTranspose], + &[NormalizationRuleImpl::LimitProjectTranspose], ) .build(); let best_plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; @@ -180,9 +185,9 @@ mod tests { let pipeline = HepOptimizerPipeline::builder() .before_batch( - "test_push_limit_through_join".to_string(), + "test_push_limit_through_join", HepBatchStrategy::once_topdown(), - vec![ + &[ NormalizationRuleImpl::LimitProjectTranspose, NormalizationRuleImpl::PushLimitThroughJoin, ], @@ -214,9 +219,9 @@ mod tests { let pipeline = HepOptimizerPipeline::builder() .before_batch( - "test_push_limit_into_table_scan".to_string(), + "test_push_limit_into_table_scan", HepBatchStrategy::once_topdown(), - vec![ + &[ NormalizationRuleImpl::LimitProjectTranspose, NormalizationRuleImpl::PushLimitIntoTableScan, ], diff --git a/src/optimizer/rule/normalization/pushdown_predicates.rs b/src/optimizer/rule/normalization/pushdown_predicates.rs index 6bdef695..ec8a2aa1 100644 --- a/src/optimizer/rule/normalization/pushdown_predicates.rs +++ b/src/optimizer/rule/normalization/pushdown_predicates.rs @@ -12,18 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::catalog::ColumnRef; use crate::errors::DatabaseError; use crate::expression::range_detacher::{DetachedPredicate, Range, RangeDetacher}; use crate::expression::visitor_mut::{ExprVisitorMut, PositionShift}; use crate::expression::{BinaryOperator, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; -use crate::optimizer::plan_utils::{replace_with_only_child, wrap_child_with}; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::optimizer::plan_utils::{predicate_scan_mut, replace_with_only_child, wrap_child_with}; use crate::planner::operator::filter::FilterOperator; use crate::planner::operator::join::extract_join_keys; use crate::planner::operator::join::{JoinCondition, JoinType}; use crate::planner::operator::{Operator, SortOption}; use crate::planner::{Childrens, ExprRef, LogicalPlan, PlanArena}; -use crate::types::index::{IndexInfo, IndexLookup, IndexMetaRef, IndexType}; +use crate::types::index::{IndexInfo, IndexLookup, IndexMetaRef}; +use crate::types::serialize::TupleValueSerializableImpl; use crate::types::value::DataValue; use crate::types::LogicalType; use std::ops::Bound; @@ -82,9 +85,19 @@ fn classify_join_filters( let mut right_filter = None; let mut common_filter = None; split_conjunctive_predicates(predicate, arena, &mut |expr, arena| { - if expr.all_referenced_columns(arena, |_, column| left_columns.contains(column))? { + let (mut left, mut right) = (true, true); + expr.all_referenced_columns(arena, |_, column, position| { + let side_holds = |columns: &[ColumnRef], position: Option| { + position.and_then(|position| columns.get(position)) == Some(column) + || columns.contains(column) + }; + left = left && side_holds(left_columns, Some(position)); + right = right && side_holds(right_columns, position.checked_sub(left_len)); + left || right + })?; + if left { append_filter(&mut left_filter, expr, arena); - } else if expr.all_referenced_columns(arena, |_, column| right_columns.contains(column))? { + } else if right { append_filter(&mut right_filter, expr, arena); } else { append_filter(&mut common_filter, expr, arena); @@ -133,163 +146,162 @@ impl NormalizationRule for PushPredicateThroughJoin { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { - let mut applied = false; - - let parent_replacement = { - let LogicalPlan { - operator, - childrens, - .. - } = plan; - let filter_op = match operator { - Operator::Filter(op) => op, - _ => return Ok(false), - }; - - let Childrens::Only(join_plan) = childrens.as_mut() else { - return Ok(false); - }; - let join_plan = join_plan.as_mut(); - - let join_type = match &join_plan.operator { - Operator::Join(op) => op.join_type, - _ => return Ok(false), - }; - - if !matches!( - join_type, - JoinType::Inner | JoinType::Cross | JoinType::LeftOuter | JoinType::RightOuter - ) { - return Ok(false); + let Operator::Filter(filter_op) = &plan.operator else { + return Ok(false); + }; + let (predicate, having) = (filter_op.predicate, filter_op.having); + let Childrens::Only(join_plan) = plan.childrens.as_mut() else { + return Ok(false); + }; + let Some((mut applied, common)) = push_through_join(join_plan, predicate, having, arena)? + else { + return Ok(false); + }; + match common { + Some(common) => { + plan.operator = Operator::Filter(filter(common, having)); + applied = true; } + None if applied => { + applied |= replace_with_only_child(plan); + } + None => {} + } + Ok(applied) + } +} - let (left_filter, mut right_filter, mut common_filter, left_len) = - classify_join_filters(filter_op.predicate, join_plan.childrens.as_mut(), arena)?; +fn filter(predicate: ExprRef, having: bool) -> FilterOperator { + FilterOperator { + predicate, + is_optimized: false, + having, + } +} - if join_type == JoinType::Cross { - if let Some(predicate) = common_filter { - let Childrens::Twins { left, right } = join_plan.childrens.as_mut() else { - return Ok(false); - }; - let left_schema = left.output_schema(arena).clone(); - let right_schema = right.output_schema(arena).clone(); - let mut keys = Vec::new(); - let mut residual = Vec::new(); - extract_join_keys( - predicate, - &mut keys, - &mut residual, - &left_schema, - &right_schema, - arena, - )?; - if !keys.is_empty() { - for (_, right) in &mut keys { - // WHERE uses joined positions; join keys use each input's local positions. - *right = right.clone_expression(arena)?; - PositionShift { - delta: -(left_len as isize), - } - .visit(right, arena)?; - } - if let Operator::Join(op) = &mut join_plan.operator { - op.join_type = JoinType::Inner; - op.on = JoinCondition::On { - on: keys, - filter: None, - }; - } - common_filter = reduce_filters(residual, filter_op.having, arena) - .map(|op| op.predicate); - applied = true; - } - } - } +fn push_through_join( + join_plan: &mut LogicalPlan, + predicate: ExprRef, + having: bool, + arena: &mut PlanArena, +) -> Result)>, DatabaseError> { + let join_type = match &join_plan.operator { + Operator::Join(op) => op.join_type, + _ => return Ok(None), + }; + if !matches!( + join_type, + JoinType::Inner | JoinType::Cross | JoinType::LeftOuter | JoinType::RightOuter + ) { + return Ok(None); + } + let mut applied = false; - let mut new_ops = (None, None, None); - match join_type { - JoinType::Inner | JoinType::Cross => { - if let Some(left_filter_op) = - reduce_filters(left_filter, filter_op.having, arena) - { - new_ops.0 = Some(Operator::Filter(left_filter_op)); - } + let (left_filter, mut right_filter, mut common_filter, left_len) = + classify_join_filters(predicate, join_plan.childrens.as_mut(), arena)?; - if let Some(expr) = &mut right_filter { - PositionShift { - delta: -(left_len as isize), - } - .visit(expr, arena)?; - } - if let Some(right_filter_op) = - reduce_filters(right_filter, filter_op.having, arena) - { - new_ops.1 = Some(Operator::Filter(right_filter_op)); - } - new_ops.2 = reduce_filters(common_filter, filter_op.having, arena) - .map(Operator::Filter); - } - JoinType::LeftOuter => { - if let Some(left_filter_op) = - reduce_filters(left_filter, filter_op.having, arena) - { - new_ops.0 = Some(Operator::Filter(left_filter_op)); + if join_type == JoinType::Cross { + if let Some(predicate) = common_filter { + let Childrens::Twins { left, right } = join_plan.childrens.as_mut() else { + return Ok(None); + }; + let left_schema = left.output_schema(arena).clone(); + let right_schema = right.output_schema(arena).clone(); + let mut keys = Vec::new(); + let mut residual = Vec::new(); + extract_join_keys( + predicate, + &mut keys, + &mut residual, + &left_schema, + &right_schema, + arena, + )?; + if !keys.is_empty() { + for (_, right) in &mut keys { + // WHERE uses joined positions; join keys use each input's local positions. + *right = right.clone_expression(arena)?; + PositionShift { + delta: -(left_len as isize), } - new_ops.2 = reduce_filters( - common_filter.into_iter().chain(right_filter), - filter_op.having, - arena, - ) - .map(Operator::Filter); + .visit(right, arena)?; } - JoinType::RightOuter => { - if let Some(expr) = &mut right_filter { - PositionShift { - delta: -(left_len as isize), - } - .visit(expr, arena)?; - } - if let Some(right_filter_op) = - reduce_filters(right_filter, filter_op.having, arena) - { - new_ops.1 = Some(Operator::Filter(right_filter_op)); - } - new_ops.2 = reduce_filters( - common_filter.into_iter().chain(left_filter), - filter_op.having, - arena, - ) - .map(Operator::Filter); + if let Operator::Join(op) = &mut join_plan.operator { + op.join_type = JoinType::Inner; + op.on = JoinCondition::On { + on: keys, + filter: None, + }; } - _ => {} - } - - if let Some(left_op) = new_ops.0 { - applied |= wrap_child_with(join_plan, 0, left_op); - } - - if let Some(right_op) = new_ops.1 { - applied |= wrap_child_with(join_plan, 1, right_op); - } - - new_ops.2 - }; - - match parent_replacement { - Some(common_op) => { - plan.operator = common_op; + common_filter = reduce_filters(residual, having, arena).map(|op| op.predicate); applied = true; } - None if applied => { - applied |= replace_with_only_child(plan); + } + } + + if let Some(expr) = &mut right_filter { + if matches!( + join_type, + JoinType::Inner | JoinType::Cross | JoinType::RightOuter + ) { + PositionShift { + delta: -(left_len as isize), } - _ => {} + .visit(expr, arena)?; } + } + let (left, right, common) = match join_type { + JoinType::Inner | JoinType::Cross => (left_filter, right_filter, common_filter), + JoinType::LeftOuter => ( + left_filter, + None, + reduce_filters(common_filter.into_iter().chain(right_filter), having, arena) + .map(|op| op.predicate), + ), + JoinType::RightOuter => ( + None, + right_filter, + reduce_filters(common_filter.into_iter().chain(left_filter), having, arena) + .map(|op| op.predicate), + ), + _ => unreachable!(), + }; + if let Some(left) = left { + applied |= push_into_child(join_plan, 0, left, having, arena)?; + } + if let Some(right) = right { + applied |= push_into_child(join_plan, 1, right, having, arena)?; + } + Ok(Some((applied, common))) +} - Ok(applied) +/// Pushes `predicate` into the `child_idx` input of `join_plan`, through it when it is a join. +fn push_into_child( + join_plan: &mut LogicalPlan, + child_idx: usize, + predicate: ExprRef, + having: bool, + arena: &mut PlanArena, +) -> Result { + let Childrens::Twins { left, right } = join_plan.childrens.as_mut() else { + return Ok(false); + }; + let child = if child_idx == 0 { left } else { right }; + let remaining = match push_through_join(child, predicate, having, arena)? { + Some((_, remaining)) => remaining, + None => Some(predicate), + }; + if let Some(remaining) = remaining { + wrap_child_with( + join_plan, + child_idx, + Operator::Filter(filter(remaining, having)), + ); } + Ok(true) } pub struct PushPredicateIntoScan; @@ -298,7 +310,8 @@ impl NormalizationRule for PushPredicateIntoScan { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let LogicalPlan { operator, @@ -312,8 +325,7 @@ impl NormalizationRule for PushPredicateIntoScan { let Childrens::Only(child) = childrens.as_mut() else { return Ok(false); }; - let child = child.as_mut(); - let Operator::TableScan(scan_op) = &mut child.operator else { + let Some((scan_op, _)) = predicate_scan_mut(child, arena) else { return Ok(false); }; @@ -338,16 +350,11 @@ impl NormalizationRule for PushPredicateIntoScan { else { return Err(DatabaseError::InvalidIndex); }; - let index_type = arena.index(*meta).ty; - let detached = match index_type { - IndexType::PrimaryKey { is_multiple: false } - | IndexType::Unique - | IndexType::Normal => { - RangeDetacher::for_index(*meta, 0, arena).detach(filter_op.predicate)? - } - IndexType::PrimaryKey { is_multiple: true } | IndexType::Composite => { - Self::composite_range(filter_op, *meta, ignore_prefix_len, arena)? - } + // Uniqueness does not determine key width: UNIQUE indexes can be composite too. + let detached = if arena.index(*meta).column_ids.len() > 1 { + Self::composite_range(filter_op, *meta, ignore_prefix_len, arena)? + } else { + RangeDetacher::for_index(*meta, 0, arena).detach(filter_op.predicate)? }; let Some(detached) = detached else { *lookup = None; @@ -358,49 +365,69 @@ impl NormalizationRule for PushPredicateIntoScan { *lookup = Some(IndexLookup::Static(detached.range)); *residual_predicate = detached.residual; - *covered_deserializers = None; - *cover_mapping = None; - - // try index covered - let mut mapping_slots = vec![usize::MAX; scan_op.columns.len()]; - let mut needs_mapping = false; - let index_meta = arena.index(*meta); - let index_column_types = match &index_meta.value_ty { - LogicalType::Tuple(tys) => tys, - ty => slice::from_ref(ty), - }; - let mut deserializers = Vec::with_capacity(index_meta.column_ids.len()); - - for (idx, column_id) in index_meta.column_ids.iter().enumerate() { - if let Some((scan_idx, column)) = - scan_op.columns.iter().enumerate().find(|(_, column)| { - arena - .column(**column) - .id() - .map(|id| id == *column_id) - .unwrap_or(false) - }) - { - mapping_slots[scan_idx] = idx; - needs_mapping |= scan_idx != idx; - deserializers.push(arena.column(*column).datatype().serializable()); - } else { - deserializers.push(index_column_types[idx].skip_serializable()); - } - } - - if mapping_slots.iter().all(|slot| *slot != usize::MAX) { - *covered_deserializers = Some(deserializers); - if needs_mapping { - *cover_mapping = Some(mapping_slots); - } - } + (*covered_deserializers, *cover_mapping) = + match index_covering(&scan_op.columns, *meta, arena) { + Some((deserializers, mapping)) => (Some(deserializers), mapping), + None => (None, None), + }; } Ok(changed) } } +/// Index-only read layout when `index` holds every scanned column: deserializers for the +/// index columns, plus a scan-to-index slot mapping when their orders differ. +pub(crate) fn index_covering( + columns: &[ColumnRef], + index: IndexMetaRef, + arena: &PlanArena, +) -> Option<(Vec, Option>)> { + if !is_index_covers(columns, index, arena) { + return None; + } + let mut mapping_slots = vec![usize::MAX; columns.len()]; + let mut needs_mapping = false; + let index_meta = arena.index(index); + let index_column_types = match &index_meta.value_ty { + LogicalType::Tuple(tys) => tys, + ty => slice::from_ref(ty), + }; + let mut deserializers = Vec::with_capacity(index_meta.column_ids.len()); + + for (idx, column_id) in index_meta.column_ids.iter().enumerate() { + if let Some((scan_idx, column)) = columns.iter().enumerate().find(|(_, column)| { + arena + .column(**column) + .id() + .map(|id| id == *column_id) + .unwrap_or(false) + }) { + mapping_slots[scan_idx] = idx; + needs_mapping |= scan_idx != idx; + deserializers.push(arena.column(*column).datatype().serializable()); + } else { + deserializers.push(index_column_types[idx].skip_serializable()); + } + } + + Some((deserializers, needs_mapping.then_some(mapping_slots))) +} + +pub(crate) fn is_index_covers( + columns: &[ColumnRef], + index: IndexMetaRef, + arena: &PlanArena, +) -> bool { + let index_meta = arena.index(index); + columns.iter().all(|column| { + arena + .column(*column) + .id() + .is_some_and(|id| index_meta.column_ids.contains(&id)) + }) +} + impl PushPredicateIntoScan { fn composite_range( op: &FilterOperator, @@ -477,7 +504,8 @@ impl NormalizationRule for PushJoinPredicateIntoScan { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let (join_type, filter_expr) = { let Operator::Join(join_op) = &mut plan.operator else { @@ -649,47 +677,206 @@ mod tests { } #[test] - fn test_push_predicate_into_scan() -> Result<(), DatabaseError> { + fn single_pass_filter_classification_preserves_reference_semantics() -> Result<(), DatabaseError> + { + use crate::expression::AliasType; + use crate::planner::operator::values::ValuesOperator; + let table_state = build_t1_table()?; let mut arena = PlanArena::new(&table_state.table_arena); - // 1 - c2 < 0 => c2 > 1 - let plan = - table_state.plan_with_arena("select * from t1 where -(1 - c2) > 0", &mut arena)?; + let table: TableName = "classification".to_string().into(); + let left_column = build_test_column(&mut arena, &table, 1, "left")?; + let right_column = build_test_column(&mut arena, &table, 2, "right")?; + let unknown_column = build_test_column(&mut arena, &table, 3, "unknown")?; + let left = cmp_predicate(&mut arena, BinaryOperator::Eq, left_column, 0, 1); + let right = cmp_predicate(&mut arena, BinaryOperator::Eq, right_column, 1, 1); + let unknown = cmp_predicate(&mut arena, BinaryOperator::Eq, unknown_column, 2, 1); + let common = arena.alloc_expression(ScalarExpression::Binary { + op: BinaryOperator::Or, + left_expr: left, + right_expr: right, + evaluator: None, + ty: LogicalType::Boolean, + }); + let constant = arena.alloc_expression(ScalarExpression::Constant(DataValue::Boolean(true))); + let alias = arena.alloc_expression(ScalarExpression::Alias { + expr: right, + alias: AliasType::Expr(left), + }); + for right_columns in [vec![right_column], vec![left_column, right_column]] { + let mut children = Childrens::Twins { + left: Box::new(LogicalPlan::new( + Operator::Values(ValuesOperator::new(Vec::new(), 0, vec![left_column])), + Childrens::None, + )), + right: Box::new(LogicalPlan::new( + Operator::Values(ValuesOperator::new(Vec::new(), 0, right_columns.clone())), + Childrens::None, + )), + }; + for expr in [left, right, unknown, common, constant, alias] { + let expected = if expr + .all_referenced_columns(&arena, |_, column, _| *column == left_column)? + { + (Some(expr), None, None, 1) + } else if expr + .all_referenced_columns(&arena, |_, column, _| right_columns.contains(column))? + { + (None, Some(expr), None, 1) + } else { + (None, None, Some(expr), 1) + }; + assert_eq!( + super::classify_join_filters(expr, &mut children, &mut arena)?, + expected + ); + } + } + Ok(()) + } - let best_plan = apply_pipeline( - plan, - HepOptimizerPipeline::builder() - .before_batch( - "simplify_filter".to_string(), - HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::SimplifyFilter], - ) - .before_batch( - "test_push_predicate_into_scan".to_string(), - HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushPredicateIntoScan], - ), - &mut arena, - )?; + #[test] + fn test_push_predicate_into_scan() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + // 1 - c2 < 0 => c2 > 1, directly and through scalar initialization. + for (sql, through_init) in [ + ("select * from t1 where -(1 - c2) > 0", false), + ( + "select * from t1 where -(1 - c2) > 0 and c1 = (select 1)", + true, + ), + ] { + let plan = table_state.plan_with_arena(sql, &mut arena)?; + let best_plan = apply_pipeline( + plan, + HepOptimizerPipeline::builder() + .before_batch( + "simplify_filter", + HepBatchStrategy::once_topdown(), + &[NormalizationRuleImpl::SimplifyFilter], + ) + .before_batch( + "test_push_predicate_into_scan", + HepBatchStrategy::once_topdown(), + &[NormalizationRuleImpl::PushPredicateIntoScan], + ), + &mut arena, + )?; - let scan_op = best_plan.childrens.pop_only().childrens.pop_only(); - if let Operator::TableScan(op) = &scan_op.operator { - let mock_range = Range::Scope { - min: Bound::Excluded(DataValue::Int32(1)), - max: Bound::Unbounded, + let mut scan = best_plan.childrens.pop_only().childrens.pop_only(); + if through_init { + assert!(matches!(scan.operator, Operator::ScalarQueryInit(_))); + scan = scan.childrens.pop_twins().0; + } + let Operator::TableScan(op) = &scan.operator else { + panic!("expected table scan") }; - assert_eq!( op.index_infos[0].lookup, - Some(IndexLookup::Static(mock_range)) + Some(IndexLookup::Static(Range::Scope { + min: Bound::Excluded(DataValue::Int32(1)), + max: Bound::Unbounded, + })) ); - } else { - unreachable!("Should be a filter operator") } Ok(()) } + #[test] + fn test_unique_composite_index_uses_tuple_ranges() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let table_name: TableName = "mock_table".to_string().into(); + let c1 = build_test_column(&mut arena, &table_name, 1, "c1")?; + let c2 = build_test_column(&mut arena, &table_name, 2, "c2")?; + let meta = arena.alloc_index(IndexMeta { + id: 0, + column_ids: vec![1, 2], + table_name: table_name.clone(), + pk_ty: LogicalType::Integer, + value_ty: LogicalType::Tuple(vec![LogicalType::Integer, LogicalType::Integer]), + name: "unique_c1_c2".to_string(), + ty: IndexType::Unique, + }); + let first_eq = cmp_predicate(&mut arena, BinaryOperator::Eq, c1, 0, 2); + let second_eq = cmp_predicate(&mut arena, BinaryOperator::Eq, c2, 1, 3); + let full_eq = and_predicate(&mut arena, first_eq, second_eq); + let first_range = cmp_predicate(&mut arena, BinaryOperator::GtEq, c1, 0, 2); + for (predicate, expected, prefix_len) in [ + ( + first_range, + Range::Scope { + min: Bound::Included(DataValue::Tuple(vec![DataValue::Int32(2)])), + max: Bound::Unbounded, + }, + 0, + ), + ( + first_eq, + Range::Scope { + min: Bound::Included(DataValue::Tuple(vec![DataValue::Int32(2)])), + max: Bound::Included(DataValue::Tuple(vec![DataValue::Int32(2)])), + }, + 1, + ), + ( + full_eq, + Range::Eq(DataValue::Tuple(vec![ + DataValue::Int32(2), + DataValue::Int32(3), + ])), + 2, + ), + ] { + let scan = LogicalPlan::new( + Operator::TableScan(TableScanOperator { + table_name: table_name.clone(), + columns: vec![c1, c2], + limit: (None, None), + index_infos: vec![IndexInfo { + meta, + sort_option: SortOption::OrderBy { + fields: vec![], + ignore_prefix_len: 0, + }, + lookup: None, + residual_predicate: None, + covered_deserializers: None, + cover_mapping: None, + sort_elimination_hint: None, + stream_aggregate_hint: None, + }], + with_pk: false, + }), + Childrens::None, + ); + let plan = apply_pipeline( + FilterOperator::build(predicate, scan, false), + HepOptimizerPipeline::builder().before_batch( + "unique_composite_range", + HepBatchStrategy::fix_point_topdown(1), + &[NormalizationRuleImpl::PushPredicateIntoScan], + ), + &mut arena, + )?; + let scan = plan.childrens.pop_only(); + let Operator::TableScan(op) = scan.operator else { + unreachable!() + }; + assert_eq!( + op.index_infos[0].lookup, + Some(IndexLookup::Static(expected)) + ); + assert_eq!(op.index_infos[0].residual_predicate, None); + assert!(matches!(op.index_infos[0].sort_option, + SortOption::OrderBy { ignore_prefix_len, .. } if ignore_prefix_len == prefix_len)); + assert_eq!(arena.index(meta).ty, IndexType::Unique); + } + Ok(()) + } + #[test] fn test_composite_range_consumes_prefix_equalities_continuously() -> Result<(), DatabaseError> { let table_name: TableName = ::std::sync::Arc::from("mock_table"); @@ -927,9 +1114,9 @@ mod tests { let best_plan = apply_pipeline( filter_plan, HepOptimizerPipeline::builder().before_batch( - "push_cover_mapping".to_string(), + "push_cover_mapping", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushPredicateIntoScan], + &[NormalizationRuleImpl::PushPredicateIntoScan], ), &mut arena, )?; @@ -990,9 +1177,9 @@ mod tests { let best_plan = apply_pipeline( plan, HepOptimizerPipeline::builder().before_batch( - "test_push_predicate_through_join".to_string(), + "test_push_predicate_through_join", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushPredicateThroughJoin], + &[NormalizationRuleImpl::PushPredicateThroughJoin], ), &mut arena, )?; @@ -1040,9 +1227,9 @@ mod tests { let best_plan = apply_pipeline( plan, HepOptimizerPipeline::builder().before_batch( - "test_push_predicate_through_join".to_string(), + "test_push_predicate_through_join", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushPredicateThroughJoin], + &[NormalizationRuleImpl::PushPredicateThroughJoin], ), &mut arena, )?; @@ -1084,15 +1271,15 @@ mod tests { for predicate in ["c1 = c3 or c2 = c4", "c1 < c3", "c1 > 0 and c3 < 9"] { let mut arena = PlanArena::new(&tables.table_arena); let plan = tables.plan_with_arena( - &format!("select * from t1, t2 where {predicate}"), + format!("select * from t1, t2 where {predicate}"), &mut arena, )?; let plan = apply_pipeline( plan, HepOptimizerPipeline::builder().before_batch( - "cross join pushdown".into(), + "cross join pushdown", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushPredicateThroughJoin], + &[NormalizationRuleImpl::PushPredicateThroughJoin], ), &mut arena, )?; @@ -1123,9 +1310,9 @@ mod tests { let best_plan = apply_pipeline( plan, HepOptimizerPipeline::builder().before_batch( - "test_push_predicate_through_join".to_string(), + "test_push_predicate_through_join", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushPredicateThroughJoin], + &[NormalizationRuleImpl::PushPredicateThroughJoin], ), &mut arena, )?; @@ -1178,9 +1365,9 @@ mod tests { let mut best_plan = apply_pipeline( plan, HepOptimizerPipeline::builder().before_batch( - "push_join_predicate_into_scan".to_string(), + "push_join_predicate_into_scan", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushJoinPredicateIntoScan], + &[NormalizationRuleImpl::PushJoinPredicateIntoScan], ), &mut arena, )?; @@ -1254,9 +1441,9 @@ mod tests { let mut best_plan = apply_pipeline( plan, HepOptimizerPipeline::builder().before_batch( - "push_join_predicate_into_scan".to_string(), + "push_join_predicate_into_scan", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushJoinPredicateIntoScan], + &[NormalizationRuleImpl::PushJoinPredicateIntoScan], ), &mut arena, )?; @@ -1302,9 +1489,9 @@ mod tests { let mut best_plan = apply_pipeline( plan, HepOptimizerPipeline::builder().before_batch( - "push_join_predicate_into_scan".to_string(), + "push_join_predicate_into_scan", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::PushJoinPredicateIntoScan], + &[NormalizationRuleImpl::PushJoinPredicateIntoScan], ), &mut arena, )?; diff --git a/src/optimizer/rule/normalization/simplification.rs b/src/optimizer/rule/normalization/simplification.rs index 9de97e4e..26a76493 100644 --- a/src/optimizer/rule/normalization/simplification.rs +++ b/src/optimizer/rule/normalization/simplification.rs @@ -16,19 +16,27 @@ use crate::errors::DatabaseError; use crate::expression::simplify::{ConstantCalculator, Simplify}; use crate::expression::visitor_mut::ExprVisitorMut; use crate::optimizer::core::rule::NormalizationRule; -use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::planner::operator::visitor_mut::{ + ChildrenVisitMut, OperatorExprVisitorMut, OperatorVisitorMut, +}; use crate::planner::operator::Operator; -use crate::planner::{Childrens, LogicalPlan}; +use crate::planner::PlanArena; +use crate::planner::{Childrens, LogicalPlan, MetaArena}; #[derive(Copy, Clone)] pub struct ConstantCalculation; -pub(crate) fn constant_calculation_current( +pub(crate) fn constant_calculation_current( plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut A, ) -> Result<(), DatabaseError> { let mut calculator = ConstantCalculator::new(arena); - OperatorExprVisitorMut::new(&mut calculator, arena).visit_operator(&mut plan.operator, None) + OperatorExprVisitorMut::new(&mut calculator, arena).visit_operator( + &mut plan.operator, + None, + &mut ChildrenVisitMut::empty(), + ) } impl ConstantCalculation { @@ -54,7 +62,8 @@ impl NormalizationRule for ConstantCalculation { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { Self::_apply(plan, arena)?; Ok(true) @@ -82,7 +91,8 @@ impl NormalizationRule for SimplifyFilter { fn apply( &self, plan: &mut LogicalPlan, - arena: &mut crate::planner::PlanArena, + arena: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { if let Operator::Filter(filter_op) = &mut plan.operator { if filter_op.is_optimized { @@ -110,25 +120,27 @@ mod test { use crate::errors::DatabaseError; use crate::expression::range_detacher::{Range, RangeDetacher}; + use crate::expression::simplify::Simplify; + use crate::expression::visitor_mut::ExprVisitorMut; use crate::expression::{BinaryOperator, ScalarExpression, UnaryOperator}; use crate::optimizer::heuristic::batch::HepBatchStrategy; use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline; use crate::optimizer::rule::normalization::NormalizationRuleImpl; use crate::planner::operator::Operator; - use crate::planner::{LogicalPlan, PlanArena}; + use crate::planner::{LogicalPlan, PlanArena, TableArenaCell}; use crate::types::value::DataValue; use crate::types::{ColumnId, LogicalType}; use std::collections::Bound; fn run_with_single_batch( plan: LogicalPlan, - name: &str, + name: &'static str, strategy: HepBatchStrategy, - rules: Vec, + rules: &[NormalizationRuleImpl], arena: &mut PlanArena, ) -> Result { HepOptimizerPipeline::builder() - .before_batch(name.to_string(), strategy, rules) + .before_batch(name, strategy, rules) .build() .instantiate(plan) .find_best(None, arena) @@ -146,9 +158,9 @@ mod test { let best_plan = HepOptimizerPipeline::builder() .before_batch( - "test_simplification".to_string(), + "test_simplification", HepBatchStrategy::once_topdown(), - vec![ + &[ NormalizationRuleImpl::SimplifyFilter, NormalizationRuleImpl::ConstantCalculation, ], @@ -202,7 +214,7 @@ mod test { plan, "test_constant_cast_elimination", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::ConstantCalculation], + &[NormalizationRuleImpl::ConstantCalculation], &mut arena, )?; @@ -234,10 +246,6 @@ mod test { #[test] fn test_constant_is_null_elimination() -> Result<(), DatabaseError> { - use crate::expression::simplify::Simplify; - use crate::expression::visitor_mut::ExprVisitorMut; - use crate::planner::TableArenaCell; - let table_arena = TableArenaCell::default(); let mut arena = PlanArena::new(&table_arena); for (value, negated, expected) in [ @@ -305,7 +313,7 @@ mod test { plan, "test_simplify_filter", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::SimplifyFilter], + &[NormalizationRuleImpl::SimplifyFilter], &mut arena, )?; @@ -377,7 +385,7 @@ mod test { plan, "test_simplify_filter", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::SimplifyFilter], + &[NormalizationRuleImpl::SimplifyFilter], &mut arena, )?; @@ -431,7 +439,7 @@ mod test { plan, "test_non_comparison_does_not_apply_pending_rearrangement", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::SimplifyFilter], + &[NormalizationRuleImpl::SimplifyFilter], &mut arena, )?; let filter = best_plan.childrens.pop_only(); @@ -473,7 +481,7 @@ mod test { plan.clone(), "test_simplify_filter", HepBatchStrategy::once_topdown(), - vec![NormalizationRuleImpl::SimplifyFilter], + &[NormalizationRuleImpl::SimplifyFilter], arena, )?; @@ -654,6 +662,64 @@ mod test { Ok(()) } + #[test] + fn constant_in_expansion_uses_original_list_length() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + for (list, kept) in [ + ("1,2,3,4,5", false), + ("1,2,3,4,5,6", true), + ("1,1,1,1,1,NULL", true), + ("NULL,NULL,NULL,NULL,NULL,NULL", true), + ] { + for negated in ["", "not "] { + let mut arena = PlanArena::new(&table_state.table_arena); + let plan = table_state.plan_with_arena( + format!("select * from t1 where c1 {negated}in ({list})"), + &mut arena, + )?; + let Operator::Filter(filter) = plan.childrens.pop_only().operator else { + unreachable!(); + }; + let mut predicate = filter.predicate; + Simplify.visit(&mut predicate, &mut arena)?; + assert_eq!( + matches!(arena.expression(predicate), ScalarExpression::In { .. }), + kept, + "{negated}IN ({list})" + ); + } + } + Ok(()) + } + + #[test] + fn unsafe_large_in_keeps_pairwise_expansion() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + for predicate in [ + "c1 in (1,2,3,4,5,2147483648)", + "c1 in (1,2,3,4,5,c1)", + "(c1 + 1) in (1,2,3,4,5,6)", + "cast(c1 as double) in (1.0,2.0,3.0,4.0,5.0,6.0)", + ] { + let mut arena = PlanArena::new(&table_state.table_arena); + let plan = table_state + .plan_with_arena(format!("select * from t1 where {predicate}"), &mut arena)?; + let Operator::Filter(filter) = plan.childrens.pop_only().operator else { + unreachable!(); + }; + let mut expression = filter.predicate; + Simplify.visit(&mut expression, &mut arena)?; + assert!( + matches!( + arena.expression(expression), + ScalarExpression::Binary { .. } + ), + "{predicate}" + ); + } + Ok(()) + } + #[test] fn test_simplify_filter_column_in() -> Result<(), DatabaseError> { let table_state = build_t1_table()?; diff --git a/src/optimizer/rule/normalization/top_k.rs b/src/optimizer/rule/normalization/top_k.rs index 73cd74a1..7b7fbc64 100644 --- a/src/optimizer/rule/normalization/top_k.rs +++ b/src/optimizer/rule/normalization/top_k.rs @@ -14,10 +14,11 @@ use crate::errors::DatabaseError; use crate::optimizer::core::rule::NormalizationRule; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; use crate::optimizer::plan_utils::{only_child_mut, replace_with_only_child}; use crate::planner::operator::top_k::TopKOperator; use crate::planner::operator::Operator; -use crate::planner::LogicalPlan; +use crate::planner::{LogicalPlan, PlanArena}; pub struct TopK; @@ -25,7 +26,8 @@ impl NormalizationRule for TopK { fn apply( &self, plan: &mut LogicalPlan, - _: &mut crate::planner::PlanArena, + _: &mut PlanArena, + _loader: Option<&StatisticMetaLoader<'_>>, ) -> Result { let (offset, limit) = match &plan.operator { Operator::Limit(op) => match op.limit { diff --git a/src/orm/README.md b/src/orm/README.md index d10e73f1..4041b524 100644 --- a/src/orm/README.md +++ b/src/orm/README.md @@ -88,7 +88,7 @@ let users = database.execute(&plan, [(1, DataValue::Int32(1))])? # Ok::<(), Box>(()) ``` -`DBTransaction::prepare` and `DBTransaction::execute` provide the same +`Transaction::prepare` and `Transaction::execute` provide the same workflow within an explicit transaction. The plan must be executed against the database that created it; DDL and ANALYZE cannot be prepared. diff --git a/src/orm/dml.rs b/src/orm/dml.rs index 4adb5cae..1cb39aa3 100644 --- a/src/orm/dml.rs +++ b/src/orm/dml.rs @@ -71,7 +71,7 @@ impl Database { } } -impl<'a, S: Storage> DBTransaction<'a, S> { +impl<'a, S: Storage> Transaction<'a, S> { /// Refreshes optimizer statistics for the model table inside the current transaction. pub fn analyze(&mut self) -> Result<(), DatabaseError> { orm_analyze::<_, M>(self) diff --git a/src/orm/dql.rs b/src/orm/dql.rs index f90b0509..fdd29195 100644 --- a/src/orm/dql.rs +++ b/src/orm/dql.rs @@ -22,7 +22,7 @@ impl Database { { let transaction = self .storage - .transaction_with_isolation(self.transaction_isolation)?; + .transaction_with_isolation(self.transaction_isolation, true)?; self.state .prepare_plan_with(params, &transaction, |binder, arena| { let mut context = OrmContext { binder, arena }; @@ -44,7 +44,7 @@ impl Database { >, ) -> Result, { - bind_orm_context(self, build) + bind_orm_context(self, false, build) } /// Explains a binder-backed plan built inside a closure. @@ -185,7 +185,7 @@ impl Database { } } -impl<'a, S: Storage> DBTransaction<'a, S> { +impl<'a, S: Storage> Transaction<'a, S> { /// Prepares an ORM-built plan for repeated execution inside this transaction. pub fn prepare( &self, @@ -228,7 +228,7 @@ impl<'a, S: Storage> DBTransaction<'a, S> { >, ) -> Result, { - bind_orm_context(self, build) + bind_orm_context(self, false, build) } /// Explains a binder-backed plan inside the current transaction. diff --git a/src/orm/mod.rs b/src/orm/mod.rs index 52c6e582..02d8c1bb 100644 --- a/src/orm/mod.rs +++ b/src/orm/mod.rs @@ -6,7 +6,7 @@ use crate::binder::{ }; use crate::catalog::{ColumnCatalog, ColumnRef, TableCatalog, TableName}; use crate::db::{ - BindSource, DBTransaction, Database, DatabaseIter, OrmIter, ResultIter, TransactionIter, + BindSource, Database, DatabaseIter, OrmIter, ResultIter, Transaction, TransactionIter, }; use crate::errors::DatabaseError; pub use crate::expression::agg::AggKind; @@ -17,8 +17,8 @@ use crate::planner::operator::join::JoinType; use crate::planner::operator::mark_apply::MarkApplyQuantifier; use crate::planner::operator::sort::SortField; use crate::planner::{ExprRef, LogicalPlan, PlanArena}; -use crate::storage::{Storage, Transaction}; -use crate::types::tuple::{SchemaView, Tuple}; +use crate::storage::{Storage, StorageTransaction}; +use crate::types::tuple::{SchemaView, TupleImpl}; use crate::types::value::DataValue; use crate::types::CharLengthUnits; use crate::types::LogicalType; @@ -71,7 +71,7 @@ impl OrmField { } } -/// One row returned by [`Database::describe`] or [`DBTransaction::describe`]. +/// One row returned by [`Database::describe`] or [`Transaction::describe`]. #[derive(Debug, Clone, PartialEq, Eq)] pub struct DescribeColumn { pub field: String, @@ -83,7 +83,10 @@ pub struct DescribeColumn { } impl FromQueryRow for DescribeColumn { - fn from_query_row(_: &SchemaView<'_, '_>, tuple: &mut Tuple) -> Result { + fn from_query_row( + _: &SchemaView<'_, '_>, + tuple: &mut TupleImpl, + ) -> Result { let field = describe_text_value(take_projected_value(tuple, 0)); let data_type = describe_text_value(take_projected_value(tuple, 1)); let len = describe_text_value(take_projected_value(tuple, 2)); @@ -219,7 +222,7 @@ impl FieldSort { #[doc(hidden)] pub trait BindOrmScalar<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_scalar( @@ -230,7 +233,7 @@ where impl<'bind, 'parent, 'arena, T, A, M, V> BindOrmScalar<'bind, 'parent, 'arena, T, A> for Field where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_scalar( @@ -243,7 +246,7 @@ where impl<'bind, 'parent, 'arena, T, A> BindOrmScalar<'bind, 'parent, 'arena, T, A> for ScalarExpression where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_scalar( @@ -257,7 +260,7 @@ where impl<'bind, 'parent, 'arena, T, A> BindOrmScalar<'bind, 'parent, 'arena, T, A> for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_scalar( @@ -271,7 +274,7 @@ where #[doc(hidden)] pub trait BindOrmSort<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_sort<'scope>( @@ -282,7 +285,7 @@ where impl<'bind, 'parent, 'arena, T, A, M, V> BindOrmSort<'bind, 'parent, 'arena, T, A> for Field where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_sort<'scope>( @@ -296,7 +299,7 @@ where impl<'bind, 'parent, 'arena, T, A, M, V> BindOrmSort<'bind, 'parent, 'arena, T, A> for FieldSort where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_sort<'scope>( @@ -312,7 +315,7 @@ where impl<'bind, 'parent, 'arena, T, A> BindOrmSort<'bind, 'parent, 'arena, T, A> for ScalarExpression where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_sort<'scope>( @@ -325,7 +328,7 @@ where impl<'bind, 'parent, 'arena, T, A> BindOrmSort<'bind, 'parent, 'arena, T, A> for SortField where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_sort<'scope>( @@ -339,7 +342,7 @@ where impl<'bind, 'parent, 'arena, T, A> BindOrmSort<'bind, 'parent, 'arena, T, A> for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_sort<'scope>( @@ -379,7 +382,7 @@ impl WindowSpec { impl<'bind, 'parent, 'arena, T, A> From> for ExprRef where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn from(expr: CtxExpression<'bind, 'parent, 'arena, T, A>) -> Self { @@ -417,7 +420,7 @@ where #[doc(hidden)] pub trait BindOrmScalarList<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn bind_scalar_list( @@ -432,7 +435,7 @@ macro_rules! impl_bind_orm_scalar_list { impl<'bind, 'parent, 'arena, Tx, Args, $($name),+> BindOrmScalarList<'bind, 'parent, 'arena, Tx, Args> for ($($name,)+) where - Tx: Transaction, + Tx: StorageTransaction, Args: AsRef<[(usize, LogicalType)]>, $($name: BindOrmScalar<'bind, 'parent, 'arena, Tx, Args>,)+ { @@ -484,7 +487,7 @@ macro_rules! impl_quantified_subquery_methods { #[allow(clippy::type_complexity)] struct ExprBindScopeHandle<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: NonNull>, @@ -494,7 +497,7 @@ where impl<'bind, 'parent, 'arena, T, A> Clone for ExprBindScopeHandle<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn clone(&self) -> Self { @@ -504,14 +507,14 @@ where impl<'bind, 'parent, 'arena, T, A> Copy for ExprBindScopeHandle<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { } impl<'bind, 'parent, 'arena, T, A> ExprBindScopeHandle<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn new<'ctx>(scope: &ExprBindScope<'ctx, 'bind, 'parent, 'arena, T, A>) -> Self { @@ -688,7 +691,7 @@ where /// internal scope handle is private. pub struct CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { expr: ExprRef, @@ -697,7 +700,7 @@ where impl<'bind, 'parent, 'arena, T, A> CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub fn into_scalar(self) -> ExprRef { @@ -968,7 +971,7 @@ where impl<'bind, 'parent, 'arena, T, A> fmt::Debug for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -978,7 +981,7 @@ where impl<'bind, 'parent, 'arena, T, A> Clone for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn clone(&self) -> Self { @@ -991,7 +994,7 @@ where impl<'bind, 'parent, 'arena, T, A> PartialEq for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn eq(&self, other: &Self) -> bool { @@ -1001,14 +1004,14 @@ where impl<'bind, 'parent, 'arena, T, A> Eq for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { } impl<'bind, 'parent, 'arena, T, A> Hash for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn hash(&self, state: &mut H) { @@ -1018,7 +1021,7 @@ where impl<'bind, 'parent, 'arena, T, A> IntoOrmExpression for CtxExpression<'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn into_orm_expression(self) -> OrmExpression { @@ -1026,7 +1029,11 @@ where } } -fn bind_orm_context<'a, E, F>(executor: E, build: F) -> Result +fn bind_orm_context<'a, E, F>( + executor: E, + read_only: bool, + build: F, +) -> Result where E: BindSource<'a>, F: for<'ctx, 'bind, 'parent, 'arena> FnOnce( @@ -1041,7 +1048,7 @@ where ) -> Result, { static EMPTY_BIND_PARAMS: &[(usize, LogicalType)] = &[]; - executor.execute(|state, tx| { + executor.execute(read_only, |state, tx| { state.build_plan(EMPTY_BIND_PARAMS, tx, |binder, arena| { let mut context = OrmContext { binder, arena }; build(&mut context) @@ -1072,12 +1079,12 @@ where /// Binder-backed ORM query context. /// -/// This context is created by [`Database::bind`] or [`DBTransaction::bind`]. Query construction inside +/// This context is created by [`Database::bind`] or [`Transaction::bind`]. Query construction inside /// the closure binds directly into [`ScalarExpression`] and [`LogicalPlan`] /// values; it does not build an ORM expression tree first. pub struct OrmContext<'ctx, 'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'ctx mut Binder<'bind, 'parent, T, A>, @@ -1087,7 +1094,7 @@ where /// Narrow expression binding scope borrowed from an [`OrmContext`]. pub struct ExprBindScope<'ctx, 'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'ctx mut Binder<'bind, 'parent, T, A>, @@ -1096,7 +1103,7 @@ where pub struct UpdateBindScope<'ctx, 'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { binder: &'ctx mut Binder<'bind, 'parent, T, A>, @@ -1107,7 +1114,7 @@ where impl<'ctx, 'bind, 'parent, 'arena, T, A> OrmContext<'ctx, 'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub fn from<'scope, M: Model>( @@ -1328,7 +1335,7 @@ where impl<'ctx, 'bind, 'parent, 'arena, T, A> ExprBindScope<'ctx, 'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { fn handle(&self) -> ExprBindScopeHandle<'bind, 'parent, 'arena, T, A> { @@ -1856,7 +1863,7 @@ where impl<'ctx, 'bind, 'parent, 'arena, T, A> UpdateBindScope<'ctx, 'bind, 'parent, 'arena, T, A> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { pub fn set_value(&mut self, field: Field, value: D) -> Result<(), DatabaseError> @@ -1952,7 +1959,7 @@ where impl<'scope_ctx, 'bind, 'parent, 'arena, T, A, M> BindPlanFrom<'scope_ctx, 'bind, 'parent, 'arena, T, A, M> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, M: Model, { @@ -2062,7 +2069,7 @@ where bind_orm_source(&mut right_binder, source, Some(join_type), self.arena)?; (right_plan, right_binder.context) }; - self.binder.extend(right_context); + self.binder.extend(right_context, self.arena)?; let on = with_query_bind_step!(self.binder, QueryBindStep::From, { let mut scope = self.expr_scope(); let handle = scope.handle(); @@ -2405,7 +2412,7 @@ where impl<'scope_ctx, 'bind, 'parent, 'arena, T, A, M> BindPlanSelectList<'scope_ctx, 'bind, 'parent, 'arena, T, A, M> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, M: Model, { @@ -2574,7 +2581,7 @@ pub trait Projection: FromQueryRow { relation: &str, ) -> Result, DatabaseError> where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>; } @@ -2592,7 +2599,7 @@ fn bind_orm_source<'bind, 'parent, 'arena, T, A>( arena: &mut PlanArena<'arena>, ) -> Result where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { let alias = orm_table_alias(&source); @@ -2606,7 +2613,7 @@ fn bind_orm_target_column<'bind, 'parent, 'arena, T, A>( arena: &mut PlanArena<'arena>, ) -> Result where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { match binder.bind_column_ref_by_name(None, column_name, Some(source_name), arena)? { @@ -2624,7 +2631,7 @@ fn bind_orm_insert_plan<'bind, 'parent, 'arena, T, A>( arena: &mut PlanArena<'arena>, ) -> Result where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, { let table_name: TableName = table_name.into(); @@ -2692,14 +2699,14 @@ fn bind_orm_insert_model<'bind, 'parent, 'arena, T, A, M>( arena: &mut PlanArena<'arena>, ) -> Result where - T: Transaction, + T: StorageTransaction, A: AsRef<[(usize, LogicalType)]>, M: Model, { let table_name: TableName = M::table_name().into(); let source = binder .context - .source_and_bind(table_name.clone(), None, None, false)? + .source_and_bind(table_name.clone(), None, None, false, arena)? .ok_or(DatabaseError::TableNotFound)?; let params = params.into_iter().collect::>(); let mut schema_ref = Vec::with_capacity(M::fields().len()); @@ -2810,7 +2817,7 @@ pub trait FromDataValue: Sized { /// ``` pub trait FromQueryTuple: Sized { /// Decodes one projected tuple into `Self`. - fn from_query_tuple(tuple: &mut Tuple) -> Result; + fn from_query_tuple(tuple: &mut TupleImpl) -> Result; } /// Conversion trait from a query result row into a Rust value. @@ -2820,7 +2827,7 @@ pub trait FromQueryRow: Sized { /// Decodes one result row into `Self`. fn from_query_row( schema: &SchemaView<'_, '_>, - tuple: &mut Tuple, + tuple: &mut TupleImpl, ) -> Result; } @@ -2994,7 +3001,7 @@ pub trait DecimalType {} #[doc(hidden)] pub fn take_value_at( - tuple: &mut Tuple, + tuple: &mut TupleImpl, index: Option, field_name: &str, ) -> Result { @@ -3002,7 +3009,7 @@ pub fn take_value_at( name: field_name.to_string(), span: None, })?; - let value = tuple.values.get_mut(idx).ok_or(DatabaseError::MisMatch( + let value = tuple.get_mut(idx).ok_or(DatabaseError::MisMatch( "the query result schema", "the query result tuple", ))?; @@ -3246,9 +3253,9 @@ macro_rules! impl_from_query_tuple { $($name: FromDataValue,)+ { #[allow(non_snake_case)] - fn from_query_tuple(tuple: &mut Tuple) -> Result { + fn from_query_tuple(tuple: &mut TupleImpl) -> Result { let expected_len = [$(stringify!($name)),+].len(); - if tuple.values.len() != expected_len { + if tuple.len() != expected_len { return Err(DatabaseError::MisMatch( "the expected tuple projection width", "the query result", @@ -3360,9 +3367,8 @@ fn invalid_from_data_value(value: &DataValue) -> DatabaseError { )) } -fn take_projected_value(tuple: &mut Tuple, index: usize) -> Option { +fn take_projected_value(tuple: &mut TupleImpl, index: usize) -> Option { tuple - .values .get_mut(index) .map(|value| std::mem::replace(value, DataValue::Null)) } @@ -3378,8 +3384,8 @@ fn extract_projected_data_value( convert_projected_value::(value) } -fn extract_value_from_tuple(tuple: &mut Tuple) -> Result { - let value = if tuple.values.len() == 1 { +fn extract_value_from_tuple(tuple: &mut TupleImpl) -> Result { + let value = if tuple.len() == 1 { take_projected_value(tuple, 0).expect("checked one projected expression") } else { return Err(DatabaseError::MisMatch( @@ -3391,13 +3397,13 @@ fn extract_value_from_tuple(tuple: &mut Tuple) -> Result(value) } -fn extract_projected_tuple(tuple: &mut Tuple) -> Result { +fn extract_projected_tuple(tuple: &mut TupleImpl) -> Result { T::from_query_tuple(tuple) } fn orm_analyze<'a, E: BindSource<'a>, M: Model>(executor: E) -> Result<(), DatabaseError> { executor - .execute(|state, tx| { + .execute(false, |state, tx| { state.build_plan(&[], tx, |binder, arena| { binder.bind_analyze(M::table_name().into(), arena) }) @@ -3411,7 +3417,7 @@ fn orm_insert<'a, E: BindSource<'a>, M: Model>( ) -> Result<(), DatabaseError> { let params = model.params(); executor - .execute(|state, tx| { + .execute(false, |state, tx| { state.build_plan(&[], tx, |binder, arena| { bind_orm_insert_model::<_, _, M>(binder, params, arena) }) @@ -3425,7 +3431,7 @@ fn orm_get<'a, E: BindSource<'a>, M: Model>( ) -> Result, DatabaseError> { let primary_key = M::primary_key_field(); let key = key.to_data_value(); - extract_optional_model(bind_orm_context(executor, |ctx| { + extract_optional_model(bind_orm_context(executor, true, |ctx| { let plan: LogicalPlan = ctx .from::()? .filter(|expr| { @@ -3443,7 +3449,7 @@ fn orm_get<'a, E: BindSource<'a>, M: Model>( fn orm_list<'a, E: BindSource<'a>, M: Model>( executor: E, ) -> Result, DatabaseError> { - Ok(bind_orm_context(executor, |ctx| { + Ok(bind_orm_context(executor, true, |ctx| { let plan: LogicalPlan = ctx.from::()?.finish()?; Ok(plan) })? @@ -3455,7 +3461,7 @@ fn orm_list<'a, E: BindSource<'a>, M: Model>( mod tests { use super::*; use crate::expression::BinaryOperator; - use crate::types::tuple::{Schema, SchemaView}; + use crate::types::tuple::{Schema, SchemaView, Tuple}; use std::sync::Arc; #[derive(Debug, PartialEq, Eq)] @@ -3515,7 +3521,7 @@ mod tests { impl FromQueryRow for OrmUnitUser { fn from_query_row( schema: &SchemaView<'_, '_>, - tuple: &mut Tuple, + tuple: &mut TupleImpl, ) -> Result { Ok(Self { id: take_value_at(tuple, schema.position("id"), "id")?, @@ -3603,7 +3609,7 @@ mod tests { impl FromQueryRow for OrmUnitOrder { fn from_query_row( schema: &SchemaView<'_, '_>, - tuple: &mut Tuple, + tuple: &mut TupleImpl, ) -> Result { Ok(Self { id: take_value_at(tuple, schema.position("id"), "id")?, @@ -3718,7 +3724,7 @@ mod tests { let arena = PlanArena::new(&table_arena); let schema: Schema = Vec::new(); let schema_view = SchemaView::new(&schema, &arena); - let mut tuple = Tuple::new( + let mut tuple: TupleImpl = Tuple::new( None, vec![ DataValue::from("id".to_string()), @@ -3728,7 +3734,8 @@ mod tests { DataValue::from("PRI".to_string()), DataValue::Null, ], - ); + ) + .into(); let column = DescribeColumn::from_query_row(&schema_view, &mut tuple)?; assert_eq!( @@ -3742,10 +3749,7 @@ mod tests { default: "NULL".to_string(), } ); - assert!(tuple - .values - .iter() - .all(|value| matches!(value, DataValue::Null))); + assert!(tuple.iter().all(|value| matches!(value, DataValue::Null))); assert_eq!(describe_text_value(None), ""); Ok(()) @@ -3810,29 +3814,29 @@ mod tests { #[test] fn tuple_projection_helpers_cast_extract_and_report_width_mismatch() -> Result<(), DatabaseError> { - let mut tuple = Tuple::new( + let mut tuple: TupleImpl = Tuple::new( None, vec![DataValue::Int8(7), DataValue::from("seven".to_string())], - ); + ) + .into(); assert_eq!(take_value_at::(&mut tuple, Some(0), "id")?, 7); - assert!(matches!(tuple.values[0], DataValue::Null)); + assert!(matches!(tuple[0], DataValue::Null)); assert!(matches!( take_value_at::(&mut tuple, None, "missing"), Err(DatabaseError::ColumnNotFound { .. }) )); - let mut tuple = Tuple::new( + let mut tuple: TupleImpl = Tuple::new( None, vec![DataValue::Int32(1), DataValue::from("one".to_string())], - ); + ) + .into(); let projected = <(i32, String) as FromQueryTuple>::from_query_tuple(&mut tuple)?; assert_eq!(projected, (1, "one".to_string())); - assert!(tuple - .values - .iter() - .all(|value| matches!(value, DataValue::Null))); + assert!(tuple.iter().all(|value| matches!(value, DataValue::Null))); - let mut too_wide = Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(2)]); + let mut too_wide: TupleImpl = + Tuple::new(None, vec![DataValue::Int32(1), DataValue::Int32(2)]).into(); assert!(matches!( extract_value_from_tuple::(&mut too_wide), Err(DatabaseError::MisMatch( @@ -3841,7 +3845,7 @@ mod tests { )) )); - let mut too_narrow = Tuple::new(None, vec![DataValue::Int32(1)]); + let mut too_narrow: TupleImpl = Tuple::new(None, vec![DataValue::Int32(1)]).into(); assert!(matches!( <(i32, i32) as FromQueryTuple>::from_query_tuple(&mut too_narrow), Err(DatabaseError::MisMatch( @@ -3886,7 +3890,7 @@ mod tests { assert_eq!(rows.len(), 1); assert_eq!(rows[0].name, expected); } - let mut tx = database.new_transaction()?; + let mut tx = database.new_transaction(false)?; let tx_plan = tx.prepare(&types, |ctx| { ctx.from::()? .filter(|e| e.column(OrmUnitUser::id())?.eq(e.param(1)?))? @@ -3903,7 +3907,9 @@ mod tests { #[test] fn database_and_transaction_orm_helpers_bind_expected_plans() -> Result<(), DatabaseError> { - let mut database = crate::db::DataBaseBuilder::path("./orm-unit-test").build_in_memory()?; + let mut database = crate::db::DataBaseBuilder::path("./orm-unit-test") + .histogram_buckets(100) + .build_in_memory()?; database .ddl("create table orm_unit_users (id int primary key, name varchar, age int null)")?; database.create_view("orm_unit_user_names", |ctx| { @@ -3960,7 +3966,7 @@ mod tests { .iter() .any(|column| column.field == "name")); - let mut tx = database.new_transaction()?; + let mut tx = database.new_transaction(false)?; tx.insert(&OrmUnitUser { id: 4, name: "Dora".to_string(), @@ -4077,8 +4083,9 @@ mod tests { concat!( "Projection [orm_unit_users.id] [Project => (Sort Option: Follow)] ", "Sort By 1 Asc Nulls Last [Sort => (Sort Option: OrderBy: (1 Asc Nulls Last) ignore_prefix_len: 0)] ", - "Filter (((orm_unit_users.id = 3) || ((orm_unit_users.id = 2) || (orm_unit_users.id = 1))) && (orm_unit_users.id != 3)), Is Having: false ", - "[Filter => (Sort Option: Follow)] TableScan orm_unit_users -> [orm_unit_users.id] [SeqScan => (Sort Option: None)]" + "Filter (orm_unit_users.id != 3), Is Having: false ", + "[Filter => (Sort Option: Follow)] TableScan orm_unit_users -> [orm_unit_users.id] ", + "[IndexScan By pk_index => 1, 2, 3 Covered => (Sort Option: OrderBy: (orm_unit_users.id Asc Nulls Last) ignore_prefix_len: 0)]" ), "{list_plan}" ); @@ -4204,10 +4211,7 @@ mod tests { let mut iter = database.run("select id,name,age from orm_unit_users where id>3 order by id")?; let mut rows = Vec::new(); - while iter - .next_tuple(|_, row| rows.push(row.values.clone()))? - .is_some() - {} + while iter.next_tuple(|_, row| rows.push(row.clone()))?.is_some() {} iter.done()?; assert_eq!( rows, @@ -4438,9 +4442,9 @@ mod tests { inner_using_plan, concat!( "Projection [orm_unit_users.id] [Project => (Sort Option: Follow)] ", - "InnerJoinApply ", - "TableScan orm_unit_users -> [orm_unit_users.id] [SeqScan => (Sort Option: None)] ", - "TableScan orm_unit_orders -> [orm_unit_orders.id] [IndexScan By pk_index => Probe ? => (Sort Option: OrderBy: (orm_unit_orders.id Asc Nulls Last) ignore_prefix_len: 0)]" + "IndexProbeJoin [IndexProbeJoin => (Sort Option: Follow)] ", + "TableScan orm_unit_orders -> [orm_unit_orders.id] [SeqScan => (Sort Option: None)] ", + "TableScan orm_unit_users -> [orm_unit_users.id] [IndexScan By pk_index => Probe ? Covered => (Sort Option: OrderBy: (orm_unit_users.id Asc Nulls Last) ignore_prefix_len: 0)]" ), "{inner_using_plan}" ); diff --git a/src/planner/arena.rs b/src/planner/arena.rs index 8df83e70..e02a6621 100644 --- a/src/planner/arena.rs +++ b/src/planner/arena.rs @@ -18,7 +18,7 @@ use crate::expression::visitor_mut::ExprVisitorMut; use crate::expression::ScalarExpression; use crate::planner::operator::scalar_query_init::ScalarQueryInitOperator; use crate::planner::operator::visitor_mut::{ - walk_mut_operator, OperatorExprVisitorMut, OperatorVisitorMut, + walk_mut_operator, ChildrenVisitMut, OperatorExprVisitorMut, OperatorVisitorMut, }; use crate::planner::operator::{Operator, PhysicalOption}; use crate::planner::{LogicalPlan, ScalarQueryRef}; @@ -27,30 +27,28 @@ use crate::types::tuple::Schema; use crate::types::value::DataValue; use crate::types::LogicalType; use std::cell::UnsafeCell; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::fmt; struct ScalarQueryRelocator { source_arena: usize, target_arena: usize, - base: usize, } impl ScalarQueryRelocator { fn relocate(&self, reference: &mut ScalarQueryRef) { if reference.arena_id == self.source_arena { reference.arena_id = self.target_arena; - reference.pos += self.base; } } } -impl ExprVisitorMut for ScalarQueryRelocator { +impl ExprVisitorMut for ScalarQueryRelocator { fn visit_outer_param( &mut self, reference: &mut ScalarQueryRef, _ty: &mut LogicalType, - _arena: &mut dyn MetaArena, + _arena: &mut A, ) -> Result<(), DatabaseError> { self.relocate(reference); Ok(()) @@ -60,7 +58,7 @@ impl ExprVisitorMut for ScalarQueryRelocator { &mut self, reference: &mut ScalarQueryRef, _ty: &mut LogicalType, - _arena: &mut dyn MetaArena, + _arena: &mut A, ) -> Result<(), DatabaseError> { self.relocate(reference); Ok(()) @@ -70,32 +68,37 @@ impl ExprVisitorMut for ScalarQueryRelocator { &mut self, reference: &mut ScalarQueryRef, _ty: &mut LogicalType, - _arena: &mut dyn MetaArena, + _arena: &mut A, ) -> Result<(), DatabaseError> { self.relocate(reference); Ok(()) } } -struct ScalarQueryPlanRelocator<'a> { +struct ScalarQueryPlanRelocator<'a, A: ?Sized> { references: ScalarQueryRelocator, - arena: &'a mut dyn MetaArena, + arena: &'a mut A, } -impl<'a> OperatorVisitorMut<'a> for ScalarQueryPlanRelocator<'_> { +impl<'a, A: MetaArena + ?Sized> OperatorVisitorMut<'a> for ScalarQueryPlanRelocator<'_, A> { fn visit_operator( &mut self, operator: &'a mut Operator, physical_option: Option<&'a mut PhysicalOption>, + children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { - OperatorExprVisitorMut::new(&mut self.references, self.arena) - .visit_operator(operator, physical_option)?; - walk_mut_operator(self, operator) + OperatorExprVisitorMut::new(&mut self.references, self.arena).visit_operator( + operator, + physical_option, + &mut ChildrenVisitMut::empty(), + )?; + walk_mut_operator(self, operator, children) } fn visit_scalar_query_init( &mut self, op: &'a mut ScalarQueryInitOperator, + _children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { for (reference, _) in &mut op.param_bindings { self.references.relocate(reference); @@ -152,7 +155,7 @@ pub struct PlanArena<'a> { plans: Vec, } -#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)] +#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq, Ord, PartialOrd)] pub struct ExprRef { pos: usize, } @@ -243,6 +246,8 @@ pub trait MetaArena { fn reserve_scalar_query_ref(&mut self, reference: ScalarQueryRef); + fn scalar_query_count(&self) -> usize; + fn init_value(&self, _reference: ScalarQueryRef) -> Option<&DataValue> { None } @@ -252,7 +257,7 @@ pub trait MetaArena { } fn same_column(&self, left: ColumnRef, right: ColumnRef) -> bool { - self.column(left).summary() == self.column(right).summary() + left == right || self.column(left).summary() == self.column(right).summary() } fn clone_column(&self, column: ColumnRef) -> ColumnCatalog { @@ -309,6 +314,9 @@ impl MetaArena for Box { fn reserve_scalar_query_ref(&mut self, reference: ScalarQueryRef) { (**self).reserve_scalar_query_ref(reference) } + fn scalar_query_count(&self) -> usize { + (**self).scalar_query_count() + } fn init_value(&self, reference: ScalarQueryRef) -> Option<&DataValue> { (**self).init_value(reference) } @@ -344,14 +352,22 @@ impl MetaArena for Box { /// Owns the metadata and scalar-query cache for one statement execution. pub struct ExecMetaArena { parent: A, - init_values: HashMap>, + init_values: Vec>, } impl ExecMetaArena { pub(crate) fn new(parent: A) -> Self { + let mut init_values = Vec::new(); + init_values.resize_with(parent.scalar_query_count(), || None); Self { parent, - init_values: HashMap::new(), + init_values, + } + } + + fn reserve_init_value(&mut self, reference: ScalarQueryRef) { + if self.init_values.len() <= reference.pos { + self.init_values.resize_with(reference.pos + 1, || None); } } @@ -370,13 +386,23 @@ impl ExecMetaArena { } pub(crate) fn set_init_value(&mut self, reference: ScalarQueryRef, value: DataValue) { - self.init_values.insert(reference, Some(value)); + let slot = &mut self.init_values[reference.pos]; + debug_assert!( + slot.as_ref().is_none_or(|(stored, _)| *stored == reference), + "scalar query {reference} shares a slot with {:?}", + slot.as_ref().map(|(stored, _)| *stored) + ); + *slot = Some((reference, value)); } pub(crate) fn take_outer_values(&mut self) -> Vec<(ScalarQueryRef, DataValue)> { self.init_values - .extract_if(|reference, _| reference.is_outer_value) - .filter_map(|(reference, value)| value.map(|value| (reference, value))) + .iter_mut() + .filter(|slot| { + slot.as_ref() + .is_some_and(|(reference, _)| reference.is_outer_value) + }) + .filter_map(Option::take) .collect() } @@ -392,13 +418,21 @@ impl MetaArena for ExecMetaArena { self.parent.arena_id() } fn alloc_scalar_query_ref(&mut self, is_outer_value: bool) -> ScalarQueryRef { - self.parent.alloc_scalar_query_ref(is_outer_value) + let reference = self.parent.alloc_scalar_query_ref(is_outer_value); + self.reserve_init_value(reference); + reference } fn reserve_scalar_query_ref(&mut self, reference: ScalarQueryRef) { - self.parent.reserve_scalar_query_ref(reference) + self.parent.reserve_scalar_query_ref(reference); + self.reserve_init_value(reference); + } + fn scalar_query_count(&self) -> usize { + self.parent.scalar_query_count() } fn init_value(&self, reference: ScalarQueryRef) -> Option<&DataValue> { - self.init_values.get(&reference).and_then(Option::as_ref) + let (stored, value) = self.init_values.get(reference.pos)?.as_ref()?; + debug_assert_eq!(*stored, reference, "scalar query slot collision"); + Some(value) } fn table_arena_cell<'a>(&self) -> &'a TableArenaCell where @@ -619,6 +653,9 @@ impl MetaArena for TableArena { assert_eq!(reference.arena_id, 0); self.scalar_query_count = self.scalar_query_count.max(reference.pos + 1); } + fn scalar_query_count(&self) -> usize { + self.scalar_query_count + } fn alloc_column(&mut self, column: ColumnCatalog) -> ColumnRef { if let Some(column_ref) = self.find_column(&column) { return column_ref; @@ -750,7 +787,7 @@ impl<'a> PlanArena<'a> { allocated_columns_len: 0, temp_table_id: 0, arena_id: table_arena.borrow().arena_id() + 1, - scalar_query_count: 0, + scalar_query_count: table_arena.borrow().scalar_query_count, columns: Vec::new(), indexes: Vec::new(), expressions: Vec::new(), @@ -824,9 +861,8 @@ impl<'a> PlanArena<'a> { let references = ScalarQueryRelocator { source_arena: self.arena_id, target_arena: target.arena_id(), - base: target.scalar_query_count, }; - target.scalar_query_count += self.scalar_query_count; + target.scalar_query_count = self.scalar_query_count; ScalarQueryPlanRelocator { references, arena: self, @@ -880,7 +916,7 @@ impl<'a> PlanArena<'a> { } pub(crate) fn same_column(&self, left: ColumnRef, right: ColumnRef) -> bool { - self.column(left).summary() == self.column(right).summary() + left == right || self.column(left).summary() == self.column(right).summary() } pub(crate) fn nullable_for_join( @@ -888,13 +924,13 @@ impl<'a> PlanArena<'a> { column: ColumnRef, nullable: bool, ) -> Option { + // A join can only make a column nullable; it never removes existing NULLs. let source = self.column(column); - if source.nullable() == nullable { + if !nullable || source.nullable() { return None; } - // FIXME let mut joined = source.clone(); - joined.set_nullable(nullable); + joined.set_nullable(true); joined.set_in_join(true); Some(self.alloc_column(joined)) } @@ -980,6 +1016,9 @@ impl MetaArena for PlanArena<'_> { assert!(reference.pos < self.table_arena.borrow().scalar_query_count); } } + fn scalar_query_count(&self) -> usize { + self.scalar_query_count + } fn table_arena_cell<'a>(&self) -> &'a TableArenaCell where Self: 'a, @@ -1104,11 +1143,14 @@ impl<'a> ParamArena<'a> { }; let id = *id; let mut value = value.clone(); - value.bind_parameters(&|id| { - params - .iter() - .find_map(|(candidate, value)| (*candidate == id).then_some(value)) - })?; + value.bind_parameters( + &|id| { + params + .iter() + .find_map(|(candidate, value)| (*candidate == id).then_some(value)) + }, + &|_| None, + )?; expressions.push((Some(id), ArenaExpr::new(ScalarExpression::Constant(value)))); } Ok(Self { @@ -1116,7 +1158,7 @@ impl<'a> ParamArena<'a> { expressions, parameter_count: parameter_expressions.len(), parent_end: parent.expression_end(), - scalar_query_count: 0, + scalar_query_count: parent.scalar_query_count, }) } } @@ -1139,6 +1181,9 @@ impl MetaArena for ParamArena<'_> { self.scalar_query_count = self.scalar_query_count.max(reference.pos + 1); } } + fn scalar_query_count(&self) -> usize { + self.scalar_query_count + } fn table_arena_cell<'a>(&self) -> &'a TableArenaCell where Self: 'a, @@ -1149,12 +1194,17 @@ impl MetaArena for ParamArena<'_> { true } fn bound_param(&self, id: usize) -> Option<&DataValue> { - self.expressions - .iter() - .find_map(|(slot, expr)| match (slot, &expr.expression) { - (Some(slot), ScalarExpression::Constant(value)) if *slot == id => Some(value), - _ => None, - }) + for (slot, expr) in &self.expressions { + let Some(slot) = slot else { + return None; + }; + if *slot == id { + if let ScalarExpression::Constant(value) = &expr.expression { + return Some(value); + } + } + } + None } fn column(&self, column: ColumnRef) -> &ColumnCatalog { self.parent.column(column) @@ -1211,6 +1261,7 @@ impl MetaArena for ParamArena<'_> { mod tests { use super::*; use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::planner::Childrens; use crate::types::index::{IndexMeta, IndexType}; use crate::types::LogicalType; use std::panic::{catch_unwind, AssertUnwindSafe}; @@ -1238,7 +1289,6 @@ mod tests { #[test] fn materializing_scalar_references_preserves_shared_catalog_ids() -> Result<(), DatabaseError> { use crate::planner::operator::project::ProjectOperator; - use crate::planner::Childrens; let catalog = TableArenaCell::default(); let persistent = catalog.borrow_mut().alloc_scalar_query_ref(false); let mut arena = PlanArena::new(&catalog); @@ -1267,9 +1317,7 @@ mod tests { .map(|expr| arena.alloc_expression(expr)) .collect::>(); let input = LogicalPlan::new( - Operator::Project(ProjectOperator { - exprs: exprs.clone(), - }), + Operator::Project(ProjectOperator::new(exprs.clone())), Childrens::Only(Box::new(LogicalPlan::new(Operator::Dummy, Childrens::None))), ); let source = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(5))); @@ -1337,7 +1385,7 @@ mod tests { ScalarQueryRef { is_outer_value: false, arena_id: 1, - pos: 0 + pos: 1 } ); assert_ne!(catalog, local); @@ -1351,7 +1399,7 @@ mod tests { ScalarQueryRef { is_outer_value: false, arena_id: 2, - pos: 0 + pos: 2 } ); let mut view = ExecMetaArena::new(params); diff --git a/src/planner/mod.rs b/src/planner/mod.rs index 612bcba5..637569c5 100644 --- a/src/planner/mod.rs +++ b/src/planner/mod.rs @@ -20,12 +20,13 @@ mod scalar_query_ref; use crate::catalog::TableName; use crate::errors::DatabaseError; use crate::expression::visitor_mut::ExprCloner; +use crate::planner::operator::join::JoinType; use crate::planner::operator::recursive_cte::{RecursiveCteOperator, RecursiveScanOperator}; use crate::planner::operator::set_membership::SetMembershipOperator; use crate::planner::operator::union::UnionOperator; use crate::planner::operator::values::ValuesOperator; use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; -use crate::planner::operator::{Operator, PhysicalOption}; +use crate::planner::operator::{ChildCosts, Operator, PhysicalOption, PlanCost, RowCountOption}; use kite_sql_serde_macros::ReferenceSerialization; use std::fmt; use std::hash::{Hash, Hasher}; @@ -36,9 +37,9 @@ pub(crate) use plan_keeper::{PlanInput, PlanKeeper}; pub use scalar_query_ref::ScalarQueryRef; pub(crate) trait Explain { - fn fmt(&self, arena: &(dyn MetaArena + '_), f: &mut fmt::Formatter<'_>) -> fmt::Result; + fn fmt(&self, arena: &A, f: &mut fmt::Formatter<'_>) -> fmt::Result; - fn explain<'a, 'p>(&'a self, arena: &'a (dyn MetaArena + 'p)) -> ExplainDisplay<'a, 'p, Self> + fn explain<'a, A: MetaArena + ?Sized>(&'a self, arena: &'a A) -> ExplainDisplay<'a, Self, A> where Self: Sized, { @@ -46,21 +47,21 @@ pub(crate) trait Explain { } } -pub(crate) struct ExplainDisplay<'a, 'p, T: ?Sized> { +pub(crate) struct ExplainDisplay<'a, T: ?Sized, A: ?Sized> { value: &'a T, - arena: &'a (dyn MetaArena + 'p), + arena: &'a A, } -impl fmt::Display for ExplainDisplay<'_, '_, T> { +impl fmt::Display for ExplainDisplay<'_, T, A> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.value.fmt(self.arena, f) } } -pub(crate) fn fmt_explain_list( +pub(crate) fn fmt_explain_list( values: &[T], separator: &str, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut fmt::Formatter<'_>, ) -> fmt::Result { for (index, value) in values.iter().enumerate() { @@ -172,9 +173,9 @@ impl LogicalPlan { } } - pub(crate) fn clone_plan( + pub(crate) fn clone_plan( &self, - arena: &mut (dyn MetaArena + '_), + arena: &mut A, ) -> Result { let mut plan = self.clone(); OperatorExprVisitorMut::new(&mut ExprCloner, arena).visit_plan(&mut plan)?; @@ -185,6 +186,65 @@ impl LogicalPlan { std::mem::replace(self, Self::new(Operator::Dummy, Childrens::None)) } + pub(crate) fn row_count_option(&self) -> Option<&PhysicalOption> { + let option = self.physical_option.as_ref()?; + match option.row_count { + RowCountOption::Estimated { .. } => Some(option), + RowCountOption::Unknown => None, + RowCountOption::Follow => match self.childrens.as_ref() { + Childrens::Only(child) => child.row_count_option(), + Childrens::Twins { left, .. } => left.row_count_option(), + Childrens::None => None, + }, + } + } + + pub(crate) fn estimated_rows(&self) -> Option { + match &self.physical_option.as_ref()?.row_count { + RowCountOption::Estimated { rows, .. } => Some(*rows), + RowCountOption::Unknown => match (&self.operator, self.childrens.as_ref()) { + (Operator::Join(op), Childrens::Twins { left, right }) + if op.join_type == JoinType::Cross && !op.limit_pushed => + { + Some( + left.estimated_rows()? + .saturating_mul(right.estimated_rows()?), + ) + } + _ => None, + }, + RowCountOption::Follow => match self.childrens.as_ref() { + Childrens::Only(child) => child.estimated_rows(), + Childrens::Twins { left, .. } => left.estimated_rows(), + Childrens::None => None, + }, + } + } + + // Pass-local memoization for NLJ orientation. New physical options start + // empty; orientation preserves child rows/startup, so their cache is reusable. + pub(crate) fn price(&mut self, arena: &mut A) -> Option { + self.physical_option.as_ref()?; + let width = self.output_schema(arena).len(); + let rows = self.estimated_rows(); + let children = match self.childrens.as_mut() { + Childrens::None => ChildCosts::None, + Childrens::Only(child) => ChildCosts::Only(child.price(arena)), + Childrens::Twins { left, right } => { + ChildCosts::Twins(left.price(arena), right.price(arena)) + } + }; + let option = self.physical_option.as_ref()?; + let mut cost = PlanCost { + rows, + width, + startup: option.calculate_startup_work(&self.operator, &children), + total: None, + }; + cost.total = option.calculate_work(&self.operator, &children, rows, &cost); + Some(cost) + } + pub fn referenced_table(&self) -> Vec { fn collect_table(plan: &LogicalPlan, results: &mut Vec) { if let Operator::TableScan(op) = &plan.operator { @@ -200,6 +260,14 @@ impl LogicalPlan { tables } + pub(crate) fn visit_referenced_columns( + &self, + arena: &A, + f: &mut impl FnMut(&A, &crate::catalog::ColumnRef) -> bool, + ) -> Result { + Operator::referenced_columns(operator::ColumnVisitRoot::Plan(self), arena, f) + } + pub(crate) fn visit_column_refs( &self, arena: &mut A, @@ -209,14 +277,10 @@ impl LogicalPlan { A: MetaArena + ?Sized, F: FnMut(&crate::catalog::ColumnRef) + ?Sized, { - self.operator - .visit_referenced_columns(arena, &mut |_, column| { - f(column); - true - })?; - for child in self.childrens.iter() { - child.visit_column_refs(arena, f)?; - } + self.visit_referenced_columns(&*arena, &mut |_, column| { + f(column); + true + })?; Ok(()) } @@ -261,17 +325,18 @@ impl LogicalPlan { }, Operator::Join(_) => match childrens { Childrens::Twins { left, right } => { - let mut schema = left.output_schema(arena).clone(); - schema.extend_from_slice(right.output_schema(arena)); + let left_schema = left.output_schema(arena); + let right_schema = right.output_schema(arena); + let mut schema = Vec::with_capacity(left_schema.len() + right_schema.len()); + schema.extend_from_slice(left_schema); + schema.extend_from_slice(right_schema); schema } _ => unreachable!(), }, - Operator::MarkApply(op) - if matches!(op.kind, operator::mark_apply::MarkApplyKind::InnerJoin) => - { + Operator::IndexProbeJoin(_) => { let Childrens::Twins { left, right } = childrens else { - unreachable!("inner apply requires two inputs") + unreachable!("index probe join requires two inputs") }; let mut schema = left.output_schema(arena).clone(); schema.extend_from_slice(right.output_schema(arena)); @@ -407,7 +472,7 @@ impl LogicalPlan { } impl Explain for LogicalPlan { - fn fmt(&self, arena: &(dyn MetaArena + '_), f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn fmt(&self, arena: &A, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.operator.fmt(arena, f)?; if let Some(physical_option) = &self.physical_option { @@ -483,7 +548,7 @@ impl crate::serdes::ReferenceSerialization for LogicalPlan { ) } - fn decode( + fn decode( reader: &mut R, context: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &crate::serdes::ReferenceTables, @@ -524,10 +589,20 @@ impl crate::serdes::ReferenceSerialization for LogicalPlan { mod tests { use super::*; use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; + use crate::expression::ScalarExpression; + use crate::optimizer::core::rule::NormalizationRule; + use crate::optimizer::heuristic::batch::HepBatchStrategy; + use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline; + use crate::optimizer::rule::normalization::NormalizationRuleImpl; use crate::planner::operator::describe::DescribeOperator; + use crate::planner::operator::filter::FilterOperator; + use crate::planner::operator::join::{JoinCondition, JoinOperator, JoinType}; use crate::planner::operator::limit::LimitOperator; + use crate::planner::operator::project::ProjectOperator; + use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::operator::{PlanImpl, SortOption}; + use crate::types::value::DataValue; use crate::types::LogicalType; use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; @@ -560,6 +635,170 @@ mod tests { .collect() } + #[test] + fn row_count_follows_input_and_stops_at_unknown() { + let mut input = table_scan("users", Vec::new()); + input.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 12, + ndv: None, + }, + )); + let mut plan = LogicalPlan::new(Operator::Dummy, Childrens::Only(Box::new(input))); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + )); + assert_eq!(plan.estimated_rows(), Some(12)); + plan.physical_option.as_mut().unwrap().row_count = RowCountOption::Unknown; + assert_eq!(plan.estimated_rows(), None); + plan.physical_option.as_mut().unwrap().row_count = + RowCountOption::Estimated { rows: 0, ndv: None }; + assert_eq!(plan.estimated_rows(), Some(0)); + plan.physical_option = None; + assert_eq!(plan.estimated_rows(), None); + } + + #[test] + fn filter_and_projection_work_preserve_input_without_blocking() -> Result<(), DatabaseError> { + for output_rows in [0, 8, 98] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let predicate = + arena.alloc_expression(ScalarExpression::Constant(DataValue::Boolean(true))); + let mut scan = table_scan("input", vec![]); + scan.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 98, + ndv: None, + }, + )); + let mut filter = FilterOperator::build(predicate, scan, false); + filter.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Estimated { + rows: output_rows, + ndv: None, + }, + )); + assert_eq!(filter.estimated_rows(), Some(output_rows)); + assert_eq!( + filter.price(&mut arena).and_then(|cost| cost.startup), + Some(0) + ); + assert_eq!( + filter + .price(&mut arena) + .and_then(|cost| cost.input()) + .map(|cost| cost.work), + Some(392) + ); + let mut project = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![])), + Childrens::Only(Box::new(filter)), + ); + project.physical_option = Some(PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + )); + assert_eq!( + project.price(&mut arena).and_then(|cost| cost.startup), + Some(0) + ); + assert_eq!( + project + .price(&mut arena) + .and_then(|cost| cost.input()) + .map(|cost| cost.work), + Some(392) + ); + } + Ok(()) + } + + #[test] + fn blocking_operator_consumes_complete_filtered_input_work() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let id = column("id", LogicalType::Integer, &mut arena); + let key = arena.alloc_expression(ScalarExpression::column_expr(id, 0)); + let predicate = + arena.alloc_expression(ScalarExpression::Constant(DataValue::Boolean(true))); + let mut scan = table_scan("input", vec![id]); + scan.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 98, + ndv: None, + }, + )); + let mut filter = FilterOperator::build(predicate, scan, false); + filter.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Estimated { rows: 8, ndv: None }, + )); + let mut sort = LogicalPlan::new( + Operator::Sort(SortOperator { + sort_fields: vec![SortField { + expr: key, + asc: true, + nulls_first: false, + }], + }), + Childrens::Only(Box::new(filter)), + ); + sort.physical_option = Some(PhysicalOption::new( + PlanImpl::Sort, + SortOption::None, + RowCountOption::Follow, + )); + // 98 one-column input rows (5 each) + 24 comparisons * 2 operator units, + // then 8 output rows. + assert_eq!( + sort.price(&mut arena).and_then(|cost| cost.startup), + Some(538) + ); + assert_eq!( + sort.price(&mut arena) + .and_then(|cost| cost.input()) + .map(|cost| cost.work), + Some(578) + ); + let mut unknown = table_scan("unknown", vec![]); + unknown.physical_option = Some(PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Unknown, + )); + let mut filter = FilterOperator::build(predicate, unknown, false); + filter.physical_option = Some(PhysicalOption::new( + PlanImpl::Filter, + SortOption::Follow, + RowCountOption::Estimated { rows: 8, ndv: None }, + )); + assert_eq!( + filter.price(&mut arena).and_then(|cost| cost.startup), + Some(0) + ); + assert_eq!( + filter + .price(&mut arena) + .and_then(|cost| cost.input()) + .map(|cost| cost.work), + None + ); + Ok(()) + } + #[test] fn childrens_iterates_and_pops_expected_plans() { let left = LogicalPlan::new(Operator::ShowTable, Childrens::None); @@ -598,7 +837,11 @@ mod tests { let id = column("id", LogicalType::Integer, &mut arena); let scan = table_scan("users", vec![id]); let mut plan = LimitOperator::build(Some(2), Some(5), scan); - plan.physical_option = Some(PhysicalOption::new(PlanImpl::Limit, SortOption::Follow)); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::Limit, + SortOption::Follow, + RowCountOption::Unknown, + )); let tables = plan .referenced_table() @@ -652,6 +895,169 @@ mod tests { assert_eq!(schema_names(&recomputed, &arena), vec!["name"]); } + #[test] + fn transpose_invalidates_changed_child_before_schema_is_read() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let a = column("a", LogicalType::Integer, &mut arena); + let b = column("b", LogicalType::Integer, &mut arena); + let expr = arena.alloc_expression(ScalarExpression::column_expr(a, 0)); + let project = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![expr])), + Childrens::Only(Box::new(table_scan("t", vec![a, b]))), + ); + let mut plan = LimitOperator::build(None, Some(1), project); + plan.populate_output_schema_recursive(&mut arena); + assert_eq!(plan.childrens.only().output_schema.as_ref().unwrap(), &[a]); + + assert!(NormalizationRuleImpl::LimitProjectTranspose.apply(&mut plan, &mut arena, None)?); + let Childrens::Only(limit) = plan.childrens.as_mut() else { + panic!("expected a limit below the projection"); + }; + assert!(limit.output_schema.is_none()); + assert_eq!( + limit.childrens.only().output_schema.as_ref().unwrap(), + &[a, b] + ); + assert_eq!(limit.output_schema(&mut arena), &[a, b]); + assert_eq!(plan.output_schema(&mut arena), &[a]); + Ok(()) + } + + #[test] + fn local_rewrite_invalidates_ancestors_but_preserves_untouched_caches( + ) -> Result<(), DatabaseError> { + for strategy in [ + HepBatchStrategy::once_topdown(), + HepBatchStrategy::once_bottomup(), + ] { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let a = column("a", LogicalType::Integer, &mut arena); + let b = column("b", LogicalType::Integer, &mut arena); + let c = column("c", LogicalType::Integer, &mut arena); + let expr = arena.alloc_expression(ScalarExpression::column_expr(a, 0)); + let project = LogicalPlan::new( + Operator::Project(ProjectOperator::new(vec![expr])), + Childrens::Only(Box::new(table_scan("left", vec![a, b]))), + ); + let mut plan = JoinOperator::build( + LimitOperator::build(None, Some(1), project), + table_scan("right", vec![c]), + JoinCondition::None, + JoinType::Cross, + false, + ); + plan.populate_output_schema_recursive(&mut arena); + assert_eq!(plan.output_schema.as_ref().unwrap(), &[a, c]); + let pipeline = HepOptimizerPipeline::builder() + .before_batch( + "Transpose", + strategy, + &[NormalizationRuleImpl::LimitProjectTranspose], + ) + .build(); + let mut plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; + assert!(plan.output_schema.is_none()); + let Childrens::Twins { left, right } = plan.childrens.as_mut() else { + panic!("expected a join"); + }; + assert!(left.output_schema.is_none()); + assert_eq!(right.output_schema.as_ref().unwrap(), &[c]); + let Childrens::Only(limit) = left.childrens.as_mut() else { + panic!("expected a limit below the projection"); + }; + assert!(limit.output_schema.is_none()); + assert_eq!( + limit.childrens.only().output_schema.as_ref().unwrap(), + &[a, b] + ); + assert_eq!(limit.output_schema(&mut arena), &[a, b]); + assert_eq!(plan.output_schema(&mut arena), &[a, c]); + } + Ok(()) + } + + #[test] + fn wrap_child_invalidates_parent_without_clearing_input() { + use crate::optimizer::plan_utils::wrap_child_with; + + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let a = column("a", LogicalType::Integer, &mut arena); + let b = column("b", LogicalType::Integer, &mut arena); + let expr = arena.alloc_expression(ScalarExpression::column_expr(a, 0)); + let mut plan = LimitOperator::build(None, Some(1), table_scan("t", vec![a, b])); + plan.populate_output_schema_recursive(&mut arena); + + assert!(wrap_child_with( + &mut plan, + 0, + Operator::Project(ProjectOperator::new(vec![expr])), + )); + assert!(plan.output_schema.is_none()); + assert!(plan + .childrens + .only() + .childrens + .only() + .output_schema + .is_some()); + assert_eq!(plan.output_schema(&mut arena), &[a]); + } + + #[test] + fn predicate_pushdown_preserves_untouched_subtree_caches() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let a = column("a", LogicalType::Integer, &mut arena); + let b = column("b", LogicalType::Integer, &mut arena); + let c = column("c", LogicalType::Integer, &mut arena); + let left = JoinOperator::build( + table_scan("a", vec![a]), + table_scan("b", vec![b]), + JoinCondition::None, + JoinType::Cross, + false, + ); + let join = JoinOperator::build( + left, + table_scan("c", vec![c]), + JoinCondition::None, + JoinType::Cross, + false, + ); + let c_expr = arena.alloc_expression(ScalarExpression::column_expr(c, 2)); + let zero = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(0))); + let predicate = arena.alloc_expression(ScalarExpression::Binary { + op: crate::expression::BinaryOperator::Gt, + left_expr: c_expr, + right_expr: zero, + evaluator: None, + ty: LogicalType::Boolean, + }); + let mut plan = FilterOperator::build(predicate, join, false); + plan.populate_output_schema_recursive(&mut arena); + let pipeline = HepOptimizerPipeline::builder() + .before_batch( + "Push Predicate", + HepBatchStrategy::once_topdown(), + &[NormalizationRuleImpl::PushPredicateThroughJoin], + ) + .build(); + let mut plan = pipeline.instantiate(plan).find_best(None, &mut arena)?; + assert!(matches!(plan.operator, Operator::Join(_))); + assert!(plan.output_schema.is_none()); + let (left, right) = plan.childrens.twins(); + assert_eq!(left.output_schema.as_ref().unwrap(), &[a, b]); + assert!(left.childrens.twins().0.output_schema.is_some()); + assert!(left.childrens.twins().1.output_schema.is_some()); + assert!(matches!(right.operator, Operator::Filter(_))); + assert_eq!(right.childrens.only().output_schema.as_ref().unwrap(), &[c]); + assert_eq!(plan.output_schema(&mut arena), &[a, b, c]); + Ok(()) + } + #[test] fn dummy_operators_produce_expected_output_schema() { let table_arena = TableArenaCell::default(); diff --git a/src/planner/operator/aggregate.rs b/src/planner/operator/aggregate.rs index 6457aeb9..28d97131 100644 --- a/src/planner/operator/aggregate.rs +++ b/src/planner/operator/aggregate.rs @@ -46,9 +46,9 @@ impl AggregateOperator { } impl Explain for AggregateOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { f.write_str("Aggregate [")?; diff --git a/src/planner/operator/alter_table/change_column.rs b/src/planner/operator/alter_table/change_column.rs index 6c8b2a33..f0b71ec1 100644 --- a/src/planner/operator/alter_table/change_column.rs +++ b/src/planner/operator/alter_table/change_column.rs @@ -43,9 +43,9 @@ pub struct ChangeColumnOperator { } impl Explain for ChangeColumnOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!( diff --git a/src/planner/operator/analyze.rs b/src/planner/operator/analyze.rs index 101c08d7..b2b4f858 100644 --- a/src/planner/operator/analyze.rs +++ b/src/planner/operator/analyze.rs @@ -17,18 +17,21 @@ use crate::planner::MetaArena; use crate::planner::{fmt_explain_list, Explain}; use crate::types::index::IndexMetaRef; use kite_sql_serde_macros::ReferenceSerialization; +use ordered_float::OrderedFloat; #[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] pub struct AnalyzeOperator { pub table_name: TableName, pub index_metas: Vec, pub histogram_buckets: Option, + pub hll_precision: Option, + pub statistics_relative_error: Option>, } impl Explain for AnalyzeOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!(f, "Analyze {} -> [", self.table_name)?; diff --git a/src/planner/operator/copy_from_file.rs b/src/planner/operator/copy_from_file.rs index a6c3418d..49b2bff4 100644 --- a/src/planner/operator/copy_from_file.rs +++ b/src/planner/operator/copy_from_file.rs @@ -27,9 +27,9 @@ pub struct CopyFromFileOperator { } impl Explain for CopyFromFileOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!(f, "Copy {} -> {} [", self.source.path.display(), self.table)?; diff --git a/src/planner/operator/create_index.rs b/src/planner/operator/create_index.rs index c367e6eb..9a31ee76 100644 --- a/src/planner/operator/create_index.rs +++ b/src/planner/operator/create_index.rs @@ -29,9 +29,9 @@ pub struct CreateIndexOperator { } impl Explain for CreateIndexOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!(f, "Create Index On {} -> [", self.table_name)?; diff --git a/src/planner/operator/filter.rs b/src/planner/operator/filter.rs index 2d76b8bd..84a0ed34 100644 --- a/src/planner/operator/filter.rs +++ b/src/planner/operator/filter.rs @@ -39,9 +39,9 @@ impl FilterOperator { } impl Explain for FilterOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!( diff --git a/src/planner/operator/index_probe_join.rs b/src/planner/operator/index_probe_join.rs new file mode 100644 index 00000000..51b3abb2 --- /dev/null +++ b/src/planner/operator/index_probe_join.rs @@ -0,0 +1,36 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::planner::ExprRef; +use kite_sql_serde_macros::ReferenceSerialization; +use std::fmt; +use std::fmt::Formatter; + +#[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] +pub struct IndexProbeJoinOperator { + pub probe: ExprRef, + pub predicates: Vec, +} + +impl IndexProbeJoinOperator { + pub(crate) fn new(probe: ExprRef, predicates: Vec) -> Self { + Self { probe, predicates } + } +} + +impl fmt::Display for IndexProbeJoinOperator { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "IndexProbeJoin") + } +} diff --git a/src/planner/operator/join.rs b/src/planner/operator/join.rs index 7ae642cc..77412c6f 100644 --- a/src/planner/operator/join.rs +++ b/src/planner/operator/join.rs @@ -91,13 +91,13 @@ impl JoinOperator { } impl Explain for JoinOperator { - fn fmt(&self, arena: &(dyn MetaArena + '_), f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn fmt(&self, arena: &A, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{} Join{}", self.join_type, self.on.explain(arena)) } } impl Explain for JoinCondition { - fn fmt(&self, arena: &(dyn MetaArena + '_), f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn fmt(&self, arena: &A, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { JoinCondition::On { on, filter } => { if !on.is_empty() { @@ -152,14 +152,9 @@ pub(crate) fn extract_join_keys( right_schema: &Schema, arena: &mut PlanArena, ) -> Result<(), DatabaseError> { - let fn_contains = |schema: &Schema, column: ColumnRef| { - let summary = arena.column(column).summary(); - schema - .iter() - .any(|candidate| arena.column(*candidate).summary() == summary) - }; + // ON keys keep the refs their inputs output, so refs identify the side. let fn_or_contains = - |column: ColumnRef| fn_contains(left_schema, column) || fn_contains(right_schema, column); + |column: ColumnRef| left_schema.contains(&column) || right_schema.contains(&column); let expr = expr.unpack_alias(arena); match arena.expression(expr) { @@ -181,12 +176,9 @@ pub(crate) fn extract_join_keys( ScalarExpression::ColumnRef { column: r, .. }, ) => { // reorder left and right joins keys to pattern: (left, right) - let key = if fn_contains(left_schema, *l) - && fn_contains(right_schema, *r) - { + let key = if left_schema.contains(l) && right_schema.contains(r) { Some((*left_expr, *right_expr)) - } else if fn_contains(left_schema, *r) && fn_contains(right_schema, *l) - { + } else if left_schema.contains(r) && right_schema.contains(l) { Some((*right_expr, *left_expr)) } else { if fn_or_contains(*l) || fn_or_contains(*r) { @@ -217,9 +209,9 @@ pub(crate) fn extract_join_keys( } _other => { // example: baz > 1 - if left_expr.all_referenced_columns(arena, |_, column| { + if left_expr.all_referenced_columns(arena, |_, column, _| { fn_or_contains(*column) - })? && right_expr.all_referenced_columns(arena, |_, column| { + })? && right_expr.all_referenced_columns(arena, |_, column, _| { fn_or_contains(*column) })? { accum_filter.push(expr); @@ -252,9 +244,9 @@ pub(crate) fn extract_join_keys( } _ => { if left_expr - .all_referenced_columns(arena, |_, column| fn_or_contains(*column))? + .all_referenced_columns(arena, |_, column, _| fn_or_contains(*column))? && right_expr - .all_referenced_columns(arena, |_, column| fn_or_contains(*column))? + .all_referenced_columns(arena, |_, column, _| fn_or_contains(*column))? { accum_filter.push(expr); } @@ -262,7 +254,7 @@ pub(crate) fn extract_join_keys( } } _ => { - if expr.all_referenced_columns(arena, |_, column| fn_or_contains(*column))? { + if expr.all_referenced_columns(arena, |_, column, _| fn_or_contains(*column))? { // example: baz > 1 accum_filter.push(expr); } diff --git a/src/planner/operator/mark_apply.rs b/src/planner/operator/mark_apply.rs index 32e056ac..0b39fa69 100644 --- a/src/planner/operator/mark_apply.rs +++ b/src/planner/operator/mark_apply.rs @@ -29,15 +29,16 @@ pub enum MarkApplyQuantifier { pub enum MarkApplyKind { Exists, Quantified(MarkApplyQuantifier), - InnerJoin, } #[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] pub struct MarkApplyOperator { pub kind: MarkApplyKind, pub predicates: Vec, - output_column: Option, + output_column: ColumnRef, pub parameterized_probe: Option, + pub(crate) materialize: bool, + pub(crate) guard: Option, } impl MarkApplyOperator { @@ -45,8 +46,10 @@ impl MarkApplyOperator { Self { kind: MarkApplyKind::Exists, predicates, - output_column: Some(output_column), + output_column, parameterized_probe: None, + materialize: false, + guard: None, } } @@ -77,8 +80,10 @@ impl MarkApplyOperator { Self { kind: MarkApplyKind::Quantified(quantifier), predicates, - output_column: Some(output_column), + output_column, parameterized_probe: None, + materialize: false, + guard: None, } } @@ -126,18 +131,7 @@ impl MarkApplyOperator { } pub fn output_column(&self) -> &ColumnRef { - self.output_column - .as_ref() - .expect("marker apply output column") - } - - pub(crate) fn new_inner_join(predicates: Vec, probe: ExprRef) -> Self { - Self { - kind: MarkApplyKind::InnerJoin, - predicates, - output_column: None, - parameterized_probe: Some(probe), - } + &self.output_column } pub fn parameterized_probe(&self) -> Option<&ExprRef> { @@ -152,7 +146,6 @@ impl MarkApplyOperator { impl fmt::Display for MarkApplyOperator { fn fmt(&self, f: &mut Formatter) -> fmt::Result { match self.kind { - MarkApplyKind::InnerJoin => write!(f, "InnerJoinApply"), MarkApplyKind::Exists => write!(f, "MarkExistsApply"), MarkApplyKind::Quantified(MarkApplyQuantifier::Any) => write!(f, "MarkAnyApply"), MarkApplyKind::Quantified(MarkApplyQuantifier::All) => write!(f, "MarkAllApply"), diff --git a/src/planner/operator/mod.rs b/src/planner/operator/mod.rs index 27d696ef..16198a88 100644 --- a/src/planner/operator/mod.rs +++ b/src/planner/operator/mod.rs @@ -33,6 +33,7 @@ pub mod drop_table; pub mod drop_view; pub mod filter; pub mod function_scan; +pub mod index_probe_join; pub mod insert; pub mod join; pub mod limit; @@ -64,6 +65,7 @@ use self::{ use crate::catalog::ColumnRef; use crate::errors::DatabaseError; use crate::expression::visitor::{walk_expr, ExprVisitor}; +use crate::expression::AliasType; use crate::planner::operator::alter_table::change_column::DefaultChange as ColumnDefaultChange; use crate::planner::operator::alter_table::drop_column::DropColumnOperator; use crate::planner::operator::analyze::AnalyzeOperator; @@ -80,6 +82,7 @@ use crate::planner::operator::drop_index::DropIndexOperator; use crate::planner::operator::drop_table::DropTableOperator; use crate::planner::operator::drop_view::DropViewOperator; use crate::planner::operator::function_scan::FunctionScanOperator; +use crate::planner::operator::index_probe_join::IndexProbeJoinOperator; use crate::planner::operator::insert::InsertOperator; use crate::planner::operator::join::JoinCondition; use crate::planner::operator::set_membership::SetMembershipOperator; @@ -89,10 +92,11 @@ use crate::planner::operator::truncate::TruncateOperator; use crate::planner::operator::union::UnionOperator; use crate::planner::operator::update::UpdateOperator; use crate::planner::operator::values::ValuesOperator; -use crate::planner::operator::visitor::OperatorVisitor; -use crate::planner::{fmt_explain_list, Explain, ExprRef, MetaArena}; +use crate::planner::operator::visitor::{walk_plan, ChildrenVisit, OperatorVisitor}; +use crate::planner::{fmt_explain_list, Explain, ExprRef, LogicalPlan, MetaArena}; use crate::types::index::{IndexInfo, IndexMetaRef}; use kite_sql_serde_macros::ReferenceSerialization; +use std::hash::{Hash, Hasher}; #[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] pub enum Operator { @@ -119,6 +123,7 @@ pub enum Operator { Union(UnionOperator), RecursiveCte(RecursiveCteOperator), RecursiveScan(RecursiveScanOperator), + IndexProbeJoin(IndexProbeJoinOperator), // DML Insert(InsertOperator), Update(UpdateOperator), @@ -156,20 +161,320 @@ pub enum SortOption { None, } -#[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] +#[derive(Debug, PartialEq, Eq, Clone, Hash)] +pub enum RowCountOption { + Follow, + Unknown, + Estimated { rows: usize, ndv: Option }, +} + +// Relative CPU units: a row costs four units, an operator/hash operation one. +// This is the same scale as row = 1 and operator = 0.25; I/O is not modeled. +pub(crate) const CPU_ROW_WORK: usize = 4; +pub(crate) const CPU_OPERATOR_WORK: usize = 1; +// Rows assumed for a table without statistics (RFC 0001 Ā§3.9). +pub(crate) const DEFAULT_TABLE_ROWS: usize = 1000; +// Rebuilding the inner executor for each outer row of an index probe. +pub(crate) const PROBE_RESET_WORK: usize = CPU_ROW_WORK; +// Hashing, inserting and later freeing one build row. Calibrated on LMDB, where a hash join +// costs ~7 scanned rows per build/probe pair (RFC 0001 Ā§4.2). +pub(crate) const HASH_BUILD_WORK: usize = 6 * CPU_ROW_WORK; + +pub(crate) fn rows_work(rows: usize, width: usize) -> usize { + rows.saturating_mul(CPU_ROW_WORK.saturating_add(width.saturating_mul(CPU_OPERATOR_WORK))) +} + +fn hash_build_work(build: &InputCost) -> usize { + rows_work(build.rows, build.width).saturating_add(build.rows.saturating_mul(HASH_BUILD_WORK)) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct InputCost { + pub(crate) rows: usize, + pub(crate) width: usize, + pub(crate) work: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum JoinImpl { + Hash, + NestLoop { full: bool }, + IndexProbe { lookup_work: usize }, +} + +impl JoinImpl { + pub(crate) fn work( + self, + left: &InputCost, + right: &InputCost, + out_rows: usize, + out_width: usize, + ) -> usize { + let join = match self { + // Build rows are stored whole; each probe row costs one hash lookup. + JoinImpl::Hash => right + .work + .saturating_add(hash_build_work(left)) + .saturating_add(right.rows.saturating_mul(CPU_OPERATOR_WORK)), + // The inner input is re-run and compared for every outer row. + JoinImpl::NestLoop { full } => left + .rows + .saturating_mul( + right + .work + .saturating_add(right.rows.saturating_mul(CPU_OPERATOR_WORK)), + ) + .saturating_add(if full { right.work } else { 0 }), + JoinImpl::IndexProbe { lookup_work } => left.rows.saturating_mul(lookup_work), + }; + left.work + .saturating_add(join) + .saturating_add(rows_work(out_rows, out_width)) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct PlanCost { + pub(crate) rows: Option, + pub(crate) width: usize, + pub(crate) startup: Option, + pub(crate) total: Option, +} + +impl PlanCost { + pub(crate) fn input(&self) -> Option { + Some(InputCost { + rows: self.rows?, + width: self.width, + work: self.total?, + }) + } +} + +/// The inputs' costs, shaped like `Childrens`; `None` marks an input that cannot be priced. +pub(crate) enum ChildCosts { + None, + Only(Option), + Twins(Option, Option), +} + +fn input(cost: &Option) -> Option { + cost.as_ref()?.input() +} + +fn startup(cost: &Option) -> Option { + cost.as_ref()?.startup +} + +#[derive(Debug, Clone)] pub struct PhysicalOption { pub plan: PlanImpl, pub(crate) sort_option: SortOption, + pub(crate) row_count: RowCountOption, + pub(crate) lookup_work: Option, + pub(crate) join_decided: bool, } impl PhysicalOption { - pub fn new(plan: PlanImpl, sort_option: SortOption) -> Self { - Self { plan, sort_option } + pub fn new(plan: PlanImpl, sort_option: SortOption, row_count: RowCountOption) -> Self { + Self { + plan, + sort_option, + row_count, + lookup_work: None, + join_decided: false, + } } pub fn sort_option(&self) -> &SortOption { &self.sort_option } + + fn index_probe_work( + &self, + left: &Option, + right: &Option, + rows: usize, + cost: &PlanCost, + ) -> Option { + let right_cost = (*right)?; + let probed = InputCost { + rows: 1, + width: right_cost.width, + work: right_cost.startup?, + }; + let lookup_work = self.lookup_work.unwrap_or_else(|| { + probed + .work + .saturating_add(PROBE_RESET_WORK) + .saturating_add(rows_work(probed.rows, probed.width)) + }); + Some(JoinImpl::IndexProbe { lookup_work }.work(&input(left)?, &probed, rows, cost.width)) + } + + /// Join formulas: RFC 0001 Ā§4.2 (`rfc/0001-greedy-join-reorder-owns-join-decisions.md`); + /// probe `lookup_work`: Ā§4.3. + pub(crate) fn calculate_work( + &self, + operator: &Operator, + children: &ChildCosts, + rows: Option, + cost: &PlanCost, + ) -> Option { + match (&self.plan, operator, children) { + (PlanImpl::Filter, Operator::Filter(_), ChildCosts::Only(child)) + | (PlanImpl::Project, Operator::Project(_), ChildCosts::Only(child)) => { + Some(input(child)?.work) + } + (PlanImpl::MarkApply, Operator::MarkApply(op), ChildCosts::Twins(left, right)) + if op.parameterized_probe().is_some() && !op.materialize => + { + self.index_probe_work(left, right, rows?, cost) + } + ( + PlanImpl::IndexProbeJoin, + Operator::IndexProbeJoin(_), + ChildCosts::Twins(left, right), + ) => self.index_probe_work(left, right, rows?, cost), + (PlanImpl::HashJoin, Operator::Join(_), ChildCosts::Twins(left, right)) => { + Some(JoinImpl::Hash.work(&input(left)?, &input(right)?, rows?, cost.width)) + } + (PlanImpl::NestLoopJoin, Operator::Join(join), ChildCosts::Twins(left, right)) => { + let (outer, inner) = if join.join_type == join::JoinType::RightOuter { + (right, left) + } else { + (left, right) + }; + Some( + JoinImpl::NestLoop { + full: join.join_type == join::JoinType::Full, + } + .work(&input(outer)?, &input(inner)?, rows?, cost.width), + ) + } + _ => Some(cost.startup?.saturating_add(rows_work(rows?, cost.width))), + } + } + + pub(crate) fn calculate_startup_work( + &self, + operator: &Operator, + children: &ChildCosts, + ) -> Option { + let work = |child: &Option| Some(input(child)?.work); + match (&self.plan, operator, children) { + ( + PlanImpl::SeqScan | PlanImpl::IndexScan(_) | PlanImpl::Values | PlanImpl::Dummy, + _, + ChildCosts::None, + ) => Some(0), + ( + PlanImpl::Project + | PlanImpl::Filter + | PlanImpl::Limit + | PlanImpl::StreamAggregate + | PlanImpl::StreamDistinct, + _, + ChildCosts::Only(child), + ) => startup(child), + (PlanImpl::HashJoin, _, ChildCosts::Twins(left, right)) => { + input(left).and_then(|build| { + Some( + build + .work + .saturating_add(hash_build_work(&build)) + .saturating_add(startup(right)?), + ) + }) + } + (PlanImpl::NestLoopJoin, _, ChildCosts::Twins(left, right)) => startup(left) + .and_then(|left| startup(right).map(|right| left.saturating_add(right))), + (PlanImpl::Sort | PlanImpl::TopK | PlanImpl::Window, _, ChildCosts::Only(child)) => { + let keys = match operator { + Operator::Sort(op) => op.sort_fields.len(), + Operator::TopK(op) => op.sort_fields.len(), + Operator::Window(op) => op.sort_fields.len(), + _ => unreachable!(), + }; + work(child).and_then(|read| { + let rows = child.as_ref()?.rows?; + let comparisons = if rows <= 1 { + 0 + } else { + ((rows as f64) * (rows as f64).log2()).ceil() as usize + }; + Some( + read.saturating_add( + comparisons.saturating_mul( + 2usize + .saturating_mul(keys) + .saturating_mul(CPU_OPERATOR_WORK), + ), + ), + ) + }) + } + ( + PlanImpl::SimpleAggregate | PlanImpl::HashAggregate, + Operator::Aggregate(op), + ChildCosts::Only(child), + ) => { + let per_row = op + .agg_calls + .len() + .saturating_mul(CPU_OPERATOR_WORK) + .saturating_add(if matches!(self.plan, PlanImpl::HashAggregate) { + CPU_ROW_WORK.saturating_add( + op.groupby_exprs.len().saturating_mul(CPU_OPERATOR_WORK), + ) + } else { + 0 + }); + work(child).and_then(|read| { + Some(read.saturating_add(child.as_ref()?.rows?.saturating_mul(per_row))) + }) + } + (PlanImpl::MarkApply, Operator::MarkApply(op), ChildCosts::Twins(left, right)) => { + startup(left).and_then(|left| { + let right = if op.materialize { + work(right) + } else { + startup(right) + }?; + Some(left.saturating_add(right)) + }) + } + ( + PlanImpl::IndexProbeJoin, + Operator::IndexProbeJoin(_), + ChildCosts::Twins(left, right), + ) => startup(left) + .and_then(|left| startup(right).map(|right| left.saturating_add(right))), + (PlanImpl::ScalarQueryInit, _, ChildCosts::Twins(left, right)) => { + startup(left).and_then(|left| work(right).map(|right| left.saturating_add(right))) + } + _ => None, + } + } +} + +// Probe lookup work and Greedy's join decision are pass-local, not plan semantics or wire data. +impl PartialEq for PhysicalOption { + fn eq(&self, other: &Self) -> bool { + self.plan == other.plan + && self.sort_option == other.sort_option + && self.row_count == other.row_count + } +} + +impl Eq for PhysicalOption {} + +impl Hash for PhysicalOption { + fn hash(&self, state: &mut H) { + self.plan.hash(state); + self.sort_option.hash(state); + self.row_count.hash(state); + } } #[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] @@ -209,12 +514,14 @@ pub enum PlanImpl { CopyToFile, Analyze, Window, + // Appended to keep the serialized indices of earlier variants stable. + IndexProbeJoin, } impl Explain for ColumnRef { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { let column = arena.column(*self); @@ -227,9 +534,9 @@ impl Explain for ColumnRef { } impl Explain for IndexMetaRef { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { f.write_str(&arena.index(*self).name) @@ -241,7 +548,7 @@ macro_rules! impl_display_explain { $( $(#[$meta])* impl Explain for $ty { - fn fmt(&self, _arena: &(dyn MetaArena + '_), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + fn fmt(&self, _arena: &A, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { std::fmt::Display::fmt(self, f) } } @@ -251,6 +558,7 @@ macro_rules! impl_display_explain { impl_display_explain!( MarkApplyOperator, + IndexProbeJoinOperator, ScalarSubqueryOperator, FunctionScanOperator, LimitOperator, @@ -270,9 +578,9 @@ impl_display_explain!( ); impl Explain for Operator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { match self { @@ -280,6 +588,7 @@ impl Explain for Operator { Operator::Aggregate(op) => Explain::fmt(op, arena, f), Operator::ScalarQueryInit(op) => Explain::fmt(op, arena, f), Operator::MarkApply(op) => Explain::fmt(op, arena, f), + Operator::IndexProbeJoin(op) => Explain::fmt(op, arena, f), Operator::Filter(op) => Explain::fmt(op, arena, f), Operator::Join(op) => Explain::fmt(op, arena, f), Operator::Project(op) => Explain::fmt(op, arena, f), @@ -321,11 +630,89 @@ impl Explain for Operator { } } +/// Where a referenced-column visit starts: one operator, or a plan and all its inputs. +pub(crate) enum ColumnVisitRoot<'a> { + Operator(&'a Operator), + Plan(&'a LogicalPlan), +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum CommandType { + DQL, + DML, + DDL, + Analyze, +} + +impl CommandType { + /// Only queries leave storage untouched. + pub fn is_read_only(self) -> bool { + matches!(self, CommandType::DQL) + } +} + impl Operator { + /// Classifies a plan by its root, matching `command_type` of SQL statements. + pub fn command_type(&self) -> CommandType { + match self { + Operator::Dummy + | Operator::Aggregate(_) + | Operator::ScalarQueryInit(_) + | Operator::MarkApply(_) + | Operator::Filter(_) + | Operator::Join(_) + | Operator::Project(_) + | Operator::ScalarSubquery(_) + | Operator::TableScan(_) + | Operator::FunctionScan(_) + | Operator::Sort(_) + | Operator::Limit(_) + | Operator::TopK(_) + | Operator::Values(_) + | Operator::ShowTable + | Operator::ShowView + | Operator::Explain + | Operator::Describe(_) + | Operator::SetMembership(_) + | Operator::Union(_) + | Operator::RecursiveCte(_) + | Operator::RecursiveScan(_) + | Operator::IndexProbeJoin(_) + | Operator::Window(_) => CommandType::DQL, + #[cfg(feature = "copy")] + Operator::CopyToFile(_) => CommandType::DQL, + Operator::Insert(_) + | Operator::Update(_) + | Operator::Delete(_) + | Operator::Truncate(_) => CommandType::DML, + #[cfg(feature = "copy")] + Operator::CopyFromFile(_) => CommandType::DML, + Operator::Analyze(_) => CommandType::Analyze, + Operator::AddColumn(_) + | Operator::ChangeColumn(_) + | Operator::DropColumn(_) + | Operator::CreateTable(_) + | Operator::CreateIndex(_) + | Operator::CreateView(_) + | Operator::DropTable(_) + | Operator::DropView(_) + | Operator::DropIndex(_) => CommandType::DDL, + } + } + pub fn visit_referenced_columns( &self, arena: &A, f: &mut impl FnMut(&A, &ColumnRef) -> bool, + ) -> Result { + Self::referenced_columns(ColumnVisitRoot::Operator(self), arena, f) + } + + /// Calls `f` for each referenced column until it returns false; returns whether it ran to the end. + pub(crate) fn referenced_columns( + root: ColumnVisitRoot<'_>, + arena: &A, + f: &mut impl FnMut(&A, &ColumnRef) -> bool, ) -> Result { struct ReferencedColumnVisitor<'a, A: ?Sized, F> { arena: &'a A, @@ -345,12 +732,29 @@ impl Operator { Ok(()) } - fn visit_column_ref(&mut self, column: &ColumnRef) -> Result<(), DatabaseError> { + fn visit_column_ref( + &mut self, + column: &ColumnRef, + _position: usize, + ) -> Result<(), DatabaseError> { if self.keep_going { self.keep_going = (self.f)(self.arena, column); } Ok(()) } + + // The alias target column is still owned by the plan. + fn visit_alias( + &mut self, + expr: ExprRef, + alias: &AliasType, + arena: &A, + ) -> Result<(), DatabaseError> { + if let AliasType::Expr(alias_expr) = alias { + self.visit(*alias_expr, arena)?; + } + self.visit(expr, arena) + } } impl<'operator, A, F> OperatorVisitor<'operator> for ReferencedColumnVisitor<'_, A, F> @@ -358,9 +762,17 @@ impl Operator { A: MetaArena + ?Sized, F: FnMut(&A, &ColumnRef) -> bool, { + fn visit_plan(&mut self, plan: &'operator LogicalPlan) -> Result<(), DatabaseError> { + if self.keep_going { + walk_plan(self, plan)?; + } + Ok(()) + } + fn visit_aggregate( &mut self, op: &'operator AggregateOperator, + _children: &mut ChildrenVisit<'operator>, ) -> Result<(), DatabaseError> { for expr in op.agg_calls.iter().chain(&op.groupby_exprs) { ExprVisitor::visit(self, *expr, self.arena)?; @@ -371,6 +783,7 @@ impl Operator { fn visit_scalar_query_init( &mut self, op: &'operator ScalarQueryInitOperator, + _children: &mut ChildrenVisit<'operator>, ) -> Result<(), DatabaseError> { for (_, expr) in &op.param_bindings { ExprVisitor::visit(self, *expr, self.arena)?; @@ -381,6 +794,7 @@ impl Operator { fn visit_mark_apply( &mut self, op: &'operator MarkApplyOperator, + _children: &mut ChildrenVisit<'operator>, ) -> Result<(), DatabaseError> { for expr in &op.predicates { ExprVisitor::visit(self, *expr, self.arena)?; @@ -388,14 +802,25 @@ impl Operator { if let Some(expr) = &op.parameterized_probe { ExprVisitor::visit(self, *expr, self.arena)?; } + if let Some(expr) = op.guard { + ExprVisitor::visit(self, expr, self.arena)?; + } Ok(()) } - fn visit_filter(&mut self, op: &'operator FilterOperator) -> Result<(), DatabaseError> { + fn visit_filter( + &mut self, + op: &'operator FilterOperator, + _children: &mut ChildrenVisit<'operator>, + ) -> Result<(), DatabaseError> { ExprVisitor::visit(self, op.predicate, self.arena) } - fn visit_join(&mut self, op: &'operator JoinOperator) -> Result<(), DatabaseError> { + fn visit_join( + &mut self, + op: &'operator JoinOperator, + _children: &mut ChildrenVisit<'operator>, + ) -> Result<(), DatabaseError> { if let JoinCondition::On { on, filter } = &op.on { for (left_expr, right_expr) in on { ExprVisitor::visit(self, *left_expr, self.arena)?; @@ -411,6 +836,7 @@ impl Operator { fn visit_project( &mut self, op: &'operator ProjectOperator, + _children: &mut ChildrenVisit<'operator>, ) -> Result<(), DatabaseError> { for expr in &op.exprs { ExprVisitor::visit(self, *expr, self.arena)?; @@ -422,8 +848,8 @@ impl Operator { &mut self, op: &'operator TableScanOperator, ) -> Result<(), DatabaseError> { - for column in &op.columns { - ExprVisitor::visit_column_ref(self, column)?; + for (position, column) in op.columns.iter().enumerate() { + ExprVisitor::visit_column_ref(self, column, position)?; } Ok(()) } @@ -438,7 +864,11 @@ impl Operator { Ok(()) } - fn visit_sort(&mut self, op: &'operator SortOperator) -> Result<(), DatabaseError> { + fn visit_sort( + &mut self, + op: &'operator SortOperator, + _children: &mut ChildrenVisit<'operator>, + ) -> Result<(), DatabaseError> { for field in &op.sort_fields { ExprVisitor::visit(self, field.expr, self.arena)?; } @@ -448,6 +878,7 @@ impl Operator { fn visit_window( &mut self, op: &'operator window::WindowOperator, + _children: &mut ChildrenVisit<'operator>, ) -> Result<(), DatabaseError> { for expr in op .sort_fields @@ -460,7 +891,11 @@ impl Operator { Ok(()) } - fn visit_top_k(&mut self, op: &'operator TopKOperator) -> Result<(), DatabaseError> { + fn visit_top_k( + &mut self, + op: &'operator TopKOperator, + _children: &mut ChildrenVisit<'operator>, + ) -> Result<(), DatabaseError> { for field in &op.sort_fields { ExprVisitor::visit(self, field.expr, self.arena)?; } @@ -468,15 +903,24 @@ impl Operator { } fn visit_values(&mut self, op: &'operator ValuesOperator) -> Result<(), DatabaseError> { - for column in &op.schema_ref { - ExprVisitor::visit_column_ref(self, column)?; + for (position, column) in op.schema_ref.iter().enumerate() { + ExprVisitor::visit_column_ref(self, column, position)?; } Ok(()) } - fn visit_union(&mut self, op: &'operator UnionOperator) -> Result<(), DatabaseError> { - for column in op.left_schema_ref.iter().chain(&op._right_schema_ref) { - ExprVisitor::visit_column_ref(self, column)?; + fn visit_union( + &mut self, + op: &'operator UnionOperator, + _children: &mut ChildrenVisit<'operator>, + ) -> Result<(), DatabaseError> { + for (position, column) in op + .left_schema_ref + .iter() + .chain(&op._right_schema_ref) + .enumerate() + { + ExprVisitor::visit_column_ref(self, column, position)?; } Ok(()) } @@ -484,9 +928,10 @@ impl Operator { fn visit_recursive_cte( &mut self, op: &'operator RecursiveCteOperator, + _children: &mut ChildrenVisit<'operator>, ) -> Result<(), DatabaseError> { - for column in &op.schema_ref { - ExprVisitor::visit_column_ref(self, column)?; + for (position, column) in op.schema_ref.iter().enumerate() { + ExprVisitor::visit_column_ref(self, column, position)?; } Ok(()) } @@ -495,8 +940,8 @@ impl Operator { &mut self, op: &'operator RecursiveScanOperator, ) -> Result<(), DatabaseError> { - for column in &op.schema_ref { - ExprVisitor::visit_column_ref(self, column)?; + for (position, column) in op.schema_ref.iter().enumerate() { + ExprVisitor::visit_column_ref(self, column, position)?; } Ok(()) } @@ -504,21 +949,35 @@ impl Operator { fn visit_set_membership( &mut self, op: &'operator SetMembershipOperator, + _children: &mut ChildrenVisit<'operator>, ) -> Result<(), DatabaseError> { - for column in op.left_schema_ref.iter().chain(&op._right_schema_ref) { - ExprVisitor::visit_column_ref(self, column)?; + for (position, column) in op + .left_schema_ref + .iter() + .chain(&op._right_schema_ref) + .enumerate() + { + ExprVisitor::visit_column_ref(self, column, position)?; } Ok(()) } - fn visit_delete(&mut self, op: &'operator DeleteOperator) -> Result<(), DatabaseError> { - for column in &op.primary_keys { - ExprVisitor::visit_column_ref(self, column)?; + fn visit_delete( + &mut self, + op: &'operator DeleteOperator, + _children: &mut ChildrenVisit<'operator>, + ) -> Result<(), DatabaseError> { + for (position, column) in op.primary_keys.iter().enumerate() { + ExprVisitor::visit_column_ref(self, column, position)?; } Ok(()) } - fn visit_update(&mut self, op: &'operator UpdateOperator) -> Result<(), DatabaseError> { + fn visit_update( + &mut self, + op: &'operator UpdateOperator, + _children: &mut ChildrenVisit<'operator>, + ) -> Result<(), DatabaseError> { for (_, expr) in &op.value_exprs { ExprVisitor::visit(self, *expr, self.arena)?; } @@ -563,7 +1022,12 @@ impl Operator { f, keep_going: true, }; - visitor.visit_operator(self)?; + match root { + ColumnVisitRoot::Operator(operator) => { + visitor.visit_operator(operator, None, &mut ChildrenVisit::empty())? + } + ColumnVisitRoot::Plan(plan) => visitor.visit_plan(plan)?, + } Ok(visitor.keep_going) } @@ -595,9 +1059,9 @@ impl Operator { } impl Explain for PlanImpl { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { match self { @@ -608,6 +1072,7 @@ impl Explain for PlanImpl { PlanImpl::StreamDistinct => f.write_str("StreamDistinct"), PlanImpl::ScalarQueryInit => f.write_str("ScalarQueryInit"), PlanImpl::MarkApply => f.write_str("MarkApply"), + PlanImpl::IndexProbeJoin => f.write_str("IndexProbeJoin"), PlanImpl::Filter => f.write_str("Filter"), PlanImpl::HashJoin => f.write_str("HashJoin"), PlanImpl::NestLoopJoin => f.write_str("NestLoopJoin"), @@ -641,9 +1106,9 @@ impl Explain for PlanImpl { } impl Explain for SortOption { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { match self { @@ -662,9 +1127,9 @@ impl Explain for SortOption { } impl Explain for PhysicalOption { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!( @@ -700,6 +1165,59 @@ mod tests { use crate::types::value::DataValue; use crate::types::LogicalType; + #[test] + fn rows_work_charges_rows_without_columns() { + assert_eq!(rows_work(3, 0), 3 * CPU_ROW_WORK); + assert_eq!(rows_work(3, 2), 3 * (CPU_ROW_WORK + 2)); + assert_eq!(rows_work(usize::MAX, 1), usize::MAX); + } + + #[test] + fn join_work_includes_inputs_and_prefers_cheap_sides() { + let small = InputCost { + rows: 10, + width: 1, + work: 50, + }; + let large = InputCost { + rows: 1000, + width: 1, + work: 5000, + }; + let output = rows_work(100, 2); + // Both inputs read once, the build side stored and hashed, every probe row looked up once. + assert_eq!( + JoinImpl::Hash.work(&small, &large, 100, 2), + 50 + 5000 + rows_work(10, 1) + 10 * HASH_BUILD_WORK + 1000 + output + ); + assert!( + JoinImpl::Hash.work(&small, &large, 100, 2) + < JoinImpl::Hash.work(&large, &small, 100, 2) + ); + // The inner input is re-run and compared for every outer row. + let nest_loop = JoinImpl::NestLoop { full: false }; + assert_eq!( + nest_loop.work(&small, &large, 100, 2), + 50 + 10 * (5000 + 1000) + output + ); + assert!(nest_loop.work(&small, &large, 100, 2) < nest_loop.work(&large, &small, 100, 2)); + // A full join reads the inner input once more for its unmatched rows. + assert_eq!( + JoinImpl::NestLoop { full: true }.work(&small, &large, 100, 2), + nest_loop.work(&small, &large, 100, 2) + 5000 + ); + // The probed input is never read whole: a few outer rows beat hashing a large input. + let probe = JoinImpl::IndexProbe { lookup_work: 30 }; + assert_eq!(probe.work(&small, &large, 100, 2), 50 + 10 * 30 + output); + assert!(probe.work(&small, &large, 100, 2) < JoinImpl::Hash.work(&small, &large, 100, 2)); + let huge = InputCost { + rows: usize::MAX, + width: 1, + work: usize::MAX, + }; + assert_eq!(nest_loop.work(&huge, &small, 0, 2), usize::MAX); + } + fn column_catalog(name: &str) -> ColumnCatalog { ColumnCatalog::new( name.to_string(), @@ -760,7 +1278,8 @@ mod tests { assert_eq!(SortOption::Follow.explain(&arena).to_string(), "Follow"); assert_eq!(SortOption::None.explain(&arena).to_string(), "None"); - let physical = PhysicalOption::new(PlanImpl::TopK, sort_option.clone()); + let physical = + PhysicalOption::new(PlanImpl::TopK, sort_option.clone(), RowCountOption::Unknown); assert_eq!( physical.explain(&arena).to_string(), "TopK => (Sort Option: OrderBy: (1 Desc Nulls First) ignore_prefix_len: 2)" @@ -780,6 +1299,7 @@ mod tests { (PlanImpl::StreamDistinct, "StreamDistinct"), (PlanImpl::ScalarQueryInit, "ScalarQueryInit"), (PlanImpl::MarkApply, "MarkApply"), + (PlanImpl::IndexProbeJoin, "IndexProbeJoin"), (PlanImpl::Filter, "Filter"), (PlanImpl::HashJoin, "HashJoin"), (PlanImpl::NestLoopJoin, "NestLoopJoin"), @@ -909,9 +1429,10 @@ mod tests { assert_eq!(referenced_columns(&join, &mut arena)?, vec![a, b, c]); assert!(!join.all_referenced_columns(&arena, |column| *column == a)?); - let project = Operator::Project(ProjectOperator { - exprs: vec![column_expr(b, 1, &mut arena), column_expr(c, 2, &mut arena)], - }); + let project = Operator::Project(ProjectOperator::new(vec![ + column_expr(b, 1, &mut arena), + column_expr(c, 2, &mut arena), + ])); assert_eq!(referenced_columns(&project, &mut arena)?, vec![b, c]); let update = Operator::Update(UpdateOperator { @@ -1024,6 +1545,8 @@ mod tests { table_name: "users".into(), index_metas: vec![IndexMetaRef::new(1)], histogram_buckets: Some(8), + hll_precision: None, + statistics_relative_error: None, }), ]; for operator in no_reference_operators { @@ -1267,6 +1790,8 @@ mod tests { table_name: "users".into(), index_metas: vec![index], histogram_buckets: Some(128), + hll_precision: None, + statistics_relative_error: None, }), "Analyze users -> [idx_users_id]", ), diff --git a/src/planner/operator/project.rs b/src/planner/operator/project.rs index 01d74216..00140a7a 100644 --- a/src/planner/operator/project.rs +++ b/src/planner/operator/project.rs @@ -21,10 +21,16 @@ pub struct ProjectOperator { pub exprs: Vec, } +impl ProjectOperator { + pub fn new(exprs: Vec) -> Self { + Self { exprs } + } +} + impl Explain for ProjectOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { f.write_str("Projection [")?; diff --git a/src/planner/operator/recursive_cte.rs b/src/planner/operator/recursive_cte.rs index 1a420655..6840dde8 100644 --- a/src/planner/operator/recursive_cte.rs +++ b/src/planner/operator/recursive_cte.rs @@ -36,9 +36,9 @@ impl RecursiveCteOperator { } impl Explain for RecursiveCteOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { f.write_str("Recursive CTE: [")?; @@ -53,9 +53,9 @@ pub struct RecursiveScanOperator { } impl Explain for RecursiveScanOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { f.write_str("Recursive Scan: [")?; diff --git a/src/planner/operator/scalar_query_init.rs b/src/planner/operator/scalar_query_init.rs index 1e3b22b8..9cbd2b3b 100644 --- a/src/planner/operator/scalar_query_init.rs +++ b/src/planner/operator/scalar_query_init.rs @@ -52,7 +52,7 @@ impl ScalarQueryInitOperator { } impl Explain for ScalarQueryInitOperator { - fn fmt(&self, arena: &dyn MetaArena, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn fmt(&self, arena: &A, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "ScalarQueryInit {}", self.value.explain(arena)) } } diff --git a/src/planner/operator/set_membership.rs b/src/planner/operator/set_membership.rs index 47b538c0..f32d3b8b 100644 --- a/src/planner/operator/set_membership.rs +++ b/src/planner/operator/set_membership.rs @@ -64,9 +64,9 @@ impl SetMembershipOperator { } impl Explain for SetMembershipOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!(f, "{}: [", self.kind.name())?; diff --git a/src/planner/operator/sort.rs b/src/planner/operator/sort.rs index d930db0e..d298cf74 100644 --- a/src/planner/operator/sort.rs +++ b/src/planner/operator/sort.rs @@ -65,9 +65,9 @@ pub struct SortOperator { } impl Explain for SortOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { f.write_str("Sort By ")?; @@ -76,9 +76,9 @@ impl Explain for SortOperator { } impl Explain for SortField { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { let direction = if self.asc { "Asc" } else { "Desc" }; diff --git a/src/planner/operator/table_scan.rs b/src/planner/operator/table_scan.rs index 599ca376..8bf22a4b 100644 --- a/src/planner/operator/table_scan.rs +++ b/src/planner/operator/table_scan.rs @@ -92,9 +92,9 @@ impl TableScanOperator { } impl Explain for TableScanOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!(f, "TableScan {} -> [", self.table_name)?; diff --git a/src/planner/operator/top_k.rs b/src/planner/operator/top_k.rs index 7eb00cd9..c8624be0 100644 --- a/src/planner/operator/top_k.rs +++ b/src/planner/operator/top_k.rs @@ -44,9 +44,9 @@ impl TopKOperator { } impl Explain for TopKOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!(f, "Top {}, ", self.limit)?; diff --git a/src/planner/operator/union.rs b/src/planner/operator/union.rs index 95af467d..9c0179ce 100644 --- a/src/planner/operator/union.rs +++ b/src/planner/operator/union.rs @@ -46,9 +46,9 @@ impl UnionOperator { } impl Explain for UnionOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { f.write_str("Union: [")?; diff --git a/src/planner/operator/update.rs b/src/planner/operator/update.rs index d35b1ff3..472f5516 100644 --- a/src/planner/operator/update.rs +++ b/src/planner/operator/update.rs @@ -24,9 +24,9 @@ pub struct UpdateOperator { } impl Explain for UpdateOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { write!(f, "Update {} set ", self.table_name)?; diff --git a/src/planner/operator/values.rs b/src/planner/operator/values.rs index 0ebc5f2e..e0781ee2 100644 --- a/src/planner/operator/values.rs +++ b/src/planner/operator/values.rs @@ -41,7 +41,7 @@ impl ValuesOperator { } impl Explain for ValuesOperator { - fn fmt(&self, arena: &(dyn MetaArena + '_), f: &mut Formatter) -> fmt::Result { + fn fmt(&self, arena: &A, f: &mut Formatter) -> fmt::Result { f.write_str("Values ")?; let width = self.schema_ref.len(); for i in 0..self.row_count { diff --git a/src/planner/operator/visitor.rs b/src/planner/operator/visitor.rs index 668f6768..e60699c3 100644 --- a/src/planner/operator/visitor.rs +++ b/src/planner/operator/visitor.rs @@ -16,47 +16,135 @@ use super::alter_table::change_column::DefaultChange; use super::*; use crate::errors::DatabaseError; use crate::expression::visitor::ExprVisitor; -use crate::planner::MetaArena; +use crate::planner::{Childrens, LogicalPlan, MetaArena}; + +/// One node's children. Taking the reference records that traversal was handled. +pub struct ChildrenVisit<'a> { + children: Option<&'a Childrens>, +} + +impl<'a> ChildrenVisit<'a> { + pub fn new(children: &'a Childrens) -> Self { + Self { + children: Some(children), + } + } + + pub fn empty() -> Self { + Self { children: None } + } + + pub fn visit( + &mut self, + walk: impl FnOnce(&'a Childrens) -> Result<(), DatabaseError>, + ) -> Result<(), DatabaseError> { + if let Some(children) = self.children.take() { + walk(children)?; + } + Ok(()) + } +} pub trait OperatorVisitor<'a>: Sized { - fn visit_operator(&mut self, operator: &'a Operator) -> Result<(), DatabaseError> { - walk_operator(self, operator) + fn visit_plan(&mut self, plan: &'a LogicalPlan) -> Result<(), DatabaseError> { + walk_plan(self, plan) + } + + fn visit_children(&mut self, childrens: &'a Childrens) -> Result<(), DatabaseError> { + match childrens { + Childrens::Only(child) => self.visit_plan(child), + Childrens::Twins { left, right } => { + self.visit_plan(left)?; + self.visit_plan(right) + } + Childrens::None => Ok(()), + } + } + + fn visit_operator( + &mut self, + operator: &'a Operator, + physical_option: Option<&'a PhysicalOption>, + children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { + walk_operator(self, operator, children)?; + if let Some(physical_option) = physical_option { + self.visit_physical_option(physical_option)?; + } + Ok(()) + } + + fn visit_physical_option( + &mut self, + _physical_option: &'a PhysicalOption, + ) -> Result<(), DatabaseError> { + Ok(()) } fn visit_dummy(&mut self) -> Result<(), DatabaseError> { Ok(()) } - fn visit_aggregate(&mut self, _op: &'a AggregateOperator) -> Result<(), DatabaseError> { + fn visit_aggregate( + &mut self, + _op: &'a AggregateOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } fn visit_scalar_query_init( &mut self, _op: &'a ScalarQueryInitOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_mark_apply(&mut self, _op: &'a MarkApplyOperator) -> Result<(), DatabaseError> { + fn visit_mark_apply( + &mut self, + _op: &'a MarkApplyOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_filter(&mut self, _op: &'a FilterOperator) -> Result<(), DatabaseError> { + fn visit_index_probe_join( + &mut self, + _op: &'a IndexProbeJoinOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_join(&mut self, _op: &'a JoinOperator) -> Result<(), DatabaseError> { + fn visit_filter( + &mut self, + _op: &'a FilterOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_project(&mut self, _op: &'a ProjectOperator) -> Result<(), DatabaseError> { + fn visit_join( + &mut self, + _op: &'a JoinOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_project( + &mut self, + _op: &'a ProjectOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } fn visit_scalar_subquery( &mut self, _op: &'a ScalarSubqueryOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { Ok(()) } @@ -69,19 +157,35 @@ pub trait OperatorVisitor<'a>: Sized { Ok(()) } - fn visit_sort(&mut self, _op: &'a SortOperator) -> Result<(), DatabaseError> { + fn visit_sort( + &mut self, + _op: &'a SortOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_window(&mut self, _op: &'a window::WindowOperator) -> Result<(), DatabaseError> { + fn visit_window( + &mut self, + _op: &'a window::WindowOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_limit(&mut self, _op: &'a LimitOperator) -> Result<(), DatabaseError> { + fn visit_limit( + &mut self, + _op: &'a LimitOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_top_k(&mut self, _op: &'a TopKOperator) -> Result<(), DatabaseError> { + fn visit_top_k( + &mut self, + _op: &'a TopKOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } @@ -97,7 +201,7 @@ pub trait OperatorVisitor<'a>: Sized { Ok(()) } - fn visit_explain(&mut self) -> Result<(), DatabaseError> { + fn visit_explain(&mut self, _children: &mut ChildrenVisit<'a>) -> Result<(), DatabaseError> { Ok(()) } @@ -108,15 +212,24 @@ pub trait OperatorVisitor<'a>: Sized { fn visit_set_membership( &mut self, _op: &'a SetMembershipOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_union(&mut self, _op: &'a UnionOperator) -> Result<(), DatabaseError> { + fn visit_union( + &mut self, + _op: &'a UnionOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_recursive_cte(&mut self, _op: &'a RecursiveCteOperator) -> Result<(), DatabaseError> { + fn visit_recursive_cte( + &mut self, + _op: &'a RecursiveCteOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } @@ -127,19 +240,35 @@ pub trait OperatorVisitor<'a>: Sized { Ok(()) } - fn visit_insert(&mut self, _op: &'a InsertOperator) -> Result<(), DatabaseError> { + fn visit_insert( + &mut self, + _op: &'a InsertOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_update(&mut self, _op: &'a UpdateOperator) -> Result<(), DatabaseError> { + fn visit_update( + &mut self, + _op: &'a UpdateOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_delete(&mut self, _op: &'a DeleteOperator) -> Result<(), DatabaseError> { + fn visit_delete( + &mut self, + _op: &'a DeleteOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_analyze(&mut self, _op: &'a AnalyzeOperator) -> Result<(), DatabaseError> { + fn visit_analyze( + &mut self, + _op: &'a AnalyzeOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } @@ -159,7 +288,11 @@ pub trait OperatorVisitor<'a>: Sized { Ok(()) } - fn visit_create_index(&mut self, _op: &'a CreateIndexOperator) -> Result<(), DatabaseError> { + fn visit_create_index( + &mut self, + _op: &'a CreateIndexOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } @@ -189,7 +322,11 @@ pub trait OperatorVisitor<'a>: Sized { } #[cfg(feature = "copy")] - fn visit_copy_to_file(&mut self, _op: &'a CopyToFileOperator) -> Result<(), DatabaseError> { + fn visit_copy_to_file( + &mut self, + _op: &'a CopyToFileOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } } @@ -215,7 +352,11 @@ impl<'a, V: ExprVisitor, A: MetaArena + ?Sized> OperatorVisitor<'a> Ok(()) } - fn visit_aggregate(&mut self, op: &'a AggregateOperator) -> Result<(), DatabaseError> { + fn visit_aggregate( + &mut self, + op: &'a AggregateOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { for expr in op.agg_calls.iter().chain(&op.groupby_exprs) { ExprVisitor::visit(self.visitor, *expr, self.arena)?; } @@ -225,6 +366,7 @@ impl<'a, V: ExprVisitor, A: MetaArena + ?Sized> OperatorVisitor<'a> fn visit_scalar_query_init( &mut self, op: &'a ScalarQueryInitOperator, + _children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { ExprVisitor::visit(self.visitor, op.value, self.arena)?; for (_, expr) in &op.param_bindings { @@ -233,21 +375,48 @@ impl<'a, V: ExprVisitor, A: MetaArena + ?Sized> OperatorVisitor<'a> Ok(()) } - fn visit_mark_apply(&mut self, op: &'a MarkApplyOperator) -> Result<(), DatabaseError> { + fn visit_mark_apply( + &mut self, + op: &'a MarkApplyOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { for expr in &op.predicates { ExprVisitor::visit(self.visitor, *expr, self.arena)?; } if let Some(expr) = &op.parameterized_probe { ExprVisitor::visit(self.visitor, *expr, self.arena)?; } + if let Some(expr) = op.guard { + ExprVisitor::visit(self.visitor, expr, self.arena)?; + } + Ok(()) + } + + fn visit_index_probe_join( + &mut self, + op: &'a IndexProbeJoinOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { + ExprVisitor::visit(self.visitor, op.probe, self.arena)?; + for expr in &op.predicates { + ExprVisitor::visit(self.visitor, *expr, self.arena)?; + } Ok(()) } - fn visit_filter(&mut self, op: &'a FilterOperator) -> Result<(), DatabaseError> { + fn visit_filter( + &mut self, + op: &'a FilterOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { ExprVisitor::visit(self.visitor, op.predicate, self.arena) } - fn visit_join(&mut self, op: &'a JoinOperator) -> Result<(), DatabaseError> { + fn visit_join( + &mut self, + op: &'a JoinOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { if let JoinCondition::On { on, filter } = &op.on { for (left_expr, right_expr) in on { ExprVisitor::visit(self.visitor, *left_expr, self.arena)?; @@ -260,7 +429,11 @@ impl<'a, V: ExprVisitor, A: MetaArena + ?Sized> OperatorVisitor<'a> Ok(()) } - fn visit_project(&mut self, op: &'a ProjectOperator) -> Result<(), DatabaseError> { + fn visit_project( + &mut self, + op: &'a ProjectOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { for expr in &op.exprs { ExprVisitor::visit(self.visitor, *expr, self.arena)?; } @@ -288,14 +461,22 @@ impl<'a, V: ExprVisitor, A: MetaArena + ?Sized> OperatorVisitor<'a> Ok(()) } - fn visit_sort(&mut self, op: &'a SortOperator) -> Result<(), DatabaseError> { + fn visit_sort( + &mut self, + op: &'a SortOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { for field in &op.sort_fields { ExprVisitor::visit(self.visitor, field.expr, self.arena)?; } Ok(()) } - fn visit_window(&mut self, op: &'a window::WindowOperator) -> Result<(), DatabaseError> { + fn visit_window( + &mut self, + op: &'a window::WindowOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { for expr in op .sort_fields .iter() @@ -307,14 +488,22 @@ impl<'a, V: ExprVisitor, A: MetaArena + ?Sized> OperatorVisitor<'a> Ok(()) } - fn visit_top_k(&mut self, op: &'a TopKOperator) -> Result<(), DatabaseError> { + fn visit_top_k( + &mut self, + op: &'a TopKOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { for field in &op.sort_fields { ExprVisitor::visit(self.visitor, field.expr, self.arena)?; } Ok(()) } - fn visit_update(&mut self, op: &'a UpdateOperator) -> Result<(), DatabaseError> { + fn visit_update( + &mut self, + op: &'a UpdateOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { for (_, expr) in &op.value_exprs { ExprVisitor::visit(self.visitor, *expr, self.arena)?; } @@ -345,42 +534,59 @@ impl<'a, V: ExprVisitor, A: MetaArena + ?Sized> OperatorVisitor<'a> } } +pub fn walk_plan<'a, V: OperatorVisitor<'a>>( + visitor: &mut V, + plan: &'a LogicalPlan, +) -> Result<(), DatabaseError> { + let LogicalPlan { + operator, + physical_option, + childrens, + .. + } = plan; + let mut children = ChildrenVisit::new(childrens); + visitor.visit_operator(operator, physical_option.as_ref(), &mut children)?; + children.visit(|children| visitor.visit_children(children)) +} + pub fn walk_operator<'a, V: OperatorVisitor<'a>>( visitor: &mut V, operator: &'a Operator, + children: &mut ChildrenVisit<'a>, ) -> Result<(), DatabaseError> { match operator { Operator::Dummy => visitor.visit_dummy(), - Operator::Aggregate(op) => visitor.visit_aggregate(op), - Operator::ScalarQueryInit(op) => visitor.visit_scalar_query_init(op), - Operator::MarkApply(op) => visitor.visit_mark_apply(op), - Operator::Filter(op) => visitor.visit_filter(op), - Operator::Join(op) => visitor.visit_join(op), - Operator::Project(op) => visitor.visit_project(op), - Operator::ScalarSubquery(op) => visitor.visit_scalar_subquery(op), + Operator::Aggregate(op) => visitor.visit_aggregate(op, children), + Operator::ScalarQueryInit(op) => visitor.visit_scalar_query_init(op, children), + Operator::MarkApply(op) => visitor.visit_mark_apply(op, children), + Operator::IndexProbeJoin(op) => visitor.visit_index_probe_join(op, children), + Operator::Filter(op) => visitor.visit_filter(op, children), + Operator::Join(op) => visitor.visit_join(op, children), + Operator::Project(op) => visitor.visit_project(op, children), + Operator::ScalarSubquery(op) => visitor.visit_scalar_subquery(op, children), Operator::TableScan(op) => visitor.visit_table_scan(op), Operator::FunctionScan(op) => visitor.visit_function_scan(op), - Operator::Sort(op) => visitor.visit_sort(op), - Operator::Limit(op) => visitor.visit_limit(op), - Operator::TopK(op) => visitor.visit_top_k(op), + Operator::Sort(op) => visitor.visit_sort(op, children), + Operator::Limit(op) => visitor.visit_limit(op, children), + Operator::TopK(op) => visitor.visit_top_k(op, children), Operator::Values(op) => visitor.visit_values(op), Operator::ShowTable => visitor.visit_show_table(), Operator::ShowView => visitor.visit_show_view(), - Operator::Explain => visitor.visit_explain(), + Operator::Explain => visitor.visit_explain(children), Operator::Describe(op) => visitor.visit_describe(op), - Operator::SetMembership(op) => visitor.visit_set_membership(op), - Operator::Union(op) => visitor.visit_union(op), - Operator::RecursiveCte(op) => visitor.visit_recursive_cte(op), + Operator::SetMembership(op) => visitor.visit_set_membership(op, children), + Operator::Union(op) => visitor.visit_union(op, children), + Operator::RecursiveCte(op) => visitor.visit_recursive_cte(op, children), Operator::RecursiveScan(op) => visitor.visit_recursive_scan(op), - Operator::Insert(op) => visitor.visit_insert(op), - Operator::Update(op) => visitor.visit_update(op), - Operator::Delete(op) => visitor.visit_delete(op), - Operator::Analyze(op) => visitor.visit_analyze(op), + Operator::Insert(op) => visitor.visit_insert(op, children), + Operator::Update(op) => visitor.visit_update(op, children), + Operator::Delete(op) => visitor.visit_delete(op, children), + Operator::Analyze(op) => visitor.visit_analyze(op, children), Operator::AddColumn(op) => visitor.visit_add_column(op), Operator::ChangeColumn(op) => visitor.visit_change_column(op), Operator::DropColumn(op) => visitor.visit_drop_column(op), Operator::CreateTable(op) => visitor.visit_create_table(op), - Operator::CreateIndex(op) => visitor.visit_create_index(op), + Operator::CreateIndex(op) => visitor.visit_create_index(op, children), Operator::CreateView(op) => visitor.visit_create_view(op), Operator::DropTable(op) => visitor.visit_drop_table(op), Operator::DropView(op) => visitor.visit_drop_view(op), @@ -389,8 +595,8 @@ pub fn walk_operator<'a, V: OperatorVisitor<'a>>( #[cfg(feature = "copy")] Operator::CopyFromFile(op) => visitor.visit_copy_from_file(op), #[cfg(feature = "copy")] - Operator::CopyToFile(op) => visitor.visit_copy_to_file(op), - Operator::Window(op) => visitor.visit_window(op), + Operator::CopyToFile(op) => visitor.visit_copy_to_file(op, children), + Operator::Window(op) => visitor.visit_window(op, children), } } @@ -412,6 +618,7 @@ pub(crate) mod tests { use crate::planner::operator::join::{JoinOperator, JoinType}; use crate::planner::operator::mark_apply::MarkApplyOperator; use crate::planner::operator::set_membership::SetMembershipKind; + use crate::planner::operator::visitor_mut::{ChildrenVisitMut, OperatorVisitorMut}; use crate::planner::test::PlanArenaTestExt; use crate::planner::ExprRef; use crate::planner::{Childrens, LogicalPlan}; @@ -468,6 +675,7 @@ pub(crate) mod tests { value: expr(0), }), Operator::MarkApply(mark_apply), + Operator::IndexProbeJoin(IndexProbeJoinOperator::new(expr(16), vec![expr(3)])), Operator::Filter(FilterOperator { predicate: expr(5), is_optimized: false, @@ -482,9 +690,7 @@ pub(crate) mod tests { force_nested_loop: false, limit_pushed: false, }), - Operator::Project(ProjectOperator { - exprs: vec![expr(9)], - }), + Operator::Project(ProjectOperator::new(vec![expr(9)])), Operator::ScalarSubquery(ScalarSubqueryOperator), Operator::TableScan(TableScanOperator { table_name: "t1".into(), @@ -553,6 +759,8 @@ pub(crate) mod tests { table_name: "t1".into(), index_metas: vec![IndexMetaRef::new(0)], histogram_buckets: Some(1), + hll_precision: None, + statistics_relative_error: None, }), Operator::AddColumn(AddColumnOperator { table_name: "t1".into(), @@ -656,6 +864,85 @@ pub(crate) mod tests { } } + #[test] + fn read_only_and_mutable_visitors_follow_the_same_order() -> Result<(), DatabaseError> { + #[derive(Default)] + struct VisitOrder(Vec<&'static str>); + + impl<'a> OperatorVisitor<'a> for VisitOrder { + fn visit_join( + &mut self, + _op: &'a JoinOperator, + _children: &mut ChildrenVisit<'a>, + ) -> Result<(), DatabaseError> { + self.0.push("operator"); + Ok(()) + } + + fn visit_dummy(&mut self) -> Result<(), DatabaseError> { + self.0.push("operator"); + Ok(()) + } + + fn visit_physical_option( + &mut self, + _option: &'a PhysicalOption, + ) -> Result<(), DatabaseError> { + self.0.push("physical"); + Ok(()) + } + } + + impl<'a> OperatorVisitorMut<'a> for VisitOrder { + fn visit_join( + &mut self, + _op: &'a mut JoinOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { + self.0.push("operator"); + Ok(()) + } + + fn visit_dummy(&mut self) -> Result<(), DatabaseError> { + self.0.push("operator"); + Ok(()) + } + + fn visit_physical_option( + &mut self, + _option: &'a mut PhysicalOption, + ) -> Result<(), DatabaseError> { + self.0.push("physical"); + Ok(()) + } + } + + let left = LogicalPlan::new(Operator::Dummy, Childrens::None); + let mut right = LogicalPlan::new(Operator::Dummy, Childrens::None); + right.physical_option = Some(PhysicalOption::new( + PlanImpl::Dummy, + SortOption::None, + RowCountOption::Unknown, + )); + let mut plan = + JoinOperator::build(left, right, JoinCondition::None, JoinType::Cross, false); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::NestLoopJoin, + SortOption::None, + RowCountOption::Unknown, + )); + let mut read_only = VisitOrder::default(); + OperatorVisitor::visit_plan(&mut read_only, &plan)?; + let mut mutable = VisitOrder::default(); + OperatorVisitorMut::visit_plan(&mut mutable, &mut plan)?; + assert_eq!( + read_only.0, + ["operator", "physical", "operator", "operator", "physical"] + ); + assert_eq!(read_only.0, mutable.0); + Ok(()) + } + #[test] fn dispatches_all_variants_and_visits_expressions() -> Result<(), DatabaseError> { struct NoopVisitor; @@ -665,15 +952,15 @@ pub(crate) mod tests { let mut arena = crate::planner::PlanArena::new(&table_arena); let operators = all_operators(&mut arena)?; for operator in &operators { - NoopVisitor.visit_operator(operator)?; + NoopVisitor.visit_operator(operator, None, &mut ChildrenVisit::empty())?; } let mut counter = ExpressionCounter::default(); let mut visitor = OperatorExprVisitor::new(&mut counter, &arena); for operator in &operators { - visitor.visit_operator(operator)?; + visitor.visit_operator(operator, None, &mut ChildrenVisit::empty())?; } - assert_eq!(counter.0, 22); // Includes the Values row expression. + assert_eq!(counter.0, 24); // Includes the Values row expression. Ok(()) } diff --git a/src/planner/operator/visitor_mut.rs b/src/planner/operator/visitor_mut.rs index 10dee0e5..4b9166be 100644 --- a/src/planner/operator/visitor_mut.rs +++ b/src/planner/operator/visitor_mut.rs @@ -23,16 +23,40 @@ use crate::expression::visitor_mut::ExprVisitorMut; use crate::planner::MetaArena; use crate::planner::{Childrens, LogicalPlan}; +/// One node's children. Taking the reference records that traversal was handled. +pub struct ChildrenVisitMut<'a> { + children: Option<&'a mut Childrens>, +} + +impl<'a> ChildrenVisitMut<'a> { + pub fn new(children: &'a mut Childrens) -> Self { + Self { + children: Some(children), + } + } + + pub fn empty() -> Self { + Self { children: None } + } + + pub fn visit( + &mut self, + walk: impl FnOnce(&'a mut Childrens) -> Result<(), DatabaseError>, + ) -> Result<(), DatabaseError> { + if let Some(children) = self.children.take() { + walk(children)?; + } + Ok(()) + } +} + pub trait OperatorVisitorMut<'a>: Sized { fn visit_plan(&mut self, plan: &'a mut LogicalPlan) -> Result<(), DatabaseError> { - let LogicalPlan { - operator, - physical_option, - childrens, - .. - } = plan; - self.visit_operator(operator, physical_option.as_mut())?; - match childrens.as_mut() { + walk_mut_plan(self, plan) + } + + fn visit_children(&mut self, childrens: &'a mut Childrens) -> Result<(), DatabaseError> { + match childrens { Childrens::Only(child) => self.visit_plan(child), Childrens::Twins { left, right } => { self.visit_plan(left)?; @@ -46,8 +70,9 @@ pub trait OperatorVisitorMut<'a>: Sized { &mut self, operator: &'a mut Operator, physical_option: Option<&'a mut PhysicalOption>, + children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { - walk_mut_operator(self, operator)?; + walk_mut_operator(self, operator, children)?; if let Some(physical_option) = physical_option { self.visit_physical_option(physical_option)?; } @@ -65,36 +90,66 @@ pub trait OperatorVisitorMut<'a>: Sized { Ok(()) } - fn visit_aggregate(&mut self, _op: &'a mut AggregateOperator) -> Result<(), DatabaseError> { + fn visit_aggregate( + &mut self, + _op: &'a mut AggregateOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } fn visit_scalar_query_init( &mut self, _op: &'a mut ScalarQueryInitOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_mark_apply( + &mut self, + _op: &'a mut MarkApplyOperator, + _children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_mark_apply(&mut self, _op: &'a mut MarkApplyOperator) -> Result<(), DatabaseError> { + fn visit_index_probe_join( + &mut self, + _op: &'a mut IndexProbeJoinOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_filter(&mut self, _op: &'a mut FilterOperator) -> Result<(), DatabaseError> { + fn visit_filter( + &mut self, + _op: &'a mut FilterOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_join(&mut self, _op: &'a mut JoinOperator) -> Result<(), DatabaseError> { + fn visit_join( + &mut self, + _op: &'a mut JoinOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_project(&mut self, _op: &'a mut ProjectOperator) -> Result<(), DatabaseError> { + fn visit_project( + &mut self, + _op: &'a mut ProjectOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } fn visit_scalar_subquery( &mut self, _op: &'a mut ScalarSubqueryOperator, + _children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { Ok(()) } @@ -110,19 +165,35 @@ pub trait OperatorVisitorMut<'a>: Sized { Ok(()) } - fn visit_sort(&mut self, _op: &'a mut SortOperator) -> Result<(), DatabaseError> { + fn visit_sort( + &mut self, + _op: &'a mut SortOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_window(&mut self, _op: &'a mut window::WindowOperator) -> Result<(), DatabaseError> { + fn visit_window( + &mut self, + _op: &'a mut window::WindowOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_limit(&mut self, _op: &'a mut LimitOperator) -> Result<(), DatabaseError> { + fn visit_limit( + &mut self, + _op: &'a mut LimitOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_top_k(&mut self, _op: &'a mut TopKOperator) -> Result<(), DatabaseError> { + fn visit_top_k( + &mut self, + _op: &'a mut TopKOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } @@ -138,7 +209,7 @@ pub trait OperatorVisitorMut<'a>: Sized { Ok(()) } - fn visit_explain(&mut self) -> Result<(), DatabaseError> { + fn visit_explain(&mut self, _children: &mut ChildrenVisitMut<'a>) -> Result<(), DatabaseError> { Ok(()) } @@ -149,17 +220,23 @@ pub trait OperatorVisitorMut<'a>: Sized { fn visit_set_membership( &mut self, _op: &'a mut SetMembershipOperator, + _children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_union(&mut self, _op: &'a mut UnionOperator) -> Result<(), DatabaseError> { + fn visit_union( + &mut self, + _op: &'a mut UnionOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } fn visit_recursive_cte( &mut self, _op: &'a mut RecursiveCteOperator, + _children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { Ok(()) } @@ -171,19 +248,35 @@ pub trait OperatorVisitorMut<'a>: Sized { Ok(()) } - fn visit_insert(&mut self, _op: &'a mut InsertOperator) -> Result<(), DatabaseError> { + fn visit_insert( + &mut self, + _op: &'a mut InsertOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_update(&mut self, _op: &'a mut UpdateOperator) -> Result<(), DatabaseError> { + fn visit_update( + &mut self, + _op: &'a mut UpdateOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_delete(&mut self, _op: &'a mut DeleteOperator) -> Result<(), DatabaseError> { + fn visit_delete( + &mut self, + _op: &'a mut DeleteOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } - fn visit_analyze(&mut self, _op: &'a mut AnalyzeOperator) -> Result<(), DatabaseError> { + fn visit_analyze( + &mut self, + _op: &'a mut AnalyzeOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } @@ -212,6 +305,7 @@ pub trait OperatorVisitorMut<'a>: Sized { fn visit_create_index( &mut self, _op: &'a mut CreateIndexOperator, + _children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { Ok(()) } @@ -245,23 +339,29 @@ pub trait OperatorVisitorMut<'a>: Sized { } #[cfg(feature = "copy")] - fn visit_copy_to_file(&mut self, _op: &'a mut CopyToFileOperator) -> Result<(), DatabaseError> { + fn visit_copy_to_file( + &mut self, + _op: &'a mut CopyToFileOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { Ok(()) } } -pub struct OperatorExprVisitorMut<'a, 'arena, V> { +pub struct OperatorExprVisitorMut<'a, V, A: ?Sized> { visitor: &'a mut V, - arena: &'a mut (dyn MetaArena + 'arena), + arena: &'a mut A, } -impl<'a, 'arena, V> OperatorExprVisitorMut<'a, 'arena, V> { - pub fn new(visitor: &'a mut V, arena: &'a mut (dyn MetaArena + 'arena)) -> Self { +impl<'a, V, A: ?Sized> OperatorExprVisitorMut<'a, V, A> { + pub fn new(visitor: &'a mut V, arena: &'a mut A) -> Self { Self { visitor, arena } } } -impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_, '_, V> { +impl<'a, V: ExprVisitorMut, A: MetaArena + ?Sized> OperatorVisitorMut<'a> + for OperatorExprVisitorMut<'_, V, A> +{ fn visit_values(&mut self, op: &'a mut ValuesOperator) -> Result<(), DatabaseError> { for expr in op.rows.iter_mut() { self.visitor.visit(expr, self.arena)?; @@ -269,7 +369,11 @@ impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_ Ok(()) } - fn visit_aggregate(&mut self, op: &'a mut AggregateOperator) -> Result<(), DatabaseError> { + fn visit_aggregate( + &mut self, + op: &'a mut AggregateOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { for expr in op.agg_calls.iter_mut().chain(&mut op.groupby_exprs) { ExprVisitorMut::visit(self.visitor, expr, self.arena)?; } @@ -279,6 +383,7 @@ impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_ fn visit_scalar_query_init( &mut self, op: &'a mut ScalarQueryInitOperator, + _children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { ExprVisitorMut::visit(self.visitor, &mut op.value, self.arena)?; for (_, expr) in &mut op.param_bindings { @@ -287,21 +392,48 @@ impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_ Ok(()) } - fn visit_mark_apply(&mut self, op: &'a mut MarkApplyOperator) -> Result<(), DatabaseError> { + fn visit_mark_apply( + &mut self, + op: &'a mut MarkApplyOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { for expr in &mut op.predicates { ExprVisitorMut::visit(self.visitor, expr, self.arena)?; } if let Some(expr) = &mut op.parameterized_probe { ExprVisitorMut::visit(self.visitor, expr, self.arena)?; } + if let Some(expr) = &mut op.guard { + ExprVisitorMut::visit(self.visitor, expr, self.arena)?; + } Ok(()) } - fn visit_filter(&mut self, op: &'a mut FilterOperator) -> Result<(), DatabaseError> { + fn visit_index_probe_join( + &mut self, + op: &'a mut IndexProbeJoinOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { + ExprVisitorMut::visit(self.visitor, &mut op.probe, self.arena)?; + for expr in &mut op.predicates { + ExprVisitorMut::visit(self.visitor, expr, self.arena)?; + } + Ok(()) + } + + fn visit_filter( + &mut self, + op: &'a mut FilterOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { ExprVisitorMut::visit(self.visitor, &mut op.predicate, self.arena) } - fn visit_join(&mut self, op: &'a mut JoinOperator) -> Result<(), DatabaseError> { + fn visit_join( + &mut self, + op: &'a mut JoinOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { if let JoinCondition::On { on, filter } = &mut op.on { for (left_expr, right_expr) in on { ExprVisitorMut::visit(self.visitor, left_expr, self.arena)?; @@ -314,7 +446,11 @@ impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_ Ok(()) } - fn visit_project(&mut self, op: &'a mut ProjectOperator) -> Result<(), DatabaseError> { + fn visit_project( + &mut self, + op: &'a mut ProjectOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { for expr in &mut op.exprs { ExprVisitorMut::visit(self.visitor, expr, self.arena)?; } @@ -345,14 +481,22 @@ impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_ Ok(()) } - fn visit_sort(&mut self, op: &'a mut SortOperator) -> Result<(), DatabaseError> { + fn visit_sort( + &mut self, + op: &'a mut SortOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { for field in &mut op.sort_fields { ExprVisitorMut::visit(self.visitor, &mut field.expr, self.arena)?; } Ok(()) } - fn visit_window(&mut self, op: &'a mut window::WindowOperator) -> Result<(), DatabaseError> { + fn visit_window( + &mut self, + op: &'a mut window::WindowOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { for expr in op .sort_fields .iter_mut() @@ -368,14 +512,22 @@ impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_ Ok(()) } - fn visit_top_k(&mut self, op: &'a mut TopKOperator) -> Result<(), DatabaseError> { + fn visit_top_k( + &mut self, + op: &'a mut TopKOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { for field in &mut op.sort_fields { ExprVisitorMut::visit(self.visitor, &mut field.expr, self.arena)?; } Ok(()) } - fn visit_update(&mut self, op: &'a mut UpdateOperator) -> Result<(), DatabaseError> { + fn visit_update( + &mut self, + op: &'a mut UpdateOperator, + _children: &mut ChildrenVisitMut<'a>, + ) -> Result<(), DatabaseError> { for (_, expr) in &mut op.value_exprs { ExprVisitorMut::visit(self.visitor, expr, self.arena)?; } @@ -409,42 +561,59 @@ impl<'a, V: ExprVisitorMut> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_ } } +pub fn walk_mut_plan<'a, V: OperatorVisitorMut<'a>>( + visitor: &mut V, + plan: &'a mut LogicalPlan, +) -> Result<(), DatabaseError> { + let LogicalPlan { + operator, + physical_option, + childrens, + .. + } = plan; + let mut children = ChildrenVisitMut::new(childrens); + visitor.visit_operator(operator, physical_option.as_mut(), &mut children)?; + children.visit(|children| visitor.visit_children(children)) +} + pub fn walk_mut_operator<'a, V: OperatorVisitorMut<'a>>( visitor: &mut V, operator: &'a mut Operator, + children: &mut ChildrenVisitMut<'a>, ) -> Result<(), DatabaseError> { match operator { Operator::Dummy => visitor.visit_dummy(), - Operator::Aggregate(op) => visitor.visit_aggregate(op), - Operator::ScalarQueryInit(op) => visitor.visit_scalar_query_init(op), - Operator::MarkApply(op) => visitor.visit_mark_apply(op), - Operator::Filter(op) => visitor.visit_filter(op), - Operator::Join(op) => visitor.visit_join(op), - Operator::Project(op) => visitor.visit_project(op), - Operator::ScalarSubquery(op) => visitor.visit_scalar_subquery(op), + Operator::Aggregate(op) => visitor.visit_aggregate(op, children), + Operator::ScalarQueryInit(op) => visitor.visit_scalar_query_init(op, children), + Operator::MarkApply(op) => visitor.visit_mark_apply(op, children), + Operator::IndexProbeJoin(op) => visitor.visit_index_probe_join(op, children), + Operator::Filter(op) => visitor.visit_filter(op, children), + Operator::Join(op) => visitor.visit_join(op, children), + Operator::Project(op) => visitor.visit_project(op, children), + Operator::ScalarSubquery(op) => visitor.visit_scalar_subquery(op, children), Operator::TableScan(op) => visitor.visit_table_scan(op), Operator::FunctionScan(op) => visitor.visit_function_scan(op), - Operator::Sort(op) => visitor.visit_sort(op), - Operator::Limit(op) => visitor.visit_limit(op), - Operator::TopK(op) => visitor.visit_top_k(op), + Operator::Sort(op) => visitor.visit_sort(op, children), + Operator::Limit(op) => visitor.visit_limit(op, children), + Operator::TopK(op) => visitor.visit_top_k(op, children), Operator::Values(op) => visitor.visit_values(op), Operator::ShowTable => visitor.visit_show_table(), Operator::ShowView => visitor.visit_show_view(), - Operator::Explain => visitor.visit_explain(), + Operator::Explain => visitor.visit_explain(children), Operator::Describe(op) => visitor.visit_describe(op), - Operator::SetMembership(op) => visitor.visit_set_membership(op), - Operator::Union(op) => visitor.visit_union(op), - Operator::RecursiveCte(op) => visitor.visit_recursive_cte(op), + Operator::SetMembership(op) => visitor.visit_set_membership(op, children), + Operator::Union(op) => visitor.visit_union(op, children), + Operator::RecursiveCte(op) => visitor.visit_recursive_cte(op, children), Operator::RecursiveScan(op) => visitor.visit_recursive_scan(op), - Operator::Insert(op) => visitor.visit_insert(op), - Operator::Update(op) => visitor.visit_update(op), - Operator::Delete(op) => visitor.visit_delete(op), - Operator::Analyze(op) => visitor.visit_analyze(op), + Operator::Insert(op) => visitor.visit_insert(op, children), + Operator::Update(op) => visitor.visit_update(op, children), + Operator::Delete(op) => visitor.visit_delete(op, children), + Operator::Analyze(op) => visitor.visit_analyze(op, children), Operator::AddColumn(op) => visitor.visit_add_column(op), Operator::ChangeColumn(op) => visitor.visit_change_column(op), Operator::DropColumn(op) => visitor.visit_drop_column(op), Operator::CreateTable(op) => visitor.visit_create_table(op), - Operator::CreateIndex(op) => visitor.visit_create_index(op), + Operator::CreateIndex(op) => visitor.visit_create_index(op, children), Operator::CreateView(op) => visitor.visit_create_view(op), Operator::DropTable(op) => visitor.visit_drop_table(op), Operator::DropView(op) => visitor.visit_drop_view(op), @@ -453,8 +622,8 @@ pub fn walk_mut_operator<'a, V: OperatorVisitorMut<'a>>( #[cfg(feature = "copy")] Operator::CopyFromFile(op) => visitor.visit_copy_from_file(op), #[cfg(feature = "copy")] - Operator::CopyToFile(op) => visitor.visit_copy_to_file(op), - Operator::Window(op) => visitor.visit_window(op), + Operator::CopyToFile(op) => visitor.visit_copy_to_file(op, children), + Operator::Window(op) => visitor.visit_window(op, children), } } @@ -463,15 +632,16 @@ pub fn walk_mut_operator<'a, V: OperatorVisitorMut<'a>>( mod tests { use super::*; use crate::planner::operator::visitor::tests::all_operators; + use crate::planner::operator::RowCountOption; use crate::types::value::DataValue; struct IncrementConstants(usize); - impl ExprVisitorMut for IncrementConstants { + impl ExprVisitorMut for IncrementConstants { fn visit_constant( &mut self, value: &mut DataValue, - _arena: &mut (dyn MetaArena + '_), + _arena: &mut A, ) -> Result<(), DatabaseError> { if let DataValue::Int32(value) = value { *value += 1; @@ -504,9 +674,17 @@ mod tests { } let mut child = LogicalPlan::new(Operator::Dummy, Childrens::None); - child.physical_option = Some(PhysicalOption::new(PlanImpl::Dummy, SortOption::None)); + child.physical_option = Some(PhysicalOption::new( + PlanImpl::Dummy, + SortOption::None, + RowCountOption::Unknown, + )); let mut plan = LogicalPlan::new(Operator::Dummy, Childrens::Only(Box::new(child))); - plan.physical_option = Some(PhysicalOption::new(PlanImpl::Dummy, SortOption::None)); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::Dummy, + SortOption::None, + RowCountOption::Unknown, + )); let mut counter = Counter { operators: 0, @@ -527,17 +705,17 @@ mod tests { let mut arena = PlanArena::new(&table_arena); let mut operators = all_operators(&mut arena)?; for operator in &mut operators { - NoopVisitor.visit_operator(operator, None)?; + NoopVisitor.visit_operator(operator, None, &mut ChildrenVisitMut::empty())?; } let mut counter = IncrementConstants(0); { let mut visitor = OperatorExprVisitorMut::new(&mut counter, &mut arena); for operator in &mut operators { - visitor.visit_operator(operator, None)?; + visitor.visit_operator(operator, None, &mut ChildrenVisitMut::empty())?; } } - assert_eq!(counter.0, 22); // Includes the Values row expression. + assert_eq!(counter.0, 24); // Includes the Values row expression. Ok(()) } diff --git a/src/planner/operator/window.rs b/src/planner/operator/window.rs index 0c200128..ac918253 100644 --- a/src/planner/operator/window.rs +++ b/src/planner/operator/window.rs @@ -42,9 +42,9 @@ impl WindowOperator { } impl Explain for WindowOperator { - fn fmt( + fn fmt( &self, - arena: &(dyn MetaArena + '_), + arena: &A, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { let (partition_by, order_by) = self.sort_fields.split_at(self.partition_by_len); diff --git a/src/planner/plan_keeper.rs b/src/planner/plan_keeper.rs index 294f69e6..fcd1309c 100644 --- a/src/planner/plan_keeper.rs +++ b/src/planner/plan_keeper.rs @@ -20,6 +20,15 @@ pub(crate) enum PlanInput<'p> { Borrowed(&'p LogicalPlan), } +impl PlanInput<'_> { + pub(crate) fn root(&self) -> &LogicalPlan { + match self { + PlanInput::Owned(plan) => plan, + PlanInput::Borrowed(plan) => plan, + } + } +} + impl From for PlanInput<'_> { fn from(plan: LogicalPlan) -> Self { PlanInput::Owned(plan) diff --git a/src/planner/scalar_query_ref.rs b/src/planner/scalar_query_ref.rs index 51e13769..53550df7 100644 --- a/src/planner/scalar_query_ref.rs +++ b/src/planner/scalar_query_ref.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; /// A definition's position in an Arena namespace, independent of execution-local slots. @@ -45,7 +45,7 @@ impl ReferenceSerialization for ScalarQueryRef { self.is_outer_value.encode(writer, direct, tables, arena) } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, tables: &ReferenceTables, diff --git a/src/python.rs b/src/python.rs index 91da11a5..3c54a28c 100644 --- a/src/python.rs +++ b/src/python.rs @@ -21,7 +21,7 @@ use crate::storage::lmdb::LmdbStorage; use crate::storage::memory::MemoryStorage; #[cfg(feature = "rocksdb")] use crate::storage::rocksdb::RocksStorage; -use crate::types::tuple::{SchemaView, Tuple}; +use crate::types::tuple::{SchemaView, TupleImpl}; use crate::types::value::DataValue; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; @@ -34,9 +34,9 @@ fn to_py_err(err: impl ToString) -> PyErr { #[allow(deprecated)] fn data_value_to_py(py: Python<'_>, value: &DataValue) -> PyResult { let object = match value { - DataValue::Parameter { id, .. } => { + DataValue::Parameter { .. } | DataValue::InitValue { .. } => { return Err(PyRuntimeError::new_err(format!( - "unbound parameter ${id} reached Python output" + "unbound parameter {value} reached Python output" ))); } DataValue::Null => py.None(), @@ -70,16 +70,15 @@ fn data_value_to_py(py: Python<'_>, value: &DataValue) -> PyResult { Ok(object) } -fn tuple_to_python_row(py: Python<'_>, tuple: &Tuple) -> PyResult { +fn tuple_to_python_row(py: Python<'_>, tuple: &TupleImpl) -> PyResult { let row = PyDict::new(py); - match tuple.pk.as_ref() { + match tuple.pk() { Some(pk) => row.set_item("pk", data_value_to_py(py, pk)?)?, None => row.set_item("pk", py.None())?, } let values = tuple - .values .iter() .map(|value| data_value_to_py(py, value)) .collect::>>()?; @@ -170,7 +169,7 @@ enum PythonResultIterInner { impl PythonResultIterInner { fn next_tuple( &mut self, - f: impl FnOnce(&mut Tuple) -> R, + f: impl FnOnce(&mut TupleImpl) -> R, ) -> Result, DatabaseError> { match self { #[cfg(feature = "lmdb")] diff --git a/src/serdes/boolean.rs b/src/serdes/boolean.rs index 7aaa6daf..244a81fe 100644 --- a/src/serdes/boolean.rs +++ b/src/serdes/boolean.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for bool { @@ -29,7 +29,7 @@ impl ReferenceSerialization for bool { if *self { 1u8 } else { 0u8 }.encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/bound.rs b/src/serdes/bound.rs index 62a8c6d8..833cf004 100644 --- a/src/serdes/bound.rs +++ b/src/serdes/bound.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; use std::ops::Bound; @@ -49,7 +49,7 @@ where Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/btree_map.rs b/src/serdes/btree_map.rs index 5a51d497..cb66922e 100644 --- a/src/serdes/btree_map.rs +++ b/src/serdes/btree_map.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::collections::BTreeMap; use std::io::{Read, Write}; @@ -40,7 +40,7 @@ where Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/char.rs b/src/serdes/char.rs index a97511d8..8d1fa46f 100644 --- a/src/serdes/char.rs +++ b/src/serdes/char.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for char { @@ -38,7 +38,7 @@ impl ReferenceSerialization for char { Ok(writer.write_all(&buf)?) } - fn decode( + fn decode( reader: &mut R, _: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, _: &ReferenceTables, diff --git a/src/serdes/char_length_units.rs b/src/serdes/char_length_units.rs index b5b7b19a..9aeb46d5 100644 --- a/src/serdes/char_length_units.rs +++ b/src/serdes/char_length_units.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::CharLengthUnits; use std::io::{Read, Write}; @@ -36,7 +36,7 @@ impl ReferenceSerialization for CharLengthUnits { Ok(()) } - fn decode( + fn decode( reader: &mut R, _: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, _: &ReferenceTables, diff --git a/src/serdes/column.rs b/src/serdes/column.rs index 071c5522..d6b2e571 100644 --- a/src/serdes/column.rs +++ b/src/serdes/column.rs @@ -16,7 +16,7 @@ use crate::catalog::{ColumnCatalog, ColumnRef, ColumnRelation}; use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::ColumnId; use std::io::{Read, Write}; @@ -33,7 +33,7 @@ impl ReferenceSerialization for ColumnRef { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, drive: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -94,7 +94,7 @@ impl ReferenceSerialization for ColumnRelation { Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -141,7 +141,7 @@ pub(crate) mod test { use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; use crate::storage::rocksdb::RocksStorage; use crate::storage::rocksdb::RocksTransaction; - use crate::storage::{Storage, Transaction}; + use crate::storage::{Storage, StorageTransaction}; use crate::types::value::DataValue; use crate::types::LogicalType; use std::io::{Cursor, Seek, SeekFrom}; diff --git a/src/serdes/data_value.rs b/src/serdes/data_value.rs index 718897f5..223b1a33 100644 --- a/src/serdes/data_value.rs +++ b/src/serdes/data_value.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::value::DataValue; use crate::types::value::Utf8Type; use crate::types::CharLengthUnits; @@ -65,7 +65,7 @@ impl ReferenceSerialization for Utf8Type { } } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -97,9 +97,9 @@ impl DataValue { writer: &mut W, ) -> Result<(), DatabaseError> { match self { - DataValue::Parameter { .. } => Err(DatabaseError::InvalidValue( - "unbound parameter cannot be serialized".to_string(), - )), + DataValue::Parameter { .. } | DataValue::InitValue { .. } => Err( + DatabaseError::InvalidValue("unbound parameter cannot be serialized".to_string()), + ), DataValue::Null => write_u8(writer, TAG_NULL), DataValue::Boolean(value) => { write_u8(writer, TAG_BOOLEAN)?; @@ -275,7 +275,7 @@ impl ReferenceSerialization for DataValue { self.encode_reference_value(writer) } - fn decode( + fn decode( reader: &mut R, _: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, _: &ReferenceTables, diff --git a/src/serdes/evaluator.rs b/src/serdes/evaluator.rs index 4edd37ff..742dce8d 100644 --- a/src/serdes/evaluator.rs +++ b/src/serdes/evaluator.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::evaluator::{ BinaryEvaluatorParams, BinaryEvaluatorRef, CastEvaluatorParams, CastEvaluatorRef, UnaryEvaluatorRef, @@ -39,7 +39,7 @@ impl ReferenceSerialization for BinaryEvaluatorParams { } } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -99,7 +99,7 @@ impl ReferenceSerialization for CastEvaluatorParams { } } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -166,7 +166,7 @@ impl ReferenceSerialization for UnaryEvaluatorRef { self.pos.encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -195,7 +195,7 @@ impl ReferenceSerialization for BinaryEvaluatorRef { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -222,7 +222,7 @@ impl ReferenceSerialization for CastEvaluatorRef { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/expression.rs b/src/serdes/expression.rs index c4ead6d9..2aeda320 100644 --- a/src/serdes/expression.rs +++ b/src/serdes/expression.rs @@ -16,7 +16,7 @@ use crate::errors::DatabaseError; use crate::expression::ScalarExpression; use crate::planner::{ExprRef, MetaArena}; use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for ExprRef { @@ -32,7 +32,7 @@ impl ReferenceSerialization for ExprRef { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/function.rs b/src/serdes/function.rs index d262a9b6..607e2aab 100644 --- a/src/serdes/function.rs +++ b/src/serdes/function.rs @@ -18,7 +18,7 @@ use crate::expression::function::table::ArcTableFunctionImpl; use crate::expression::function::FunctionSummary; use crate::planner::MetaArena; use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for ArcScalarFunctionImpl { @@ -33,7 +33,7 @@ impl ReferenceSerialization for ArcScalarFunctionImpl { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -69,7 +69,7 @@ impl ReferenceSerialization for ArcTableFunctionImpl { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/hasher.rs b/src/serdes/hasher.rs index 737fbbfa..ea2ef0d2 100644 --- a/src/serdes/hasher.rs +++ b/src/serdes/hasher.rs @@ -16,7 +16,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::stable_hash::StableHasher; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; impl ReferenceSerialization for StableHasher { fn encode( @@ -31,7 +31,7 @@ impl ReferenceSerialization for StableHasher { key1.encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/index.rs b/src/serdes/index.rs index 9aef83a9..dda12f95 100644 --- a/src/serdes/index.rs +++ b/src/serdes/index.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use crate::types::index::{IndexMeta, IndexMetaRef}; use std::io::{Read, Write}; @@ -32,7 +32,7 @@ impl ReferenceSerialization for IndexMetaRef { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, drive: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/mod.rs b/src/serdes/mod.rs index 3f4a4626..e2092b6b 100644 --- a/src/serdes/mod.rs +++ b/src/serdes/mod.rs @@ -29,6 +29,7 @@ mod option; mod pair; mod path_buf; mod phantom; +mod physical_option; mod ptr; mod slice; pub(crate) mod stable_hash; @@ -40,7 +41,7 @@ use crate::catalog::TableName; use crate::db::{ScalaFunctions, TableFunctions}; use crate::errors::DatabaseError; use crate::planner::MetaArena; -use crate::storage::{TableCache, Transaction}; +use crate::storage::{StorageTransaction, TableCache}; use std::io; use std::io::{Read, Write}; @@ -53,7 +54,7 @@ pub trait ReferenceSerialization { arena: &A, ) -> Result<(), DatabaseError>; - fn decode( + fn decode( reader: &mut R, context: Option<&ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -63,13 +64,13 @@ pub trait ReferenceSerialization { Self: Sized; } -pub struct ReferenceDecodeContext<'a, T: Transaction> { +pub struct ReferenceDecodeContext<'a, T: StorageTransaction> { drive: Option<(&'a T, &'a TableCache)>, scala_functions: Option<&'a ScalaFunctions>, table_functions: Option<&'a TableFunctions>, } -impl<'a, T: Transaction> ReferenceDecodeContext<'a, T> { +impl<'a, T: StorageTransaction> ReferenceDecodeContext<'a, T> { pub fn new(drive: Option<(&'a T, &'a TableCache)>) -> Self { Self { drive, diff --git a/src/serdes/num.rs b/src/serdes/num.rs index c2792807..df66d6bb 100644 --- a/src/serdes/num.rs +++ b/src/serdes/num.rs @@ -15,7 +15,8 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; +use ordered_float::OrderedFloat; use std::io::Read; use std::io::Write; use std::mem::size_of; @@ -36,7 +37,7 @@ macro_rules! implement_num_serialization { Ok(()) } - fn decode( + fn decode( reader: &mut R, _: Option<&$crate::serdes::ReferenceDecodeContext<'_, T>>, _: &ReferenceTables, @@ -63,6 +64,27 @@ implement_num_serialization!(u64); implement_num_serialization!(f32); implement_num_serialization!(f64); +impl ReferenceSerialization for OrderedFloat { + fn encode( + &self, + writer: &mut W, + is_direct: bool, + reference_tables: &mut ReferenceTables, + arena: &A, + ) -> Result<(), DatabaseError> { + self.0.encode(writer, is_direct, reference_tables, arena) + } + + fn decode( + reader: &mut R, + drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, + reference_tables: &ReferenceTables, + arena: &mut A, + ) -> Result { + f64::decode(reader, drive, reference_tables, arena).map(OrderedFloat) + } +} + impl ReferenceSerialization for usize { fn encode( &self, @@ -74,7 +96,7 @@ impl ReferenceSerialization for usize { (*self as u32).encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/option.rs b/src/serdes/option.rs index 82e2427f..26770c6c 100644 --- a/src/serdes/option.rs +++ b/src/serdes/option.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for Option @@ -40,7 +40,7 @@ where Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/pair.rs b/src/serdes/pair.rs index 3fd86c02..9b8b2b99 100644 --- a/src/serdes/pair.rs +++ b/src/serdes/pair.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for (A, B) @@ -37,7 +37,7 @@ where Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/path_buf.rs b/src/serdes/path_buf.rs index 072dac6f..eda11243 100644 --- a/src/serdes/path_buf.rs +++ b/src/serdes/path_buf.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::path::PathBuf; use std::str::FromStr; @@ -33,7 +33,7 @@ impl ReferenceSerialization for PathBuf { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/phantom.rs b/src/serdes/phantom.rs index 33e01b8e..980a0acd 100644 --- a/src/serdes/phantom.rs +++ b/src/serdes/phantom.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; use std::marker::PhantomData; @@ -30,7 +30,7 @@ impl ReferenceSerialization for PhantomData { Ok(()) } - fn decode( + fn decode( _: &mut R, _: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, _: &ReferenceTables, diff --git a/src/serdes/physical_option.rs b/src/serdes/physical_option.rs new file mode 100644 index 00000000..fd887ca1 --- /dev/null +++ b/src/serdes/physical_option.rs @@ -0,0 +1,230 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::planner::operator::{PhysicalOption, PlanImpl, RowCountOption, SortOption}; +use crate::planner::MetaArena; +use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; +use crate::storage::StorageTransaction; +use std::io::{Read, Write}; + +impl ReferenceSerialization for RowCountOption { + fn encode( + &self, + writer: &mut W, + is_direct: bool, + reference_tables: &mut ReferenceTables, + arena: &A, + ) -> Result<(), DatabaseError> { + match self { + Self::Follow => writer.write_all(&[0])?, + Self::Unknown => writer.write_all(&[1])?, + Self::Estimated { rows, .. } => { + writer.write_all(&[2])?; + rows.encode(writer, is_direct, reference_tables, arena)?; + } + } + Ok(()) + } + + fn decode( + reader: &mut R, + context: Option<&ReferenceDecodeContext<'_, T>>, + reference_tables: &ReferenceTables, + arena: &mut A, + ) -> Result { + let mut tag = [0]; + reader.read_exact(&mut tag)?; + Ok(match tag[0] { + 0 => Self::Follow, + 1 => Self::Unknown, + 2 => Self::Estimated { + rows: usize::decode(reader, context, reference_tables, arena)?, + ndv: None, + }, + _ => return Err(DatabaseError::InvalidValue("invalid row count tag".into())), + }) + } +} + +impl ReferenceSerialization for PhysicalOption { + fn encode( + &self, + writer: &mut W, + is_direct: bool, + reference_tables: &mut ReferenceTables, + arena: &A, + ) -> Result<(), DatabaseError> { + self.plan + .encode(writer, is_direct, reference_tables, arena)?; + self.sort_option + .encode(writer, is_direct, reference_tables, arena)?; + self.row_count + .encode(writer, is_direct, reference_tables, arena) + } + + fn decode( + reader: &mut R, + context: Option<&ReferenceDecodeContext<'_, T>>, + reference_tables: &ReferenceTables, + arena: &mut A, + ) -> Result { + Ok(Self { + plan: PlanImpl::decode(reader, context, reference_tables, arena)?, + sort_option: SortOption::decode(reader, context, reference_tables, arena)?, + row_count: RowCountOption::decode(reader, context, reference_tables, arena)?, + lookup_work: None, + join_decided: false, + }) + } +} + +// GRCOV_EXCL_START +#[cfg(all( + test, + not(target_arch = "wasm32"), + any(feature = "rocksdb", feature = "lmdb") +))] +mod tests { + use super::*; + use crate::expression::ScalarExpression; + use crate::planner::operator::sort::SortField; + use crate::planner::{PlanArena, TableArenaCell}; + #[cfg(all(not(feature = "rocksdb"), feature = "lmdb"))] + use crate::storage::lmdb::LmdbTransaction as TestTransaction; + #[cfg(feature = "rocksdb")] + use crate::storage::rocksdb::OptimisticRocksTransaction as TestTransaction; + use crate::types::value::DataValue; + + #[test] + fn physical_option_preserves_fields_in_roundtrip() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + for option in [ + PhysicalOption::new( + PlanImpl::SeqScan, + SortOption::None, + RowCountOption::Estimated { + rows: 12, + ndv: None, + }, + ), + PhysicalOption::new( + PlanImpl::Project, + SortOption::Follow, + RowCountOption::Follow, + ), + PhysicalOption::new(PlanImpl::Sort, SortOption::None, RowCountOption::Unknown), + PhysicalOption::new( + PlanImpl::Values, + SortOption::None, + RowCountOption::Estimated { rows: 0, ndv: None }, + ), + ] { + for is_direct in [false, true] { + let mut encoded = Vec::new(); + let mut references = ReferenceTables::new(); + option.encode(&mut encoded, is_direct, &mut references, &arena)?; + let mut reader = encoded.as_slice(); + let decoded = PhysicalOption::decode::, _, _>( + &mut reader, + None, + &references, + &mut arena, + )?; + assert_eq!(decoded, option); + assert!(reader.is_empty()); + } + } + Ok(()) + } + + #[test] + fn physical_option_keeps_wire_order() -> Result<(), DatabaseError> { + let tables = TableArenaCell::default(); + let mut arena = PlanArena::new(&tables); + let key = arena.alloc_expression(ScalarExpression::Constant(DataValue::Int32(1))); + let mut option = PhysicalOption::new( + PlanImpl::Sort, + SortOption::OrderBy { + fields: vec![SortField { + expr: key, + asc: false, + nulls_first: true, + }], + ignore_prefix_len: 1, + }, + RowCountOption::Estimated { + rows: 42, + ndv: None, + }, + ); + option.row_count = RowCountOption::Estimated { + rows: 42, + ndv: Some(37), + }; + for is_direct in [false, true] { + // The wire layout is exactly three ordered fields. + let mut expected = Vec::new(); + let mut reference_tables = ReferenceTables::new(); + option + .plan + .encode(&mut expected, is_direct, &mut reference_tables, &arena)?; + option + .sort_option + .encode(&mut expected, is_direct, &mut reference_tables, &arena)?; + option + .row_count + .encode(&mut expected, is_direct, &mut reference_tables, &arena)?; + let mut encoded = Vec::new(); + let mut references = ReferenceTables::new(); + option.encode(&mut encoded, is_direct, &mut references, &arena)?; + assert_eq!(encoded, expected); + assert_eq!(references, reference_tables); + let mut reader = encoded.as_slice(); + let decoded = PhysicalOption::decode::, _, _>( + &mut reader, + None, + &references, + &mut arena, + )?; + assert_eq!(decoded.plan, PlanImpl::Sort); + assert_eq!( + decoded.row_count, + RowCountOption::Estimated { + rows: 42, + ndv: None + } + ); + let SortOption::OrderBy { + fields, + ignore_prefix_len, + } = decoded.sort_option + else { + unreachable!() + }; + assert_eq!(ignore_prefix_len, 1); + assert_eq!(fields.len(), 1); + assert!(!fields[0].asc); + assert!(fields[0].nulls_first); + assert_eq!( + arena.expression(fields[0].expr), + &ScalarExpression::Constant(DataValue::Int32(1)) + ); + assert!(reader.is_empty()); + } + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/serdes/ptr.rs b/src/serdes/ptr.rs index e5301158..6837f517 100644 --- a/src/serdes/ptr.rs +++ b/src/serdes/ptr.rs @@ -13,7 +13,7 @@ // limitations under the License. use crate::serdes::DatabaseError; -use crate::serdes::Transaction; +use crate::serdes::StorageTransaction; use crate::serdes::{ReferenceSerialization, ReferenceTables}; use std::io::{Read, Write}; use std::sync::Arc; @@ -36,7 +36,7 @@ macro_rules! implement_ptr_serialization { .encode(writer, is_direct, reference_tables, arena) } - fn decode( + fn decode( reader: &mut R, drive: Option<&$crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/slice.rs b/src/serdes/slice.rs index 1d4f567b..488c1bf1 100644 --- a/src/serdes/slice.rs +++ b/src/serdes/slice.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for [V; 2] @@ -35,7 +35,7 @@ where Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/string.rs b/src/serdes/string.rs index 0110b277..ee3cd754 100644 --- a/src/serdes/string.rs +++ b/src/serdes/string.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; use std::sync::Arc; @@ -34,7 +34,7 @@ impl ReferenceSerialization for String { Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, @@ -63,7 +63,7 @@ impl ReferenceSerialization for Arc { Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/serdes/trim.rs b/src/serdes/trim.rs index afbb42e1..f966dcd0 100644 --- a/src/serdes/trim.rs +++ b/src/serdes/trim.rs @@ -16,7 +16,7 @@ use crate::errors::DatabaseError; use crate::expression::TrimWhereField; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for TrimWhereField { @@ -37,7 +37,7 @@ impl ReferenceSerialization for TrimWhereField { Ok(()) } - fn decode( + fn decode( reader: &mut R, _: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, _: &ReferenceTables, diff --git a/src/serdes/vec.rs b/src/serdes/vec.rs index 5cbabd6c..c704dc6c 100644 --- a/src/serdes/vec.rs +++ b/src/serdes/vec.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::planner::MetaArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; -use crate::storage::Transaction; +use crate::storage::StorageTransaction; use std::io::{Read, Write}; impl ReferenceSerialization for Vec @@ -37,7 +37,7 @@ where Ok(()) } - fn decode( + fn decode( reader: &mut R, drive: Option<&crate::serdes::ReferenceDecodeContext<'_, T>>, reference_tables: &ReferenceTables, diff --git a/src/storage/lmdb.rs b/src/storage/lmdb.rs index 94621ca8..935284d9 100644 --- a/src/storage/lmdb.rs +++ b/src/storage/lmdb.rs @@ -16,11 +16,11 @@ use crate::errors::DatabaseError; use crate::storage::table_codec::Bytes; use crate::storage::{ bounds_contain, bytes_bound_as_slice, owned_bound, InnerIter, KeyValueRef, Storage, - Transaction, TransactionIsolationLevel, + StorageTransaction, TransactionIsolationLevel, }; use lmdb::{ - Cursor, Database, DatabaseFlags, Environment, EnvironmentFlags, RoCursor, RwTransaction, - Transaction as _, WriteFlags, + Cursor, Database, DatabaseFlags, Environment, EnvironmentFlags, RoCursor, RoTransaction, + RwTransaction, Transaction as _, WriteFlags, }; use std::collections::Bound; use std::fmt::{self, Display, Formatter}; @@ -99,7 +99,9 @@ impl LmdbStorage { let mut builder = Environment::new(); builder.set_map_size(config.map_size); - builder.set_flags(config.flags); + // Read transactions may outlive the statement call (result iterators) and + // several may be open on one thread, so they must not be bound to TLS slots. + builder.set_flags(config.flags | EnvironmentFlags::NO_TLS); if let Some(max_readers) = config.max_readers { builder.set_max_readers(max_readers); } @@ -128,10 +130,14 @@ impl Storage for LmdbStorage { fn transaction_with_isolation( &self, isolation: TransactionIsolationLevel, + read_only: bool, ) -> Result, DatabaseError> { self.validate_transaction_isolation(isolation)?; - let tx = self.env.begin_rw_txn()?; - + let tx = if read_only { + LmdbTxn::Ro(self.env.begin_ro_txn()?) + } else { + LmdbTxn::Rw(self.env.begin_rw_txn()?) + }; Ok(LmdbTransaction { tx, db: self.db, @@ -162,11 +168,49 @@ impl Storage for LmdbStorage { } pub struct LmdbTransaction<'env> { - tx: RwTransaction<'env>, + tx: LmdbTxn<'env>, db: Database, statements: u64, } +/// Read transactions do not take LMDB's single-writer lock. +enum LmdbTxn<'env> { + Ro(RoTransaction<'env>), + Rw(RwTransaction<'env>), +} + +impl<'env> LmdbTxn<'env> { + fn raw(&self) -> *mut lmdb_sys::MDB_txn { + match self { + LmdbTxn::Ro(tx) => tx.txn(), + LmdbTxn::Rw(tx) => tx.txn(), + } + } + + fn get(&self, db: Database, key: &[u8]) -> Result<&[u8], lmdb::Error> { + match self { + LmdbTxn::Ro(tx) => tx.get(db, &key), + LmdbTxn::Rw(tx) => tx.get(db, &key), + } + } + + fn open_ro_cursor(&self, db: Database) -> Result, lmdb::Error> { + match self { + LmdbTxn::Ro(tx) => tx.open_ro_cursor(db), + LmdbTxn::Rw(tx) => tx.open_ro_cursor(db), + } + } + + fn writable(&mut self) -> Result<&mut RwTransaction<'env>, DatabaseError> { + match self { + LmdbTxn::Ro(_) => Err(DatabaseError::UnsupportedStmt( + "cannot write in a read-only transaction".into(), + )), + LmdbTxn::Rw(tx) => Ok(tx), + } + } +} + pub struct LmdbIter<'txn> { cursor: RoCursor<'txn>, scan: Scan, @@ -276,7 +320,7 @@ fn cursor_get<'txn, C: Cursor<'txn>>( } } -impl Transaction for LmdbTransaction<'_> { +impl StorageTransaction for LmdbTransaction<'_> { type BorrowedBytes<'a> = &'a [u8] where @@ -300,7 +344,7 @@ impl Transaction for LmdbTransaction<'_> { const STATEMENT_BITS: u32 = 20; self.statements += 1; // SAFETY: `self.tx` is a live transaction handle owned by this value. - let txn_id = unsafe { lmdb_sys::mdb_txn_id(self.tx.txn()) } as u64; + let txn_id = unsafe { lmdb_sys::mdb_txn_id(self.tx.raw()) } as u64; if self.statements >= 1 << STATEMENT_BITS || txn_id >= 1 << (u64::BITS - STATEMENT_BITS) { return Err(DatabaseError::InvalidValue( "statement stamps exhausted".into(), @@ -313,7 +357,7 @@ impl Transaction for LmdbTransaction<'_> { &'a self, key: &[u8], ) -> Result>, DatabaseError> { - match self.tx.get(self.db, &key) { + match self.tx.get(self.db, key) { Ok(value) => Ok(Some(value)), Err(lmdb::Error::NotFound) => Ok(None), Err(err) => Err(err.into()), @@ -338,12 +382,13 @@ impl Transaction for LmdbTransaction<'_> { fn set(&mut self, key: &[u8], value: &[u8]) -> Result<(), DatabaseError> { self.tx + .writable()? .put(self.db, &key, &value, lmdb::WriteFlags::empty())?; Ok(()) } fn remove(&mut self, key: &[u8]) -> Result<(), DatabaseError> { - match self.tx.del(self.db, &key, None) { + match self.tx.writable()?.del(self.db, &key, None) { Ok(()) | Err(lmdb::Error::NotFound) => Ok(()), Err(err) => Err(err.into()), } @@ -372,7 +417,7 @@ impl Transaction for LmdbTransaction<'_> { } fn remove_range(&mut self, min: Bound<&[u8]>, max: Bound<&[u8]>) -> Result<(), DatabaseError> { - let mut cursor = self.tx.open_rw_cursor(self.db)?; + let mut cursor = self.tx.writable()?.open_rw_cursor(self.db)?; let mut scan = Scan::new(min, max, false); while scan.next(&cursor)?.is_some() { cursor.del(WriteFlags::empty())?; @@ -381,7 +426,10 @@ impl Transaction for LmdbTransaction<'_> { } fn commit(self) -> Result<(), DatabaseError> { - self.tx.commit()?; + match self.tx { + LmdbTxn::Ro(tx) => tx.commit()?, + LmdbTxn::Rw(tx) => tx.commit()?, + } Ok(()) } } @@ -414,7 +462,7 @@ mod tests { let mut iter = kite_sql.run("select b from t1 where a = 2").unwrap(); let tuple = iter.next_tuple(|_, tuple| tuple.clone()).unwrap().unwrap(); - assert_eq!(tuple.values[0].to_string(), "20"); + assert_eq!(tuple[0].to_string(), "20"); iter.done().unwrap(); } @@ -443,6 +491,60 @@ mod tests { crate::storage::check_explicit_transaction_does_not_rescan_own_writes(db).unwrap(); } + #[test] + fn read_transactions_skip_the_writer_lock() { + use crate::storage::{Storage, StorageTransaction}; + let temp_dir = TempDir::new().expect("unable to create temporary working directory"); + let storage = LmdbStorage::new(temp_dir.path().join("kite_sql.lmdb")).unwrap(); + let isolation = storage.default_transaction_isolation(); + + // A writer holds LMDB's single write lock; readers still open on the same thread. + let mut writer = storage.transaction().unwrap(); + writer.set(b"k", b"v").unwrap(); + let mut reader = storage.transaction_with_isolation(isolation, true).unwrap(); + let second = storage.transaction_with_isolation(isolation, true).unwrap(); + assert!(reader.get_borrowed(b"k").unwrap().is_none()); + assert!(second.get_borrowed(b"k").unwrap().is_none()); + assert!(reader.set(b"k", b"x").is_err()); + assert!(reader.remove(b"k").is_err()); + writer.commit().unwrap(); + drop((reader, second)); + + let reader = storage.transaction_with_isolation(isolation, true).unwrap(); + assert_eq!(reader.get_borrowed(b"k").unwrap(), Some(&b"v"[..])); + } + + #[test] + fn read_statements_and_plans_use_read_transactions() { + let temp_dir = TempDir::new().expect("unable to create temporary working directory"); + let mut db = DataBaseBuilder::path(temp_dir.path()).build_lmdb().unwrap(); + db.ddl("create table t (a int primary key)").unwrap(); + db.run("insert into t values (1)").unwrap().done().unwrap(); + + // An open write transaction does not block reads. + let mut writer = db.new_transaction(false).unwrap(); + writer + .run("insert into t values (2)") + .unwrap() + .done() + .unwrap(); + let select = db.prepare_sql("select count(*) from t", &[]).unwrap(); + let insert = db.prepare_sql("insert into t values (3)", &[]).unwrap(); + assert!(select.read_only && !insert.read_only); + let mut first = db.run("select a from t").unwrap(); + let mut second = db.execute(&select, []).unwrap(); + let count = second.next_tuple(|_, t| t[0].to_string()).unwrap(); + assert_eq!(count.as_deref(), Some("1")); + assert!(first.next_tuple(|_, _| ()).unwrap().is_some()); + // An explicit read-only transaction reads beside the writer and rejects writes. + let mut reader = db.new_transaction(true).unwrap(); + assert!(reader.run("select a from t").unwrap().done().is_ok()); + assert!(reader.run("insert into t values (4)").is_err()); + assert!(reader.execute(&insert, []).is_err()); + drop((first, second, reader)); + writer.commit().unwrap(); + } + #[test] fn build_with_lmdb_storage() { let temp_dir = TempDir::new().expect("unable to create temporary working directory"); diff --git a/src/storage/memory.rs b/src/storage/memory.rs index bdf5e794..fde58037 100644 --- a/src/storage/memory.rs +++ b/src/storage/memory.rs @@ -15,7 +15,7 @@ use crate::errors::DatabaseError; use crate::storage::table_codec::Bytes; use crate::storage::{ - bytes_bound_as_slice, owned_bound, EmptyStorageMetrics, InnerIter, Storage, Transaction, + bytes_bound_as_slice, owned_bound, EmptyStorageMetrics, InnerIter, Storage, StorageTransaction, TransactionIsolationLevel, }; use std::cell::{Cell, Ref, RefCell}; @@ -45,6 +45,7 @@ impl Storage for MemoryStorage { fn transaction_with_isolation( &self, isolation: TransactionIsolationLevel, + _read_only: bool, ) -> Result, DatabaseError> { self.validate_transaction_isolation(isolation)?; Ok(MemoryTransaction { @@ -147,7 +148,7 @@ impl InnerIter for MemoryIter<'_> { } } -impl Transaction for MemoryTransaction { +impl StorageTransaction for MemoryTransaction { type BorrowedBytes<'a> = MemoryValue<'a> where @@ -275,7 +276,8 @@ mod wasm_tests { &Tuple::new( Some(DataValue::Int32(1)), vec![DataValue::Int32(1), DataValue::Boolean(true)], - ), + ) + .into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), @@ -288,7 +290,8 @@ mod wasm_tests { &Tuple::new( Some(DataValue::Int32(2)), vec![DataValue::Int32(2), DataValue::Boolean(true)], - ), + ) + .into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), @@ -440,7 +443,8 @@ mod native_tests { &Tuple::new( Some(DataValue::Int32(1)), vec![DataValue::Int32(1), DataValue::Boolean(true)], - ), + ) + .into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), @@ -453,7 +457,8 @@ mod native_tests { &Tuple::new( Some(DataValue::Int32(2)), vec![DataValue::Int32(2), DataValue::Boolean(true)], - ), + ) + .into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), diff --git a/src/storage/mod.rs b/src/storage/mod.rs index bcc537cc..92dfe441 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -29,14 +29,15 @@ use crate::iter_ext::Itertools; use crate::optimizer::core::cm_sketch::{ CountMinSketch, CountMinSketchPage, COUNT_MIN_SKETCH_STORAGE_PAGE_LEN, }; -use crate::optimizer::core::statistics_meta::StatisticsMeta; +use crate::optimizer::core::histogram::HistogramMeta; +use crate::optimizer::core::statistics_meta::IndexStatistics; use crate::planner::operator::alter_table::change_column::{DefaultChange, NotNullChange}; use crate::planner::{MetaArena, TableArenaCell}; use crate::serdes::ReferenceTables; -use crate::storage::table_codec::{Bytes, StatisticsCodecType, TableCodec, BOUND_MAX_TAG}; -use crate::types::index::{Index, IndexId, IndexMeta, IndexMetaRef, IndexType}; +use crate::storage::table_codec::{Bytes, IndexStatisticsCodecType, TableCodec, BOUND_MAX_TAG}; +use crate::types::index::{Index, IndexId, IndexMeta, IndexMetaRef, IndexType, RuntimeIndexProbe}; use crate::types::serialize::TupleValueSerializableImpl; -use crate::types::tuple::{Tuple, TupleId}; +use crate::types::tuple::{Tuple, TupleId, TupleImpl, TupleLike}; use crate::types::value::{DataValue, OrderedMapping, TupleMappingRef}; use crate::types::{ColumnId, LogicalType}; use std::borrow::{Borrow, Cow}; @@ -49,7 +50,18 @@ use std::path::Path; pub type KeyValueRef<'a> = (&'a [u8], &'a [u8]); -pub(crate) type StatisticsMetaCache = HashMap<(TableName, IndexId), StatisticsMeta>; +#[derive(Default)] +pub(crate) struct StatisticsMetaCache { + pub(crate) indexes: HashMap<(TableName, IndexId), IndexStatistics>, + pub(crate) columns: HashMap<(TableName, ColumnId), HistogramMeta>, +} + +impl StatisticsMetaCache { + pub(crate) fn remove_table(&mut self, table_name: &TableName) { + self.indexes.retain(|(name, _), _| name != table_name); + self.columns.retain(|(name, _), _| name != table_name); + } +} pub(crate) type TableCache = HashMap; pub(crate) type ViewCache = HashMap; @@ -109,17 +121,18 @@ impl Display for EmptyStorageMetrics { pub trait Storage: Clone { type Metrics: Display; - type TransactionType<'a>: Transaction + type TransactionType<'a>: StorageTransaction where Self: 'a; fn transaction(&self) -> Result, DatabaseError> { - self.transaction_with_isolation(self.default_transaction_isolation()) + self.transaction_with_isolation(self.default_transaction_isolation(), false) } fn transaction_with_isolation( &self, isolation: TransactionIsolationLevel, + read_only: bool, ) -> Result, DatabaseError>; fn default_transaction_isolation(&self) -> TransactionIsolationLevel { @@ -153,7 +166,7 @@ pub trait CheckpointableStorage: Storage { /// Optional bounds of the reader, of the form (offset, limit). pub(crate) type Bounds = (Option, Option); -pub trait Transaction: Sized { +pub trait StorageTransaction: Sized { type BorrowedBytes<'a>: AsRef<[u8]> where Self: 'a; @@ -251,7 +264,14 @@ pub trait Transaction: Sized { LogicalType::Tuple(tys) => tys.len(), _ => 1, }; - let cover_mapping = mapping.map(|slots| TupleMapping::new(slots, tuple_len)); + let cover_mapping = match mapping { + Some(slots) => Some(TupleMapping::new(slots.iter().copied(), tuple_len)), + None if columns.len() < tuple_len => { + // Omitted mapping denotes an ordered prefix, not every key field. + Some(TupleMapping::new(0..columns.len(), tuple_len)) + } + None => None, + }; ( IndexImplEnum::Covered(CoveredIndexImpl), @@ -361,11 +381,11 @@ pub trait Transaction: Sized { Ok((table, index_id)) } - fn add_index( + fn add_index( &mut self, table_codec: &mut TableCodec, table_name: &str, - index: Index, + index: Index<'_, V>, tuple_id: &TupleId, ) -> Result<(), DatabaseError> { if matches!(index.ty, IndexType::PrimaryKey { .. }) { @@ -387,11 +407,11 @@ pub trait Transaction: Sized { }) } - fn del_index( + fn del_index( &mut self, table_codec: &mut TableCodec, table_name: &str, - index: &Index, + index: &Index<'_, V>, tuple_id: &TupleId, ) -> Result<(), DatabaseError> { if matches!(index.ty, IndexType::PrimaryKey { .. }) { @@ -406,7 +426,7 @@ pub trait Transaction: Sized { &mut self, table_codec: &mut TableCodec, table_name: &str, - tuple: &Tuple, + tuple: &TupleImpl, serializers: I, is_overwrite: bool, ) -> Result<(), DatabaseError> @@ -414,14 +434,13 @@ pub trait Transaction: Sized { I: IntoIterator, S: Borrow, { - let tuple_id = tuple.pk.as_ref().ok_or(DatabaseError::PrimaryKeyNotFound)?; + let tuple_id = tuple.pk().ok_or(DatabaseError::PrimaryKeyNotFound)?; let mut serializers = serializers.into_iter(); - let mut write_value = - |tuple: &Tuple, value: &mut Bytes| tuple.serialize_to(&mut serializers, value); + let mut write_value = |value: &mut Bytes| tuple.serialize_to(&mut serializers, value); table_codec.with_tuple( table_name, tuple_id, - Some((tuple, &mut write_value)), + Some(&mut write_value), |key, value| { if !is_overwrite && self.exists(key)? { return Err(DatabaseError::DuplicatePrimaryKey); @@ -632,6 +651,7 @@ pub trait Transaction: Sized { column_id }; + self.remove_column_statistics(table_codec, table_name, column_id, plan_arena)?; for index_meta in table_catalog.indexes.iter() { let index_meta = plan_arena.index(*index_meta); if !index_meta.column_ids.contains(&column_id) { @@ -649,7 +669,7 @@ pub trait Transaction: Sized { self.remove_range(Bound::Included(min), Bound::Included(max)) })?; - self.remove_statistics_meta(table_codec, table_name, index_meta.id)?; + self.remove_index_statistics(table_codec, table_name, index_meta.id)?; } self.load_table(table_codec, plan_arena, table_name.clone())? .ok_or(DatabaseError::TableNotFound) @@ -800,7 +820,7 @@ pub trait Transaction: Sized { self.remove_range(Bound::Included(min), Bound::Included(max)) })?; - self.remove_statistics_meta(table_codec, &table_name, index_id)?; + self.remove_index_statistics(table_codec, &table_name, index_id)?; self.load_table(table_codec, plan_arena, table_name.clone())? .map(|table| (table, index_id)) @@ -901,19 +921,10 @@ pub trait Transaction: Sized { fn views<'a>( &'a self, table_codec: &mut TableCodec, - table_cache: &'a TableCache, - table_arena: &'a TableArenaCell, - scala_functions: &'a ScalaFunctions, - table_functions: &'a TableFunctions, ) -> Result, DatabaseError> { table_codec.with_view_bound(|min, max| { Ok(ViewIter { iter: self.range(Bound::Included(min), Bound::Included(max))?, - transaction: self, - table_cache, - table_arena, - scala_functions, - table_functions, }) }) } @@ -950,16 +961,61 @@ pub trait Transaction: Sized { }) } - fn save_statistics_meta( + fn save_column_statistics( + &mut self, + table_codec: &mut TableCodec, + table_name: &TableName, + column_id: ColumnId, + meta: &HistogramMeta, + arena: &(impl MetaArena + ?Sized), + ) -> Result<(), DatabaseError> { + table_codec.with_column_statistics( + table_name, + column_id, + Some(meta), + arena, + |key, value| self.set(key, value), + ) + } + + fn column_statistics( + &self, + table_codec: &mut TableCodec, + table_name: &str, + column_id: ColumnId, + arena: &mut (impl MetaArena + ?Sized), + ) -> Result, DatabaseError> { + let value = + table_codec.with_column_statistics(table_name, column_id, None, arena, |key, _| { + self.get_borrowed(key) + })?; + value + .map(|bytes| TableCodec::decode_column_statistics::(bytes.as_ref(), arena)) + .transpose() + } + + fn remove_column_statistics( + &mut self, + table_codec: &mut TableCodec, + table_name: &TableName, + column_id: ColumnId, + arena: &(impl MetaArena + ?Sized), + ) -> Result<(), DatabaseError> { + table_codec.with_column_statistics(table_name, column_id, None, arena, |key, _| { + self.remove(key) + }) + } + + fn save_index_statistics( &mut self, table_codec: &mut TableCodec, table_name: &TableName, - statistics_meta: StatisticsMeta, + index_statistics: IndexStatistics, arena: &(impl MetaArena + ?Sized), ) -> Result<(), DatabaseError> { - let index_id = statistics_meta.index_id(); - let (root, buckets, cm_sketch, top_n) = statistics_meta.into_parts(); - table_codec.with_statistics_meta( + let index_id = index_statistics.index_id(); + let (root, buckets, cm_sketch, top_n) = index_statistics.into_parts(); + table_codec.with_index_statistics( table_name.as_ref(), index_id, Some(&root), @@ -968,7 +1024,7 @@ pub trait Transaction: Sized { )?; let (sketch_meta, sketch_pages) = cm_sketch.into_storage_parts(COUNT_MIN_SKETCH_STORAGE_PAGE_LEN); - table_codec.with_statistics_sketch_meta( + table_codec.with_index_statistics_sketch_meta( table_name.as_ref(), index_id, Some(&sketch_meta), @@ -976,7 +1032,7 @@ pub trait Transaction: Sized { |key, value| self.set(key, value), )?; for sketch_page in sketch_pages { - table_codec.with_statistics_sketch_page( + table_codec.with_index_statistics_sketch_page( table_name.as_ref(), index_id, &sketch_page, @@ -987,7 +1043,7 @@ pub trait Transaction: Sized { } for (ordinal, bucket) in buckets.iter().enumerate() { - table_codec.with_statistics_bucket( + table_codec.with_index_statistics_bucket( table_name.as_ref(), index_id, ordinal as u32, @@ -996,7 +1052,7 @@ pub trait Transaction: Sized { |key, value| self.set(key, value), )?; } - table_codec.with_statistics_top_n( + table_codec.with_index_statistics_top_n( table_name.as_ref(), index_id, Some(&top_n), @@ -1007,14 +1063,14 @@ pub trait Transaction: Sized { Ok(()) } - fn statistics_meta( + fn index_statistics( &self, table_codec: &mut TableCodec, table_name: &str, index_id: IndexId, arena: &mut (impl MetaArena + ?Sized), - ) -> Result, DatabaseError> { - table_codec.with_statistics_index_bound(table_name, index_id, |min, max| { + ) -> Result, DatabaseError> { + table_codec.with_index_statistics_bound(table_name, index_id, |min, max| { let mut iter = self.range(Bound::Included(min), Bound::Included(max))?; let mut root = None; let mut buckets = Vec::new(); @@ -1023,25 +1079,29 @@ pub trait Transaction: Sized { let mut top_n = None; while let Some((key, value)) = iter.try_next()? { - match TableCodec::decode_statistics_codec_type(key)? { - StatisticsCodecType::Root => { - root = Some(TableCodec::decode_statistics_meta::(value, arena)?); + match TableCodec::decode_index_statistics_codec_type(key)? { + IndexStatisticsCodecType::Root => { + root = Some(TableCodec::decode_index_statistics::(value, arena)?); } - StatisticsCodecType::SketchMeta => { - sketch_meta = Some(TableCodec::decode_statistics_sketch_meta::( + IndexStatisticsCodecType::SketchMeta => { + sketch_meta = Some( + TableCodec::decode_index_statistics_sketch_meta::(value, arena)?, + ); + } + IndexStatisticsCodecType::SketchPage => { + sketch_pages.push(TableCodec::decode_index_statistics_sketch_page::( value, arena, )?); } - StatisticsCodecType::SketchPage => { - sketch_pages.push(TableCodec::decode_statistics_sketch_page::( + IndexStatisticsCodecType::Bucket => { + buckets.push(TableCodec::decode_index_statistics_bucket::( value, arena, )?); } - StatisticsCodecType::Bucket => { - buckets.push(TableCodec::decode_statistics_bucket::(value, arena)?); - } - StatisticsCodecType::TopN => { - top_n = Some(TableCodec::decode_statistics_top_n::(value, arena)?); + IndexStatisticsCodecType::TopN => { + top_n = Some(TableCodec::decode_index_statistics_top_n::( + value, arena, + )?); } } } @@ -1049,7 +1109,7 @@ pub trait Transaction: Sized { match (root, sketch_meta) { (Some(root), Some(sketch_meta)) => { let sketch = CountMinSketch::from_storage_parts(sketch_meta, sketch_pages)?; - StatisticsMeta::from_parts(root, buckets, sketch, top_n).map(Some) + IndexStatistics::from_parts(root, buckets, sketch, top_n).map(Some) } (None, None) if buckets.is_empty() && sketch_pages.is_empty() && top_n.is_none() => @@ -1063,13 +1123,13 @@ pub trait Transaction: Sized { }) } - fn remove_statistics_meta( + fn remove_index_statistics( &mut self, table_codec: &mut TableCodec, table_name: &TableName, index_id: IndexId, ) -> Result<(), DatabaseError> { - table_codec.with_statistics_index_bound(table_name.as_ref(), index_id, |min, max| { + table_codec.with_index_statistics_bound(table_name.as_ref(), index_id, |min, max| { self.remove_range(Bound::Included(min), Bound::Included(max)) })?; @@ -1300,7 +1360,7 @@ fn encode_bound<'a>( bound: &Bound, is_upper: bool, buffer: &'a mut Bytes, - params: &IndexImplParams<'_, impl Transaction>, + params: &IndexImplParams<'_, impl StorageTransaction>, inner: &IndexImplEnum, ) -> Result, DatabaseError> { match bound { @@ -1316,7 +1376,7 @@ fn encode_bound<'a>( } } -trait IndexImpl { +trait IndexImpl { fn index_lookup_into( &self, table_codec: &mut TableCodec, @@ -1377,15 +1437,17 @@ struct TupleMapping { } impl TupleMapping { - fn new(scan_to_index: &[usize], tuple_len: usize) -> Self { + fn new(scan_to_index: impl Iterator, tuple_len: usize) -> Self { let mut index_to_scan = vec![usize::MAX; tuple_len]; + let mut target_len = 0; - for (scan_idx, index_idx) in scan_to_index.iter().enumerate() { - index_to_scan[*index_idx] = scan_idx; + for index_idx in scan_to_index { + index_to_scan[index_idx] = target_len; + target_len += 1; } TupleMapping { index_to_scan, - target_len: scan_to_index.len(), + target_len, } } @@ -1394,7 +1456,7 @@ impl TupleMapping { } } -struct IndexImplParams<'a, T: Transaction> { +struct IndexImplParams<'a, T: StorageTransaction> { index_meta: IndexMetaRef, meta_arena: &'a dyn MetaArena, table_name: TableName, @@ -1407,7 +1469,7 @@ struct IndexImplParams<'a, T: Transaction> { reverse: bool, } -impl IndexImplParams<'_, T> { +impl IndexImplParams<'_, T> { #[inline] pub(crate) fn index_meta(&self) -> &IndexMeta { self.meta_arena.index(self.index_meta) @@ -1441,18 +1503,18 @@ impl IndexImplParams<'_, T> { } } -enum IndexResult<'a, T: Transaction + 'a> { +enum IndexResult<'a, T: StorageTransaction + 'a> { Hit, Miss, Scope(RangeIter<'a, T>), } -pub enum RangeIter<'a, T: Transaction + 'a> { +pub enum RangeIter<'a, T: StorageTransaction + 'a> { Forward(T::IterType<'a>), Reverse(T::RevIterType<'a>), } -impl<'a, T: Transaction + 'a> RangeIter<'a, T> { +impl<'a, T: StorageTransaction + 'a> RangeIter<'a, T> { fn open( tx: &'a T, min: Bound<&[u8]>, @@ -1467,7 +1529,7 @@ impl<'a, T: Transaction + 'a> RangeIter<'a, T> { } } -impl<'a, T: Transaction + 'a> InnerIter for RangeIter<'a, T> { +impl<'a, T: StorageTransaction + 'a> InnerIter for RangeIter<'a, T> { fn try_next(&mut self) -> Result>, DatabaseError> { match self { RangeIter::Forward(iter) => iter.try_next(), @@ -1476,7 +1538,7 @@ impl<'a, T: Transaction + 'a> InnerIter for RangeIter<'a, T> { } } -impl IndexImpl for IndexImplEnum { +impl IndexImpl for IndexImplEnum { fn index_lookup_into( &self, table_codec: &mut TableCodec, @@ -1560,7 +1622,7 @@ impl IndexImpl for IndexImplEnum { } } -impl IndexImpl for PrimaryKeyIndexImpl { +impl IndexImpl for PrimaryKeyIndexImpl { fn index_lookup_into( &self, _: &mut TableCodec, @@ -1569,11 +1631,14 @@ impl IndexImpl for PrimaryKeyIndexImpl { value: &[u8], params: &IndexImplParams, ) -> Result<(), DatabaseError> { - let tuple_id = TableCodec::decode_tuple_key(key, ¶ms.index_meta().pk_ty)?; + let tuple_id = params + .with_pk + .then(|| TableCodec::decode_tuple_key(key, ¶ms.index_meta().pk_ty)) + .transpose()?; TableCodec::decode_tuple_into( tuple, ¶ms.deserializers[..], - Some(tuple_id), + tuple_id, value, params.total_len, ) @@ -1604,7 +1669,7 @@ impl IndexImpl for PrimaryKeyIndexImpl { TableCodec::decode_tuple_into( tuple, ¶ms.deserializers[..], - Some(tuple_id.clone()), + params.with_pk.then(|| tuple_id.clone()), bytes.as_ref(), params.total_len, )?; @@ -1634,7 +1699,7 @@ impl IndexImpl for PrimaryKeyIndexImpl { } #[inline(always)] -fn secondary_index_lookup( +fn secondary_index_lookup( table_codec: &mut TableCodec, tuple: &mut Tuple, bytes: &[u8], @@ -1648,7 +1713,7 @@ fn secondary_index_lookup( } } -impl IndexImpl for UniqueIndexImpl { +impl IndexImpl for UniqueIndexImpl { fn index_lookup_into( &self, table_codec: &mut TableCodec, @@ -1694,7 +1759,7 @@ impl IndexImpl for UniqueIndexImpl { table_codec: &mut TableCodec, params: &IndexImplParams, value: &DataValue, - _: bool, + is_upper: bool, out: &mut Bytes, ) -> Result<(), DatabaseError> { let index = Index::new( @@ -1704,14 +1769,19 @@ impl IndexImpl for UniqueIndexImpl { ); table_codec.with_index(params.table_name.as_ref(), &index, None, |key, _| { - out.clear(); - out.extend_from_slice(key); + if index.values.len() < params.index_meta().column_ids.len() { + // A partial UNIQUE key bounds a prefix, not one exact index entry. + encode_bound_key(out, key, is_upper); + } else { + out.clear(); + out.extend_from_slice(key); + } Ok(()) }) } } -impl IndexImpl for NormalIndexImpl { +impl IndexImpl for NormalIndexImpl { fn index_lookup_into( &self, table_codec: &mut TableCodec, @@ -1763,7 +1833,7 @@ impl IndexImpl for NormalIndexImpl { } } -impl IndexImpl for CompositeIndexImpl { +impl IndexImpl for CompositeIndexImpl { fn index_lookup_into( &self, table_codec: &mut TableCodec, @@ -1815,7 +1885,7 @@ impl IndexImpl for CompositeIndexImpl { } } -impl IndexImpl for CoveredIndexImpl { +impl IndexImpl for CoveredIndexImpl { fn index_lookup_into( &self, _: &mut TableCodec, @@ -1824,12 +1894,14 @@ impl IndexImpl for CoveredIndexImpl { value: &[u8], params: &IndexImplParams, ) -> Result<(), DatabaseError> { - let key = match ¶ms.cover_mapping { + match ¶ms.cover_mapping { Some(mapping) => { - TableCodec::decode_index_key(key, params.value_ty(), &mapping.as_ref())? + TableCodec::decode_index_key_into(tuple, key, params.value_ty(), &mapping.as_ref())? } - None => TableCodec::decode_index_key(key, params.value_ty(), &OrderedMapping)?, - }; + None => { + TableCodec::decode_index_key_into(tuple, key, params.value_ty(), &OrderedMapping)? + } + } let tuple_id = if params.with_pk { Some(TableCodec::decode_index(value)?) @@ -1837,7 +1909,6 @@ impl IndexImpl for CoveredIndexImpl { None }; tuple.pk = tuple_id; - tuple.values = key; Ok(()) } @@ -1882,7 +1953,7 @@ impl IndexImpl for CoveredIndexImpl { } #[inline(always)] -fn eq_to_res_scope<'a, T: Transaction + 'a>( +fn eq_to_res_scope<'a, T: StorageTransaction + 'a>( table_codec: &mut TableCodec, tuple: &mut Tuple, index_impl: &impl IndexImpl, @@ -1904,7 +1975,7 @@ fn eq_to_res_scope<'a, T: Transaction + 'a>( Ok(IndexResult::Scope(iter)) } -pub struct TupleIter<'a, T: Transaction + 'a> { +pub struct TupleIter<'a, T: StorageTransaction + 'a> { bounds: IterBounds, pk_ty: Option, deserializers: Vec, @@ -1912,7 +1983,23 @@ pub struct TupleIter<'a, T: Transaction + 'a> { iter: T::IterType<'a>, } -impl<'a, T: Transaction + 'a> Iter for TupleIter<'a, T> { +impl<'a, T: StorageTransaction + 'a> TupleIter<'a, T> { + pub(crate) fn reset( + &mut self, + tx: &'a T, + table_codec: &mut TableCodec, + table_name: &TableName, + bounds: Bounds, + ) -> Result<(), DatabaseError> { + self.iter = table_codec.with_tuple_bound(table_name, |min, max| { + tx.range(Bound::Included(min), Bound::Included(max)) + })?; + self.bounds = IterBounds::new(bounds.0.unwrap_or(0), bounds.1); + Ok(()) + } +} + +impl<'a, T: StorageTransaction + 'a> Iter for TupleIter<'a, T> { fn next_tuple_into( &mut self, table_codec: &mut TableCodec, @@ -2049,7 +2136,7 @@ impl IterBounds { } } -pub struct IndexIter<'a, T: Transaction> { +pub struct IndexIter<'a, T: StorageTransaction> { bounds: IterBounds, params: IndexImplParams<'a, T>, inner: IndexImplEnum, @@ -2059,13 +2146,35 @@ pub struct IndexIter<'a, T: Transaction> { encode_max_buffer: Bytes, } -pub enum IndexIterState<'a, T: Transaction + 'a> { +pub enum IndexIterState<'a, T: StorageTransaction + 'a> { Init, Range(RangeIter<'a, T>), Over, } -impl Iter for IndexIter<'_, T> { +impl<'a, T: StorageTransaction> IndexIter<'a, T> { + pub(crate) fn reset(&mut self, ranges: IndexRanges<'a>, bounds: Bounds) { + self.params.reverse = ranges.reverse; + self.ranges = ranges; + self.rewind(bounds); + } + + pub(crate) fn reset_probe(&mut self, probe: RuntimeIndexProbe, bounds: Bounds) { + match (&mut self.ranges.source, probe) { + (RangeSource::Owned(Range::Eq(value)), RuntimeIndexProbe::Eq(probe)) => *value = probe, + (source, probe) => *source = RangeSource::Owned(probe.into()), + } + self.rewind(bounds); + } + + pub(crate) fn rewind(&mut self, bounds: Bounds) { + self.state = IndexIterState::Init; + self.ranges.next = 0; + self.bounds = IterBounds::new(bounds.0.unwrap_or(0), bounds.1); + } +} + +impl Iter for IndexIter<'_, T> { fn next_tuple_into( &mut self, table_codec: &mut TableCodec, @@ -2147,7 +2256,7 @@ impl Iter for IndexIter<'_, T> { self.bounds.consume_limit(); return Ok(true); } - IndexResult::Miss => return Ok(false), + IndexResult::Miss => continue, IndexResult::Scope(iter) => { self.state = IndexIterState::Range(iter); } @@ -2192,11 +2301,11 @@ pub trait Iter { ) -> Result; } -pub struct TableIter<'a, T: Transaction + 'a> { +pub struct TableIter<'a, T: StorageTransaction + 'a> { iter: T::IterType<'a>, } -impl TableIter<'_, T> { +impl TableIter<'_, T> { pub fn try_next( &mut self, arena: &mut (impl MetaArena + ?Sized), @@ -2209,34 +2318,26 @@ impl TableIter<'_, T> { } } -pub struct ViewIter<'a, T: Transaction + 'a> { +pub struct ViewIter<'a, T: StorageTransaction + 'a> { iter: T::IterType<'a>, - transaction: &'a T, - table_cache: &'a TableCache, - table_arena: &'a TableArenaCell, - scala_functions: &'a ScalaFunctions, - table_functions: &'a TableFunctions, } -impl ViewIter<'_, T> { - pub fn try_next(&mut self) -> Result, DatabaseError> { +impl ViewIter<'_, T> { + pub fn try_next( + &mut self, + arena: &mut (impl MetaArena + ?Sized), + ) -> Result, DatabaseError> { let Some((_, value)) = self.iter.try_next()? else { return Ok(None); }; - Ok(Some(TableCodec::decode_view( - value, - (self.transaction, self.table_cache), - self.scala_functions, - self.table_functions, - self.table_arena.borrow_mut(), - )?)) + Ok(Some(TableCodec::decode_view_name::(value, arena)?)) } } // GRCOV_EXCL_START #[cfg(test)] -pub(crate) fn check_range_rev_matches_range( +pub(crate) fn check_range_rev_matches_range( tx: &mut T, ) -> Result<(), DatabaseError> { fn keys(mut iter: I) -> Result>, DatabaseError> { @@ -2315,7 +2416,7 @@ pub(crate) fn check_explicit_transaction_does_not_rescan_own_writes( mut db: crate::db::Database, ) -> Result<(), DatabaseError> { db.ddl("create table t (id int primary key, v int not null)")?; - let mut tx = db.new_transaction()?; + let mut tx = db.new_transaction(false)?; tx.run("insert into t values (1,1), (2,2), (3,3), (4,4), (5,5)")? .done()?; tx.run("update t set id = id + 1000 where id > 3 and id < 3000")? @@ -2324,7 +2425,7 @@ pub(crate) fn check_explicit_transaction_does_not_rescan_own_writes( .done()?; let mut iter = tx.run("select id from t order by id")?; let mut ids = Vec::new(); - while let Some(id) = iter.next_tuple(|_, tuple| tuple.values[0].clone())? { + while let Some(id) = iter.next_tuple(|_, tuple| tuple[0].clone())? { ids.push(id); } iter.done()?; @@ -2337,7 +2438,7 @@ pub(crate) fn check_explicit_transaction_does_not_rescan_own_writes( } #[cfg(test)] -pub(crate) fn check_remove_range(tx: &mut T) -> Result<(), DatabaseError> { +pub(crate) fn check_remove_range(tx: &mut T) -> Result<(), DatabaseError> { let data: Vec<[u8; 4]> = (0u32..3000).map(u32::to_be_bytes).collect(); let [k0, k1, k10, k2500, k2999, k5000] = [0u32, 1, 10, 2500, 2999, 5000].map(u32::to_be_bytes); type Case<'a> = (Bound<&'a [u8]>, Bound<&'a [u8]>); @@ -2396,7 +2497,7 @@ mod test { use crate::storage::rocksdb::{RocksStorage, RocksTransaction}; use crate::storage::table_codec::TableCodec; use crate::storage::{ - IndexIter, InnerIter, Storage, TableCache, Transaction, TransactionIsolationLevel, + IndexIter, InnerIter, Storage, StorageTransaction, TableCache, TransactionIsolationLevel, }; use crate::types::index::{Index, IndexType}; use crate::types::tuple::Tuple; @@ -2561,11 +2662,11 @@ mod test { let plan_arena = PlanArena::new(&table_arena); let tuples = build_tuples(); - for tuple in tuples.iter() { + for tuple in build_tuples() { transaction.append_tuple( &mut table_codec, "t1", - tuple, + &tuple.into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), @@ -2876,11 +2977,11 @@ mod test { for (tuple_id, index) in indexes.iter().cloned() { transaction.add_index(&mut table_codec, "t1", index, &tuple_id)?; } - for tuple in tuples.iter() { + for tuple in build_tuples() { transaction.append_tuple( &mut table_codec, "t1", - tuple, + &tuple.into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), @@ -2955,22 +3056,24 @@ mod test { LogicalType::Integer.serializable(), ]; - let initial_tuple = Tuple::new( + let initial_tuple: crate::types::tuple::TupleImpl = Tuple::new( Some(DataValue::Int32(0)), vec![ DataValue::Int32(0), DataValue::Boolean(true), DataValue::Int32(0), ], - ); - let updated_tuple = Tuple::new( + ) + .into(); + let updated_tuple: crate::types::tuple::TupleImpl = Tuple::new( Some(DataValue::Int32(0)), vec![ DataValue::Int32(0), DataValue::Boolean(true), DataValue::Int32(1), ], - ); + ) + .into(); let table_arena = TableArenaCell::default(); let mut plan_arena = PlanArena::new(&table_arena); @@ -2998,10 +3101,10 @@ mod test { "t1", Index::new( index_id, - std::slice::from_ref(&initial_tuple.values[2]), + std::slice::from_ref(&initial_tuple[2]), IndexType::Normal, ), - initial_tuple.pk.as_ref().unwrap(), + initial_tuple.pk().unwrap(), )?; setup_tx.append_tuple(&mut table_codec, "t1", &initial_tuple, &serializers, false)?; setup_tx.commit()?; @@ -3009,7 +3112,7 @@ mod test { index_id }; let reader_tx = - storage.transaction_with_isolation(TransactionIsolationLevel::ReadCommitted)?; + storage.transaction_with_isolation(TransactionIsolationLevel::ReadCommitted, true)?; let tuple_id = { let mut index_iter = table_codec.with_index_bound("t1", index_id, |min, max| { reader_tx.range(Bound::Included(min), Bound::Included(max)) @@ -3040,20 +3143,20 @@ mod test { "t1", &Index::new( index_id, - std::slice::from_ref(&initial_tuple.values[2]), + std::slice::from_ref(&initial_tuple[2]), IndexType::Normal, ), - initial_tuple.pk.as_ref().unwrap(), + initial_tuple.pk().unwrap(), )?; writer_tx.add_index( &mut table_codec, "t1", Index::new( index_id, - std::slice::from_ref(&updated_tuple.values[2]), + std::slice::from_ref(&updated_tuple[2]), IndexType::Normal, ), - updated_tuple.pk.as_ref().unwrap(), + updated_tuple.pk().unwrap(), )?; writer_tx.append_tuple(&mut table_codec, "t1", &updated_tuple, &serializers, true)?; writer_tx.commit()?; diff --git a/src/storage/rocksdb.rs b/src/storage/rocksdb.rs index 11e23f45..3ee90be6 100644 --- a/src/storage/rocksdb.rs +++ b/src/storage/rocksdb.rs @@ -14,7 +14,7 @@ use crate::errors::DatabaseError; use crate::storage::{ - CheckpointableStorage, InnerIter, Storage, Transaction, TransactionIsolationLevel, + CheckpointableStorage, InnerIter, Storage, StorageTransaction, TransactionIsolationLevel, }; #[cfg(feature = "unsafe_txdb_checkpoint")] use librocksdb_sys as ffi; @@ -446,6 +446,7 @@ impl Storage for OptimisticRocksStorage { fn transaction_with_isolation( &self, isolation: TransactionIsolationLevel, + _read_only: bool, ) -> Result, DatabaseError> { self.validate_transaction_isolation(isolation)?; let tx = if self.config.sync { @@ -502,6 +503,7 @@ impl Storage for RocksStorage { fn transaction_with_isolation( &self, isolation: TransactionIsolationLevel, + _read_only: bool, ) -> Result, DatabaseError> { self.validate_transaction_isolation(isolation)?; let tx = if self.config.sync { @@ -636,7 +638,7 @@ fn build_read_options( #[macro_export] macro_rules! impl_transaction { ($tx:ident, $iter:ident) => { - impl<'storage> Transaction for $tx<'storage> { + impl<'storage> StorageTransaction for $tx<'storage> { type BorrowedBytes<'a> = DBPinnableSlice<'a> where @@ -899,7 +901,7 @@ mod test { use crate::storage::table_codec::TableCodec; use crate::storage::{ IndexImplEnum, IndexImplParams, IndexIter, IndexIterState, InnerIter, IterBounds, - PrimaryKeyIndexImpl, Storage, Transaction, + PrimaryKeyIndexImpl, Storage, StorageTransaction, }; use crate::types::index::IndexType; use crate::types::tuple::Tuple; @@ -1044,7 +1046,8 @@ mod test { &Tuple::new( Some(DataValue::Int32(1)), vec![DataValue::Int32(1), DataValue::Boolean(true)], - ), + ) + .into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), @@ -1057,7 +1060,8 @@ mod test { &Tuple::new( Some(DataValue::Int32(2)), vec![DataValue::Int32(2), DataValue::Boolean(true)], - ), + ) + .into(), &[ LogicalType::Integer.serializable(), LogicalType::Boolean.serializable(), @@ -1392,6 +1396,37 @@ mod test { ); drop(iter); + // An identity prefix projection still skips the remaining key fields; + // omitted mapping must not widen the output back to the full index. + let prefix_deserializers = [ + plan_arena.column(b_cover_column).datatype().serializable(), + plan_arena + .column(a_cover_column) + .datatype() + .skip_serializable(), + ]; + let mut iter = transaction.read_by_index( + kite_sql.state.table_cache(), + &plan_arena, + "t1".to_string().into(), + (None, None), + &[b_cover_column], + composite_index, + &[Range::Scope { + min: Bound::Unbounded, + max: Bound::Unbounded, + }][..], + false, + Some(&prefix_deserializers), + None, + )?; + for value in [0, 1, 2, 4] { + let tuple = crate::storage::next_tuple_for_test(&mut iter)?.unwrap(); + assert_eq!(tuple.values, vec![DataValue::Int32(value)]); + } + assert!(crate::storage::next_tuple_for_test(&mut iter)?.is_none()); + drop(iter); + let target_pk = DataValue::Int32(3); let covered_value = DataValue::Int32(4); let covered_ranges = [Range::Eq(covered_value.clone())]; diff --git a/src/storage/table_codec.rs b/src/storage/table_codec.rs index 4d9672e4..83a3bd58 100644 --- a/src/storage/table_codec.rs +++ b/src/storage/table_codec.rs @@ -13,22 +13,23 @@ // limitations under the License. use crate::catalog::view::View; -use crate::catalog::{ColumnCatalog, ColumnRelation, TableMeta}; +use crate::catalog::{ColumnCatalog, ColumnRelation, TableMeta, TableName}; use crate::db::{ScalaFunctions, TableFunctions}; use crate::errors::DatabaseError; use crate::optimizer::core::cm_sketch::{CountMinSketchMeta, CountMinSketchPage}; use crate::optimizer::core::histogram::Bucket; -use crate::optimizer::core::statistics_meta::StatisticsMetaRoot; +use crate::optimizer::core::histogram::HistogramMeta; +use crate::optimizer::core::statistics_meta::IndexStatisticsRoot; use crate::optimizer::core::top_n::ColumnTopN; use crate::planner::MetaArena; use crate::serdes::stable_hash::{StableHasher, TABLE_NAME_HASH_KEYS}; use crate::serdes::{ReferenceDecodeContext, ReferenceSerialization, ReferenceTables}; -use crate::storage::{TableCache, Transaction}; +use crate::storage::{StorageTransaction, TableCache}; use crate::types::index::{Index, IndexId, IndexMeta, IndexType, INDEX_ID_LEN}; use crate::types::serialize::TupleValueSerializableImpl; -use crate::types::tuple::{Tuple, TupleId}; +use crate::types::tuple::{Tuple, TupleId, TupleLike}; use crate::types::value::{DataValue, IndexKeyMapping}; -use crate::types::LogicalType; +use crate::types::{ColumnId, LogicalType}; use std::borrow::Borrow; use std::hash::{Hash, Hasher}; use std::io::{Cursor, Read, Seek, SeekFrom}; @@ -53,7 +54,7 @@ static EMPTY_REFERENCE_TABLES: LazyLock = LazyLock::new(Referen pub type Bytes = Vec; pub type BumpBytes<'bump> = bumpalo::collections::Vec<'bump, u8>; -type TupleValueWriter<'a> = &'a mut dyn FnMut(&Tuple, &mut Bytes) -> Result<(), DatabaseError>; +type TupleValueWriter<'a> = &'a mut dyn FnMut(&mut Bytes) -> Result<(), DatabaseError>; #[derive(Default)] pub struct TableCodec { @@ -90,8 +91,11 @@ enum CodecType { Hash, } +const INDEX_STATISTICS_TAG: u8 = b'i'; +const COLUMN_STATISTICS_TAG: u8 = b'c'; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum StatisticsCodecType { +pub(crate) enum IndexStatisticsCodecType { Root, SketchMeta, SketchPage, @@ -99,24 +103,24 @@ pub(crate) enum StatisticsCodecType { TopN, } -impl StatisticsCodecType { +impl IndexStatisticsCodecType { fn tag(self) -> u8 { match self { - StatisticsCodecType::Root => b'0', - StatisticsCodecType::SketchMeta => b'1', - StatisticsCodecType::SketchPage => b'2', - StatisticsCodecType::Bucket => b'3', - StatisticsCodecType::TopN => b'4', + IndexStatisticsCodecType::Root => b'0', + IndexStatisticsCodecType::SketchMeta => b'1', + IndexStatisticsCodecType::SketchPage => b'2', + IndexStatisticsCodecType::Bucket => b'3', + IndexStatisticsCodecType::TopN => b'4', } } fn from_tag(tag: u8) -> Result { match tag { - b'0' => Ok(StatisticsCodecType::Root), - b'1' => Ok(StatisticsCodecType::SketchMeta), - b'2' => Ok(StatisticsCodecType::SketchPage), - b'3' => Ok(StatisticsCodecType::Bucket), - b'4' => Ok(StatisticsCodecType::TopN), + b'0' => Ok(IndexStatisticsCodecType::Root), + b'1' => Ok(IndexStatisticsCodecType::SketchMeta), + b'2' => Ok(IndexStatisticsCodecType::SketchPage), + b'3' => Ok(IndexStatisticsCodecType::Bucket), + b'4' => Ok(IndexStatisticsCodecType::TopN), _ => Err(DatabaseError::InvalidValue(format!( "invalid statistics codec tag: {tag}" ))), @@ -290,7 +294,7 @@ impl TableCodec { &mut self, table_name: &str, tuple_id: &TupleId, - tuple_value: Option<(&Tuple, TupleValueWriter<'_>)>, + tuple_value: Option>, f: impl FnOnce(&[u8], &[u8]) -> Result, ) -> Result { Self::check_primary_key(tuple_id, 0)?; @@ -302,7 +306,7 @@ impl TableCodec { &mut self, table_name: &str, tuple_id: &TupleId, - tuple_value: Option<(&Tuple, TupleValueWriter<'_>)>, + tuple_value: Option>, f: impl FnOnce(&[u8], &[u8]) -> Result, ) -> Result { self.clear_buffers(); @@ -312,8 +316,8 @@ impl TableCodec { lower.push(BOUND_MIN_TAG); tuple_id.memcomparable_encode(lower)?; - if let Some((tuple, serializers)) = tuple_value { - serializers(tuple, value)?; + if let Some(write_value) = tuple_value { + write_value(value)?; value.extend_from_slice(&stamp.to_le_bytes()); } @@ -426,10 +430,10 @@ impl TableCodec { /// /// Unique index key: /// `{TableName}{INDEX_TAG}{BOUND_MIN_TAG}{IndexID}{BOUND_MIN_TAG}{DataValue}` - pub fn with_index( + pub fn with_index( &mut self, table_name: &str, - index: &Index, + index: &Index<'_, V>, tuple_id: Option<&TupleId>, f: impl FnOnce(&[u8], &[u8]) -> Result, ) -> Result { @@ -440,8 +444,8 @@ impl TableCodec { lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&index.id.to_le_bytes()); lower.push(BOUND_MIN_TAG); - for value in index.values { - value.memcomparable_encode(lower)?; + for i in 0..index.values.len() { + index.values.value_at(i).memcomparable_encode(lower)?; } if let Some(tuple_id) = tuple_id { @@ -548,7 +552,7 @@ impl TableCodec { } /// Range bounds covering all statistics keys for one index. - pub fn with_statistics_index_bound( + pub fn with_index_statistics_bound( &mut self, table_name: &str, index_id: IndexId, @@ -559,6 +563,7 @@ impl TableCodec { Self::write_key_prefix(lower, CodecType::Statistics, table_hash); lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&index_id.to_le_bytes()); + lower.push(INDEX_STATISTICS_TAG); upper.clear(); upper.extend_from_slice(lower); @@ -568,32 +573,79 @@ impl TableCodec { }) } - /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{ROOT_TAG}`. - pub fn with_statistics_meta( + fn with_statistics( &mut self, table_name: &str, - index_id: IndexId, - statistics_meta: Option<&StatisticsMetaRoot>, + id: &[u8], + suffix: &[u8], + statistics: Option<&V>, arena: &(impl MetaArena + ?Sized), f: impl FnOnce(&[u8], &[u8]) -> Result, ) -> Result { self.clear_buffers(); - self.with_table_hash_buffers(table_name, |lower, table_hash, value, refs| { - Self::write_key_prefix(lower, CodecType::Statistics, table_hash); - lower.push(BOUND_MIN_TAG); - lower.extend_from_slice(&index_id.to_le_bytes()); - lower.push(StatisticsCodecType::Root.tag()); - - if let Some(statistics_meta) = statistics_meta { - Self::encode_statistics_meta_value_into(statistics_meta, refs, value, arena)?; + self.with_table_hash_buffers(table_name, |key, table_hash, value, refs| { + Self::write_key_prefix(key, CodecType::Statistics, table_hash); + key.push(BOUND_MIN_TAG); + key.extend_from_slice(id); + key.extend_from_slice(suffix); + if let Some(statistics) = statistics { + statistics.encode(value, true, refs, arena)?; } - - f(lower.as_slice(), value.as_slice()) + f(key.as_slice(), value.as_slice()) }) } - /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{SKETCH_META_TAG}`. - pub fn with_statistics_sketch_meta( + pub fn with_column_statistics( + &mut self, + table_name: &str, + column_id: ColumnId, + meta: Option<&HistogramMeta>, + arena: &(impl MetaArena + ?Sized), + f: impl FnOnce(&[u8], &[u8]) -> Result, + ) -> Result { + self.with_statistics( + table_name, + &column_id.to_le_bytes(), + &[COLUMN_STATISTICS_TAG], + meta, + arena, + f, + ) + } + + pub fn decode_column_statistics( + bytes: &[u8], + arena: &mut (impl MetaArena + ?Sized), + ) -> Result { + HistogramMeta::decode::( + &mut Cursor::new(bytes), + None, + &ReferenceTables::new(), + arena, + ) + } + + /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{INDEX_STATISTICS_TAG}{ROOT_TAG}`. + pub fn with_index_statistics( + &mut self, + table_name: &str, + index_id: IndexId, + index_statistics: Option<&IndexStatisticsRoot>, + arena: &(impl MetaArena + ?Sized), + f: impl FnOnce(&[u8], &[u8]) -> Result, + ) -> Result { + self.with_statistics( + table_name, + &index_id.to_le_bytes(), + &[INDEX_STATISTICS_TAG, IndexStatisticsCodecType::Root.tag()], + index_statistics, + arena, + f, + ) + } + + /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{INDEX_STATISTICS_TAG}{SKETCH_META_TAG}`. + pub fn with_index_statistics_sketch_meta( &mut self, table_name: &str, index_id: IndexId, @@ -606,7 +658,8 @@ impl TableCodec { Self::write_key_prefix(lower, CodecType::Statistics, table_hash); lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&index_id.to_le_bytes()); - lower.push(StatisticsCodecType::SketchMeta.tag()); + lower.push(INDEX_STATISTICS_TAG); + lower.push(IndexStatisticsCodecType::SketchMeta.tag()); if let Some(sketch_meta) = sketch_meta { Self::encode_statistics_sketch_meta_value_into(sketch_meta, refs, value, arena)?; @@ -616,8 +669,8 @@ impl TableCodec { }) } - /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{SKETCH_PAGE_TAG}{BOUND_MIN_TAG}{ROW_ID}{BOUND_MIN_TAG}{PAGE_ID}`. - pub fn with_statistics_sketch_page( + /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{INDEX_STATISTICS_TAG}{SKETCH_PAGE_TAG}{BOUND_MIN_TAG}{ROW_ID}{BOUND_MIN_TAG}{PAGE_ID}`. + pub fn with_index_statistics_sketch_page( &mut self, table_name: &str, index_id: IndexId, @@ -631,7 +684,8 @@ impl TableCodec { Self::write_key_prefix(lower, CodecType::Statistics, table_hash); lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&index_id.to_le_bytes()); - lower.push(StatisticsCodecType::SketchPage.tag()); + lower.push(INDEX_STATISTICS_TAG); + lower.push(IndexStatisticsCodecType::SketchPage.tag()); lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&(sketch_page.row_idx() as u32).to_be_bytes()); lower.push(BOUND_MIN_TAG); @@ -645,8 +699,8 @@ impl TableCodec { }) } - /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{BUCKET_TAG}{BOUND_MIN_TAG}{ORDINAL}`. - pub fn with_statistics_bucket( + /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{INDEX_STATISTICS_TAG}{BUCKET_TAG}{BOUND_MIN_TAG}{ORDINAL}`. + pub fn with_index_statistics_bucket( &mut self, table_name: &str, index_id: IndexId, @@ -660,7 +714,8 @@ impl TableCodec { Self::write_key_prefix(lower, CodecType::Statistics, table_hash); lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&index_id.to_le_bytes()); - lower.push(StatisticsCodecType::Bucket.tag()); + lower.push(INDEX_STATISTICS_TAG); + lower.push(IndexStatisticsCodecType::Bucket.tag()); lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&ordinal.to_be_bytes()); @@ -672,8 +727,8 @@ impl TableCodec { }) } - /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{TOP_N_TAG}`. - pub fn with_statistics_top_n( + /// Key: `{TableName}{STATISTICS_TAG}{BOUND_MIN_TAG}{INDEX_ID}{INDEX_STATISTICS_TAG}{TOP_N_TAG}`. + pub fn with_index_statistics_top_n( &mut self, table_name: &str, index_id: IndexId, @@ -686,7 +741,8 @@ impl TableCodec { Self::write_key_prefix(lower, CodecType::Statistics, table_hash); lower.push(BOUND_MIN_TAG); lower.extend_from_slice(&index_id.to_le_bytes()); - lower.push(StatisticsCodecType::TopN.tag()); + lower.push(INDEX_STATISTICS_TAG); + lower.push(IndexStatisticsCodecType::TopN.tag()); if let Some(top_n) = top_n { Self::encode_statistics_top_n_value_into(top_n, refs, value, arena)?; @@ -813,7 +869,7 @@ impl TableCodec { index_meta.encode(value, true, reference_tables, arena) } - pub fn decode_index_meta( + pub fn decode_index_meta( bytes: &[u8], arena: &mut (impl MetaArena + ?Sized), ) -> Result { @@ -825,26 +881,37 @@ impl TableCodec { ) } - pub fn decode_index_key( + pub fn decode_index_key_into( + tuple: &mut Tuple, bytes: &[u8], ty: &LogicalType, mapping: &M, - ) -> Result, DatabaseError> { + ) -> Result<(), DatabaseError> { // Hash + TypeTag + Bound Min + Index Id Len + Bound Min let start = TUPLE_KEY_PREFIX_LEN + INDEX_ID_LEN + KEY_BOUND_LEN; - let mut reader = Cursor::new(&bytes[start..]); + let bytes = bytes.get(start..).ok_or(DatabaseError::InvalidIndex)?; + let mut reader = Cursor::new(bytes); let types = match ty { LogicalType::Tuple(types) => types.as_slice(), ty => std::slice::from_ref(ty), }; - let mut values = vec![DataValue::Null; mapping.target_len(types.len())]; + tuple.values.clear(); + tuple + .values + .resize(mapping.target_len(types.len()), DataValue::Null); for (index_pos, ty) in types.iter().enumerate() { - let value = DataValue::memcomparable_decode(&mut reader, ty)?; if let Some(scan_pos) = mapping.scan_index(index_pos) { - values[scan_pos] = value; + let value = DataValue::memcomparable_decode(&mut reader, ty)?; + let target = tuple + .values + .get_mut(scan_pos) + .ok_or(DatabaseError::InvalidIndex)?; + *target = value; + } else { + DataValue::memcomparable_skip(&mut reader, ty)?; } } - Ok(values) + Ok(()) } pub fn decode_index(bytes: &[u8]) -> Result { @@ -860,7 +927,7 @@ impl TableCodec { col.encode(value, true, reference_tables, arena) } - pub fn decode_column( + pub fn decode_column( reader: &mut R, reference_tables: &ReferenceTables, arena: &mut (impl MetaArena + ?Sized), @@ -869,20 +936,11 @@ impl TableCodec { ColumnCatalog::decode::(reader, None, reference_tables, arena) } - fn encode_statistics_meta_value_into( - statistics_meta: &StatisticsMetaRoot, - reference_tables: &mut ReferenceTables, - value: &mut Bytes, - arena: &(impl MetaArena + ?Sized), - ) -> Result<(), DatabaseError> { - statistics_meta.encode(value, true, reference_tables, arena) - } - - pub fn decode_statistics_meta( + pub fn decode_index_statistics( bytes: &[u8], arena: &mut (impl MetaArena + ?Sized), - ) -> Result { - StatisticsMetaRoot::decode::( + ) -> Result { + IndexStatisticsRoot::decode::( &mut Cursor::new(bytes), None, &EMPTY_REFERENCE_TABLES, @@ -899,7 +957,7 @@ impl TableCodec { sketch_meta.encode(value, true, reference_tables, arena) } - pub fn decode_statistics_sketch_meta( + pub fn decode_index_statistics_sketch_meta( bytes: &[u8], arena: &mut (impl MetaArena + ?Sized), ) -> Result { @@ -920,7 +978,7 @@ impl TableCodec { sketch_page.encode(value, true, reference_tables, arena) } - pub fn decode_statistics_sketch_page( + pub fn decode_index_statistics_sketch_page( bytes: &[u8], arena: &mut (impl MetaArena + ?Sized), ) -> Result { @@ -941,20 +999,20 @@ impl TableCodec { bucket.encode(value, true, reference_tables, arena) } - pub(crate) fn decode_statistics_codec_type( + pub(crate) fn decode_index_statistics_codec_type( key: &[u8], - ) -> Result { - let prefix_len = TUPLE_KEY_PREFIX_LEN + INDEX_ID_LEN; + ) -> Result { + let prefix_len = TUPLE_KEY_PREFIX_LEN + INDEX_ID_LEN + 1; let Some(tag) = key.get(prefix_len).copied() else { return Err(DatabaseError::InvalidValue( "statistics key is too short".to_string(), )); }; - StatisticsCodecType::from_tag(tag) + IndexStatisticsCodecType::from_tag(tag) } - pub fn decode_statistics_bucket( + pub fn decode_index_statistics_bucket( bytes: &[u8], arena: &mut (impl MetaArena + ?Sized), ) -> Result { @@ -975,7 +1033,7 @@ impl TableCodec { top_n.encode(value, true, reference_tables, arena) } - pub fn decode_statistics_top_n( + pub fn decode_index_statistics_top_n( bytes: &[u8], arena: &mut (impl MetaArena + ?Sized), ) -> Result { @@ -987,7 +1045,7 @@ impl TableCodec { ) } - pub fn decode_statistics_bucket_ordinal(bytes: &[u8]) -> Result { + pub fn decode_index_statistics_bucket_ordinal(bytes: &[u8]) -> Result { let len = bytes.len(); let Some(ordinal_bytes) = bytes.get(len - STATISTICS_BUCKET_ORD_LEN..) else { return Err(DatabaseError::InvalidValue( @@ -1018,7 +1076,7 @@ impl TableCodec { Ok(()) } - pub fn decode_view( + pub fn decode_view( bytes: &[u8], drive: (&T, &TableCache), scala_functions: &ScalaFunctions, @@ -1040,6 +1098,15 @@ impl TableCodec { View::decode(&mut cursor, Some(&context), &reference_tables, arena) } + pub fn decode_view_name( + bytes: &[u8], + arena: &mut (impl MetaArena + ?Sized), + ) -> Result { + let mut cursor = Cursor::new(bytes); + cursor.seek(SeekFrom::Start(4))?; + TableName::decode::(&mut cursor, None, &EMPTY_REFERENCE_TABLES, arena) + } + fn encode_root_table_value_into( meta: &TableMeta, reference_tables: &mut ReferenceTables, @@ -1049,7 +1116,7 @@ impl TableCodec { meta.encode(value, true, reference_tables, arena) } - pub fn decode_root_table( + pub fn decode_root_table( bytes: &[u8], arena: &mut (impl MetaArena + ?Sized), ) -> Result { @@ -1067,8 +1134,11 @@ mod tests { use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRelation, TableCatalog, TableMeta}; use crate::errors::DatabaseError; use crate::iter_ext::Itertools; - use crate::optimizer::core::histogram::{HistogramBuilder, ANALYZE_STATISTICS_RELATIVE_ERROR}; - use crate::optimizer::core::statistics_meta::StatisticsMeta; + use crate::optimizer::core::histogram::{ + HistogramBuilder, HistogramMetaBuilder, DEFAULT_HLL_PRECISION, + DEFAULT_STATISTICS_RELATIVE_ERROR, + }; + use crate::optimizer::core::statistics_meta::IndexStatistics; use crate::planner::{PlanArena, TableArenaCell}; use crate::serdes::ReferenceTables; use crate::storage::rocksdb::RocksTransaction; @@ -1076,9 +1146,9 @@ mod tests { use crate::storage::Storage; use crate::types::index::{Index, IndexMeta, IndexType}; use crate::types::tuple::Tuple; - use crate::types::value::DataValue; - use crate::types::ColumnId; + use crate::types::value::{DataValue, OrderedMapping, TupleMappingRef, Utf8Type}; use crate::types::LogicalType; + use crate::types::{CharLengthUnits, ColumnId}; use rust_decimal::Decimal; use std::collections::BTreeSet; use std::io::Cursor; @@ -1101,26 +1171,78 @@ mod tests { TableCatalog::new("t1".to_string().into(), columns, table_arena.borrow_mut()).unwrap() } + #[test] + fn covered_key_decode_uses_skip_configuration_and_output_mapping() -> Result<(), DatabaseError> + { + let varchar = LogicalType::Varchar(None, CharLengthUnits::Characters); + let char_type = LogicalType::Char(16, CharLengthUnits::Octets); + let values = vec![ + DataValue::Int32(7), + DataValue::Null, + DataValue::Utf8 { + value: "skippedę±‰å­—abcdefgh".into(), + ty: Utf8Type::Variable(None), + unit: CharLengthUnits::Characters, + }, + DataValue::Utf8 { + value: "0123456789abcdef".into(), + ty: Utf8Type::Fixed(16), + unit: CharLengthUnits::Octets, + }, + DataValue::Int32(99), + ]; + let index = Index::new(7, values.as_slice(), IndexType::Normal); + let mut codec = TableCodec::default(); + let key = codec.with_index("cover", &index, Some(&DataValue::Int32(1)), |key, _| { + Ok(key.to_vec()) + })?; + let index_ty = LogicalType::Tuple(vec![ + LogicalType::Integer, + varchar.clone(), + varchar, + char_type, + LogicalType::Integer, + ]); + let mapping = TupleMappingRef::new(&[1, usize::MAX, usize::MAX, usize::MAX, 0], 2); + let mut tuple = Tuple::new(None, Vec::with_capacity(8)); + let allocation = (tuple.values.as_ptr(), tuple.values.capacity()); + for _ in 0..2 { + TableCodec::decode_index_key_into(&mut tuple, &key, &index_ty, &mapping)?; + assert_eq!( + tuple.values, + vec![DataValue::Int32(99), DataValue::Int32(7)] + ); + assert_eq!((tuple.values.as_ptr(), tuple.values.capacity()), allocation); + } + + let mapping = TupleMappingRef::new(&[usize::MAX, 0, usize::MAX, usize::MAX, usize::MAX], 1); + TableCodec::decode_index_key_into(&mut tuple, &key, &index_ty, &mapping)?; + assert_eq!(tuple.values, vec![DataValue::Null]); + assert!( + TableCodec::decode_index_key_into(&mut tuple, &[], &index_ty, &OrderedMapping).is_err() + ); + Ok(()) + } + #[test] fn test_table_codec_tuple() -> Result<(), DatabaseError> { let table_arena = TableArenaCell::default(); let table_catalog = build_table_codec(&table_arena); let plan_arena = PlanArena::new(&table_arena); - let expected = Tuple::new( + let expected = crate::types::tuple::TupleImpl::from(Tuple::new( Some(DataValue::Int32(0)), vec![DataValue::Int32(0), DataValue::Decimal(Decimal::new(1, 0))], - ); + )); let serializers = [ LogicalType::Integer.serializable(), LogicalType::Decimal(None, None).serializable(), ]; - let mut write_value = - |tuple: &Tuple, value: &mut Bytes| tuple.serialize_to(&serializers, value); + let mut write_value = |value: &mut Bytes| expected.serialize_to(&serializers, value); let bytes = TableCodec::default().with_tuple( "t1", &DataValue::Int32(0), - Some((&expected, &mut write_value)), + Some(&mut write_value), |_, value| Ok(value.to_vec()), )?; let deserializers = table_catalog @@ -1166,7 +1288,7 @@ mod tests { } #[test] - fn test_table_codec_statistics_meta() -> Result<(), DatabaseError> { + fn test_table_codec_column_and_index_statistics() -> Result<(), DatabaseError> { let mut table_codec = TableCodec::default(); let table_arena = TableArenaCell::default(); let mut plan_arena = PlanArena::new(&table_arena); @@ -1179,22 +1301,63 @@ mod tests { name: "pk_c1".to_string(), ty: IndexType::PrimaryKey { is_multiple: false }, }; - let mut builder = HistogramBuilder::new(&index_meta, ANALYZE_STATISTICS_RELATIVE_ERROR)?; + let mut builder = HistogramBuilder::new( + &index_meta, + DEFAULT_STATISTICS_RELATIVE_ERROR, + DEFAULT_HLL_PRECISION, + )?; for value in 0..4 { builder.append(DataValue::Int32(value))?; } let (histogram, sketch, top_n) = builder.build(2)?; let (root, buckets, _, top_n) = - StatisticsMeta::new(histogram, sketch.clone(), top_n).into_parts(); + IndexStatistics::new(histogram, sketch.clone(), top_n).into_parts(); let root_bytes = - table_codec.with_statistics_meta("t1", 0, Some(&root), &plan_arena, |_, value| { + table_codec.with_index_statistics("t1", 0, Some(&root), &plan_arena, |_, value| { Ok::<_, DatabaseError>(value.to_vec()) })?; let decoded_root = - TableCodec::decode_statistics_meta::(&root_bytes, &mut plan_arena)?; + TableCodec::decode_index_statistics::(&root_bytes, &mut plan_arena)?; assert_eq!(decoded_root.index_id(), root.index_id()); + let index_key = + table_codec + .with_index_statistics("t1", 0, None, &plan_arena, |key, _| Ok(key.to_vec()))?; + let mut builder = HistogramMetaBuilder::new(DEFAULT_HLL_PRECISION)?; + for value in [ + DataValue::Int32(1), + DataValue::Int32(1), + DataValue::Int32(2), + DataValue::Null, + ] { + builder.append(&value); + } + let column_meta = builder.build(); + let (column_key, bytes) = table_codec.with_column_statistics( + "t1", + 0, + Some(&column_meta), + &plan_arena, + |key, value| Ok((key.to_vec(), value.to_vec())), + )?; + assert_ne!(column_key, index_key); + table_codec.with_index_statistics_bound("t1", 0, |min, max| { + assert!(index_key.as_slice() >= min && index_key.as_slice() <= max); + assert!(column_key.as_slice() < min || column_key.as_slice() > max); + Ok(()) + })?; + table_codec.with_statistics_bound("t1", |min, max| { + assert!(index_key.as_slice() >= min && index_key.as_slice() <= max); + assert!(column_key.as_slice() >= min && column_key.as_slice() <= max); + Ok(()) + })?; + let decoded = + TableCodec::decode_column_statistics::(&bytes, &mut plan_arena)?; + assert_eq!(decoded, column_meta); + assert_eq!(decoded.distinct_values_len(), 2); + assert_eq!(decoded.null_count(), 1); + assert_eq!(decoded.buckets_len(), 0); assert_eq!( decoded_root.histogram_meta().values_len(), root.histogram_meta().values_len() @@ -1205,22 +1368,21 @@ mod tests { ); let (sketch_meta, mut sketch_pages) = sketch.into_storage_parts(1); - let sketch_meta_bytes = table_codec.with_statistics_sketch_meta( + let sketch_meta_bytes = table_codec.with_index_statistics_sketch_meta( "t1", 0, Some(&sketch_meta), &plan_arena, |_, value| Ok::<_, DatabaseError>(value.to_vec()), )?; - let decoded_sketch_meta = TableCodec::decode_statistics_sketch_meta::( - &sketch_meta_bytes, - &mut plan_arena, - )?; + let decoded_sketch_meta = TableCodec::decode_index_statistics_sketch_meta::< + RocksTransaction, + >(&sketch_meta_bytes, &mut plan_arena)?; assert_eq!(decoded_sketch_meta.width(), sketch_meta.width()); assert_eq!(decoded_sketch_meta.k_num(), sketch_meta.k_num()); let first_sketch_page = sketch_pages.next().unwrap(); - let sketch_page_bytes = table_codec.with_statistics_sketch_page( + let sketch_page_bytes = table_codec.with_index_statistics_sketch_page( "t1", 0, &first_sketch_page, @@ -1228,30 +1390,29 @@ mod tests { &plan_arena, |_, value| Ok::<_, DatabaseError>(value.to_vec()), )?; - let decoded_sketch_page = TableCodec::decode_statistics_sketch_page::( - &sketch_page_bytes, - &mut plan_arena, - )?; + let decoded_sketch_page = TableCodec::decode_index_statistics_sketch_page::< + RocksTransaction, + >(&sketch_page_bytes, &mut plan_arena)?; assert_eq!(decoded_sketch_page.counters(), first_sketch_page.counters()); let bucket0_key = - table_codec.with_statistics_bucket("t1", 0, 0, None, &plan_arena, |key, _| { + table_codec.with_index_statistics_bucket("t1", 0, 0, None, &plan_arena, |key, _| { Ok::<_, DatabaseError>(key.to_vec()) })?; let bucket1_key = - table_codec.with_statistics_bucket("t1", 0, 1, None, &plan_arena, |key, _| { + table_codec.with_index_statistics_bucket("t1", 0, 1, None, &plan_arena, |key, _| { Ok::<_, DatabaseError>(key.to_vec()) })?; assert!(bucket0_key < bucket1_key); let (bucket0_min, bucket0_max) = - table_codec.with_statistics_index_bound("t1", 0, |min, max| { + table_codec.with_index_statistics_bound("t1", 0, |min, max| { Ok::<_, DatabaseError>((min.to_vec(), max.to_vec())) })?; assert!(bucket0_key.as_slice() >= bucket0_min.as_slice()); assert!(bucket1_key.as_slice() <= bucket0_max.as_slice()); - let bucket_bytes = table_codec.with_statistics_bucket( + let bucket_bytes = table_codec.with_index_statistics_bucket( "t1", 0, 0, @@ -1259,33 +1420,38 @@ mod tests { &plan_arena, |_, value| Ok::<_, DatabaseError>(value.to_vec()), )?; - let decoded_bucket = TableCodec::decode_statistics_bucket::( + let decoded_bucket = TableCodec::decode_index_statistics_bucket::( &bucket_bytes, &mut plan_arena, )?; assert_eq!(decoded_bucket, buckets[0]); assert_eq!( - TableCodec::decode_statistics_bucket_ordinal(&bucket0_key)?, + TableCodec::decode_index_statistics_bucket_ordinal(&bucket0_key)?, 0 ); assert_eq!( - TableCodec::decode_statistics_bucket_ordinal(&bucket1_key)?, + TableCodec::decode_index_statistics_bucket_ordinal(&bucket1_key)?, 1 ); let top_n_key = - table_codec.with_statistics_top_n("t1", 0, None, &plan_arena, |key, _| { + table_codec.with_index_statistics_top_n("t1", 0, None, &plan_arena, |key, _| { Ok::<_, DatabaseError>(key.to_vec()) })?; assert!(top_n_key.as_slice() >= bucket0_min.as_slice()); assert!(top_n_key.as_slice() <= bucket0_max.as_slice()); - let top_n_bytes = - table_codec.with_statistics_top_n("t1", 0, Some(&top_n), &plan_arena, |_, value| { - Ok::<_, DatabaseError>(value.to_vec()) - })?; - let decoded_top_n = - TableCodec::decode_statistics_top_n::(&top_n_bytes, &mut plan_arena)?; + let top_n_bytes = table_codec.with_index_statistics_top_n( + "t1", + 0, + Some(&top_n), + &plan_arena, + |_, value| Ok::<_, DatabaseError>(value.to_vec()), + )?; + let decoded_top_n = TableCodec::decode_index_statistics_top_n::( + &top_n_bytes, + &mut plan_arena, + )?; assert_eq!(decoded_top_n, top_n); Ok(()) @@ -1321,7 +1487,7 @@ mod tests { #[test] fn test_table_codec_index() -> Result<(), DatabaseError> { let tuple_id = DataValue::Int32(0); - let index = Index::new(0, &[], IndexType::Unique); + let index = Index::new(0, &[][..], IndexType::Unique); let bytes = TableCodec::default() .with_index("t1", &index, Some(&tuple_id), |_, value| Ok(value.to_vec()))?; @@ -1647,7 +1813,7 @@ mod tests { let value = Arc::new(value); let index = Index::new( index_id as u32, - std::slice::from_ref(&value), + std::slice::from_ref(value.as_ref()), IndexType::PrimaryKey { is_multiple: false }, ); @@ -1700,7 +1866,7 @@ mod tests { let value = Arc::new(value); let index = Index::new( index_id as u32, - std::slice::from_ref(&value), + std::slice::from_ref(value.as_ref()), IndexType::PrimaryKey { is_multiple: false }, ); diff --git a/src/types/bitset.rs b/src/types/bitset.rs new file mode 100644 index 00000000..295b7052 --- /dev/null +++ b/src/types/bitset.rs @@ -0,0 +1,141 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use kite_sql_serde_macros::ReferenceSerialization; + +#[derive(Debug, Default, Clone, PartialEq, Eq, Hash, ReferenceSerialization)] +pub struct BitSet { + blocks: Vec, + len: usize, +} + +impl BitSet { + pub fn with_capacity(len: usize) -> Self { + let bits_per_block = usize::BITS as usize; + let blocks = len.div_ceil(bits_per_block); + BitSet { + blocks: vec![0; blocks], + len, + } + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn clear(&mut self) { + self.blocks.fill(0); + } + + pub fn insert(&mut self, index: usize) { + assert!( + index < self.len, + "insert at index {index} exceeds bitmap size {}", + self.len + ); + let bits_per_block = usize::BITS as usize; + self.blocks[index / bits_per_block] |= 1usize << (index % bits_per_block); + } + + pub fn contains(&self, index: usize) -> bool { + if index >= self.len { + return false; + } + let bits_per_block = usize::BITS as usize; + (self.blocks[index / bits_per_block] & (1usize << (index % bits_per_block))) != 0 + } +} + +// GRCOV_EXCL_START +#[cfg(test)] +mod tests { + use super::BitSet; + + #[test] + fn bit_set_contains_false_for_empty_bitmap() { + let bitmap = BitSet::with_capacity(0); + + assert!(!bitmap.contains(0)); + assert!(!bitmap.contains(usize::MAX)); + } + + #[test] + fn bit_set_contains_false_for_unset_bits_within_capacity() { + let bitmap = BitSet::with_capacity(usize::BITS as usize + 1); + + assert!(!bitmap.contains(0)); + assert!(!bitmap.contains(usize::BITS as usize)); + } + + #[test] + fn bit_set_insert_marks_bits_across_blocks() { + let boundary = usize::BITS as usize; + let mut bitmap = BitSet::with_capacity(boundary + 2); + + bitmap.insert(0); + bitmap.insert(boundary - 1); + bitmap.insert(boundary); + bitmap.insert(boundary + 1); + + assert!(bitmap.contains(0)); + assert!(bitmap.contains(boundary - 1)); + assert!(bitmap.contains(boundary)); + assert!(bitmap.contains(boundary + 1)); + assert!(!bitmap.contains(1)); + } + + #[test] + fn bit_set_contains_returns_false_past_capacity() { + let mut bitmap = BitSet::with_capacity(2); + + bitmap.insert(1); + + assert!(!bitmap.contains(2)); + assert!(!bitmap.contains(usize::MAX)); + } + + #[test] + fn bit_set_clear_preserves_width_and_allocation() { + let width = usize::BITS as usize + 2; + let mut bitmap = BitSet::with_capacity(width); + let blocks = bitmap.blocks.as_ptr(); + bitmap.insert(0); + bitmap.insert(width - 1); + + bitmap.clear(); + assert_eq!(bitmap.len(), width); + assert_eq!(bitmap.blocks.as_ptr(), blocks); + assert!((0..width).all(|position| !bitmap.contains(position))); + bitmap.insert(1); + bitmap.clear(); + assert!(!bitmap.contains(1)); + + let mut empty = BitSet::with_capacity(0); + empty.clear(); + assert!(empty.is_empty()); + } + + #[test] + #[should_panic(expected = "exceeds bitmap size")] + fn bit_set_insert_panics_past_capacity() { + let mut bitmap = BitSet::with_capacity(1); + + bitmap.insert(1); + } +} +// GRCOV_EXCL_STOP diff --git a/src/types/index.rs b/src/types/index.rs index 7a0ff9c4..7b9753d5 100644 --- a/src/types/index.rs +++ b/src/types/index.rs @@ -21,6 +21,7 @@ use crate::planner::Explain; use crate::planner::ExprRef; use crate::planner::MetaArena; use crate::types::serialize::TupleValueSerializableImpl; +use crate::types::tuple::TupleLike; use crate::types::value::DataValue; use crate::types::{ColumnId, LogicalType}; use kite_sql_serde_macros::ReferenceSerialization; @@ -64,6 +65,15 @@ pub enum RuntimeIndexProbe { }, } +impl From for Range { + fn from(probe: RuntimeIndexProbe) -> Self { + match probe { + RuntimeIndexProbe::Eq(value) => Range::Eq(value), + RuntimeIndexProbe::Scope { min, max } => Range::Scope { min, max }, + } + } +} + #[derive(Debug, Clone, Eq, PartialEq, Hash, ReferenceSerialization)] pub enum IndexLookup { Static(Range), @@ -130,15 +140,23 @@ impl IndexMeta { } } -#[derive(Debug, Clone, Copy)] -pub struct Index<'a> { +#[derive(Debug)] +pub struct Index<'a, V: TupleLike + ?Sized = [DataValue]> { pub id: IndexId, - pub values: &'a [DataValue], + pub values: &'a V, pub ty: IndexType, } -impl<'a> Index<'a> { - pub fn new(id: IndexId, values: &'a [DataValue], ty: IndexType) -> Self { +impl Clone for Index<'_, V> { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Index<'_, V> {} + +impl<'a, V: TupleLike + ?Sized> Index<'a, V> { + pub fn new(id: IndexId, values: &'a V, ty: IndexType) -> Self { Index { id, values, ty } } } @@ -162,7 +180,7 @@ impl IndexInfo { } impl Explain for IndexInfo { - fn fmt(&self, arena: &dyn MetaArena, f: &mut Formatter<'_>) -> fmt::Result { + fn fmt(&self, arena: &A, f: &mut Formatter<'_>) -> fmt::Result { write!(f, "{} => ", self.meta.explain(arena))?; match &self.lookup { Some(IndexLookup::Static(range)) => write!(f, "{range}")?, diff --git a/src/types/mod.rs b/src/types/mod.rs index 9f6439d0..f73c0f5f 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod bitset; pub mod evaluator; pub mod index; pub mod serialize; @@ -154,6 +155,27 @@ impl LogicalType { } } + pub fn fixed_encoded_len(&self) -> Option { + let payload = match self { + LogicalType::Boolean | LogicalType::Tinyint | LogicalType::UTinyint => 1, + LogicalType::Smallint | LogicalType::USmallint => 2, + LogicalType::Integer + | LogicalType::UInteger + | LogicalType::Float + | LogicalType::Date => 4, + LogicalType::Bigint + | LogicalType::UBigint + | LogicalType::Double + | LogicalType::DateTime => 8, + LogicalType::Time(_) => 12, + LogicalType::TimeStamp(..) => 17, + #[cfg(feature = "decimal")] + LogicalType::Decimal(..) => 16, + _ => return None, + }; + Some(1 + payload) // One-byte variant tag. + } + pub fn numeric() -> Vec { vec![ LogicalType::Tinyint, @@ -495,6 +517,49 @@ pub(crate) mod test { use std::borrow::Cow; use std::io::{Cursor, Seek, SeekFrom}; + #[test] + #[cfg(feature = "spill")] + fn fixed_encoded_len_matches_values() -> Result<(), DatabaseError> { + use crate::execution::spill::SpillCodec; + use crate::types::value::DataValue; + + for ty in [ + LogicalType::Boolean, + LogicalType::Tinyint, + LogicalType::UTinyint, + LogicalType::Smallint, + LogicalType::USmallint, + LogicalType::Integer, + LogicalType::UInteger, + LogicalType::Bigint, + LogicalType::UBigint, + LogicalType::Float, + LogicalType::Double, + LogicalType::Date, + LogicalType::DateTime, + LogicalType::Time(Some(3)), + LogicalType::TimeStamp(Some(6), true), + #[cfg(feature = "decimal")] + LogicalType::Decimal(None, None), + ] { + assert_eq!( + ty.fixed_encoded_len(), + Some(DataValue::init(&ty).encoded_len()?), + "{ty:?}" + ); + } + for ty in [ + LogicalType::SqlNull, + LogicalType::Varchar(None, CharLengthUnits::Characters), + LogicalType::Char(2, CharLengthUnits::Octets), + LogicalType::Varbinary(None), + LogicalType::Tuple(vec![LogicalType::Integer]), + ] { + assert_eq!(ty.fixed_encoded_len(), None); + } + Ok(()) + } + #[test] fn test_char_length_units_display() { assert_eq!(CharLengthUnits::Characters.to_string(), "CHARACTERS"); diff --git a/src/types/tuple.rs b/src/types/tuple.rs index da734600..e1621587 100644 --- a/src/types/tuple.rs +++ b/src/types/tuple.rs @@ -19,7 +19,10 @@ use crate::planner::MetaArena; use crate::types::serialize::{TupleValueSerializable, TupleValueSerializableImpl}; use crate::types::value::DataValue; use std::borrow::Borrow; +use std::borrow::Cow; use std::io::Cursor; +use std::mem::MaybeUninit; +use std::ptr; const BITS_MAX_INDEX: usize = 8; @@ -75,6 +78,35 @@ impl<'a, 'p> SchemaView<'a, 'p> { } } +pub trait Input { + fn width(&self) -> usize; + + fn column(&self, position: usize) -> Cow<'_, DataValue>; +} + +impl Input for T { + fn width(&self) -> usize { + TupleLike::len(self) + } + + fn column(&self, position: usize) -> Cow<'_, DataValue> { + Cow::Borrowed(self.value_at(position)) + } +} + +// Sized adapter for passing any Input through the object-safe function interface. +pub(crate) struct InputRef<'a, T: ?Sized>(pub &'a T); + +impl Input for InputRef<'_, T> { + fn width(&self) -> usize { + self.0.width() + } + + fn column(&self, position: usize) -> Cow<'_, DataValue> { + self.0.column(position) + } +} + pub trait TupleLike { fn len(&self) -> usize; @@ -186,6 +218,486 @@ impl TupleLike for SplitTupleRef<' } } +#[derive(Clone, Debug)] +pub enum TupleImpl { + Tuple(Tuple), + Concat(Box), +} + +#[derive(Clone, Debug)] +pub struct ConcatTuple { + pk: Option, + left: TupleImpl, + right: TupleImpl, + left_len: usize, +} + +impl ConcatTuple { + #[inline] + fn route(&self, position: &mut usize) -> bool { + if *position < self.left_len { + true + } else { + *position -= self.left_len; + false + } + } +} + +#[allow(clippy::vec_box)] +#[derive(Debug, Default)] +pub(crate) struct ConcatPool(Vec>>); + +impl ConcatPool { + const CAPACITY: usize = 64; + + pub(crate) fn concat( + &mut self, + left: TupleImpl, + right: TupleImpl, + pk: Option, + ) -> TupleImpl { + let concat = ConcatTuple { + pk, + left_len: left.len(), + left, + right, + }; + TupleImpl::Concat(match self.0.pop() { + Some(slot) => Box::write(slot, concat), + None => Box::new(concat), + }) + } + + pub(crate) fn retain_leftmost<'a>(&mut self, tuple: &'a mut TupleImpl) -> &'a mut Tuple { + self.retain(tuple, true) + } + + pub(crate) fn retain_rightmost<'a>(&mut self, tuple: &'a mut TupleImpl) -> &'a mut Tuple { + self.retain(tuple, false) + } + + fn retain<'a>(&mut self, tuple: &'a mut TupleImpl, keep_left: bool) -> &'a mut Tuple { + loop { + match tuple { + TupleImpl::Tuple(tuple) => return tuple, + TupleImpl::Concat(_) => *tuple = self.split(std::mem::take(tuple), keep_left), + } + } + } + + fn split(&mut self, tuple: TupleImpl, keep_left: bool) -> TupleImpl { + match tuple { + TupleImpl::Tuple(tuple) => TupleImpl::Tuple(tuple), + TupleImpl::Concat(boxed) => { + let ConcatTuple { left, right, .. } = self.recycle(boxed); + let (keep, discard) = if keep_left { + (left, right) + } else { + (right, left) + }; + self.release(discard); + keep + } + } + } + + pub(crate) fn release(&mut self, tuple: TupleImpl) { + if let TupleImpl::Concat(boxed) = tuple { + let ConcatTuple { left, right, .. } = self.recycle(boxed); + self.release(left); + self.release(right); + } + } + + fn recycle(&mut self, boxed: Box) -> ConcatTuple { + let raw = Box::into_raw(boxed); + let (concat, slot) = unsafe { (ptr::read(raw), Box::from_raw(raw.cast())) }; + if self.0.len() < Self::CAPACITY { + self.0.push(slot); + } + concat + } +} + +impl Default for TupleImpl { + fn default() -> Self { + TupleImpl::Tuple(Tuple::default()) + } +} + +impl From for TupleImpl { + fn from(tuple: Tuple) -> Self { + TupleImpl::Tuple(tuple) + } +} + +impl TupleImpl { + pub fn concat(left: TupleImpl, right: TupleImpl, pk: Option) -> Self { + let left_len = left.len(); + TupleImpl::Concat(Box::new(ConcatTuple { + pk, + left, + right, + left_len, + })) + } + + #[inline] + pub fn pk(&self) -> Option<&TupleId> { + self.pk_slot().as_ref() + } + + /// The stored `Option` itself, for callers that encode it as a whole. + #[inline] + pub(crate) fn pk_slot(&self) -> &Option { + match self { + TupleImpl::Tuple(tuple) => &tuple.pk, + TupleImpl::Concat(concat) => &concat.pk, + } + } + + #[inline] + pub fn take_pk(&mut self) -> Option { + match self { + TupleImpl::Tuple(tuple) => tuple.pk.take(), + TupleImpl::Concat(concat) => concat.pk.take(), + } + } + + #[inline] + pub fn set_pk(&mut self, pk: Option) { + match self { + TupleImpl::Tuple(tuple) => tuple.pk = pk, + TupleImpl::Concat(concat) => concat.pk = pk, + } + } + + #[inline] + pub fn len(&self) -> usize { + TupleLike::len(self) + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn get(&self, index: usize) -> Option<&DataValue> { + (index < self.len()).then(|| self.value_at(index)) + } + + #[inline] + pub fn get_mut(&mut self, index: usize) -> Option<&mut DataValue> { + if index < self.len() { + Some(&mut self[index]) + } else { + None + } + } + + pub fn push(&mut self, value: DataValue) { + self.rightmost_mut().values.push(value); + } + + pub fn remove(&mut self, mut index: usize) -> DataValue { + let mut node = self; + loop { + match node { + TupleImpl::Tuple(tuple) => return tuple.values.remove(index), + TupleImpl::Concat(concat) => { + node = if concat.route(&mut index) { + concat.left_len -= 1; + &mut concat.left + } else { + &mut concat.right + }; + } + } + } + } + + pub fn prepend(&mut self, values: impl IntoIterator) { + match self { + TupleImpl::Tuple(tuple) => { + let len = tuple.values.len(); + tuple.values.extend(values); + let count = tuple.values.len() - len; + tuple.values.rotate_right(count); + } + TupleImpl::Concat(concat) => { + concat.left.prepend(values); + concat.left_len = concat.left.len(); + } + } + } + + pub fn resize(&mut self, len: usize, value: DataValue) { + let prefix: usize = self.len() - self.rightmost_mut().values.len(); + debug_assert!(len >= prefix, "resize cannot shrink into the left prefix"); + self.rightmost_mut() + .values + .resize(len.saturating_sub(prefix), value); + } + + pub(crate) fn rightmost_mut(&mut self) -> &mut Tuple { + let mut node = self; + loop { + match node { + TupleImpl::Tuple(tuple) => return tuple, + TupleImpl::Concat(concat) => node = &mut concat.right, + } + } + } + + fn value_at_mut(&mut self, mut index: usize) -> &mut DataValue { + let mut node = self; + loop { + match node { + TupleImpl::Tuple(tuple) => return &mut tuple.values[index], + TupleImpl::Concat(concat) => { + node = if concat.route(&mut index) { + &mut concat.left + } else { + &mut concat.right + }; + } + } + } + } + + /// e.g.: bits(u8)..|data_0(len for utf8_1)|utf8_0|data_1| + /// Tips: all len is u32 + pub fn serialize_to( + &self, + serializers: I, + bytes: &mut Vec, + ) -> Result<(), DatabaseError> + where + I: IntoIterator, + S: Borrow, + { + fn flip_bit(bits: u8, i: usize) -> u8 { + bits | (1 << (7 - i)) + } + + let values_len = self.len(); + let bits_len = (values_len + BITS_MAX_INDEX) / BITS_MAX_INDEX; + let values_bytes_len = self + .iter() + .map(DataValue::serialized_len_hint) + .sum::(); + bytes.clear(); + bytes.reserve(bits_len + values_bytes_len); + bytes.resize(bits_len, 0u8); + + for (i, (value, serializer)) in self.iter().zip(serializers).enumerate() { + if value.is_null() { + bytes[i / BITS_MAX_INDEX] = flip_bit(bytes[i / BITS_MAX_INDEX], i % BITS_MAX_INDEX); + } else { + serializer.borrow().to_raw(value, bytes)?; + } + } + + Ok(()) + } + + pub fn iter(&self) -> TupleImplIter<'_> { + match self.as_slice() { + Some(values) => TupleImplIter::Flat(values.iter()), + None => TupleImplIter::Joined { + tuple: self, + position: 0, + len: self.len(), + }, + } + } + + pub fn take_values(&mut self) -> Vec { + match self { + TupleImpl::Tuple(tuple) => std::mem::take(&mut tuple.values), + TupleImpl::Concat(_) => { + let mut values = Vec::with_capacity(self.len()); + self.drain_values(|value| values.push(value)); + values + } + } + } + + pub fn drain_values(&mut self, mut f: impl FnMut(DataValue)) { + self.drain_leaves(&mut f); + self.output_mut().clear(); + } + + fn drain_leaves(&mut self, f: &mut impl FnMut(DataValue)) { + match self { + TupleImpl::Tuple(tuple) => tuple.values.drain(..).for_each(f), + TupleImpl::Concat(concat) => { + concat.left.drain_leaves(f); + concat.right.drain_leaves(f); + } + } + } + + pub fn values_eq(&self, other: &TupleImpl) -> bool { + match (self.as_slice(), other.as_slice()) { + (Some(left), Some(right)) => left == right, + _ => self.len() == other.len() && self.iter().eq(other.iter()), + } + } + + pub fn output_mut(&mut self) -> &mut Tuple { + loop { + match self { + TupleImpl::Tuple(tuple) => return tuple, + TupleImpl::Concat(concat) => *self = std::mem::take(&mut concat.left), + } + } + } +} + +impl Extend for TupleImpl { + fn extend>(&mut self, values: I) { + self.rightmost_mut().values.extend(values); + } +} + +impl std::ops::Index for TupleImpl { + type Output = DataValue; + + #[inline] + fn index(&self, index: usize) -> &DataValue { + self.value_at(index) + } +} + +impl std::ops::IndexMut for TupleImpl { + #[inline] + fn index_mut(&mut self, index: usize) -> &mut DataValue { + self.value_at_mut(index) + } +} + +pub enum TupleImplIter<'a> { + Flat(std::slice::Iter<'a, DataValue>), + Joined { + tuple: &'a TupleImpl, + position: usize, + len: usize, + }, +} + +impl<'a> Iterator for TupleImplIter<'a> { + type Item = &'a DataValue; + + #[inline] + fn next(&mut self) -> Option { + match self { + TupleImplIter::Flat(iter) => iter.next(), + TupleImplIter::Joined { + tuple, + position, + len, + } => { + if *position >= *len { + return None; + } + let value = tuple.value_at(*position); + *position += 1; + Some(value) + } + } + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = match self { + TupleImplIter::Flat(iter) => iter.len(), + TupleImplIter::Joined { position, len, .. } => len - position, + }; + (remaining, Some(remaining)) + } +} + +impl ExactSizeIterator for TupleImplIter<'_> {} + +impl PartialEq for TupleImpl { + fn eq(&self, other: &Self) -> bool { + self.pk() == other.pk() && self.values_eq(other) + } +} + +impl Eq for TupleImpl {} + +impl PartialEq<[DataValue]> for TupleImpl { + fn eq(&self, other: &[DataValue]) -> bool { + self.iter().eq(other) + } +} + +impl PartialEq> for TupleImpl { + fn eq(&self, other: &Vec) -> bool { + *self == **other + } +} + +impl PartialEq<[DataValue; N]> for TupleImpl { + fn eq(&self, other: &[DataValue; N]) -> bool { + *self == other[..] + } +} + +impl PartialEq for TupleImpl { + fn eq(&self, other: &Tuple) -> bool { + self.pk() == other.pk.as_ref() && *self == other.values + } +} + +impl std::hash::Hash for TupleImpl { + fn hash(&self, state: &mut H) { + self.pk().hash(state); + state.write_usize(self.len()); + for value in self.iter() { + value.hash(state); + } + } +} + +impl TupleLike for TupleImpl { + #[inline] + fn len(&self) -> usize { + match self { + TupleImpl::Tuple(tuple) => tuple.values.len(), + TupleImpl::Concat(concat) => concat.left_len + concat.right.len(), + } + } + + #[inline] + fn value_at(&self, mut index: usize) -> &DataValue { + let mut node = self; + loop { + match node { + TupleImpl::Tuple(tuple) => return &tuple.values[index], + TupleImpl::Concat(concat) => { + node = if concat.route(&mut index) { + &concat.left + } else { + &concat.right + }; + } + } + } + } + + #[inline] + fn as_slice(&self) -> Option<&[DataValue]> { + match self { + TupleImpl::Tuple(tuple) => Some(tuple.values.as_slice()), + TupleImpl::Concat(_) => None, + } + } +} + impl<'a> From<&'a Tuple> for &'a [DataValue] { fn from(val: &'a Tuple) -> Self { val.values.as_slice() @@ -193,6 +705,11 @@ impl<'a> From<&'a Tuple> for &'a [DataValue] { } impl Tuple { + pub fn clear(&mut self) { + self.pk = None; + self.values.clear(); + } + pub fn new(pk: Option, values: Vec) -> Self { Tuple { pk, values } } @@ -238,48 +755,16 @@ impl Tuple { Ok(()) } - /// e.g.: bits(u8)..|data_0(len for utf8_1)|utf8_0|data_1| - /// Tips: all len is u32 - pub fn serialize_to( - &self, - serializers: I, - bytes: &mut Vec, - ) -> Result<(), DatabaseError> - where - I: IntoIterator, - S: Borrow, - { - fn flip_bit(bits: u8, i: usize) -> u8 { - bits | (1 << (7 - i)) - } - - let values_len = self.values.len(); - let bits_len = (values_len + BITS_MAX_INDEX) / BITS_MAX_INDEX; - let values_bytes_len = self - .values - .iter() - .map(DataValue::serialized_len_hint) - .sum::(); - bytes.clear(); - bytes.reserve(bits_len + values_bytes_len); - bytes.resize(bits_len, 0u8); - - for (i, (value, serializer)) in self.values.iter().zip(serializers).enumerate() { - if value.is_null() { - bytes[i / BITS_MAX_INDEX] = flip_bit(bytes[i / BITS_MAX_INDEX], i % BITS_MAX_INDEX); - } else { - serializer.borrow().to_raw(value, bytes)?; - } - } - - Ok(()) - } - - pub fn primary_projection(pk_indices: &[usize], values: &[DataValue]) -> TupleId { + pub fn primary_projection(pk_indices: &[usize], values: &V) -> TupleId { if pk_indices.len() > 1 { - DataValue::Tuple(pk_indices.iter().map(|i| values[*i].clone()).collect_vec()) + DataValue::Tuple( + pk_indices + .iter() + .map(|i| values.value_at(*i).clone()) + .collect_vec(), + ) } else { - values[pk_indices[0]].clone() + values.value_at(pk_indices[0]).clone() } } } @@ -292,7 +777,7 @@ mod tests { use crate::iter_ext::Itertools; use crate::planner::{PlanArena, TableArenaCell}; use crate::types::serialize::TupleValueSerializable; - use crate::types::tuple::Tuple; + use crate::types::tuple::{Tuple, TupleImpl}; use crate::types::value::{DataValue, Utf8Type}; use crate::types::CharLengthUnits; use crate::types::LogicalType; @@ -313,6 +798,23 @@ mod tests { } } + #[test] + fn clear_removes_pk_and_values_without_reallocating() { + let mut values = Vec::with_capacity(8); + values.push(DataValue::Int32(1)); + values.push(DataValue::Null); + let mut tuple = Tuple::new(Some(DataValue::Int32(1)), values); + let allocation = (tuple.values.as_ptr(), tuple.values.capacity()); + for _ in 0..2 { + tuple.clear(); + assert!(tuple.pk.is_none()); + assert!(tuple.values.is_empty()); + assert_eq!((tuple.values.as_ptr(), tuple.values.capacity()), allocation); + } + tuple.values.push(DataValue::Int32(2)); + assert_eq!((tuple.values.as_ptr(), tuple.values.capacity()), allocation); + } + #[test] fn test_tuple_serialize_to_and_deserialize_from() { let columns = Arc::new(vec![ @@ -498,7 +1000,8 @@ mod tests { DataValue::Null, ], ), - ]; + ] + .map(TupleImpl::from); let serializers = columns .iter() .map(|column| column.datatype().serializable()) @@ -506,7 +1009,7 @@ mod tests { let mut bytes = Vec::new(); { let mut tuple_0 = Tuple { - pk: tuples[0].pk.clone(), + pk: tuples[0].pk().cloned(), values: Vec::with_capacity(serializers.len()), }; tuples[0].serialize_to(&serializers, &mut bytes).unwrap(); @@ -514,11 +1017,11 @@ mod tests { .deserialize_from_into(&serializers, &bytes, columns.len()) .unwrap(); - assert_eq!(tuples[0], tuple_0); + assert_eq!(tuples[0], TupleImpl::from(tuple_0)); } { let mut tuple_1 = Tuple { - pk: tuples[1].pk.clone(), + pk: tuples[1].pk().cloned(), values: Vec::with_capacity(serializers.len()), }; tuples[1].serialize_to(&serializers, &mut bytes).unwrap(); @@ -526,7 +1029,7 @@ mod tests { .deserialize_from_into(&serializers, &bytes, columns.len()) .unwrap(); - assert_eq!(tuples[1], tuple_1); + assert_eq!(tuples[1], TupleImpl::from(tuple_1)); } // projection { @@ -537,7 +1040,7 @@ mod tests { columns[3].datatype().serializable(), ]; let mut tuple_2 = Tuple { - pk: tuples[0].pk.clone(), + pk: tuples[0].pk().cloned(), values: Vec::with_capacity(2), }; tuples[0].serialize_to(&serializers, &mut bytes).unwrap(); @@ -555,7 +1058,7 @@ mod tests { // A NULL in a skipped column must not shift later projected columns. let mut null_skipped = tuples[0].clone(); - null_skipped.values[1] = DataValue::Null; + null_skipped[1] = DataValue::Null; null_skipped.serialize_to(&serializers, &mut bytes).unwrap(); tuple_2 .deserialize_from_into(&projection_serializers, &bytes, columns.len()) @@ -574,13 +1077,11 @@ mod tests { .map(|column| column.datatype().serializable()) .collect_vec(); let mut multi_pk_tuple = tuples[0].clone(); - multi_pk_tuple.pk = Some(DataValue::Tuple(vec![ - multi_pk_tuple.values[4].clone(), - multi_pk_tuple.values[2].clone(), - ])); + let pk = DataValue::Tuple(vec![multi_pk_tuple[4].clone(), multi_pk_tuple[2].clone()]); + multi_pk_tuple.set_pk(Some(pk)); let mut tuple_3 = Tuple { - pk: multi_pk_tuple.pk.clone(), + pk: multi_pk_tuple.pk().cloned(), values: Vec::with_capacity(serializers.len()), }; multi_pk_tuple @@ -593,7 +1094,7 @@ mod tests { assert_eq!( tuple_3, Tuple { - pk: multi_pk_tuple.pk.clone(), + pk: multi_pk_tuple.pk().cloned(), values: vec![ DataValue::Int32(0), DataValue::UInt32(1), @@ -735,13 +1236,150 @@ mod tests { assert_eq!(tuple.values, vec![DataValue::Null, DataValue::Int32(7)]); assert_eq!( - Tuple::primary_projection(&[1], &tuple.values), + Tuple::primary_projection(&[1], tuple.values.as_slice()), DataValue::Int32(7) ); assert_eq!( - Tuple::primary_projection(&[0, 1], &tuple.values), + Tuple::primary_projection(&[0, 1], tuple.values.as_slice()), DataValue::Tuple(vec![DataValue::Null, DataValue::Int32(7)]) ); } + + fn ints(pk: Option, values: &[i32]) -> super::TupleImpl { + Tuple::new( + pk.map(DataValue::Int32), + values.iter().copied().map(DataValue::Int32).collect(), + ) + .into() + } + + fn values_of(tuple: &super::TupleImpl) -> Vec { + (0..tuple.len()) + .map(|i| tuple.value_at(i).clone()) + .collect() + } + + #[test] + fn tuple_impl_concat_reads_as_contiguous() { + use super::TupleImpl; + + let left = TupleImpl::concat(ints(Some(9), &[0, 1]), ints(None, &[2]), None); + let right = TupleImpl::concat(ints(None, &[3]), ints(None, &[4, 5]), None); + let mut tuple = TupleImpl::concat(left, right, Some(DataValue::Int32(7))); + + assert_eq!(tuple.len(), 6); + assert!(tuple.as_slice().is_none()); + assert_eq!( + values_of(&tuple), + (0..6).map(DataValue::Int32).collect_vec() + ); + assert_eq!(tuple.pk(), Some(&DataValue::Int32(7))); + + tuple.push(DataValue::Int32(6)); + assert_eq!(tuple.len(), 7); + assert_eq!(tuple.value_at(6), &DataValue::Int32(6)); + + assert_eq!(tuple, (0..7).map(DataValue::Int32).collect_vec()); + assert_eq!(tuple[6], DataValue::Int32(6)); + assert_eq!(tuple.get(7), None); + + tuple[1] = DataValue::Int32(10); + assert_eq!(tuple.get(1), Some(&DataValue::Int32(10))); + tuple.set_pk(None); + assert_eq!(tuple.pk(), None); + + let tuple = TupleImpl::concat(ints(None, &[]), ints(None, &[1, 2]), None); + assert_eq!(tuple, [1, 2].map(DataValue::Int32)); + } + + #[test] + fn tuple_impl_positional_access_ignores_shape() { + use super::TupleImpl; + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + + let hash = |tuple: &TupleImpl| { + let mut hasher = DefaultHasher::new(); + tuple.hash(&mut hasher); + hasher.finish() + }; + let mut joined = TupleImpl::concat( + TupleImpl::concat(ints(None, &[0]), ints(None, &[1, 2]), None), + ints(None, &[3]), + Some(DataValue::Int32(9)), + ); + let flat = ints(Some(9), &[0, 1, 2, 3]); + + assert_eq!(joined.iter().len(), 4); + assert_eq!(joined, (0..4).map(DataValue::Int32).collect_vec()); + assert!(joined.values_eq(&flat)); + assert_eq!(joined, flat); + assert_eq!(hash(&joined), hash(&flat)); + assert_ne!(joined, ints(None, &[0, 1, 2, 3])); + assert!(!joined.values_eq(&ints(None, &[0, 1, 2]))); + + *joined.value_at_mut(2) = DataValue::Int32(20); + assert_eq!(joined.value_at(2), &DataValue::Int32(20)); + joined.resize(6, DataValue::Null); + assert_eq!(joined.len(), 6); + assert_eq!(joined.value_at(5), &DataValue::Null); + + let mut drained = Vec::new(); + joined.drain_values(|value| drained.push(value)); + assert_eq!( + drained, + [0, 1, 20, 3] + .map(DataValue::Int32) + .into_iter() + .chain([DataValue::Null, DataValue::Null]) + .collect_vec() + ); + assert!(matches!(joined, TupleImpl::Tuple(_))); + assert_eq!(joined.len(), 0); + assert_eq!(joined.pk(), None); + } + + #[test] + fn tuple_impl_output_keeps_leftmost_buffer() { + use super::TupleImpl; + + let mut tuple = TupleImpl::concat(ints(None, &[0, 1, 2, 3]), ints(None, &[4]), None); + let output = tuple.output_mut(); + assert!(output.values.capacity() >= 4); + assert_eq!(output.values.len(), 4); + assert!(matches!(tuple, TupleImpl::Tuple(_))); + } + + #[test] + fn concat_pool_recycles_and_reuses_boxes() { + use super::{ConcatPool, TupleImpl}; + + let mut pool = ConcatPool::default(); + let inner = pool.concat(ints(None, &[0]), ints(None, &[1]), None); + let mut tuple = pool.concat(inner, ints(None, &[2, 3]), Some(DataValue::Int32(9))); + assert_eq!(tuple, [0, 1, 2, 3].map(DataValue::Int32)); + + assert_eq!( + pool.retain_rightmost(&mut tuple).values, + [2, 3].map(DataValue::Int32) + ); + assert_eq!(pool.0.len(), 2); + + let inner = pool.concat(ints(None, &[4]), ints(None, &[5]), None); + let mut tuple = pool.concat(ints(None, &[6]), inner, None); + assert!(pool.0.is_empty()); + assert_eq!(tuple, [6, 4, 5].map(DataValue::Int32)); + + assert_eq!( + pool.retain_leftmost(&mut tuple).values, + [6].map(DataValue::Int32) + ); + assert!(matches!(tuple, TupleImpl::Tuple(_))); + assert_eq!(pool.0.len(), 2); + + let dropped = pool.concat(ints(None, &[7]), ints(None, &[8]), None); + drop(dropped); + assert_eq!(pool.0.len(), 1); + } } // GRCOV_EXCL_STOP diff --git a/src/types/tuple_builder.rs b/src/types/tuple_builder.rs index b5d9cd9e..08cd6bc1 100644 --- a/src/types/tuple_builder.rs +++ b/src/types/tuple_builder.rs @@ -75,7 +75,7 @@ impl<'a> TupleBuilder<'a> { tuple.pk = self .pk_indices - .map(|indices| Tuple::primary_projection(indices, &tuple.values)); + .map(|indices| Tuple::primary_projection(indices, tuple.values.as_slice())); Ok(()) } diff --git a/src/types/value.rs b/src/types/value.rs index a1e8fe0f..0ff77c2f 100644 --- a/src/types/value.rs +++ b/src/types/value.rs @@ -15,6 +15,7 @@ use super::LogicalType; use crate::errors::DatabaseError; use crate::iter_ext::Itertools; +use crate::planner::ScalarQueryRef; use crate::storage::table_codec::{BumpBytes, NOTNULL_TAG, NULL_TAG}; use crate::types::evaluator::cast::cast_create; use crate::types::CharLengthUnits; @@ -137,6 +138,10 @@ pub enum DataValue { id: usize, ty: LogicalType, }, + InitValue { + id: ScalarQueryRef, + ty: LogicalType, + }, Null, Boolean(bool), Float32(OrderedFloat), @@ -363,6 +368,17 @@ impl PartialEq for DataValue { }, ) => left == right && left_ty == right_ty, (Parameter { .. }, _) => false, + ( + InitValue { + id: left, + ty: left_ty, + }, + InitValue { + id: right, + ty: right_ty, + }, + ) => left == right && left_ty == right_ty, + (InitValue { .. }, _) => false, (Boolean(v1), Boolean(v2)) => v1.eq(v2), (Boolean(_), _) => false, (Float32(v1), Float32(v2)) => v1.eq(v2), @@ -456,7 +472,7 @@ impl PartialOrd for DataValue { fn partial_cmp(&self, other: &Self) -> Option { use DataValue::*; match (self, other) { - (Parameter { .. }, _) => None, + (Parameter { .. }, _) | (InitValue { .. }, _) => None, (Boolean(v1), Boolean(v2)) => v1.partial_cmp(v2), (Boolean(_), _) => None, (Float32(v1), Float32(v2)) => v1.partial_cmp(v2), @@ -528,6 +544,11 @@ impl Hash for DataValue { id.hash(state); ty.hash(state); } + InitValue { id, ty } => { + 20u8.hash(state); + id.hash(state); + ty.hash(state); + } Null => 0u8.hash(state), Boolean(v) => { 1u8.hash(state); @@ -676,30 +697,46 @@ impl DataValue { digits.parse().ok() } - pub(crate) fn bind_parameters<'p, F>(&mut self, get: &F) -> Result<(), DatabaseError> + pub(crate) fn bind_parameters<'p, F, G>( + &mut self, + params: &F, + init_values: &G, + ) -> Result where F: Fn(usize) -> Option<&'p DataValue>, + G: Fn(ScalarQueryRef) -> Option<&'p DataValue>, { match self { DataValue::Parameter { id, ty } => { - let Some(input) = get(*id) else { + let Some(input) = params(*id) else { return Err(DatabaseError::parameter_not_found(format!("${id}"))); }; *self = input.clone().cast(ty)?; + Ok(false) + } + DataValue::InitValue { id, ty } => { + let Some(input) = init_values(*id) else { + return Err(DatabaseError::InvalidValue(format!( + "scalar query {id} is not initialized" + ))); + }; + *self = input.clone().cast(ty)?; + Ok(self.is_null()) } DataValue::Tuple(values) => { + let mut has_null = false; for value in values { - value.bind_parameters(get)?; + has_null |= value.bind_parameters(params, init_values)?; } + Ok(has_null) } - _ => {} + _ => Ok(false), } - Ok(()) } pub(crate) fn has_parameter(&self) -> bool { match self { - DataValue::Parameter { .. } => true, + DataValue::Parameter { .. } | DataValue::InitValue { .. } => true, DataValue::Tuple(values) => values.iter().any(Self::has_parameter), _ => false, } @@ -707,7 +744,7 @@ impl DataValue { pub(crate) fn serialized_len_hint(&self) -> usize { match self { - DataValue::Parameter { .. } => 0, + DataValue::Parameter { .. } | DataValue::InitValue { .. } => 0, DataValue::Null => 0, DataValue::Boolean(_) | DataValue::Int8(_) | DataValue::UInt8(_) => 1, DataValue::Int16(_) | DataValue::UInt16(_) => 2, @@ -884,6 +921,14 @@ impl DataValue { matches!(self, DataValue::Null) } + #[inline] + pub fn has_null(&self) -> bool { + match self { + DataValue::Tuple(values) => values.iter().any(DataValue::has_null), + value => value.is_null(), + } + } + #[inline] pub fn init(logic_type: &LogicalType) -> DataValue { match logic_type { @@ -984,7 +1029,7 @@ impl DataValue { #[inline] pub fn logical_type(&self) -> LogicalType { match self { - DataValue::Parameter { ty, .. } => ty.clone(), + DataValue::Parameter { ty, .. } | DataValue::InitValue { ty, .. } => ty.clone(), DataValue::Null => LogicalType::SqlNull, DataValue::Boolean(_) => LogicalType::Boolean, DataValue::Float32(_) => LogicalType::Float, @@ -1073,7 +1118,14 @@ impl DataValue { #[inline] fn decode_string(reader: &mut R) -> std::io::Result> { let mut result = Vec::new(); + Self::read_memcomparable_string(reader, |group| result.extend_from_slice(group))?; + Ok(result) + } + fn read_memcomparable_string( + reader: &mut R, + mut append: impl FnMut(&[u8]), + ) -> std::io::Result<()> { loop { let mut group = [0u8; ENCODE_GROUP_SIZE]; reader.read_exact(&mut group)?; @@ -1092,14 +1144,40 @@ impl DataValue { } let data_len = ENCODE_GROUP_SIZE - pad_count; - result.extend_from_slice(&group[..data_len]); + append(&group[..data_len]); if pad_count != 0 { - break; + return Ok(()); } } + } - Ok(result) + fn skip_memcomparable_string(reader: &mut R) -> std::io::Result<()> { + Self::read_memcomparable_string(reader, |_| {}) + } + + fn skip_memcomparable_decimal(reader: &mut R) -> Result<(), DatabaseError> { + let flag = reader.read_u8()?; + match flag { + 0x08 | 0x14 | 0x16 | 0x22 => { + reader.read_u8()?; + } + 0x09..=0x13 | 0x17..=0x21 => {} + 0x15 => return Ok(()), + _ => { + return Err(DatabaseError::InvalidValue(format!( + "invalid decimal exponent: {flag}" + ))) + } + } + let negative = (0x07..0x15).contains(&flag); + loop { + let byte = reader.read_u8()?; + let byte = if negative { !byte } else { byte }; + if byte & 1 == 0 { + return Ok(()); + } + } } #[inline] @@ -1125,7 +1203,7 @@ impl DataValue { b.push_byte(not_null_tag); match self { - DataValue::Parameter { .. } => { + DataValue::Parameter { .. } | DataValue::InitValue { .. } => { return Err(DatabaseError::InvalidValue( "unbound parameter reached storage encoding".to_string(), )); @@ -1189,6 +1267,44 @@ impl DataValue { self.memcomparable_encode_with_null_order(b, false) } + pub(crate) fn memcomparable_skip( + reader: &mut R, + ty: &LogicalType, + ) -> Result<(), DatabaseError> { + if reader.read_u8()? == NULL_TAG { + return Ok(()); + } + let size = match ty { + LogicalType::SqlNull => return Ok(()), + LogicalType::Boolean | LogicalType::Tinyint | LogicalType::UTinyint => 1, + LogicalType::Smallint | LogicalType::USmallint => 2, + LogicalType::Integer + | LogicalType::UInteger + | LogicalType::Float + | LogicalType::Date + | LogicalType::Time(_) => 4, + LogicalType::Bigint + | LogicalType::UBigint + | LogicalType::Double + | LogicalType::DateTime + | LogicalType::TimeStamp(..) => 8, + LogicalType::Char(..) | LogicalType::Varchar(..) | LogicalType::Varbinary(_) => { + Self::skip_memcomparable_string(reader)?; + return Ok(()); + } + LogicalType::Decimal(..) => return Self::skip_memcomparable_decimal(reader), + LogicalType::Tuple(types) => { + for ty in types { + Self::memcomparable_skip(reader, ty)?; + } + return Ok(()); + } + }; + let mut discard = [0u8; 8]; + reader.read_exact(&mut discard[..size])?; + Ok(()) + } + pub fn memcomparable_decode( reader: &mut R, ty: &LogicalType, @@ -1501,6 +1617,9 @@ impl DataValue { (DataValue::Parameter { id, .. }, ty) => { Ok(DataValue::Parameter { id, ty: ty.clone() }) } + (DataValue::InitValue { id, .. }, ty) => { + Ok(DataValue::InitValue { id, ty: ty.clone() }) + } (DataValue::Null, _) => Ok(DataValue::Null), (value, _) => { let evaluator = cast_create(&from, to)?; @@ -1823,6 +1942,7 @@ impl fmt::Display for DataValue { fn fmt(&self, f: &mut Formatter) -> fmt::Result { match self { DataValue::Parameter { id, .. } => write!(f, "${id}")?, + DataValue::InitValue { id, .. } => write!(f, "InitValue({id})")?, DataValue::Boolean(e) => write!(f, "{e}")?, DataValue::Float32(e) => format_float_option!(f, e)?, DataValue::Float64(e) => format_float_option!(f, e)?, @@ -1900,6 +2020,7 @@ impl fmt::Debug for DataValue { fn fmt(&self, f: &mut Formatter) -> fmt::Result { match self { DataValue::Parameter { ty, .. } => write!(f, "Parameter({self}: {ty})"), + DataValue::InitValue { ty, .. } => write!(f, "{self}: {ty}"), DataValue::Boolean(_) => write!(f, "Boolean({self})"), DataValue::Float32(_) => write!(f, "Float32({self})"), DataValue::Float64(_) => write!(f, "Float64({self})"), @@ -2351,6 +2472,121 @@ mod test { ); } + #[test] + fn memcomparable_skip_stops_at_the_value_boundary() -> Result<(), DatabaseError> { + // Fixed-width scalars, at their initial value. + let mut cases: Vec<_> = [ + LogicalType::Boolean, + LogicalType::Tinyint, + LogicalType::Smallint, + LogicalType::Integer, + LogicalType::Bigint, + LogicalType::UTinyint, + LogicalType::USmallint, + LogicalType::UInteger, + LogicalType::UBigint, + LogicalType::Float, + LogicalType::Double, + LogicalType::Date, + LogicalType::DateTime, + LogicalType::Time(Some(3)), + LogicalType::TimeStamp(Some(6), true), + ] + .into_iter() + .map(|ty| (DataValue::init(&ty), ty)) + .collect(); + // Variable payloads around the 8-byte group boundary. + for length in [0, 7, 8, 9, 65] { + cases.push(( + DataValue::Binary { + value: vec![0xff; length], + len: None, + }, + LogicalType::Varbinary(None), + )); + } + cases.extend([ + ( + DataValue::Utf8 { + value: "ę±‰å­—abcdefgh".into(), + ty: Utf8Type::Variable(None), + unit: CharLengthUnits::Characters, + }, + LogicalType::Varchar(None, CharLengthUnits::Characters), + ), + ( + DataValue::Utf8 { + value: "0123456789abcdef".into(), + ty: Utf8Type::Fixed(16), + unit: CharLengthUnits::Octets, + }, + LogicalType::Char(16, CharLengthUnits::Octets), + ), + ( + DataValue::Tuple(vec![DataValue::Int32(1), DataValue::Null]), + LogicalType::Tuple(vec![ + LogicalType::Integer, + LogicalType::Varchar(None, CharLengthUnits::Characters), + ]), + ), + ]); + #[cfg(feature = "decimal")] + for text in ["0", "-0.001", "79228162514264337593543950335"] { + cases.push(( + DataValue::Decimal(text.parse::().unwrap()), + LogicalType::Decimal(None, None), + )); + } + + for (value, ty) in cases { + for value in [value, DataValue::Null] { + let mut bytes = Vec::new(); + value.memcomparable_encode(&mut bytes)?; + let boundary = bytes.len(); + DataValue::Int32(99).memcomparable_encode(&mut bytes)?; + let mut reader = Cursor::new(bytes.as_slice()); + DataValue::memcomparable_skip(&mut reader, &ty)?; + assert_eq!(reader.position() as usize, boundary, "{ty:?}"); + assert_eq!( + DataValue::memcomparable_decode(&mut reader, &LogicalType::Integer)?, + DataValue::Int32(99) + ); + } + } + Ok(()) + } + + #[test] + fn memcomparable_skip_rejects_truncation_and_invalid_markers() { + assert!(DataValue::memcomparable_skip( + &mut Cursor::new(&[0, 1, 2][..]), + &LogicalType::Integer + ) + .is_err()); + assert!( + DataValue::memcomparable_skip(&mut Cursor::new(&[][..]), &LogicalType::Integer) + .is_err() + ); + for bytes in [ + vec![0, 1, 2], + vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 246], + vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 255], + ] { + assert!(DataValue::memcomparable_skip( + &mut Cursor::new(bytes.as_slice()), + &LogicalType::Varbinary(None) + ) + .is_err()); + } + for bytes in [vec![0, 7], vec![0, 0x22], vec![0, 0x18, 1]] { + assert!(DataValue::memcomparable_skip( + &mut Cursor::new(bytes.as_slice()), + &LogicalType::Decimal(None, None) + ) + .is_err()); + } + } + #[test] fn test_memcomparable_decode_rejects_invalid_strings() { let mut invalid_marker = vec![NOTNULL_TAG]; diff --git a/src/wasm.rs b/src/wasm.rs index c8986d12..c4123ea5 100644 --- a/src/wasm.rs +++ b/src/wasm.rs @@ -16,7 +16,7 @@ use crate::db::{DataBaseBuilder, Database, DatabaseIter}; use crate::storage::memory::MemoryStorage; -use crate::types::tuple::Tuple; +use crate::types::tuple::TupleImpl; use crate::types::value::{DataValue, Utf8Type}; use crate::types::CharLengthUnits; use js_sys::{Array, Object, Reflect, Uint8Array}; @@ -33,8 +33,8 @@ fn set_prop(object: &Object, key: &str, value: JsValue) -> Result<(), JsValue> { fn data_value_to_js(value: &DataValue) -> Result { match value { - DataValue::Parameter { id, .. } => Err(to_js_err(format!( - "unbound parameter ${id} reached WASM output" + DataValue::Parameter { .. } | DataValue::InitValue { .. } => Err(to_js_err(format!( + "unbound parameter {value} reached WASM output" ))), DataValue::Null => Ok(JsValue::NULL), DataValue::Boolean(value) => Ok(JsValue::from_bool(*value)), @@ -108,7 +108,9 @@ fn char_length_units_to_js(unit: CharLengthUnits) -> JsValue { }) } -fn data_values_to_js(values: &[DataValue]) -> Result { +fn data_values_to_js<'v>( + values: impl IntoIterator, +) -> Result { let array = Array::new(); for value in values { array.push(&data_value_to_js(value)?); @@ -116,19 +118,18 @@ fn data_values_to_js(values: &[DataValue]) -> Result { Ok(array.into()) } -fn tuple_to_wasm_row(tuple: &Tuple) -> Result { +fn tuple_to_wasm_row(tuple: &TupleImpl) -> Result { let object = Object::new(); set_prop( &object, "pk", tuple - .pk - .as_ref() + .pk() .map(data_value_to_js) .transpose()? .unwrap_or(JsValue::NULL), )?; - set_prop(&object, "values", data_values_to_js(&tuple.values)?)?; + set_prop(&object, "values", data_values_to_js(tuple.iter())?)?; Ok(object.into()) } diff --git a/tests/kitesql-testing b/tests/kitesql-testing new file mode 160000 index 00000000..54eed09e --- /dev/null +++ b/tests/kitesql-testing @@ -0,0 +1 @@ +Subproject commit 54eed09e149d10da06ba4ae7ffd04b743e0b7363 diff --git a/tests/macros-test/src/main.rs b/tests/macros-test/src/main.rs index ec399e45..95721683 100644 --- a/tests/macros-test/src/main.rs +++ b/tests/macros-test/src/main.rs @@ -29,7 +29,7 @@ mod test { use kite_sql::planner::{MetaArena, PlanArena, TableArena, TableArenaCell}; use kite_sql::storage::rocksdb::RocksStorage; use kite_sql::types::evaluator::binary_create; - use kite_sql::types::tuple::{Schema, SchemaView, Tuple}; + use kite_sql::types::tuple::{Schema, SchemaView, Tuple, TupleImpl}; use kite_sql::types::value::{DataValue, Utf8Type}; use kite_sql::types::{CharLengthUnits, LogicalType}; use kite_sql::{from_tuple, scala_function, table_function, Model, Projection}; @@ -68,16 +68,9 @@ mod test { (Tuple::new(None, values), schema) } - fn tuple_owned(tuple: &Tuple) -> Tuple { - Tuple { - pk: tuple.pk.clone(), - values: tuple.values.clone(), - } - } - - fn collect_result_tuples(mut iter: I) -> Result, DatabaseError> { + fn collect_result_tuples(mut iter: I) -> Result, DatabaseError> { let mut rows = Vec::new(); - while let Some(row) = iter.next_tuple(|_, tuple| tuple_owned(tuple))? { + while let Some(row) = iter.next_tuple(|_, tuple| tuple.clone())? { rows.push(row); } iter.done()?; @@ -357,10 +350,11 @@ mod test { fn test_from_tuple() { let table_arena = TableArenaCell::default(); let mut plan_arena = PlanArena::new(&table_arena); - let (mut tuple, schema) = build_tuple(&mut plan_arena); + let (tuple, schema) = build_tuple(&mut plan_arena); let schema = SchemaView::new(&schema, &plan_arena); let my_struct = - ::from_query_row(&schema, &mut tuple).unwrap(); + ::from_query_row(&schema, &mut tuple.into()) + .unwrap(); println!("{:?}", my_struct); @@ -386,7 +380,7 @@ mod test { true, ColumnDesc::new(LogicalType::Integer, None, true, None).unwrap(), ))); - let mut tuple = Tuple::new( + let tuple = Tuple::new( None, vec![ DataValue::Int32(9), @@ -400,9 +394,11 @@ mod test { ); let schema = SchemaView::new(&schema, &plan_arena); - let derived = - ::from_query_row(&schema, &mut tuple) - .unwrap(); + let derived = ::from_query_row( + &schema, + &mut tuple.into(), + ) + .unwrap(); assert_eq!(derived.c1, 9); assert_eq!(derived.name, "LOL"); @@ -483,13 +479,13 @@ mod test { let balance = rows .iter() - .find(|row| match &row.values[0] { + .find(|row| match &row[0] { DataValue::Utf8 { value, .. } => value == "balance", _ => false, }) .expect("balance column should exist"); - match &balance.values[1] { + match &balance[1] { DataValue::Utf8 { value, .. } => assert_eq!(value, "Decimal(Some(10), Some(2))"), other => panic!("unexpected describe datatype: {other:?}"), } @@ -509,13 +505,13 @@ mod test { let code = rows .iter() - .find(|row| match &row.values[0] { + .find(|row| match &row[0] { DataValue::Utf8 { value, .. } => value == "code", _ => false, }) .expect("code column should exist"); - match &code.values[1] { + match &code[1] { DataValue::Utf8 { value, .. } => assert_eq!(value, "Char(2, CHARACTERS)"), other => panic!("unexpected describe datatype: {other:?}"), } @@ -547,7 +543,7 @@ mod test { let describe_rows = collect_result_tuples(database.run("describe migrating_users")?)?; let column_names = describe_rows .iter() - .filter_map(|row| match row.values.first() { + .filter_map(|row| match row.get(0) { Some(DataValue::Utf8 { value, .. }) => Some(value.as_str()), _ => None, }) @@ -1495,7 +1491,7 @@ mod test { vec![(1, Some(100)), (1, Some(200)), (2, Some(300)), (3, None)] ); - let mut tx = database.new_transaction()?; + let mut tx = database.new_transaction(false)?; let in_tx = tx .bind(|ctx| { ctx.from::()? @@ -2530,7 +2526,7 @@ mod test { collect_result_tuples(database.run("explain select age from users where age = 1050")?)?; let explain_plan = explain_rows .iter() - .filter_map(|row| match row.values.first() { + .filter_map(|row| match row.get(0) { Some(DataValue::Utf8 { value, .. }) => Some(value.as_str()), _ => None, }) @@ -2844,7 +2840,7 @@ mod test { }); assert_eq!(set_query_ages, vec![None, Some(30), Some(18), Some(18)]); - let mut tx = database.new_transaction()?; + let mut tx = database.new_transaction(false)?; tx.insert_many([User { id: 4, name: "Dora".to_string(), @@ -2950,7 +2946,7 @@ mod test { })?; assert_eq!( plan, - "Projection [users.user_name] [Project => (Sort Option: Follow)] Filter (users.id = 1), Is Having: false [Filter => (Sort Option: Follow)] TableScan users -> [users.id, users.user_name] [SeqScan => (Sort Option: None)]" + "Projection [users.user_name] [Project => (Sort Option: Follow)] TableScan users -> [users.id, users.user_name] [IndexScan By pk_index => 1 => (Sort Option: OrderBy: (users.id Asc Nulls Last) ignore_prefix_len: 0)]" ); let set_plan = database.explain(|ctx| { @@ -2965,7 +2961,7 @@ mod test { "Aggregate [] -> Group By [users.id] [HashAggregate => (Sort Option: None)] Union: [users.id] Projection [users.id] [Project => (Sort Option: Follow)] TableScan users -> [users.id] [SeqScan => (Sort Option: None)] Projection [wallets.id] [Project => (Sort Option: Follow)] TableScan wallets -> [wallets.id] [SeqScan => (Sort Option: None)]" ); - let mut tx = database.new_transaction()?; + let mut tx = database.new_transaction(false)?; let tx_tables = tx.show_tables()?.collect::, _>>()?; assert!(tx_tables.iter().any(|name| name == "users")); let tx_describe = tx.describe::()?.collect::, _>>()?; diff --git a/tests/slt/aggregation.slt b/tests/slt/aggregation.slt index 811b89c6..c9529f93 100644 --- a/tests/slt/aggregation.slt +++ b/tests/slt/aggregation.slt @@ -105,17 +105,21 @@ create table ss(id int primary key, v varchar); statement ok insert into ss values (1, '10'), (2, 'x'); -statement error -select sum(v) from ss; - -statement error -select sum(distinct v) from ss; +# Text is summed by its numeric prefix (0 without one), as in SQLite. +query RRRRR +select sum(v), sum(distinct v), avg(v), sum('a'), total(v) from ss; +---- +10.0 10.0 5.0 0.0 10.0 -statement error -select avg(v) from ss; +query RT +select total(v), group_concat(v) from ss where id > 2; +---- +0.0 NULL -statement error -select sum('a') from ss; +query TT +select group_concat(v), group_concat(distinct id % 2) from ss; +---- +10,x 1,0 statement ok drop table ss; diff --git a/tests/slt/any_all_quantifier.slt b/tests/slt/any_all_quantifier.slt index caab46ec..2b83405e 100644 --- a/tests/slt/any_all_quantifier.slt +++ b/tests/slt/any_all_quantifier.slt @@ -48,3 +48,159 @@ drop table any_all_outer; statement ok drop table any_all_inner; + +# Independent guards skip RHSs that cannot affect their own AND/OR branches. +statement ok +CREATE TABLE guard_left (x INTEGER PRIMARY KEY) + +statement ok +INSERT INTO guard_left VALUES (1), (2), (3), (4), (5), (6) + +statement ok +CREATE TABLE guard_bad (id INTEGER PRIMARY KEY, s VARCHAR) + +statement ok +INSERT INTO guard_bad VALUES (1, 'not an integer') + +# IN2 is masked by its ordinary OR sibling or by x > 3. +# IN3's AND condition is false for every input row. +# Neither CAST is allowed to execute, including after IN1 returns FALSE. +query I rowsort +SELECT x FROM guard_left WHERE x IN (SELECT x FROM guard_left WHERE x = 5) OR (x <= 3 AND (x < 4 OR x IN (SELECT CAST(s AS INTEGER) FROM guard_bad))) OR (x < 0 AND x IN (SELECT CAST(s AS INTEGER) FROM guard_bad)) +---- +1 +2 +3 +5 + +# NOT preserves the same local demand conditions and three-valued semantics. +query I rowsort +SELECT x FROM guard_left WHERE NOT (x IN (SELECT x FROM guard_left WHERE x = 5) OR (x <= 3 AND (x < 4 OR x IN (SELECT CAST(s AS INTEGER) FROM guard_bad))) OR (x < 0 AND x IN (SELECT CAST(s AS INTEGER) FROM guard_bad))) +---- +4 +6 + +statement ok +CREATE TABLE guard_flags (id INTEGER PRIMARY KEY, b BOOLEAN) + +statement ok +INSERT INTO guard_flags VALUES (1, FALSE) + +# IN1's marker is consumed by IN2, even though it is absent from the final Filter. +# It must not be replaced with NULL before the second comparison. +query I rowsort +SELECT x FROM guard_left WHERE (x IN (SELECT x FROM guard_left WHERE x = 5)) IN (SELECT b FROM guard_flags) +---- +1 +2 +3 +4 +6 + +statement ok +drop table guard_flags; + +statement ok +drop table guard_bad; + +statement ok +drop table guard_left; + +# Repeated outer rows materialize the RHS keys; more than 1024 keys spill into +# several sorted runs. Membership uses binary search, other ops scan the keys. +statement ok +create table mat_outer(id int primary key, a int) + +statement ok +create table mat_inner(id int primary key, b int) + +statement ok +create table mat_inner_null(id int primary key, b int) + +statement ok +insert into mat_outer select cast(number as int), cast(number as int) from table(numbers(2000)) + +statement ok +insert into mat_outer values (2000, null) + +statement ok +insert into mat_inner select cast(number as int), cast(number * 2 as int) from table(numbers(1500)) + +statement ok +insert into mat_inner_null select id, b from mat_inner + +statement ok +insert into mat_inner_null values (1500, null) + +statement ok +analyze table mat_outer + +statement ok +analyze table mat_inner + +statement ok +analyze table mat_inner_null + +query I +select count(*) from mat_outer where a in (select b from mat_inner) +---- +1000 + +query I +select count(*) from mat_outer where a not in (select b from mat_inner) +---- +1000 + +query I +select count(*) from mat_outer where a <> all(select b from mat_inner) +---- +1000 + +query I +select count(*) from mat_outer where a in (select b from mat_inner_null) +---- +1000 + +query I +select count(*) from mat_outer where a not in (select b from mat_inner_null) +---- +0 + +query I +select count(*) from mat_outer where (a in (select b from mat_inner_null)) is null +---- +1001 + +query I +select count(*) from mat_outer where a > any(select b from mat_inner) +---- +1999 + +query I +select count(*) from mat_outer where a <= all(select b from mat_inner) +---- +1 + +query I +select count(*) from mat_outer where a > all(select b from mat_inner_null) +---- +0 + +query I +select count(*) from mat_outer where a >= any(select b from mat_inner_null) +---- +2000 + +query I +select count(*) from mat_outer where (a <= all(select b from mat_inner_null)) is null +---- +2 + +statement ok +drop table mat_outer + +statement ok +drop table mat_inner + +statement ok +drop table mat_inner_null diff --git a/tests/slt/crdb/join.slt b/tests/slt/crdb/join.slt index ee22890c..41b5820f 100644 --- a/tests/slt/crdb/join.slt +++ b/tests/slt/crdb/join.slt @@ -1211,7 +1211,7 @@ insert into ij_c values (1,1),(2,2); query T explain select a.id,b.id from ij_a a,ij_b b where b.k=a.k and a.v>15 and b.v<250 and a.v (Sort Option: Follow)] Filter (a.v < b.v), Is Having: false [Filter => (Sort Option: Follow)] Inner Join On cast (a.k as Bigint) = b.k [HashJoin => (Sort Option: None)] Filter ((ij_a.v) as (a.v) > 15), Is Having: false [Filter => (Sort Option: Follow)] TableScan ij_a -> [ij_a.id, ij_a.k, ij_a.v] [SeqScan => (Sort Option: None)] Filter ((ij_b.v) as (b.v) < 250), Is Having: false [Filter => (Sort Option: Follow)] TableScan ij_b -> [ij_b.id, ij_b.k, ij_b.v] [SeqScan => (Sort Option: None)] +Projection [a.id, b.id] [Project => (Sort Option: Follow)] Filter (a.v < b.v), Is Having: false [Filter => (Sort Option: Follow)] Inner Join On cast (a.k as Bigint) = b.k [HashJoin => (Sort Option: None)] Projection [(ij_a.id) as (a.id), (ij_a.k) as (a.k), (ij_a.v) as (a.v)] [Project => (Sort Option: Follow)] Filter ((ij_a.v) as (a.v) > 15), Is Having: false [Filter => (Sort Option: Follow)] TableScan ij_a -> [ij_a.id, ij_a.k, ij_a.v] [SeqScan => (Sort Option: None)] Projection [(ij_b.id) as (b.id), (ij_b.k) as (b.k), (ij_b.v) as (b.v)] [Project => (Sort Option: Follow)] Filter ((ij_b.v) as (b.v) < 250), Is Having: false [Filter => (Sort Option: Follow)] TableScan ij_b -> [ij_b.id, ij_b.k, ij_b.v] [SeqScan => (Sort Option: None)] query II select a.id,b.id from ij_a a,ij_b b where b.k=a.k and a.v>15 and b.v<250 and a.v=4 and v3<=14 order by v3; +---- +4 +9 +14 + +query I +select v3 from t where v2=2 order by v3 limit 3; +---- +2 +7 +12 + +query I +select v3 from t where v2=2 and v3=7; +---- +7 + +# Fetching a non-indexed column must use the same prefix bounds as covered reads. +query I +select id from t where v2=2 and v3<=7 order by id; +---- +2 +7 + +query I +select id from t where v2>=4 and v3<=14 order by id; +---- +4 +9 +14 + +query I +select id from t where v2>2 and v3<=4 order by id; +---- +3 +4 + +query I +select id from t where v2<=2 and v3<=7 order by id; +---- +0 +1 +2 +5 +6 +7 + +query I +select id from t where v2 in (2,4) and v3<=7 order by id; +---- +2 +4 +7 + +query I +select id from t where v2=2 and v3=7; +---- +7 + +# A missing equality key must not terminate the remaining index ranges. +query I +select id from t where id in (-1,2,7,101) order by id; +---- +2 +7 + +query I +select id from t where id in (-1,2,7,101) order by id limit 1 offset 1; +---- +7 + +query I +select id from t where id=-1 or id>=99 order by id; +---- +99 +100 + +query I +select id from t where v2=2 and v3 in (-1,2,3,7) order by id; +---- +2 +7 + +query I +select id from t where v2=2 and v3 in (3,4) order by id; +---- + statement ok drop index t.index_1 diff --git a/tests/slt/cte.slt b/tests/slt/cte.slt index 7977ffb8..65bd1207 100644 --- a/tests/slt/cte.slt +++ b/tests/slt/cte.slt @@ -99,7 +99,7 @@ EXPLAIN WITH RECURSIVE recursive_cte(value) AS ( ) SELECT value FROM recursive_cte ---- -Projection [recursive_cte.value] [Project => (Sort Option: Follow)] Recursive CTE: [recursive_cte.value] Projection [(1) as (recursive_cte.value)] [Project => (Sort Option: Follow)] Projection [1] [Project => (Sort Option: Follow)] Dummy [Dummy => (Sort Option: None)] Projection [(recursive_cte.value + 1)] [Project => (Sort Option: Follow)] Filter (recursive_cte.value < 3), Is Having: false [Filter => (Sort Option: Follow)] Recursive Scan: [recursive_cte.value] +Projection [recursive_cte.value] [Project => (Sort Option: Follow)] Recursive CTE: [recursive_cte.value] Projection [(1) as (recursive_cte.value)] [Project => (Sort Option: Follow)] Dummy [Dummy => (Sort Option: None)] Projection [(recursive_cte.value + 1)] [Project => (Sort Option: Follow)] Filter (recursive_cte.value < 3), Is Having: false [Filter => (Sort Option: Follow)] Recursive Scan: [recursive_cte.value] query I WITH RECURSIVE recursive_cte(value) AS ( diff --git a/tests/slt/join.slt b/tests/slt/join.slt index c65945a9..c0f264ac 100644 --- a/tests/slt/join.slt +++ b/tests/slt/join.slt @@ -206,6 +206,41 @@ select * from a as aliased(only_one) statement error select * from a join b using (missing) +# One FROM namespace cannot bind the same visible name twice. +statement error table name: `a` specified more than once +select * from a, a + +statement error table name: `a` specified more than once +select * from a join a on true + +statement error table name: `x` specified more than once +select * from a x join b x on x.v1 = x.v3 + +statement error table name: `b` specified more than once +select * from a, a b, b + +statement error table name: `x` specified more than once +select * from a x, (select * from b) x + +statement error table name: `c` specified more than once +with c as (select * from a) select * from c, c + +# Distinct aliases and nested scopes may reuse a table. +query I +select count(*) from a x, a y +---- +9 + +query I +with c as (select * from a) select count(*) from c x, c y +---- +9 + +query I +select count(*) from a where exists (select 1 from a where a.v1 = 1) +---- +3 + # Parameterized inner join statement ok @@ -223,14 +258,28 @@ insert into lookup_outer values (1,10),(2,20),(3,null),(4,99),(5,10); statement ok insert into lookup_inner values (1,10,7),(2,10,9),(3,20,3),(4,null,8); -# A plain SeqScan outer input is eligible; no size/limit heuristic is required. +# A non-unique index stays on the ordinary join path, preserving all matches. query T explain select lookup_outer.id,lookup_inner.id from lookup_outer join lookup_inner on lookup_outer.k=lookup_inner.k; ---- -Projection [lookup_outer.id, lookup_inner.id] [Project => (Sort Option: Follow)] InnerJoinApply TableScan lookup_outer -> [lookup_outer.id, lookup_outer.k] [SeqScan => (Sort Option: None)] TableScan lookup_inner -> [lookup_inner.id, lookup_inner.k] [IndexScan By lookup_inner_k => Probe ? => (Sort Option: OrderBy: (lookup_inner.k Asc Nulls Last) ignore_prefix_len: 0)] +Projection [lookup_outer.id, lookup_inner.id] [Project => (Sort Option: Follow)] Inner Join On lookup_outer.k = lookup_inner.k [HashJoin => (Sort Option: None)] TableScan lookup_outer -> [lookup_outer.id, lookup_outer.k] [SeqScan => (Sort Option: None)] TableScan lookup_inner -> [lookup_inner.id, lookup_inner.k] [SeqScan => (Sort Option: None)] + +# A complete primary-key equality uses the single-match path. +query III +select o.id,i.id,i.v from lookup_outer o join lookup_inner i on o.id=i.id order by o.id; +---- +1 1 7 +2 2 9 +3 3 3 +4 4 8 + +# Extra equalities can reject the unique hit, including UNKNOWN from NULL. +query II +select o.id,i.id from lookup_outer o join lookup_inner i on o.id=i.id and o.k=i.k order by o.id; +---- +1 1 -# Resume the same non-unique probe across calls, then rebuild after NULL and -# missing keys. In particular, NULL = NULL must not produce the inner NULL row. +# Non-unique keys still emit every match. NULL = NULL must not match. query II select lookup_outer.id,lookup_inner.id from lookup_outer join lookup_inner on lookup_outer.k=lookup_inner.k @@ -575,6 +624,126 @@ select * from full_using_l full join full_using_r using(x) 1 2 20 5 200 NULL 3 NULL 6 300 +# Greedy reorder must preserve output slots, duplicate matches and NULL equality. +statement ok +create table greedy_a(id int primary key, k int, v int) + +statement ok +create table greedy_b(id int primary key, k bigint, v int) + +statement ok +create table greedy_c(id int primary key, k int, v int) + +statement ok +insert into greedy_a values (1,1,10),(2,2,20),(3,3,30),(4,NULL,40),(5,5,50),(6,6,60) + +statement ok +insert into greedy_b values (1,1,100),(2,2,200) + +statement ok +insert into greedy_c values (1,1,1000),(2,1,1100),(3,2,2000) + +statement ok +insert into greedy_a select cast(number+100 as int),cast(-number-1000 as int),0 from table(numbers(300)) + +statement ok +insert into greedy_b select cast(number+100 as int),-number-2000,0 from table(numbers(100)) + +statement ok +insert into greedy_c select cast(number+100 as int),cast(-number-3000 as int),0 from table(numbers(200)) + +statement ok +analyze table greedy_a + +statement ok +analyze table greedy_b + +statement ok +analyze table greedy_c + +query T +explain select a.v,b.v,c.v from greedy_a a join greedy_b b on a.k=b.k join greedy_c c on b.k=c.k +---- +Projection [a.v, b.v, c.v] [Project => (Sort Option: Follow)] Inner Join On cast (a.k as Bigint) = b.k [HashJoin => (Sort Option: None)] Projection [(greedy_a.k) as (a.k), (greedy_a.v) as (a.v)] [Project => (Sort Option: Follow)] TableScan greedy_a -> [greedy_a.k, greedy_a.v] [SeqScan => (Sort Option: None)] Inner Join On b.k = cast (c.k as Bigint) [HashJoin => (Sort Option: None)] Projection [(greedy_b.k) as (b.k), (greedy_b.v) as (b.v)] [Project => (Sort Option: Follow)] TableScan greedy_b -> [greedy_b.k, greedy_b.v] [SeqScan => (Sort Option: None)] Projection [(greedy_c.k) as (c.k), (greedy_c.v) as (c.v)] [Project => (Sort Option: Follow)] TableScan greedy_c -> [greedy_c.k, greedy_c.v] [SeqScan => (Sort Option: None)] + +query III rowsort +select a.v,b.v,c.v from greedy_a a join greedy_b b on a.k=b.k join greedy_c c on b.k=c.k +---- +10 100 1000 +10 100 1100 +20 200 2000 + +# A materializing derived input can drive a CrossJoin without changing output slots. +query III rowsort +select a.v,h.bv,h.cv from greedy_a a cross join (select b.v as bv,c.v as cv from greedy_b b join greedy_c c on b.k=c.k) h where a.id<=2 +---- +10 100 1000 +10 100 1100 +10 200 2000 +20 100 1000 +20 100 1100 +20 200 2000 + +query III +select a.v,h.bv,h.cv from greedy_a a cross join (select b.v as bv,c.v as cv from greedy_b b join greedy_c c on b.k=c.k) h where a.id<=2 order by a.v desc,h.cv desc limit 3 +---- +20 200 2000 +20 100 1100 +20 100 1000 + +# The smallest inputs b and c have no direct key; apply both keys when a joins. +query III rowsort +select a.v,b.v,c.v from greedy_a a join greedy_b b on a.k=b.k join greedy_c c on a.k=c.k +---- +10 100 1000 +10 100 1100 +20 200 2000 + +query III rowsort +select c.v,a.v,b.v from greedy_a a join greedy_b b on a.k=b.k join greedy_c c on b.k=c.k where a.v+c.v>1050 +---- +1100 10 100 +2000 20 200 + +query III rowsort +select a.v,b.v,c.v from greedy_a a left join greedy_b b on a.k=b.k left join greedy_c c on b.k=c.k where a.id<100 +---- +10 100 1000 +10 100 1100 +20 200 2000 +30 NULL NULL +40 NULL NULL +50 NULL NULL +60 NULL NULL + +# A comma join initially contains a CrossJoin without a direct a/b key. +query III rowsort +select greedy_a.v,greedy_b.v,greedy_c.v from greedy_a,greedy_b,greedy_c where greedy_a.k=greedy_c.k and greedy_b.k=greedy_c.k +---- +10 100 1000 +10 100 1100 +20 200 2000 + +# A non-unique index added after reorder must preserve the ordinary join results. +statement ok +create index greedy_b_k on greedy_b(k) + +query III rowsort +select greedy_a.v,greedy_b.v,greedy_c.v from greedy_a join greedy_b on greedy_a.k=greedy_b.k join greedy_c on greedy_b.k=greedy_c.k +---- +10 100 1000 +10 100 1100 +20 200 2000 + +statement ok +drop table greedy_a + +statement ok +drop table greedy_b + +statement ok +drop table greedy_c + statement ok drop table full_using_l @@ -589,3 +758,64 @@ drop table a statement ok drop table b + +# Greedy chooses the join implementation: a few outer rows probe a large unique key, +# a large outer side hashes a small input, and a cross join drives from the smaller side. +statement ok +create table jp_big(id int primary key, c1 int, c2 int); + +statement ok +copy jp_big from 'tests/data/row_20000.csv' ( DELIMITER '|' ); + +statement ok +create table jp_small(id int primary key, c1 int, c2 int); + +statement ok +insert into jp_small values (0,3,1),(1,6,1),(2,9,1),(3,12,1),(4,15,1),(5,18,1),(6,21,1),(7,24,1),(8,27,1),(9,30,1); + +statement ok +analyze table jp_big; + +statement ok +analyze table jp_small; + +query T +explain select jp_small.id, jp_big.c2 from jp_small join jp_big on jp_small.c1 = jp_big.id; +---- +Projection [jp_small.id, jp_big.c2] [Project => (Sort Option: Follow)] IndexProbeJoin [IndexProbeJoin => (Sort Option: Follow)] TableScan jp_small -> [jp_small.id, jp_small.c1] [SeqScan => (Sort Option: None)] TableScan jp_big -> [jp_big.id, jp_big.c2] [IndexScan By pk_index => Probe ? => (Sort Option: OrderBy: (jp_big.id Asc Nulls Last) ignore_prefix_len: 0)] + +query II rowsort +select jp_small.id, jp_big.c2 from jp_small join jp_big on jp_small.c1 = jp_big.id; +---- +0 5 +1 8 +2 11 +3 14 +4 17 +5 20 +6 23 +7 26 +8 29 +9 32 + +query T +explain select jp_big.id, jp_small.c2 from jp_big join jp_small on jp_big.c1 = jp_small.id; +---- +Projection [jp_big.id, jp_small.c2] [Project => (Sort Option: Follow)] Inner Join On jp_small.id = jp_big.c1 [HashJoin => (Sort Option: None)] TableScan jp_small -> [jp_small.id, jp_small.c2] [SeqScan => (Sort Option: None)] TableScan jp_big -> [jp_big.id, jp_big.c1] [SeqScan => (Sort Option: None)] + +query T +explain select jp_big.id, jp_small.id from jp_big, jp_small; +---- +Projection [jp_big.id, jp_small.id] [Project => (Sort Option: Follow)] Cross Join Nothing [NestLoopJoin => (Sort Option: None)] TableScan jp_small -> [jp_small.id] [SeqScan => (Sort Option: None)] TableScan jp_big -> [jp_big.id] [SeqScan => (Sort Option: None)] + +# FORCE_NEST_LOOP_JOIN regions are still reordered, but only with nested loops. +query T +explain select /*+ FORCE_NEST_LOOP_JOIN */ jp_big.id, jp_small.c2 from jp_big join jp_small on jp_big.c1 = jp_small.id; +---- +Projection [jp_big.id, jp_small.c2] [Project => (Sort Option: Follow)] Inner Join On jp_small.id = jp_big.c1 [NestLoopJoin => (Sort Option: None)] TableScan jp_small -> [jp_small.id, jp_small.c2] [SeqScan => (Sort Option: None)] TableScan jp_big -> [jp_big.id, jp_big.c1] [SeqScan => (Sort Option: None)] + +statement ok +drop table jp_big + +statement ok +drop table jp_small diff --git a/tests/slt/parameterized_subquery.slt b/tests/slt/parameterized_subquery.slt index 83c1c6e1..eaee110e 100644 --- a/tests/slt/parameterized_subquery.slt +++ b/tests/slt/parameterized_subquery.slt @@ -26,22 +26,22 @@ insert into in_inner_nn values (0, 2); query T explain select id from in_outer where a in (select v from in_inner where in_inner.v = in_outer.a); ---- -Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner.v] [Project => (Sort Option: Follow)] TableScan in_inner -> [in_inner.v] [IndexScan By in_inner_v_index => Probe ? => (Sort Option: OrderBy: (in_inner.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner.v] [Project => (Sort Option: Follow)] TableScan in_inner -> [in_inner.v] [IndexScan By in_inner_v_index => Probe ? => (Sort Option: OrderBy: (in_inner.v Asc Nulls Last) ignore_prefix_len: 0)] query T explain select id from in_outer where a not in (select v from in_inner where in_inner.v = in_outer.a); ---- -Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner.v] [Project => (Sort Option: Follow)] TableScan in_inner -> [in_inner.v] [IndexScan By in_inner_v_index => Probe ? => (Sort Option: OrderBy: (in_inner.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner.v] [Project => (Sort Option: Follow)] TableScan in_inner -> [in_inner.v] [IndexScan By in_inner_v_index => Probe ? => (Sort Option: OrderBy: (in_inner.v Asc Nulls Last) ignore_prefix_len: 0)] query T explain select id from in_outer where a in (select v from in_inner_nn where in_inner_nn.v = in_outer.a); ---- -Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner_nn.v] [Project => (Sort Option: Follow)] TableScan in_inner_nn -> [in_inner_nn.v] [IndexScan By in_inner_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_nn.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner_nn.v] [Project => (Sort Option: Follow)] TableScan in_inner_nn -> [in_inner_nn.v] [IndexScan By in_inner_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_nn.v Asc Nulls Last) ignore_prefix_len: 0)] query T explain select id from in_outer where a not in (select v from in_inner_nn where in_inner_nn.v = in_outer.a); ---- -Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner_nn.v] [Project => (Sort Option: Follow)] TableScan in_inner_nn -> [in_inner_nn.v] [IndexScan By in_inner_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_nn.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer -> [in_outer.id, in_outer.a] [SeqScan => (Sort Option: None)] Projection [in_inner_nn.v] [Project => (Sort Option: Follow)] TableScan in_inner_nn -> [in_inner_nn.v] [IndexScan By in_inner_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_nn.v Asc Nulls Last) ignore_prefix_len: 0)] query I rowsort select id from in_outer where a in (select v from in_inner where in_inner.v = in_outer.a) order by id; @@ -96,22 +96,22 @@ insert into in_inner_flag_nn values (0, 2, 1); query T explain select id from in_outer_flag where a in (select v from in_inner_flag where in_inner_flag.flag = in_outer_flag.b); ---- -Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag.v, in_inner_flag.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag -> [in_inner_flag.v, in_inner_flag.flag] [IndexScan By in_inner_flag_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag.v, in_inner_flag.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag -> [in_inner_flag.v, in_inner_flag.flag] [IndexScan By in_inner_flag_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag.v Asc Nulls Last) ignore_prefix_len: 0)] query T explain select id from in_outer_flag where a not in (select v from in_inner_flag where in_inner_flag.flag = in_outer_flag.b); ---- -Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag.v, in_inner_flag.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag -> [in_inner_flag.v, in_inner_flag.flag] [IndexScan By in_inner_flag_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag.v, in_inner_flag.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag -> [in_inner_flag.v, in_inner_flag.flag] [IndexScan By in_inner_flag_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag.v Asc Nulls Last) ignore_prefix_len: 0)] query T explain select id from in_outer_flag where a in (select v from in_inner_flag_nn where in_inner_flag_nn.flag = in_outer_flag.b); ---- -Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag_nn.v, in_inner_flag_nn.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag_nn -> [in_inner_flag_nn.v, in_inner_flag_nn.flag] [IndexScan By in_inner_flag_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag_nn.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter _temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag_nn.v, in_inner_flag_nn.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag_nn -> [in_inner_flag_nn.v, in_inner_flag_nn.flag] [IndexScan By in_inner_flag_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag_nn.v Asc Nulls Last) ignore_prefix_len: 0)] query T explain select id from in_outer_flag where a not in (select v from in_inner_flag_nn where in_inner_flag_nn.flag = in_outer_flag.b); ---- -Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag_nn.v, in_inner_flag_nn.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag_nn -> [in_inner_flag_nn.v, in_inner_flag_nn.flag] [IndexScan By in_inner_flag_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag_nn.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [in_outer_flag.id] [Project => (Sort Option: Follow)] Filter !_temp_table_1_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkAnyApply [MarkApply => (Sort Option: Follow)] TableScan in_outer_flag -> [in_outer_flag.id, in_outer_flag.a, in_outer_flag.b] [SeqScan => (Sort Option: None)] Projection [in_inner_flag_nn.v, in_inner_flag_nn.flag] [Project => (Sort Option: Follow)] TableScan in_inner_flag_nn -> [in_inner_flag_nn.v, in_inner_flag_nn.flag] [IndexScan By in_inner_flag_nn_v_index => Probe ? => (Sort Option: OrderBy: (in_inner_flag_nn.v Asc Nulls Last) ignore_prefix_len: 0)] query I rowsort select id from in_outer_flag where a in (select v from in_inner_flag where in_inner_flag.flag = in_outer_flag.b) order by id; @@ -153,12 +153,12 @@ insert into exists_inner values (0, 1, 1), (1, 1, null), (2, 2, 1); query T explain select id from exists_outer where exists (select 1 from exists_inner where exists_inner.v = exists_outer.a and exists_inner.flag = exists_outer.b); ---- -Projection [exists_outer.id] [Project => (Sort Option: Follow)] Filter _temp_table_0_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkExistsApply TableScan exists_outer -> [exists_outer.id, exists_outer.a, exists_outer.b] [SeqScan => (Sort Option: None)] TableScan exists_inner -> [exists_inner.id, exists_inner.v, exists_inner.flag] [IndexScan By exists_inner_v_index => Probe ? => (Sort Option: OrderBy: (exists_inner.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [exists_outer.id] [Project => (Sort Option: Follow)] Filter _temp_table_0_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkExistsApply [MarkApply => (Sort Option: Follow)] TableScan exists_outer -> [exists_outer.id, exists_outer.a, exists_outer.b] [SeqScan => (Sort Option: None)] TableScan exists_inner -> [exists_inner.id, exists_inner.v, exists_inner.flag] [IndexScan By exists_inner_v_index => Probe ? => (Sort Option: OrderBy: (exists_inner.v Asc Nulls Last) ignore_prefix_len: 0)] query T explain select id from exists_outer where not exists (select 1 from exists_inner where exists_inner.v = exists_outer.a and exists_inner.flag = exists_outer.b); ---- -Projection [exists_outer.id] [Project => (Sort Option: Follow)] Filter !_temp_table_0_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkExistsApply TableScan exists_outer -> [exists_outer.id, exists_outer.a, exists_outer.b] [SeqScan => (Sort Option: None)] TableScan exists_inner -> [exists_inner.id, exists_inner.v, exists_inner.flag] [IndexScan By exists_inner_v_index => Probe ? => (Sort Option: OrderBy: (exists_inner.v Asc Nulls Last) ignore_prefix_len: 0)] +Projection [exists_outer.id] [Project => (Sort Option: Follow)] Filter !_temp_table_0_.true, Is Having: false [Filter => (Sort Option: Follow)] MarkExistsApply [MarkApply => (Sort Option: Follow)] TableScan exists_outer -> [exists_outer.id, exists_outer.a, exists_outer.b] [SeqScan => (Sort Option: None)] TableScan exists_inner -> [exists_inner.id, exists_inner.v, exists_inner.flag] [IndexScan By exists_inner_v_index => Probe ? => (Sort Option: OrderBy: (exists_inner.v Asc Nulls Last) ignore_prefix_len: 0)] query I rowsort select id from exists_outer where exists (select 1 from exists_inner where exists_inner.v = exists_outer.a and exists_inner.flag = exists_outer.b) order by id; diff --git a/tests/slt/subquery.slt b/tests/slt/subquery.slt index 0e3628a8..95fa41d1 100644 --- a/tests/slt/subquery.slt +++ b/tests/slt/subquery.slt @@ -37,6 +37,29 @@ select x.a from (select -a as a from t1) as x; -1 -3 +# A derived table does not see sources bound earlier in the same FROM. +query IIIII rowsort +select * from t1 x, (select a, b from t1) y; +---- +0 1 2 1 2 +0 1 2 3 4 +1 3 4 1 2 +1 3 4 3 4 + +query I rowsort +select y.* from t1 x, (select id from t1) y; +---- +0 +0 +1 +1 + +statement error +select * from t1 x, (select a from t1 where t1.a = x.a) y; + +statement error LATERAL derived tables are not supported +select * from t1 x, lateral (select a from t1) y; + query III select * from t1 where a <= (select 4) and a > (select 1) ---- @@ -827,7 +850,7 @@ insert into corr_rows values (1,10),(2,20),(3,30),(4,null),(5,20); query T explain select id,(select 7) from corr_rows where b<(select 25); ---- -Projection [corr_rows.id, InitValue(1:0)] [Project => (Sort Option: Follow)] ScalarQueryInit InitValue(1:0) Filter (corr_rows.b < InitValue(1:1)), Is Having: false [Filter => (Sort Option: Follow)] ScalarQueryInit InitValue(1:1) TableScan corr_rows -> [corr_rows.id, corr_rows.b] [SeqScan => (Sort Option: None)] ScalarSubquery [ScalarSubquery => (Sort Option: Follow)] Projection [25] [Project => (Sort Option: Follow)] Dummy [Dummy => (Sort Option: None)] ScalarSubquery [ScalarSubquery => (Sort Option: Follow)] Projection [7] [Project => (Sort Option: Follow)] Dummy [Dummy => (Sort Option: None)] +Projection [corr_rows.id, InitValue(1:)] [Project => (Sort Option: Follow)] ScalarQueryInit InitValue(1:) Filter (corr_rows.b < InitValue(1:)), Is Having: false [Filter => (Sort Option: Follow)] ScalarQueryInit InitValue(1:) TableScan corr_rows -> [corr_rows.id, corr_rows.b] [SeqScan => (Sort Option: None)] ScalarSubquery [ScalarSubquery => (Sort Option: Follow)] Projection [25] [Project => (Sort Option: Follow)] Dummy [Dummy => (Sort Option: None)] ScalarSubquery [ScalarSubquery => (Sort Option: Follow)] Projection [7] [Project => (Sort Option: Follow)] Dummy [Dummy => (Sort Option: None)] query II rowsort select id,(select 7) from corr_rows where b<(select 25); @@ -1080,6 +1103,25 @@ select id,b in (select b from corr_rows where id=2) as matched,b not in (select 3 false false 4 NULL NULL +# Each row compares its own value with the same cached, nested RHS. +query ITTT +select id,b in (select b from corr_rows where b in (select b from corr_rows where id=2)),ball(select b from corr_rows where false) from corr_rows order by id; +---- +1 false true +2 false true +3 false true +4 false true +5 false true + statement error correlated quantified subqueries in SELECT are not supported yet select b in (select x.b from corr_rows x where x.id15 or id in (select x.id from corr_rows x where x.id=4) order by id; +---- +2 +3 +4 +5 + +# A marker visible in SELECT must remain real even when the OR branch is true. +query IT +select id,b in (select x.b from corr_rows x where x.id=1) as matched from corr_rows where b>15 or matched order by id; +---- +1 true +2 false +3 false +5 false + +# Mark definitions are attached at their first consumer, not their SELECT declaration. +query ITT +select id,b in (select x.b from corr_rows x where x.id=1) as first_match,b in (select x.b from corr_rows x where x.id=2) as second_match from corr_rows where second_match or first_match order by id; +---- +1 true false +2 false true +5 false true + +query IT +select id,b not in (select x.b from corr_rows x where x.id=1) as unmatched from corr_rows where unmatched order by id; +---- +2 true +3 true +5 true + +query IT +select id,b=all(select x.b from corr_rows x where x.id=1) as matched from corr_rows where matched is null or matched order by id; +---- +1 true +4 NULL + +# The outer marker consumes an earlier SELECT marker; neither is produced twice. +query ITT +select id,b in (select x.b from corr_rows x where x.id=1) as first_match,first_match in (select true) as second_match from corr_rows where second_match order by id; +---- +1 true true + +# A marker exposed by a derived table is read above the filter, so it cannot be guarded. +query IT +select id,m from (select id,b in (select x.b from corr_rows x where x.id=1) as m from corr_rows) where id>0 or m order by id; +---- +1 true +2 false +3 false +4 NULL +5 false + +statement error aggregate functions are not allowed in WHERE +select count(*) in (select x.id from corr_rows x) as matched from corr_rows where matched; + +statement ok +create table guard_bad_cast(id int primary key, txt varchar); + +statement ok +insert into guard_bad_cast values(1,'not an integer'); + +# A shared numeric guard protects every marker in a consecutive Apply segment. +query I +select id from corr_rows where id<0 and (b in (select cast(txt as int) from guard_bad_cast) or b in (select cast(txt as int) from guard_bad_cast)); +---- + +# An outer filter merged through a derived table also guards the inner Apply. +query I +select id from (select id,b from corr_rows where b in (select cast(txt as int) from guard_bad_cast)) x where x.id<0; +---- + +# Sort between the projection and the filter does not hide the guard. +query I +select id from corr_rows where id<0 and b in (select cast(txt as int) from guard_bad_cast) order by id; +---- + +query I +select id from corr_rows where id>0 or (b in (select cast(txt as int) from guard_bad_cast) and b in (select cast(txt as int) from guard_bad_cast)) order by id; +---- +1 +2 +3 +4 +5 + +statement error invalid digit found in string +select id from corr_rows where id>0 and b in (select cast(txt as int) from guard_bad_cast); + +statement error invalid digit found in string +select id from corr_rows where b>0 or id in (select cast(txt as int) from guard_bad_cast); + +# A scalar index probe must use the comparison's common type, not the raw +# integer literal or narrower outer column. NULL remains three-valued. +statement ok +create table probe_key(id int primary key, k bigint unique); + +statement ok +insert into probe_key values (1, 2), (2, 3), (3, 4), (4, NULL); + +query III +select 2 in (select k from probe_key), 2 not in (select k from probe_key), 9 in (select k from probe_key); +---- +true false NULL + +query I rowsort +select id from corr_rows where id in (select k from probe_key); +---- +2 +3 +4 + +# Casting the indexed side would invalidate a direct lookup. These use the +# normal comparison path rather than truncate, parse, or overflow a probe. +query III +select 2.0 in (select k from probe_key), 2.5 in (select k from probe_key), '2' in (select k from probe_key); +---- +true NULL true + +query III +select 1099511627776 in (select id from corr_rows), 1.5 in (select id from corr_rows), 'hello' in (select id from corr_rows); +---- +false false false + +statement ok +drop table probe_key; + +statement ok +drop table guard_bad_cast; + statement ok drop table corr_lookup; diff --git a/tests/slt/where.slt b/tests/slt/where.slt index 82c6ee46..0331c930 100644 --- a/tests/slt/where.slt +++ b/tests/slt/where.slt @@ -421,3 +421,544 @@ select cast(-2147483649.0 as int); statement ok drop table alias_input; + +# Name lookup order is independent of FROM layout and wildcard expansion. +statement ok +create table scope_a(id int primary key, v int); + +statement ok +create table scope_b(id int primary key, w int); + +statement ok +create table scope_c(id int primary key, v int, z int); + +statement ok +insert into scope_a values(1,10); + +statement ok +insert into scope_b values(2,10); + +statement ok +insert into scope_c values(1,20,30); + +query IIII +select * from scope_a a join scope_b b on v=b.w; +---- +1 10 2 10 + +query IIII +select b.*,a.* from scope_a a join scope_b b on v=b.w; +---- +2 10 1 10 + +query III +select a.v,w,z from scope_a a join scope_b b on v=b.w join scope_c c on a.id=c.id; +---- +10 10 30 + +statement error is ambiguous +select v from scope_a a join scope_b b on v=b.w join scope_c c on a.id=c.id; + +query III +select w,(select w+z from scope_c where scope_c.id=a.id),(select v from scope_c where scope_c.id=a.id) from scope_a a join scope_b b on v=b.w; +---- +10 40 20 + +statement ok +drop table scope_a; + +statement ok +drop table scope_b; + +statement ok +drop table scope_c; + +# Native point-list predicates cross the 5/6 expansion boundary. Repeat them +# before and after index creation so range elimination agrees with full scans. +statement ok +create table native_in_points(id int primary key, v int null, payload int); + +statement ok +insert into native_in_points values +(1,NULL,10),(2,-1,20),(3,1,30),(4,1,40),(5,2,50),(6,3,60), +(7,4,70),(8,5,80),(9,6,90),(10,7,100),(11,8,110),(12,9,120), +(13,NULL,130),(14,NULL,140),(15,NULL,150),(16,NULL,160), +(17,NULL,170),(18,NULL,180),(19,NULL,190),(20,NULL,200), +(21,NULL,210),(22,NULL,220),(23,NULL,230),(24,NULL,240), +(25,NULL,250),(26,NULL,260),(27,NULL,270),(28,NULL,280), +(29,NULL,290),(30,NULL,300),(31,NULL,310),(32,NULL,320); + +statement ok +insert into native_in_points values +(33,10,330),(34,11,340),(35,12,350),(36,13,360),(37,14,370), +(38,15,380),(39,16,390),(40,17,400),(41,18,410),(42,19,420), +(43,20,430),(44,21,440),(45,22,450),(46,23,460),(47,24,470), +(48,25,480),(49,26,490),(50,27,500),(51,28,510),(52,29,520), +(53,30,530),(54,31,540),(55,32,550),(56,33,560),(57,34,570), +(58,35,580),(59,36,590),(60,37,600),(61,38,610),(62,39,620), +(63,40,630),(64,41,640),(65,42,650),(66,43,660),(67,44,670), +(68,45,680),(69,46,690),(70,47,700),(71,48,710),(72,49,720), +(73,50,730),(74,51,740),(75,52,750),(76,53,760),(77,54,770), +(78,55,780),(79,56,790),(80,57,800),(81,58,810),(82,59,820), +(83,60,830),(84,61,840),(85,62,850),(86,63,860),(87,64,870), +(88,65,880),(89,66,890),(90,67,900),(91,68,910),(92,69,920), +(93,70,930),(94,71,940),(95,72,950),(96,73,960),(97,74,970), +(98,75,980),(99,76,990),(100,77,1000),(101,78,1010),(102,79,1020), +(103,80,1030),(104,81,1040),(105,82,1050),(106,83,1060),(107,84,1070), +(108,85,1080),(109,86,1090),(110,87,1100),(111,88,1110),(112,89,1120), +(113,90,1130),(114,91,1140),(115,92,1150),(116,93,1160),(117,94,1170), +(118,95,1180),(119,96,1190),(120,97,1200),(121,98,1210),(122,99,1220), +(123,100,1230),(124,101,1240),(125,102,1250),(126,103,1260),(127,104,1270), +(128,105,1280),(129,106,1290),(130,107,1300),(131,108,1310),(132,109,1320); + +statement ok +analyze table native_in_points; + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5); +---- +3 +4 +5 +6 +7 +8 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6); +---- +3 +4 +5 +6 +7 +8 +9 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,6) and v < 10; +---- +10 +11 +12 +2 + +query I rowsort +select id from native_in_points where v in (9,NULL,1,1,1,2); +---- +12 +3 +4 +5 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,NULL); +---- + +query I rowsort +select id from native_in_points where v in (NULL,NULL,NULL,NULL,NULL,NULL); +---- + +query I rowsort +select id from native_in_points where v not in (NULL,NULL,NULL,NULL,NULL,NULL); +---- + +query I rowsort +select id from native_in_points where v in (-999,-998,-997,6,1,6); +---- +3 +4 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) and v > 3; +---- +7 +8 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) and payload > 50; +---- +6 +7 +8 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) or payload = 10; +---- +1 +3 +4 +5 +6 +7 +8 +9 + +query I +select count(*) from native_in_points where v not in (1,2,3,4,5,6); +---- +104 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,6) limit 2 offset 1; +---- +10 +11 + +statement ok +create index native_in_v on native_in_points(v); + +statement ok +analyze table native_in_points; + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5); +---- +3 +4 +5 +6 +7 +8 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6); +---- +3 +4 +5 +6 +7 +8 +9 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,6) and v < 10; +---- +10 +11 +12 +2 + +query I rowsort +select id from native_in_points where v in (9,NULL,1,1,1,2); +---- +12 +3 +4 +5 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,NULL); +---- + +query I rowsort +select id from native_in_points where v in (NULL,NULL,NULL,NULL,NULL,NULL); +---- + +query I rowsort +select id from native_in_points where v not in (NULL,NULL,NULL,NULL,NULL,NULL); +---- + +query I rowsort +select id from native_in_points where v in (-999,-998,-997,6,1,6); +---- +3 +4 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) and v > 3; +---- +7 +8 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) and payload > 50; +---- +6 +7 +8 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) or payload = 10; +---- +1 +3 +4 +5 +6 +7 +8 +9 + +query I +select count(*) from native_in_points where v not in (1,2,3,4,5,6); +---- +104 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,6) limit 2 offset 1; +---- +10 +11 + +statement ok +drop index native_in_v on native_in_points; + +statement ok +create unique index native_in_prefix on native_in_points(v,payload); + +statement ok +analyze table native_in_points; + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5); +---- +3 +4 +5 +6 +7 +8 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6); +---- +3 +4 +5 +6 +7 +8 +9 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,6) and v < 10; +---- +10 +11 +12 +2 + +query I rowsort +select id from native_in_points where v in (9,NULL,1,1,1,2); +---- +12 +3 +4 +5 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,NULL); +---- + +query I rowsort +select id from native_in_points where v in (NULL,NULL,NULL,NULL,NULL,NULL); +---- + +query I rowsort +select id from native_in_points where v not in (NULL,NULL,NULL,NULL,NULL,NULL); +---- + +query I rowsort +select id from native_in_points where v in (-999,-998,-997,6,1,6); +---- +3 +4 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) and v > 3; +---- +7 +8 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) and payload > 50; +---- +6 +7 +8 +9 + +query I rowsort +select id from native_in_points where v in (1,2,3,4,5,6) or payload = 10; +---- +1 +3 +4 +5 +6 +7 +8 +9 + +query I +select count(*) from native_in_points where v not in (1,2,3,4,5,6); +---- +104 + +query I rowsort +select id from native_in_points where v not in (1,2,3,4,5,6) limit 2 offset 1; +---- +10 +11 + +query IIIIII rowsort +select id, v in (1,2,3,4,5), v in (1,2,3,4,5,6), + v not in (1,2,3,4,5,6), v in (9,NULL,1,1,1,2), v not in (9,NULL,1,1,1,2) +from native_in_points where id <= 12; +---- +1 NULL NULL NULL NULL NULL +10 false false true NULL NULL +11 false false true NULL NULL +12 false false true true false +2 false false true NULL NULL +3 true true false true false +4 true true false true false +5 true true false true false +6 true true false NULL NULL +7 true true false NULL NULL +8 true true false NULL NULL +9 false true false NULL NULL + +statement ok +drop table native_in_points; + +# Constant propagation: `a = b AND a = c` derives `b = c`, so the derived side can use an index. + +statement ok +create table cp_a(id int primary key, c1 int, c2 int); + +statement ok +create table cp_b(id int primary key, c1 int, c2 int); + +statement ok +create table cp_c(id int primary key, c1 int, c2 int); + +statement ok +copy cp_a from 'tests/data/row_20000.csv' ( DELIMITER '|' ); + +statement ok +copy cp_b from 'tests/data/row_20000.csv' ( DELIMITER '|' ); + +statement ok +copy cp_c from 'tests/data/row_20000.csv' ( DELIMITER '|' ); + +statement ok +create index cp_b_c1 on cp_b (c1); + +statement ok +create index cp_c_c1 on cp_c (c1); + +statement ok +analyze table cp_a; + +statement ok +analyze table cp_b; + +statement ok +analyze table cp_c; + +# Inner join: a constant on the ON key's left column reaches the right scan. +query T +explain select cp_a.id, cp_b.id from cp_a join cp_b on cp_a.c1 = cp_b.c1 where cp_a.c1 = 4; +---- +Projection [cp_a.id, cp_b.id] [Project => (Sort Option: Follow)] Inner Join On cp_b.c1 = cp_a.c1 [HashJoin => (Sort Option: None)] TableScan cp_b -> [cp_b.id, cp_b.c1] [IndexScan By cp_b_c1 => 4 => (Sort Option: OrderBy: (cp_b.c1 Asc Nulls Last) ignore_prefix_len: 0)] Filter (cp_a.c1 = 4), Is Having: false [Filter => (Sort Option: Follow)] TableScan cp_a -> [cp_a.id, cp_a.c1] [SeqScan => (Sort Option: None)] + +query II +select cp_a.id, cp_b.id from cp_a join cp_b on cp_a.c1 = cp_b.c1 where cp_a.c1 = 4; +---- +3 3 + +query T +explain select cp_a.id, cp_b.id from cp_a, cp_b where cp_a.c1 = cp_b.c1 and cp_a.c1 = 4; +---- +Projection [cp_a.id, cp_b.id] [Project => (Sort Option: Follow)] Inner Join On cp_b.c1 = cp_a.c1 [HashJoin => (Sort Option: None)] TableScan cp_b -> [cp_b.id, cp_b.c1] [IndexScan By cp_b_c1 => 4 => (Sort Option: OrderBy: (cp_b.c1 Asc Nulls Last) ignore_prefix_len: 0)] Filter (cp_a.c1 = 4), Is Having: false [Filter => (Sort Option: Follow)] TableScan cp_a -> [cp_a.id, cp_a.c1] [SeqScan => (Sort Option: None)] + +query II +select cp_a.id, cp_b.id from cp_a, cp_b where cp_a.c1 = cp_b.c1 and cp_a.c1 = 4; +---- +3 3 + +# Chains close transitively inside one filter. +query T +explain select cp_a.id, cp_b.id, cp_c.id from cp_a, cp_b, cp_c where cp_a.c1 = cp_b.c1 and cp_b.c1 = cp_c.c1 and cp_a.c1 = 4; +---- +Projection [cp_a.id, cp_b.id, cp_c.id] [Project => (Sort Option: Follow)] Inner Join On cp_c.c1 = cp_b.c1 [HashJoin => (Sort Option: None)] TableScan cp_c -> [cp_c.id, cp_c.c1] [IndexScan By cp_c_c1 => 4 => (Sort Option: OrderBy: (cp_c.c1 Asc Nulls Last) ignore_prefix_len: 0)] Inner Join On cp_b.c1 = cp_a.c1 [HashJoin => (Sort Option: None)] TableScan cp_b -> [cp_b.id, cp_b.c1] [IndexScan By cp_b_c1 => 4 => (Sort Option: OrderBy: (cp_b.c1 Asc Nulls Last) ignore_prefix_len: 0)] Filter (cp_a.c1 = 4), Is Having: false [Filter => (Sort Option: Follow)] TableScan cp_a -> [cp_a.id, cp_a.c1] [SeqScan => (Sort Option: None)] + +query III +select cp_a.id, cp_b.id, cp_c.id from cp_a, cp_b, cp_c where cp_a.c1 = cp_b.c1 and cp_b.c1 = cp_c.c1 and cp_a.c1 = 4; +---- +3 3 3 + +# Single table. +query T +explain select id from cp_b where c1 = c2 and c2 = 4; +---- +Projection [cp_b.id] [Project => (Sort Option: Follow)] Filter ((cp_b.c1 = cp_b.c2) && (cp_b.c2 = 4)), Is Having: false [Filter => (Sort Option: Follow)] TableScan cp_b -> [cp_b.id, cp_b.c1, cp_b.c2] [IndexScan By cp_b_c1 => 4 => (Sort Option: OrderBy: (cp_b.c1 Asc Nulls Last) ignore_prefix_len: 0)] + +query I +select id from cp_b where c1 = c2 and c2 = 4; +---- + + +# Left join: the preserved left side may constrain the right side. +query T +explain select cp_a.id, cp_b.id from cp_a left join cp_b on cp_a.c1 = cp_b.c1 where cp_a.c1 = 4; +---- +Projection [cp_a.id, cp_b.id] [Project => (Sort Option: Follow)] LeftOuter Join On cp_a.c1 = cp_b.c1 [HashJoin => (Sort Option: None)] Filter (cp_a.c1 = 4), Is Having: false [Filter => (Sort Option: Follow)] TableScan cp_a -> [cp_a.id, cp_a.c1] [SeqScan => (Sort Option: None)] TableScan cp_b -> [cp_b.id, cp_b.c1] [IndexScan By cp_b_c1 => 4 => (Sort Option: OrderBy: (cp_b.c1 Asc Nulls Last) ignore_prefix_len: 0)] + +query II +select cp_a.id, cp_b.id from cp_a left join cp_b on cp_a.c1 = cp_b.c1 where cp_a.c1 = 4; +---- +3 3 + +# Left join: a right-side constant must not filter the preserved left side. +query T +explain select count(*), count(cp_b.id) from cp_a left join cp_b on cp_a.c1 = cp_b.c1 and cp_b.c1 = 4; +---- +Projection [Count(*), Count(cp_b.id)] [Project => (Sort Option: Follow)] Aggregate [Count(*), Count(cp_b.id)] [SimpleAggregate => (Sort Option: None)] LeftOuter Join On cp_a.c1 = cp_b.c1 [HashJoin => (Sort Option: None)] TableScan cp_a -> [cp_a.c1] [SeqScan => (Sort Option: None)] TableScan cp_b -> [cp_b.id, cp_b.c1] [IndexScan By cp_b_c1 => 4 => (Sort Option: OrderBy: (cp_b.c1 Asc Nulls Last) ignore_prefix_len: 0)] + +query II +select count(*), count(cp_b.id) from cp_a left join cp_b on cp_a.c1 = cp_b.c1 and cp_b.c1 = 4; +---- +20000 1 + +# Right join: a left-side constant must not filter the preserved right side. +query T +explain select count(*), count(cp_a.id) from cp_a right join cp_b on cp_a.c1 = cp_b.c1 and cp_a.c1 = 4; +---- +Projection [Count(*), Count(cp_a.id)] [Project => (Sort Option: Follow)] Aggregate [Count(*), Count(cp_a.id)] [SimpleAggregate => (Sort Option: None)] RightOuter Join On cp_a.c1 = cp_b.c1 [HashJoin => (Sort Option: None)] Filter (cp_a.c1 = 4), Is Having: false [Filter => (Sort Option: Follow)] TableScan cp_a -> [cp_a.id, cp_a.c1] [SeqScan => (Sort Option: None)] TableScan cp_b -> [cp_b.c1] [SeqScan => (Sort Option: None)] + +query II +select count(*), count(cp_a.id) from cp_a right join cp_b on cp_a.c1 = cp_b.c1 and cp_a.c1 = 4; +---- +20000 1 + +# NULL constants carry nothing to propagate. +query T +explain select count(*) from cp_a join cp_b on cp_a.c1 = cp_b.c1 where cp_a.c1 = null; +---- +Projection [Count(*)] [Project => (Sort Option: Follow)] Aggregate [Count(*)] [SimpleAggregate => (Sort Option: None)] Inner Join On cp_a.c1 = cp_b.c1 [HashJoin => (Sort Option: None)] Filter (cp_a.c1 = cast (NULL as Integer)), Is Having: false [Filter => (Sort Option: Follow)] TableScan cp_a -> [cp_a.c1] [SeqScan => (Sort Option: None)] TableScan cp_b -> [cp_b.c1] [SeqScan => (Sort Option: None)] + +query I +select count(*) from cp_a join cp_b on cp_a.c1 = cp_b.c1 where cp_a.c1 = null; +---- +0 + +statement ok +drop table cp_a; + +statement ok +drop table cp_b; + +statement ok +drop table cp_c; diff --git a/tests/slt/where_by_index.slt b/tests/slt/where_by_index.slt index 02d78e9b..dac9b7df 100644 --- a/tests/slt/where_by_index.slt +++ b/tests/slt/where_by_index.slt @@ -362,6 +362,17 @@ select k from tuple_bounds where w = 1 and k >= 1 and k < 3 order by k; 1 2 +# Scalar initialization must not block the static index prefix or erase its residual. +query T +explain select k from tuple_bounds where w = 1 and k = (select max(k) from tuple_bounds where w = 1); +---- +Projection [tuple_bounds.k] [Project => (Sort Option: Follow)] ScalarQueryInit InitValue(1:) TableScan tuple_bounds -> [tuple_bounds.w, tuple_bounds.k] [IndexScan By pk_index => (1, InitValue(1:)) Covered => (Sort Option: OrderBy: (tuple_bounds.w Asc Nulls Last, tuple_bounds.k Asc Nulls Last) ignore_prefix_len: 2)] ScalarSubquery [ScalarSubquery => (Sort Option: Follow)] Projection [Max(tuple_bounds.k)] [Project => (Sort Option: Follow)] Aggregate [Max(tuple_bounds.k)] [SimpleAggregate => (Sort Option: None)] Top 1, Sort By tuple_bounds.k Desc Nulls Last [TopK => (Sort Option: OrderBy: (tuple_bounds.k Desc Nulls Last) ignore_prefix_len: 0)] TableScan tuple_bounds -> [tuple_bounds.w, tuple_bounds.k] [IndexScan By pk_index => [(1), (1)] Covered => (Sort Option: OrderBy: (tuple_bounds.w Asc Nulls Last, tuple_bounds.k Asc Nulls Last) ignore_prefix_len: 1)] + +query I +select k from tuple_bounds where w = 1 and k = (select max(k) from tuple_bounds where w = 1); +---- +3 + statement ok drop table tuple_bounds; diff --git a/tests/slt/where_by_index_explain.slt b/tests/slt/where_by_index_explain.slt index b24ada38..fc1f0452 100644 --- a/tests/slt/where_by_index_explain.slt +++ b/tests/slt/where_by_index_explain.slt @@ -23,7 +23,7 @@ analyze table t1; query T explain select a.c2, b.c2 from t1 a join t1 b on a.id + 3 = b.id where a.id = 3 and b.id = 6; ---- -Projection [a.c2, b.c2] [Project => (Sort Option: Follow)] Inner Join Where ((a.id + 3) = b.id) [NestLoopJoin => (Sort Option: None)] TableScan t1 -> [t1.id, t1.c2] [IndexScan By pk_index => 3 => (Sort Option: OrderBy: (t1.id Asc Nulls Last) ignore_prefix_len: 0)] Filter ((t1.id) as (b.id) = 6), Is Having: false [Filter => (Sort Option: Follow)] TableScan t1 -> [t1.id, t1.c2] [IndexScan By pk_index => 6 => (Sort Option: OrderBy: (t1.id Asc Nulls Last) ignore_prefix_len: 0)] +Projection [a.c2, b.c2] [Project => (Sort Option: Follow)] Inner Join Where ((a.id + 3) = b.id) [NestLoopJoin => (Sort Option: None)] Projection [(t1.id) as (a.id), (t1.c2) as (a.c2)] [Project => (Sort Option: Follow)] TableScan t1 -> [t1.id, t1.c2] [IndexScan By pk_index => 3 => (Sort Option: OrderBy: (t1.id Asc Nulls Last) ignore_prefix_len: 0)] Projection [(t1.id) as (b.id), (t1.c2) as (b.c2)] [Project => (Sort Option: Follow)] TableScan t1 -> [t1.id, t1.c2] [IndexScan By pk_index => 6 => (Sort Option: OrderBy: (t1.id Asc Nulls Last) ignore_prefix_len: 0)] query II select a.c2, b.c2 from t1 a join t1 b on a.id + 3 = b.id where a.id = 3 and b.id = 6; diff --git a/tests/sqllogictest/src/bin/sqlite.rs b/tests/sqllogictest/src/bin/sqlite.rs new file mode 100644 index 00000000..5f9419fd --- /dev/null +++ b/tests/sqllogictest/src/bin/sqlite.rs @@ -0,0 +1,317 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Runs the SQLite sqllogictest corpus converted by `kitesql-testing/scripts/convert.py`. +//! +//! Unlike the main harness, every file gets its own database, values are rendered the way the +//! SQLite reference runner does (by the column type declared in the query), queries that exceed +//! the parser recursion limit are skipped, and failures are counted instead of aborting the run. + +use clap::Parser; +use kite_sql::db::{DataBaseBuilder, Database}; +use kite_sql::errors::DatabaseError; +use kite_sql::storage::lmdb::LmdbStorage; +use kite_sql::types::value::DataValue; +use sqllogictest::{DBOutput, DefaultColumnType, QueryExpect, Record, Runner, DB}; +use sqllogictest_test::run_sql; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; +use tempfile::TempDir; + +const STACK_SIZE: usize = 64 * 1024 * 1024; +const MAX_ERROR_LINES: usize = 30; + +#[derive(Parser, Debug)] +struct Args { + #[clap(long, default_value = "tests/kitesql-testing/data/sqlite/**/*.slt")] + path: String, + /// Failures printed per file; the rest are only counted. + #[clap(long, default_value_t = 5)] + max_errors: usize, + /// Print the location of every record before running it, to find a hanging query. + #[clap(long)] + trace: bool, + /// Include select1ā€“5, which are excluded while correlated scalar queries and join costs are investigated. + #[clap(long)] + include_select: bool, + /// Override the ANALYZE bucket count for small corpus fixtures. + #[clap(long)] + histogram_buckets: Option, + /// Override SQL parser recursion depth for deeply nested corpus queries. + #[clap(long)] + parser_recursion_limit: Option, + /// Override the HyperLogLog precision (4..=26) used by ANALYZE for NDV. + #[clap(long)] + hll_precision: Option, + /// Override the relative error (0, 1) of ANALYZE's KLL and Count-Min sketches. + #[clap(long)] + statistics_relative_error: Option, +} + +/// Database settings shared by every file in a run. +#[derive(Clone, Copy)] +struct DbOptions { + histogram_buckets: Option, + parser_recursion_limit: Option, + hll_precision: Option, + statistics_relative_error: Option, +} + +impl DbOptions { + fn builder(self, path: PathBuf) -> DataBaseBuilder { + // Always NO_SYNC, matching the reference SQLite runner's `PRAGMA synchronous=OFF`; + // otherwise every autocommit statement pays an fdatasync that dominates the run. + let mut builder = DataBaseBuilder::path(path).lmdb_no_sync(true); + if let Some(buckets) = self.histogram_buckets { + builder = builder.histogram_buckets(buckets); + } + if let Some(limit) = self.parser_recursion_limit { + builder = builder.parser_recursion_limit(limit); + } + if let Some(precision) = self.hll_precision { + builder = builder.hll_precision(precision); + } + if let Some(relative_error) = self.statistics_relative_error { + builder = builder.statistics_relative_error(relative_error); + } + builder + } +} + +struct SqliteDB { + db: Database, + /// Column types declared by the query being run. + types: Arc>>, +} + +impl DB for SqliteDB { + type Error = DatabaseError; + type ColumnType = DefaultColumnType; + + fn run(&mut self, sql: &str) -> Result, Self::Error> { + let types = self.types.lock().unwrap().clone(); + run_sql(&mut self.db, sql, |i, value| { + format_value(types.get(i), value) + }) + } +} + +/// Renders like the SQLite reference runner: `I` as `%d`, `R` as `%.3f`, NULL as `NULL`, an +/// empty string as `(empty)` and unprintable characters as `@`. +fn format_value(ty: Option<&DefaultColumnType>, value: &DataValue) -> String { + if value.is_null() { + return "NULL".to_string(); + } + match (ty, value) { + (Some(DefaultColumnType::Integer), DataValue::Boolean(b)) => (*b as i64).to_string(), + ( + Some(DefaultColumnType::Integer), + DataValue::Int8(_) + | DataValue::Int16(_) + | DataValue::Int32(_) + | DataValue::Int64(_) + | DataValue::UInt8(_) + | DataValue::UInt16(_) + | DataValue::UInt32(_) + | DataValue::UInt64(_), + ) => value.to_string(), + (Some(DefaultColumnType::Integer), _) => (as_f64(value) as i64).to_string(), + (Some(DefaultColumnType::FloatingPoint), _) => format!("{:.3}", as_f64(value)), + _ => { + let text = value.to_string(); + if text.is_empty() { + return "(empty)".to_string(); + } + text.chars() + .map(|c| if (' '..='~').contains(&c) { c } else { '@' }) + .collect() + } + } +} + +fn as_f64(value: &DataValue) -> f64 { + match value { + DataValue::Boolean(b) => *b as i64 as f64, + DataValue::Float32(f) => f.0 as f64, + DataValue::Float64(f) => f.0, + DataValue::Utf8 { value, .. } => value.trim().parse().unwrap_or(0.0), + other => other.to_string().parse().unwrap_or(0.0), + } +} + +#[derive(Default, Clone, Copy)] +struct Summary { + passed: usize, + failed: usize, + skipped: usize, + /// A statement failed, so the rest of the file was not run. + halted: bool, +} + +fn run_file( + path: &Path, + max_errors: usize, + trace: bool, + options: DbOptions, +) -> Result { + let script = std::fs::read_to_string(path).map_err(|e| e.to_string())?; + let records = + sqllogictest::parse_with_name::(&script, path.display().to_string()) + .map_err(|e| e.to_string())?; + let dir = TempDir::new().map_err(|e| e.to_string())?; + let db_path = dir.path().to_path_buf(); + let types = Arc::new(Mutex::new(Vec::new())); + let db_types = types.clone(); + let mut runner = Runner::new(move || { + let types = db_types.clone(); + let db = options.builder(db_path.clone()).build_lmdb(); + std::future::ready(db.map(|db| SqliteDB { db, types })) + }); + // The corpus was generated with this threshold, even in files that never declare it. + runner.with_hash_threshold(8); + + let mut summary = Summary::default(); + for record in records { + let is_statement = match &record { + Record::Statement { loc, .. } => { + if trace { + println!(" run {loc}"); + } + true + } + Record::Query { loc, expected, .. } => { + if trace { + println!(" run {loc}"); + } + if let QueryExpect::Results { + types: expected, .. + } = expected + { + *types.lock().unwrap() = expected.clone(); + } + false + } + Record::Halt { .. } => break, + _ => { + runner.run(record).map_err(|e| e.to_string())?; + continue; + } + }; + let Err(err) = runner.run(record) else { + summary.passed += 1; + continue; + }; + let message = err.to_string(); + if message.contains("recursion limit exceeded") { + summary.skipped += 1; + continue; + } + summary.failed += 1; + if summary.failed <= max_errors { + let lines: Vec<_> = message.lines().take(MAX_ERROR_LINES).collect(); + println!(" {}", lines.join("\n ")); + } + if is_statement { + summary.halted = true; + break; + } + } + Ok(summary) +} + +fn main() { + let args = Args::parse(); + std::env::set_current_dir(Path::new(env!("CARGO_MANIFEST_DIR")).join("../..")).unwrap(); + let files: Vec = glob::glob(&args.path) + .expect("invalid --path pattern") + .filter_map(Result::ok) + .collect(); + if files.is_empty() { + eprintln!( + "no file matches {} (run `git submodule update --init`?)", + args.path + ); + std::process::exit(2); + } + + let start = Instant::now(); + let (mut total, mut broken, mut excluded) = (Summary::default(), 0, 0); + for path in files { + if !args.include_select + && matches!( + path.file_stem().and_then(|name| name.to_str()), + Some("select1" | "select2" | "select3" | "select4" | "select5") + ) + { + println!( + "-> {}: excluded (select1ā€“5; use --include-select to run)", + path.display() + ); + excluded += 1; + continue; + } + println!("-> {}", path.display()); + let file_start = Instant::now(); + let (file, max_errors, trace) = (path.clone(), args.max_errors, args.trace); + let options = DbOptions { + histogram_buckets: args.histogram_buckets, + parser_recursion_limit: args.parser_recursion_limit, + hll_precision: args.hll_precision, + statistics_relative_error: args.statistics_relative_error, + }; + let result = std::thread::Builder::new() + .stack_size(STACK_SIZE) + .spawn(move || run_file(&file, max_errors, trace, options)) + .expect("failed to spawn test thread") + .join(); + match result { + Ok(Ok(s)) => { + let halted = if s.halted { + ", halted on a failed statement" + } else { + "" + }; + println!( + " passed {} failed {} skipped {}{halted} ({:?})", + s.passed, + s.failed, + s.skipped, + file_start.elapsed() + ); + total.passed += s.passed; + total.failed += s.failed; + total.skipped += s.skipped; + } + Ok(Err(err)) => { + broken += 1; + println!(" error: {err}"); + } + Err(_) => { + broken += 1; + println!(" panicked"); + } + } + } + println!( + "total: passed {} failed {} skipped {}, {broken} file(s) not run, {excluded} file(s) excluded ({:?})", + total.passed, + total.failed, + total.skipped, + start.elapsed() + ); + if total.failed > 0 || broken > 0 { + std::process::exit(1); + } +} diff --git a/tests/sqllogictest/src/lib.rs b/tests/sqllogictest/src/lib.rs index dc125001..852b3130 100644 --- a/tests/sqllogictest/src/lib.rs +++ b/tests/sqllogictest/src/lib.rs @@ -49,7 +49,6 @@ pub fn run_sql( let mut rows = Vec::new(); while let Some(row) = iter.next_tuple(|_, tuple| { tuple - .values .iter() .enumerate() .map(|(i, value)| format(i, value)) diff --git a/tpcc/README.md b/tpcc/README.md index b60481d6..ca9936e1 100644 --- a/tpcc/README.md +++ b/tpcc/README.md @@ -50,15 +50,14 @@ Local 720-second comparison on the machine above: | Backend | TpmC | New-Order p90 (Āµs) | Payment p90 (Āµs) | Order-Status p90 (Āµs) | Delivery p90 (Āµs) | Stock-Level p90 (Āµs) | | --- | ---: | ---: | ---: | ---: | ---: | ---: | -| KiteSQL LMDB | 159616 | 315 | 81 | 44 | 435 | 332 | -| KiteSQL RocksDB | 48035 | 459 | 258 | 200 | 11055 | 704 | -| SQLite balanced | 67102 | 283 | 67 | 48 | 338 | 473 | -| SQLite practical | 67153 | 336 | 66 | 41 | 427 | 339 | +| KiteSQL LMDB | 165956 | 303 | 75 | 43 | 436 | 318 | +| KiteSQL RocksDB | 50739 | 392 | 188 | 198 | 11743 | 812 | +| SQLite | 191484 | 274 | 41 | 33 | 343 | 268 | -- Run dates: `2026-10-01`; results: `2026-10-01_17-12-31`. Latency is measured in microseconds and includes commit. +- Run dates: `2026-10-25`; results: `2026-10-25_16-30-15`. Latency is measured in microseconds and includes commit. - Every row is pinned to `cpu8`, the highest-frequency P-core on this machine (5.4 GHz). - All rows use `--num-ware 1`, `--max-retry 5`, and TPCC's default 720-second measure time. -- SQLite rows use the `balanced` and `practical` profiles respectively. +- All rows run without `--sync`, so commits are not fsynced; SQLite uses WAL with `synchronous = OFF`. ## Historical raw outputs ā€” 2026-07-11 diff --git a/tpcc/src/backend/dual.rs b/tpcc/src/backend/dual.rs index 959357c3..f87b4223 100644 --- a/tpcc/src/backend/dual.rs +++ b/tpcc/src/backend/dual.rs @@ -19,7 +19,7 @@ use super::{ SimpleExecutor, StatementSpec, }; use crate::TpccError; -use kite_sql::types::tuple::Tuple; +use kite_sql::types::tuple::TupleImpl; use std::borrow::Cow; pub struct DualBackend { @@ -135,7 +135,7 @@ impl<'a> BackendTransaction for DualTransaction<'a> { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { self.with_query_nth(statement, params, 0, visitor) } @@ -144,7 +144,7 @@ impl<'a> BackendTransaction for DualTransaction<'a> { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut rows = Vec::new(); self.kitesql @@ -155,12 +155,12 @@ impl<'a> BackendTransaction for DualTransaction<'a> { let mut sqlite_rows = Vec::new(); self.sqlite .with_query_all(&mut statement.sqlite, params, &mut |tuple| { - sqlite_rows.push(tuple.values.clone()); + sqlite_rows.push(tuple.clone()); Ok(()) })?; // SQL without ORDER BY may return the same rows in different orders. for row in &rows { - let Some(index) = sqlite_rows.iter().position(|values| *values == row.values) else { + let Some(index) = sqlite_rows.iter().position(|values| row.values_eq(values)) else { return Err(TpccError::BackendMismatch(format!( "Result mismatch for SQL: {}", statement.spec.sql @@ -185,7 +185,7 @@ impl<'a> BackendTransaction for DualTransaction<'a> { statement: &mut Self::PreparedStatement, params: &[DbParam], n: usize, - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let spec = statement.spec.clone(); @@ -223,17 +223,17 @@ fn drain_sqlite_iter(mut iter: SqliteResult<'_, '_>) -> Result<(), TpccError> { Ok(()) } -fn drain_kitesql_iter( +fn drain_kitesql_iter( mut iter: KiteSqlTxnResult<'_, T>, ) -> Result<(), TpccError> { while iter.skip_next_tuple()? {} Ok(()) } -fn with_kitesql_nth( +fn with_kitesql_nth( mut iter: KiteSqlTxnResult<'_, T>, n: usize, - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { for _ in 0..n { if !iter.skip_next_tuple()? { @@ -244,12 +244,12 @@ fn with_kitesql_nth( .ok_or(TpccError::EmptyTuples) } -fn with_ordered_nth( +fn with_ordered_nth( mut kitesql_iter: KiteSqlTxnResult<'_, T>, mut sqlite_iter: SqliteResult<'_, '_>, sql: &'static str, n: usize, - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut visited = false; let mut index = 0usize; @@ -265,9 +265,12 @@ fn with_ordered_nth( ))) } }; - if kitesql_tuple.values != sqlite_tuple.values { + if !kitesql_tuple.iter().eq(&sqlite_tuple.values) { println!("[Dual] mismatch SQL: {}", sql); - println!(" KiteSQL row: {:?}", kitesql_tuple.values); + println!( + " KiteSQL row: {:?}", + kitesql_tuple.iter().collect::>() + ); println!(" SQLite row: {:?}", sqlite_tuple.values); return Err(TpccError::BackendMismatch(format!( "Result mismatch for SQL: {}", @@ -275,7 +278,7 @@ fn with_ordered_nth( ))); } if index == n { - visitor(&sqlite_tuple)?; + visitor(&sqlite_tuple.into())?; visited = true; } index += 1; @@ -300,7 +303,7 @@ fn with_ordered_nth( } } -fn drain_and_compare_ordered( +fn drain_and_compare_ordered( mut kitesql_iter: KiteSqlTxnResult<'_, T>, mut sqlite_iter: SqliteResult<'_, '_>, sql: &'static str, @@ -316,9 +319,12 @@ fn drain_and_compare_ordered( ))) } }; - if kitesql_tuple.values != sqlite_tuple.values { + if !kitesql_tuple.iter().eq(&sqlite_tuple.values) { println!("[Dual] mismatch SQL: {}", sql); - println!(" KiteSQL row: {:?}", kitesql_tuple.values); + println!( + " KiteSQL row: {:?}", + kitesql_tuple.iter().collect::>() + ); println!(" SQLite row: {:?}", sqlite_tuple.values); return Err(TpccError::BackendMismatch(format!( "Result mismatch for SQL: {}", diff --git a/tpcc/src/backend/kitesql_lmdb.rs b/tpcc/src/backend/kitesql_lmdb.rs index 712ff077..892aebd2 100644 --- a/tpcc/src/backend/kitesql_lmdb.rs +++ b/tpcc/src/backend/kitesql_lmdb.rs @@ -18,9 +18,9 @@ use super::{ StatementSpec, }; use crate::TpccError; -use kite_sql::db::{DBTransaction, DataBaseBuilder, Database}; +use kite_sql::db::{DataBaseBuilder, Database, Transaction}; use kite_sql::storage::lmdb::{LmdbStorage, LmdbTransaction as KiteSqlLmdbTransaction}; -use kite_sql::types::tuple::Tuple; +use kite_sql::types::tuple::TupleImpl; pub struct KiteSqlLmdbBackend { database: Database, @@ -53,7 +53,7 @@ impl KiteSqlLmdbBackend { fn start_transaction(&self) -> Result, TpccError> { Ok(KiteSqlLmdbTransactionWrapper { - inner: self.database.new_transaction()?, + inner: self.database.new_transaction(false)?, }) } } @@ -88,7 +88,7 @@ impl SimpleExecutor for KiteSqlLmdbBackend { } pub struct KiteSqlLmdbTransactionWrapper<'a> { - inner: DBTransaction<'a, LmdbStorage>, + inner: Transaction<'a, LmdbStorage>, } impl<'a> KiteSqlLmdbTransactionWrapper<'a> { @@ -120,7 +120,7 @@ impl<'a> BackendTransaction for KiteSqlLmdbTransactionWrapper<'a> { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; iter.with_next_tuple(|tuple| visitor(tuple))? @@ -131,7 +131,7 @@ impl<'a> BackendTransaction for KiteSqlLmdbTransactionWrapper<'a> { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; while iter.with_next_tuple(|tuple| visitor(tuple))?.is_some() {} @@ -143,7 +143,7 @@ impl<'a> BackendTransaction for KiteSqlLmdbTransactionWrapper<'a> { statement: &mut Self::PreparedStatement, params: &[DbParam], n: usize, - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; for _ in 0..n { diff --git a/tpcc/src/backend/kitesql_rocksdb.rs b/tpcc/src/backend/kitesql_rocksdb.rs index f052d5d6..ed03f0bc 100644 --- a/tpcc/src/backend/kitesql_rocksdb.rs +++ b/tpcc/src/backend/kitesql_rocksdb.rs @@ -18,10 +18,10 @@ use super::{ }; use crate::TpccError; use kite_sql::binder::{command_type, CommandType}; -use kite_sql::db::{DBTransaction, DataBaseBuilder, Database, Statement, TransactionIter}; +use kite_sql::db::{DataBaseBuilder, Database, Statement, Transaction, TransactionIter}; use kite_sql::storage::rocksdb::{OptimisticRocksStorage, RocksStorage}; -use kite_sql::storage::{Storage, Transaction}; -use kite_sql::types::tuple::Tuple; +use kite_sql::storage::{Storage, StorageTransaction}; +use kite_sql::types::tuple::TupleImpl; pub type KiteSqlRocksDbBackend = KiteSqlRocksBackend; pub type KiteSqlOptimisticRocksDbBackend = KiteSqlRocksBackend; @@ -74,7 +74,7 @@ impl KiteSqlRocksBackend { fn start_transaction(&self) -> Result, TpccError> { Ok(KiteSqlRocksTransaction { - inner: self.database.new_transaction()?, + inner: self.database.new_transaction(false)?, }) } } @@ -154,14 +154,14 @@ pub(crate) fn execute_kitesql_batch( return Ok(()); } - let mut transaction = database.new_transaction()?; + let mut transaction = database.new_transaction(false)?; transaction.run(sql)?.done()?; transaction.commit()?; Ok(()) } pub struct KiteSqlRocksTransaction<'a, S: Storage> { - inner: DBTransaction<'a, S>, + inner: Transaction<'a, S>, } impl<'a, S: Storage> KiteSqlRocksTransaction<'a, S> { @@ -193,7 +193,7 @@ impl<'a, S: Storage> BackendTransaction for KiteSqlRocksTransaction<'a, S> { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; iter.with_next_tuple(|tuple| visitor(tuple))? @@ -204,7 +204,7 @@ impl<'a, S: Storage> BackendTransaction for KiteSqlRocksTransaction<'a, S> { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; while iter.with_next_tuple(|tuple| visitor(tuple))?.is_some() {} @@ -216,7 +216,7 @@ impl<'a, S: Storage> BackendTransaction for KiteSqlRocksTransaction<'a, S> { statement: &mut Self::PreparedStatement, params: &[DbParam], n: usize, - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; for _ in 0..n { @@ -234,9 +234,9 @@ impl<'a, S: Storage> BackendTransaction for KiteSqlRocksTransaction<'a, S> { } } -pub struct KiteSqlTxnResult<'a, T: Transaction + 'a>(TransactionIter<'a, T>); +pub struct KiteSqlTxnResult<'a, T: StorageTransaction + 'a>(TransactionIter<'a, T>); -impl<'a, T: Transaction + 'a> KiteSqlTxnResult<'a, T> { +impl<'a, T: StorageTransaction + 'a> KiteSqlTxnResult<'a, T> { pub(crate) fn new(iter: TransactionIter<'a, T>) -> Self { Self(iter) } @@ -251,7 +251,7 @@ impl<'a, T: Transaction + 'a> KiteSqlTxnResult<'a, T> { pub(crate) fn with_next_tuple( &mut self, - f: impl FnOnce(&Tuple) -> Result, + f: impl FnOnce(&TupleImpl) -> Result, ) -> Result, TpccError> { match self .0 diff --git a/tpcc/src/backend/mod.rs b/tpcc/src/backend/mod.rs index 09688b0a..fbb13314 100644 --- a/tpcc/src/backend/mod.rs +++ b/tpcc/src/backend/mod.rs @@ -19,7 +19,7 @@ pub mod sqlite; use crate::TpccError; use kite_sql::db::PreparedPlan; -use kite_sql::types::tuple::Tuple; +use kite_sql::types::tuple::TupleImpl; use kite_sql::types::value::DataValue; use kite_sql::types::LogicalType; @@ -62,14 +62,14 @@ pub trait BackendTransaction { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError>; fn with_query_all( &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError>; fn with_query_nth( @@ -77,7 +77,7 @@ pub trait BackendTransaction { statement: &mut Self::PreparedStatement, params: &[DbParam], n: usize, - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError>; fn commit(self) -> Result<(), TpccError>; diff --git a/tpcc/src/backend/sqlite.rs b/tpcc/src/backend/sqlite.rs index c2b56eaa..a25f261a 100644 --- a/tpcc/src/backend/sqlite.rs +++ b/tpcc/src/backend/sqlite.rs @@ -17,7 +17,7 @@ use super::{ }; use crate::TpccError; use chrono::{NaiveDateTime, TimeZone, Utc}; -use kite_sql::types::tuple::Tuple; +use kite_sql::types::tuple::{Tuple, TupleImpl}; use kite_sql::types::value::{DataValue, Utf8Type}; use kite_sql::types::{CharLengthUnits, LogicalType}; use rust_decimal::Decimal; @@ -156,22 +156,22 @@ impl<'a> BackendTransaction for SqliteTransaction<'a> { &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; let tuple = iter.next().transpose()?.ok_or(TpccError::EmptyTuples)?; - visitor(&tuple) + visitor(&tuple.into()) } fn with_query_all( &mut self, statement: &mut Self::PreparedStatement, params: &[DbParam], - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let iter = self.execute_raw(statement, params)?; for row in iter { - visitor(&row?)?; + visitor(&row?.into())?; } Ok(()) } @@ -181,7 +181,7 @@ impl<'a> BackendTransaction for SqliteTransaction<'a> { statement: &mut Self::PreparedStatement, params: &[DbParam], n: usize, - visitor: &mut dyn FnMut(&Tuple) -> Result<(), TpccError>, + visitor: &mut dyn FnMut(&TupleImpl) -> Result<(), TpccError>, ) -> Result<(), TpccError> { let mut iter = self.execute_raw(statement, params)?; for _ in 0..n { @@ -190,7 +190,7 @@ impl<'a> BackendTransaction for SqliteTransaction<'a> { } } let tuple = iter.next().transpose()?.ok_or(TpccError::EmptyTuples)?; - visitor(&tuple) + visitor(&tuple.into()) } fn commit(mut self) -> Result<(), TpccError> { @@ -209,7 +209,9 @@ fn bind_params(statement: &mut SqliteStatement<'_>, params: &[DbParam]) -> Resul fn convert_value(value: &DataValue) -> Result { Ok(match value { - DataValue::Parameter { .. } => return Err(TpccError::InvalidParameter), + DataValue::Parameter { .. } | DataValue::InitValue { .. } => { + return Err(TpccError::InvalidParameter) + } DataValue::Null => Value::Null, DataValue::Boolean(v) => Value::Integer(*v as i64), DataValue::Float32(v) => Value::Float(v.0 as f64), diff --git a/tpcc/src/delivery.rs b/tpcc/src/delivery.rs index 23fd022c..86e3a4a3 100644 --- a/tpcc/src/delivery.rs +++ b/tpcc/src/delivery.rs @@ -55,7 +55,7 @@ impl TpccTransaction for Delivery { (2, DataValue::Int16(args.w_id as i16)), ], &mut |tuple| { - no_o_id = tuple.values[0].i32().unwrap(); + no_o_id = tuple[0].i32().unwrap(); Ok(()) }, )?; @@ -82,7 +82,7 @@ impl TpccTransaction for Delivery { (3, DataValue::Int16(args.w_id as i16)), ], &mut |tuple| { - c_id = tuple.values[0].i32().unwrap(); + c_id = tuple[0].i32().unwrap(); Ok(()) }, )?; @@ -116,7 +116,7 @@ impl TpccTransaction for Delivery { (3, DataValue::Int16(args.w_id as i16)), ], &mut |tuple| { - ol_total = tuple.values[0].decimal().unwrap(); + ol_total = tuple[0].decimal().unwrap(); Ok(()) }, )?; diff --git a/tpcc/src/main.rs b/tpcc/src/main.rs index 9398f9be..e105fc5a 100644 --- a/tpcc/src/main.rs +++ b/tpcc/src/main.rs @@ -864,7 +864,7 @@ impl_from_tpcc_error!(std::io::Error, Io); fn explain_tpcc() -> Result<(), DatabaseError> { use kite_sql::db::CatalogKind; use kite_sql::db::{DataBaseBuilder, ResultIter}; - use kite_sql::types::tuple::Tuple; + use kite_sql::types::tuple::TupleImpl; fn create_table(mut iter: I) -> Result { let mut output = iter.schema(|schema| { @@ -880,7 +880,6 @@ fn explain_tpcc() -> Result<(), DatabaseError> { while let Some(()) = iter.next_tuple(|_, tuple| { output.push_str( &tuple - .values .iter() .map(|value| value.to_string()) .collect::>() @@ -894,7 +893,7 @@ fn explain_tpcc() -> Result<(), DatabaseError> { fn with_next_tuple( mut iter: I, - f: impl FnOnce(&Tuple) -> R, + f: impl FnOnce(&TupleImpl) -> R, ) -> Result { let value = iter .next_tuple(|_, tuple| f(tuple))? @@ -917,53 +916,41 @@ fn explain_tpcc() -> Result<(), DatabaseError> { ] { database.load(CatalogKind::Table(table_name.to_string().into()))?; } - let mut tx = database.new_transaction()?; + let mut tx = database.new_transaction(false)?; let (c_w_id, c_d_id, c_id, c_last, c_balance, c_data) = with_next_tuple( tx.run("SELECT c_w_id, c_d_id, c_id, c_last, c_balance, c_data FROM customer limit 1")?, |tuple| { ( - tuple.values[0].clone(), - tuple.values[1].clone(), - tuple.values[2].clone(), - tuple.values[3].clone(), - tuple.values[4].clone(), - tuple.values[5].clone(), + tuple[0].clone(), + tuple[1].clone(), + tuple[2].clone(), + tuple[3].clone(), + tuple[4].clone(), + tuple[5].clone(), ) }, )?; let (d_id, d_w_id, d_next_o_id) = with_next_tuple( tx.run("SELECT d_id, d_w_id, d_next_o_id FROM district limit 1")?, - |tuple| { - ( - tuple.values[0].clone(), - tuple.values[1].clone(), - tuple.values[2].clone(), - ) - }, + |tuple| (tuple[0].clone(), tuple[1].clone(), tuple[2].clone()), )?; let i_id = with_next_tuple(tx.run("SELECT i_id FROM item limit 1")?, |tuple| { - tuple.values[0].clone() + tuple[0].clone() })?; let (s_i_id, s_w_id, s_quantity) = with_next_tuple( tx.run("SELECT s_i_id, s_w_id, s_quantity FROM stock limit 1")?, - |tuple| { - ( - tuple.values[0].clone(), - tuple.values[1].clone(), - tuple.values[2].clone(), - ) - }, + |tuple| (tuple[0].clone(), tuple[1].clone(), tuple[2].clone()), )?; let (o_w_id, o_d_id, o_c_id, o_id, o_carrier_id) = with_next_tuple( tx.run("SELECT o_w_id, o_d_id, o_c_id, o_id, o_carrier_id FROM orders limit 1")?, |tuple| { ( - tuple.values[0].clone(), - tuple.values[1].clone(), - tuple.values[2].clone(), - tuple.values[3].clone(), - tuple.values[4].clone(), + tuple[0].clone(), + tuple[1].clone(), + tuple[2].clone(), + tuple[3].clone(), + tuple[4].clone(), ) }, )?; @@ -971,22 +958,16 @@ fn explain_tpcc() -> Result<(), DatabaseError> { tx.run("SELECT ol_w_id, ol_d_id, ol_o_id, ol_delivery_d FROM order_line limit 1")?, |tuple| { ( - tuple.values[0].clone(), - tuple.values[1].clone(), - tuple.values[2].clone(), - tuple.values[3].clone(), + tuple[0].clone(), + tuple[1].clone(), + tuple[2].clone(), + tuple[3].clone(), ) }, )?; let (no_d_id, no_w_id, no_o_id) = with_next_tuple( tx.run("SELECT no_d_id, no_w_id, no_o_id FROM new_orders limit 1")?, - |tuple| { - ( - tuple.values[0].clone(), - tuple.values[1].clone(), - tuple.values[2].clone(), - ) - }, + |tuple| (tuple[0].clone(), tuple[1].clone(), tuple[2].clone()), )?; // ORDER { diff --git a/tpcc/src/new_ord.rs b/tpcc/src/new_ord.rs index ce2d44fd..7836c219 100644 --- a/tpcc/src/new_ord.rs +++ b/tpcc/src/new_ord.rs @@ -94,10 +94,10 @@ impl TpccTransaction for NewOrd { (4, DataValue::Int64(args.c_id as i64)), ], &mut |tuple| { - c_discount = tuple.values[0].decimal().unwrap(); - c_last = tuple.values[1].utf8().unwrap().to_string(); - c_credit = tuple.values[2].utf8().unwrap().to_string(); - w_tax = tuple.values[3].decimal().unwrap(); + c_discount = tuple[0].decimal().unwrap(); + c_last = tuple[1].utf8().unwrap().to_string(); + c_credit = tuple[2].utf8().unwrap().to_string(); + w_tax = tuple[3].decimal().unwrap(); Ok(()) }, )?; @@ -116,9 +116,9 @@ impl TpccTransaction for NewOrd { (3, DataValue::Int32(args.c_id as i32)), ], &mut |tuple| { - c_discount = tuple.values[0].decimal().unwrap(); - c_last = tuple.values[1].utf8().unwrap().to_string(); - c_credit = tuple.values[2].utf8().unwrap().to_string(); + c_discount = tuple[0].decimal().unwrap(); + c_last = tuple[1].utf8().unwrap().to_string(); + c_credit = tuple[2].utf8().unwrap().to_string(); Ok(()) }, )?; @@ -128,7 +128,7 @@ impl TpccTransaction for NewOrd { &mut statements[2], &[(1, DataValue::Int16(args.w_id as i16))], &mut |tuple| { - w_tax = tuple.values[0].decimal().unwrap(); + w_tax = tuple[0].decimal().unwrap(); Ok(()) }, )?; @@ -145,8 +145,8 @@ impl TpccTransaction for NewOrd { (2, DataValue::Int16(args.w_id as i16)), ], &mut |tuple| { - d_next_o_id = tuple.values[0].i32().unwrap(); - d_tax = tuple.values[1].decimal().unwrap(); + d_next_o_id = tuple[0].i32().unwrap(); + d_tax = tuple[1].decimal().unwrap(); Ok(()) }, )?; @@ -215,9 +215,9 @@ impl TpccTransaction for NewOrd { let mut i_name = String::new(); let mut i_data = String::new(); tx.with_query_one(&mut statements[7], ¶ms, &mut |tuple| { - i_price = tuple.values[0].decimal().unwrap(); - i_name = tuple.values[1].utf8().unwrap().to_string(); - i_data = tuple.values[2].utf8().unwrap().to_string(); + i_price = tuple[0].decimal().unwrap(); + i_name = tuple[1].utf8().unwrap().to_string(); + i_data = tuple[2].utf8().unwrap().to_string(); Ok(()) })?; @@ -242,18 +242,18 @@ impl TpccTransaction for NewOrd { let mut s_dist_09 = String::new(); let mut s_dist_10 = String::new(); tx.with_query_one(&mut statements[8], ¶ms, &mut |tuple| { - s_quantity = tuple.values[0].i16().unwrap(); - s_data = tuple.values[1].utf8().unwrap().to_string(); - s_dist_01 = tuple.values[2].utf8().unwrap().to_string(); - s_dist_02 = tuple.values[3].utf8().unwrap().to_string(); - s_dist_03 = tuple.values[4].utf8().unwrap().to_string(); - s_dist_04 = tuple.values[5].utf8().unwrap().to_string(); - s_dist_05 = tuple.values[6].utf8().unwrap().to_string(); - s_dist_06 = tuple.values[7].utf8().unwrap().to_string(); - s_dist_07 = tuple.values[8].utf8().unwrap().to_string(); - s_dist_08 = tuple.values[9].utf8().unwrap().to_string(); - s_dist_09 = tuple.values[10].utf8().unwrap().to_string(); - s_dist_10 = tuple.values[11].utf8().unwrap().to_string(); + s_quantity = tuple[0].i16().unwrap(); + s_data = tuple[1].utf8().unwrap().to_string(); + s_dist_01 = tuple[2].utf8().unwrap().to_string(); + s_dist_02 = tuple[3].utf8().unwrap().to_string(); + s_dist_03 = tuple[4].utf8().unwrap().to_string(); + s_dist_04 = tuple[5].utf8().unwrap().to_string(); + s_dist_05 = tuple[6].utf8().unwrap().to_string(); + s_dist_06 = tuple[7].utf8().unwrap().to_string(); + s_dist_07 = tuple[8].utf8().unwrap().to_string(); + s_dist_08 = tuple[9].utf8().unwrap().to_string(); + s_dist_09 = tuple[10].utf8().unwrap().to_string(); + s_dist_10 = tuple[11].utf8().unwrap().to_string(); Ok(()) })?; diff --git a/tpcc/src/order_stat.rs b/tpcc/src/order_stat.rs index 5243c599..6bbd0dc4 100644 --- a/tpcc/src/order_stat.rs +++ b/tpcc/src/order_stat.rs @@ -68,7 +68,7 @@ impl TpccTransaction for OrderStat { (3, DataValue::from(args.c_last.clone())), ], &mut |tuple| { - name_cnt = tuple.values[0].i32().unwrap() as usize; + name_cnt = tuple[0].i32().unwrap() as usize; Ok(()) }, )?; @@ -87,10 +87,10 @@ impl TpccTransaction for OrderStat { let mut c_middle = String::new(); let mut c_last = String::new(); tx.with_query_nth(&mut statements[1], ¶ms, target, &mut |tuple| { - c_balance = tuple.values[0].decimal().unwrap(); - c_first = tuple.values[1].utf8().unwrap().to_string(); - c_middle = tuple.values[2].utf8().unwrap().to_string(); - c_last = tuple.values[3].utf8().unwrap().to_string(); + c_balance = tuple[0].decimal().unwrap(); + c_first = tuple[1].utf8().unwrap().to_string(); + c_middle = tuple[2].utf8().unwrap().to_string(); + c_last = tuple[3].utf8().unwrap().to_string(); Ok(()) })?; (c_balance, c_first, c_middle, c_last) @@ -108,10 +108,10 @@ impl TpccTransaction for OrderStat { (3, DataValue::Int32(args.c_id as i32)), ], &mut |tuple| { - c_balance = tuple.values[0].decimal().unwrap(); - c_first = tuple.values[1].utf8().unwrap().to_string(); - c_middle = tuple.values[2].utf8().unwrap().to_string(); - c_last = tuple.values[3].utf8().unwrap().to_string(); + c_balance = tuple[0].decimal().unwrap(); + c_first = tuple[1].utf8().unwrap().to_string(); + c_middle = tuple[2].utf8().unwrap().to_string(); + c_last = tuple[3].utf8().unwrap().to_string(); Ok(()) }, )?; @@ -128,7 +128,7 @@ impl TpccTransaction for OrderStat { ]; let mut o_id = 0; tx.with_query_one(&mut statements[3], ¶ms, &mut |tuple| { - o_id = tuple.values[0].i32().unwrap(); + o_id = tuple[0].i32().unwrap(); Ok(()) })?; // "SELECT ol_i_id, ol_supply_w_id, ol_quantity, ol_amount, ol_delivery_d FROM order_line WHERE ol_w_id = ? AND ol_d_id = ? AND ol_o_id = ?" @@ -138,11 +138,11 @@ impl TpccTransaction for OrderStat { (3, DataValue::Int32(o_id)), ]; tx.with_query_one(&mut statements[4], ¶ms, &mut |_| Ok(()))?; - // let ol_i_id = tuple.values[0].i32(); - // let ol_supply_w_id = tuple.values[1].i16(); - // let ol_quantity = tuple.values[2].i8(); - // let ol_amount = tuple.values[3].decimal(); - // let ol_delivery_d = tuple.values[4].datetime(); + // let ol_i_id = tuple[0].i32(); + // let ol_supply_w_id = tuple[1].i16(); + // let ol_quantity = tuple[2].i8(); + // let ol_amount = tuple[3].decimal(); + // let ol_delivery_d = tuple[4].datetime(); Ok(()) } diff --git a/tpcc/src/payment.rs b/tpcc/src/payment.rs index 80ebe1fd..27eed290 100644 --- a/tpcc/src/payment.rs +++ b/tpcc/src/payment.rs @@ -90,12 +90,12 @@ impl TpccTransaction for Payment { &mut statements[1], &[(1, DataValue::Int16(args.w_id as i16))], &mut |tuple| { - w_street_1 = tuple.values[0].utf8().unwrap().to_string(); - w_street_2 = tuple.values[1].utf8().unwrap().to_string(); - w_city = tuple.values[2].utf8().unwrap().to_string(); - w_state = tuple.values[3].utf8().unwrap().to_string(); - w_zip = tuple.values[4].utf8().unwrap().to_string(); - w_name = tuple.values[5].utf8().unwrap().to_string(); + w_street_1 = tuple[0].utf8().unwrap().to_string(); + w_street_2 = tuple[1].utf8().unwrap().to_string(); + w_city = tuple[2].utf8().unwrap().to_string(); + w_state = tuple[3].utf8().unwrap().to_string(); + w_zip = tuple[4].utf8().unwrap().to_string(); + w_name = tuple[5].utf8().unwrap().to_string(); Ok(()) }, )?; @@ -124,12 +124,12 @@ impl TpccTransaction for Payment { (2, DataValue::Int8(args.d_id as i8)), ], &mut |tuple| { - d_street_1 = tuple.values[0].utf8().unwrap().to_string(); - d_street_2 = tuple.values[1].utf8().unwrap().to_string(); - d_city = tuple.values[2].utf8().unwrap().to_string(); - d_state = tuple.values[3].utf8().unwrap().to_string(); - d_zip = tuple.values[4].utf8().unwrap().to_string(); - d_name = tuple.values[5].utf8().unwrap().to_string(); + d_street_1 = tuple[0].utf8().unwrap().to_string(); + d_street_2 = tuple[1].utf8().unwrap().to_string(); + d_city = tuple[2].utf8().unwrap().to_string(); + d_state = tuple[3].utf8().unwrap().to_string(); + d_zip = tuple[4].utf8().unwrap().to_string(); + d_name = tuple[5].utf8().unwrap().to_string(); Ok(()) }, )?; @@ -146,7 +146,7 @@ impl TpccTransaction for Payment { (3, DataValue::from(args.c_last.clone())), ], &mut |tuple| { - name_cnt = tuple.values[0].i32().unwrap(); + name_cnt = tuple[0].i32().unwrap(); Ok(()) }, )?; @@ -161,7 +161,7 @@ impl TpccTransaction for Payment { } let target = name_cnt as usize / 2 - 1; tx.with_query_nth(&mut statements[5], ¶ms, target, &mut |tuple| { - c_id = tuple.values[0].i32().unwrap(); + c_id = tuple[0].i32().unwrap(); Ok(()) })?; } @@ -188,20 +188,20 @@ impl TpccTransaction for Payment { (3, DataValue::Int32(c_id)), ], &mut |tuple| { - c_first = tuple.values[0].utf8().unwrap().to_string(); - c_middle = tuple.values[1].utf8().unwrap().to_string(); - c_last = tuple.values[2].utf8().unwrap().to_string(); - c_street_1 = tuple.values[3].utf8().unwrap().to_string(); - c_street_2 = tuple.values[4].utf8().unwrap().to_string(); - c_city = tuple.values[5].utf8().unwrap().to_string(); - c_state = tuple.values[6].utf8().unwrap().to_string(); - c_zip = tuple.values[7].utf8().unwrap().to_string(); - c_phone = tuple.values[8].utf8().unwrap().to_string(); - c_credit = tuple.values[9].utf8().map(ToString::to_string); - c_credit_lim = tuple.values[10].i64().unwrap(); - c_discount = tuple.values[11].decimal().unwrap(); - c_balance = tuple.values[12].decimal().unwrap(); - c_since = tuple.values[13].datetime().unwrap(); + c_first = tuple[0].utf8().unwrap().to_string(); + c_middle = tuple[1].utf8().unwrap().to_string(); + c_last = tuple[2].utf8().unwrap().to_string(); + c_street_1 = tuple[3].utf8().unwrap().to_string(); + c_street_2 = tuple[4].utf8().unwrap().to_string(); + c_city = tuple[5].utf8().unwrap().to_string(); + c_state = tuple[6].utf8().unwrap().to_string(); + c_zip = tuple[7].utf8().unwrap().to_string(); + c_phone = tuple[8].utf8().unwrap().to_string(); + c_credit = tuple[9].utf8().map(ToString::to_string); + c_credit_lim = tuple[10].i64().unwrap(); + c_discount = tuple[11].decimal().unwrap(); + c_balance = tuple[12].decimal().unwrap(); + c_since = tuple[13].datetime().unwrap(); Ok(()) }, )?; @@ -219,7 +219,7 @@ impl TpccTransaction for Payment { (3, DataValue::Int32(c_id)), ], &mut |tuple| { - c_data = tuple.values[0].utf8().unwrap().to_string(); + c_data = tuple[0].utf8().unwrap().to_string(); Ok(()) }, )?; diff --git a/tpcc/src/slev.rs b/tpcc/src/slev.rs index 5fc5a773..6c54f7de 100644 --- a/tpcc/src/slev.rs +++ b/tpcc/src/slev.rs @@ -52,7 +52,7 @@ impl TpccTransaction for Slev { (2, DataValue::Int16(args.w_id as i16)), ], &mut |tuple| { - d_next_o_id = tuple.values[0].i32().unwrap(); + d_next_o_id = tuple[0].i32().unwrap(); Ok(()) }, )?; @@ -67,7 +67,7 @@ impl TpccTransaction for Slev { (4, DataValue::Int32(d_next_o_id)), ], &mut |tuple| { - item_ids.push(tuple.values[0].i32().unwrap()); + item_ids.push(tuple[0].i32().unwrap()); Ok(()) }, )?; @@ -82,7 +82,7 @@ impl TpccTransaction for Slev { (3, DataValue::Int16(args.level as i16)), ], &mut |tuple| { - _low_stock += tuple.values[0].i32().unwrap(); + _low_stock += tuple[0].i32().unwrap(); Ok(()) }, )?;