From 9a57e1195d95b77ddfa03cfe73a360b3fb896ec3 Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Sat, 10 Oct 2026 08:31:18 +0000 Subject: [PATCH 01/46] feat(json2): json_get respect type hints (#9349) Signed-off-by: fys --- Cargo.lock | 70 ++-- Cargo.toml | 30 +- src/query/src/datafusion/json_expr_planner.rs | 236 ++++++++++++-- .../json/json2_query_respect_type_hint.result | 307 +++++++++++++++--- .../json/json2_query_respect_type_hint.sql | 106 +++++- 5 files changed, 620 insertions(+), 129 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 055bf712394..3b66b13e949 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3722,7 +3722,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3774,7 +3774,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3798,7 +3798,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3821,7 +3821,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3847,7 +3847,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "futures", "log", @@ -3857,7 +3857,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "async-compression", @@ -3893,7 +3893,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3917,7 +3917,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3940,7 +3940,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3963,7 +3963,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3995,12 +3995,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" [[package]] name = "datafusion-execution" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4025,7 +4025,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4049,7 +4049,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4060,7 +4060,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4091,7 +4091,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4111,7 +4111,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4122,7 +4122,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-ord 59.2.0", @@ -4146,7 +4146,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4161,7 +4161,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4177,7 +4177,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4186,7 +4186,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "datafusion-doc", "quote", @@ -4196,7 +4196,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "chrono", @@ -4244,7 +4244,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4267,7 +4267,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4281,7 +4281,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "chrono", @@ -4298,7 +4298,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4317,7 +4317,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "arrow-data 59.2.0", @@ -4354,7 +4354,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-catalog", @@ -4380,7 +4380,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4390,7 +4390,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -4400,7 +4400,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4415,7 +4415,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow-schema 59.2.0", "async-trait", @@ -4429,7 +4429,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "arrow 59.2.0", "bigdecimal 0.4.8", @@ -4448,7 +4448,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index cce9b502928..549412f369a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -363,21 +363,21 @@ rev = "f743c3a5fe8c1f57363eac5866c356b7d91f9b39" [patch.crates-io] substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" } -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } diff --git a/src/query/src/datafusion/json_expr_planner.rs b/src/query/src/datafusion/json_expr_planner.rs index 4ec1b729007..4837403c73f 100644 --- a/src/query/src/datafusion/json_expr_planner.rs +++ b/src/query/src/datafusion/json_expr_planner.rs @@ -16,11 +16,12 @@ use std::sync::{Arc, LazyLock}; use arrow_schema::Field; use arrow_schema::extension::ExtensionType; -use common_function::scalars::json::json_get::JsonGetWithType; +use common_function::scalars::json::json_get::{JsonGetWithType, parse_json_get_path}; use common_function::scalars::udf::create_udf; use datafusion_common::arrow::datatypes::DataType; use datafusion_common::{ - Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference, plan_datafusion_err, + Column, DFSchema, DataFusionError, ExprSchema, Result, ScalarValue, TableReference, + plan_datafusion_err, }; use datafusion_expr::expr::{BinaryExpr, ScalarFunction}; use datafusion_expr::planner::{ @@ -35,6 +36,7 @@ use datatypes::extension::json::{ Json2ExtensionType, is_json2_extension_type, parse_legacy_json2_settings, }; use datatypes::types::json_type::JsonNativeType; +use jsonb::jsonpath::Path; use sqlparser::ast::BinaryOperator; /// Rewrites JSON-aware SQL expressions into DataFusion expressions. @@ -44,6 +46,8 @@ use sqlparser::ast::BinaryOperator; /// For example, `select a.b.c` => `select json_get(a, '$.b.c')`. /// - Extends a JSON path with list indexes and fields following an index. /// For example, `select a.b[0].c` => `select json_get(a, '$.b[0].c')`. +/// - Applies JSON2 path type hints to both compound identifiers and explicit `json_get` calls +/// before inferring types from their enclosing expressions. /// - Pushes an "expected type" argument into the `json_get` function when it participates in a /// binary operator. So that `json_get` knows the wanted data type when dealing with variant /// JSON values. @@ -152,12 +156,7 @@ impl ExprPlanner for JsonExprPlanner { Expr::Column(Column::from((qualifier, field))), Expr::Literal(ScalarValue::Utf8(Some(path)), None), ]; - if let Some(json_type) = json_type_hint(field, nested_names)? { - args.push(Expr::Literal( - ScalarValue::try_new_null(&json_type.as_arrow_type())?, - None, - )); - } + apply_json_type_hint(field, nested_names, &mut args)?; Ok(PlannerResult::Planned(Expr::ScalarFunction( ScalarFunction::new_udf(json_get, args), @@ -168,7 +167,27 @@ impl ExprPlanner for JsonExprPlanner { /// /// `Original` carries the possibly modified raw expression to subsequent planners and then /// DataFusion's default function construction. Returning `Planned` would short-circuit both. - fn plan_scalar(&self, mut expr: RawScalarExpr) -> Result> { + fn plan_scalar( + &self, + mut expr: RawScalarExpr, + schema: &DFSchema, + ) -> Result> { + if expr.func.name().eq_ignore_ascii_case(JsonGetWithType::NAME) && expr.args.len() == 2 { + let field = match &expr.args[0] { + // Resolving an unqualified column may require JOIN USING information unavailable + // in this schema. Skip hint injection if lookup fails and let DataFusion resolve + // the column or report the error later. + Expr::Column(column) => schema.field_from_column(column).ok().cloned(), + // Outer references carry their field metadata from the outer query's schema. + Expr::OuterReferenceColumn(field, _) => Some(Arc::clone(field)), + _ => None, + }; + if let Some(field) = field + && let Some(path) = json_get_object_path(&expr.args[1]) + { + apply_json_type_hint(&field, &path, &mut expr.args)?; + } + } push_function_arg_types(expr.func.as_ref(), &mut expr.args)?; Ok(PlannerResult::Original(expr)) } @@ -196,6 +215,46 @@ impl ExprPlanner for JsonExprPlanner { } } +/// Sets the native read type for an exact hinted object path. Both SQL syntaxes use this before +/// an enclosing expression can infer a type or push down a cast. +fn apply_json_type_hint(field: &Field, path: &[String], args: &mut Vec) -> Result<()> { + if !is_json2_extension_type(field) || args.len() != 2 { + return Ok(()); + } + if let Some(json_type) = json_type_hint(field, path)? { + args.push(Expr::Literal( + ScalarValue::try_new_null(&json_type.as_arrow_type())?, + None, + )); + } + Ok(()) +} + +/// Resolves `json_get`'s second argument expression (the path) to object field names for type hint +/// lookup. For example, the string literal `'$.user.age'` resolves to `["user", "age"]`. +/// +/// Returns `None` for non-string literals, non-literal expressions, invalid JSONPaths, or paths +/// containing array indexes, wildcards, or filters. +/// Returning `None` only skips hint injection; the original `json_get` expression is preserved. +fn json_get_object_path(expr: &Expr) -> Option> { + let path = expr.as_literal()?.try_as_str().flatten()?; + // Invalid paths skip hint lookup here. Later planning still validates the path and reports + // the error, preserving the existing JSONPath error behavior. + let json_path = parse_json_get_path(path).ok()?; + let mut names = Vec::new(); + for segment in json_path.paths { + match segment { + Path::Root => {} + Path::DotField(name) | Path::ColonField(name) | Path::ObjectField(name) => { + names.push(name.into_owned()); + } + // Array indexes, wildcards and filters do not identify a hinted object path. + _ => return None, + } + } + Some(names) +} + /// Returns the configured native type for an exact JSON2 object path. fn json_type_hint(field: &Field, path: &[String]) -> Result> { let settings = if field.extension_type_name() == Some(Json2ExtensionType::NAME) { @@ -705,15 +764,138 @@ mod tests { Ok(()) } + #[test] + fn test_plan_json_get_applies_type_hint() -> Result<()> { + let settings = JsonSettings::try_new( + vec![ + JsonTypeHint { + path: vec!["payload".to_string(), "cpu".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }, + JsonTypeHint { + path: vec!["payload.cpu".to_string()], + data_type: ConcreteDataType::string_datatype(), + inverted_index: false, + }, + ], + None, + ) + .unwrap(); + let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))), + ); + let relation = TableReference::bare("events"); + let schema = DFSchema::new_with_metadata( + vec![ + (Some(relation.clone()), Arc::new(field.clone())), + ( + Some(TableReference::bare("other")), + Arc::new(Field::new("j", DataType::Binary, true)), + ), + ], + Default::default(), + )?; + let base = Expr::Column(Column::new(Some(relation.clone()), "j")); + + for (path, names, expected) in [ + ("payload.cpu", vec!["payload", "cpu"], Some(DataType::Int64)), + ( + "$.payload.cpu", + vec!["payload", "cpu"], + Some(DataType::Int64), + ), + ( + r#"$."payload.cpu""#, + vec!["payload.cpu"], + Some(DataType::Utf8View), + ), + ("missing", vec!["missing"], None), + ] { + let Expr::ScalarFunction(function) = json_get_expr(base.clone(), path) else { + unreachable!(); + }; + let PlannerResult::Original(scalar) = JsonExprPlanner.plan_scalar( + RawScalarExpr { + func: function.func, + args: function.args, + }, + &schema, + )? + else { + unreachable!(); + }; + let direct = Expr::ScalarFunction(ScalarFunction::new_udf(scalar.func, scalar.args)); + let names = names.into_iter().map(String::from).collect::>(); + let PlannerResult::Planned(dotted) = + JsonExprPlanner.plan_compound_identifier(&field, Some(&relation), &names)? + else { + unreachable!(); + }; + assert_eq!(expected, extract_json_get_type(&direct), "{path}"); + assert_eq!( + extract_json_get_type(&dotted), + extract_json_get_type(&direct) + ); + + // Outer function inference must not replace a declared hint. + if expected.is_some() { + let PlannerResult::Original(outer) = JsonExprPlanner.plan_scalar( + RawScalarExpr { + func: coalesce(), + args: vec![ + direct.clone(), + Expr::Literal(ScalarValue::Float64(Some(1.0)), None), + ], + }, + &schema, + )? + else { + unreachable!(); + }; + assert_eq!(direct, outer.args[0]); + } + } + + // Non-object paths cannot match an object type hint. A different table's unhinted + // column must not inherit the same-named JSON2 column's hints either. + for (base, path) in [ + (base.clone(), "$"), + (base.clone(), "$.payload.*"), + (base.clone(), "$.payload.cpu[0]"), + (base.clone(), "$.payload["), + (Expr::Column(Column::new(Some("other"), "j")), "payload.cpu"), + ] { + let Expr::ScalarFunction(function) = json_get_expr(base, path) else { + unreachable!(); + }; + let PlannerResult::Original(scalar) = JsonExprPlanner.plan_scalar( + RawScalarExpr { + func: function.func, + args: function.args, + }, + &schema, + )? + else { + unreachable!(); + }; + assert_eq!(2, scalar.args.len(), "{path}"); + } + Ok(()) + } + #[test] fn test_plan_functions() -> Result<()> { let planner = JsonExprPlanner; let json_get = || json_get_expr(Expr::Column(Column::new_unqualified("j")), "a.b"); - let PlannerResult::Original(scalar) = planner.plan_scalar(RawScalarExpr { - func: abs(), - args: vec![json_get()], - })? + let PlannerResult::Original(scalar) = planner.plan_scalar( + RawScalarExpr { + func: abs(), + args: vec![json_get()], + }, + &DFSchema::empty(), + )? else { unreachable!(); }; @@ -722,13 +904,16 @@ mod tests { extract_json_get_type(&scalar.args[0]) ); - let PlannerResult::Original(scalar) = planner.plan_scalar(RawScalarExpr { - func: power(), - args: vec![ - json_get(), - Expr::Column(Column::new_unqualified("exponent")), - ], - })? + let PlannerResult::Original(scalar) = planner.plan_scalar( + RawScalarExpr { + func: power(), + args: vec![ + json_get(), + Expr::Column(Column::new_unqualified("exponent")), + ], + }, + &DFSchema::empty(), + )? else { unreachable!(); }; @@ -793,10 +978,13 @@ mod tests { let mut typed = json_get(); push_json_get_type_arg(&mut typed, &DataType::Float64)?; - let PlannerResult::Original(scalar) = planner.plan_scalar(RawScalarExpr { - func: coalesce(), - args: vec![json_get(), typed], - })? + let PlannerResult::Original(scalar) = planner.plan_scalar( + RawScalarExpr { + func: coalesce(), + args: vec![json_get(), typed], + }, + &DFSchema::empty(), + )? else { unreachable!(); }; diff --git a/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result index dcbffbf6058..d022976cb22 100644 --- a/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result +++ b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result @@ -13,9 +13,43 @@ Affected Rows: 0 INSERT INTO json2_query_respect_type_hint VALUES (1, '{"a":9007199254740993,"b":"three"}', 2), - (2, '{"a":7,"b":"7"}', 2); + (2, '{"a":7,"b":"7"}', 2), + (3, '{}', 2), + (4, '{"a":null,"b":null}', 2); -Affected Rows: 2 +Affected Rows: 4 + +-- Both syntaxes use the hint when reading the memtable, including missing/null paths. +SELECT j.a AS dotted, json_get(j, 'a') AS direct, + arrow_typeof(j.a) AS dotted_type, arrow_typeof(json_get(j, 'a')) AS direct_type +FROM json2_query_respect_type_hint +ORDER BY ts; + ++------------------+------------------+-------------+-------------+ +| dotted | direct | dotted_type | direct_type | ++------------------+------------------+-------------+-------------+ +| 9007199254740993 | 9007199254740993 | Int64 | Int64 | +| 7 | 7 | Int64 | Int64 | +| | | Int64 | Int64 | +| | | Int64 | Int64 | ++------------------+------------------+-------------+-------------+ + +-- An explicit third argument takes precedence over the BIGINT path hint in memtable reads. +SELECT json_get(j, 'a') AS hinted_bigint, + json_get(j, 'a', NULL::DOUBLE) AS explicit_double, + arrow_typeof(json_get(j, 'a')) AS hinted_type, + arrow_typeof(json_get(j, 'a', NULL::DOUBLE)) AS explicit_type +FROM json2_query_respect_type_hint +ORDER BY ts; + ++------------------+--------------------+-------------+---------------+ +| hinted_bigint | explicit_double | hinted_type | explicit_type | ++------------------+--------------------+-------------+---------------+ +| 9007199254740993 | 9007199254740992.0 | Int64 | Float64 | +| 7 | 7.0 | Int64 | Float64 | +| | | Int64 | Float64 | +| | | Int64 | Float64 | ++------------------+--------------------+-------------+---------------+ ADMIN FLUSH_TABLE('json2_query_respect_type_hint'); @@ -25,83 +59,172 @@ ADMIN FLUSH_TABLE('json2_query_respect_type_hint'); | 0 | +----------------------------------------------------+ +-- The explicit read type also overrides the hint in the plan and SST reads. -- SQLNESS REPLACE (peers.*) REDACTED -EXPLAIN SELECT j.a AS hinted_bigint +EXPLAIN SELECT json_get(j, 'a', NULL::DOUBLE) AS explicit_double FROM json2_query_respect_type_hint; -+---------------+--------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+--------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | Projection: json_get(json2_query_respect_type_hint.j, Utf8("$.a"), Int64(NULL)) AS hinted_bigint | -| | TableScan: json2_query_respect_type_hint | -| | ]] | -| physical_plan | CooperativeExec | ++---------------+----------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: json_get(json2_query_respect_type_hint.j, Utf8("a"), Float64(NULL)) AS explicit_double | +| | TableScan: json2_query_respect_type_hint | +| | ]] | +| physical_plan | CooperativeExec | | | MergeScanExec: REDACTED -| | | -+---------------+--------------------------------------------------------------------------------------------------+ +| | | ++---------------+----------------------------------------------------------------------------------------------------+ + +SELECT json_get(j, 'a') AS hinted_bigint, + json_get(j, 'a', NULL::DOUBLE) AS explicit_double, + arrow_typeof(json_get(j, 'a')) AS hinted_type, + arrow_typeof(json_get(j, 'a', NULL::DOUBLE)) AS explicit_type +FROM json2_query_respect_type_hint +ORDER BY ts; + ++------------------+--------------------+-------------+---------------+ +| hinted_bigint | explicit_double | hinted_type | explicit_type | ++------------------+--------------------+-------------+---------------+ +| 9007199254740993 | 9007199254740992.0 | Int64 | Float64 | +| 7 | 7.0 | Int64 | Float64 | +| | | Int64 | Float64 | +| | | Int64 | Float64 | ++------------------+--------------------+-------------+---------------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT j.a AS hinted_bigint, json_get(j, 'a') AS direct_bigint +FROM json2_query_respect_type_hint; + ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: json_get(json2_query_respect_type_hint.j, Utf8("$.a"), Int64(NULL)) AS hinted_bigint, json_get(json2_query_respect_type_hint.j, Utf8("a"), Int64(NULL)) AS direct_bigint | +| | TableScan: json2_query_respect_type_hint | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -- Uncast paths use their JSON2 type hints. SELECT j.a AS hinted_bigint, - j.b AS hinted_string + json_get(j, 'a') AS direct_bigint, + json_get(j, '$.a') AS rooted_bigint, + j.b AS hinted_string, + json_get(j, 'b') AS direct_string FROM json2_query_respect_type_hint ORDER BY ts; -+------------------+---------------+ -| hinted_bigint | hinted_string | -+------------------+---------------+ -| 9007199254740993 | three | -| 7 | 7 | -+------------------+---------------+ ++------------------+------------------+------------------+---------------+---------------+ +| hinted_bigint | direct_bigint | rooted_bigint | hinted_string | direct_string | ++------------------+------------------+------------------+---------------+---------------+ +| 9007199254740993 | 9007199254740993 | 9007199254740993 | three | three | +| 7 | 7 | 7 | 7 | 7 | +| | | | | | +| | | | | | ++------------------+------------------+------------------+---------------+---------------+ + +-- The SST read keeps the same result types as the memtable read. +SELECT arrow_typeof(j.a) AS dotted_type, arrow_typeof(json_get(j, 'a')) AS direct_type, + arrow_typeof(j.b) AS dotted_string_type, arrow_typeof(json_get(j, 'b')) AS direct_string_type +FROM json2_query_respect_type_hint LIMIT 1; + ++-------------+-------------+--------------------+--------------------+ +| dotted_type | direct_type | dotted_string_type | direct_string_type | ++-------------+-------------+--------------------+--------------------+ +| Int64 | Int64 | Utf8View | Utf8View | ++-------------+-------------+--------------------+--------------------+ -- A JSON2 type hint takes precedence over an explicit SQL cast for the native -- read type. -- SQLNESS REPLACE (peers.*) REDACTED -EXPLAIN SELECT j.a::STRING +EXPLAIN SELECT j.a::STRING, json_get(j, 'a')::STRING FROM json2_query_respect_type_hint; -+---------------+-----------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+-----------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | Projection: CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.a"), Int64(NULL)) AS Utf8) | -| | TableScan: json2_query_respect_type_hint | -| | ]] | -| physical_plan | CooperativeExec | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.a"), Int64(NULL)) AS Utf8), CAST(json_get(json2_query_respect_type_hint.j, Utf8("a"), Int64(NULL)) AS Utf8) | +| | TableScan: json2_query_respect_type_hint | +| | ]] | +| physical_plan | CooperativeExec | | | MergeScanExec: REDACTED -| | | -+---------------+-----------------------------------------------------------------------------------------------+ +| | | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -- An explicit CAST changes the expression result type. SELECT - j.a::STRING AS cast_string + j.a::STRING AS cast_string, + json_get(j, 'a')::STRING AS direct_cast_string, + arrow_typeof(json_get(j, 'a')::STRING) AS cast_type FROM json2_query_respect_type_hint ORDER BY ts; -+------------------+ -| cast_string | -+------------------+ -| 9007199254740993 | -| 7 | -+------------------+ ++------------------+--------------------+-----------+ +| cast_string | direct_cast_string | cast_type | ++------------------+--------------------+-----------+ +| 9007199254740993 | 9007199254740993 | Utf8 | +| 7 | 7 | Utf8 | +| | | Utf8 | +| | | Utf8 | ++------------------+--------------------+-----------+ -- A JSON2 type hint takes precedence over expression-context inference. -- SQLNESS REPLACE (peers.*) REDACTED -EXPLAIN SELECT coalesce(j.b, exponent) +EXPLAIN SELECT coalesce(j.b, exponent), coalesce(json_get(j, 'b'), exponent) FROM json2_query_respect_type_hint; -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | Projection: CASE WHEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.b"), Utf8View(NULL)) AS Float64) IS NOT NULL THEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.b"), Utf8View(NULL)) AS Float64) ELSE json2_query_respect_type_hint.exponent END AS coalesce(json_get(json2_query_respect_type_hint.j,Utf8("$.b"),Utf8View(NULL)),json2_query_respect_type_hint.exponent) | -| | TableScan: json2_query_respect_type_hint | -| | ]] | -| physical_plan | CooperativeExec | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: CASE WHEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.b"), Utf8View(NULL)) AS Float64) IS NOT NULL THEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.b"), Utf8View(NULL)) AS Float64) ELSE json2_query_respect_type_hint.exponent END AS coalesce(json_get(json2_query_respect_type_hint.j,Utf8("$.b"),Utf8View(NULL)),json2_query_respect_type_hint.exponent), CASE WHEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("b"), Utf8View(NULL)) AS Float64) IS NOT NULL THEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("b"), Utf8View(NULL)) AS Float64) ELSE json2_query_respect_type_hint.exponent END AS coalesce(json_get(json2_query_respect_type_hint.j,Utf8("b"),Utf8View(NULL)),json2_query_respect_type_hint.exponent) | +| | TableScan: json2_query_respect_type_hint | +| | ]] | +| physical_plan | CooperativeExec | | | MergeScanExec: REDACTED -| | | -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| | | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- Hint injection precedes scalar, aggregate, window and binary-expression inference. +SELECT abs(j.a) AS dotted_abs, abs(json_get(j, 'a')) AS direct_abs, + j.a + 1 AS dotted_plus, json_get(j, 'a') + 1 AS direct_plus +FROM json2_query_respect_type_hint ORDER BY ts; + ++------------------+------------------+------------------+------------------+ +| dotted_abs | direct_abs | dotted_plus | direct_plus | ++------------------+------------------+------------------+------------------+ +| 9007199254740993 | 9007199254740993 | 9007199254740994 | 9007199254740994 | +| 7 | 7 | 8 | 8 | +| | | | | +| | | | | ++------------------+------------------+------------------+------------------+ + +SELECT sum(j.a) AS dotted_sum, sum(json_get(j, 'a')) AS direct_sum +FROM json2_query_respect_type_hint; + ++------------------+------------------+ +| dotted_sum | direct_sum | ++------------------+------------------+ +| 9007199254741000 | 9007199254741000 | ++------------------+------------------+ + +SELECT sum(j.a) OVER (ORDER BY ts) AS dotted_sum, + sum(json_get(j, 'a')) OVER (ORDER BY ts) AS direct_sum +FROM json2_query_respect_type_hint ORDER BY ts; + ++------------------+------------------+ +| dotted_sum | direct_sum | ++------------------+------------------+ +| 9007199254740993 | 9007199254740993 | +| 9007199254741000 | 9007199254741000 | +| 9007199254741000 | 9007199254741000 | +| 9007199254741000 | 9007199254741000 | ++------------------+------------------+ -- A STRING type hint is respected, so numeric arithmetic requires an explicit cast. SELECT j.b + exponent @@ -109,6 +232,11 @@ FROM json2_query_respect_type_hint; Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Cannot coerce arithmetic expression Utf8View + Float64 to valid types +SELECT json_get(j, 'b') + exponent +FROM json2_query_respect_type_hint; + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Cannot coerce arithmetic expression Utf8View + Float64 to valid types + -- An explicit cast changes the expression type while preserving the STRING native read type. -- The invalid string value still reports a normal SQL cast error. SELECT j.b::DOUBLE + exponent @@ -116,7 +244,88 @@ FROM json2_query_respect_type_hint; Error: 3001(EngineExecuteQuery), Cast error: Cannot cast string 'three' to value of Float64 type +SELECT json_get(j, 'b')::DOUBLE + exponent +FROM json2_query_respect_type_hint; + +Error: 3001(EngineExecuteQuery), Cast error: Cannot cast string 'three' to value of Float64 type + +-- Resolve the input field through table aliases and a renamed derived-table column. +SELECT t.j.a AS dotted, json_get(t.j, 'a') AS direct +FROM json2_query_respect_type_hint AS t ORDER BY t.ts; + ++------------------+------------------+ +| dotted | direct | ++------------------+------------------+ +| 9007199254740993 | 9007199254740993 | +| 7 | 7 | +| | | +| | | ++------------------+------------------+ + +SELECT q.payload.a AS dotted, json_get(q.payload, 'a') AS direct +FROM (SELECT ts, j AS payload FROM json2_query_respect_type_hint) AS q ORDER BY q.ts; + ++------------------+------------------+ +| dotted | direct | ++------------------+------------------+ +| 9007199254740993 | 9007199254740993 | +| 7 | 7 | +| | | +| | | ++------------------+------------------+ + DROP TABLE json2_query_respect_type_hint; Affected Rows: 0 +CREATE TABLE json2_hint_paths ( + ts TIMESTAMP TIME INDEX, + j JSON2 (nested.a BIGINT, "nested.a" BOOLEAN) +) WITH ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +INSERT INTO json2_hint_paths VALUES + (1, '{"nested":{"a":42},"nested.a":true,"unhinted":7}'); + +Affected Rows: 1 + +-- Quoted keys are literal names, not nested paths; unhinted paths retain their default type. +SELECT j.nested.a AS dotted_nested, json_get(j, '$.nested.a') AS direct_nested, + j."nested.a" AS dotted_key, json_get(j, '$."nested.a"') AS direct_key, + arrow_typeof(json_get(j, 'nested.a')) AS nested_type, + arrow_typeof(json_get(j, '$."nested.a"')) AS key_type, + arrow_typeof(j.unhinted) AS dotted_unhinted_type, + arrow_typeof(json_get(j, 'unhinted')) AS direct_unhinted_type +FROM json2_hint_paths; + ++---------------+---------------+------------+------------+-------------+----------+----------------------+----------------------+ +| dotted_nested | direct_nested | dotted_key | direct_key | nested_type | key_type | dotted_unhinted_type | direct_unhinted_type | ++---------------+---------------+------------+------------+-------------+----------+----------------------+----------------------+ +| 42 | 42 | true | true | Int64 | Boolean | Utf8View | Utf8View | ++---------------+---------------+------------+------------+-------------+----------+----------------------+----------------------+ + +DROP TABLE json2_hint_paths; + +Affected Rows: 0 + +-- Hint lookup must defer unqualified JOIN USING columns to DataFusion, even for legacy JSON. +SELECT json_get(j, 'a') AS value +FROM (SELECT parse_json('{"a":"x"}') AS j) AS l +JOIN (SELECT parse_json('{"a":"x"}') AS j) AS r USING (j); + ++-------+ +| value | ++-------+ +| x | ++-------+ + +-- Without USING, the unqualified column is genuinely ambiguous and must still fail. +SELECT json_get(j, 'a') AS value +FROM (SELECT parse_json('{"a":"x"}') AS j) AS l +JOIN (SELECT parse_json('{"a":"x"}') AS j) AS r ON l.j = r.j; + +Error: 3000(PlanQuery), Failed to plan SQL: Ambiguous reference to unqualified field j + diff --git a/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql index 604f7d1f31b..eb91536b40b 100644 --- a/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql +++ b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql @@ -11,45 +11,139 @@ CREATE TABLE json2_query_respect_type_hint ( INSERT INTO json2_query_respect_type_hint VALUES (1, '{"a":9007199254740993,"b":"three"}', 2), - (2, '{"a":7,"b":"7"}', 2); + (2, '{"a":7,"b":"7"}', 2), + (3, '{}', 2), + (4, '{"a":null,"b":null}', 2); + +-- Both syntaxes use the hint when reading the memtable, including missing/null paths. +SELECT j.a AS dotted, json_get(j, 'a') AS direct, + arrow_typeof(j.a) AS dotted_type, arrow_typeof(json_get(j, 'a')) AS direct_type +FROM json2_query_respect_type_hint +ORDER BY ts; + +-- An explicit third argument takes precedence over the BIGINT path hint in memtable reads. +SELECT json_get(j, 'a') AS hinted_bigint, + json_get(j, 'a', NULL::DOUBLE) AS explicit_double, + arrow_typeof(json_get(j, 'a')) AS hinted_type, + arrow_typeof(json_get(j, 'a', NULL::DOUBLE)) AS explicit_type +FROM json2_query_respect_type_hint +ORDER BY ts; ADMIN FLUSH_TABLE('json2_query_respect_type_hint'); +-- The explicit read type also overrides the hint in the plan and SST reads. -- SQLNESS REPLACE (peers.*) REDACTED -EXPLAIN SELECT j.a AS hinted_bigint +EXPLAIN SELECT json_get(j, 'a', NULL::DOUBLE) AS explicit_double +FROM json2_query_respect_type_hint; + +SELECT json_get(j, 'a') AS hinted_bigint, + json_get(j, 'a', NULL::DOUBLE) AS explicit_double, + arrow_typeof(json_get(j, 'a')) AS hinted_type, + arrow_typeof(json_get(j, 'a', NULL::DOUBLE)) AS explicit_type +FROM json2_query_respect_type_hint +ORDER BY ts; + +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT j.a AS hinted_bigint, json_get(j, 'a') AS direct_bigint FROM json2_query_respect_type_hint; -- Uncast paths use their JSON2 type hints. SELECT j.a AS hinted_bigint, - j.b AS hinted_string + json_get(j, 'a') AS direct_bigint, + json_get(j, '$.a') AS rooted_bigint, + j.b AS hinted_string, + json_get(j, 'b') AS direct_string FROM json2_query_respect_type_hint ORDER BY ts; +-- The SST read keeps the same result types as the memtable read. +SELECT arrow_typeof(j.a) AS dotted_type, arrow_typeof(json_get(j, 'a')) AS direct_type, + arrow_typeof(j.b) AS dotted_string_type, arrow_typeof(json_get(j, 'b')) AS direct_string_type +FROM json2_query_respect_type_hint LIMIT 1; + -- A JSON2 type hint takes precedence over an explicit SQL cast for the native -- read type. -- SQLNESS REPLACE (peers.*) REDACTED -EXPLAIN SELECT j.a::STRING +EXPLAIN SELECT j.a::STRING, json_get(j, 'a')::STRING FROM json2_query_respect_type_hint; -- An explicit CAST changes the expression result type. SELECT - j.a::STRING AS cast_string + j.a::STRING AS cast_string, + json_get(j, 'a')::STRING AS direct_cast_string, + arrow_typeof(json_get(j, 'a')::STRING) AS cast_type FROM json2_query_respect_type_hint ORDER BY ts; -- A JSON2 type hint takes precedence over expression-context inference. -- SQLNESS REPLACE (peers.*) REDACTED -EXPLAIN SELECT coalesce(j.b, exponent) +EXPLAIN SELECT coalesce(j.b, exponent), coalesce(json_get(j, 'b'), exponent) FROM json2_query_respect_type_hint; +-- Hint injection precedes scalar, aggregate, window and binary-expression inference. +SELECT abs(j.a) AS dotted_abs, abs(json_get(j, 'a')) AS direct_abs, + j.a + 1 AS dotted_plus, json_get(j, 'a') + 1 AS direct_plus +FROM json2_query_respect_type_hint ORDER BY ts; + +SELECT sum(j.a) AS dotted_sum, sum(json_get(j, 'a')) AS direct_sum +FROM json2_query_respect_type_hint; + +SELECT sum(j.a) OVER (ORDER BY ts) AS dotted_sum, + sum(json_get(j, 'a')) OVER (ORDER BY ts) AS direct_sum +FROM json2_query_respect_type_hint ORDER BY ts; + -- A STRING type hint is respected, so numeric arithmetic requires an explicit cast. SELECT j.b + exponent FROM json2_query_respect_type_hint; +SELECT json_get(j, 'b') + exponent +FROM json2_query_respect_type_hint; + -- An explicit cast changes the expression type while preserving the STRING native read type. -- The invalid string value still reports a normal SQL cast error. SELECT j.b::DOUBLE + exponent FROM json2_query_respect_type_hint; +SELECT json_get(j, 'b')::DOUBLE + exponent +FROM json2_query_respect_type_hint; + +-- Resolve the input field through table aliases and a renamed derived-table column. +SELECT t.j.a AS dotted, json_get(t.j, 'a') AS direct +FROM json2_query_respect_type_hint AS t ORDER BY t.ts; + +SELECT q.payload.a AS dotted, json_get(q.payload, 'a') AS direct +FROM (SELECT ts, j AS payload FROM json2_query_respect_type_hint) AS q ORDER BY q.ts; + DROP TABLE json2_query_respect_type_hint; + +CREATE TABLE json2_hint_paths ( + ts TIMESTAMP TIME INDEX, + j JSON2 (nested.a BIGINT, "nested.a" BOOLEAN) +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO json2_hint_paths VALUES + (1, '{"nested":{"a":42},"nested.a":true,"unhinted":7}'); + +-- Quoted keys are literal names, not nested paths; unhinted paths retain their default type. +SELECT j.nested.a AS dotted_nested, json_get(j, '$.nested.a') AS direct_nested, + j."nested.a" AS dotted_key, json_get(j, '$."nested.a"') AS direct_key, + arrow_typeof(json_get(j, 'nested.a')) AS nested_type, + arrow_typeof(json_get(j, '$."nested.a"')) AS key_type, + arrow_typeof(j.unhinted) AS dotted_unhinted_type, + arrow_typeof(json_get(j, 'unhinted')) AS direct_unhinted_type +FROM json2_hint_paths; + +DROP TABLE json2_hint_paths; + +-- Hint lookup must defer unqualified JOIN USING columns to DataFusion, even for legacy JSON. +SELECT json_get(j, 'a') AS value +FROM (SELECT parse_json('{"a":"x"}') AS j) AS l +JOIN (SELECT parse_json('{"a":"x"}') AS j) AS r USING (j); + +-- Without USING, the unqualified column is genuinely ambiguous and must still fail. +SELECT json_get(j, 'a') AS value +FROM (SELECT parse_json('{"a":"x"}') AS j) AS l +JOIN (SELECT parse_json('{"a":"x"}') AS j) AS r ON l.j = r.j; From ad83232ad1da0eaf992bee70654dbdee15e76d56 Mon Sep 17 00:00:00 2001 From: sun Date: Sat, 10 Oct 2026 09:21:56 +0000 Subject: [PATCH 02/46] feat: update dashboard to v0.14.0 (#9510) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index 2878fc78d46..4a29f93bbe5 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.15 +v0.14.0 From 32e8467cda0704d79cfae7bca3d0698d24540cee Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 10 Oct 2026 11:18:46 +0000 Subject: [PATCH 03/46] fix(query): stop reporting region scans as hash partitioned (#9491) * fix(query): stop reporting region scans as hash partitioned Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): use current DataFusion plan downcasts Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): match null rendering in join regression Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): refresh union plan for unknown scan partitioning Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): refresh shared plans after distribution rule removal Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): cover semi joins with remote dynamic filtering Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): strengthen partitioned merge scan regressions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): remove unused merge scan test helpers Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(perf): measure merge scan hash repartition cost Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): adapt merge scan deserialization to two-arg constructor Adapt the production and test serializer callers added by #9206 and remove the obsolete partition_cols assertion. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): adapt nested merge scan integration constructor Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/analyzer.rs | 9 +- src/query/src/dist_plan/analyzer/fallback.rs | 72 +--- src/query/src/dist_plan/merge_scan.rs | 221 +---------- src/query/src/dist_plan/planner.rs | 1 - src/query/src/metrics.rs | 3 +- src/query/src/optimizer.rs | 1 - .../src/optimizer/json_schema_concretize.rs | 3 +- src/query/src/optimizer/pass_distribution.rs | 362 ------------------ .../src/query_engine/default_serializer.rs | 13 +- src/query/src/query_engine/state.rs | 9 +- tests-integration/src/tests/instance_test.rs | 284 ++++++++++++++ .../nested_merge_scan_capability_test.rs | 2 +- .../distributed/explain/subqueries.result | 3 +- .../common/join/merge_scan_left_join.result | 148 +++++++ .../common/join/merge_scan_left_join.sql | 81 ++++ .../common/join/merge_scan_semi_join.result | 104 +++++ .../common/join/merge_scan_semi_join.sql | 49 +++ .../common/order/order_by_exceptions.result | 2 +- .../common/tql-explain-analyze/explain.result | 3 - .../standalone/common/tql/tql-cte.result | 2 +- .../mergescan_shuffle_9456/case.toml | 152 ++++++++ 21 files changed, 851 insertions(+), 673 deletions(-) delete mode 100644 src/query/src/optimizer/pass_distribution.rs create mode 100644 tests/cases/standalone/common/join/merge_scan_left_join.result create mode 100644 tests/cases/standalone/common/join/merge_scan_left_join.sql create mode 100644 tests/cases/standalone/common/join/merge_scan_semi_join.result create mode 100644 tests/cases/standalone/common/join/merge_scan_semi_join.sql create mode 100644 tests/perf/query_cases/mergescan_shuffle_9456/case.toml diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index da26d88b433..38d386cc818 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -768,14 +768,7 @@ impl PlanRewriter { ); // add merge scan as the new root - let mut node = MergeScanLogicalPlan::new( - on_node.clone(), - false, - // at this stage, the partition cols should be set - // treat it as non-partitioned if None - self.partition_cols.clone().unwrap_or_default(), - ) - .into_logical_plan(); + let mut node = MergeScanLogicalPlan::new(on_node.clone(), false).into_logical_plan(); // expand stages for new_stage in self.stage.drain(..) { diff --git a/src/query/src/dist_plan/analyzer/fallback.rs b/src/query/src/dist_plan/analyzer/fallback.rs index 86c7f50628d..0aabdd6cfd8 100644 --- a/src/query/src/dist_plan/analyzer/fallback.rs +++ b/src/query/src/dist_plan/analyzer/fallback.rs @@ -17,18 +17,11 @@ //! This is a temporary solution, and will be removed once we have a more robust plan rewriter //! -use std::collections::BTreeSet; - -use common_telemetry::debug; -use datafusion::datasource::DefaultTableSource; use datafusion_common::Result as DfResult; use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; use datafusion_expr::LogicalPlan; -use table::metadata::TableType; -use table::table::adapter::DfTableProviderAdapter; use crate::dist_plan::MergeScanLogicalPlan; -use crate::dist_plan::analyzer::{AliasMapping, OTHER_PHY_PART_COL_PLACEHOLDER}; /// FallbackPlanRewriter is a plan rewriter that will only push down table scan node /// This is used when `PlanRewriter` produce errors when trying to rewrite the plan @@ -44,69 +37,8 @@ impl TreeNodeRewriter for FallbackPlanRewriter { &mut self, plan: Self::Node, ) -> DfResult> { - if let LogicalPlan::TableScan(table_scan) = &plan { - let partition_cols = if let Some(source) = - table_scan.source.downcast_ref::() - { - if let Some(provider) = source - .table_provider - .downcast_ref::() - { - if provider.table().table_type() == TableType::Base { - let info = provider.table().table_info(); - let partition_key_indices = info.meta.partition_key_indices.clone(); - let schema = info.meta.schema.clone(); - let partition_cols = partition_key_indices - .iter() - .map(|index| schema.column_name_by_index(*index).to_string()) - .collect::>(); - debug!( - "FallbackPlanRewriter: loaded table partition metadata, table: {}, table_id: {}, partition_key_indices: {:?}, partition_columns: {:?}", - info.name, - info.ident.table_id, - info.meta.partition_key_indices, - partition_cols, - ); - Some(partition_cols - .into_iter() - .map(|c| { - if c == OTHER_PHY_PART_COL_PLACEHOLDER { - // for placeholder, just return a empty alias - return Ok((c.clone(), BTreeSet::new())); - } - let index = - plan.schema().index_of_column_by_name(None, &c).ok_or_else(|| { - datafusion_common::DataFusionError::Internal( - format!( - "PlanRewriter: maybe_set_partitions: column {c} not found in schema of plan: {plan}" - ), - ) - })?; - let column = plan.schema().columns().get(index).cloned().ok_or_else(|| { - datafusion_common::DataFusionError::Internal(format!( - "PlanRewriter: maybe_set_partitions: column index {index} out of bounds in schema of plan: {plan}" - )) - })?; - Ok((c.clone(), BTreeSet::from([column]))) - }) - .collect::>()?) - } else { - None - } - } else { - None - } - } else { - None - }; - let node = MergeScanLogicalPlan::new( - plan, - false, - // at this stage, the partition cols should be set - // treat it as non-partitioned if None - partition_cols.clone().unwrap_or_default(), - ) - .into_logical_plan(); + if matches!(plan, LogicalPlan::TableScan(_)) { + let node = MergeScanLogicalPlan::new(plan, false).into_logical_plan(); Ok(Transformed::yes(node)) } else { Ok(Transformed::no(plan)) diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index ee4e9574280..041fea160cf 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -17,7 +17,7 @@ use std::cell::Cell; use std::sync::{Arc, Mutex}; use std::time::Duration; -use ahash::{HashMap, HashSet}; +use ahash::HashMap; use arrow_schema::{ ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, SortOptions, @@ -44,10 +44,9 @@ use datafusion::physical_plan::{ }; use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{Column as ColumnExpr, DFSchemaRef, DataFusionError, Result, Statistics}; +use datafusion_common::{DFSchemaRef, DataFusionError, Result, Statistics}; use datafusion_expr::{Expr, Extension, FetchType, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion_physical_expr::expressions::Column; -use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalSortExpr}; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalSortExpr}; use datatypes::extension::json::is_any_json_extension_type; use futures_util::StreamExt; use greptime_proto::v1::region::RegionRequestHeader; @@ -64,7 +63,6 @@ use tokio::time; use tokio::time::Instant; use tracing::{Instrument, Span}; -use crate::dist_plan::analyzer::AliasMapping; use crate::dist_plan::analyzer::utils::patch_batch_timezone; use crate::dist_plan::dyn_filter_bridge::{ CapturedDynFilter, capture_remote_dyn_filters_for_pushdown, @@ -339,7 +337,6 @@ pub struct MergeScanLogicalPlan { output_schema: DFSchemaRef, /// If this plan is a placeholder is_placeholder: bool, - partition_cols: AliasMapping, /// Assigned after dist-plan rewriting so rewriters only deal with plan shape. remote_dyn_filter_producer_id: Option, } @@ -350,29 +347,20 @@ impl PartialOrd for MergeScanLogicalPlan { input, output_schema, is_placeholder, - partition_cols, remote_dyn_filter_producer_id, } = self; let Self { input: other_input, output_schema: other_output_schema, is_placeholder: other_is_placeholder, - partition_cols: other_partition_cols, remote_dyn_filter_producer_id: other_remote_dyn_filter_producer_id, } = other; - let ordering = ( - input, - is_placeholder, - partition_cols, - remote_dyn_filter_producer_id, - ) - .partial_cmp(&( - other_input, - other_is_placeholder, - other_partition_cols, - other_remote_dyn_filter_producer_id, - )); + let ordering = (input, is_placeholder, remote_dyn_filter_producer_id).partial_cmp(&( + other_input, + other_is_placeholder, + other_remote_dyn_filter_producer_id, + )); match ordering { Some(std::cmp::Ordering::Equal) if output_schema != other_output_schema => None, ordering => ordering, @@ -419,12 +407,11 @@ impl UserDefinedLogicalNodeCore for MergeScanLogicalPlan { impl MergeScanLogicalPlan { /// Creates a merge scan with the input plan's schema. - pub fn new(input: LogicalPlan, is_placeholder: bool, partition_cols: AliasMapping) -> Self { + pub fn new(input: LogicalPlan, is_placeholder: bool) -> Self { Self { output_schema: input.schema().clone(), input, is_placeholder, - partition_cols, remote_dyn_filter_producer_id: None, } } @@ -462,10 +449,6 @@ impl MergeScanLogicalPlan { &self.input } - pub fn partition_cols(&self) -> &AliasMapping { - &self.partition_cols - } - pub fn remote_dyn_filter_producer_id(&self) -> Option { self.remote_dyn_filter_producer_id } @@ -489,7 +472,6 @@ pub struct MergeScanExec { remote_dyn_filter_producer_id: Option, captured_remote_dyn_filters: Arc>>, target_partition: usize, - partition_cols: AliasMapping, enable_per_region_metrics: bool, } @@ -514,7 +496,6 @@ impl MergeScanExec { region_query_handler: RegionQueryHandlerRef, query_ctx: QueryContextRef, target_partition: usize, - partition_cols: AliasMapping, remote_dyn_filter_producer_id: Option, enable_per_region_metrics: bool, ) -> Result { @@ -551,24 +532,8 @@ impl MergeScanExec { EquivalenceProperties::new(arrow_schema.clone()) }; - let partition_exprs = partition_cols - .iter() - .filter_map(|col| { - if let Some(first_alias) = col.1.first() { - session_state - .create_physical_expr( - Expr::Column(ColumnExpr::new_unqualified( - first_alias.name().to_string(), - )), - plan.schema(), - ) - .ok() - } else { - None - } - }) - .collect(); - let partitioning = Partitioning::Hash(partition_exprs, output_partition_count); + // Region stripes are positional, not hash buckets; do not advertise a hash distribution. + let partitioning = Partitioning::UnknownPartitioning(output_partition_count); let properties = Arc::new(PlanProperties::new( eq_properties, @@ -590,7 +555,6 @@ impl MergeScanExec { remote_dyn_filter_producer_id, captured_remote_dyn_filters: Arc::default(), target_partition, - partition_cols, enable_per_region_metrics, }) } @@ -912,78 +876,6 @@ impl MergeScanExec { ))) } - pub fn try_with_new_distribution(&self, distribution: Distribution) -> Option { - let Distribution::KeyPartitioned(hash_exprs) = distribution else { - // not applicable - return None; - }; - - if let Partitioning::Hash(curr_dist, _) = &self.properties.partitioning - && curr_dist == &hash_exprs - { - // No need to change the distribution - return None; - } - - let hash_expr_col_names: HashSet<_> = hash_exprs - .iter() - .filter_map(|expr| { - expr.downcast_ref::() - .map(|col_expr| col_expr.name()) - }) - .collect(); - - let covers_all_partition_cols = self.partition_cols.values().all(|aliases| { - aliases - .iter() - .any(|col| hash_expr_col_names.contains(col.name())) - }); - if !covers_all_partition_cols { - return None; - } - - let all_partition_col_aliases: HashSet<_> = self - .partition_cols - .values() - .flat_map(|aliases| aliases.iter().map(|c| c.name())) - .collect(); - let overlaps: Vec<_> = hash_exprs - .iter() - .filter(|expr| { - expr.downcast_ref::() - .is_some_and(|col_expr| all_partition_col_aliases.contains(col_expr.name())) - }) - .cloned() - .collect(); - - if overlaps.is_empty() { - return None; - } - - Some(Self { - table: self.table.clone(), - regions: self.regions.clone(), - plan: self.plan.clone(), - arrow_schema: self.arrow_schema.clone(), - region_query_handler: self.region_query_handler.clone(), - metric: self.metric.clone(), - properties: Arc::new(PlanProperties::new( - self.properties.eq_properties.clone(), - Partitioning::Hash(overlaps, self.partition_count()), - self.properties.emission_type, - self.properties.boundedness, - )), - sub_stage_metrics: self.sub_stage_metrics.clone(), - partition_metrics: self.partition_metrics.clone(), - query_ctx: self.query_ctx.clone(), - remote_dyn_filter_producer_id: self.remote_dyn_filter_producer_id, - captured_remote_dyn_filters: self.captured_remote_dyn_filters.clone(), - target_partition: self.target_partition, - partition_cols: self.partition_cols.clone(), - enable_per_region_metrics: self.enable_per_region_metrics, - }) - } - fn captured_remote_dyn_filters(&self) -> Vec { self.captured_remote_dyn_filters.lock().unwrap().clone() } @@ -1037,13 +929,6 @@ impl MergeScanExec { } } -#[cfg(test)] -impl MergeScanExec { - fn remote_dyn_filter_producer_id(&self) -> Option { - self.remote_dyn_filter_producer_id - } -} - /// Metrics for a region of a partition. #[derive(Debug, Clone)] struct RegionMetrics { @@ -1420,7 +1305,7 @@ impl MergeScanMetric { #[cfg(test)] mod tests { - use std::collections::{BTreeSet, HashMap as StdHashMap}; + use std::collections::HashMap as StdHashMap; use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; @@ -1441,15 +1326,15 @@ mod tests { }; use datafusion::config::ConfigOptions; use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::Partitioning; use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::{StatisticsArgs, StatisticsContext}; - use datafusion_common::TableReference; use datafusion_expr::{LogicalPlanBuilder, col, lit}; + use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{ Column, DynamicFilterPhysicalExpr, lit as physical_lit, }; - use datafusion_physical_expr::{Distribution, PhysicalExpr}; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector}; @@ -1517,7 +1402,6 @@ mod tests { Arc::new(TestRegionQueryHandler::default()), QueryContext::arc(), target_partition, - AliasMapping::new(), None, false, ) @@ -1560,7 +1444,6 @@ mod tests { handler, query_ctx, 1, - AliasMapping::new(), Some(RemoteDynFilterProducerId::new(42)), false, ) @@ -1661,6 +1544,10 @@ mod tests { exec.properties().output_partitioning().partition_count(), expected ); + assert!(matches!( + exec.properties().output_partitioning(), + Partitioning::UnknownPartitioning(count) if *count == expected + )); } } @@ -1874,7 +1761,6 @@ mod tests { handler.clone(), query_ctx, 1, - AliasMapping::new(), Some(RemoteDynFilterProducerId::new(42)), false, ) @@ -2682,7 +2568,6 @@ mod tests { handler, QueryContext::arc(), target_partition, - AliasMapping::new(), None, false, ) @@ -3288,73 +3173,6 @@ mod tests { assert_int64_batch(&batches[0], (11, 12)); } - #[test] - fn try_with_new_distribution_preserves_remote_dyn_filter_producer_id() { - let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42); - - // Build a plan whose schema contains "col1" - let plan = LogicalPlanBuilder::empty(true) - .project(vec![lit(1i32).alias("col1")]) - .unwrap() - .build() - .unwrap(); - - let schema = plan.schema().as_arrow().clone(); - let table = TableName::new("catalog", "schema", "table"); - let regions = vec![RegionId::new(1024, 1)]; - let query_ctx = QueryContext::arc(); - - // Non-empty partition_cols so try_with_new_distribution can detect an overlap - let mut partition_cols = AliasMapping::new(); - partition_cols.insert( - "col1".to_string(), - BTreeSet::from([ColumnExpr::new(Some(TableReference::bare("table")), "col1")]), - ); - - let session_state = SessionStateBuilder::new().build(); - - let handler = Arc::new(TestRegionQueryHandler::default()); - let target_partition = 2; - - let exec = MergeScanExec::new( - &session_state, - table, - regions, - plan, - &schema, - handler, - query_ctx, - target_partition, - partition_cols, - Some(remote_dyn_filter_producer_id), - false, - ) - .unwrap(); - - assert_eq!( - exec.remote_dyn_filter_producer_id(), - Some(remote_dyn_filter_producer_id) - ); - - // A distribution that differs from the current partitioning but shares a - // column name present in partition_cols, so try_with_new_distribution - // produces a clone instead of returning None. - let new_dist = Distribution::KeyPartitioned(vec![ - Arc::new(Column::new("col1", 0)), - Arc::new(Column::new("col2", 1)), - ]); - - let cloned = exec - .try_with_new_distribution(new_dist) - .expect("expected a cloned exec with overlapping partition col"); - - assert_eq!( - cloned.remote_dyn_filter_producer_id(), - Some(remote_dyn_filter_producer_id), - "try_with_new_distribution must preserve remote dynamic filter producer id" - ); - } - #[test] fn merge_scan_apply_expressions_exposes_remote_dyn_filter_id() { let query_ctx = QueryContext::arc(); @@ -3397,7 +3215,6 @@ mod tests { handler, query_ctx, 1, - AliasMapping::new(), Some(remote_dyn_filter_producer_id), false, ) @@ -3503,7 +3320,6 @@ mod tests { Arc::new(TestRegionQueryHandler::default()), QueryContext::arc(), 1, - AliasMapping::new(), None, true, ) @@ -3572,7 +3388,6 @@ mod tests { Arc::new(TestRegionQueryHandler::default()), QueryContext::arc(), 1, - AliasMapping::new(), None, true, ) diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index b2a2ade9b4f..88b13caf388 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -228,7 +228,6 @@ impl ExtensionPlanner for DistExtensionPlanner { self.region_query_handler.clone(), query_ctx, session.config().target_partitions(), - merge_scan.partition_cols().clone(), merge_scan.remote_dyn_filter_producer_id(), self.enable_per_region_metrics, )?; diff --git a/src/query/src/metrics.rs b/src/query/src/metrics.rs index d6d62c99c4f..3c0043dcc09 100644 --- a/src/query/src/metrics.rs +++ b/src/query/src/metrics.rs @@ -407,7 +407,7 @@ fn finalize_region_watermarks(merged: BTreeMap) -> Vec>::new(), Some(RemoteDynFilterProducerId::new(0)), false, ) diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index fec001e7506..03d6ce3aa46 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -22,7 +22,6 @@ pub(crate) mod insert_assignment; pub(crate) mod json_schema_concretize; pub(crate) mod json_type_concretize; pub mod parallelize_scan; -pub mod pass_distribution; pub mod promql_tsid_narrow_join; pub mod remove_duplicate; pub mod scan_hint; diff --git a/src/query/src/optimizer/json_schema_concretize.rs b/src/query/src/optimizer/json_schema_concretize.rs index 12f71d44c31..9ea3531853d 100644 --- a/src/query/src/optimizer/json_schema_concretize.rs +++ b/src/query/src/optimizer/json_schema_concretize.rs @@ -192,8 +192,7 @@ mod tests { produce_one_row: false, schema, }); - let merge_scan = - MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + let merge_scan = MergeScanLogicalPlan::new(input, false).into_logical_plan(); let plan = LogicalPlanBuilder::from(merge_scan) .project(vec![col("j")])? .build()?; diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs deleted file mode 100644 index 732168e962f..00000000000 --- a/src/query/src/optimizer/pass_distribution.rs +++ /dev/null @@ -1,362 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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 std::sync::Arc; - -use datafusion::config::ConfigOptions; -use datafusion::physical_optimizer::PhysicalOptimizerRule; -use datafusion::physical_plan::projection::ProjectionExec; -use datafusion::physical_plan::repartition::RepartitionExec; -use datafusion::physical_plan::{ - ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, -}; -use datafusion_common::Result as DfResult; -use datafusion_physical_expr::Distribution; -use datafusion_physical_expr::utils::map_columns_before_projection; - -use crate::dist_plan::MergeScanExec; - -/// This is a [`PhysicalOptimizerRule`] to pass distribution requirement to -/// [`MergeScanExec`] to avoid unnecessary shuffling. -/// -/// This rule is expected to be run before [`EnforceDistribution`]. -/// -/// [`EnforceDistribution`]: datafusion::physical_optimizer::enforce_distribution::EnforceDistribution -/// [`MergeScanExec`]: crate::dist_plan::MergeScanExec -#[derive(Debug)] -pub struct PassDistribution; - -impl PhysicalOptimizerRule for PassDistribution { - fn optimize( - &self, - plan: Arc, - config: &ConfigOptions, - ) -> DfResult> { - Self::do_optimize(plan, config) - } - - fn name(&self) -> &str { - "PassDistributionRule" - } - - fn schema_check(&self) -> bool { - false - } -} - -impl PassDistribution { - fn do_optimize( - plan: Arc, - _config: &ConfigOptions, - ) -> DfResult> { - // Start from root with no requirement - Self::rewrite_with_distribution(plan, None) - } - - /// Top-down rewrite that propagates distribution requirements to children. - fn rewrite_with_distribution( - plan: Arc, - current_req: Option, - ) -> DfResult> { - // If this is a MergeScanExec, try to apply the current requirement. - if let Some(merge_scan) = plan.downcast_ref::() - && let Some(Distribution::KeyPartitioned(hash_exprs)) = current_req.as_ref() - { - if let Partitioning::Hash(current_hash_exprs, _) = &merge_scan.properties().partitioning - && *current_hash_exprs == *hash_exprs - { - return Ok(plan); - } - - if let Some(new_plan) = merge_scan - .try_with_new_distribution(Distribution::KeyPartitioned(hash_exprs.clone())) - { - // Leaf node; no children to process - return Ok(Arc::new(new_plan) as _); - } - - let partitioning = Partitioning::Hash( - hash_exprs.clone(), - merge_scan.properties().partitioning.partition_count(), - ); - return Ok(Arc::new(RepartitionExec::try_new(plan, partitioning)?) as _); - } - - // Compute per-child requirements from the current node. - let children = plan.children(); - if children.is_empty() { - return Ok(plan); - } - - let required = plan.input_distribution_requirements(); - let mut new_children = Vec::with_capacity(children.len()); - for (idx, child) in children.into_iter().enumerate() { - let child_req = match required.child_distribution(idx) { - Some(Distribution::UnspecifiedDistribution) if idx == 0 => { - Self::map_hash_requirement_through_projection(plan.as_ref(), ¤t_req) - } - Some(Distribution::UnspecifiedDistribution) => None, - None => current_req.clone(), - Some(req) => Some(req.clone()), - }; - let new_child = Self::rewrite_with_distribution(child.clone(), child_req)?; - new_children.push(new_child); - } - - // Rebuild the node only if any child changed (pointer inequality) - let unchanged = plan - .children() - .into_iter() - .zip(new_children.iter()) - .all(|(old, new)| Arc::ptr_eq(old, new)); - if unchanged { - Ok(plan) - } else { - plan.replace_children( - new_children, - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - } - } - - fn map_hash_requirement_through_projection( - plan: &dyn ExecutionPlan, - current_req: &Option, - ) -> Option { - let Some(Distribution::KeyPartitioned(required_exprs)) = current_req else { - return None; - }; - - let projection = plan.downcast_ref::()?; - let proj_exprs = projection - .expr() - .iter() - .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())) - .collect::>(); - let mapped = map_columns_before_projection(required_exprs, &proj_exprs); - - (mapped.len() == required_exprs.len()).then_some(Distribution::KeyPartitioned(mapped)) - } -} - -#[cfg(test)] -mod tests { - use std::collections::{BTreeMap, BTreeSet}; - - use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; - use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; - use async_trait::async_trait; - use common_query::request::QueryRequest; - use datafusion::common::NullEquality; - use datafusion::execution::SessionStateBuilder; - use datafusion::physical_optimizer::PhysicalOptimizerRule; - use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; - use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; - use datafusion::physical_plan::{ExecutionPlanProperties, Partitioning}; - use datafusion_expr::{JoinType, LogicalPlanBuilder}; - use datafusion_physical_expr::PhysicalExpr; - use datafusion_physical_expr::expressions::Column as PhysicalColumn; - use session::ReadPreference; - use session::context::QueryContext; - use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; - use store_api::storage::RegionId; - use table::table_name::TableName; - - use super::*; - use crate::dist_plan::RemoteDynFilterProducerId; - use crate::error::Result as QueryResult; - use crate::region_query::RegionQueryHandler; - - struct NoopRegionQueryHandler; - - #[async_trait] - impl RegionQueryHandler for NoopRegionQueryHandler { - async fn select_target( - &self, - _read_preference: ReadPreference, - _region_id: RegionId, - ) -> QueryResult { - unreachable!("pass distribution tests should not execute remote queries") - } - - async fn do_get( - &self, - _target: &crate::region_query::RegionQueryTarget, - _request: QueryRequest, - ) -> QueryResult { - unreachable!("pass distribution tests should not execute remote queries") - } - - async fn handle_remote_dyn_filter_update( - &self, - _target: &crate::region_query::RegionQueryTarget, - _query_id: String, - _update: RemoteDynFilterUpdate, - ) -> QueryResult<()> { - unreachable!("pass distribution tests should not send remote dyn filter updates") - } - - async fn handle_remote_dyn_filter_unregister( - &self, - _target: &crate::region_query::RegionQueryTarget, - _query_id: String, - _unregister: RemoteDynFilterUnregister, - ) -> QueryResult<()> { - unreachable!("pass distribution tests should not send remote dyn filter unregisters") - } - } - - #[test] - fn passes_hash_requirement_through_projection_to_merge_scan() { - let schema = test_schema(); - let left_merge_scan = Arc::new(test_merge_scan_exec(schema.clone())); - let right_merge_scan = Arc::new(test_merge_scan_exec(schema.clone())); - let left_projection = Arc::new( - ProjectionExec::try_new( - vec![ - ProjectionExpr::new(partition_column("greptime_value", 3), "greptime_value"), - ProjectionExpr::new( - partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1), - DATA_SCHEMA_TSID_COLUMN_NAME, - ), - ProjectionExpr::new( - partition_column("greptime_timestamp", 2), - "greptime_timestamp", - ), - ], - left_merge_scan, - ) - .unwrap(), - ) as Arc; - let join = Arc::new( - HashJoinExec::try_new( - left_projection, - right_merge_scan, - vec![ - ( - partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1), - partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1), - ), - ( - partition_column("greptime_timestamp", 2), - partition_column("greptime_timestamp", 2), - ), - ], - None, - &JoinType::Inner, - None, - PartitionMode::Partitioned, - NullEquality::NullEqualsNull, - false, - ) - .unwrap(), - ) as Arc; - - let optimized = PassDistribution - .optimize(join, &ConfigOptions::default()) - .unwrap(); - let hash_join = optimized.downcast_ref::().unwrap(); - let left_projection = hash_join.left().downcast_ref::().unwrap(); - let left_partitioning = left_projection.input().output_partitioning(); - let right_partitioning = hash_join.right().output_partitioning(); - - let Partitioning::Hash(left_exprs, left_count) = left_partitioning else { - panic!("expected left merge scan hash partitioning"); - }; - let Partitioning::Hash(right_exprs, right_count) = right_partitioning else { - panic!("expected right merge scan hash partitioning"); - }; - - assert_eq!(*left_count, 2); - assert_eq!(*right_count, 2); - assert_eq!( - column_names(left_exprs), - vec![DATA_SCHEMA_TSID_COLUMN_NAME, "greptime_timestamp"] - ); - assert_eq!( - column_names(right_exprs), - vec![DATA_SCHEMA_TSID_COLUMN_NAME, "greptime_timestamp"] - ); - } - - #[test] - fn merge_scan_rejects_hash_requirement_on_partition_key_subset() { - let merge_scan = test_merge_scan_exec(test_schema()); - - let new_plan = merge_scan.try_with_new_distribution(Distribution::KeyPartitioned(vec![ - partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1), - ])); - - assert!( - new_plan.is_none(), - "partitioning by a subset of multi-column partition keys is not sufficient" - ); - } - - fn test_merge_scan_exec(schema: SchemaRef) -> MergeScanExec { - let session_state = SessionStateBuilder::new().with_default_features().build(); - let partition_cols = BTreeMap::from([ - ( - DATA_SCHEMA_TSID_COLUMN_NAME.to_string(), - BTreeSet::from([datafusion_common::Column::from_name( - DATA_SCHEMA_TSID_COLUMN_NAME, - )]), - ), - ( - "greptime_timestamp".to_string(), - BTreeSet::from([datafusion_common::Column::from_name("greptime_timestamp")]), - ), - ]); - let plan = LogicalPlanBuilder::empty(false).build().unwrap(); - - MergeScanExec::new( - &session_state, - TableName::new("greptime", "public", "test"), - vec![RegionId::new(1, 0), RegionId::new(1, 1)], - plan, - schema.as_ref(), - Arc::new(NoopRegionQueryHandler), - QueryContext::arc(), - 32, - partition_cols, - Some(RemoteDynFilterProducerId::new(1)), - false, - ) - .unwrap() - } - - fn test_schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new(DATA_SCHEMA_TSID_COLUMN_NAME, DataType::UInt64, false), - Field::new( - "greptime_timestamp", - DataType::Timestamp(TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", DataType::Float64, true), - ])) - } - - fn partition_column(name: &str, index: usize) -> Arc { - Arc::new(PhysicalColumn::new(name, index)) - } - - fn column_names(exprs: &[Arc]) -> Vec<&str> { - exprs - .iter() - .map(|expr| expr.downcast_ref::().unwrap().name()) - .collect() - } -} diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index c924a57a68e..63ee0856aaf 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -220,11 +220,7 @@ impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { let input = self.decode_payload(merge_scan.input).await?; - // `PbMergeScan` lacks `partition_cols`; decoded plans lose this optimization and may repartition. - Ok( - MergeScanLogicalPlan::new(input, merge_scan.is_placeholder, Default::default()) - .into_logical_plan(), - ) + Ok(MergeScanLogicalPlan::new(input, merge_scan.is_placeholder).into_logical_plan()) } // Keep child dispatch on this consumer; delegating the parent bypasses async MergeScan decoding. @@ -739,8 +735,6 @@ mod tests { .expect("Expect a MergeScan plan node"); assert!(merge_scan.is_placeholder()); assert_eq!(merge_scan.input().to_string(), numbers_scan().to_string()); - // `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping. - assert!(merge_scan.partition_cols().is_empty()); } /// Payload table binding uses the engine catalog, not the request catalog. @@ -979,8 +973,7 @@ mod tests { .unwrap() .build() .unwrap(); - let plan = - MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + let plan = MergeScanLogicalPlan::new(input, false).into_logical_plan(); let bytes = DFLogicalSubstraitConvertor .encode(&plan, DefaultSerializer) @@ -1041,7 +1034,7 @@ mod tests { } fn merge_scan(input: LogicalPlan, is_placeholder: bool) -> LogicalPlan { - MergeScanLogicalPlan::new(input, is_placeholder, Default::default()).into_logical_plan() + MergeScanLogicalPlan::new(input, is_placeholder).into_logical_plan() } fn encode_plan(plan: &LogicalPlan) -> Bytes { diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index aa3b9417b14..e6db7db10f2 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -70,7 +70,6 @@ use crate::optimizer::global_limit::EnsureGlobalLimitForFetch; use crate::optimizer::json_schema_concretize::JsonSchemaConcretizeRule; use crate::optimizer::json_type_concretize::JsonTypeConcretizeRule; use crate::optimizer::parallelize_scan::ParallelizeScan; -use crate::optimizer::pass_distribution::PassDistribution; use crate::optimizer::promql_tsid_narrow_join::PromqlTsidNarrowJoin; use crate::optimizer::remove_duplicate::RemoveDuplicate; use crate::optimizer::scan_hint::ScanHintRule; @@ -232,17 +231,13 @@ impl QueryEngineState { physical_optimizer .rules .insert(5, Arc::new(ParallelizeScan)); - // Pass distribution requirement to MergeScanExec to avoid unnecessary shuffling - physical_optimizer - .rules - .insert(6, Arc::new(PassDistribution)); // Prefer collecting narrow PromQL build sides over repartitioning wide label streams. physical_optimizer .rules - .insert(7, Arc::new(PromqlTsidNarrowJoin)); + .insert(6, Arc::new(PromqlTsidNarrowJoin)); // Re-enforce sorting after custom rules update scan partitioning and distribution. // Keep it immediately before DataFusion's default EnsureRequirements. - physical_optimizer.rules.insert(8, Arc::new(EnforceSorting)); + physical_optimizer.rules.insert(7, Arc::new(EnforceSorting)); // Add rule for windowed sort physical_optimizer .rules diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index a1ccac27328..89706b53764 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -29,6 +29,7 @@ use frontend::error::Error; use frontend::instance::Instance; use operator::error::Error as OperatorError; use query::datafusion::QUERY_PARALLELISM_HINT; +use query::options::QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN; use rstest::rstest; use rstest_reuse::apply; use servers::error as server_error; @@ -257,6 +258,289 @@ WHERE greptime_timestamp = ( ); } +#[apply(both_instances_cases)] +async fn test_partitioned_merge_scan_left_join_repartitions(instance: Arc) { + let frontend = instance.frontend(); + execute_sql( + &frontend, + r#" +CREATE TABLE join_repro ( + k STRING, + kind STRING, + v DOUBLE, + ts TIMESTAMP NOT NULL, + TIME INDEX (ts) +) +PARTITION ON COLUMNS (k) ( + k < '1', + k >= '1' AND k < '2', + k >= '2' AND k < '3', + k >= '3' AND k < '4', + k >= '4' AND k < '5', + k >= '5' AND k < '6', + k >= '6' AND k < '7', + k >= '7' AND k < '8', + k >= '8' AND k < '9', + k >= '9' AND k < 'a', + k >= 'a' AND k < 'b', + k >= 'b' AND k < 'c', + k >= 'c' AND k < 'd', + k >= 'd' AND k < 'e', + k >= 'e' AND k < 'f', + k >= 'f' +) +ENGINE = mito +WITH (append_mode = 'true') +"#, + ) + .await; + + let mut rows = Vec::new(); + for key in [ + '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f', 'z', + ] { + rows.push(format!("('{key}', 'root', 1.0, '2024-01-01 00:00:00')")); + if key != 'z' { + for _ in 0..3 { + rows.push(format!("('{key}', 'child', 1.0, '2024-01-01 00:00:00')")); + } + } + } + execute_sql( + &frontend, + &format!("INSERT INTO join_repro VALUES {}", rows.join(",")), + ) + .await; + + let aggregate_sql = r#"SELECT count(*) roots, count(c.n) roots_with_children, sum(c.n) children +FROM (SELECT k FROM join_repro WHERE kind = 'root') r +LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' GROUP BY k) c ON c.k = r.k"#; + for parallelism in [16, 8] { + let output = execute_sql_with_query_parallelism(&frontend, aggregate_sql, parallelism) + .await + .data + .pretty_print() + .await; + assert_eq!( + output, + "+-------+---------------------+----------+\n| roots | roots_with_children | children |\n+-------+---------------------+----------+\n| 17 | 16 | 48 |\n+-------+---------------------+----------+" + ); + + let explain = execute_sql_with_query_parallelism( + &frontend, + &format!("EXPLAIN {aggregate_sql}"), + parallelism, + ) + .await + .data + .pretty_print() + .await; + assert!( + explain.contains("HashJoinExec: mode=Partitioned"), + "query_parallelism={parallelism} should use a partitioned hash join; explain:\n{explain}" + ); + assert!( + explain.contains("RepartitionExec: partitioning=Hash([k@0]"), + "query_parallelism={parallelism} should repartition on the join key; explain:\n{explain}" + ); + + let ordered_sql = r#"SELECT r.k, c.n +FROM (SELECT k FROM join_repro WHERE kind = 'root') r +LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' GROUP BY k) c ON c.k = r.k +ORDER BY r.k"#; + let ordered = execute_sql_with_query_parallelism(&frontend, ordered_sql, parallelism) + .await + .data + .pretty_print() + .await; + assert_eq!( + ordered, + r#"+---+---+ +| k | n | ++---+---+ +| 0 | 3 | +| 1 | 3 | +| 2 | 3 | +| 3 | 3 | +| 4 | 3 | +| 5 | 3 | +| 6 | 3 | +| 7 | 3 | +| 8 | 3 | +| 9 | 3 | +| a | 3 | +| b | 3 | +| c | 3 | +| d | 3 | +| e | 3 | +| f | 3 | +| z | | ++---+---+"# + ); + + let pruned_sql = r#"SELECT count(*) roots, count(c.n) roots_with_children, sum(c.n) children +FROM (SELECT k FROM join_repro WHERE kind = 'root' AND k >= '1' AND k < '9') r +LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' AND k < '8' GROUP BY k) c ON c.k = r.k"#; + let pruned = execute_sql_with_query_parallelism(&frontend, pruned_sql, parallelism) + .await + .data + .pretty_print() + .await; + assert_eq!( + pruned, + "+-------+---------------------+----------+\n| roots | roots_with_children | children |\n+-------+---------------------+----------+\n| 8 | 7 | 21 |\n+-------+---------------------+----------+" + ); + + let pruned_explain = execute_sql_with_query_parallelism( + &frontend, + &format!("EXPLAIN {pruned_sql}"), + parallelism, + ) + .await + .data + .pretty_print() + .await; + assert!( + pruned_explain.contains("HashJoinExec: mode=Partitioned"), + "query_parallelism={parallelism} pruned join should remain partitioned; explain:\n{pruned_explain}" + ); + let hash_repartition_lines = pruned_explain + .lines() + .filter(|line| line.contains("RepartitionExec: partitioning=Hash([k@0]")) + .count(); + assert_eq!( + hash_repartition_lines, 2, + "query_parallelism={parallelism} pruned join should repartition both inputs on k; explain:\n{pruned_explain}" + ); + } +} + +#[apply(both_instances_cases)] +async fn test_partitioned_merge_scan_semi_join_repartitions(instance: Arc) { + let frontend = instance.frontend(); + execute_sql( + &frontend, + r#" +CREATE TABLE semi_join_repro ( + ts TIMESTAMP TIME INDEX, + key_id INT, + is_match BOOLEAN, + PRIMARY KEY (key_id) +) +PARTITION ON COLUMNS (key_id) ( + key_id < 10, + key_id >= 10 AND key_id < 20, + key_id >= 20 AND key_id < 30, + key_id >= 30 +) +ENGINE = mito +"#, + ) + .await; + execute_sql( + &frontend, + r#"INSERT INTO semi_join_repro VALUES +('2026-10-09 00:00:00', 1, false), +('2026-10-09 00:00:01', 1, true), +('2026-10-09 00:00:00', 11, false), +('2026-10-09 00:00:01', 11, true), +('2026-10-09 00:00:00', 21, false), +('2026-10-09 00:00:01', 21, true), +('2026-10-09 00:00:00', 31, false), +('2026-10-09 00:00:01', 31, true)"#, + ) + .await; + + let query = r#"SELECT t.key_id, MIN(t.ts) AS first_ts +FROM semi_join_repro t +WHERE t.key_id IN ( + SELECT s.key_id FROM semi_join_repro s WHERE s.is_match = true +) +GROUP BY t.key_id +ORDER BY t.key_id"#; + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, "2"); + query_ctx.set_extension(QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN, "true"); + let query_ctx = Arc::new(query_ctx); + let output = execute_sql_with(&frontend, query, query_ctx.clone()) + .await + .data + .pretty_print() + .await; + let expected = r#"+--------+---------------------+ +| key_id | first_ts | ++--------+---------------------+ +| 1 | 2026-10-09T00:00:00 | +| 11 | 2026-10-09T00:00:00 | +| 21 | 2026-10-09T00:00:00 | +| 31 | 2026-10-09T00:00:00 | ++--------+---------------------+"#; + assert_eq!(output, expected); + + let explain = execute_sql_with(&frontend, &format!("EXPLAIN {query}"), query_ctx.clone()) + .await + .data + .pretty_print() + .await; + assert!( + explain.contains("HashJoinExec: mode=Partitioned, join_type=LeftSemi"), + "parallel IN query should use a partitioned left-semi join; explain:\n{explain}" + ); + assert!( + explain.contains("RepartitionExec: partitioning=Hash([key_id@1], 2)") + && explain.contains("RepartitionExec: partitioning=Hash([key_id@0], 2)"), + "parallel left-semi join should repartition both inputs on key_id; explain:\n{explain}" + ); + + let analyze_explain = execute_sql_with( + &frontend, + &format!("EXPLAIN ANALYZE VERBOSE {query}"), + query_ctx.clone(), + ) + .await + .data + .pretty_print() + .await; + let has_remote_dyn_filter = |explain: &str| { + explain.lines().any(|line| { + line.contains("SeqScan: region=") + && line.contains("\"dyn_filters\"") + && line.contains("DynamicFilter") + }) + }; + assert!( + has_remote_dyn_filter(&analyze_explain), + "expected remote region SeqScan dyn_filters to include DynamicFilter; explain:\n{analyze_explain}" + ); + + let mut query_ctx_without_rdf = QueryContext::with_db_name(None); + query_ctx_without_rdf.set_extension(QUERY_PARALLELISM_HINT, "2"); + query_ctx_without_rdf.set_extension(QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN, "false"); + let query_ctx_without_rdf = Arc::new(query_ctx_without_rdf); + let output_without_rdf = execute_sql_with(&frontend, query, query_ctx_without_rdf.clone()) + .await + .data + .pretty_print() + .await; + assert_eq!(output_without_rdf, expected); + + let analyze_explain_without_rdf = execute_sql_with( + &frontend, + &format!("EXPLAIN ANALYZE VERBOSE {query}"), + query_ctx_without_rdf, + ) + .await + .data + .pretty_print() + .await; + assert!( + !analyze_explain_without_rdf + .lines() + .any(|line| line.contains("SeqScan: region=") && line.contains("\"dyn_filters\"")), + "expected no region SeqScan dyn_filters line when remote dynamic filter pushdown is disabled; explain:\n{analyze_explain_without_rdf}" + ); +} + #[apply(both_instances_cases)] async fn test_show_create_table(instance: Arc) { let frontend = instance.frontend(); diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 2c86589112d..6ad57b7a14c 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -547,7 +547,7 @@ fn nested_plan_of(probe: &LogicalPlan, build: LogicalPlan) -> LogicalPlan { info!("probe side of the nested plan:\n{probe}"); info!("build side of the nested plan:\n{build}"); - let inner = MergeScanLogicalPlan::new(build, false, Default::default()).into_logical_plan(); + let inner = MergeScanLogicalPlan::new(build, false).into_logical_plan(); let plan = LogicalPlanBuilder::from(probe.clone()) .join( inner, diff --git a/tests/cases/distributed/explain/subqueries.result b/tests/cases/distributed/explain/subqueries.result index 28fe75df516..3d3f24777db 100644 --- a/tests/cases/distributed/explain/subqueries.result +++ b/tests/cases/distributed/explain/subqueries.result @@ -517,9 +517,10 @@ EXPLAIN SELECT a FROM ((SELECT a FROM t1) UNION ALL (SELECT a FROM t2)) u ORDER |_|_TableScan: t2_| |_| ]]_| | physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST]_| +|_|_UnionExec_| |_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true]_| -|_|_InterleaveExec_| |_|_MergeScanExec: REDACTED +|_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/standalone/common/join/merge_scan_left_join.result b/tests/cases/standalone/common/join/merge_scan_left_join.result new file mode 100644 index 00000000000..b50ff0f5bd5 --- /dev/null +++ b/tests/cases/standalone/common/join/merge_scan_left_join.result @@ -0,0 +1,148 @@ +-- Regression coverage for issues 9456/9492 (PR 9491): preserve merge-scan +-- results for a pruned left join over partitioned append data. +CREATE TABLE merge_left_data ( + k STRING, + kind STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (k, kind) +) +PARTITION ON COLUMNS (k) ( + k < '1', k >= '1' AND k < '2', k >= '2' AND k < '3', + k >= '3' AND k < '4', k >= '4' AND k < '5', k >= '5' AND k < '6', + k >= '6' AND k < '7', k >= '7' AND k < '8', k >= '8' AND k < '9', + k >= '9' AND k < 'a', k >= 'a' AND k < 'b', k >= 'b' AND k < 'c', + k >= 'c' AND k < 'd', k >= 'd' AND k < 'e', k >= 'e' AND k < 'f', + k >= 'f' +) +ENGINE=mito WITH (append_mode = 'true'); + +Affected Rows: 0 + +INSERT INTO merge_left_data VALUES + ('0', 'root', '2026-10-09 00:00:00'), + ('1', 'root', '2026-10-09 00:00:00'), + ('2', 'root', '2026-10-09 00:00:00'), + ('3', 'root', '2026-10-09 00:00:00'), + ('4', 'root', '2026-10-09 00:00:00'), + ('5', 'root', '2026-10-09 00:00:00'), + ('6', 'root', '2026-10-09 00:00:00'), + ('7', 'root', '2026-10-09 00:00:00'), + ('8', 'root', '2026-10-09 00:00:00'), + ('9', 'root', '2026-10-09 00:00:00'), + ('a', 'root', '2026-10-09 00:00:00'), + ('b', 'root', '2026-10-09 00:00:00'), + ('c', 'root', '2026-10-09 00:00:00'), + ('d', 'root', '2026-10-09 00:00:00'), + ('e', 'root', '2026-10-09 00:00:00'), + ('f', 'root', '2026-10-09 00:00:00'), + ('z', 'root', '2026-10-09 00:00:00'), + ('0', 'child', '2026-10-09 00:00:00'), ('0', 'child', '2026-10-09 00:00:00'), ('0', 'child', '2026-10-09 00:00:00'), + ('1', 'child', '2026-10-09 00:00:00'), ('1', 'child', '2026-10-09 00:00:00'), ('1', 'child', '2026-10-09 00:00:00'), + ('2', 'child', '2026-10-09 00:00:00'), ('2', 'child', '2026-10-09 00:00:00'), ('2', 'child', '2026-10-09 00:00:00'), + ('3', 'child', '2026-10-09 00:00:00'), ('3', 'child', '2026-10-09 00:00:00'), ('3', 'child', '2026-10-09 00:00:00'), + ('4', 'child', '2026-10-09 00:00:00'), ('4', 'child', '2026-10-09 00:00:00'), ('4', 'child', '2026-10-09 00:00:00'), + ('5', 'child', '2026-10-09 00:00:00'), ('5', 'child', '2026-10-09 00:00:00'), ('5', 'child', '2026-10-09 00:00:00'), + ('6', 'child', '2026-10-09 00:00:00'), ('6', 'child', '2026-10-09 00:00:00'), ('6', 'child', '2026-10-09 00:00:00'), + ('7', 'child', '2026-10-09 00:00:00'), ('7', 'child', '2026-10-09 00:00:00'), ('7', 'child', '2026-10-09 00:00:00'), + ('8', 'child', '2026-10-09 00:00:00'), ('8', 'child', '2026-10-09 00:00:00'), ('8', 'child', '2026-10-09 00:00:00'), + ('9', 'child', '2026-10-09 00:00:00'), ('9', 'child', '2026-10-09 00:00:00'), ('9', 'child', '2026-10-09 00:00:00'), + ('a', 'child', '2026-10-09 00:00:00'), ('a', 'child', '2026-10-09 00:00:00'), ('a', 'child', '2026-10-09 00:00:00'), + ('b', 'child', '2026-10-09 00:00:00'), ('b', 'child', '2026-10-09 00:00:00'), ('b', 'child', '2026-10-09 00:00:00'), + ('c', 'child', '2026-10-09 00:00:00'), ('c', 'child', '2026-10-09 00:00:00'), ('c', 'child', '2026-10-09 00:00:00'), + ('d', 'child', '2026-10-09 00:00:00'), ('d', 'child', '2026-10-09 00:00:00'), ('d', 'child', '2026-10-09 00:00:00'), + ('e', 'child', '2026-10-09 00:00:00'), ('e', 'child', '2026-10-09 00:00:00'), ('e', 'child', '2026-10-09 00:00:00'), + ('f', 'child', '2026-10-09 00:00:00'), ('f', 'child', '2026-10-09 00:00:00'), ('f', 'child', '2026-10-09 00:00:00'); + +Affected Rows: 65 + +-- Full join: 17 roots, 16 matched roots, 48 child rows. +SELECT r.k, c.n +FROM (SELECT k FROM merge_left_data WHERE kind = 'root') r +LEFT JOIN ( + SELECT k, COUNT(*) AS n FROM merge_left_data WHERE kind = 'child' GROUP BY k +) c ON r.k = c.k +ORDER BY r.k; + ++---+---+ +| k | n | ++---+---+ +| 0 | 3 | +| 1 | 3 | +| 2 | 3 | +| 3 | 3 | +| 4 | 3 | +| 5 | 3 | +| 6 | 3 | +| 7 | 3 | +| 8 | 3 | +| 9 | 3 | +| a | 3 | +| b | 3 | +| c | 3 | +| d | 3 | +| e | 3 | +| f | 3 | +| z | | ++---+---+ + +-- Assert partitioned hash join and both hash repartition keys at SQLness N=4. +-- SQLNESS REPLACE (peers.*) peers=REDACTED +EXPLAIN SELECT r.k, c.n +FROM (SELECT k FROM merge_left_data WHERE kind = 'root' AND k >= '1' AND k < '5') r +LEFT JOIN ( + SELECT k, COUNT(*) AS n FROM merge_left_data + WHERE kind = 'child' AND k < '4' GROUP BY k +) c ON r.k = c.k +ORDER BY r.k; + ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: r.k ASC NULLS LAST | +| | Projection: r.k, c.n | +| | Left Join: r.k = c.k | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: r | +| | Projection: merge_left_data.k | +| | Filter: merge_left_data.kind = Utf8("root") AND merge_left_data.k >= Utf8("1") AND merge_left_data.k < Utf8("5") | +| | TableScan: merge_left_data, partial_filters=[merge_left_data.kind = Utf8("root"), merge_left_data.k >= Utf8("1"), merge_left_data.k < Utf8("5")] | +| | ]] | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: c | +| | Projection: merge_left_data.k, count(Int64(1)) AS count(*) AS n | +| | Aggregate: groupBy=[[merge_left_data.k]], aggr=[[count(merge_left_data.ts) AS count(Int64(1))]] | +| | Filter: merge_left_data.kind = Utf8("child") AND merge_left_data.k < Utf8("4") AND merge_left_data.k IS NOT NULL | +| | TableScan: merge_left_data, partial_filters=[merge_left_data.kind = Utf8("child"), merge_left_data.k < Utf8("4"), merge_left_data.k IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [k@0 ASC NULLS LAST] | +| | SortExec: expr=[k@0 ASC NULLS LAST], preserve_partitioning=[true] | +| | HashJoinExec: mode=Partitioned, join_type=Left, on=[(k@0, k@0)], projection=[k@0, n@2] | +| | RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=4 | +| | MergeScanExec: peers=REDACTED +| | RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=4 | +| | MergeScanExec: peers=REDACTED +| | | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- Four roots, three matches, nine child rows; key 4 remains with NULL n. +SELECT r.k, c.n +FROM (SELECT k FROM merge_left_data WHERE kind = 'root' AND k >= '1' AND k < '5') r +LEFT JOIN ( + SELECT k, COUNT(*) AS n FROM merge_left_data + WHERE kind = 'child' AND k < '4' GROUP BY k +) c ON r.k = c.k +ORDER BY r.k; + ++---+---+ +| k | n | ++---+---+ +| 1 | 3 | +| 2 | 3 | +| 3 | 3 | +| 4 | | ++---+---+ + +DROP TABLE merge_left_data; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/join/merge_scan_left_join.sql b/tests/cases/standalone/common/join/merge_scan_left_join.sql new file mode 100644 index 00000000000..72b5f75db7b --- /dev/null +++ b/tests/cases/standalone/common/join/merge_scan_left_join.sql @@ -0,0 +1,81 @@ +-- Regression coverage for issues 9456/9492 (PR 9491): preserve merge-scan +-- results for a pruned left join over partitioned append data. +CREATE TABLE merge_left_data ( + k STRING, + kind STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (k, kind) +) +PARTITION ON COLUMNS (k) ( + k < '1', k >= '1' AND k < '2', k >= '2' AND k < '3', + k >= '3' AND k < '4', k >= '4' AND k < '5', k >= '5' AND k < '6', + k >= '6' AND k < '7', k >= '7' AND k < '8', k >= '8' AND k < '9', + k >= '9' AND k < 'a', k >= 'a' AND k < 'b', k >= 'b' AND k < 'c', + k >= 'c' AND k < 'd', k >= 'd' AND k < 'e', k >= 'e' AND k < 'f', + k >= 'f' +) +ENGINE=mito WITH (append_mode = 'true'); + +INSERT INTO merge_left_data VALUES + ('0', 'root', '2026-10-09 00:00:00'), + ('1', 'root', '2026-10-09 00:00:00'), + ('2', 'root', '2026-10-09 00:00:00'), + ('3', 'root', '2026-10-09 00:00:00'), + ('4', 'root', '2026-10-09 00:00:00'), + ('5', 'root', '2026-10-09 00:00:00'), + ('6', 'root', '2026-10-09 00:00:00'), + ('7', 'root', '2026-10-09 00:00:00'), + ('8', 'root', '2026-10-09 00:00:00'), + ('9', 'root', '2026-10-09 00:00:00'), + ('a', 'root', '2026-10-09 00:00:00'), + ('b', 'root', '2026-10-09 00:00:00'), + ('c', 'root', '2026-10-09 00:00:00'), + ('d', 'root', '2026-10-09 00:00:00'), + ('e', 'root', '2026-10-09 00:00:00'), + ('f', 'root', '2026-10-09 00:00:00'), + ('z', 'root', '2026-10-09 00:00:00'), + ('0', 'child', '2026-10-09 00:00:00'), ('0', 'child', '2026-10-09 00:00:00'), ('0', 'child', '2026-10-09 00:00:00'), + ('1', 'child', '2026-10-09 00:00:00'), ('1', 'child', '2026-10-09 00:00:00'), ('1', 'child', '2026-10-09 00:00:00'), + ('2', 'child', '2026-10-09 00:00:00'), ('2', 'child', '2026-10-09 00:00:00'), ('2', 'child', '2026-10-09 00:00:00'), + ('3', 'child', '2026-10-09 00:00:00'), ('3', 'child', '2026-10-09 00:00:00'), ('3', 'child', '2026-10-09 00:00:00'), + ('4', 'child', '2026-10-09 00:00:00'), ('4', 'child', '2026-10-09 00:00:00'), ('4', 'child', '2026-10-09 00:00:00'), + ('5', 'child', '2026-10-09 00:00:00'), ('5', 'child', '2026-10-09 00:00:00'), ('5', 'child', '2026-10-09 00:00:00'), + ('6', 'child', '2026-10-09 00:00:00'), ('6', 'child', '2026-10-09 00:00:00'), ('6', 'child', '2026-10-09 00:00:00'), + ('7', 'child', '2026-10-09 00:00:00'), ('7', 'child', '2026-10-09 00:00:00'), ('7', 'child', '2026-10-09 00:00:00'), + ('8', 'child', '2026-10-09 00:00:00'), ('8', 'child', '2026-10-09 00:00:00'), ('8', 'child', '2026-10-09 00:00:00'), + ('9', 'child', '2026-10-09 00:00:00'), ('9', 'child', '2026-10-09 00:00:00'), ('9', 'child', '2026-10-09 00:00:00'), + ('a', 'child', '2026-10-09 00:00:00'), ('a', 'child', '2026-10-09 00:00:00'), ('a', 'child', '2026-10-09 00:00:00'), + ('b', 'child', '2026-10-09 00:00:00'), ('b', 'child', '2026-10-09 00:00:00'), ('b', 'child', '2026-10-09 00:00:00'), + ('c', 'child', '2026-10-09 00:00:00'), ('c', 'child', '2026-10-09 00:00:00'), ('c', 'child', '2026-10-09 00:00:00'), + ('d', 'child', '2026-10-09 00:00:00'), ('d', 'child', '2026-10-09 00:00:00'), ('d', 'child', '2026-10-09 00:00:00'), + ('e', 'child', '2026-10-09 00:00:00'), ('e', 'child', '2026-10-09 00:00:00'), ('e', 'child', '2026-10-09 00:00:00'), + ('f', 'child', '2026-10-09 00:00:00'), ('f', 'child', '2026-10-09 00:00:00'), ('f', 'child', '2026-10-09 00:00:00'); + +-- Full join: 17 roots, 16 matched roots, 48 child rows. +SELECT r.k, c.n +FROM (SELECT k FROM merge_left_data WHERE kind = 'root') r +LEFT JOIN ( + SELECT k, COUNT(*) AS n FROM merge_left_data WHERE kind = 'child' GROUP BY k +) c ON r.k = c.k +ORDER BY r.k; + +-- Assert partitioned hash join and both hash repartition keys at SQLness N=4. +-- SQLNESS REPLACE (peers.*) peers=REDACTED +EXPLAIN SELECT r.k, c.n +FROM (SELECT k FROM merge_left_data WHERE kind = 'root' AND k >= '1' AND k < '5') r +LEFT JOIN ( + SELECT k, COUNT(*) AS n FROM merge_left_data + WHERE kind = 'child' AND k < '4' GROUP BY k +) c ON r.k = c.k +ORDER BY r.k; + +-- Four roots, three matches, nine child rows; key 4 remains with NULL n. +SELECT r.k, c.n +FROM (SELECT k FROM merge_left_data WHERE kind = 'root' AND k >= '1' AND k < '5') r +LEFT JOIN ( + SELECT k, COUNT(*) AS n FROM merge_left_data + WHERE kind = 'child' AND k < '4' GROUP BY k +) c ON r.k = c.k +ORDER BY r.k; + +DROP TABLE merge_left_data; diff --git a/tests/cases/standalone/common/join/merge_scan_semi_join.result b/tests/cases/standalone/common/join/merge_scan_semi_join.result new file mode 100644 index 00000000000..5eb70b1241b --- /dev/null +++ b/tests/cases/standalone/common/join/merge_scan_semi_join.result @@ -0,0 +1,104 @@ +-- Regression coverage for issues 9456/9492 (PR 9491): preserve IN/semi-join +-- results when merge-scan inputs are distributed across hash partitions. +CREATE TABLE merge_semi_data ( + ts TIMESTAMP TIME INDEX, + key_id INT, + is_match BOOLEAN, + PRIMARY KEY (key_id) +) +PARTITION ON COLUMNS (key_id) ( + key_id < 10, + key_id >= 10 AND key_id < 20, + key_id >= 20 AND key_id < 30, + key_id >= 30 +) +ENGINE=mito; + +Affected Rows: 0 + +-- Three keys in each of four ranges; each key has a false row then a true row. +INSERT INTO merge_semi_data VALUES + ('2026-10-09 00:00:00', 1, false), ('2026-10-09 00:00:01', 1, true), + ('2026-10-09 00:00:00', 2, false), ('2026-10-09 00:00:01', 2, true), + ('2026-10-09 00:00:00', 3, false), ('2026-10-09 00:00:01', 3, true), + ('2026-10-09 00:00:00', 11, false), ('2026-10-09 00:00:01', 11, true), + ('2026-10-09 00:00:00', 12, false), ('2026-10-09 00:00:01', 12, true), + ('2026-10-09 00:00:00', 13, false), ('2026-10-09 00:00:01', 13, true), + ('2026-10-09 00:00:00', 21, false), ('2026-10-09 00:00:01', 21, true), + ('2026-10-09 00:00:00', 22, false), ('2026-10-09 00:00:01', 22, true), + ('2026-10-09 00:00:00', 23, false), ('2026-10-09 00:00:01', 23, true), + ('2026-10-09 00:00:00', 31, false), ('2026-10-09 00:00:01', 31, true), + ('2026-10-09 00:00:00', 32, false), ('2026-10-09 00:00:01', 32, true), + ('2026-10-09 00:00:00', 33, false), ('2026-10-09 00:00:01', 33, true); + +Affected Rows: 24 + +-- Check the plan retains the partitioned left-semi join and hash keys (N=4). +-- SQLNESS REPLACE (peers.*) peers=REDACTED +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: merge_semi_data.key_id ASC NULLS LAST | +| | Projection: merge_semi_data.key_id, min(merge_semi_data.ts) AS first_ts | +| | Aggregate: groupBy=[[merge_semi_data.key_id]], aggr=[[min(merge_semi_data.ts)]] | +| | LeftSemi Join: merge_semi_data.key_id = __correlated_sq_1.key_id | +| | Projection: merge_semi_data.ts, merge_semi_data.key_id | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | Filter: merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: __correlated_sq_1 | +| | Projection: merge_semi_data.key_id | +| | Filter: merge_semi_data.is_match AND merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.is_match, merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [key_id@0 ASC NULLS LAST] | +| | ProjectionExec: expr=[key_id@0 as key_id, min(merge_semi_data.ts)@1 as first_ts] | +| | SortExec: expr=[key_id@0 ASC NULLS LAST], preserve_partitioning=[true] | +| | AggregateExec: mode=SinglePartitioned, gby=[key_id@1 as key_id], aggr=[min(merge_semi_data.ts)] | +| | HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(key_id@1, key_id@0)] | +| | RepartitionExec: partitioning=Hash([key_id@1], 4), input_partitions=4 | +| | ProjectionExec: expr=[ts@0 as ts, key_id@1 as key_id] | +| | MergeScanExec: peers=REDACTED +| | RepartitionExec: partitioning=Hash([key_id@0], 4), input_partitions=4 | +| | MergeScanExec: peers=REDACTED +| | | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ + +-- Every key matches and reports its earliest (false-row) timestamp. +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++--------+---------------------+ +| key_id | first_ts | ++--------+---------------------+ +| 1 | 2026-10-09T00:00:00 | +| 2 | 2026-10-09T00:00:00 | +| 3 | 2026-10-09T00:00:00 | +| 11 | 2026-10-09T00:00:00 | +| 12 | 2026-10-09T00:00:00 | +| 13 | 2026-10-09T00:00:00 | +| 21 | 2026-10-09T00:00:00 | +| 22 | 2026-10-09T00:00:00 | +| 23 | 2026-10-09T00:00:00 | +| 31 | 2026-10-09T00:00:00 | +| 32 | 2026-10-09T00:00:00 | +| 33 | 2026-10-09T00:00:00 | ++--------+---------------------+ + +DROP TABLE merge_semi_data; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/join/merge_scan_semi_join.sql b/tests/cases/standalone/common/join/merge_scan_semi_join.sql new file mode 100644 index 00000000000..34e7b226592 --- /dev/null +++ b/tests/cases/standalone/common/join/merge_scan_semi_join.sql @@ -0,0 +1,49 @@ +-- Regression coverage for issues 9456/9492 (PR 9491): preserve IN/semi-join +-- results when merge-scan inputs are distributed across hash partitions. +CREATE TABLE merge_semi_data ( + ts TIMESTAMP TIME INDEX, + key_id INT, + is_match BOOLEAN, + PRIMARY KEY (key_id) +) +PARTITION ON COLUMNS (key_id) ( + key_id < 10, + key_id >= 10 AND key_id < 20, + key_id >= 20 AND key_id < 30, + key_id >= 30 +) +ENGINE=mito; + +-- Three keys in each of four ranges; each key has a false row then a true row. +INSERT INTO merge_semi_data VALUES + ('2026-10-09 00:00:00', 1, false), ('2026-10-09 00:00:01', 1, true), + ('2026-10-09 00:00:00', 2, false), ('2026-10-09 00:00:01', 2, true), + ('2026-10-09 00:00:00', 3, false), ('2026-10-09 00:00:01', 3, true), + ('2026-10-09 00:00:00', 11, false), ('2026-10-09 00:00:01', 11, true), + ('2026-10-09 00:00:00', 12, false), ('2026-10-09 00:00:01', 12, true), + ('2026-10-09 00:00:00', 13, false), ('2026-10-09 00:00:01', 13, true), + ('2026-10-09 00:00:00', 21, false), ('2026-10-09 00:00:01', 21, true), + ('2026-10-09 00:00:00', 22, false), ('2026-10-09 00:00:01', 22, true), + ('2026-10-09 00:00:00', 23, false), ('2026-10-09 00:00:01', 23, true), + ('2026-10-09 00:00:00', 31, false), ('2026-10-09 00:00:01', 31, true), + ('2026-10-09 00:00:00', 32, false), ('2026-10-09 00:00:01', 32, true), + ('2026-10-09 00:00:00', 33, false), ('2026-10-09 00:00:01', 33, true); + +-- Check the plan retains the partitioned left-semi join and hash keys (N=4). +-- SQLNESS REPLACE (peers.*) peers=REDACTED +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +-- Every key matches and reports its earliest (false-row) timestamp. +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +DROP TABLE merge_semi_data; diff --git a/tests/cases/standalone/common/order/order_by_exceptions.result b/tests/cases/standalone/common/order/order_by_exceptions.result index bdf215f8f26..dc964ebfe09 100644 --- a/tests/cases/standalone/common/order/order_by_exceptions.result +++ b/tests/cases/standalone/common/order/order_by_exceptions.result @@ -83,7 +83,7 @@ EXPLAIN SELECT a % 2, b FROM test UNION SELECT a % 2 AS k, b FROM test ORDER BY | | RepartitionExec: REDACTED | | AggregateExec: mode=Partial, gby=[test.a % Int64(2)@0 as test.a % Int64(2), b@1 as b], aggr=[] | | | RepartitionExec: REDACTED -| | InterleaveExec | +| | UnionExec | | | MergeScanExec: REDACTED | | MergeScanExec: REDACTED | | | diff --git a/tests/cases/standalone/common/tql-explain-analyze/explain.result b/tests/cases/standalone/common/tql-explain-analyze/explain.result index fed75f0448e..3e8b0ec2133 100644 --- a/tests/cases/standalone/common/tql-explain-analyze/explain.result +++ b/tests/cases/standalone/common/tql-explain-analyze/explain.result @@ -184,7 +184,6 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test; | physical_plan after LimitedDistinctAggregation_| SAME TEXT AS ABOVE_| | physical_plan after FilterPushdown_| SAME TEXT AS ABOVE_| | physical_plan after parallelize_scan_| SAME TEXT AS ABOVE_| -| physical_plan after PassDistributionRule_| SAME TEXT AS ABOVE_| | physical_plan after PromqlTsidNarrowJoin_| SAME TEXT AS ABOVE_| | physical_plan after EnforceSorting_| SAME TEXT AS ABOVE_| | physical_plan after WindowTopN_| SAME TEXT AS ABOVE_| @@ -341,7 +340,6 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test AS series; | physical_plan after LimitedDistinctAggregation_| SAME TEXT AS ABOVE_| | physical_plan after FilterPushdown_| SAME TEXT AS ABOVE_| | physical_plan after parallelize_scan_| SAME TEXT AS ABOVE_| -| physical_plan after PassDistributionRule_| SAME TEXT AS ABOVE_| | physical_plan after PromqlTsidNarrowJoin_| SAME TEXT AS ABOVE_| | physical_plan after EnforceSorting_| SAME TEXT AS ABOVE_| | physical_plan after WindowTopN_| SAME TEXT AS ABOVE_| @@ -528,7 +526,6 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; | physical_plan after LimitedDistinctAggregation_| SAME TEXT AS ABOVE_| | physical_plan after FilterPushdown_| SAME TEXT AS ABOVE_| | physical_plan after parallelize_scan_| SAME TEXT AS ABOVE_| -| physical_plan after PassDistributionRule_| SAME TEXT AS ABOVE_| | physical_plan after PromqlTsidNarrowJoin_| SAME TEXT AS ABOVE_| | physical_plan after EnforceSorting_| SAME TEXT AS ABOVE_| | physical_plan after WindowTopN_| SAME TEXT AS ABOVE_| diff --git a/tests/cases/standalone/common/tql/tql-cte.result b/tests/cases/standalone/common/tql/tql-cte.result index 7fbc621319c..c2bc149e66c 100644 --- a/tests/cases/standalone/common/tql/tql-cte.result +++ b/tests/cases/standalone/common/tql/tql-cte.result @@ -596,7 +596,7 @@ SELECT 'host2' as source, ts, cpu FROM host2_data; | | Filter: labels.host = Utf8("host2") AND labels.ts >= TimestampMillisecond(-299999, None) AND labels.ts <= TimestampMillisecond(40000, None) | | | TableScan: labels, partial_filters=[labels.host = Utf8("host2"), labels.ts >= TimestampMillisecond(-299999, None), labels.ts <= TimestampMillisecond(40000, None)] | | | ]] | -| physical_plan | InterleaveExec | +| physical_plan | UnionExec | | | CooperativeExec | | | MergeScanExec: REDACTED | | CooperativeExec | diff --git a/tests/perf/query_cases/mergescan_shuffle_9456/case.toml b/tests/perf/query_cases/mergescan_shuffle_9456/case.toml new file mode 100644 index 00000000000..d96bdc9eb01 --- /dev/null +++ b/tests/perf/query_cases/mergescan_shuffle_9456/case.toml @@ -0,0 +1,152 @@ +# Opt-in comparison of Hash RepartitionExec cost on a multi-region PromQL join. +# ANALYZE records plans and metrics; EVAL preserves complete results for comparison. + +[case] +name = "mergescan_shuffle_9456" +description = "Opt-in multi-region Partitioned PromQL shuffle cost measurement" +issue = "https://github.com/GreptimeTeam/greptimedb/issues/9456" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "mergescan_shuffle_9456" +physical_table = "mergescan_shuffle_phys" +series_count = 4096 +samples_per_series = 120 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 128 +timeout_seconds = 300 +visibility_timeout_seconds = 300 + +base_setup_sql = [ + """ +CREATE TABLE mergescan_shuffle_phys ( + greptime_timestamp TIMESTAMP(3) TIME INDEX, + host STRING, + instance STRING, + greptime_value DOUBLE, + PRIMARY KEY (host, instance) +) +PARTITION ON COLUMNS (instance) ( + instance < 'instance000128', + instance >= 'instance000128' AND instance < 'instance000256', + instance >= 'instance000256' AND instance < 'instance000384', + instance >= 'instance000384' AND instance < 'instance000512', + instance >= 'instance000512' AND instance < 'instance000640', + instance >= 'instance000640' AND instance < 'instance000768', + instance >= 'instance000768' AND instance < 'instance000896', + instance >= 'instance000896' AND instance < 'instance001024', + instance >= 'instance001024' AND instance < 'instance001152', + instance >= 'instance001152' AND instance < 'instance001280', + instance >= 'instance001280' AND instance < 'instance001408', + instance >= 'instance001408' AND instance < 'instance001536', + instance >= 'instance001536' AND instance < 'instance001664', + instance >= 'instance001664' AND instance < 'instance001792', + instance >= 'instance001792' AND instance < 'instance001920', + instance >= 'instance001920' AND instance < 'instance002048', + instance >= 'instance002048' AND instance < 'instance002176', + instance >= 'instance002176' AND instance < 'instance002304', + instance >= 'instance002304' AND instance < 'instance002432', + instance >= 'instance002432' AND instance < 'instance002560', + instance >= 'instance002560' AND instance < 'instance002688', + instance >= 'instance002688' AND instance < 'instance002816', + instance >= 'instance002816' AND instance < 'instance002944', + instance >= 'instance002944' AND instance < 'instance003072', + instance >= 'instance003072' AND instance < 'instance003200', + instance >= 'instance003200' AND instance < 'instance003328', + instance >= 'instance003328' AND instance < 'instance003456', + instance >= 'instance003456' AND instance < 'instance003584', + instance >= 'instance003584' AND instance < 'instance003712', + instance >= 'instance003712' AND instance < 'instance003840', + instance >= 'instance003840' AND instance < 'instance003968', + instance >= 'instance003968' +) ENGINE=metric WITH ('physical_metric_table'='true') +""", + """ +CREATE TABLE mergescan_shuffle_9456 ( + greptime_timestamp TIMESTAMP(3) TIME INDEX, + host STRING, + instance STRING, + greptime_value DOUBLE, + PRIMARY KEY (host, instance) +) ENGINE=metric WITH ('on_physical_table'='mergescan_shuffle_phys') +""", +] +candidate_setup_sql = [ + """ +CREATE TABLE mergescan_shuffle_phys ( + greptime_timestamp TIMESTAMP(3) TIME INDEX, + host STRING, + instance STRING, + greptime_value DOUBLE, + PRIMARY KEY (host, instance) +) +PARTITION ON COLUMNS (instance) ( + instance < 'instance000128', + instance >= 'instance000128' AND instance < 'instance000256', + instance >= 'instance000256' AND instance < 'instance000384', + instance >= 'instance000384' AND instance < 'instance000512', + instance >= 'instance000512' AND instance < 'instance000640', + instance >= 'instance000640' AND instance < 'instance000768', + instance >= 'instance000768' AND instance < 'instance000896', + instance >= 'instance000896' AND instance < 'instance001024', + instance >= 'instance001024' AND instance < 'instance001152', + instance >= 'instance001152' AND instance < 'instance001280', + instance >= 'instance001280' AND instance < 'instance001408', + instance >= 'instance001408' AND instance < 'instance001536', + instance >= 'instance001536' AND instance < 'instance001664', + instance >= 'instance001664' AND instance < 'instance001792', + instance >= 'instance001792' AND instance < 'instance001920', + instance >= 'instance001920' AND instance < 'instance002048', + instance >= 'instance002048' AND instance < 'instance002176', + instance >= 'instance002176' AND instance < 'instance002304', + instance >= 'instance002304' AND instance < 'instance002432', + instance >= 'instance002432' AND instance < 'instance002560', + instance >= 'instance002560' AND instance < 'instance002688', + instance >= 'instance002688' AND instance < 'instance002816', + instance >= 'instance002816' AND instance < 'instance002944', + instance >= 'instance002944' AND instance < 'instance003072', + instance >= 'instance003072' AND instance < 'instance003200', + instance >= 'instance003200' AND instance < 'instance003328', + instance >= 'instance003328' AND instance < 'instance003456', + instance >= 'instance003456' AND instance < 'instance003584', + instance >= 'instance003584' AND instance < 'instance003712', + instance >= 'instance003712' AND instance < 'instance003840', + instance >= 'instance003840' AND instance < 'instance003968', + instance >= 'instance003968' +) ENGINE=metric WITH ('physical_metric_table'='true') +""", + """ +CREATE TABLE mergescan_shuffle_9456 ( + greptime_timestamp TIMESTAMP(3) TIME INDEX, + host STRING, + instance STRING, + greptime_value DOUBLE, + PRIMARY KEY (host, instance) +) ENGINE=metric WITH ('on_physical_table'='mergescan_shuffle_phys') +""", +] + +[[scenario.queries]] +name = "partitioned_shuffle_cost" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704074340, '1m') rate(mergescan_shuffle_9456[5m]) / ignoring(host) rate(mergescan_shuffle_9456[5m])" +warmup = 1 +iterations = 7 + +[[scenario.queries]] +name = "collect_left_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704074340, '1m') rate(mergescan_shuffle_9456[5m]) / rate(mergescan_shuffle_9456[5m])" +warmup = 1 +iterations = 7 + +[[scenario.queries]] +name = "partitioned_complete_results" +kind = "tql" +query = "TQL EVAL (1704067200, 1704074340, '1m') rate(mergescan_shuffle_9456[5m]) / ignoring(host) rate(mergescan_shuffle_9456[5m])" +warmup = 0 +iterations = 1 From 28ec3ff5aad8d274d83ae42f96dc54f838fb6b20 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 10 Oct 2026 12:28:47 +0000 Subject: [PATCH 04/46] fix(table): require complete primary keys for hash partitioning (#9507) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/table/src/table/scan.rs | 100 +++++++++++++++++- .../aggregate/projected_pk_group_by.result | 69 ++++++++++++ .../aggregate/projected_pk_group_by.sql | 33 ++++++ 3 files changed, 200 insertions(+), 2 deletions(-) create mode 100644 tests/cases/standalone/common/aggregate/projected_pk_group_by.result create mode 100644 tests/cases/standalone/common/aggregate/projected_pk_group_by.sql diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 2db051a1e9c..3e0eeab84a9 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -213,7 +213,10 @@ impl RegionScanExec { Ok(tsid_col) => { Partitioning::Hash(vec![Arc::new(tsid_col) as _], num_output_partition) } - Err(_) => Partitioning::Hash(pk_columns.clone(), num_output_partition), + Err(_) if pk_columns.len() == pk_names.len() => { + Partitioning::Hash(pk_columns.clone(), num_output_partition) + } + Err(_) => Partitioning::UnknownPartitioning(num_output_partition), } } Some(TimeSeriesDistribution::TimeWindowed) | None => { @@ -694,7 +697,7 @@ mod test { use store_api::region_engine::{ PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties, SinglePartitionScanner, }; - use store_api::storage::RegionId; + use store_api::storage::{RegionId, TimeSeriesDistribution}; use super::*; @@ -857,6 +860,99 @@ mod test { ) } + #[test] + fn test_per_series_partitioning_requires_all_primary_keys() { + let full_schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("marker", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ])); + let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5683)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "marker", + ConcreteDataType::int32_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "host", + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 3, + }) + .primary_key(vec![1, 2]); + let metadata = Arc::new(builder.build().unwrap()); + + for (schema, expected) in [ + ( + full_schema.clone(), + Partitioning::Hash( + vec![ + Arc::new( + Column::new_with_schema("marker", full_schema.arrow_schema()).unwrap(), + ) as _, + Arc::new( + Column::new_with_schema("host", full_schema.arrow_schema()).unwrap(), + ) as _, + ], + 2, + ), + ), + ( + Arc::new(Schema::new(vec![ + ColumnSchema::new("marker", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ])), + Partitioning::UnknownPartitioning(2), + ), + ] { + let batches = RecordBatches::try_new(schema, vec![]).unwrap(); + let scanner = Box::new(RepeatableScanner { + batches, + properties: ScannerProperties::new(vec![vec![], vec![]], false, 0), + metadata: metadata.clone(), + dynamic_filters: Arc::new(Mutex::new(Vec::new())), + }); + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }; + let plan = RegionScanExec::new(scanner, request, None).unwrap(); + let actual = plan.properties().output_partitioning(); + match (&actual, &expected) { + ( + Partitioning::UnknownPartitioning(actual_count), + Partitioning::UnknownPartitioning(expected_count), + ) => assert_eq!(actual_count, expected_count), + _ => assert_eq!(actual, &expected), + } + } + } + #[test] fn test_count_statistics_require_exact_source_rows() { for (append_mode, exact, expected) in [ diff --git a/tests/cases/standalone/common/aggregate/projected_pk_group_by.result b/tests/cases/standalone/common/aggregate/projected_pk_group_by.result new file mode 100644 index 00000000000..7ca3f1e5826 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/projected_pk_group_by.result @@ -0,0 +1,69 @@ +CREATE TABLE projected_pk_group_by ( + marker TINYINT UNSIGNED, + host STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (marker, host) +); + +Affected Rows: 0 + +INSERT INTO projected_pk_group_by VALUES + (0, 'a', 1000), + (0, 'a', 2000), + (0, 'b', 1000), + (0, 'b', 2000), + (0, 'c', 1000), + (1, NULL, 1000); + +Affected Rows: 6 + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + ++--------+-----------+ +| marker | row_count | ++--------+-----------+ +| 0 | 5 | +| 1 | 1 | ++--------+-----------+ + +SELECT marker, host, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker, host +ORDER BY marker, host; + ++--------+------+-----------+ +| marker | host | row_count | ++--------+------+-----------+ +| 0 | a | 2 | +| 0 | b | 2 | +| 0 | c | 1 | +| 1 | | 1 | ++--------+------+-----------+ + +ADMIN FLUSH_TABLE('projected_pk_group_by'); + ++--------------------------------------------+ +| ADMIN FLUSH_TABLE('projected_pk_group_by') | ++--------------------------------------------+ +| 0 | ++--------------------------------------------+ + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + ++--------+-----------+ +| marker | row_count | ++--------+-----------+ +| 0 | 5 | +| 1 | 1 | ++--------+-----------+ + +DROP TABLE projected_pk_group_by; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/projected_pk_group_by.sql b/tests/cases/standalone/common/aggregate/projected_pk_group_by.sql new file mode 100644 index 00000000000..d525e16bc55 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/projected_pk_group_by.sql @@ -0,0 +1,33 @@ +CREATE TABLE projected_pk_group_by ( + marker TINYINT UNSIGNED, + host STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (marker, host) +); + +INSERT INTO projected_pk_group_by VALUES + (0, 'a', 1000), + (0, 'a', 2000), + (0, 'b', 1000), + (0, 'b', 2000), + (0, 'c', 1000), + (1, NULL, 1000); + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + +SELECT marker, host, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker, host +ORDER BY marker, host; + +ADMIN FLUSH_TABLE('projected_pk_group_by'); + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + +DROP TABLE projected_pk_group_by; From e6f3ca10a560b195e947f78320cf8b704eb6b908 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 12 Oct 2026 03:32:46 +0000 Subject: [PATCH 05/46] fix(deps): bump OpenDAL to 0.59.4 for OSS RRSA credentials (#9520) Signed-off-by: WenyXu --- Cargo.lock | 76 ++++++++++++++++++------------------- src/object-store/Cargo.toml | 4 +- 2 files changed, 40 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3b66b13e949..095d6d4acec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9686,9 +9686,9 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" [[package]] name = "opendal" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fe43e16d96bed57937eb7c4a28559bf4c5fc4a3951656505d1fdd4f856349b4" +checksum = "3cbd193aec6d62efaff87611dfae52c132aed650bb3cb0c684c0a9394f09a42a" dependencies = [ "ctor", "opendal-core", @@ -9711,9 +9711,9 @@ dependencies = [ [[package]] name = "opendal-core" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de4566a412776e3f65d53dd9fceced5b432a0a9c9a1e8d50ccf58823795e5c5b" +checksum = "54b5555d753dd6a6383dc0891254b72f3cdc40c011ba0372b560e0ae5aa29f5c" dependencies = [ "anyhow", "asyncband", @@ -9737,9 +9737,9 @@ dependencies = [ [[package]] name = "opendal-http-transport-reqwest" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f772ce6c137ab43647726116d505a649a88d41928fcd813438eb30b02312bcbb" +checksum = "4d6588678f3040a3eab1819364dcbc6a746a8286d1ceb23a70a6614faa01dc6c" dependencies = [ "bytes", "futures", @@ -9751,9 +9751,9 @@ dependencies = [ [[package]] name = "opendal-layer-concurrent-limit" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3734c648617e5ad724d49175dde8323550624f9614f50e9e81b7aafae30b1734" +checksum = "db8574309e68bf6d49c5378eaef9d0b3f6786971650b061c69e808254d71cb05" dependencies = [ "asyncband", "futures", @@ -9763,9 +9763,9 @@ dependencies = [ [[package]] name = "opendal-layer-logging" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a8daab5a91eac76f94516e7c71e677195c9221840ce8f5e71815e3fc8207742" +checksum = "4ad43246d2cb7b7ad32b10e22bfa7b6ced8a6677cb6d6f005ae7e1d94cb19c6a" dependencies = [ "log", "opendal-core", @@ -9773,9 +9773,9 @@ dependencies = [ [[package]] name = "opendal-layer-observe-metrics-common" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fbe664be117c293cfa3980a0d2fc238bbb38eae5f6c191f376e88ccb830d26b" +checksum = "3b8bc97d5a00eb41d50f1f6077c4bd5396c1eefea495614084ee9a37202ae06d" dependencies = [ "futures", "http 1.5.0", @@ -9784,9 +9784,9 @@ dependencies = [ [[package]] name = "opendal-layer-prometheus" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b32078d8248555ee6271332ec480b78306a44f0d902ea4721ccac00d9b866a21" +checksum = "404e766cc2e4e4c9225ce08605ae151f768a05a41a448a235e016b1aea9ff98f" dependencies = [ "opendal-core", "opendal-layer-observe-metrics-common", @@ -9795,9 +9795,9 @@ dependencies = [ [[package]] name = "opendal-layer-retry" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb8e22cb1f86ef9af771a8947d3ccf380c8f5b76236f628314dd41c9ba73331a" +checksum = "4aeb54be4940c68609483591934a813f7bdd97d747826783a393914d3139e363" dependencies = [ "backon", "log", @@ -9806,9 +9806,9 @@ dependencies = [ [[package]] name = "opendal-layer-timeout" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08ac6c77fe0d1b18e0064a02c19b98dbebcab78bf5e7753ad1f3abbc18850ff4" +checksum = "1d11d53e67473e8914eacb2e14d586779ae57e1a84759ad25232ef339db77d8a" dependencies = [ "opendal-core", "tokio", @@ -9816,9 +9816,9 @@ dependencies = [ [[package]] name = "opendal-layer-tracing" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "603807927a799ea981f62a3b0012e17181e39f468793b78ca8bedda66d19813c" +checksum = "bbe2ebfc795bf08da77d20bed13d149b478cc8e5002df44d237e460134f00be4" dependencies = [ "futures", "http 1.5.0", @@ -9828,9 +9828,9 @@ dependencies = [ [[package]] name = "opendal-service-azblob" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6cf31b02e7b44191d97c07254eea109b5da9c17210a29e478c8d0dc26dffb1f" +checksum = "c32c3fc2e507c854ce53615b28fc2208f5e00cbe456c4d344ffdfbf4c9cb976d" dependencies = [ "base64 0.23.1", "bytes", @@ -9849,9 +9849,9 @@ dependencies = [ [[package]] name = "opendal-service-azure-common" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8107324c3f3cb9970fe20526c2c341b7c5769acbd4699bab73bd525f5fd63d62" +checksum = "cc93c30bd347a36d93fd46d47f3e2bb6494dcf9458938aa501fe046ff8f7562f" dependencies = [ "http 1.5.0", "opendal-core", @@ -9859,9 +9859,9 @@ dependencies = [ [[package]] name = "opendal-service-fs" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0e7d8974aebd33be899200c9d0d50b2141fd5fa5c462b29e20653e3ddbfe339" +checksum = "ec5c16fe91b77c0364f66546a1c32b2999e9c3c3c75c2e908a282268e0890bac" dependencies = [ "bytes", "log", @@ -9873,9 +9873,9 @@ dependencies = [ [[package]] name = "opendal-service-gcs" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d80eca6f9227bdf65d43e5b8d5fe75dca8cfab1da07571442a0420400847bea" +checksum = "e1b5e42c6606ea1e37fa9cdedf07dc998073a45eb5044317c0fa6f3153e98d6e" dependencies = [ "async-trait", "bytes", @@ -9895,9 +9895,9 @@ dependencies = [ [[package]] name = "opendal-service-hdfs-native" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b84340e11a3bcd7f1f0b55df032d040b2907483c546c3021485aba749c3072c" +checksum = "20fb97b2ac06a87e6720f0f93d5107ac046117deb8ac957030b0962ba4ecc3c5" dependencies = [ "bytes", "futures", @@ -9909,9 +9909,9 @@ dependencies = [ [[package]] name = "opendal-service-http" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ac1553d6b57ca79e200e5c7ded9ab708f2016dadb28ede0f2da71e02fd0b960" +checksum = "65fdde60d1abf14e47eaf19360c08ea1b89c14f03769db448e016781c2e0e32b" dependencies = [ "http 1.5.0", "log", @@ -9921,9 +9921,9 @@ dependencies = [ [[package]] name = "opendal-service-mysql" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e15592be32b105cd152ab398b3db847e22c8c5f230d9b34ff095ea3d0dc6fd97" +checksum = "9ba30e2a0b38aa55579b69bf57f607df7aa613f4d964496407b7998f2dbaefcc" dependencies = [ "asyncband", "opendal-core", @@ -9933,9 +9933,9 @@ dependencies = [ [[package]] name = "opendal-service-oss" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0ebfa355b1d4edb0e5b2a5b86d8de9f811fae1d7baca6d65246b230b2afc0bd" +checksum = "d64d2909b444008e42f00137a50df6b55ba599eeb6c8b46c71a2a84948f91880" dependencies = [ "bytes", "http 1.5.0", @@ -9950,9 +9950,9 @@ dependencies = [ [[package]] name = "opendal-service-s3" -version = "0.59.2" +version = "0.59.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2942f2d8d3d4953c0e8d07cd1244879628d084948e48ce6b6d808c5938f7858d" +checksum = "1d83e716c25a81fe2ca701894b8e14b9966acca3812af250e50132accb19da02" dependencies = [ "base64 0.23.1", "bytes", diff --git a/src/object-store/Cargo.toml b/src/object-store/Cargo.toml index 15e2069e3bd..96785fbf243 100644 --- a/src/object-store/Cargo.toml +++ b/src/object-store/Cargo.toml @@ -24,7 +24,7 @@ common-telemetry.workspace = true derive_builder = { workspace = true, optional = true } hdfs-native = { workspace = true, optional = true } humantime-serde.workspace = true -opendal = { version = "0.59.2", features = [ +opendal = { version = "0.59.4", features = [ "layers-tracing", "layers-prometheus", "services-azblob", @@ -34,7 +34,7 @@ opendal = { version = "0.59.2", features = [ "services-oss", "services-s3", ] } -opendal-http-transport-reqwest = "0.59.2" +opendal-http-transport-reqwest = "0.59.4" prometheus.workspace = true reqwest.workspace = true serde.workspace = true From 845fc0af9c09a66f5a19ce0175ec67be3454b1d5 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 12 Oct 2026 05:10:07 +0000 Subject: [PATCH 06/46] fix(promql): default a stepless subquery to a one-minute step (#9481) * fix(promql): default a stepless subquery to a one-minute step `x[5m:]` evaluated the inner expression at the query step, which is 1s for instant queries. Prometheus uses the evaluation interval, one minute by default. Signed-off-by: Dennis Zhuang * test(promql): cover the default subquery step Signed-off-by: Dennis Zhuang * test(promql): cover the default subquery step in a range query Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/promql/planner.rs | 7 +-- src/query/src/promql/planner/test.rs | 8 --- .../promql/subquery_default_step.result | 53 +++++++++++++++++++ .../common/promql/subquery_default_step.sql | 25 +++++++++ 4 files changed, 82 insertions(+), 11 deletions(-) create mode 100644 tests/cases/standalone/common/promql/subquery_default_step.result create mode 100644 tests/cases/standalone/common/promql/subquery_default_step.sql diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5867301f972..972f04b9ccf 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -587,9 +587,10 @@ impl PromPlanner { }; let current_interval = self.ctx.interval; - if let Some(step) = step { - self.ctx.interval = step.as_millis() as _; - } + self.ctx.interval = match step { + Some(step) => step.as_millis() as _, + None => DEFAULT_SUBQUERY_STEP_MS, + }; ensure!(self.ctx.interval > 0, ZeroRangeSelectorSnafu); let current_start = self.ctx.start; let current_end = self.ctx.end; diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 19de1098288..f78162d8fdb 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -5682,14 +5682,6 @@ async fn subquery_rejects_invalid_range_and_zero_step() { matches!(&err, crate::promql::error::Error::ZeroRangeSelector { .. }), "{err}" ); - // An omitted child step inherits the zero parent step. - let err = plan_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:])", 60, 60, 0) - .await - .unwrap_err(); - assert!( - matches!(&err, crate::promql::error::Error::ZeroRangeSelector { .. }), - "{err}" - ); } /// Rejects nonempty child windows whose derived millisecond bounds are unrepresentable. diff --git a/tests/cases/standalone/common/promql/subquery_default_step.result b/tests/cases/standalone/common/promql/subquery_default_step.result new file mode 100644 index 00000000000..47008660cad --- /dev/null +++ b/tests/cases/standalone/common/promql/subquery_default_step.result @@ -0,0 +1,53 @@ +-- A subquery without a step uses the evaluation interval (one minute by +-- default in Prometheus), not the query step. Expected results are from Prometheus 3. +create table subquery_default_step ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +insert into subquery_default_step values + (0, 1.0, 'a'), + (60000, 1.0, 'a'), + (120000, 1.0, 'a'), + (180000, 1.0, 'a'), + (240000, 1.0, 'a'), + (300000, 1.0, 'a'); + +Affected Rows: 6 + +tql eval (300, 300, '1s') count_over_time(subquery_default_step[5m:]); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:05:00 | 5.0 | a | ++---------------------+------------------------------------+------+ + +tql eval (300, 300, '1s') count_over_time(subquery_default_step[5m:30s]); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:05:00 | 10.0 | a | ++---------------------+------------------------------------+------+ + +-- The outer range query keeps its own step around the subquery's default step. +tql eval (300, 360, '15s') count_over_time(subquery_default_step[5m:]); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:05:00 | 5.0 | a | +| 1970-01-01T00:05:15 | 5.0 | a | +| 1970-01-01T00:05:30 | 5.0 | a | +| 1970-01-01T00:05:45 | 5.0 | a | +| 1970-01-01T00:06:00 | 5.0 | a | ++---------------------+------------------------------------+------+ + +drop table subquery_default_step; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/subquery_default_step.sql b/tests/cases/standalone/common/promql/subquery_default_step.sql new file mode 100644 index 00000000000..fb1acb5f45b --- /dev/null +++ b/tests/cases/standalone/common/promql/subquery_default_step.sql @@ -0,0 +1,25 @@ +-- A subquery without a step uses the evaluation interval (one minute by +-- default in Prometheus), not the query step. Expected results are from Prometheus 3. + +create table subquery_default_step ( + ts timestamp time index, + val double, + host string primary key +); + +insert into subquery_default_step values + (0, 1.0, 'a'), + (60000, 1.0, 'a'), + (120000, 1.0, 'a'), + (180000, 1.0, 'a'), + (240000, 1.0, 'a'), + (300000, 1.0, 'a'); + +tql eval (300, 300, '1s') count_over_time(subquery_default_step[5m:]); + +tql eval (300, 300, '1s') count_over_time(subquery_default_step[5m:30s]); + +-- The outer range query keeps its own step around the subquery's default step. +tql eval (300, 360, '15s') count_over_time(subquery_default_step[5m:]); + +drop table subquery_default_step; From f3a0841f4f44c72534a55269e5dc3d7c94c5737c Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 12 Oct 2026 05:10:18 +0000 Subject: [PATCH 07/46] fix(promql): follow IEEE 754 float semantics in comparisons and aggregations (#9479) * fix(promql): follow IEEE 754 float semantics in operators and aggregations - Comparisons: Arrow orders floats totally (NaN == NaN, NaN > +Inf); PromQL comparisons with NaN are false except !=. - min/max: skip NaN unless the whole group is NaN. DataFusion's grouped min/max let NaN replace and be replaced by neighbours, so the result depended on input order. - topk/bottomk rank NaN last. - round() rounds ties up and returns NaN for a zero step, like Prometheus. - sqrt() of a negative number is NaN instead of an error. Signed-off-by: Dennis Zhuang * fix(promql): round() without a step rounds to integers The planner passed a zero step for single-argument round(), which now yields NaN like Prometheus; pass 1 instead. Found by the compatibility harness. Signed-off-by: Dennis Zhuang * fix(promql): apply the sqrt guard to builtin functions and register min/max for plan decoding sqrt resolves through the session's builtin functions, not the UDF branch. Distributed plans are decoded with the serializer's session, which did not know prom_min/prom_max. Signed-off-by: Dennis Zhuang * test(promql): cover NaN and rounding semantics Signed-off-by: Dennis Zhuang * fix(promql): keep the nested count rewrite for PromQL min/max count(min(x) by (l)) was rewritten by count_nest_aggr only for the builtin min/max names. prom_min/prom_max return NULL for the same groups, so the rewrite stays valid. Update plan snapshots for the IEEE comparison, sqrt, topk and min/max changes. Signed-off-by: Dennis Zhuang * fix(promql): register prom_min/prom_max in the function registry Distributed plans split an aggregation into state and merge steps only for functions whose state function exists in FUNCTION_REGISTRY; prom_min and prom_max lived in the promql crate, so min/max stopped being pushed down to datanodes. Move them to common-function and register them before the state functions are derived. Signed-off-by: Dennis Zhuang * test: update sqlness results for PromQL float semantics Plans and column names now show prom_min/prom_max, the NaN guards of comparisons and sqrt, and the NaN sort key of topk. round() results follow Prometheus' formula. Distributed min/max keep the state/merge split. Signed-off-by: Dennis Zhuang * fix(promql): keep column names and NULLs in IEEE comparisons and sqrt Rewriting comparisons as `a op b AND NOT (isnan(a) OR isnan(b))` changed TQL output column names and turned NULL operands into true for `!=` and false for the bool forms. Implement comparisons and sqrt as PromQL UDFs that propagate NULL and report the column name of the expression they replace; prom_min/prom_max report min/max column names the same way. Signed-off-by: Dennis Zhuang * test: restore column names in sqlness results for PromQL float semantics Signed-off-by: Dennis Zhuang * test(promql): share the prom_min/prom_max groups accumulator with tests Signed-off-by: Dennis Zhuang * fix(promql): keep distinct names for prom_min/prom_max Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/aggrs.rs | 1 + src/common/function/src/aggrs/extremum.rs | 259 +++++++++++++++++ src/common/function/src/function_registry.rs | 4 + src/promql/src/functions.rs | 2 + src/promql/src/functions/ieee_float.rs | 265 ++++++++++++++++++ src/promql/src/functions/round.rs | 62 ++-- src/query/src/optimizer/count_nest_aggr.rs | 19 +- src/query/src/promql/planner.rs | 115 ++++++-- src/query/src/promql/planner/test.rs | 55 +++- .../src/query_engine/default_serializer.rs | 6 +- .../explain/step_aggr_advance.result | 108 +++---- .../common/promql/count_values.result | 2 +- .../common/promql/delta_temporality.result | 2 +- .../common/promql/dotted_column_names.result | 14 +- .../common/promql/float_semantics.result | 115 ++++++++ .../common/promql/float_semantics.sql | 41 +++ .../standalone/common/promql/functions.result | 12 +- .../promql/tsid_binary_join_regression.result | 12 +- 18 files changed, 954 insertions(+), 140 deletions(-) create mode 100644 src/common/function/src/aggrs/extremum.rs create mode 100644 src/promql/src/functions/ieee_float.rs create mode 100644 tests/cases/standalone/common/promql/float_semantics.result create mode 100644 tests/cases/standalone/common/promql/float_semantics.sql diff --git a/src/common/function/src/aggrs.rs b/src/common/function/src/aggrs.rs index e618921a18d..bad38a075d4 100644 --- a/src/common/function/src/aggrs.rs +++ b/src/common/function/src/aggrs.rs @@ -14,6 +14,7 @@ pub mod approximate; pub mod count_hash; +pub mod extremum; #[cfg(feature = "geo")] pub mod geo; pub mod vector; diff --git a/src/common/function/src/aggrs/extremum.rs b/src/common/function/src/aggrs/extremum.rs new file mode 100644 index 00000000000..668c307ff36 --- /dev/null +++ b/src/common/function/src/aggrs/extremum.rs @@ -0,0 +1,259 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! PromQL `min` and `max` aggregation operators. +//! +//! PromQL ignores NaN samples unless every sample of a group is NaN. DataFusion's +//! `min`/`max` order NaN differently depending on the code path (its grouped +//! accumulator lets NaN replace the current value and be replaced by the next one), +//! so the result depends on the order of the input. + +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, AsArray}; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Float64Type}; +use datafusion::error::Result as DfResult; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, GroupsAccumulator, Signature, Volatility, +}; +use datafusion_common::ScalarValue; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator; + +use crate::function_registry::FunctionRegistry; + +/// `max` or `min` with PromQL NaN semantics. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct Extremum { + name: &'static str, + is_max: bool, + signature: Signature, +} + +impl Extremum { + pub const MAX_NAME: &'static str = "prom_max"; + pub const MIN_NAME: &'static str = "prom_min"; + + /// Registers both functions; the registry also derives their state functions, which + /// distributed plans need to split the aggregation between datanodes and the frontend. + pub fn register(registry: &FunctionRegistry) { + registry.register_aggr(Self::max_udaf()); + registry.register_aggr(Self::min_udaf()); + } + + pub fn max_udaf() -> AggregateUDF { + AggregateUDF::from(Self::new(Self::MAX_NAME, true)) + } + + pub fn min_udaf() -> AggregateUDF { + AggregateUDF::from(Self::new(Self::MIN_NAME, false)) + } + + fn new(name: &'static str, is_max: bool) -> Self { + Self { + name, + is_max, + signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable), + } + } +} + +/// Folds `value` into `current` like Prometheus' aggregation: a NaN is replaced by any +/// sample, and a NaN sample never replaces a number. +fn fold(is_max: bool, current: &mut f64, value: f64) { + let better = if is_max { + value > *current + } else { + value < *current + }; + if current.is_nan() || better { + *current = value; + } +} + +impl AggregateUDFImpl for Extremum { + fn name(&self) -> &str { + self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(DataType::Float64) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> DfResult> { + Ok(Box::new(ExtremumAccumulator { + is_max: self.is_max, + value: None, + })) + } + + fn state_fields(&self, args: StateFieldsArgs) -> DfResult> { + Ok(vec![Arc::new(Field::new( + format!("{}[{}]", args.name, self.name), + DataType::Float64, + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + _args: AccumulatorArgs, + ) -> DfResult> { + Ok(groups_accumulator(self.is_max)) + } +} + +fn groups_accumulator(is_max: bool) -> Box { + Box::new( + PrimitiveGroupsAccumulator::::new(&DataType::Float64, move |cur, v| { + fold(is_max, cur, v) + }) + .with_starting_value(f64::NAN), + ) +} + +#[derive(Debug)] +struct ExtremumAccumulator { + is_max: bool, + value: Option, +} + +impl ExtremumAccumulator { + fn update(&mut self, values: &ArrayRef) { + for v in values.as_primitive::().iter().flatten() { + match &mut self.value { + Some(current) => fold(self.is_max, current, v), + None => self.value = Some(v), + } + } + } +} + +impl Accumulator for ExtremumAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> { + self.update(&values[0]); + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> { + self.update(&states[0]); + Ok(()) + } + + fn state(&mut self) -> DfResult> { + Ok(vec![ScalarValue::Float64(self.value)]) + } + + fn evaluate(&mut self) -> DfResult { + Ok(ScalarValue::Float64(self.value)) + } + + fn size(&self) -> usize { + size_of_val(self) + } +} + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::Float64Array; + use datafusion::logical_expr::EmitTo; + + use super::*; + + fn groups_result( + is_max: bool, + values: Vec>, + groups: Vec, + ) -> Vec> { + let num_groups = groups.iter().max().map_or(0, |g| g + 1); + let mut acc = groups_accumulator(is_max); + let values: ArrayRef = Arc::new(Float64Array::from(values)); + acc.update_batch(&[values], &groups, None, num_groups) + .unwrap(); + let out = acc.evaluate(EmitTo::All).unwrap(); + out.as_primitive::().iter().collect() + } + + fn is_nan(v: Option) -> bool { + v.is_some_and(f64::is_nan) + } + + #[test] + fn nan_is_ignored_unless_the_group_is_all_nan() { + // Group 0 has NaN on both sides of a number; group 1 is all NaN; group 2 has only NULL. + let values = vec![ + Some(f64::NAN), + Some(3.0), + Some(f64::NAN), + Some(f64::NAN), + None, + ]; + let groups = vec![0, 0, 0, 1, 2]; + for is_max in [true, false] { + let out = groups_result(is_max, values.clone(), groups.clone()); + assert_eq!(out[0], Some(3.0)); + assert!(is_nan(out[1])); + assert_eq!(out[2], None); + } + } + + #[test] + fn result_does_not_depend_on_input_order() { + let forward = vec![Some(1.0), Some(f64::NAN), Some(4.0), Some(2.0)]; + let mut backward = forward.clone(); + backward.reverse(); + assert_eq!( + groups_result(true, forward.clone(), vec![0; 4]), + vec![Some(4.0)] + ); + assert_eq!( + groups_result(true, backward.clone(), vec![0; 4]), + vec![Some(4.0)] + ); + assert_eq!(groups_result(false, forward, vec![0; 4]), vec![Some(1.0)]); + assert_eq!(groups_result(false, backward, vec![0; 4]), vec![Some(1.0)]); + } + + #[test] + fn merged_partial_states_keep_nan_semantics() { + let mut partial_nan = ExtremumAccumulator { + is_max: true, + value: None, + }; + partial_nan.update(&(Arc::new(Float64Array::from(vec![f64::NAN])) as ArrayRef)); + let mut partial_num = ExtremumAccumulator { + is_max: true, + value: None, + }; + partial_num.update(&(Arc::new(Float64Array::from(vec![2.0, 5.0])) as ArrayRef)); + + let mut merged = ExtremumAccumulator { + is_max: true, + value: None, + }; + for partial in [&mut partial_nan, &mut partial_num] { + let state = partial.state().unwrap()[0].to_array().unwrap(); + merged.merge_batch(&[state]).unwrap(); + } + assert_eq!(merged.evaluate().unwrap(), ScalarValue::Float64(Some(5.0))); + } +} diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index 6583f7ba607..8657636fe6b 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -25,6 +25,7 @@ use crate::admin::AdminFunction; use crate::aggrs::aggr_wrapper::StateMergeHelper; use crate::aggrs::approximate::ApproximateFunction; use crate::aggrs::count_hash::CountHash; +use crate::aggrs::extremum::Extremum; use crate::aggrs::vector::VectorFunction as VectorAggrFunction; use crate::function::{Function, FunctionRef}; use crate::function_factory::ScalarFunctionFactory; @@ -260,6 +261,9 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| // CountHash function CountHash::register(&function_registry); + // PromQL min/max; registered before the state functions are derived. + Extremum::register(&function_registry); + // state function of supported aggregate functions StateMergeHelper::register(&function_registry); diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index 34618f0a953..6703505c277 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -19,6 +19,7 @@ mod double_exponential_smoothing; mod edge_count; mod extrapolate_rate; mod idelta; +mod ieee_float; mod native_histogram; mod predict_linear; mod quantile; @@ -45,6 +46,7 @@ pub use deriv::Deriv; pub use double_exponential_smoothing::DoubleExponentialSmoothing; pub use extrapolate_rate::{Delta, Increase, Rate}; pub use idelta::IDelta; +pub use ieee_float::{IeeeComparison, IeeeSqrt}; pub use native_histogram::{ MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, diff --git a/src/promql/src/functions/ieee_float.rs b/src/promql/src/functions/ieee_float.rs new file mode 100644 index 00000000000..d13480d8071 --- /dev/null +++ b/src/promql/src/functions/ieee_float.rs @@ -0,0 +1,265 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! PromQL comparison operators and `sqrt()`, which follow IEEE 754. +//! +//! Arrow compares floats by total order, where NaN equals NaN and is greater than `+Inf`, and +//! DataFusion's `sqrt` rejects negative input. In PromQL every comparison involving NaN is false +//! except `!=`, and the square root of a negative number is NaN. NULL input gives NULL output. +//! +//! The functions report the column name of the expression they replace (`a < b`, `sqrt(a)`), so +//! query output columns are unchanged. + +use std::sync::Arc; + +use datafusion::arrow::array::{AsArray, BooleanArray, Float64Array}; +use datafusion::arrow::datatypes::{DataType, Float64Type}; +use datafusion::error::Result as DfResult; +use datafusion::logical_expr::{ + ColumnarValue, Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; +use datafusion_expr::Expr; + +/// A PromQL comparison operator on two float operands. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct IeeeComparison { + op: Operator, + signature: Signature, +} + +impl IeeeComparison { + /// The comparison UDF for `op`, which must be a comparison operator. + pub fn scalar_udf(op: Operator) -> ScalarUDF { + ScalarUDF::new_from_impl(Self { + op, + signature: Signature::uniform(2, vec![DataType::Float64], Volatility::Immutable), + }) + } + + /// All comparison UDFs, for registering them where plans are decoded. + pub fn all() -> Vec { + [ + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + ] + .into_iter() + .map(Self::scalar_udf) + .collect() + } + + fn compare(&self, lhs: f64, rhs: f64) -> bool { + match self.op { + Operator::Eq => lhs == rhs, + Operator::NotEq => lhs != rhs, + Operator::Lt => lhs < rhs, + Operator::LtEq => lhs <= rhs, + Operator::Gt => lhs > rhs, + Operator::GtEq => lhs >= rhs, + _ => unreachable!("not a comparison operator: {}", self.op), + } + } +} + +impl ScalarUDFImpl for IeeeComparison { + fn name(&self) -> &str { + match self.op { + Operator::Eq => "prom_eq", + Operator::NotEq => "prom_ne", + Operator::Lt => "prom_lt", + Operator::LtEq => "prom_le", + Operator::Gt => "prom_gt", + Operator::GtEq => "prom_ge", + _ => unreachable!("not a comparison operator: {}", self.op), + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(DataType::Boolean) + } + + fn schema_name(&self, args: &[Expr]) -> DfResult { + Ok(format!( + "{} {} {}", + args[0].schema_name(), + self.op, + args[1].schema_name() + )) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let lhs = arrays[0].as_primitive::(); + let rhs = arrays[1].as_primitive::(); + let result: BooleanArray = lhs + .iter() + .zip(rhs.iter()) + .map(|(lhs, rhs)| Some(self.compare(lhs?, rhs?))) + .collect(); + Ok(ColumnarValue::Array(Arc::new(result))) + } +} + +/// PromQL `sqrt()`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct IeeeSqrt { + signature: Signature, +} + +impl IeeeSqrt { + pub fn scalar_udf() -> ScalarUDF { + ScalarUDF::new_from_impl(Self { + signature: Signature::uniform(1, vec![DataType::Float64], Volatility::Immutable), + }) + } +} + +impl ScalarUDFImpl for IeeeSqrt { + fn name(&self) -> &str { + "prom_sqrt" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(DataType::Float64) + } + + fn schema_name(&self, args: &[Expr]) -> DfResult { + Ok(format!("sqrt({})", args[0].schema_name())) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let input = arrays[0].as_primitive::(); + let result: Float64Array = input.unary(f64::sqrt); + Ok(ColumnarValue::Array(Arc::new(result))) + } +} + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::{Array, ArrayRef}; + use datafusion::arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; + + use super::*; + + fn invoke(udf: &ScalarUDF, args: Vec, return_type: DataType) -> ArrayRef { + let number_rows = args[0].len(); + let arg_fields = args + .iter() + .map(|array| Arc::new(Field::new("a", array.data_type().clone(), true))) + .collect(); + let result = udf + .invoke_with_args(ScalarFunctionArgs { + args: args.into_iter().map(ColumnarValue::Array).collect(), + arg_fields, + number_rows, + return_field: Arc::new(Field::new("r", return_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap(); + match result { + ColumnarValue::Array(array) => array, + ColumnarValue::Scalar(_) => unreachable!(), + } + } + + #[test] + fn comparisons_follow_ieee_and_propagate_null() { + let lhs: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.0), + Some(f64::NAN), + Some(f64::NAN), + None, + Some(f64::INFINITY), + ])); + let rhs: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.0), + Some(f64::NAN), + Some(1.0), + Some(f64::NAN), + Some(f64::NAN), + ])); + let expected = [ + ( + Operator::Eq, + [Some(true), Some(false), Some(false), None, Some(false)], + ), + ( + Operator::NotEq, + [Some(false), Some(true), Some(true), None, Some(true)], + ), + ( + Operator::Lt, + [Some(false), Some(false), Some(false), None, Some(false)], + ), + ( + Operator::LtEq, + [Some(true), Some(false), Some(false), None, Some(false)], + ), + ( + Operator::Gt, + [Some(false), Some(false), Some(false), None, Some(false)], + ), + ( + Operator::GtEq, + [Some(true), Some(false), Some(false), None, Some(false)], + ), + ]; + for (op, expected) in expected { + let result = invoke( + &IeeeComparison::scalar_udf(op), + vec![lhs.clone(), rhs.clone()], + DataType::Boolean, + ); + let result: Vec<_> = result.as_boolean().iter().collect(); + assert_eq!(result, expected, "{op}"); + } + } + + #[test] + fn sqrt_of_negative_is_nan() { + let input: ArrayRef = Arc::new(Float64Array::from(vec![Some(4.0), Some(-1.0), None])); + let result = invoke(&IeeeSqrt::scalar_udf(), vec![input], DataType::Float64); + let result = result.as_primitive::(); + assert_eq!(result.value(0), 2.0); + assert!(result.value(1).is_nan()); + assert!(result.is_null(2)); + } + + #[test] + fn column_names_match_the_replaced_expressions() { + let a = datafusion_expr::col("a"); + let b = datafusion_expr::lit(1.5_f64); + let comparison = + IeeeComparison::scalar_udf(Operator::LtEq).call(vec![a.clone(), b.clone()]); + assert_eq!( + comparison.schema_name().to_string(), + a.clone().lt_eq(b).schema_name().to_string() + ); + let sqrt = IeeeSqrt::scalar_udf().call(vec![a.clone()]); + assert_eq!(sqrt.schema_name().to_string(), "sqrt(a)"); + } +} diff --git a/src/promql/src/functions/round.rs b/src/promql/src/functions/round.rs index 72736ef6e2e..c5ba4c218c0 100644 --- a/src/promql/src/functions/round.rs +++ b/src/promql/src/functions/round.rs @@ -66,13 +66,7 @@ impl Round { let null_array = Float64Array::new_null(value_array.len()); return Ok(ColumnarValue::Array(Arc::new(null_array))); }; - let op = |a: f64| { - if nearest == 0.0 { - a.round() - } else { - (a / nearest).round() * nearest - } - }; + let op = |a: f64| round_to_nearest(a, nearest); let result: PrimitiveArray = value_array.as_primitive::().unary(op); Ok(ColumnarValue::Array(Arc::new(result) as _)) @@ -90,14 +84,10 @@ impl Round { )?; let result: PrimitiveArray = - datatypes::arrow::compute::binary(value_array, nearest_array, |a, nearest| { - if nearest == 0.0 { - a.round() - } else { - (a / nearest).round() * nearest - } - }) - .map_err(|err: ArrowError| DataFusionError::ArrowError(Box::new(err), None))?; + datatypes::arrow::compute::binary(value_array, nearest_array, round_to_nearest) + .map_err(|err: ArrowError| { + DataFusionError::ArrowError(Box::new(err), None) + })?; Ok(ColumnarValue::Array(Arc::new(result) as _)) } @@ -105,6 +95,13 @@ impl Round { } } +/// Rounds like Prometheus: ties round up (towards +Inf), and multiplying by the inverse +/// of `nearest` follows its float behavior, including NaN for `nearest == 0`. +fn round_to_nearest(value: f64, nearest: f64) -> f64 { + let inverse = 1.0 / nearest; + (value * inverse + 0.5).floor() / inverse +} + #[cfg(test)] mod tests { use datafusion_common::config::ConfigOptions; @@ -143,14 +140,43 @@ mod tests { #[test] fn test_round() { + // Expected values are Prometheus 3.15 `round()` results. test_round_f64(vec![123.456], 0.001, vec![123.456]); - test_round_f64(vec![123.456], 0.01, vec![123.46000000000001]); + test_round_f64(vec![123.456], 0.01, vec![123.46]); test_round_f64(vec![123.456], 0.1, vec![123.5]); - test_round_f64(vec![123.456], 0.0, vec![123.0]); test_round_f64(vec![123.456], 1.0, vec![123.0]); test_round_f64(vec![123.456], 10.0, vec![120.0]); test_round_f64(vec![123.456], 100.0, vec![100.0]); - test_round_f64(vec![123.456], 105.0, vec![105.0]); + test_round_f64(vec![123.456], 105.0, vec![104.99999999999999]); test_round_f64(vec![123.456], 1000.0, vec![0.0]); } + + #[test] + fn test_round_ties_towards_positive_infinity() { + test_round_f64(vec![-0.5], 1.0, vec![0.0]); + test_round_f64(vec![-2.5], 1.0, vec![-2.0]); + test_round_f64(vec![2.5], 1.0, vec![3.0]); + } + + #[test] + fn test_round_to_zero_is_nan() { + let round_udf = Round::scalar_udf(); + let input = vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![123.456]))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))), + ]; + let arg_fields = vec![ + Arc::new(Field::new("a", input[0].data_type(), false)), + Arc::new(Field::new("b", input[1].data_type(), false)), + ]; + let args = ScalarFunctionArgs { + args: input, + arg_fields, + number_rows: 1, + return_field: Arc::new(Field::new("c", DataType::Float64, false)), + config_options: Arc::new(ConfigOptions::default()), + }; + let result = extract_array(&round_udf.invoke_with_args(args).unwrap()).unwrap(); + assert!(result.as_primitive::().value(0).is_nan()); + } } diff --git a/src/query/src/optimizer/count_nest_aggr.rs b/src/query/src/optimizer/count_nest_aggr.rs index 89ba4260742..6e3e9cdba05 100644 --- a/src/query/src/optimizer/count_nest_aggr.rs +++ b/src/query/src/optimizer/count_nest_aggr.rs @@ -15,6 +15,8 @@ use std::collections::HashSet; use std::sync::Arc; +use common_function::aggrs::aggr_wrapper::get_aggr_func; +use common_function::aggrs::extremum::Extremum; use datafusion::config::ConfigOptions; use datafusion::functions_aggregate::count::count_udaf; use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder, Sort}; @@ -236,9 +238,8 @@ impl CountNestAggrRule { where F: FnOnce(&str) -> bool, { - let Expr::AggregateFunction(func) = expr else { - return None; - }; + // The PromQL planner aliases `prom_min`/`prom_max` to the builtin names. + let func = get_aggr_func(expr)?; let name = func.func.name(); if !accept_name(name) || func.params.filter.is_some() @@ -253,9 +254,19 @@ impl CountNestAggrRule { } fn is_supported_inner_aggregate(name: &str) -> bool { + // PromQL `min`/`max` plan to `prom_min`/`prom_max`, which return NULL for the same + // groups as the builtin ones. matches!( name, - "count" | "sum" | "avg" | "min" | "max" | "stddev_pop" | "var_pop" + "count" + | "sum" + | "avg" + | "min" + | "max" + | Extremum::MIN_NAME + | Extremum::MAX_NAME + | "stddev_pop" + | "var_pop" ) } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 972f04b9ccf..afe52c7b560 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -27,6 +27,7 @@ use async_recursion::async_recursion; use catalog::table_source::DfTableSourceProvider; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_function::aggrs::extremum::Extremum; use common_function::function::FunctionContext; use common_query::native_histogram::native_histogram_value_type; use common_query::prelude::{ @@ -69,18 +70,18 @@ use promql::extension_plan::{ }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MatchGroupViolation, MaxOverTime, MinOverTime, MixedRange, - NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, - NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, - NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, - NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramIDelta, - NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, - NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, - NativeHistogramPresentOverTime, NativeHistogramRate, NativeHistogramResets, - NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, - NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, PredictLinear, - PresentOverTime, PromqlFloatToString, QuantileOverTime, Rate, Resets, Round, StddevOverTime, - StdvarOverTime, SumOverTime, UniqueMatchGroup, quantile_udaf, + IDelta, IeeeComparison, IeeeSqrt, Increase, LastOverTime, MatchGroupViolation, MaxOverTime, + MinOverTime, MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, + NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, + NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, + NativeHistogramDelta, NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, + NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, + NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, + NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramRate, + NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, + NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, + PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, Rate, Resets, Round, + StddevOverTime, StdvarOverTime, SumOverTime, UniqueMatchGroup, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -3963,10 +3964,14 @@ impl PromPlanner { } "round" => { if other_input_exprs.is_empty() { - other_input_exprs.push_front(0.0f64.lit()); + other_input_exprs.push_front(1.0f64.lit()); } ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf())) } + // DataFusion's `sqrt` rejects negative input; PromQL returns NaN for it. + "sqrt" if !self.all_field_columns_are_native_histograms(input_schema) => { + ScalarFunc::DataFusionUdf(Arc::new(IeeeSqrt::scalar_udf())) + } "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => { ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name)) } @@ -4674,19 +4679,41 @@ impl PromPlanner { vec![] }; - // update value column name according to the aggregators, - let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len()); + // The alias is built from normalized columns so it matches the builtin's + // qualified name. + let exprs = normalize_cols(exprs, input_plan) + .context(DataFusionPlanningSnafu)? + .into_iter() + .map(Self::alias_extremum_as_builtin) + .collect::>(); - let normalized_exprs = - normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?; - for expr in normalized_exprs { - new_field_columns.push(expr.schema_name().to_string()); - } - self.ctx.field_columns = new_field_columns; + // update value column name according to the aggregators, + self.ctx.field_columns = exprs + .iter() + .map(|expr| expr.schema_name().to_string()) + .collect(); Ok((exprs, prev_field_exprs)) } + /// `prom_min`/`prom_max` keep their own names so SQL can use them next to `min`/`max` + /// in one aggregate. PromQL output columns keep the builtin names. + fn alias_extremum_as_builtin(expr: DfExpr) -> DfExpr { + let DfExpr::AggregateFunction(aggr) = &expr else { + return expr; + }; + let builtin = match aggr.func.name() { + Extremum::MIN_NAME => min_udaf(), + Extremum::MAX_NAME => max_udaf(), + _ => return expr, + }; + let name = builtin + .call(aggr.params.args.clone()) + .schema_name() + .to_string(); + expr.alias(name) + } + fn create_numeric_aggregate_expr( op: TokenType, param: &Option>, @@ -4704,8 +4731,8 @@ impl PromPlanner { } token::T_AVG => avg_udaf().call(vec![input]), token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]), - token::T_MIN => min_udaf().call(vec![input]), - token::T_MAX => max_udaf().call(vec![input]), + token::T_MIN => Extremum::min_udaf().call(vec![input]), + token::T_MAX => Extremum::max_udaf().call(vec![input]), // PromQL's `group()` aggregator produces 1 for each group. // Use `max(1.0)` (per-group) to match semantics and output type (Float64). token::T_GROUP => max_udaf().call(vec![lit(1_f64)]), @@ -5049,7 +5076,15 @@ impl PromPlanner { .field_columns .iter() .map(|col| { - let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1); + let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2); + // PromQL ranks NaN last for both `topk` and `bottomk`, while DataFusion + // orders NaN above every number. + sort_exprs.push( + datafusion_functions::math::expr_fn::isnan(DfExpr::Column(Column::from_name( + col, + ))) + .sort(true, true), + ); // Order by value in the specific order sort_exprs.push(DfExpr::Column(Column::from_name(col)).sort(asc, true)); // Then tags if the values are equal, @@ -5240,12 +5275,24 @@ impl PromPlanner { token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| { Ok(cast_float(lhs) % cast_float(rhs)) })), - token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))), - token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))), - token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))), - token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))), - token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))), - token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))), + token::T_EQLC => Ok(Box::new(|lhs, rhs| { + Ok(Self::ieee_comparison(lhs, Operator::Eq, rhs)) + })), + token::T_NEQ => Ok(Box::new(|lhs, rhs| { + Ok(Self::ieee_comparison(lhs, Operator::NotEq, rhs)) + })), + token::T_GTR => Ok(Box::new(|lhs, rhs| { + Ok(Self::ieee_comparison(lhs, Operator::Gt, rhs)) + })), + token::T_LSS => Ok(Box::new(|lhs, rhs| { + Ok(Self::ieee_comparison(lhs, Operator::Lt, rhs)) + })), + token::T_GTE => Ok(Box::new(|lhs, rhs| { + Ok(Self::ieee_comparison(lhs, Operator::GtEq, rhs)) + })), + token::T_LTE => Ok(Box::new(|lhs, rhs| { + Ok(Self::ieee_comparison(lhs, Operator::LtEq, rhs)) + })), token::T_POW => Ok(Box::new(move |lhs, rhs| { Ok(DfExpr::ScalarFunction(ScalarFunction { func: datafusion_functions::math::power(), @@ -5262,6 +5309,14 @@ impl PromPlanner { } } + /// Compares two expressions with IEEE 754 semantics, as PromQL does; see [`IeeeComparison`]. + fn ieee_comparison(lhs: DfExpr, op: Operator, rhs: DfExpr) -> DfExpr { + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(IeeeComparison::scalar_udf(op)), + args: vec![lhs, rhs], + }) + } + /// Check if the given op is a [comparison operator](https://prometheus.io/docs/prometheus/latest/querying/operators/#comparison-binary-operators). fn is_token_a_comparison_op(token: TokenType) -> bool { matches!( diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index f78162d8fdb..ae40640b50b 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -1679,7 +1679,18 @@ async fn single_sort_desc() { #[tokio::test] async fn single_sqrt() { - do_single_instant_function_call("sqrt", "sqrt").await; + // Negative input is mapped to NaN before DataFusion's `sqrt`, which rejects it. + indie_query_plan_compare( + r#"sqrt(some_metric{tag_0!="bar"})"#, + "Filter: sqrt(field_0) IS NOT NULL [timestamp:Timestamp(ms), sqrt(field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_sqrt(some_metric.field_0) AS sqrt(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), sqrt(field_0):Float64;N, tag_0:Utf8]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ) + .await; } #[tokio::test] @@ -1811,6 +1822,19 @@ async fn single_rad() { // }, // }, async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { + do_aggregate_expr_plan_with_display(fn_name, plan_name, plan_name).await; +} + +/// `display_name` is the aggregate function as the plan shows it, `plan_name` the name of the +/// output column; they differ for PromQL `min`/`max`. +async fn do_aggregate_expr_plan_with_display(fn_name: &str, plan_name: &str, display_name: &str) { + let display = |field: &str| { + if display_name == plan_name { + format!("{plan_name}({field})") + } else { + format!("{display_name}({field}) AS {plan_name}({field})") + } + }; let prom_expr = parser::parse(&format!( "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})", )) @@ -1837,13 +1861,16 @@ async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { .unwrap(); let expected_no_without = String::from( "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[DISPLAY_0, DISPLAY_1]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); + ) + .replace("TEMPLATE", plan_name) + .replace("DISPLAY_0", &display("some_metric.field_0")) + .replace("DISPLAY_1", &display("some_metric.field_1")); assert_eq!( plan.display_indent_schema().to_string(), expected_no_without @@ -1866,13 +1893,15 @@ async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { .unwrap(); let expected_without = String::from( "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[DISPLAY_0, DISPLAY_1]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); + ).replace("TEMPLATE", plan_name) + .replace("DISPLAY_0", &display("some_metric.field_0")) + .replace("DISPLAY_1", &display("some_metric.field_1")); assert_eq!(plan.display_indent_schema().to_string(), expected_without); } @@ -3898,12 +3927,12 @@ async fn aggregate_count() { #[tokio::test] async fn aggregate_min() { - do_aggregate_expr_plan("min", "min").await; + do_aggregate_expr_plan_with_display("min", "min", "prom_min").await; } #[tokio::test] async fn aggregate_max() { - do_aggregate_expr_plan("max", "max").await; + do_aggregate_expr_plan_with_display("max", "max", "prom_max").await; } #[tokio::test] @@ -4081,7 +4110,7 @@ async fn binary_op_literal_literal() { async fn simple_bool_grammar() { let query = "some_metric != bool 1.2345"; let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(prom_ne(some_metric.field_0, Float64(1.2345)) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ @@ -4096,7 +4125,7 @@ async fn simple_bool_grammar() { async fn bool_with_additional_arithmetic() { let query = "some_metric + (1 == bool 2)"; let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(prom_eq(Float64(1), Float64(2)) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ @@ -5064,7 +5093,7 @@ async fn native_histogram_mixed_field_table_behaves() { async fn less_filter_on_value() { let query = "some_metric < 1.2345"; let expected = String::from( - "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + "Filter: prom_lt(some_metric.field_0, Float64(1.2345)) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ @@ -6642,9 +6671,9 @@ async fn test_topk_expr() { .await .unwrap(); let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\ - \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [isnan(sum(prometheus_tsdb_head_series.greptime_value)) ASC NULLS FIRST, sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [isnan(sum(prometheus_tsdb_head_series.greptime_value)) ASC NULLS FIRST, sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [isnan(sum(prometheus_tsdb_head_series.greptime_value)) ASC NULLS FIRST, sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [isnan(sum(prometheus_tsdb_head_series.greptime_value)) ASC NULLS FIRST, sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [isnan(sum(prometheus_tsdb_head_series.greptime_value)) ASC NULLS FIRST, sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [isnan(sum(prometheus_tsdb_head_series.greptime_value)) ASC NULLS FIRST, sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 63ee0856aaf..ced633382d3 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -32,7 +32,7 @@ use datafusion_expr::{Expr, UserDefinedLogicalNode}; use greptime_proto::substrait_extension::MergeScan as PbMergeScan; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, + IDelta, IeeeComparison, IeeeSqrt, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, @@ -334,6 +334,10 @@ fn register_greptime_functions( } let _ = session_state.register_udaf(quantile_udaf()); + for udf in IeeeComparison::all() { + let _ = session_state.register_udf(Arc::new(udf)); + } + let _ = session_state.register_udf(Arc::new(IeeeSqrt::scalar_udf())); let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index 348021360ff..11adf22f866 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -21,26 +21,26 @@ Affected Rows: 0 -- SQLNESS REPLACE (Hash.*) REDACTED tql explain (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr_optimize_not [2m])); -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeSort: aggr_optimize_not.a ASC NULLS LAST, aggr_optimize_not.b ASC NULLS LAST, aggr_optimize_not.c ASC NULLS LAST, aggr_optimize_not.greptime_timestamp ASC NULLS LAST | -| | MergeScan [is_placeholder=false, remote_input=[ | -| | Sort: aggr_optimize_not.a ASC NULLS LAST, aggr_optimize_not.b ASC NULLS LAST, aggr_optimize_not.c ASC NULLS LAST, aggr_optimize_not.greptime_timestamp ASC NULLS LAST | -| | Aggregate: groupBy=[[aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.greptime_timestamp]], aggr=[[max(prom_max_over_time(greptime_timestamp_range,greptime_value))]] | -| | Filter: prom_max_over_time(greptime_timestamp_range,greptime_value) IS NOT NULL | -| | Projection: aggr_optimize_not.greptime_timestamp, prom_max_over_time(greptime_timestamp_range, greptime_value) AS prom_max_over_time(greptime_timestamp_range,greptime_value), aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d | -| | PromRangeManipulate: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp], values=["greptime_value"] | -| | PromSeriesNormalize: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | -| | PromSeriesDivide: tags=["a", "b", "c", "d"] | -| | Sort: aggr_optimize_not.a ASC NULLS FIRST, aggr_optimize_not.b ASC NULLS FIRST, aggr_optimize_not.c ASC NULLS FIRST, aggr_optimize_not.d ASC NULLS FIRST, aggr_optimize_not.greptime_timestamp ASC NULLS FIRST | -| | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | -| | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | -| | ]] | -| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeSort: aggr_optimize_not.a ASC NULLS LAST, aggr_optimize_not.b ASC NULLS LAST, aggr_optimize_not.c ASC NULLS LAST, aggr_optimize_not.greptime_timestamp ASC NULLS LAST | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | Sort: aggr_optimize_not.a ASC NULLS LAST, aggr_optimize_not.b ASC NULLS LAST, aggr_optimize_not.c ASC NULLS LAST, aggr_optimize_not.greptime_timestamp ASC NULLS LAST | +| | Aggregate: groupBy=[[aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.greptime_timestamp]], aggr=[[prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value)) AS max(prom_max_over_time(greptime_timestamp_range,greptime_value))]] | +| | Filter: prom_max_over_time(greptime_timestamp_range,greptime_value) IS NOT NULL | +| | Projection: aggr_optimize_not.greptime_timestamp, prom_max_over_time(greptime_timestamp_range, greptime_value) AS prom_max_over_time(greptime_timestamp_range,greptime_value), aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d | +| | PromRangeManipulate: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp], values=["greptime_value"] | +| | PromSeriesNormalize: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | +| | PromSeriesDivide: tags=["a", "b", "c", "d"] | +| | Sort: aggr_optimize_not.a ASC NULLS FIRST, aggr_optimize_not.b ASC NULLS FIRST, aggr_optimize_not.c ASC NULLS FIRST, aggr_optimize_not.d ASC NULLS FIRST, aggr_optimize_not.greptime_timestamp ASC NULLS FIRST | +| | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | +| | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | +| | ]] | +| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | | | MergeScanExec: REDACTED -| | | -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| | | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED @@ -57,11 +57,12 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr | 0_| 0_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| -| 1_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED +| 1_| 0_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value))@4 as max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED -|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_FilterExec: prom_max_over_time(greptime_timestamp_range,greptime_value)@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_max_over_time(greptime_timestamp_range@6, greptime_value@5) as prom_max_over_time(greptime_timestamp_range,greptime_value), a@0 as a, b@1 as b, c@2 as c] REDACTED |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED @@ -69,11 +70,12 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED |_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| -| 1_| 1_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED +| 1_| 1_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value))@4 as max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED -|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_FilterExec: prom_max_over_time(greptime_timestamp_range,greptime_value)@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_max_over_time(greptime_timestamp_range@6, greptime_value@5) as prom_max_over_time(greptime_timestamp_range,greptime_value), a@0 as a, b@1 as b, c@2 as c] REDACTED |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED @@ -305,31 +307,31 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time -- SQLNESS REPLACE (Hash.*) REDACTED tql explain (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr_optimize_not [2m])); -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | Sort: aggr_optimize_not.b ASC NULLS LAST, aggr_optimize_not.c ASC NULLS LAST, aggr_optimize_not.d ASC NULLS LAST, aggr_optimize_not.greptime_timestamp ASC NULLS LAST | -| | Aggregate: groupBy=[[aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d, aggr_optimize_not.greptime_timestamp]], aggr=[[__min_merge(__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) AS min(prom_max_over_time(greptime_timestamp_range,greptime_value))]] | -| | MergeScan [is_placeholder=false, remote_input=[ | -| | Aggregate: groupBy=[[aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d, aggr_optimize_not.greptime_timestamp]], aggr=[[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))]] | -| | Filter: prom_max_over_time(greptime_timestamp_range,greptime_value) IS NOT NULL | -| | Projection: aggr_optimize_not.greptime_timestamp, prom_max_over_time(greptime_timestamp_range, greptime_value) AS prom_max_over_time(greptime_timestamp_range,greptime_value), aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d | -| | PromRangeManipulate: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp], values=["greptime_value"] | -| | PromSeriesNormalize: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | -| | PromSeriesDivide: tags=["a", "b", "c", "d"] | -| | Sort: aggr_optimize_not.a ASC NULLS FIRST, aggr_optimize_not.b ASC NULLS FIRST, aggr_optimize_not.c ASC NULLS FIRST, aggr_optimize_not.d ASC NULLS FIRST, aggr_optimize_not.greptime_timestamp ASC NULLS FIRST | -| | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | -| | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | -| | ]] | -| physical_plan | SortPreservingMergeExec: [b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | -| | SortExec: expr=[b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | -| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_merge(__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: aggr_optimize_not.b ASC NULLS LAST, aggr_optimize_not.c ASC NULLS LAST, aggr_optimize_not.d ASC NULLS LAST, aggr_optimize_not.greptime_timestamp ASC NULLS LAST | +| | Aggregate: groupBy=[[aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d, aggr_optimize_not.greptime_timestamp]], aggr=[[__prom_min_merge(__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) AS min(prom_max_over_time(greptime_timestamp_range,greptime_value))]] | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | Aggregate: groupBy=[[aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d, aggr_optimize_not.greptime_timestamp]], aggr=[[__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))]] | +| | Filter: prom_max_over_time(greptime_timestamp_range,greptime_value) IS NOT NULL | +| | Projection: aggr_optimize_not.greptime_timestamp, prom_max_over_time(greptime_timestamp_range, greptime_value) AS prom_max_over_time(greptime_timestamp_range,greptime_value), aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.c, aggr_optimize_not.d | +| | PromRangeManipulate: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp], values=["greptime_value"] | +| | PromSeriesNormalize: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | +| | PromSeriesDivide: tags=["a", "b", "c", "d"] | +| | Sort: aggr_optimize_not.a ASC NULLS FIRST, aggr_optimize_not.b ASC NULLS FIRST, aggr_optimize_not.c ASC NULLS FIRST, aggr_optimize_not.d ASC NULLS FIRST, aggr_optimize_not.greptime_timestamp ASC NULLS FIRST | +| | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | +| | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | +| | SortExec: expr=[b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | +| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__prom_min_merge(__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | | | RepartitionExec: partitioning=REDACTED -| | AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_merge(__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__prom_min_merge(__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | | | RepartitionExec: partitioning=REDACTED | | MergeScanExec: REDACTED -| | | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| | | ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED @@ -345,15 +347,15 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_merge(__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__prom_min_merge(__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED -|_|_|_AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_merge(__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__prom_min_merge(__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))) as min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| -| 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +| 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED -|_|_|_AggregateExec: mode=Partial, gby=[b@2 as b, c@3 as c, d@4 as d, greptime_timestamp@0 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[b@2 as b, c@3 as c, d@4 as d, greptime_timestamp@0 as greptime_timestamp], aggr=[__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_FilterExec: prom_max_over_time(greptime_timestamp_range,greptime_value)@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_max_over_time(greptime_timestamp_range@6, greptime_value@5) as prom_max_over_time(greptime_timestamp_range,greptime_value), b@1 as b, c@2 as c, d@3 as d] REDACTED |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED @@ -361,9 +363,9 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED |_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| -| 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +| 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED -|_|_|_AggregateExec: mode=Partial, gby=[b@2 as b, c@3 as c, d@4 as d, greptime_timestamp@0 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[b@2 as b, c@3 as c, d@4 as d, greptime_timestamp@0 as greptime_timestamp], aggr=[__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_FilterExec: prom_max_over_time(greptime_timestamp_range,greptime_value)@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_max_over_time(greptime_timestamp_range@6, greptime_value@5) as prom_max_over_time(greptime_timestamp_range,greptime_value), b@1 as b, c@2 as c, d@3 as d] REDACTED |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED diff --git a/tests/cases/standalone/common/promql/count_values.result b/tests/cases/standalone/common/promql/count_values.result index 886c3a143b7..1d04301d73b 100644 --- a/tests/cases/standalone/common/promql/count_values.result +++ b/tests/cases/standalone/common/promql/count_values.result @@ -84,7 +84,7 @@ TQL EVAL (0, 15, '5s') abs(count_values("v", http_requests)); TQL EVAL (0, 15, '5s') round(count_values("v", http_requests)); +---------------------+-------------------------------------------------+-----+ -| ts | prom_round(count(http_requests.val),Float64(0)) | v | +| ts | prom_round(count(http_requests.val),Float64(1)) | v | +---------------------+-------------------------------------------------+-----+ | 1970-01-01T00:00:00 | 1.0 | 401 | | 1970-01-01T00:00:00 | 3.0 | 200 | diff --git a/tests/cases/standalone/common/promql/delta_temporality.result b/tests/cases/standalone/common/promql/delta_temporality.result index 787e72e8dc8..ac32a907cb9 100644 --- a/tests/cases/standalone/common/promql/delta_temporality.result +++ b/tests/cases/standalone/common/promql/delta_temporality.result @@ -162,7 +162,7 @@ TQL EVAL (180, 180, '1m') round(sum(rate(delta_temporality[3m])), 0.000001); +---------------------+----------------------------------------------------------------------------------------+ | ts | prom_round(sum(prom_rate(ts_range,greptime_value,ts,Int64(180000))),Float64(0.000001)) | +---------------------+----------------------------------------------------------------------------------------+ -| 1970-01-01T00:03:00 | 0.45555599999999996 | +| 1970-01-01T00:03:00 | 0.455556 | +---------------------+----------------------------------------------------------------------------------------+ CREATE TABLE delta_marker_only ( diff --git a/tests/cases/standalone/common/promql/dotted_column_names.result b/tests/cases/standalone/common/promql/dotted_column_names.result index c855a5ba5c7..bb4912bc94e 100644 --- a/tests/cases/standalone/common/promql/dotted_column_names.result +++ b/tests/cases/standalone/common/promql/dotted_column_names.result @@ -424,9 +424,9 @@ TQL EXPLAIN (0, 10, '5s') topk(1, {"otel.m"}); | plan_type_| plan_| +-+-+ | logical_plan_| Projection: otel.m.v.val, otel.m.service.name, otel.m.Host, otel.m.host, otel.m.ts.time_| -|_|_Sort: otel.m.ts.time ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST_| -|_|_Filter: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS Float64) <= Float64(1)_| -|_|_WindowAggr: windowExpr=[[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]_| +|_|_Sort: otel.m.ts.time ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [isnan(otel.m.v.val) ASC NULLS FIRST, otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST_| +|_|_Filter: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [isnan(otel.m.v.val) ASC NULLS FIRST, otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS Float64) <= Float64(1)_| +|_|_WindowAggr: windowExpr=[[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [isnan(otel.m.v.val) ASC NULLS FIRST, otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]_| |_|_MergeScan [is_placeholder=false, remote_input=[_| |_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| |_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| @@ -435,10 +435,10 @@ TQL EXPLAIN (0, 10, '5s') topk(1, {"otel.m"}); |_|_TableScan: otel.m, partial_filters=[otel.m.ts.time >= TimestampMillisecond(-299999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)]_| |_| ]]_| | physical_plan | ProjectionExec: expr=[v.val@4 as v.val, service.name@1 as service.name, Host@2 as Host, host@3 as host, ts.time@0 as ts.time]_| -|_|_SortPreservingMergeExec: [ts.time@0 ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 ASC NULLS LAST]_| -|_|_FilterExec: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 AS Float64) <= 1_| -|_|_BoundedWindowAggExec: wdw=[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] | -|_|_SortExec: expr=[ts.time@0 ASC NULLS LAST, v.val@4 DESC, service.name@1 DESC, Host@2 DESC, host@3 DESC], preserve_REDACTED +|_|_SortPreservingMergeExec: [ts.time@0 ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [isnan(otel.m.v.val) ASC NULLS FIRST, otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 ASC NULLS LAST]_| +|_|_FilterExec: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [isnan(otel.m.v.val) ASC NULLS FIRST, otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 AS Float64) <= 1_| +|_|_BoundedWindowAggExec: wdw=[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [isnan(otel.m.v.val) ASC NULLS FIRST, otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [otel.m.ts.time] ORDER BY [isnan(otel.m.v.val) ASC NULLS FIRST, otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] | +|_|_SortExec: expr=[ts.time@0 ASC NULLS LAST, isnan(v.val@4) ASC, v.val@4 DESC, service.name@1 DESC, Host@2 DESC, host@3 DESC], preserve_REDACTED |_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| diff --git a/tests/cases/standalone/common/promql/float_semantics.result b/tests/cases/standalone/common/promql/float_semantics.result new file mode 100644 index 00000000000..e56db8b0ba6 --- /dev/null +++ b/tests/cases/standalone/common/promql/float_semantics.result @@ -0,0 +1,115 @@ +-- NaN and rounding follow Prometheus (IEEE 754) rather than SQL ordering. +-- Expected results are from Prometheus 3. +create table float_semantics ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +insert into float_semantics values + (0, 1.0, 'a'), + (0, 2.0, 'b'), + (0, 'NaN'::double, 'c'), + (0, -2.5, 'd'); + +Affected Rows: 4 + +tql eval (0, 0, '1s') max(float_semantics); + ++---------------------+--------------------------+ +| ts | max(float_semantics.val) | ++---------------------+--------------------------+ +| 1970-01-01T00:00:00 | 2.0 | ++---------------------+--------------------------+ + +tql eval (0, 0, '1s') min(float_semantics); + ++---------------------+--------------------------+ +| ts | min(float_semantics.val) | ++---------------------+--------------------------+ +| 1970-01-01T00:00:00 | -2.5 | ++---------------------+--------------------------+ + +tql eval (0, 0, '1s') topk(1, float_semantics); + ++-----+------+---------------------+ +| val | host | ts | ++-----+------+---------------------+ +| 2.0 | b | 1970-01-01T00:00:00 | ++-----+------+---------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') float_semantics <= NaN; + ++----+-----+------+ +| ts | val | host | ++----+-----+------+ ++----+-----+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') float_semantics != bool NaN; + ++------+---------------------+---------------------+ +| host | ts | val != Float64(NaN) | ++------+---------------------+---------------------+ +| a | 1970-01-01T00:00:00 | 1.0 | +| b | 1970-01-01T00:00:00 | 1.0 | +| c | 1970-01-01T00:00:00 | 1.0 | +| d | 1970-01-01T00:00:00 | 1.0 | ++------+---------------------+---------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') round(float_semantics); + ++---------------------+----------------------------+------+ +| ts | prom_round(val,Float64(1)) | host | ++---------------------+----------------------------+------+ +| 1970-01-01T00:00:00 | -2.0 | d | +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 2.0 | b | +| 1970-01-01T00:00:00 | NaN | c | ++---------------------+----------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') sqrt(float_semantics); + ++---------------------+--------------------+------+ +| ts | sqrt(val) | host | ++---------------------+--------------------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 1.4142135623730951 | b | +| 1970-01-01T00:00:00 | NaN | c | +| 1970-01-01T00:00:00 | NaN | d | ++---------------------+--------------------+------+ + +-- prom_min/prom_max have their own column names, so SQL can use them next to min/max. +select min(val) as sql_lo, prom_min(val) as prom_lo from float_semantics; + ++--------+---------+ +| sql_lo | prom_lo | ++--------+---------+ +| -2.5 | -2.5 | ++--------+---------+ + +select max(val) as sql_hi, prom_max(val) as prom_hi from float_semantics; + ++--------+---------+ +| sql_hi | prom_hi | ++--------+---------+ +| NaN | 2.0 | ++--------+---------+ + +select prom_min(val) as lo from float_semantics having min(val) < 0; + ++------+ +| lo | ++------+ +| -2.5 | ++------+ + +drop table float_semantics; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/float_semantics.sql b/tests/cases/standalone/common/promql/float_semantics.sql new file mode 100644 index 00000000000..f23e95141c4 --- /dev/null +++ b/tests/cases/standalone/common/promql/float_semantics.sql @@ -0,0 +1,41 @@ +-- NaN and rounding follow Prometheus (IEEE 754) rather than SQL ordering. +-- Expected results are from Prometheus 3. + +create table float_semantics ( + ts timestamp time index, + val double, + host string primary key +); + +insert into float_semantics values + (0, 1.0, 'a'), + (0, 2.0, 'b'), + (0, 'NaN'::double, 'c'), + (0, -2.5, 'd'); + +tql eval (0, 0, '1s') max(float_semantics); + +tql eval (0, 0, '1s') min(float_semantics); + +tql eval (0, 0, '1s') topk(1, float_semantics); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') float_semantics <= NaN; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') float_semantics != bool NaN; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') round(float_semantics); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (0, 0, '1s') sqrt(float_semantics); + +-- prom_min/prom_max have their own column names, so SQL can use them next to min/max. +select min(val) as sql_lo, prom_min(val) as prom_lo from float_semantics; + +select max(val) as sql_hi, prom_max(val) as prom_hi from float_semantics; + +select prom_min(val) as lo from float_semantics having min(val) < 0; + +drop table float_semantics; diff --git a/tests/cases/standalone/common/promql/functions.result b/tests/cases/standalone/common/promql/functions.result index 325fbccd17d..1867413aec8 100644 --- a/tests/cases/standalone/common/promql/functions.result +++ b/tests/cases/standalone/common/promql/functions.result @@ -109,7 +109,7 @@ tql eval (4, 4, '1s') quantile_over_time(0.4 + 0.1, prom_series_q[4s]); tql eval (3, 3, '1s') round(prom_series_q); +---------------------+----------------------------+------+ -| ts | prom_round(val,Float64(0)) | host | +| ts | prom_round(val,Float64(1)) | host | +---------------------+----------------------------+------+ | 1970-01-01T00:00:03 | 346.0 | q | +---------------------+----------------------------+------+ @@ -133,8 +133,8 @@ tql eval (1, 4, '1s') round(prom_series_q, 0.05 + 0.05); | ts | prom_round(val,Float64(0.05) + Float64(0.05)) | host | +---------------------+-----------------------------------------------+------+ | 1970-01-01T00:00:01 | 123.5 | q | -| 1970-01-01T00:00:02 | 234.60000000000002 | q | -| 1970-01-01T00:00:03 | 345.70000000000005 | q | +| 1970-01-01T00:00:02 | 234.6 | q | +| 1970-01-01T00:00:03 | 345.7 | q | | 1970-01-01T00:00:04 | 456.8 | q | +---------------------+-----------------------------------------------+------+ @@ -156,10 +156,10 @@ tql eval (1, 4, '1s') round(prom_series_q, 100.0 + 3.0); +---------------------+-------------------------------------------+------+ | ts | prom_round(val,Float64(100) + Float64(3)) | host | +---------------------+-------------------------------------------+------+ -| 1970-01-01T00:00:01 | 103.0 | q | -| 1970-01-01T00:00:02 | 206.0 | q | +| 1970-01-01T00:00:01 | 103.00000000000001 | q | +| 1970-01-01T00:00:02 | 206.00000000000003 | q | | 1970-01-01T00:00:03 | 309.0 | q | -| 1970-01-01T00:00:04 | 412.0 | q | +| 1970-01-01T00:00:04 | 412.00000000000006 | q | +---------------------+-------------------------------------------+------+ -- SQLNESS SORT_RESULT 3 1 diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index fc4bc6f01c2..0288a045390 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -336,7 +336,7 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__tsid@3, __tsid@1), (ts@4, ts@2)], filter=greptime_value@1 < greptime_value@0, projection=[ts@4, greptime_value@0, host@1, job@2, __tsid@3], NullsEqual: true REDACTED +| 0_| 0_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__tsid@3, __tsid@1), (ts@4, ts@2)], filter=prom_gt(greptime_value@0, greptime_value@1), projection=[ts@4, greptime_value@0, host@1, job@2, __tsid@3], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([__tsid@3, ts@4],REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_RepartitionExec: partitioning=Hash([__tsid@1, ts@2],REDACTED @@ -374,8 +374,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[host@0 as host, job@1 as job, ts@2 as ts, __tsid@3 as __tsid, CAST(greptime_value@4 < greptime_value@5 AS Float64) as tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@2, ts@4)], projection=[host@4, job@5, ts@7, __tsid@6, greptime_value@3, greptime_value@0], NullsEqual: true REDACTED +| 0_| 0_|_ProjectionExec: expr=[host@0 as host, job@1 as job, ts@2 as ts, __tsid@3 as __tsid, CAST(prom_gt(greptime_value@4, greptime_value@5) AS Float64) as tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value] REDACTED +|_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@2, ts@4)], projection=[host@4, job@5, ts@7, __tsid@6, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED @@ -416,7 +416,7 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi | 0_| 0_|_ProjectionExec: expr=[host@0 as host, job@1 as job, ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 / greptime_value@5 * 100 as .greptime_value / tsid_binary_join_left.greptime_value * Float64(100)] REDACTED |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@2, __tsid@3), (ts@0, ts@4)], projection=[host@4, job@5, ts@7, __tsid@6, greptime_value@1, greptime_value@3], NullsEqual: true REDACTED |_|_|_CoalescePartitionsExec REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__tsid@1, __tsid@1), (ts@2, ts@2)], filter=greptime_value@1 < greptime_value@0, projection=[ts@2, greptime_value@0, __tsid@1], NullsEqual: true REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__tsid@1, __tsid@1), (ts@2, ts@2)], filter=prom_gt(greptime_value@0, greptime_value@1), projection=[ts@2, greptime_value@0, __tsid@1], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED @@ -467,8 +467,8 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_CoalescePartitionsExec REDACTED |_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value@2 + greptime_value@3 as tsid_binary_join_right.tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value + tsid_binary_join_left.greptime_value] REDACTED |_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__tsid@1, __tsid@1), (ts@0, ts@2)], projection=[ts@5, __tsid@4, tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value@2, greptime_value@3], NullsEqual: true REDACTED -|_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, CAST(greptime_value@2 < greptime_value@3 AS Float64) as tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__tsid@1, __tsid@1), (ts@2, ts@2)], projection=[ts@5, __tsid@4, greptime_value@3, greptime_value@0], NullsEqual: true REDACTED +|_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, CAST(prom_gt(greptime_value@2, greptime_value@3) AS Float64) as tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value] REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__tsid@1, __tsid@1), (ts@2, ts@2)], projection=[ts@5, __tsid@4, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED From 2d99a5cf69eb5c0eb7f7dc54335c3006f2d46c00 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 12 Oct 2026 05:11:03 +0000 Subject: [PATCH 08/46] feat(jaeger): add query endpoints for Jaeger UI v2 (#9409) * feat(jaeger): support v3 HTTP query APIs Signed-off-by: Dennis Zhuang * fix(jaeger): preserve full traces with v3 time hints Signed-off-by: Dennis Zhuang * docs(jaeger): clarify validation of v3 lookup hints Signed-off-by: Dennis Zhuang * fix(jaeger): bound v3 searches and preserve v1 binary attributes Signed-off-by: Dennis Zhuang * fix(jaeger): match resource attributes in v0 tag filters Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance.rs | 2 +- src/frontend/src/instance/jaeger.rs | 91 ++-- src/servers/src/http.rs | 20 + src/servers/src/http/jaeger.rs | 16 +- src/servers/src/http/jaeger/v3.rs | 542 ++++++++++++++++++++++ src/servers/src/http/jaeger/v3/otlp.rs | 294 ++++++++++++ src/servers/src/http/jaeger/v3/summary.rs | 142 ++++++ tests-integration/tests/http.rs | 319 +++++++++++++ 8 files changed, 1364 insertions(+), 62 deletions(-) create mode 100644 src/servers/src/http/jaeger/v3.rs create mode 100644 src/servers/src/http/jaeger/v3/otlp.rs create mode 100644 src/servers/src/http/jaeger/v3/summary.rs diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index b91290d2a55..df30106c54a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -3477,7 +3477,7 @@ mod tests { &instance, ctx.clone(), servers::http::jaeger::QueryTraceParams { - service_name: "service".to_string(), + service_name: Some("service".to_string()), ..Default::default() }, ) diff --git a/src/frontend/src/instance/jaeger.rs b/src/frontend/src/instance/jaeger.rs index c9b5bd9e988..ac59698678b 100644 --- a/src/frontend/src/instance/jaeger.rs +++ b/src/frontend/src/instance/jaeger.rs @@ -202,8 +202,9 @@ impl JaegerQueryHandler for Instance { let mut filters = vec![]; - // `service_name` is already validated in `from_jaeger_query_params()`, so no additional check needed here. - filters.push(col(SERVICE_NAME_COLUMN).eq(lit(query_params.service_name))); + if let Some(service_name) = query_params.service_name { + filters.push(col(SERVICE_NAME_COLUMN).eq(lit(service_name))); + } if let Some(operation_name) = query_params.operation_name { filters.push(col(SPAN_NAME_COLUMN).eq(lit(operation_name))); @@ -254,6 +255,7 @@ impl JaegerQueryHandler for Instance { .await?; // Get all traces that match the trace ids from the previous query. + // API v3 applies the time window only to the search, not to the returned spans. // It's equivalent to the following SQL query: // // ``` @@ -276,11 +278,15 @@ impl JaegerQueryHandler for Instance { ), ]; - if let Some(start_time) = query_params.start_time { + if !query_params.fetch_full_trace + && let Some(start_time) = query_params.start_time + { filters.push(col(TIMESTAMP_COLUMN).gt_eq(lit_timestamp_nano(start_time))); } - if let Some(end_time) = query_params.end_time { + if !query_params.fetch_full_trace + && let Some(end_time) = query_params.end_time + { filters.push(col(TIMESTAMP_COLUMN).lt_eq(lit_timestamp_nano(end_time))); } @@ -540,64 +546,37 @@ fn json_tag_filters( dataframe: &DataFrame, tags: HashMap, ) -> ServerResult> { + let has_resource_attributes = dataframe + .schema() + .has_column_with_unqualified_name(RESOURCE_ATTRIBUTES_COLUMN); let mut filters = vec![]; // NOTE: The key of the tags may contain `.`, for example: `http.status_code`, so we need to use `["http.status_code"]` in json path to access the value. for (key, value) in tags.iter() { - if let JsonValue::String(value) = value { - filters.push( - dataframe - .registry() - .udf(JsonGetString::NAME) - .context(DataFusionSnafu)? - .call(vec![ - col(SPAN_ATTRIBUTES_COLUMN), - lit(format!("[\"{}\"]", key)), - ]) - .eq(lit(value)), - ); - } - if let JsonValue::Number(value) = value { - if value.is_i64() { - filters.push( - dataframe - .registry() - .udf(JsonGetInt::NAME) - .context(DataFusionSnafu)? - .call(vec![ - col(SPAN_ATTRIBUTES_COLUMN), - lit(format!("[\"{}\"]", key)), - ]) - .eq(lit(value.as_i64().unwrap())), - ); + let (udf, value) = match value { + JsonValue::String(value) => (JsonGetString::NAME, lit(value)), + JsonValue::Number(value) if value.is_i64() => { + (JsonGetInt::NAME, lit(value.as_i64().unwrap())) } - if value.is_f64() { - filters.push( - dataframe - .registry() - .udf(JsonGetFloat::NAME) - .context(DataFusionSnafu)? - .call(vec![ - col(SPAN_ATTRIBUTES_COLUMN), - lit(format!("[\"{}\"]", key)), - ]) - .eq(lit(value.as_f64().unwrap())), - ); + JsonValue::Number(value) if value.is_f64() => { + (JsonGetFloat::NAME, lit(value.as_f64().unwrap())) } - } - if let JsonValue::Bool(value) = value { - filters.push( - dataframe - .registry() - .udf(JsonGetBool::NAME) - .context(DataFusionSnafu)? - .call(vec![ - col(SPAN_ATTRIBUTES_COLUMN), - lit(format!("[\"{}\"]", key)), - ]) - .eq(lit(*value)), - ); - } + JsonValue::Bool(value) => (JsonGetBool::NAME, lit(*value)), + _ => continue, + }; + let get = dataframe.registry().udf(udf).context(DataFusionSnafu)?; + let path = lit(format!("[\"{}\"]", key)); + let span_value = get.call(vec![col(SPAN_ATTRIBUTES_COLUMN), path.clone()]); + // Same precedence as `json2_tag_filters`: a span attribute shadows a resource attribute. + let attribute = if has_resource_attributes { + coalesce(vec![ + span_value, + get.call(vec![col(RESOURCE_ATTRIBUTES_COLUMN), path]), + ]) + } else { + span_value + }; + filters.push(attribute.eq(value)); } Ok(filters) diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 2653ceddbe5..210fe0cee00 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -1576,6 +1576,26 @@ impl HttpServer { "/api/traces/{trace_id}", routing::get(jaeger::handle_get_trace), ) + .route( + "/api/v3/services", + routing::get(jaeger::v3::handle_get_services), + ) + .route( + "/api/v3/operations", + routing::get(jaeger::v3::handle_get_operations), + ) + .route( + "/api/v3/traces", + routing::get(jaeger::v3::handle_find_traces), + ) + .route( + "/api/v3/trace-summaries", + routing::get(jaeger::v3::handle_find_trace_summaries), + ) + .route( + "/api/v3/traces/{trace_id}", + routing::get(jaeger::v3::handle_get_trace), + ) .with_state(handler) } diff --git a/src/servers/src/http/jaeger.rs b/src/servers/src/http/jaeger.rs index 16dbaf76e53..f636ccb3df8 100644 --- a/src/servers/src/http/jaeger.rs +++ b/src/servers/src/http/jaeger.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod v3; + use std::collections::{BTreeMap, HashMap}; use std::fmt; use std::str::FromStr; @@ -289,9 +291,9 @@ fn update_query_context(query_ctx: &mut QueryContext, table_name: Option impl QueryTraceParams { fn from_jaeger_query_params(query_params: JaegerQueryParams) -> Result { let mut internal_query_params: QueryTraceParams = QueryTraceParams { - service_name: query_params.service_name.context(InvalidJaegerQuerySnafu { + service_name: Some(query_params.service_name.context(InvalidJaegerQuerySnafu { reason: "service_name is required".to_string(), - })?, + })?), operation_name: query_params.operation_name, // Convert start time from microseconds to nanoseconds. start_time: query_params.start.map(|start| start * 1000), @@ -347,7 +349,7 @@ impl QueryTraceParams { #[derive(Debug, Default, PartialEq)] pub struct QueryTraceParams { - pub service_name: String, + pub service_name: Option, pub operation_name: Option, // The limit of the number of traces to return. @@ -362,6 +364,9 @@ pub struct QueryTraceParams { pub min_duration: Option, pub max_duration: Option, + /// Fetch all spans of matching traces, including spans outside the search window. + pub fetch_full_trace: bool, + // The user agent of the trace query, mainly find traces pub user_agent: TraceUserAgent, } @@ -1647,7 +1652,7 @@ mod tests { ..Default::default() }, QueryTraceParams { - service_name: "test-service-0".to_string(), + service_name: Some("test-service-0".to_string()), ..Default::default() }, ), @@ -1664,7 +1669,7 @@ mod tests { ..Default::default() }, QueryTraceParams { - service_name: "test-service-0".to_string(), + service_name: Some("test-service-0".to_string()), operation_name: Some("access-mysql".to_string()), start_time: Some(1738726754492422000), end_time: Some(1738726754642422000), @@ -1679,6 +1684,7 @@ mod tests { ("http.path".to_string(), JsonValue::String("/api/v1/users".to_string())), ])), user_agent: TraceUserAgent::Jaeger, + fetch_full_trace: false, }, ), ]; diff --git a/src/servers/src/http/jaeger/v3.rs b/src/servers/src/http/jaeger/v3.rs new file mode 100644 index 00000000000..ccb8a235fd3 --- /dev/null +++ b/src/servers/src/http/jaeger/v3.rs @@ -0,0 +1,542 @@ +// Copyright 2023 Greptime Team +// +// 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. + +mod otlp; +mod summary; + +use std::collections::HashMap; +use std::sync::Arc; + +use axum::extract::rejection::QueryRejection; +use axum::extract::{Path, Query, State}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::{Extension, Json}; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode as GreptimeStatusCode; +use common_query::Output; +use common_telemetry::error; +use rust_decimal::Decimal; +use rust_decimal::prelude::ToPrimitive; +use serde::Deserialize; +use serde_json::{Value, json}; +use session::context::QueryContext; + +use crate::error::{Error, status_code_to_http_status}; +use crate::http::extractor::TraceTableName; +use crate::http::jaeger::{ + QueryTraceParams, convert_string_to_boolean, convert_string_to_number, covert_to_records, + empty_string_as_none, operations_from_records, services_from_records, update_query_context, +}; +use crate::metrics::METRIC_JAEGER_QUERY_ELAPSED; +use crate::query_handler::JaegerQueryHandlerRef; + +/// Jaeger API v3 uses an error envelope rather than the legacy `errors` array. +#[derive(Debug)] +pub struct ApiError { + status: StatusCode, + message: String, +} + +impl ApiError { + fn invalid(message: impl Into) -> Self { + Self { + status: StatusCode::BAD_REQUEST, + message: message.into(), + } + } + + fn not_found() -> Self { + Self { + status: StatusCode::NOT_FOUND, + message: "No traces found".into(), + } + } +} + +impl From for ApiError { + fn from(err: Error) -> Self { + error!(err; "Jaeger v3 query failed"); + Self { + status: status_code_to_http_status(&err.status_code()), + message: err.output_msg(), + } + } +} + +impl From for ApiError { + fn from(err: QueryRejection) -> Self { + Self::invalid(err.body_text()) + } +} + +impl IntoResponse for ApiError { + fn into_response(self) -> Response { + ( + self.status, + Json(json!({"error": {"httpCode": self.status.as_u16(), "message": self.message}})), + ) + .into_response() + } +} + +#[derive(Debug, Default, Deserialize)] +pub struct OperationsParams { + service: String, + #[serde(rename = "spanKind", alias = "span_kind")] + span_kind: Option, +} + +#[derive(Debug, Default, Deserialize)] +pub struct GetTraceParams { + #[serde(rename = "startTime", alias = "start_time")] + #[serde(default, deserialize_with = "empty_string_as_none")] + start_time: Option, + #[serde(rename = "endTime", alias = "end_time")] + #[serde(default, deserialize_with = "empty_string_as_none")] + end_time: Option, + // GreptimeDB returns stored spans without Jaeger's clock-skew enrichment. + #[serde(rename = "rawTraces", alias = "raw_traces")] + #[serde(default, deserialize_with = "empty_string_as_none")] + _raw_traces: Option, +} + +#[derive(Debug, Default, Deserialize)] +pub struct FindTracesParams { + #[serde(rename = "query.serviceName", alias = "query.service_name")] + service_name: Option, + #[serde(rename = "query.operationName", alias = "query.operation_name")] + operation_name: Option, + #[serde(rename = "query.startTimeMin", alias = "query.start_time_min")] + start_time_min: String, + #[serde(rename = "query.startTimeMax", alias = "query.start_time_max")] + start_time_max: String, + #[serde(rename = "query.durationMin", alias = "query.duration_min")] + #[serde(default, deserialize_with = "empty_string_as_none")] + duration_min: Option, + #[serde(rename = "query.durationMax", alias = "query.duration_max")] + #[serde(default, deserialize_with = "empty_string_as_none")] + duration_max: Option, + #[serde( + rename = "query.searchDepth", + alias = "query.search_depth", + alias = "query.num_traces" + )] + #[serde(default, deserialize_with = "empty_string_as_none")] + search_depth: Option, + #[serde(rename = "query.attributes")] + #[serde(default, deserialize_with = "empty_string_as_none")] + attributes: Option, + #[serde(rename = "query.rawTraces", alias = "query.raw_traces")] + #[serde(default, deserialize_with = "empty_string_as_none")] + _raw_traces: Option, + #[serde(rename = "query.filter")] + filter: Option, +} + +impl FindTracesParams { + fn into_query(self) -> Result { + if self.filter.is_some() { + return Err(ApiError::invalid("query.filter is not supported")); + } + let start = parse_timestamp(&self.start_time_min)?; + let end = parse_timestamp(&self.start_time_max)?; + if start >= end { + return Err(ApiError::invalid( + "query.startTimeMin must be before query.startTimeMax", + )); + } + let limit = match self.search_depth.unwrap_or(0) { + 0 => 100, + limit @ 1..=10_000 => limit, + _ => { + return Err(ApiError::invalid( + "query.searchDepth must be between 0 and 10000", + )); + } + }; + let tags = self + .attributes + .map(|attributes| { + let values = serde_json::from_str::>(&attributes) + .map_err(|err| ApiError::invalid(format!("Invalid query.attributes: {err}")))?; + Ok::<_, ApiError>( + values + .into_iter() + .map(|(key, value)| { + let value = Value::String(value); + let value = convert_string_to_number(&value) + .or_else(|| convert_string_to_boolean(&value)) + .unwrap_or(value); + (key, value) + }) + .collect(), + ) + }) + .transpose()?; + let query = QueryTraceParams { + service_name: self.service_name.filter(|name| !name.is_empty()), + operation_name: self.operation_name.filter(|name| !name.is_empty()), + start_time: Some(start), + // The backend uses an inclusive bound; API v3 specifies an exclusive bound. + end_time: Some(end - 1), + min_duration: self + .duration_min + .as_deref() + .map(parse_duration) + .transpose()?, + max_duration: self + .duration_max + .as_deref() + .map(parse_duration) + .transpose()?, + limit: Some(limit as usize), + tags, + fetch_full_trace: true, + ..Default::default() + }; + if let (Some(min), Some(max)) = (query.min_duration, query.max_duration) + && min > max + { + return Err(ApiError::invalid( + "query.durationMin must not exceed query.durationMax", + )); + } + Ok(query) + } +} + +fn parse_timestamp(value: &str) -> Result { + chrono::DateTime::parse_from_rfc3339(value) + .map_err(|err| ApiError::invalid(format!("Invalid RFC3339 timestamp '{value}': {err}")))? + .timestamp_nanos_opt() + .ok_or_else(|| ApiError::invalid("Timestamp is outside the supported nanosecond range")) +} + +fn parse_duration(value: &str) -> Result { + let invalid = || ApiError::invalid(format!("Invalid Go duration '{value}'")); + let mut rest = value.strip_prefix('+').unwrap_or(value); + if rest == "0" { + return Ok(0); + } + if rest.is_empty() { + return Err(invalid()); + } + let mut total = 0_u64; + while !rest.is_empty() { + let number_end = rest + .find(|ch: char| !ch.is_ascii_digit() && ch != '.') + .ok_or_else(invalid)?; + let amount: Decimal = rest[..number_end].parse().map_err(|_| invalid())?; + rest = &rest[number_end..]; + let unit_end = rest + .find(|ch: char| ch.is_ascii_digit() || ch == '.') + .unwrap_or(rest.len()); + let unit = match &rest[..unit_end] { + "ns" => 1_u64, + "us" | "µs" | "μs" => 1_000, + "ms" => 1_000_000, + "s" => 1_000_000_000, + "m" => 60_000_000_000, + "h" => 3_600_000_000_000, + _ => return Err(invalid()), + }; + let nanos = amount + .checked_mul(Decimal::from(unit)) + .and_then(|amount| amount.trunc().to_u64()) + .ok_or_else(invalid)?; + total = total + .checked_add(nanos) + .filter(|total| *total <= i64::MAX as u64) + .ok_or_else(invalid)?; + rest = &rest[unit_end..]; + } + Ok(total) +} + +pub async fn handle_get_services( + State(handler): State, + Extension(mut ctx): Extension, + TraceTableName(table): TraceTableName, +) -> Result, ApiError> { + update_query_context(&mut ctx, table); + let _timer = METRIC_JAEGER_QUERY_ELAPSED + .with_label_values(&[&ctx.get_db_string(), "/api/v3/services"]) + .start_timer(); + let services = match query_records(handler.get_services(Arc::new(ctx)).await).await? { + Some(records) => services_from_records(records)?, + None => vec![], + }; + Ok(Json(json!({"services": services}))) +} + +pub async fn handle_get_operations( + State(handler): State, + params: Result, QueryRejection>, + Extension(mut ctx): Extension, + TraceTableName(table): TraceTableName, +) -> Result, ApiError> { + let Query(params) = params?; + if params.service.is_empty() { + return Err(ApiError::invalid("service is required")); + } + update_query_context(&mut ctx, table); + let _timer = METRIC_JAEGER_QUERY_ELAPSED + .with_label_values(&[&ctx.get_db_string(), "/api/v3/operations"]) + .start_timer(); + let operations = match query_records( + handler + .get_operations( + Arc::new(ctx), + ¶ms.service, + params.span_kind.as_deref().filter(|kind| !kind.is_empty()), + ) + .await, + ) + .await? + { + Some(records) => operations_from_records(records, true)?, + None => vec![], + }; + let operations: Vec<_> = operations.into_iter().map(|operation| json!({ + "name": operation.name, + "spanKind": operation.span_kind.filter(|kind| !kind.is_empty()).unwrap_or_else(|| "internal".into()), + })).collect(); + Ok(Json(json!({"operations": operations}))) +} + +pub async fn handle_get_trace( + State(handler): State, + Path(trace_id): Path, + params: Result, QueryRejection>, + Extension(mut ctx): Extension, + TraceTableName(table): TraceTableName, +) -> Result, ApiError> { + let Query(params) = params?; + if !matches!(trace_id.len(), 16 | 32) || !trace_id.bytes().all(|ch| ch.is_ascii_hexdigit()) { + return Err(ApiError::invalid( + "trace_id must be a 64-bit or 128-bit hexadecimal ID", + )); + } + let trace_id = format!("{:0>32}", trace_id.to_ascii_lowercase()); + let start = params + .start_time + .as_deref() + .map(parse_timestamp) + .transpose()?; + let end = params + .end_time + .as_deref() + .map(parse_timestamp) + .transpose()?; + if let (Some(start), Some(end)) = (start, end) + && start > end + { + return Err(ApiError::invalid("startTime must not exceed endTime")); + } + update_query_context(&mut ctx, table); + let _timer = METRIC_JAEGER_QUERY_ELAPSED + .with_label_values(&[&ctx.get_db_string(), "/api/v3/traces/{trace_id}"]) + .start_timer(); + // API v3 time bounds are lookup hints. Applying them as span filters would truncate the trace, + // so they are only validated. + trace_response( + handler + .get_trace(Arc::new(ctx), &trace_id, None, None, None) + .await, + ) + .await +} + +pub async fn handle_find_traces( + State(handler): State, + params: Result, QueryRejection>, + Extension(mut ctx): Extension, + TraceTableName(table): TraceTableName, +) -> Result, ApiError> { + let Query(params) = params?; + let query = params.into_query()?; + update_query_context(&mut ctx, table); + let _timer = METRIC_JAEGER_QUERY_ELAPSED + .with_label_values(&[&ctx.get_db_string(), "/api/v3/traces"]) + .start_timer(); + trace_response(handler.find_traces(Arc::new(ctx), query).await).await +} + +pub async fn handle_find_trace_summaries( + State(handler): State, + params: Result, QueryRejection>, + Extension(mut ctx): Extension, + TraceTableName(table): TraceTableName, +) -> Result, ApiError> { + let Query(params) = params?; + let query = params.into_query()?; + update_query_context(&mut ctx, table); + let _timer = METRIC_JAEGER_QUERY_ELAPSED + .with_label_values(&[&ctx.get_db_string(), "/api/v3/trace-summaries"]) + .start_timer(); + let summaries = match query_records(handler.find_traces(Arc::new(ctx), query).await).await? { + Some(records) => summary::summarize(otlp::traces_from_records(records)?), + None => vec![], + }; + Ok(Json(json!({"summaries": summaries}))) +} + +async fn query_records( + output: crate::error::Result, +) -> Result, ApiError> { + match output { + Ok(output) => covert_to_records(output) + .await? + .map(Some) + .ok_or_else(|| ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: "Unexpected Jaeger query output".into(), + }), + Err(err) if err.status_code() == GreptimeStatusCode::TableNotFound => Ok(None), + Err(err) => Err(err.into()), + } +} + +async fn trace_response(output: crate::error::Result) -> Result, ApiError> { + let records = query_records(output) + .await? + .ok_or_else(ApiError::not_found)?; + if records.rows.is_empty() { + return Err(ApiError::not_found()); + } + Ok(Json(json!({"result": otlp::traces_from_records(records)?}))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_find_traces_parameters() { + for (start, end, service, operation, depth, duration) in [ + ( + "start_time_min", + "start_time_max", + "service_name", + "operation_name", + "num_traces", + "duration_min", + ), + ( + "startTimeMin", + "startTimeMax", + "serviceName", + "operationName", + "searchDepth", + "durationMin", + ), + ] { + let uri = format!("/api/v3/traces?query.{start}=2026-01-01T00:00:00.000000001Z&query.{end}=2026-01-01T00:00:01Z&query.{service}=checkout&query.{operation}=GET&query.{depth}=7&query.{duration}=1.5ms&query.attributes=%7B%22error%22%3A%22true%22%7D").parse().unwrap(); + let Query(params) = Query::::try_from_uri(&uri).unwrap(); + let query = params.into_query().unwrap(); + assert_eq!(query.service_name.as_deref(), Some("checkout")); + assert_eq!(query.operation_name.as_deref(), Some("GET")); + assert_eq!(query.limit, Some(7)); + assert_eq!(query.min_duration, Some(1_500_000)); + assert_eq!(query.start_time, Some(1767225600000000001)); + assert_eq!(query.end_time, Some(1767225600999999999)); + assert_eq!(query.tags.unwrap()["error"], json!(true)); + assert!(query.fetch_full_trace); + } + let uri = "/api/v3/traces?query.startTimeMin=2026-01-01T00:00:00Z&query.startTimeMax=2026-01-01T01:00:00Z&query.durationMin=&query.durationMax=&query.searchDepth=&query.rawTraces=&query.attributes=".parse().unwrap(); + let Query(params) = Query::::try_from_uri(&uri).unwrap(); + let query = params.into_query().unwrap(); + assert_eq!(query.service_name, None); + assert_eq!(query.limit, Some(100)); + for name in [ + "query.searchDepth", + "query.search_depth", + "query.num_traces", + ] { + for (depth, expected) in [(0, 100), (1, 1), (10_000, 10_000)] { + let uri = format!("/api/v3/traces?query.startTimeMin=2026-01-01T00:00:00Z&query.startTimeMax=2026-01-01T01:00:00Z&{name}={depth}").parse().unwrap(); + let Query(params) = Query::::try_from_uri(&uri).unwrap(); + assert_eq!(params.into_query().unwrap().limit, Some(expected)); + } + } + } + + #[test] + fn test_go_durations() { + for (value, expected) in [ + ("0", 0), + ("1.5ms", 1_500_000), + ("1h2m3.5s", 3_723_500_000_000), + ("+.5us", 500), + ("1.µs", 1_000), + ("1μs", 1_000), + ("0.1ns", 0), + ("9223372036854775807ns", i64::MAX as u64), + ] { + assert_eq!(parse_duration(value).unwrap(), expected, "{value}"); + } + for value in [ + "", + "1", + "-1s", + "1.2.3s", + "s", + "1 s", + "1d", + "1s-1s", + "9223372036854775808ns", + ] { + assert!(parse_duration(value).is_err(), "{value}"); + } + } + + #[test] + fn test_invalid_find_traces_parameters() { + for extra in [ + "query.searchDepth=-1", + "query.searchDepth=10001", + "query.searchDepth=2147483647", + "query.durationMin=-1s", + "query.durationMin=2s&query.durationMax=1s", + "query.durationMax=999999999999h", + "query.attributes=%7B%22error%22%3Atrue%7D", + "query.attributes=invalid", + "query.filter=%7B%7D", + ] { + let uri = format!("/api/v3/traces?query.startTimeMin=2026-01-01T00:00:00Z&query.startTimeMax=2026-01-01T01:00:00Z&{extra}").parse().unwrap(); + let Query(params) = Query::::try_from_uri(&uri).unwrap(); + assert_eq!( + params.into_query().unwrap_err().status, + StatusCode::BAD_REQUEST, + "{extra}" + ); + } + for (start, end) in [ + ("invalid", "2026-01-01T00:00:00Z"), + ("2500-01-01T00:00:00Z", "2500-01-02T00:00:00Z"), + ("2026-01-01T00:00:00Z", "2026-01-01T00:00:00Z"), + ] { + assert!( + FindTracesParams { + start_time_min: start.into(), + start_time_max: end.into(), + ..Default::default() + } + .into_query() + .is_err() + ); + } + } +} diff --git a/src/servers/src/http/jaeger/v3/otlp.rs b/src/servers/src/http/jaeger/v3/otlp.rs new file mode 100644 index 00000000000..4d09924ad39 --- /dev/null +++ b/src/servers/src/http/jaeger/v3/otlp.rs @@ -0,0 +1,294 @@ +// Copyright 2023 Greptime Team +// +// 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 std::collections::{HashMap, HashSet}; + +use axum::http::StatusCode; +use opentelemetry_proto::tonic::common::v1::{ + AnyValue, ArrayValue, InstrumentationScope, KeyValue, KeyValueList, any_value, +}; +use opentelemetry_proto::tonic::resource::v1::Resource; +use opentelemetry_proto::tonic::trace::v1::span::{Event, Link, SpanKind}; +use opentelemetry_proto::tonic::trace::v1::status::StatusCode as SpanStatusCode; +use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans, Span, Status, TracesData}; +use serde_json::Value; + +use crate::http::HttpRecordsOutput; +use crate::http::jaeger::v3::ApiError; +use crate::otlp::trace::{ + KEY_SERVICE_NAME, PARENT_SPAN_ID_COLUMN, RESOURCE_ATTRIBUTES_COLUMN, SCOPE_ATTRIBUTES_COLUMN, + SCOPE_NAME_COLUMN, SCOPE_VERSION_COLUMN, SERVICE_NAME_COLUMN, SPAN_ATTRIBUTES_COLUMN, + SPAN_EVENTS_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_COLUMN, SPAN_LINKS_COLUMN, SPAN_NAME_COLUMN, + SPAN_STATUS_CODE, SPAN_STATUS_MESSAGE_COLUMN, TIMESTAMP_COLUMN, TIMESTAMP_END_COLUMN, + TRACE_ID_COLUMN, TRACE_STATE_COLUMN, +}; + +/// Converts stored trace rows directly to OTLP, without the legacy microsecond conversion. +pub fn traces_from_records(records: HttpRecordsOutput) -> Result { + let binary_columns: HashSet<_> = records + .schema + .column_schemas + .iter() + .filter(|column| column.data_type == "Binary") + .map(|column| column.name.as_str()) + .collect(); + let mut resource_spans = Vec::with_capacity(records.rows.len()); + for row in &records.rows { + let cells: HashMap<_, _> = records + .schema + .column_schemas + .iter() + .zip(row) + .map(|(column, cell)| (column.name.as_str(), cell)) + .collect(); + let string = |name| cells.get(name).and_then(|v| v.as_str()).unwrap_or_default(); + let timestamp = |name| { + cells + .get(name) + .and_then(|v| v.as_u64()) + .ok_or_else(|| invalid_data(format!("Invalid timestamp column: {name}"))) + }; + let span = Span { + trace_id: decode_id(string(TRACE_ID_COLUMN), 16)?, + span_id: decode_id(string(SPAN_ID_COLUMN), 8)?, + parent_span_id: if string(PARENT_SPAN_ID_COLUMN).is_empty() { + vec![] + } else { + decode_id(string(PARENT_SPAN_ID_COLUMN), 8)? + }, + trace_state: string(TRACE_STATE_COLUMN).into(), + name: string(SPAN_NAME_COLUMN).into(), + kind: SpanKind::from_str_name(string(SPAN_KIND_COLUMN)).unwrap_or(SpanKind::Unspecified) + as i32, + start_time_unix_nano: timestamp(TIMESTAMP_COLUMN)?, + end_time_unix_nano: timestamp(TIMESTAMP_END_COLUMN)?, + attributes: attributes(&cells, SPAN_ATTRIBUTES_COLUMN, &binary_columns)?, + events: events(cells.get(SPAN_EVENTS_COLUMN).copied())?, + links: links(cells.get(SPAN_LINKS_COLUMN).copied())?, + status: Some(Status { + code: SpanStatusCode::from_str_name(string(SPAN_STATUS_CODE)) + .unwrap_or(SpanStatusCode::Unset) as i32, + message: string(SPAN_STATUS_MESSAGE_COLUMN).into(), + }), + ..Default::default() + }; + let mut resource_attributes = + attributes(&cells, RESOURCE_ATTRIBUTES_COLUMN, &binary_columns)?; + if !resource_attributes + .iter() + .any(|attr| attr.key == KEY_SERVICE_NAME) + && !string(SERVICE_NAME_COLUMN).is_empty() + { + resource_attributes.push(KeyValue { + key: KEY_SERVICE_NAME.into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue( + string(SERVICE_NAME_COLUMN).into(), + )), + }), + ..Default::default() + }); + } + // Each row retains its own resource and scope. Service name alone does not identify + // a resource: different instances of one service can have different attributes. + resource_spans.push(ResourceSpans { + resource: Some(Resource { + attributes: resource_attributes, + ..Default::default() + }), + scope_spans: vec![ScopeSpans { + scope: Some(InstrumentationScope { + name: string(SCOPE_NAME_COLUMN).into(), + version: string(SCOPE_VERSION_COLUMN).into(), + attributes: attributes(&cells, SCOPE_ATTRIBUTES_COLUMN, &binary_columns)?, + ..Default::default() + }), + spans: vec![span], + ..Default::default() + }], + ..Default::default() + }); + } + Ok(TracesData { resource_spans }) +} + +fn invalid_data(message: impl Into) -> ApiError { + ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: message.into(), + } +} + +fn decode_id(value: &str, length: usize) -> Result, ApiError> { + let bytes = hex::decode(value).map_err(|_| invalid_data("Invalid stored span or trace ID"))?; + if bytes.len() != length { + return Err(invalid_data("Invalid stored span or trace ID length")); + } + Ok(bytes) +} + +fn attributes( + cells: &HashMap<&str, &Value>, + name: &str, + binary_columns: &HashSet<&str>, +) -> Result, ApiError> { + if let Some(value) = cells.get(name) { + return object_attributes(value); + } + let prefix = format!("{name}."); + let mut attrs = cells + .iter() + .filter_map(|(column, value)| { + column + .strip_prefix(&prefix) + .filter(|_| !value.is_null()) + .map(|key| (key, *value, binary_columns.contains(column))) + }) + .map(|(key, value, is_binary)| { + let value = if is_binary { + AnyValue { + value: Some(any_value::Value::BytesValue( + serde_json::from_value(value.clone()) + .map_err(|_| invalid_data("Invalid stored binary attribute"))?, + )), + } + } else { + any_value(value)? + }; + Ok(KeyValue { + key: key.into(), + value: Some(value), + ..Default::default() + }) + }) + .collect::, ApiError>>()?; + attrs.sort_unstable_by(|a, b| a.key.cmp(&b.key)); + Ok(attrs) +} + +fn object_attributes(value: &Value) -> Result, ApiError> { + match value { + Value::Null => Ok(vec![]), + Value::Object(values) => values + .iter() + .map(|(key, value)| { + Ok(KeyValue { + key: key.clone(), + value: Some(any_value(value)?), + ..Default::default() + }) + }) + .collect(), + _ => Err(invalid_data("Stored attributes must be a JSON object")), + } +} + +fn any_value(value: &Value) -> Result { + use any_value::Value as OtlpValue; + + let value = match value { + Value::Null => None, + Value::String(value) => Some(OtlpValue::StringValue(value.clone())), + Value::Bool(value) => Some(OtlpValue::BoolValue(*value)), + Value::Number(value) => Some(if let Some(value) = value.as_i64() { + OtlpValue::IntValue(value) + } else { + OtlpValue::DoubleValue( + value + .as_f64() + .ok_or_else(|| invalid_data("Invalid numeric attribute"))?, + ) + }), + Value::Array(values) => Some(OtlpValue::ArrayValue(ArrayValue { + values: values.iter().map(any_value).collect::>()?, + })), + Value::Object(_) => Some(OtlpValue::KvlistValue(KeyValueList { + values: object_attributes(value)?, + })), + }; + Ok(AnyValue { value }) +} + +fn events(value: Option<&Value>) -> Result, ApiError> { + let Some(value) = value.filter(|value| !value.is_null()) else { + return Ok(vec![]); + }; + let values = value + .as_array() + .ok_or_else(|| invalid_data("Invalid stored span events"))?; + values + .iter() + .map(|event| { + let time = event["time"] + .as_str() + .ok_or_else(|| invalid_data("Missing span event time"))?; + let time = chrono::DateTime::parse_from_str(time, "%Y-%m-%d %H:%M:%S%.f%z") + .ok() + .and_then(|time| time.timestamp_nanos_opt()) + .and_then(|time| u64::try_from(time).ok()) + .ok_or_else(|| invalid_data("Invalid stored span event time"))?; + Ok(Event { + time_unix_nano: time, + name: event["name"] + .as_str() + .ok_or_else(|| invalid_data("Missing span event name"))? + .into(), + attributes: object_attributes(&event["attributes"])?, + ..Default::default() + }) + }) + .collect() +} + +fn links(value: Option<&Value>) -> Result, ApiError> { + let Some(value) = value.filter(|value| !value.is_null()) else { + return Ok(vec![]); + }; + let values = value + .as_array() + .ok_or_else(|| invalid_data("Invalid stored span links"))?; + values + .iter() + .map(|link| { + Ok(Link { + trace_id: decode_id(link["trace_id"].as_str().unwrap_or_default(), 16)?, + span_id: decode_id(link["span_id"].as_str().unwrap_or_default(), 8)?, + trace_state: link["trace_state"].as_str().unwrap_or_default().into(), + attributes: object_attributes(&link["attributes"])?, + ..Default::default() + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn test_event_timestamp_precision() { + for (time, expected) in [ + ("1970-01-01 00:00:01+0000", 1_000_000_000), + ("1970-01-01 00:00:01.123+0000", 1_123_000_000), + ("1970-01-01 00:00:01.123456+0000", 1_123_456_000), + ("1970-01-01 00:00:01.123456789+0000", 1_123_456_789), + ] { + let value = json!([{"name":"event", "time":time, "attributes":{}}]); + let events = events(Some(&value)).unwrap(); + assert_eq!(events[0].time_unix_nano, expected); + } + } +} diff --git a/src/servers/src/http/jaeger/v3/summary.rs b/src/servers/src/http/jaeger/v3/summary.rs new file mode 100644 index 00000000000..14273c3a705 --- /dev/null +++ b/src/servers/src/http/jaeger/v3/summary.rs @@ -0,0 +1,142 @@ +// Copyright 2023 Greptime Team +// +// 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 std::collections::{BTreeMap, HashSet}; + +use opentelemetry_proto::tonic::common::v1::any_value; +use opentelemetry_proto::tonic::trace::v1::status::StatusCode; +use opentelemetry_proto::tonic::trace::v1::{Span, TracesData}; +use serde_json::{Value, json}; + +use crate::otlp::trace::KEY_SERVICE_NAME; + +pub fn summarize(traces: TracesData) -> Vec { + let mut by_trace: BTreeMap<&[u8], Vec<(&str, &Span)>> = BTreeMap::new(); + for resource in &traces.resource_spans { + let service = resource + .resource + .as_ref() + .and_then(|resource| { + resource + .attributes + .iter() + .find(|attr| attr.key == KEY_SERVICE_NAME) + .and_then(|attr| attr.value.as_ref()) + .and_then(|value| match &value.value { + Some(any_value::Value::StringValue(name)) => Some(name.as_str()), + _ => None, + }) + }) + .unwrap_or_default(); + for scope in &resource.scope_spans { + for span in &scope.spans { + by_trace + .entry(&span.trace_id) + .or_default() + .push((service, span)); + } + } + } + by_trace + .into_iter() + .map(|(trace_id, spans)| { + let ids: HashSet<_> = spans + .iter() + .map(|(_, span)| span.span_id.as_slice()) + .collect(); + let mut services: BTreeMap<&str, (usize, usize)> = BTreeMap::new(); + let mut min_start = u64::MAX; + let mut max_end = 0; + let mut errors = 0; + let mut orphans = 0; + let mut root: Option<(&str, &Span)> = None; + for &(service, span) in &spans { + let stats = services.entry(service).or_default(); + stats.0 += 1; + if span + .status + .as_ref() + .is_some_and(|status| status.code == StatusCode::Error as i32) + { + stats.1 += 1; + errors += 1; + } + min_start = min_start.min(span.start_time_unix_nano); + max_end = max_end.max(span.end_time_unix_nano); + if span.parent_span_id.iter().all(|byte| *byte == 0) { + if root.is_none_or(|(_, root)| { + span.start_time_unix_nano < root.start_time_unix_nano + }) { + root = Some((service, span)); + } + } else if !ids.contains(span.parent_span_id.as_slice()) { + orphans += 1; + } + } + let (root_service, root_operation) = + root.map_or(("", ""), |(service, span)| (service, span.name.as_str())); + let services: Vec<_> = services + .into_iter() + .map(|(name, (spans, errors))| { + json!({ + "name": name, "spanCount": spans, "errorSpanCount": errors, + }) + }) + .collect(); + json!({ + "traceId": hex::encode(trace_id), + "rootServiceName": root_service, + "rootOperationName": root_operation, + "minStartTimeUnixNano": min_start.to_string(), + "maxEndTimeUnixNano": max_end.to_string(), + "spanCount": spans.len(), + "errorSpanCount": errors, + "orphanSpanCount": orphans, + "services": services, + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_summary_roots_and_orphans() { + let traces: TracesData = serde_json::from_value(json!({"resourceSpans":[ + {"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"b"}}]},"scopeSpans":[{"spans":[ + {"traceId":"00000000000000000000000000000001","spanId":"0000000000000001","name":"late-root","startTimeUnixNano":"20","endTimeUnixNano":"50"}, + {"traceId":"00000000000000000000000000000001","spanId":"0000000000000002","parentSpanId":"0000000000000099","name":"orphan","startTimeUnixNano":"1","endTimeUnixNano":"60","status":{"code":2}}, + {"traceId":"00000000000000000000000000000002","spanId":"0000000000000003","parentSpanId":"0000000000000001","startTimeUnixNano":"1","endTimeUnixNano":"2"} + ]}]}, + {"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"a"}}]},"scopeSpans":[{"spans":[ + {"traceId":"00000000000000000000000000000001","spanId":"0000000000000004","parentSpanId":"0000000000000000","name":"early-root","startTimeUnixNano":"10","endTimeUnixNano":"20"}, + {"traceId":"00000000000000000000000000000001","spanId":"0000000000000005","parentSpanId":"0000000000000004","startTimeUnixNano":"11","endTimeUnixNano":"15","status":{"code":2}} + ]}]} + ]})).unwrap(); + let summaries = summarize(traces); + assert_eq!( + summaries[0], + json!({ + "traceId":"00000000000000000000000000000001", "rootServiceName":"a", "rootOperationName":"early-root", + "minStartTimeUnixNano":"1", "maxEndTimeUnixNano":"60", "spanCount":4, "errorSpanCount":2, "orphanSpanCount":1, + "services":[{"name":"a","spanCount":2,"errorSpanCount":1},{"name":"b","spanCount":2,"errorSpanCount":1}] + }) + ); + assert_eq!(summaries[1]["orphanSpanCount"], 1); + assert_eq!(summaries[1]["rootServiceName"], ""); + assert_eq!(summaries[1]["rootOperationName"], ""); + } +} diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 91979ddbf59..dd2fc51a9b9 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -227,6 +227,7 @@ macro_rules! http_tests { test_splunk_raw, test_log_query, test_jaeger_query_api, + test_jaeger_v3_query_api, test_jaeger_query_api_for_trace_v1, test_jaeger_query_api_for_trace_v2, @@ -11490,6 +11491,324 @@ pub async fn test_log_query(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_jaeger_v3_query_api( + store_type: StorageType, +) -> Result<(), Box> { + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_jaeger_v3_query_api").await; + let client = TestClient::new(app).await; + for (path, expected_status, expected) in [ + ("services", StatusCode::OK, json!({"services":[]})), + ( + "trace-summaries?query.startTimeMin=2026-01-01T00:00:00Z&query.startTimeMax=2026-01-02T00:00:00Z", + StatusCode::OK, + json!({"summaries":[]}), + ), + ( + "operations?service=checkout", + StatusCode::OK, + json!({"operations":[]}), + ), + ( + "traces/00000000000000000000000000000001", + StatusCode::NOT_FOUND, + json!({"error":{"httpCode":404,"message":"No traces found"}}), + ), + ] { + let response = client + .get(&format!("/v1/jaeger/api/v3/{path}")) + .send() + .await; + assert_eq!(response.status(), expected_status, "{path}"); + assert_eq!( + serde_json::from_str::(&response.text().await)?, + expected + ); + } + for path in [ + "operations", + "traces", + "traces/not-hex", + "traces/0000000000000001?startTime=invalid", + "traces/0000000000000001?rawTraces=invalid", + "traces?query.startTimeMin=2026-01-01T00:00:00Z&query.startTimeMax=2026-01-02T00:00:00Z&query.searchDepth=10001", + "trace-summaries?query.startTimeMin=2026-01-01T00:00:00Z&query.startTimeMax=2026-01-02T00:00:00Z&query.searchDepth=10001", + ] { + let response = client + .get(&format!("/v1/jaeger/api/v3/{path}")) + .send() + .await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{path}"); + let body: Value = serde_json::from_str(&response.text().await)?; + assert_eq!(body["error"]["httpCode"], 400); + } + + // The parent starts before the search window. A matching child must return the whole trace. + let start = chrono::Utc::now() + .timestamp_nanos_opt() + .ok_or("Timestamp out of range")? as u64; + let start = start / 1_000_000_000 * 1_000_000_000 + 123; + let timestamp = |ns: u64| { + chrono::DateTime::from_timestamp_nanos(ns as i64) + .to_rfc3339_opts(chrono::SecondsFormat::Nanos, true) + }; + let request: ExportTraceServiceRequest = serde_json::from_value(json!({"resourceSpans":[ + {"resource":{"attributes":[make_string_attr("service.name","checkout"), make_string_attr("service.instance.id","one")]}, + "scopeSpans":[{"scope":{"name":"sdk","version":"1.0","attributes":[make_string_attr("scope.key","value"), {"key":"scope.bytes","value":{"bytesValue":"AQID/w=="}}]},"spans":[{ + "traceId":"000000000000000000000000000000ab", "spanId":"0000000000000001", "name":"parent", "kind":2, + "startTimeUnixNano":start.to_string(), "endTimeUnixNano":(start+2000).to_string() + }]}]}, + {"resource":{"attributes":[make_string_attr("service.name","checkout"), make_string_attr("service.instance.id","two"), {"key":"resource.bytes","value":{"bytesValue":"AQID/w=="}}]}, + "scopeSpans":[{"scope":{"name":"sdk","version":"1.0","attributes":[make_string_attr("scope.key","value"), {"key":"scope.bytes","value":{"bytesValue":"AQID/w=="}}]},"spans":[{ + "traceId":"000000000000000000000000000000ab", "spanId":"0000000000000002", "parentSpanId":"0000000000000001", "name":"child", "kind":3, + "traceState":"vendor=value", "startTimeUnixNano":(start+1000).to_string(), "endTimeUnixNano":(start+1999).to_string(), + "attributes":[make_int_attr("http.status_code",500),make_bool_attr("failed",true), + {"key":"bytes","value":{"bytesValue":"AQID/w=="}}, + {"key":"integers","value":{"arrayValue":{"values":[{"intValue":"1"},{"intValue":"2"},{"intValue":"3"},{"intValue":"255"}]}}}, + {"key":"nested","value":{"arrayValue":{"values":[{"intValue":"9223372036854775807"},{"kvlistValue":{"values":[make_string_attr("key","value")]}}]}}}], + "events":[{"name":"exception","timeUnixNano":(start+1011).to_string(),"attributes":[make_string_attr("message","failure")]}], + "links":[{"traceId":"000000000000000000000000000000cd","spanId":"0000000000000003","traceState":"link=value","attributes":[make_bool_attr("remote",true)]}], + "status":{"code":2,"message":"failure"} + }, { + "traceId":"000000000000000000000000000000cd", "spanId":"0000000000000003", "name":"boundary", "kind":2, + "startTimeUnixNano":(start+2000).to_string(), "endTimeUnixNano":(start+3000).to_string() + }]}]} + ]}))?; + for pipeline in [ + "greptime_trace_v0", + GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME, + GREPTIME_INTERNAL_TRACE_PIPELINE_V2_NAME, + ] { + let table = format!("jaeger_api_v3_{pipeline}"); + let response = send_req( + &client, + vec![ + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ), + ( + HeaderName::from_static("x-greptime-pipeline-name"), + HeaderValue::from_str(pipeline)?, + ), + ( + HeaderName::from_static("x-greptime-trace-table-name"), + HeaderValue::from_str(&table)?, + ), + ], + "/v1/otlp/v1/traces", + request.encode_to_vec(), + false, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + for (path, expected) in [ + ("services", json!({"services":["checkout"]})), + ( + "operations?service=checkout&span_kind=client", + json!({"operations":[{"name":"child","spanKind":"client"}]}), + ), + ( + "operations?service=checkout&spanKind=client", + json!({"operations":[{"name":"child","spanKind":"client"}]}), + ), + ] { + let response = client + .get(&format!("/v1/jaeger/api/v3/{path}")) + .header("x-greptime-trace-table-name", &table) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + serde_json::from_str::(&response.text().await)?, + expected + ); + } + for flushed in [false, true] { + if flushed { + let response = client + .post("/v1/sql") + .form(&[("sql", format!("ADMIN FLUSH_TABLE('{table}')"))]) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + } + let response = client.get(&format!( + "/v1/jaeger/api/v3/trace-summaries?query.serviceName=checkout&query.startTimeMin={}&query.startTimeMax={}", + timestamp(start+1000), timestamp(start+2000), + )).header("x-greptime-trace-table-name", &table).send().await; + assert_eq!(response.status(), StatusCode::OK); + let body: Value = serde_json::from_str(&response.text().await)?; + assert_eq!( + body, + json!({"summaries":[{ + "traceId":"000000000000000000000000000000ab", + "rootServiceName":"checkout", "rootOperationName":"parent", + "minStartTimeUnixNano":start.to_string(), "maxEndTimeUnixNano":(start+2000).to_string(), + "spanCount":2, "errorSpanCount":1, "orphanSpanCount":0, + "services":[{"name":"checkout", "spanCount":2, "errorSpanCount":1}], + }]}) + ); + for path in [ + "traces/00000000000000AB".to_string(), + format!( + "traces/00000000000000AB?startTime={}", + timestamp(start + 1000) + ), + format!("traces/00000000000000AB?endTime={}", timestamp(start + 500)), + format!( + "traces/00000000000000AB?start_time={}&end_time={}", + timestamp(start + 500), + timestamp(start + 999) + ), + format!( + "traces?query.startTimeMin={}&query.startTimeMax={}&query.searchDepth=0", + timestamp(start + 1000), + timestamp(start + 2000) + ), + format!( + "traces?query.service_name=checkout&query.operation_name=child&query.start_time_min={}&query.start_time_max={}&query.num_traces=1&query.duration_min=999ns&query.duration_max=999ns&query.attributes={}", + timestamp(start + 1000), + timestamp(start + 2000), + encode(r#"{"http.status_code":"500","failed":"true"}"#) + ), + format!( + "traces?query.startTimeMin={}&query.startTimeMax={}&query.attributes={}", + timestamp(start + 1000), + timestamp(start + 2000), + encode(r#"{"service.instance.id":"two"}"#) + ), + ] { + let response = client + .get(&format!("/v1/jaeger/api/v3/{path}")) + .header("x-greptime-trace-table-name", &table) + .send() + .await; + let status = response.status(); + let body: Value = serde_json::from_str(&response.text().await)?; + assert_eq!(status, StatusCode::OK, "{pipeline}: {path}: {body}"); + let resources = body["result"]["resourceSpans"] + .as_array() + .ok_or("Missing resource spans")?; + let spans: Vec<_> = resources + .iter() + .flat_map(|r| r["scopeSpans"][0]["spans"].as_array().unwrap()) + .collect(); + assert_eq!(spans.len(), 2, "{body}"); + assert!( + spans + .iter() + .all(|span| span["traceId"] == "000000000000000000000000000000ab") + ); + let child = spans + .iter() + .find(|span| span["name"] == "child") + .ok_or("Missing child span")?; + assert_eq!(child["startTimeUnixNano"], (start + 1000).to_string()); + assert_eq!(child["endTimeUnixNano"], (start + 1999).to_string()); + assert_eq!(child["parentSpanId"], "0000000000000001"); + assert_eq!(child["kind"], 3); + assert_eq!(child["status"]["code"], 2); + assert_eq!(child["traceState"], "vendor=value"); + assert_eq!( + child["events"][0]["timeUnixNano"], + (start + 1011).to_string() + ); + assert_eq!(child["links"][0]["traceState"], "link=value"); + assert_eq!( + child["links"][0]["attributes"][0]["value"]["boolValue"], + true + ); + let nested = child["attributes"] + .as_array() + .unwrap() + .iter() + .find(|attr| attr["key"] == "nested") + .unwrap(); + assert_eq!( + nested["value"]["arrayValue"]["values"][0]["intValue"], + "9223372036854775807" + ); + let attribute = |values: &Value, key: &str| { + values + .as_array() + .unwrap() + .iter() + .find(|attr| attr["key"] == key) + .unwrap()["value"] + .clone() + }; + assert_eq!( + attribute(&child["attributes"], "integers"), + json!({"arrayValue":{"values":[{"intValue":"1"},{"intValue":"2"},{"intValue":"3"},{"intValue":"255"}]}}) + ); + if pipeline == GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME { + assert_eq!( + attribute(&child["attributes"], "bytes"), + json!({"bytesValue":"AQID/w=="}) + ); + } + for resource in resources { + let span = &resource["scopeSpans"][0]["spans"][0]; + let attrs = resource["resource"]["attributes"].as_array().unwrap(); + let instance = attrs + .iter() + .find(|attr| attr["key"] == "service.instance.id") + .unwrap(); + assert_eq!( + instance["value"]["stringValue"], + if span["name"] == "parent" { + "one" + } else { + "two" + } + ); + let scope_attrs = &resource["scopeSpans"][0]["scope"]["attributes"]; + assert_eq!( + attribute(scope_attrs, "scope.key"), + json!({"stringValue":"value"}) + ); + if pipeline == GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME { + assert_eq!( + attribute(scope_attrs, "scope.bytes"), + json!({"bytesValue":"AQID/w=="}) + ); + if span["name"] == "child" { + assert_eq!( + attribute(&resource["resource"]["attributes"], "resource.bytes"), + json!({"bytesValue":"AQID/w=="}) + ); + } + } + } + } + // Instance `one` is on the parent, which starts before the search window. + let response = client + .get(&format!( + "/v1/jaeger/api/v3/traces?query.startTimeMin={}&query.startTimeMax={}&query.attributes={}", + timestamp(start + 1000), + timestamp(start + 2000), + encode(r#"{"service.instance.id":"one"}"#) + )) + .header("x-greptime-trace-table-name", &table) + .send() + .await; + assert_eq!(response.status(), StatusCode::NOT_FOUND, "{pipeline}"); + } + // Legacy clients continue receiving Jaeger JSON at the old URL. + let response = client + .get("/v1/jaeger/api/traces/000000000000000000000000000000ab") + .header("x-greptime-trace-table-name", &table) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + let body: Value = serde_json::from_str(&response.text().await)?; + assert_eq!(body["data"][0]["spans"].as_array().unwrap().len(), 2); + } + guard.remove_all().await; + Ok(()) +} + pub async fn test_jaeger_query_api(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = From 483d1e49a3b984fbe7b6e299ee9e3189bb75c4e9 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 12 Oct 2026 06:12:56 +0000 Subject: [PATCH 09/46] fix(query): keep filter fetch when rewriting matches term (#9521) * test: cover matches term limit dropped by constant rewrite Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): keep filter fetch when rewriting matches term Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/optimizer/constant_term.rs | 12 +- tests-integration/src/tests/instance_test.rs | 119 +++++++++++++++++++ 2 files changed, 127 insertions(+), 4 deletions(-) diff --git a/src/query/src/optimizer/constant_term.rs b/src/query/src/optimizer/constant_term.rs index 70ddcc2f647..6b77a8aa7f6 100644 --- a/src/query/src/optimizer/constant_term.rs +++ b/src/query/src/optimizer/constant_term.rs @@ -201,10 +201,14 @@ impl PhysicalOptimizerRule for MatchesConstantTermOptimizer { })?; if new_pred.transformed { - let exec = FilterExecBuilder::new(new_pred.data, filter.input().clone()) - .with_default_selectivity(filter.default_selectivity()) - .apply_projection_by_ref(filter.projection().as_ref()) - .and_then(|x| x.build())?; + // `FilterExecBuilder::from(&FilterExec)` already carries over the + // input, projection, default selectivity, batch size and fetch, so + // only the predicate has to be replaced. Re-applying the projection + // is not only pointless but wrong: `apply_projection` composes with + // the projection copied from the builder. + let exec = FilterExecBuilder::from(filter) + .with_predicate(new_pred.data) + .build()?; return Ok(Transformed::yes(Arc::new(exec) as _)); } } diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 89706b53764..1a9cc1790f7 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -415,6 +415,125 @@ LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' AND k < '8' } } +/// Regression test for `MatchesConstantTermOptimizer` dropping `FilterExec`'s `fetch`. +/// +/// On a single-partition plan the `LIMIT` is enforced by the `fetch` that the physical +/// `LimitPushdown` rule leaves on the `FilterExec`, so a rule rebuilding that `FilterExec` +/// has to carry the `fetch` over. While it does not, the limit is silently lifted for +/// `matches_term` filters and +/// `SELECT count(*) FROM (SELECT ts, msg FROM t WHERE matches_term(msg, 'foo') LIMIT 1) s` +/// returns the number of matching rows (2) instead of 1. +/// +/// `query_parallelism=4` is expected to stay correct, because the limit is then materialized +/// above the scan instead of being delegated to the `FilterExec`; the test pins down that +/// contrast. The metric-engine logical table covers the same filter on a shared physical +/// region. +/// +/// Both matching rows share `ts` and `msg`, so the limited output does not depend on the +/// order in which the scan returns them. +#[apply(both_instances_cases)] +async fn test_matches_term_subquery_limit_count(instance: Arc) { + let frontend = instance.frontend(); + + execute_sql( + &frontend, + r#" +CREATE TABLE f07_mito_msgs ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + msg STRING, + val DOUBLE +) +ENGINE = mito +"#, + ) + .await; + execute_sql( + &frontend, + r#"INSERT INTO f07_mito_msgs (ts, host, msg, val) VALUES + (0, 'host1', 'foo bar', 1.0), + (0, 'host2', 'foo bar', 2.0), + (0, 'host3', 'other baz', 3.0)"#, + ) + .await; + + execute_sql( + &frontend, + r#" +CREATE TABLE f07_metric_phy ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + msg STRING, + PRIMARY KEY (host, msg) +) +ENGINE = metric +WITH ('physical_metric_table' = '') +"#, + ) + .await; + execute_sql( + &frontend, + r#" +CREATE TABLE f07_metric_msgs ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + msg STRING, + PRIMARY KEY (host, msg) +) +ENGINE = metric +WITH ('on_physical_table' = 'f07_metric_phy') +"#, + ) + .await; + execute_sql( + &frontend, + r#"INSERT INTO f07_metric_msgs (ts, host, msg, val) VALUES + (0, 'host1', 'foo bar', 1.0), + (0, 'host2', 'foo bar', 2.0), + (0, 'host3', 'other baz', 3.0)"#, + ) + .await; + + for table in ["f07_mito_msgs", "f07_metric_msgs"] { + let limited = format!("SELECT ts, msg FROM {table} WHERE matches_term(msg, 'foo') LIMIT 1"); + let count = format!("SELECT count(*) FROM ({limited}) s"); + + let limit_rows = execute_sql_with_query_parallelism(&frontend, &limited, 1) + .await + .data + .pretty_print() + .await; + assert_eq!( + limit_rows, + r#"+---------------------+---------+ +| ts | msg | ++---------------------+---------+ +| 1970-01-01T00:00:00 | foo bar | ++---------------------+---------+"#, + "query_parallelism=1: `{limited}` must return exactly the single limited row" + ); + + for parallelism in [1, 4] { + let count_rows = execute_sql_with_query_parallelism(&frontend, &count, parallelism) + .await + .data + .pretty_print() + .await; + assert_eq!( + count_rows, + r#"+----------+ +| count(*) | ++----------+ +| 1 | ++----------+"#, + "query_parallelism={parallelism}: the LIMIT 1 of the subquery must cap `{count}` to one row" + ); + } + } +} + #[apply(both_instances_cases)] async fn test_partitioned_merge_scan_semi_join_repartitions(instance: Arc) { let frontend = instance.frontend(); From 18726c760257047af4029692ce4ce147fdc3d35c Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 12 Oct 2026 06:17:03 +0000 Subject: [PATCH 10/46] fix: bound metric snapshot memory and reduce export overhead (#9505) * fix(mito): bound retained scan memory Compact unused tag dictionary values before merge retention and avoid splitting a single scan source into redundant batches. Include regression coverage. Signed-off-by: jeremyhi * perf(snapshot): bound writes and reduce logical export overhead Bound metadata writes to 8 MiB and abort failed writers. Preserve compatible string dictionaries, bound conversion work, combine Parquet row groups and filter single-table exports before routing. Include resource-boundary and packed-restore regression coverage. Signed-off-by: jeremyhi * docs: clarify metric snapshot acceptance criteria Specify populated scale fixtures, controlled S3 request delay, native macOS validation, paging exclusions and alternating comparison rounds. Signed-off-by: jeremyhi * fix(mito): split batches only for merged readers Signed-off-by: jeremyhi * test(snapshot): cover supported types in both layouts Signed-off-by: jeremyhi * docs: consolidate metric snapshot experiment conclusions Signed-off-by: jeremyhi * fix(mito): honor series statistics when splitting L0 scans Apply the existing series-density heuristic to L0 files with known statistics and include them in the batch-size estimate. Preserve the fallback for older files without statistics. Cover short-series overlapping L0 files with an ordered latest-value regression. Signed-off-by: jeremyhi * docs(mito): explain memory risk of splitting short L0 series Signed-off-by: jeremyhi * refactor: address snapshot export review feedback Signed-off-by: jeremyhi * fix(cli): accept source errors in lazy snapshot contexts Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 1 + Cargo.toml | 1 + docs/rfcs/2026-09-11-metric-export-import.md | 65 +++- .../experiment-results.md | 34 -- src/cli/src/data/snapshot_storage.rs | 121 +++++- src/common/datasource/src/packed_writer.rs | 48 +-- src/common/datasource/src/parquet_writer.rs | 25 +- src/mito2/Cargo.toml | 1 + src/mito2/src/engine/scan_test.rs | 137 +++++++ src/mito2/src/read/flat_projection.rs | 69 +++- src/mito2/src/read/scan_util.rs | 34 +- src/mito2/src/read/seq_scan.rs | 12 +- src/object-store/src/layers/mock.rs | 46 ++- .../src/statement/export_logical_tables.rs | 145 +++++-- .../statement/export_logical_tables/tests.rs | 359 ++++++++++++++++-- .../tests/export_logical_tables.rs | 30 +- 16 files changed, 952 insertions(+), 176 deletions(-) delete mode 100644 docs/rfcs/2026-09-11-metric-export-import/experiment-results.md diff --git a/Cargo.lock b/Cargo.lock index 095d6d4acec..61a9c246ed1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8702,6 +8702,7 @@ dependencies = [ "api", "aquamarine", "arrow-schema 59.2.0", + "arrow-select 59.2.0", "async-channel 1.9.0", "async-stream", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 549412f369a..f47907aa797 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -107,6 +107,7 @@ arrow-buffer = "=59.2.0" arrow-flight = "=59.2.0" arrow-ipc = { version = "=59.2.0", default-features = false, features = ["lz4", "zstd"] } arrow-schema = { version = "=59.2.0", features = ["serde"] } +arrow-select = "=59.2.0" async-stream = "0.3" async-trait = "0.1" # Remember to update axum-extra, axum-macros when updating axum diff --git a/docs/rfcs/2026-09-11-metric-export-import.md b/docs/rfcs/2026-09-11-metric-export-import.md index 8ddc31f40b0..57212ea53cb 100644 --- a/docs/rfcs/2026-09-11-metric-export-import.md +++ b/docs/rfcs/2026-09-11-metric-export-import.md @@ -31,13 +31,21 @@ includes reports of exports taking days for roughly 100,000 logical tables with only several GiB of data. Table count, physical schema width and data volume must be measured separately; total database bytes do not describe the transfer cost. -The [experiment conclusions](2026-09-11-metric-export-import/experiment-results.md) -support shared scans, batch DDL, explicit-file COPY and packed-object transfer. -Packing with a bounded range reader substantially reduces storage requests and -improves full export and restore for many small logical tables. Benefits depend -on request latency, and memory use increases. Remaining CREATE work makes batch -DDL part of the first release. These local experiments establish direction; -production V2 recovery and release performance still require acceptance tests. +Local standalone experiments with development builds, local files and Docker +MinIO support shared scans, batch DDL, explicit-file COPY and packed-object +transfer. Packing with a bounded range reader reduces storage requests and +improves full export and restore for many small logical tables; low-latency +export gains are smaller. CPU and request overhead fell while peak process +memory increased. Remaining CREATE work makes batch DDL part of the first +release. Explicit-file COPY removes repeated directory traversal while retaining +ordinary per-table insertion, independently of merged database writes. + +The exercised fixtures preserved schemas, DDL and typed values. Standalone-file +experiments used the existing importer; packing used an independent driver and +prototype range reader. These experiments establish design direction, not +production packed-format compatibility, V2 recovery safety or release performance. +The 8 MiB write/part limit was a design choice, not a measured optimum. Current +acceptance status is maintained in the [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120). # Terminology @@ -552,18 +560,39 @@ batching, still writing standalone table objects; (C) B plus packing and its reader, with identical DDL handling; (D) C plus batch DDL. If general restore request deduplication is available, include it in B as well. -The proposed release target is at least 3x faster full export and full restore -for D versus A on 10,000 and 100,000 small logical tables with MinIO request -latency increased by 5 ms. Fix this target during RFC review, before running -release measurements. The independent local C-versus-B experiment establishes -neither this comparison nor proof of that target. +The release target is at least 3x faster full export and full restore for D +versus A on both 10,000 and 100,000 small logical tables at both +50 ms and ++100 ms added MinIO request latency. Populate every selected table with the same +deterministic data across controls, initially 32 rows per table. Include schema +export, metadata discovery and CLI startup/shutdown in complete wall duration. +The independent local C-versus-B experiment establishes neither this comparison +nor proof of that target. Keep production improvements already shared by all +controls in A, and identify which increments can actually be isolated. -Use Linux release builds and at least three alternating rounds for primary -comparisons. Report medians/spread, schema/data/full durations, CPU, RSS, bytes, -objects, request counts and queue/cache peaks. Include zero-added-latency, -local-file, mixed small/large-table and wide-schema controls. Investigate and -remeasure full-duration regressions exceeding 5% versus A; do not generalize the -latency-injected speedup to these cases. Check values and schemas outside timing. +Start with native optimized release builds on macOS. Record hardware, OS, +MinIO placement, container/VM CPU and memory limits, and connection/storage +settings. Linux deployment-side performance confirmation can follow separately; +Linux root/non-root and Windows correctness checks remain distinct requirements. +Earlier development-build diagnostics, runs with paging interference and sparse +fixtures with only one populated table do not establish this acceptance. + +Calibrate one added delay per object-store HTTP request, including HEAD, GET +(and range GET), PUT, LIST, DELETE and each multipart operation. Do not delay +individual packets, streaming body chunks or database SQL/DDL. Route CLI and +server object access through the same controlled endpoint and report measured +per-method latency distributions. Synthetic delay is not a cloud provider profile. + +Use a small calibration pilot before each scale and at least three alternating +A/B/C/D rounds for primary comparisons, with fresh restore namespaces, equal +concurrency and comparable cache/compaction conditions. Report medians/spread, +schema/data/full durations, CPU, RSS, swap activity, bytes, objects, requests and +queue/cache peaks. Check RAM, disk and MinIO health before large runs; do not +compile concurrently. Fix interfering paging/load or report the scale unverified. +Retain local-file and MinIO 0/+5 ms controls, mixed small/large tables, wide +physical column unions and multiple physical groups. Add +20 ms only if needed +to explain sensitivity. Investigate and remeasure full-duration regressions +exceeding 5% versus A. Check table membership, schemas and typed values outside +timing, and report B/A, C/B and D/C independently where comparable. Validate actual write/part sizes at or below 8 MiB, aggregate concurrency across groups/chunks, and resource release after success, cancellation and failure. diff --git a/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md b/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md deleted file mode 100644 index 9cc7b46464d..00000000000 --- a/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md +++ /dev/null @@ -1,34 +0,0 @@ -# Metric export/import experiment conclusions - -The experiments support shared physical scans, explicit-file COPY, batched -logical DDL, and packing small logical-table Parquet streams into shared objects. - -- **Shared scans reduce repeated query work.** Projecting each logical schema - before dictionary expansion preserves its columns and types, but does not - remove the memory cost of a wide physical schema or distributed ordering. -- **Explicit-file COPY removes repeated directory traversal.** This improvement - uses ordinary per-table insertion; it is independent of database merged writes. -- **Packing and its range reader reduce small-object overhead.** They improve - complete export and restore and sharply reduce object-store requests. Benefits - grow with table count and request latency; low-latency export gains are smaller. - Some reader request savings can also benefit standalone files. -- **Batch DDL remains necessary.** Once data restore is faster, logical-table - creation becomes the main remaining restore cost. The batch-DDL experiment used - an internal interface; production transport and retry integration remain work. -- **Bounded writes allow larger objects.** The packing experiment kept writes - and upload parts within 8 MiB without imposing that limit on whole objects or - splitting logical tables. This was the chosen size, not a measured optimum. -- **Packing has a memory trade-off.** CPU and request overhead fell, while peak - process memory increased. Encoding buffers, upload concurrency, range windows - and table/index metadata need separate accounting. - -The exercised fixtures preserved schemas, DDL and typed values after restore. -Earlier standalone-file experiments used the existing importer; the packing -experiment used an independent driver and prototype range reader. The latter -does not establish production packed-format compatibility or V2 resume safety. - -These were local standalone experiments using development builds, local files -and Docker MinIO. Release acceptance still needs the actual packed V2 path, -request-owned reader isolation, old-reader rejection, large-table throughput, -distributed ordering, cancellation and recovery. See the [RFC](../2026-09-11-metric-export-import.md) -and [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120). diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index 17008706066..f8dd22f0307 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -21,6 +21,7 @@ use std::collections::BTreeSet; use std::path::Component; use async_trait::async_trait; +use common_datasource::packed_writer::WRITE_BYTES; use futures::TryStreamExt; use object_store::services::{Azblob, Fs, Gcs, Oss, S3}; use object_store::util::{with_instrument_layers, with_retry_layers}; @@ -570,13 +571,31 @@ impl OpenDalStorage { /// Writes bytes to a file. async fn write_file(&self, path: &str, data: Vec) -> Result<()> { - self.object_store - .write(path, data) + let mut writer = self + .object_store + .writer_with(path) + .chunk(WRITE_BYTES) + .concurrent(1) .await - .map(|_| ()) - .context(StorageOperationSnafu { - operation: format!("write {}", path), - }) + .with_context(|_| StorageOperationSnafu { + operation: format!("open writer {}", path), + })?; + let data = object_store::Buffer::from(data); + let result = async { + for chunk in data.chunks(WRITE_BYTES) { + writer.write(chunk).await?; + } + writer.close().await.map(|_| ()) + } + .await; + if result.is_err() + && let Err(error) = writer.abort().await + { + common_telemetry::warn!(error; "Failed to abort snapshot metadata write"); + } + result.with_context(|_| StorageOperationSnafu { + operation: format!("write {}", path), + }) } /// Checks if a file exists using stat. @@ -837,6 +856,96 @@ mod tests { ) } + #[tokio::test] + async fn snapshot_metadata_submissions_are_bounded_and_invoke_abort_on_failure() { + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use object_store::layers::mock::{Metadata, ObservedWriter, WriterObservations, oio}; + + struct FailWriter { + inner: oio::Writer, + writes: usize, + fail: Arc, + } + impl oio::Write for FailWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.writes += 1; + if self.fail.load(Ordering::SeqCst) == 1 && self.writes == 2 { + return Err(object_store::Error::new( + ErrorKind::Unexpected, + "injected metadata write failure", + )); + } + self.inner.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + if self.fail.load(Ordering::SeqCst) == 2 { + return Err(object_store::Error::new( + ErrorKind::Unexpected, + "injected metadata close failure", + )); + } + self.inner.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.abort().await + } + } + + let dir = tempdir().unwrap(); + let mut storage = make_storage_with_rooted_fs(dir.path()); + let observations = Arc::new(Mutex::new(WriterObservations::default())); + let fail = Arc::new(AtomicUsize::new(0)); + storage.object_store = storage.object_store.layer( + MockLayerBuilder::default() + .writer_factory(Arc::new({ + let observations = observations.clone(); + let fail = fail.clone(); + move |_, _, inner| { + Box::new(ObservedWriter::new( + Box::new(FailWriter { + inner, + writes: 0, + fail: fail.clone(), + }), + observations.clone(), + )) + } + })) + .build() + .unwrap(), + ); + for length in [2 * WRITE_BYTES + 17, 1, 0] { + *observations.lock().unwrap() = WriterObservations::default(); + let data = vec![42; length]; + storage.write_file("metadata", data.clone()).await.unwrap(); + assert_eq!(storage.read_file("metadata").await.unwrap(), data); + let observed = observations.lock().unwrap(); + assert_eq!(observed.write_sizes.iter().sum::(), length); + assert!(observed.write_sizes.iter().all(|size| *size <= WRITE_BYTES)); + assert_eq!(observed.close_calls, 1); + assert_eq!(observed.abort_calls, 0); + } + for failure in [1, 2] { + *observations.lock().unwrap() = WriterObservations::default(); + fail.store(failure, Ordering::SeqCst); + assert!( + storage + .write_file("failed", vec![42; 2 * WRITE_BYTES + 17]) + .await + .is_err() + ); + assert_eq!( + observations.lock().unwrap().write_sizes.len(), + if failure == 1 { 2 } else { 3 } + ); + let observed = observations.lock().unwrap(); + assert_eq!(observed.close_calls, usize::from(failure == 2)); + assert_eq!(observed.abort_calls, 1); + } + } + #[test] fn test_validate_uri_valid() { assert_eq!(validate_uri("s3://bucket/path").unwrap(), StorageScheme::S3); diff --git a/src/common/datasource/src/packed_writer.rs b/src/common/datasource/src/packed_writer.rs index 635a1a7a771..d0560fa06f7 100644 --- a/src/common/datasource/src/packed_writer.rs +++ b/src/common/datasource/src/packed_writer.rs @@ -410,7 +410,9 @@ mod tests { use arrow::array::{ArrayRef, Int64Array, StringArray}; use arrow::record_batch::RecordBatch; use futures::TryStreamExt; - use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + use object_store::layers::mock::{ + Metadata, MockLayerBuilder, MockWriterFactory, ObservedWriter, WriterObservations, oio, + }; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; @@ -424,10 +426,10 @@ mod tests { let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) .unwrap() .build_operator(); - let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let observations = Arc::new(std::sync::Mutex::new(WriterObservations::default())); let factory: MockWriterFactory = Arc::new({ - let sizes = sizes.clone(); - move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + let observations = observations.clone(); + move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone())) }); let store = store.layer( MockLayerBuilder::default() @@ -483,7 +485,14 @@ mod tests { .validate_membership(["table0", "table1", "table2"]) .unwrap(); assert_eq!(inventory.len(), 3); - assert!(sizes.lock().unwrap().iter().all(|n| *n <= WRITE_BYTES)); + assert!( + observations + .lock() + .unwrap() + .write_sizes + .iter() + .all(|n| *n <= WRITE_BYTES) + ); assert_eq!( index .objects @@ -540,27 +549,13 @@ mod tests { } } - struct Observe(oio::Writer, Arc>>); - impl oio::Write for Observe { - async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { - self.1.lock().unwrap().push(bytes.len()); - self.0.write(bytes).await - } - async fn close(&mut self) -> object_store::Result { - self.0.close().await - } - async fn abort(&mut self) -> object_store::Result<()> { - self.0.abort().await - } - } - #[tokio::test] async fn rolls_at_table_boundary_and_aborts_standalone_at_part_limit() { let directory = common_test_util::temp_dir::create_temp_dir("packed-limit"); - let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let observations = Arc::new(std::sync::Mutex::new(WriterObservations::default())); let factory: MockWriterFactory = Arc::new({ - let sizes = sizes.clone(); - move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + let observations = observations.clone(); + move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone())) }); let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) .unwrap() @@ -605,7 +600,14 @@ mod tests { .iter() .all(|t| t.offset == 0 && t.length == 600) ); - assert!(sizes.lock().unwrap().iter().all(|n| *n <= 512)); + assert!( + observations + .lock() + .unwrap() + .write_sizes + .iter() + .all(|n| *n <= 512) + ); } struct FailClose(oio::Writer); diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index 243a46b4e1a..157733e856e 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -500,28 +500,17 @@ mod tests { async fn large_footer_flush_uses_bounded_submissions_in_one_parquet_stream() { use std::sync::Mutex; - use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; - struct ObservedWriter(oio::Writer, Arc>>); - impl oio::Write for ObservedWriter { - async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { - self.1.lock().unwrap().push(bytes.len()); - self.0.write(bytes).await - } - async fn close(&mut self) -> object_store::Result { - self.0.close().await - } - async fn abort(&mut self) -> object_store::Result<()> { - self.0.abort().await - } - } + use object_store::layers::mock::{ + MockLayerBuilder, MockWriterFactory, ObservedWriter, WriterObservations, + }; let directory = common_test_util::temp_dir::create_temp_dir("bounded_parquet"); let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) .unwrap() .build_operator(); - let sizes = Arc::new(Mutex::new(Vec::new())); + let observations = Arc::new(Mutex::new(WriterObservations::default())); let factory: MockWriterFactory = Arc::new({ - let sizes = sizes.clone(); - move |_, _, writer| Box::new(ObservedWriter(writer, sizes.clone())) + let observations = observations.clone(); + move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone())) }); let store = store.layer( MockLayerBuilder::default() @@ -548,7 +537,7 @@ mod tests { writer.sink = ParquetSink::Object(store.writer("large.parquet").await.unwrap()); writer.write(batch.clone(), None).await.unwrap(); writer.finish(None).await.unwrap(); - let sizes = sizes.lock().unwrap().clone(); + let sizes = observations.lock().unwrap().write_sizes.clone(); let limit = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; assert!(sizes.iter().sum::() > limit); assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}"); diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 39dd96f687b..02c134e964c 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -19,6 +19,7 @@ workspace = true api.workspace = true aquamarine.workspace = true arrow-schema.workspace = true +arrow-select.workspace = true async-channel = "1.9" async-stream.workspace = true async-trait.workspace = true diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index f86b327735d..f55f084ebd7 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -4498,6 +4498,143 @@ async fn test_range_cache_key_separates_sequence_ranges() { ); } +#[tokio::test] +async fn test_single_source_ranges_preserve_batches() { + for (file_count, row_group_size) in [(1, 64), (1, 16), (2, 16)] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .insert_option("max_row_group_row_count", &row_group_size.to_string()) + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows_per_file = 64 / file_count; + for file in 0..file_count { + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: test_util::build_rows(file * rows_per_file, (file + 1) * rows_per_file), + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + let Scanner::Seq(scanner) = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + else { + panic!("expected a sequential scan"); + }; + assert_eq!(file_count, scanner.input().files.len()); + assert_eq!( + 64 / row_group_size, + scanner.properties().partitions[0].len() + ); + let stream = scanner + .scan_partition(&Default::default(), &ExecutionPlanMetricsSet::default(), 0) + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + assert_eq!( + vec![row_group_size; 64 / row_group_size], + batches + .iter() + .map(|batch| batch.num_rows()) + .collect::>() + ); + let mut values = scan_field_values(&batches); + values.sort_by(f64::total_cmp); + assert_eq!((0..64).map(|i| i as f64).collect::>(), values); + } +} + +#[tokio::test] +async fn test_short_series_level_zero_merge_preserves_batches_and_latest_values() { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .insert_option("max_row_group_row_count", "65536") + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + for (start, end, offset) in [(0, 12_000, 0), (6_000, 18_000, 100_000)] { + let rows = (start..end) + .map(|i| { + row(vec![ + ValueData::StringValue(format!("{i:05}")), + ValueData::F64Value((i + offset) as f64), + ValueData::TimestampMillisecondValue(0), + ]) + }) + .collect(); + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + let Scanner::Seq(scanner) = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + else { + panic!("expected a sequential scan"); + }; + assert_eq!(2, scanner.input().files.len()); + for file in &scanner.input().files { + assert_eq!(0, file.meta_ref().level); + assert_eq!(12_000, file.meta_ref().num_series); + } + let batches = scanner + .scan_partition(&Default::default(), &ExecutionPlanMetricsSet::default(), 0) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!( + batches.len() < 64, + "short series became {} batches", + batches.len() + ); + assert_eq!( + (0..18_000) + .map(|i| if i < 6_000 { + i as f64 + } else { + (i + 100_000) as f64 + }) + .collect::>(), + scan_field_values(&batches), + ); +} + #[tokio::test] async fn test_reset_state_discards_dynamic_field_pruning() { for (append_mode, expected_scanner_name) in [(false, "SeqScan"), (true, "UnorderedScan")] { diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index 2abf6aa48af..1734ca364af 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -22,7 +22,7 @@ use common_recordbatch::error::{ ArrowComputeSnafu, DataTypesSnafu, ExternalSnafu, NewDfRecordBatchSnafu, }; use common_recordbatch::{DfRecordBatch, RecordBatch}; -use datatypes::arrow::array::Array; +use datatypes::arrow::array::{Array, AsArray}; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; @@ -320,6 +320,16 @@ impl FlatProjectionMapper { let mut arrays = Vec::with_capacity(self.output_schema.num_columns()); for (output_idx, index) in self.batch_indices.iter().enumerate() { let mut array = batch.column(*index).clone(); + if matches!( + self.output_schema.arrow_schema().fields()[output_idx].data_type(), + ArrowDataType::Dictionary(_, _) + ) && let Some(dictionary) = array.as_any_dictionary_opt() + && dictionary.values().len() > dictionary.len() + { + // Small scan slices must not retain an entire SST's tag dictionary in merges. + array = arrow_select::dictionary::garbage_collect_any_dictionary(dictionary) + .context(ArrowComputeSnafu)?; + } // Cast dictionary values to the target type. if let ArrowDataType::Dictionary(_key_type, value_type) = array.data_type() && !matches!( @@ -667,6 +677,63 @@ mod tests { ); } + #[test] + fn test_tag_projection_compacts_dictionary_slices() { + use datatypes::arrow::array::{ + DictionaryArray, StringArray, TimestampMillisecondArray, UInt32Array, + }; + use datatypes::arrow::datatypes::UInt32Type; + + let mut builder = RegionMetadataBuilder::new(RegionId::new(1024, 0)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let mapper = FlatProjectionMapper::new(&metadata, [0, 1]) + .unwrap() + .with_pk_dictionary_encoding(); + let values = StringArray::from_iter( + (0..1024).map(|i| (i != 1).then(|| format!("{i}{}", "x".repeat(4096)))), + ); + let dictionary = DictionaryArray::::try_new( + UInt32Array::from(vec![Some(1023), None, Some(1), Some(1023)]), + Arc::new(values), + ) + .unwrap(); + let expected = datatypes::arrow::compute::cast(&dictionary, &ArrowDataType::Utf8).unwrap(); + let batch = DfRecordBatch::try_from_iter([ + ("tag", Arc::new(dictionary) as Arc), + ( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])), + ), + ]) + .unwrap(); + let output = mapper.convert(&batch, &CacheStrategy::Disabled).unwrap(); + let actual = output.column(0); + assert!(actual.get_array_memory_size() < 64 * 1024); + assert_eq!(actual.data_type(), batch.column(0).data_type()); + assert_eq!( + datatypes::arrow::compute::cast(actual.as_ref(), &ArrowDataType::Utf8) + .unwrap() + .as_ref(), + expected.as_ref(), + ); + } + #[test] fn test_single_value_dictionary_preserves_string_type() { use datatypes::arrow::array::{DictionaryArray, TimestampMillisecondArray, UInt32Array}; diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index c5eed8853ae..7224143f2c4 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1234,7 +1234,15 @@ pub(crate) fn should_split_flat_batches_for_merge( let file = &stream_ctx.input.files[file_index]; let file_meta = file.meta_ref(); if file_meta.level == 0 { - // Always split level 0 files. + // Short series can leave too many tiny batches retained by the merge. + if file_meta.num_series > 0 { + if !can_split_series(file_meta.num_rows, file_meta.num_series) { + return None; + } + total_rows += file_meta.num_rows; + total_series += file_meta.num_series; + } + // Older L0 files may lack series statistics; keep their fallback behavior. num_files_to_split += 1; continue; } else if file_meta.num_rows < SPLIT_ROW_THRESHOLD || file_meta.num_series == 0 { @@ -1913,6 +1921,30 @@ mod tests { ); } + #[tokio::test] + async fn test_should_split_flat_batches_for_merge_uses_level_zero_series_stats() { + for (num_rows, num_series, expected) in [ + (70_144, 35_083, None), + (100_000, 1_000, Some(100)), + (100_000, 0, Some(DEFAULT_READ_BATCH_SIZE / 4)), + ] { + let mut meta = new_test_file(num_rows, num_series).meta_ref().clone(); + meta.level = 0; + let file = FileHandle::new(meta, Arc::new(NoopFilePurger)); + let stream_ctx = new_test_stream_ctx(vec![file], vec![]).await; + let range_meta = new_test_range_meta(smallvec![RowGroupIndex { + index: 0, + row_group_index: 0, + }]); + + assert_eq!( + expected, + should_split_flat_batches_for_merge(&stream_ctx, &range_meta), + "rows={num_rows}, series={num_series}", + ); + } + } + #[tokio::test] async fn test_should_split_flat_batches_for_merge_falls_back_to_memtables() { let stream_ctx = new_test_stream_ctx(vec![], vec![new_test_memtable(5_000, 100)]).await; diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 323704c31e2..bbf6ff1bf92 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -662,8 +662,6 @@ pub(crate) async fn build_flat_sources( return Ok(None); } - let split_batch_size = should_split_flat_batches_for_merge(stream_ctx, range_meta); - let should_split = split_batch_size.is_some(); sources.reserve(num_indices); let mut ordered_sources = Vec::with_capacity(num_indices); ordered_sources.resize_with(num_indices, || None); @@ -737,6 +735,16 @@ pub(crate) async fn build_flat_sources( } } + // Queries merge one range at a time; compaction merges sources from all ranges. + let needs_merge = + ordered_sources.iter().flatten().count() > 1 || (compaction && stream_ctx.ranges.len() > 1); + let split_batch_size = if needs_merge { + should_split_flat_batches_for_merge(stream_ctx, range_meta) + } else { + None + }; + let should_split = split_batch_size.is_some(); + for stream in ordered_sources.into_iter().flatten() { if should_split { sources.push(Box::pin(SplitRecordBatchStream::new(stream))); diff --git a/src/object-store/src/layers/mock.rs b/src/object-store/src/layers/mock.rs index 270f776f935..4807a401c95 100644 --- a/src/object-store/src/layers/mock.rs +++ b/src/object-store/src/layers/mock.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::fmt::Debug; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use derive_builder::Builder; pub use oio::*; @@ -225,6 +225,50 @@ impl oio::Read for MockReader { } } +/// Writer calls observed before forwarding, including calls that fail. +#[derive(Debug, Default)] +pub struct WriterObservations { + pub write_sizes: Vec, + pub close_calls: usize, + pub abort_calls: usize, +} + +/// Records calls without changing the inner writer's completion or abort semantics. +pub struct ObservedWriter { + inner: oio::Writer, + observations: Arc>, +} + +impl ObservedWriter { + pub fn new(inner: oio::Writer, observations: Arc>) -> Self { + Self { + inner, + observations, + } + } +} + +impl oio::Write for ObservedWriter { + async fn write(&mut self, bytes: Buffer) -> Result<()> { + self.observations + .lock() + .unwrap() + .write_sizes + .push(bytes.len()); + self.inner.write(bytes).await + } + + async fn close(&mut self) -> Result { + self.observations.lock().unwrap().close_calls += 1; + self.inner.close().await + } + + async fn abort(&mut self) -> Result<()> { + self.observations.lock().unwrap().abort_calls += 1; + self.inner.abort().await + } +} + pub struct MockWriter { inner: oio::Writer, } diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index 368335f8450..8789e633995 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -41,7 +41,7 @@ use common_recordbatch::SendableRecordBatchStream; use common_time::range::TimestampRange; use datafusion::datasource::DefaultTableSource; use datafusion_common::TableReference as DfTableReference; -use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col}; +use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col, lit}; use futures::StreamExt; use object_store::ObjectStore; use session::context::QueryContextRef; @@ -53,6 +53,7 @@ use store_api::metric_engine_consts::{ use table::TableRef; use table::metadata::TableId; use table::table::adapter::DfTableProviderAdapter; +use tokio::sync::OwnedSemaphorePermit; use tokio_util::sync::CancellationToken; use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result}; @@ -79,7 +80,7 @@ impl Default for LogicalTableExportLimits { input_batch_bytes: 64 * 1024 * 1024, conversion_bytes: 1024 * 1024, writer: ParquetWriterLimits { - row_group_rows: 8192, + row_group_rows: 1024 * 1024, flush_threshold_bytes: 8 * 1024 * 1024, max_row_groups: 4096, }, @@ -255,7 +256,7 @@ impl LogicalTableExport { fn build_plan(&self, time_range: Option<&TimestampRange>) -> Result { let info = self.physical_table.table_info(); - let filters = self + let mut filters = self .physical_table .schema() .timestamp_column() @@ -264,6 +265,11 @@ impl LogicalTableExport { }) .into_iter() .collect::>(); + if self.logical_tables.len() == 1 + && let Some((&id, _)) = self.logical_tables.first_key_value() + { + filters.push(col(TABLE_ID).eq(lit(id))); + } let source = Arc::new(DefaultTableSource::new(Arc::new( DfTableProviderAdapter::new(self.physical_table.clone()), ))); @@ -283,11 +289,13 @@ impl LogicalTableExport { .filter(filter) .context(error::BuildDfLogicalPlanSnafu)?; } - builder - .sort(vec![col(TABLE_ID).sort(true, false)]) - .context(error::BuildDfLogicalPlanSnafu)? - .build() - .context(error::BuildDfLogicalPlanSnafu) + // A single selected table has a constant routing key after filtering. + if self.logical_tables.len() > 1 { + builder = builder + .sort(vec![col(TABLE_ID).sort(true, false)]) + .context(error::BuildDfLogicalPlanSnafu)?; + } + builder.build().context(error::BuildDfLogicalPlanSnafu) } } @@ -613,6 +621,46 @@ async fn expand_bounded_slice( batch.num_columns(), requested, )?; + // Parquet accepts compatible string dictionaries without expanding their values. + let fields = schema + .fields() + .iter() + .zip(batch.columns()) + .map(|(field, array)| { + if let DataType::Dictionary(_, value) = array.data_type() + && value.as_ref() == field.data_type() + && matches!(value.as_ref(), DataType::Utf8 | DataType::LargeUtf8) + { + Arc::new( + field + .as_ref() + .clone() + .with_data_type(array.data_type().clone()), + ) + } else { + field.clone() + } + }) + .collect::>(); + let schema = Arc::new(arrow::datatypes::Schema::new_with_metadata( + fields, + schema.metadata().clone(), + )); + // Sharing compatible arrays needs no CPU conversion. Keep the size check on + // this task only when its dictionary inspection is bounded. + if start < end + && batch + .columns() + .iter() + .zip(schema.fields()) + .all(|(array, field)| array.data_type() == field.data_type()) + && let Some(estimated) = estimate_batch_bytes_upper_bound(&batch, end - start, &[])? + && estimated <= conversion + { + let reservation = retained.saturating_add(estimated.saturating_mul(4)); + let permit = budget.reserve(reservation, cancellation).await?; + return convert_slice_payload(batch, schema, start, end - start, reservation, permit); + } let input = batch.clone(); let (len, estimated) = common_runtime::spawn_blocking_global(move || { rows_within_budget(&input, start, end, conversion, &[]) @@ -622,26 +670,36 @@ async fn expand_bounded_slice( let reservation = retained.saturating_add(estimated.saturating_mul(4)); let permit = budget.reserve(reservation, cancellation).await?; common_runtime::spawn_blocking_global(move || { - let slice = batch.slice(start, len); - let expanded = expand_export_batch(&slice, schema)?; - ensure!( - expanded.get_array_memory_size() <= reservation, - LogicalTableExportResourceSnafu { - reason: "converted backing buffers exceed reservation" - } - ); - Ok(( - Payload { - batch: expanded, - permit, - }, - len, - )) + convert_slice_payload(batch, schema, start, len, reservation, permit) }) .await .context(error::JoinTaskSnafu)? } +fn convert_slice_payload( + batch: RecordBatch, + schema: SchemaRef, + start: usize, + len: usize, + reservation: usize, + permit: OwnedSemaphorePermit, +) -> Result<(Payload, usize)> { + let expanded = expand_export_batch(&batch.slice(start, len), schema)?; + ensure!( + expanded.get_array_memory_size() <= reservation, + LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok(( + Payload { + batch: expanded, + permit, + }, + len, + )) +} + /// Limits nested casts to selected children: Arrow's List cast otherwise expands /// the entire values array, even when the parent has been sliced. pub(crate) fn expand_export_batch(batch: &RecordBatch, schema: SchemaRef) -> Result { @@ -788,6 +846,41 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result { Ok(bytes.saturating_add(16)) } +/// Estimates a conservative byte upper bound with bounded work on an async task. +/// `None` means this path cannot estimate the batch, not that it exceeds a budget. +fn estimate_batch_bytes_upper_bound( + batch: &RecordBatch, + rows: usize, + json_columns: &[usize], +) -> Result> { + let mut dictionary_values = 1024usize; + let mut max_row_bytes = 0usize; + for (index, array) in batch.columns().iter().enumerate() { + let size = if let Some(width) = array.data_type().primitive_width() { + (width + 16).max(32) + } else if let Some(dictionary) = array.as_any_dictionary_opt() + && dictionary.values().len() < rows + && dictionary.values().len() <= dictionary_values + && matches!( + dictionary.values().data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) + { + let values = dictionary.values(); + dictionary_values -= values.len(); + let max = (0..values.len()).try_fold(0usize, |max, row| { + Ok::<_, error::Error>(max.max(estimate_value_size(values.as_ref(), row)?)) + })?; + max.saturating_add(16).max(32) + } else { + return Ok(None); + }; + let expansion = if json_columns.contains(&index) { 8 } else { 1 }; + max_row_bytes = max_row_bytes.saturating_add(size.saturating_mul(expansion)); + } + Ok(max_row_bytes.checked_mul(rows)) +} + pub(crate) fn rows_within_budget( batch: &RecordBatch, start: usize, @@ -795,6 +888,12 @@ pub(crate) fn rows_within_budget( budget: usize, json_columns: &[usize], ) -> Result<(usize, usize)> { + if start < end + && let Some(bytes) = estimate_batch_bytes_upper_bound(batch, end - start, json_columns)? + && bytes <= budget + { + return Ok((end - start, bytes)); + } let mut bytes = 0usize; let mut row = start; while row < end { diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 7aa6daab127..6513d7e43fd 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -127,6 +127,188 @@ async fn read(store: &ObjectStore, path: &str) -> (SchemaRef, Vec) ) } +#[tokio::test] +async fn single_table_plan_filters_other_routing_ids_before_export() { + let mut unit = unit(); + unit.logical_tables.retain(|id, _| *id == 1025); + let schema = unit.physical_table.schema(); + let input = expand_export_batch( + &batch( + vec![Some(1027), Some(1025), Some(1024), Some(1025), Some(1026)], + vec![ + Some("other"), + Some("selected-a"), + None, + Some("selected-b"), + None, + ], + ), + schema.arrow_schema().clone(), + ) + .unwrap(); + unit.physical_table = table::test_util::MemTable::table( + "phy", + GreptimeRecordBatch::from_df_record_batch(schema, input), + ); + let context = datafusion::prelude::SessionContext::new_with_config( + datafusion::prelude::SessionConfig::new().with_target_partitions(4), + ); + let batches = context + .execute_logical_plan(unit.build_plan(None).unwrap()) + .await + .unwrap() + .collect() + .await + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let summary = export_stream( + &unit, + stream(batches), + &store, + LogicalTableExportLimits::default(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 2); + assert_eq!(summary.files, 1); + let (_, output) = read(&store, "cpu.v1.parquet").await; + let hosts = output + .iter() + .flat_map(|batch| batch.column(0).as_string::().iter()) + .collect::>(); + assert_eq!(hosts, vec![Some("selected-a"), Some("selected-b")]); +} + +#[tokio::test] +async fn dictionary_null_values_write_the_logical_parquet_schema() { + let input = batch(vec![Some(1025); 4], vec![Some("placeholder"); 4]); + let mut columns = input.columns().to_vec(); + columns[2] = Arc::new( + DictionaryArray::::try_new( + UInt32Array::from(vec![Some(0), Some(1), None, Some(0)]), + Arc::new(StringArray::from(vec![Some("host"), None])), + ) + .unwrap(), + ); + let input = RecordBatch::try_new(input.schema(), columns).unwrap(); + let unit = unit(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let summary = export_stream( + &unit, + stream(vec![input]), + &store, + LogicalTableExportLimits::default(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 4); + let (schema, output) = read(&store, "cpu.v1.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); + let hosts = output + .iter() + .flat_map(|batch| batch.column(0).as_string::().iter()) + .collect::>(); + assert_eq!(hosts, vec![Some("host"), None, None, Some("host")]); +} + +#[tokio::test] +async fn dictionary_slices_share_backing_and_respect_budget() { + let input = batch(vec![Some(1025); 8192], vec![Some("host"); 8192]) + .project(&[2, 3, 1]) + .unwrap(); + let schema = unit().logical_tables[&1025].schema.clone(); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let conversion = expand_bounded_slice( + input.clone(), + schema.clone(), + 7, + 8192, + 1024 * 1024, + &budget, + &token, + ); + tokio::pin!(conversion); + // Sharing an admitted batch completes without waiting for a conversion worker. + let std::task::Poll::Ready(result) = futures::poll!(&mut conversion) else { + panic!("compatible slice dispatched a conversion worker"); + }; + let (payload, rows) = result.unwrap(); + assert_eq!(rows, 8185); + let input_values = input.column(0).as_dictionary::().values(); + let output_values = payload + .batch + .column(0) + .as_dictionary::() + .values(); + assert!(Arc::ptr_eq(input_values, output_values)); + assert!(64 * 1024 * 1024 - budget.available().1 >= input.get_array_memory_size()); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + + // A batch which does not fit still takes the bounded conversion path. + let (payload, rows) = + expand_bounded_slice(input.clone(), schema.clone(), 7, 8192, 128, &budget, &token) + .await + .unwrap(); + assert_eq!(rows, 1); + drop(payload); + assert!( + expand_bounded_slice(input, schema, 7, 8192, 1, &budget, &token) + .await + .is_err() + ); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); +} + +#[test] +fn dictionary_size_bound_limits_inline_work() { + let values = (0..1025).map(|i| format!("host-{i}")).collect::>(); + let dictionary = DictionaryArray::::new( + UInt32Array::from(vec![0; 2048]), + Arc::new(StringArray::from_iter_values(values.iter())), + ); + let input = RecordBatch::try_from_iter([("host", Arc::new(dictionary) as ArrayRef)]).unwrap(); + assert!( + estimate_batch_bytes_upper_bound(&input, 2048, &[]) + .unwrap() + .is_none() + ); + assert_eq!( + rows_within_budget(&input, 0, 2048, 128, &[]).unwrap(), + (3, 114) + ); +} + +#[tokio::test] +async fn narrow_batches_share_row_groups() { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let batches = (0..8) + .map(|_| batch(vec![Some(1025); 8192], vec![Some("host"); 8192])) + .collect(); + let summary = export_stream( + &unit(), + stream(batches), + &store, + LogicalTableExportLimits::default(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 65536); + let bytes = store.read("cpu.v1.parquet").await.unwrap().to_bytes(); + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + assert_eq!(reader.metadata().num_row_groups(), 1); + let rows: usize = reader + .build() + .unwrap() + .map(|batch| batch.unwrap().num_rows()) + .sum(); + assert_eq!(rows, summary.rows); +} + #[tokio::test] async fn routes_across_batches_and_writes_empty_files() { let unit = unit(); @@ -249,6 +431,19 @@ async fn existing_outputs_are_not_overwritten() { ); } +#[test] +fn repeated_dictionary_budget_preserves_null_and_offset_charges() { + let host = "x".repeat(64); + let input = batch(vec![Some(1025); 8], vec![Some(host.as_str()); 8]) + .project(&[2, 4, 1]) + .unwrap(); + // Each row charges 96 bytes for the dictionary value, 32 for null, and 24 for time. + assert_eq!(rows_within_budget(&input, 0, 8, 1215, &[]).unwrap().0, 7); + let (rows, estimated) = rows_within_budget(&input, 0, 8, 1280, &[]).unwrap(); + assert_eq!(rows, 8); + assert!((1216..=1280).contains(&estimated)); +} + #[test] fn dictionary_and_nested_histogram_values_are_bounded_before_expansion() { let dictionary = DictionaryArray::::try_new( @@ -291,6 +486,25 @@ fn validates_selected_schemas_and_projects_only_selected_columns() { .unwrap(); assert_eq!(one.scan_projection, vec![0, 3]); assert_eq!(one.logical_tables[&1025].projection, vec![1]); + let LogicalPlan::Filter(filter) = one.build_plan(None).unwrap() else { + panic!("single-table export must filter the routing key without sorting"); + }; + let datafusion_expr::Expr::BinaryExpr(predicate) = filter.predicate else { + panic!("expected table-id equality"); + }; + assert_eq!(predicate.op, datafusion_expr::Operator::Eq); + assert!(matches!( + predicate.left.as_ref(), + datafusion_expr::Expr::Column(column) if column.name == TABLE_ID + )); + assert!(matches!( + predicate.right.as_ref(), + datafusion_expr::Expr::Literal(datafusion_common::ScalarValue::UInt32(Some(1025)), _) + )); + assert!(matches!( + unit.build_plan(None).unwrap(), + LogicalPlan::Sort(_) + )); assert!( LogicalTableExport::try_new(unit.physical_table.clone(), &[selected.clone(), selected]) .is_err() @@ -333,7 +547,7 @@ async fn cancellation_drops_input_and_aborts_active_upload() { } #[tokio::test] -async fn native_histogram_parquet_roundtrip() { +async fn supported_metric_types_parquet_roundtrip() { fn sample(data_type: &DataType) -> ArrayRef { match data_type { DataType::Struct(fields) => Arc::new(StructArray::new( @@ -359,44 +573,102 @@ async fn native_histogram_parquet_roundtrip() { other => panic!("unexpected histogram field {other}"), } } - let data_type = common_query::native_histogram::native_histogram_arrow_type(); - let field = Field::new("histogram", data_type.clone(), true); - let unit = LogicalTableExport::try_new( - table( - 1024, - "phy", - vec![Field::new(TABLE_ID, DataType::UInt32, false), field.clone()], - true, - ), - &[table(1025, "histogram", vec![field], false)], - ) - .unwrap(); - let batch = RecordBatch::try_from_iter_with_nullable([ - ( - TABLE_ID, - Arc::new(UInt32Array::from(vec![1025; 2])) as ArrayRef, - false, - ), - ("histogram", sample(&data_type), true), - ]) - .unwrap(); - let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); - let summary = export_stream( - &unit, - stream(vec![batch.clone()]), - &store, - export_limits(), - &CancellationToken::new(), - ) - .await - .unwrap(); - assert_eq!(summary.rows, 2); - let (schema, batches) = read(&store, "histogram.parquet").await; - assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); - assert_eq!( - arrow::compute::concat_batches(&schema, &batches).unwrap(), - batch.project(&[1]).unwrap() - ); + use common_datasource::file_format::parquet::packed_reader::{ + PackReadWindows, PackedParquetReader, + }; + use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; + use common_datasource::packed_writer::PackedWriter; + use datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder; + use futures::TryStreamExt; + + let timestamps: [ArrayRef; 4] = [ + Arc::new(arrow::array::TimestampSecondArray::from(vec![4, 5])), + Arc::new(arrow::array::TimestampMillisecondArray::from(vec![4, 5])), + Arc::new(arrow::array::TimestampMicrosecondArray::from(vec![4, 5])), + Arc::new(arrow::array::TimestampNanosecondArray::from(vec![4, 5])), + ]; + let values = [ + Arc::new(Float64Array::from(vec![Some(3.5), None])) as ArrayRef, + sample(&common_query::native_histogram::native_histogram_arrow_type()), + ]; + for timestamp in timestamps { + for value in &values { + let logical = RecordBatch::try_from_iter_with_nullable([ + ("ts", timestamp.clone(), false), + ( + "host", + Arc::new(StringArray::from(vec![Some(""), None])) as ArrayRef, + true, + ), + ("value", value.clone(), true), + ]) + .unwrap(); + let mut fields = vec![Field::new(TABLE_ID, DataType::UInt32, false)]; + fields.extend(logical.schema().fields().iter().map(|f| f.as_ref().clone())); + let unit = LogicalTableExport::try_new( + table(1024, "phy", fields.clone(), true), + &[table(1025, "metric", fields[1..].to_vec(), false)], + ) + .unwrap(); + let mut columns = vec![Arc::new(UInt32Array::from(vec![1025; 2])) as ArrayRef]; + columns.extend_from_slice(logical.columns()); + let input = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + for packed in [false, true] { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let writer = packed.then(|| PackedWriter::new(store.clone()).unwrap()); + let summary = export_stream_managed( + &unit, + stream(vec![input.clone()]), + &store, + export_limits(), + &CancellationToken::new(), + ExportWriteBudget::new(1), + writer.clone(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 2); + let batches = if let Some(writer) = writer { + writer + .lock() + .await + .finish(&CancellationToken::new()) + .await + .unwrap(); + let index: PackIndex = serde_json::from_slice( + &store.read(PACK_INDEX_FILE).await.unwrap().to_bytes(), + ) + .unwrap(); + index.validate_membership(["metric"]).unwrap(); + let entry = &index.tables[0]; + let object = &index.objects[0]; + let reader = PackedParquetReader::new( + PackReadWindows::new(store.clone()), + object.path.clone(), + object.length, + entry.offset, + entry.length, + ) + .unwrap(); + ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .build() + .unwrap() + .try_collect::>() + .await + .unwrap() + } else { + read(&store, "metric.parquet").await.1 + }; + assert_eq!(batches[0].schema().fields(), logical.schema().fields()); + assert_eq!( + arrow::compute::concat_batches(&logical.schema(), &batches).unwrap(), + logical + ); + } + } + } } struct PausedFileWriter { @@ -748,9 +1020,12 @@ async fn retained_backing_is_reserved_before_conversion_and_until_payload_drop() drop(blocker); let (payload, rows) = convert.await.unwrap(); assert_eq!(rows, 1); - assert!( - 64 * 1024 * 1024 - budget.available().1 >= full + payload.batch.get_array_memory_size() - ); + assert!(matches!( + payload.batch.column(0).data_type(), + DataType::Dictionary(_, _) + )); + // The retained dictionary shares its input buffers, so they are charged once. + assert!(64 * 1024 * 1024 - budget.available().1 >= full); drop(payload); assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); assert!(budget.reserve(64 * 1024 * 1024 + 1, &token).await.is_err()); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 8e3057bb7aa..13f9384cb5c 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1426,6 +1426,7 @@ struct ImportRequests { async fn import_http( instance: Arc, older: bool, + packed: bool, ) -> (String, tokio::task::JoinHandle<()>, Arc) { use axum::response::IntoResponse; use std::sync::atomic::Ordering; @@ -1446,6 +1447,12 @@ async fn import_http( if request.uri().path() == "/v1/capabilities" { observed.capabilities.fetch_add(1, Ordering::SeqCst); if older { + if packed { + return axum::Json(serde_json::json!({ + "metric_packed_import": 1 + })) + .into_response(); + } return axum::http::StatusCode::NOT_FOUND.into_response(); } } @@ -1528,6 +1535,11 @@ async fn packed_export_v2_cli_local_roundtrip() { metric_export_v2_cli_roundtrip(false, true, false).await; } +#[tokio::test(flavor = "multi_thread")] +async fn packed_import_v2_without_batch_ddl_roundtrip() { + metric_export_v2_cli_roundtrip(false, true, true).await; +} + #[tokio::test(flavor = "multi_thread")] async fn packed_export_v2_cli_s3_roundtrip() { if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) @@ -1606,7 +1618,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { &format!("CREATE VIEW \"{quoted_special}\" AS SELECT * FROM audit"), ) .await; - if packed && !s3 { + if packed && !s3 && !older { for id in 0..128 { let name = format!("batch_{id:03}"); sql(instance, &format!("CREATE TABLE {name} (host STRING PRIMARY KEY, val DOUBLE, ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='v2_a')")).await; @@ -1630,7 +1642,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { "CREATE TABLE bulk (host STRING PRIMARY KEY, payload STRING, ts TIMESTAMP TIME INDEX)", ) .await; - for batch in 0..64 { + for batch in 0..96 { let rows = (0..64) .map(|row| { let payload = (0..128) @@ -1813,7 +1825,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .await .unwrap() .content_length() - > 5 * 1024 * 1024 + > common_datasource::packed_writer::WRITE_BYTES as u64 ); } let mut preserved = Vec::new(); @@ -1923,7 +1935,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .await; } let (target_addr, target_server, requests) = - import_http(target.fe_instance().clone(), older).await; + import_http(target.fe_instance().clone(), older, packed).await; let state = destination.path().join("restore-state.json"); let mut import = vec![ "import-v2", @@ -1959,7 +1971,11 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { requests.capabilities.load(Ordering::SeqCst), if older { 1 } else { 2 } ); - assert_eq!(requests.batches.load(Ordering::SeqCst) > 1, !older); + if older { + assert_eq!(requests.batches.load(Ordering::SeqCst), 0); + } else { + assert!(requests.batches.load(Ordering::SeqCst) > 1); + } assert!(!state.exists()); for name in names { let source_table = table(instance, &name).await; @@ -2003,7 +2019,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .column_schemas() ); } - if packed { + if packed && !older { let schema_uri = format!("{uri}/schema-only"); let mut schema_args = vec![ "export-v2", @@ -2039,7 +2055,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .build() .await; let (schema_addr, schema_server, schema_requests) = - import_http(schema_target.fe_instance().clone(), false).await; + import_http(schema_target.fe_instance().clone(), false, false).await; schema_requests.fail_batch.store(false, Ordering::SeqCst); let schema_state = destination.path().join("schema-only-state.json"); let mut schema_import = vec![ From 29d5854e6a648281b04c36a4bdcc94b3ab47789d Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 12 Oct 2026 07:17:44 +0000 Subject: [PATCH 11/46] test(log-store): drive recovery, read, delete and stall failures through the fault matrix (#9488) * test(log-store): drive recovery, read, delete and stall failures through the fault matrix Add the second half of the object store WAL fault matrix: the failures of the object store access and of the store state that the conditional-create rows did not drive. Each case asserts the failed waiters and their error, whether the store is poisoned, the sequence, both indexed gauges at the failure, the exact movement of every counter and what a read sees after. Recovery: a listing, a trailer fetch and a footer fetch that fail, the catalog rejecting an object after every footer was read, and an object gone by the time it is fetched, below the boundary of the tip (ignored) and at it (construction fails); each leaves the gauges and the recovery metrics untouched. Reads: a segment fetch that fails, without a hole or a read timer sample, and a later read of the region that succeeds. Collection: a delete that fails and is attempted again at the next collection, and a stop that waits for a delete in flight, with the catalog agreeing with the listing in both outcomes. Store state: per-region entry id exhaustion refusing the append without touching the batch it sealed, and the durable stop paths with several creates in flight that succeed, fail transiently or conflict after stop began. Stalled appends released by admission, by a stop, by a terminal error and by the teardown of the actor, the first two timed as a pair. The recovery and read timers observed around parked fetches, twice each, with holds that cannot overlap. The cases use the existing doubles, extended with a failing listing and a failing range read at an offset on RecordingIo and parked range reads on ParkedIo. Case::row is now optional, and the assertion that every row of the failure matrix has a case is unchanged. Signed-off-by: jeremyhi * test(log-store): lose a wide object at its header and footer reads in the fault matrix The vanished-object case now runs three ways: a short object read whole, and a wide object deleted before its header and tail window reads or before its separate footer read, each below and at the boundary of the tip. The entry id exhaustion case reads the other region of the sealed batch back, the teardown case asserts the indexed gauges before and after the actor is gone, and two comments that restated their names are gone. Signed-off-by: jeremyhi * test(log-store): prove the stop-path ordering and share the wide-footer fixture in the fault matrix The case of creates that succeed after stop began keeps object 1 parked while objects 4 to 2 complete, round-trips a command through the actor and asserts that no append is acknowledged and nothing is indexed before object 1 is released. The stop-with-delete case replaces its yields with the same command barrier, queued behind the stop command. The wide-footer fixture moves to the tests module and is shared with the recovery test that fetches a footer longer than the tail window. The gauge helper is named for its side effect and the vanished-object stage enum for what it names. Signed-off-by: jeremyhi * test(log-store): keep the four-create scenario doc on its fault matrix case Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/object_store_wal/store.rs | 1155 ++++++++++++++++++- 1 file changed, 1118 insertions(+), 37 deletions(-) diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index d3a3acc6968..98b0c8cad6b 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -3292,15 +3292,7 @@ mod tests { async fn test_store_recovery_fetches_a_footer_longer_than_the_tail_window() { let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); - let regions = (RECOVERY_TAIL_WINDOW / FOOTER_ENTRY_LEN + 100) as u32; - let records = (1..=regions) - .map(|number| Record { - region_id: region(number), - entry_id: 1, - payload: Bytes::from_static(b"wide"), - }) - .collect::>(); - put_records(&object_store, 0, &records).await; + put_wide_object(&object_store).await; put_object(&object_store, 1, region(1), &[id(1, 1)]).await; let (io, reads) = RecordingIo::over(object_store.clone()); @@ -3312,7 +3304,7 @@ mod tests { let recovered = recover(io.as_ref()).await.unwrap(); let expected = recover_by_decoding(io.as_ref()).await.unwrap(); assert_same_recovery(&expected, &recovered); - assert_eq!(regions as usize, recovered.durable_entry_ids.len()); + assert_eq!(WIDE_REGIONS as usize, recovered.durable_entry_ids.len()); // The wide object took the header, the tail window and the footer; // the narrow one was read whole. @@ -3345,7 +3337,7 @@ mod tests { let store = open(object_store, &eager()).await; assert_eq!(id(1, 1), latest(&store, region(1))); - assert_eq!(1, latest(&store, region(regions))); + assert_eq!(1, latest(&store, region(WIDE_REGIONS))); store.stop().await.unwrap(); } @@ -3366,6 +3358,33 @@ mod tests { put_records(object_store, object_seq, &records).await; } + /// Regions enough for a footer longer than the window recovery reads + /// from the end of an object, so that recovery fetches the header, that + /// window and the footer of such an object separately. + const WIDE_REGIONS: u32 = (RECOVERY_TAIL_WINDOW / FOOTER_ENTRY_LEN + 100) as u32; + + /// Writes object 0 with a footer longer than the tail window and returns + /// the offsets of the window and of the footer. + async fn put_wide_object(object_store: &ObjectStore) -> (u64, u64) { + let records = (1..=WIDE_REGIONS) + .map(|number| Record { + region_id: region(number), + entry_id: 1, + payload: Bytes::from_static(b"wide"), + }) + .collect::>(); + put_records(object_store, 0, &records).await; + let bytes = object_store + .read(&object_path(object_store, 0)) + .await + .unwrap() + .to_vec(); + let (trailer, _) = footer_of(&bytes); + let tail_start = (bytes.len() - RECOVERY_TAIL_WINDOW) as u64; + assert!(trailer.footer_offset < tail_start); + (tail_start, trailer.footer_offset) + } + #[tokio::test] async fn test_store_resumes_sequence_and_durable_ids() { let _serialized = SERIALIZED.lock().await; @@ -6870,9 +6889,9 @@ mod tests { } /// Object access whose whole-object reads, or conditional creates, and - /// on request deletes, park until the test releases them, counting how - /// many are in flight. A release of false fails the operation with a - /// transient error before it reaches the object store. + /// on request deletes and range reads, park until the test releases + /// them, counting how many are in flight. A release of false fails the + /// operation with a transient error before it reaches the object store. struct ParkedIo { inner: ObjectStoreIo, parked: mpsc::UnboundedSender<(u64, oneshot::Sender)>, @@ -6881,6 +6900,7 @@ mod tests { creates_parked: AtomicBool, /// Whether deletes park. deletes_parked: AtomicBool, + range_reads_parked: AtomicBool, in_flight: AtomicUsize, max_in_flight: AtomicUsize, } @@ -6919,6 +6939,7 @@ mod tests { park_creates, creates_parked: AtomicBool::new(false), deletes_parked: AtomicBool::new(false), + range_reads_parked: AtomicBool::new(false), in_flight: AtomicUsize::new(0), max_in_flight: AtomicUsize::new(0), }), @@ -6959,6 +6980,9 @@ mod tests { } async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + if self.range_reads_parked.load(Ordering::SeqCst) && !self.park(object_seq).await { + return injected_failure("read", self.object_path(object_seq)); + } self.inner.get_range(object_seq, offset, len).await } @@ -6984,7 +7008,8 @@ mod tests { /// length) and, on request, reports the next conditional create as failed /// after it wrote the object, fails it with a given error before it /// writes, stores the same content before it writes, or damages the bytes - /// the next range read returns. + /// the next range read returns; it also fails the next listing, or the + /// next range read of an object at an offset, with a transient error. struct RecordingIo { inner: ObjectStoreIo, reads: RangeReads, @@ -6992,6 +7017,8 @@ mod tests { fail_next_put: Mutex>, damage_next_range_read: AtomicBool, duplicate_next_put: AtomicBool, + fail_next_list: AtomicBool, + fail_range_read_at: Mutex>, } impl RecordingIo { @@ -7004,6 +7031,8 @@ mod tests { fail_next_put: Mutex::new(None), damage_next_range_read: AtomicBool::new(false), duplicate_next_put: AtomicBool::new(false), + fail_next_list: AtomicBool::new(false), + fail_range_read_at: Mutex::new(None), }; (Arc::new(io), reads) } @@ -7053,6 +7082,15 @@ mod tests { async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { self.reads.lock().unwrap().push((object_seq, offset, len)); + if self + .fail_range_read_at + .lock() + .unwrap() + .take_if(|at| *at == (object_seq, offset)) + .is_some() + { + return injected_failure("read", self.object_path(object_seq)); + } let bytes = self.inner.get_range(object_seq, offset, len).await?; if self.damage_next_range_read.swap(false, Ordering::Relaxed) { let mut damaged = bytes.to_vec(); @@ -7069,6 +7107,9 @@ mod tests { } async fn list(&self) -> Result> { + if self.fail_next_list.swap(false, Ordering::Relaxed) { + return injected_failure("list", PREFIX.to_string()); + } self.inner.list().await } @@ -7077,8 +7118,9 @@ mod tests { } } - /// Drives every outcome of a conditional create through the store, one - /// case per fault, and checks that every row of the *Failure matrix* has one. + /// Drives every outcome of a conditional create and every failure of the + /// object store access through the store, one case per fault, and checks + /// that every row of the *Failure matrix* has one. mod fault_matrix { use object_store::layers::mock::{self, MockLayerBuilder, oio}; use prometheus::Histogram; @@ -7126,10 +7168,11 @@ mod tests { Row::StartFails, ]; - /// One failure of the object store access and the row it exercises. + /// One failure of the object store access and the row it exercises, + /// when the failure is an outcome of a conditional create. struct Case { fault: &'static str, - row: Row, + row: Option, run: fn() -> BoxFuture<'static, ()>, } @@ -7362,89 +7405,169 @@ mod tests { let cases = [ Case { fault: "a create that succeeds", - row: Row::Created, + row: Some(Row::Created), run: || Box::pin(create_succeeds()), }, Case { fault: "a create that finds an object with identical content", - row: Row::Created, + row: Some(Row::Created), run: || Box::pin(create_finds_an_identical_object()), }, Case { fault: "a create whose response is lost", - row: Row::Created, + row: Some(Row::Created), run: || Box::pin(create_response_is_lost()), }, Case { fault: "a create that fails transiently with later objects created and in flight, `durable` mode", - row: Row::TransientDurable, + row: Some(Row::TransientDurable), run: || Box::pin(create_fails_transiently_durable()), }, Case { fault: "a create that fails transiently, `enqueued` mode", - row: Row::TransientEnqueued, + row: Some(Row::TransientEnqueued), run: || Box::pin(create_fails_transiently_enqueued()), }, Case { fault: "a create that fails transiently after `stop` began, `enqueued` mode", - row: Row::TransientEnqueuedAfterStop, + row: Some(Row::TransientEnqueuedAfterStop), run: || Box::pin(create_fails_transiently_after_stop_began()), }, Case { fault: "a create that finds an object of an earlier epoch, `durable` mode", - row: Row::EarlierEpochDurable, + row: Some(Row::EarlierEpochDurable), run: || Box::pin(create_finds_an_earlier_epoch()), }, Case { fault: "a create that finds an object of an earlier epoch, `enqueued` mode", - row: Row::Conflict, + row: Some(Row::Conflict), run: || Box::pin(create_finds_an_earlier_epoch_enqueued()), }, Case { fault: "a create that finds an object of the same or a later epoch", - row: Row::Conflict, + row: Some(Row::Conflict), run: || Box::pin(create_finds_another_writer()), }, Case { fault: "a create that finds an object whose header cannot be read", - row: Row::Conflict, + row: Some(Row::Conflict), run: || Box::pin(create_finds_an_unreadable_object()), }, Case { fault: "a create that completes after a conflict poisoned the store", - row: Row::Conflict, + row: Some(Row::Conflict), run: || Box::pin(create_completes_after_the_store_was_poisoned()), }, Case { fault: "a created object the catalog rejects", - row: Row::CatalogError, + row: Some(Row::CatalogError), run: || Box::pin(catalog_rejects_the_created_object()), }, Case { fault: "a create at the last representable object sequence", - row: Row::LastSequence, + row: Some(Row::LastSequence), run: || Box::pin(create_takes_the_last_object_sequence()), }, Case { fault: "a start object that meets an object of an earlier epoch", - row: Row::StartMovesPast, + row: Some(Row::StartMovesPast), run: || Box::pin(start_object_meets_an_earlier_epoch()), }, Case { fault: "a start object that meets an object of the same or a later epoch", - row: Row::StartFails, + row: Some(Row::StartFails), run: || Box::pin(start_object_meets_another_writer()), }, Case { fault: "a start object that meets an object with identical content", - row: Row::StartFails, + row: Some(Row::StartFails), run: || Box::pin(start_object_meets_an_identical_object()), }, Case { fault: "a start object whose create has an unknown outcome", - row: Row::StartFails, + row: Some(Row::StartFails), run: || Box::pin(start_object_create_has_an_unknown_outcome()), }, + Case { + fault: "a listing that fails at recovery", + row: None, + run: || Box::pin(listing_fails_at_recovery()), + }, + Case { + fault: "a trailer fetch that fails at recovery", + row: None, + run: || Box::pin(trailer_fetch_fails_at_recovery()), + }, + Case { + fault: "a footer fetch that fails at recovery", + row: None, + run: || Box::pin(footer_fetch_fails_at_recovery()), + }, + Case { + fault: "a recovery the catalog rejects after every footer was read", + row: None, + run: || Box::pin(catalog_rejects_the_recovery()), + }, + Case { + fault: "an object gone by the time recovery fetches it", + row: None, + run: || Box::pin(object_vanishes_before_its_fetch()), + }, + Case { + fault: "a segment fetch that fails at a read", + row: None, + run: || Box::pin(segment_fetch_fails_at_a_read()), + }, + Case { + fault: "a delete that fails at a collection", + row: None, + run: || Box::pin(delete_fails_at_a_collection()), + }, + Case { + fault: "a stop with a delete in flight", + row: None, + run: || Box::pin(stop_waits_for_a_delete_in_flight()), + }, + Case { + fault: "an append past the entry id range of its region", + row: None, + run: || Box::pin(region_exhausts_its_entry_ids()), + }, + Case { + fault: "creates that succeed after `stop` began, `durable` mode", + row: None, + run: || Box::pin(creates_succeed_after_stop_began()), + }, + Case { + fault: "a create that fails transiently after `stop` began, `durable` mode", + row: None, + run: || Box::pin(create_fails_transiently_after_stop_began_durable()), + }, + Case { + fault: "a create that conflicts after `stop` began", + row: None, + run: || Box::pin(create_conflicts_after_stop_began()), + }, + Case { + fault: "a stalled append released by admission, then one released by a stop", + row: None, + run: || Box::pin(stalled_appends_released_by_admission_and_by_a_stop()), + }, + Case { + fault: "a stalled append released by a terminal error", + row: None, + run: || Box::pin(stalled_append_released_by_a_terminal_error()), + }, + Case { + fault: "a stalled append released by the teardown of the actor", + row: None, + run: || Box::pin(stalled_append_released_by_the_teardown()), + }, + Case { + fault: "an object store that answers slowly", + row: None, + run: || Box::pin(object_store_answers_slowly()), + }, ]; common_telemetry::init_default_ut_logging(); @@ -7454,7 +7577,7 @@ mod tests { } for row in ROWS { assert!( - cases.iter().any(|case| case.row == row), + cases.iter().any(|case| case.row == Some(row)), "no case covers the failure matrix row {row:?}" ); } @@ -8226,5 +8349,963 @@ mod tests { ) .await; } + + async fn assert_appends_stopped( + appends: Vec>>, + ) { + for append in appends { + assert_stopped(&timeout(WAIT, append).await.unwrap().unwrap().unwrap_err()); + } + } + + /// Opens and stops a store on an empty prefix, which sets both + /// process-wide indexed gauges to its one start object and moves the + /// recovery and create counters, so a case that abandons a recovery + /// can tell the gauges standing still from a publication of what its + /// prefix holds. It runs before the case samples `Counters`. + async fn initialize_indexed_gauges() -> (i64, i64) { + let store = open(memory_store(), &eager()).await; + store.stop().await.unwrap(); + let gauges = indexed_gauges(); + assert_eq!((1, object_len(&[]) as i64), gauges); + gauges + } + + /// Opens a store over `io` and asserts that the recovery was + /// abandoned: `open` fails with an error `recognise` accepts and + /// `retry_hint`, no start object is written so the prefix still holds + /// `present`, both gauges stand still and no metric moves. + async fn assert_recovery_abandoned( + io: Arc, + present: &[u64], + recognise: impl Fn(&Error) -> bool, + retry_hint: RetryHint, + ) { + let gauges = initialize_indexed_gauges().await; + let before = Counters::sample(); + let error = ObjectStoreLogStore::open(io.clone(), &eager(), PREFIX.to_string()) + .await + .unwrap_err(); + assert!(recognise(&error), "unexpected error: {error:?}"); + assert_eq!(retry_hint, error.retry_hint()); + assert_eq!(present, object_seqs(io.as_ref()).await.as_slice()); + assert_eq!(gauges, indexed_gauges()); + assert_eq!(Counters::default(), Counters::since(before)); + } + + /// Writes object `object_seq` holding one entry of region 1 with + /// `data`, chained to the object below it. + async fn put_entry(object_store: &ObjectStore, object_seq: u64, data: &str) { + let record = Record { + region_id: region(1), + entry_id: id(object_seq, 1), + payload: Bytes::from(data.as_bytes().to_vec()), + }; + put_records(object_store, object_seq, &[record]).await; + } + + async fn listing_fails_at_recovery() { + let object_store = memory_store(); + put_entry(&object_store, 0, "a0").await; + put_entry(&object_store, 1, "a1").await; + let (io, _) = RecordingIo::over(object_store); + io.fail_next_list.store(true, Ordering::Relaxed); + assert_recovery_abandoned( + io.clone(), + &[0, 1], + |error| { + matches!( + error, + Error::WalObjectStore { + operation: "list", + .. + } + ) + }, + RetryHint::Retryable, + ) + .await; + assert!(!io.fail_next_list.load(Ordering::Relaxed)); + + let before = Counters::sample(); + let store = open_over(io.clone(), &eager()).await; + assert_indexed(io.as_ref(), &[0, 1, 2]).await; + assert_eq!( + expected_entries(region(1), &[(id(0, 1), "a0"), (id(1, 1), "a1")]), + read_entries(&store, region(1), 1).await + ); + assert_eq!( + Counters { + created_objects: 1, + recovered_objects: 2, + recoveries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn trailer_fetch_fails_at_recovery() { + fetch_fails_at_recovery(false).await; + } + + async fn footer_fetch_fails_at_recovery() { + fetch_fails_at_recovery(true).await; + } + + /// The range read of the tail window, or of the footer, of the one + /// object fails; the next open recovers it. + async fn fetch_fails_at_recovery(footer: bool) { + let object_store = memory_store(); + let (tail_start, footer_offset) = put_wide_object(&object_store).await; + let (io, _) = RecordingIo::over(object_store); + let offset = if footer { footer_offset } else { tail_start }; + *io.fail_range_read_at.lock().unwrap() = Some((0, offset)); + let path = io.object_path(0); + assert_recovery_abandoned( + io.clone(), + &[0], + |error| { + matches!(error, Error::WalObjectStore { operation: "read", path: actual, .. } if *actual == path) + }, + RetryHint::Retryable, + ) + .await; + assert!(io.fail_range_read_at.lock().unwrap().is_none()); + + let before = Counters::sample(); + let store = open_over(io.clone(), &eager()).await; + assert_indexed(io.as_ref(), &[0, 1]).await; + assert_eq!( + expected_entries(region(1), &[(1, "wide")]), + read_entries(&store, region(1), 1).await + ); + assert_eq!( + Counters { + created_objects: 1, + recovered_objects: 1, + recoveries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// The entry ranges of the region do not increase with the sequence, + /// which the catalog refuses once every footer is read. + async fn catalog_rejects_the_recovery() { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[id(3, 1), id(3, 2)]).await; + put_object(&object_store, 1, region(1), &[id(2, 1)]).await; + let io = Arc::new(ObjectStoreIo::new(object_store, PREFIX).unwrap()); + let path = io.object_path(1); + assert_recovery_abandoned( + io.clone(), + &[0, 1], + |error| { + matches!(error, Error::InvalidWalObject { path: actual, source, .. } + if *actual == path && matches!(**source, Error::CorruptedWalObject { .. })) + }, + RetryHint::NonRetryable, + ) + .await; + } + + /// Writes the chain 0 to 3 of epoch 1, one entry of region 1 each, + /// and `WIDE_REGIONS` more regions per object when `wide`; objects 2 + /// and 3 carry reclaim boundary 2, so object 1 is below the boundary + /// of the tip and object 2 is not. + async fn put_bounded_chain(object_store: &ObjectStore, wide: bool) { + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + for object_seq in 0..=3 { + let header = Header { + reclaim_boundary: if object_seq >= 2 { 2 } else { 0 }, + ..chain_header(object_seq, 1, (object_seq > 0).then(|| (object_seq - 1, 1))) + }; + let mut records = vec![Record { + region_id: region(1), + entry_id: id(object_seq, 1), + payload: Bytes::from(format!("a{object_seq}")), + }]; + if wide { + records.extend((2..=WIDE_REGIONS).map(|number| Record { + region_id: region(number), + entry_id: id(object_seq, 1), + payload: Bytes::from_static(b"wide"), + })); + } + put_object_with_header(&io, header, &records).await; + } + } + + /// The fetch of recovery at which an object is found gone. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum RecoveryFetchStage { + /// The read of a short object as a whole. + WholeObject, + /// The reads of the header and of the tail window of a wide object. + Header, + /// The separate read of the footer of a wide object. + Footer, + } + + /// Waits for `count` parked operations and returns their releases. + async fn parked_releases( + parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + count: usize, + ) -> Vec> { + let mut releases = Vec::with_capacity(count); + for _ in 0..count { + releases.push(next_parked_operation(parked).await.1); + } + releases + } + + /// Every fetch of the listed objects parks and one object is deleted + /// before one of its reads: below the boundary of the tip the object + /// is ignored, at or above it the construction fails. + async fn object_vanishes_before_its_fetch() { + for (vanished, fails) in [(1, false), (2, true)] { + for stage in [ + RecoveryFetchStage::WholeObject, + RecoveryFetchStage::Header, + RecoveryFetchStage::Footer, + ] { + object_vanishes_at(vanished, fails, stage).await; + } + } + } + + async fn object_vanishes_at(vanished: u64, fails: bool, stage: RecoveryFetchStage) { + let object_store = memory_store(); + put_bounded_chain(&object_store, stage != RecoveryFetchStage::WholeObject).await; + let (io, mut parked) = ParkedIo::over(object_store.clone()); + io.range_reads_parked.store(true, Ordering::SeqCst); + let gauges = initialize_indexed_gauges().await; + let before = Counters::sample(); + let opening = { + let io = io.clone(); + tokio::spawn(async move { + ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await + }) + }; + let path = object_path(&object_store, vanished); + let delete = || object_store.delete(&path); + let release_all = |releases: Vec>| { + for release in releases { + release.send(true).unwrap(); + } + }; + // A short object parks once, a wide one twice, at its header and + // at its tail window, and then once more at its footer. + let whole = stage == RecoveryFetchStage::WholeObject; + let first = parked_releases(&mut parked, if whole { 4 } else { 8 }).await; + if stage != RecoveryFetchStage::Footer { + delete().await.unwrap(); + } + release_all(first); + match stage { + RecoveryFetchStage::WholeObject => {} + RecoveryFetchStage::Header => release_all(parked_releases(&mut parked, 3).await), + RecoveryFetchStage::Footer => { + let footers = parked_releases(&mut parked, 4).await; + delete().await.unwrap(); + release_all(footers); + } + } + let opened = timeout(WAIT, opening).await.unwrap().unwrap(); + io.range_reads_parked.store(false, Ordering::SeqCst); + if fails { + let error = opened.unwrap_err(); + assert!( + is_not_found(&error) + && matches!(&error, Error::WalObjectStore { path, .. } if *path == io.object_path(2)), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0, 1, 3], object_seqs(io.as_ref()).await); + assert_eq!(gauges, indexed_gauges()); + assert_eq!(Counters::default(), Counters::since(before)); + return; + } + let store = opened.unwrap(); + assert_eq!(vec![0, 2, 3, 4], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[2, 3, 4]).await; + assert_eq!( + expected_entries(region(1), &[(id(2, 1), "a2"), (id(3, 1), "a3")]), + read_entries(&store, region(1), 1).await + ); + assert_eq!( + Counters { + created_objects: 1, + recovered_objects: 3, + recoveries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn segment_fetch_fails_at_a_read() { + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + let (path, segment) = store + .segment_location(&provider(region_id), 1) + .unwrap() + .unwrap(); + *io.fail_range_read_at.lock().unwrap() = Some((1, segment.start)); + let error = store + .read(&provider(region_id), 1, None) + .await + .unwrap() + .try_collect::>() + .await + .unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { operation: "read", path: actual, .. } if *actual == path), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + assert!(io.fail_range_read_at.lock().unwrap().is_none()); + assert!(store.wal_holes(&provider(region_id)).unwrap().is_empty()); + assert!(terminal(&store.terminal_error).is_none()); + assert_indexed(io.as_ref(), &[1, 2]).await; + let written = Counters { + created_objects: 2, + seal_to_durable: 2, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + acknowledged_appends: 2, + deleted_objects: 1, + ..Counters::default() + }; + assert_eq!(written, Counters::since(before)); + + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(2, 1), "a2")]), + read_entries(&store, region_id, 1).await + ); + assert_eq!( + Counters { + reads: 1, + ..written + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// The delete of the start object fails at the collection that + /// indexing object 1 starts; a watermark that leaves the boundary + /// where it is starts the next collection, which deletes it. + async fn delete_fails_at_a_collection() { + let (io, mut parked) = ParkedIo::parking_creates(memory_store()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let before = Counters::sample(); + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + release.send(false).unwrap(); + timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(); + assert!(store.catalog.read().unwrap().contains_object(0)); + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0, 1]).await; + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + let written = Counters { + created_objects: 1, + seal_to_durable: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + acknowledged_appends: 1, + reads: 1, + failed_deletes: 1, + ..Counters::default() + }; + assert_eq!(written, Counters::since(before)); + + store + .obsolete(&provider(region_id), region_id, 0) + .await + .unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + release.send(true).unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert!(!store.catalog.read().unwrap().contains_object(0)); + assert_indexed(io.as_ref(), &[1]).await; + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + assert_eq!( + Counters { + reads: 2, + deleted_objects: 1, + ..written + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// `stop` returns only after the delete in flight completed, and the + /// catalog then holds exactly what the prefix holds, whether the + /// delete succeeded or failed. + async fn stop_waits_for_a_delete_in_flight() { + for deleted in [true, false] { + let (io, mut parked) = ParkedIo::parking_creates(memory_store()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let before = Counters::sample(); + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + // The first poll queues the stop command; the command queued + // behind it is answered once the stop was handled, and the + // delete in flight still holds the answer to the stop back. + let mut stop = Box::pin(store.stop()); + assert!((&mut stop).now_or_never().is_none()); + round_trip_stopping_actor(&store).await; + assert!((&mut stop).now_or_never().is_none()); + release.send(deleted).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap(); + + let present = if deleted { vec![1] } else { vec![0, 1] }; + assert_eq!(present, object_seqs(io.as_ref()).await); + let indexed = store + .catalog + .read() + .unwrap() + .objects_in_order() + .map(|(object_seq, _)| object_seq) + .collect::>(); + assert_eq!(present, indexed); + assert_indexed(io.as_ref(), &present).await; + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + acknowledged_appends: 1, + reads: 1, + deleted_objects: deleted as u64, + failed_deletes: (!deleted) as u64, + ..Counters::default() + }, + Counters::since(before) + ); + } + } + + /// An append that alone runs past the position range of its region + /// is refused; the open batch it sealed first, which holds another + /// region too, is acknowledged, and the store keeps serving. + async fn region_exhausts_its_entry_ids() { + let store = open(memory_store(), &manual()).await; + let before = Counters::sample(); + let (region_id, other) = (region(1), region(2)); + let shared = spawn_append_batch( + &store, + vec![entry(&store, other, "b1"), entry(&store, region_id, "a1")], + ); + store.wait_for_admitted_appends(1).await.unwrap(); + let exhausting = (0..POSITION_LIMIT) + .map(|_| entry(&store, region_id, "")) + .collect::>(); + let error = store.append_batch(exhausting).await.unwrap_err(); + assert!( + matches!(error, Error::WalEntryPositionExhausted { region_id: actual, .. } if actual == region_id), + "unexpected error: {error:?}" + ); + let response = timeout(WAIT, shared).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(other, id(1, 1)), (region_id, id(1, 1))]), + response.last_entry_ids + ); + assert!(terminal(&store.terminal_error).is_none()); + + let next = spawn_append_batch(&store, vec![entry(&store, region_id, "a2")]); + store.wait_for_admitted_appends(2).await.unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, next).await.unwrap().unwrap().unwrap(); + assert_eq!(Some(&id(2, 1)), response.last_entry_ids.get(®ion_id)); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_indexed(store.io.as_ref(), &[1, 2]).await; + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(2, 1), "a2")]), + read_entries(&store, region_id, 1).await + ); + assert_eq!( + expected_entries(other, &[(id(1, 1), "b1")]), + read_entries(&store, other, 1).await + ); + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + sealed_objects: 2, + object_bytes: object_bytes_of(&[(other, "b1"), (region_id, "a1")]) + + one_entry_object(), + object_entries: 3, + acknowledged_appends: 2, + reads: 2, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Round-trips a command through the actor once `stop` began, when + /// the seal command is refused: the collection that `obsolete_all` of + /// a region without entries requests is answered by the actor, and + /// hides nothing of the regions a case reads. + async fn round_trip_stopping_actor(store: &ObjectStoreLogStore) { + let region_id = region(9); + store + .obsolete_all(&provider(region_id), region_id) + .await + .unwrap(); + } + + /// Four creates are in flight and a fifth batch waits for a slot when + /// `stop` begins: the fifth is dropped; objects 4 to 2 complete while + /// object 1 stays parked and nothing is acknowledged or indexed; then + /// object 1 releases the four acknowledgements in order. + async fn creates_succeed_after_stop_began() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let mut appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 1).await; + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + let dropped = appends.pop().unwrap(); + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + assert_stopped(&timeout(WAIT, dropped).await.unwrap().unwrap().unwrap_err()); + for object_seq in (2..=MAX_IN_FLIGHT_CREATES as u64).rev() { + releases.remove(&object_seq).unwrap().send(true).unwrap(); + } + wait_until(|| { + Counters::since(before).created_objects == MAX_IN_FLIGHT_CREATES as u64 - 1 + }) + .await; + round_trip_stopping_actor(&store).await; + assert!(appends.iter().all(|append| !append.is_finished())); + assert_eq!(0, latest(&store, region_id)); + assert_indexed(io.as_ref(), &[0]).await; + + releases.remove(&1).unwrap().send(true).unwrap(); + for (index, append) in appends.into_iter().enumerate() { + let response = timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + assert_eq!( + Some(&id(index as u64 + 1, 1)), + response.last_entry_ids.get(®ion_id) + ); + } + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + + // Nothing is collected once stop began. + assert_eq!(vec![0, 1, 2, 3, 4], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0, 1, 2, 3, 4]).await; + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!( + expected_entries( + region_id, + &[ + (id(1, 1), "a1"), + (id(2, 1), "a2"), + (id(3, 1), "a3"), + (id(4, 1), "a4") + ] + ), + read_entries(&store, region_id, 1).await + ); + assert_eq!( + Counters { + created_objects: 4, + seal_to_durable: 4, + sealed_objects: 5, + object_bytes: 5 * one_entry_object(), + object_entries: 5, + acknowledged_appends: 4, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + } + + /// Object 2 is created and object 3 in flight when object 1 fails + /// after `stop` began: every batch receives the stopped error, `stop` + /// waits for object 3, whose late create is an orphan. + async fn create_fails_transiently_after_stop_began_durable() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 3).await; + let mut releases = parked_operations(&mut parked, 3).await; + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + wait_until(|| store.stopped.load(Ordering::Acquire)).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + releases.remove(&1).unwrap().send(false).unwrap(); + assert_appends_stopped(appends).await; + assert!(!stop.is_finished()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 3, + object_bytes: 3 * one_entry_object(), + object_entries: 3, + ..Counters::default() + }, + Counters::since(before) + ); + + releases.remove(&3).unwrap().send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![0, 2, 3], object_seqs(io.as_ref()).await); + assert!(terminal(&store.terminal_error).is_none()); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + create_failures: 1, + sealed_objects: 3, + object_bytes: 3 * one_entry_object(), + object_entries: 3, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + } + + /// Object 2 is created and object 3 in flight when object 1 meets + /// another writer after `stop` began: the store is poisoned, every + /// batch receives the stopped error and `stop` still returns. + async fn create_conflicts_after_stop_began() { + let object_store = memory_store(); + let (store, io, mut parked) = + open_parking_creates(object_store.clone(), &eager()).await; + put_foreign(&object_store, 1, 1).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 3).await; + let mut releases = parked_operations(&mut parked, 3).await; + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + wait_until(|| store.stopped.load(Ordering::Acquire)).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + releases.remove(&1).unwrap().send(true).unwrap(); + assert_appends_stopped(appends).await; + assert!(matches!( + terminal(&store.terminal_error).as_deref(), + Some(Error::WalObjectConflict { .. }) + )); + assert_read_is_terminal(&store, region_id).await; + assert!(!stop.is_finished()); + assert_indexed(io.as_ref(), &[0]).await; + + releases.remove(&3).unwrap().send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + create_conflicts: 1, + poisoned: 1, + sealed_objects: 3, + object_bytes: 3 * one_entry_object(), + object_entries: 3, + ..Counters::default() + }, + Counters::since(before) + ); + } + + /// The `enqueued` configuration under which the second append stalls + /// behind the create of the first. + fn stalling() -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(1), + ..enqueued(manual()) + } + } + + /// What `stall_second_append` leaves in the metrics. + fn stalled_once() -> Counters { + Counters { + created_objects: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + recoveries: 1, + stalled_appends: 1, + ..Counters::default() + } + } + + /// The first stalled append is admitted once the upload completes; + /// the entry it admits puts the backlog at the threshold again, so + /// the next append stalls behind the create it seals, and the stop + /// releases it with the stopped error. Both waits are timed. + async fn stalled_appends_released_by_admission_and_by_a_stop() { + const SHORT: Duration = Duration::from_millis(20); + let region_id = region(1); + let before = Counters::sample(); + let stalled_at = Instant::now(); + let (store, io, mut parked, stalled, release) = stall_second_append(stalling()).await; + assert_eq!(stalled_once(), Counters::since(before)); + + let timer = timer_before(&METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS); + let holding_since = Instant::now(); + tokio::time::sleep(SHORT).await; + let held = holding_since.elapsed(); + release.send(true).unwrap(); + let response = timeout(WAIT, stalled).await.unwrap().unwrap().unwrap(); + assert_eq!(Some(&id(2, 1)), response.last_entry_ids.get(®ion_id)); + let admitted = Timed::since( + &METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS, + timer, + held, + stalled_at.elapsed(), + ); + assert_eq!(vec![1], collected_object_seqs(&store).await); + + let stalled_at = Instant::now(); + let refused = spawn_append_batch(&store, vec![entry(&store, region_id, "a3")]); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(2, object_seq); + let timer = timer_before(&METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS); + let holding_since = Instant::now(); + tokio::time::sleep(admitted.elapsed + SHORT).await; + let held = holding_since.elapsed(); + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + assert_stopped(&timeout(WAIT, refused).await.unwrap().unwrap().unwrap_err()); + let refused = Timed::since( + &METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS, + timer, + held, + stalled_at.elapsed(), + ); + Timed::assert_pair("stalled append", admitted, refused); + + release.send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![1, 2], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[1, 2]).await; + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(2, 1), "a2")]), + read_entries(&store, region_id, 1).await + ); + assert_eq!( + Counters { + created_objects: 3, + seal_to_durable: 2, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + recoveries: 1, + reads: 1, + deleted_objects: 1, + stalled_appends: 2, + stalled_waits: 2, + ..Counters::default() + }, + Counters::since(before) + ); + } + + /// The create the stalled append waits for meets another writer, + /// which poisons the `enqueued` store: the append is released with + /// the terminal error and its wait is recorded. + async fn stalled_append_released_by_a_terminal_error() { + let region_id = region(1); + let before = Counters::sample(); + let (store, io, _parked, stalled, release) = stall_second_append(stalling()).await; + assert_eq!(stalled_once(), Counters::since(before)); + let foreign = Header { + object_seq: 1, + epoch: 1, + prev: None, + reclaim_boundary: 0, + }; + put_object_with_header(&io.inner, foreign, &[]).await; + release.send(true).unwrap(); + let error = timeout(WAIT, stalled).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + assert_read_is_terminal(&store, region_id).await; + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + stalled_waits: 1, + ..stalled_once() + }, + Counters::since(before) + ); + } + + /// The store is dropped without `stop` while the append is stalled: + /// the actor's teardown releases it and records its wait, and the + /// parked create never runs. + async fn stalled_append_released_by_the_teardown() { + let before = Counters::sample(); + let (store, io, _parked, stalled, _release) = stall_second_append(stalling()).await; + assert_eq!(stalled_once(), Counters::since(before)); + assert_indexed(io.as_ref(), &[0]).await; + + stalled.abort(); + drop(store); + wait_until(|| LIVE_ACTORS.load(Ordering::SeqCst) == 0).await; + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + stalled_waits: 1, + ..stalled_once() + }, + Counters::since(before) + ); + } + + /// Every fetch of recovery and of a read parks until the case releases + /// it, which is how the recovery and read timers are read for the + /// durations they recorded. Each is observed around two holds, the + /// second longer than anything the first observations took. + async fn object_store_answers_slowly() { + const SHORT: Duration = Duration::from_millis(20); + let object_store = memory_store(); + put_entry(&object_store, 0, "a0").await; + let (io, mut parked) = ParkedIo::over(object_store); + let before = Counters::sample(); + let region_id = region(1); + let mut recoveries = Vec::new(); + let mut reads = Vec::new(); + let mut hold = SHORT; + for objects in 1..=2 { + let timer = timer_before(&METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS); + let started_at = Instant::now(); + let opening = { + let io = io.clone(); + tokio::spawn(async move { + ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await + }) + }; + // Every object is short enough to be read whole, and every + // such read parks. + let releases = parked_operations(&mut parked, objects).await; + let parked_at = Instant::now(); + tokio::time::sleep(hold).await; + let held = parked_at.elapsed(); + for release in releases.into_values() { + release.send(true).unwrap(); + } + let store = timeout(WAIT, opening).await.unwrap().unwrap().unwrap(); + recoveries.push(Timed::since( + &METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS, + timer, + held, + started_at.elapsed(), + )); + + io.range_reads_parked.store(true, Ordering::SeqCst); + let timer = timer_before(&METRIC_OBJECT_STORE_WAL_READ_SECONDS); + let started_at = Instant::now(); + let reading = { + let store = store.clone(); + tokio::spawn(async move { read_entries(&store, region_id, 1).await }) + }; + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + let parked_at = Instant::now(); + tokio::time::sleep(hold).await; + let held = parked_at.elapsed(); + release.send(true).unwrap(); + let entries = timeout(WAIT, reading).await.unwrap().unwrap(); + reads.push(Timed::since( + &METRIC_OBJECT_STORE_WAL_READ_SECONDS, + timer, + held, + started_at.elapsed(), + )); + io.range_reads_parked.store(false, Ordering::SeqCst); + assert_eq!(expected_entries(region_id, &[(id(0, 1), "a0")]), entries); + store.stop().await.unwrap(); + hold = recoveries[0].elapsed.max(reads[0].elapsed) + SHORT; + } + Timed::assert_pair("recovery", recoveries[0], recoveries[1]); + Timed::assert_pair("read", reads[0], reads[1]); + assert_indexed(io.as_ref(), &[0, 1, 2]).await; + assert_eq!( + Counters { + created_objects: 2, + recovered_objects: 3, + recoveries: 2, + reads: 2, + ..Counters::default() + }, + Counters::since(before) + ); + } } } From f50f42140accfb9ead325bf359d45240fb1cb95a Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 12 Oct 2026 08:08:21 +0000 Subject: [PATCH 12/46] build: reduce checkout-specific inputs for cache reuse (#9198) * test(mito2): cover regex inverted index pruning Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * build: keep build inputs checkout-relative for cross-worktree cache reuse Two small changes so compilations are byte-identical across worktrees of the same commit and can share a build cache (mbx, sccache): - src/common/version/build.rs: derive the workspace root relative to CARGO_MANIFEST_DIR instead of the CARGO_WORKSPACE_DIR injected by .cargo/config.toml (an absolute path). The env var remains as an explicit override for the submodule case; the [env] injection in .cargo/config.toml is removed so the relative derivation is used by default. - Replace three env!("CARGO_MANIFEST_DIR") uses in test/bench code (servers prom_decode bench + decode test, common-meta test_util) with a runtime std::env::var lookup, so the compiled artifact no longer embeds the checkout path. No behavior change: same files are read, same version info emitted. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * build: address cross-worktree cache review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .cargo/config.toml | 5 +++-- src/common/meta/AGENTS.md | 9 +++++++++ src/common/meta/src/test_util.rs | 8 ++++---- src/common/version/build.rs | 19 ++++++++++++++----- src/servers/AGENTS.md | 10 ++++++++++ src/servers/benches/prom_decode.rs | 4 +++- src/servers/src/prom_remote_write/decode.rs | 4 +++- 7 files changed, 46 insertions(+), 13 deletions(-) diff --git a/.cargo/config.toml b/.cargo/config.toml index fd372c95b84..8291bf467e4 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -21,5 +21,6 @@ checkout = true list_files = true internal_use_git2 = false -[env] -CARGO_WORKSPACE_DIR = { value = "", relative = true } +# Avoid globally injecting a checkout-specific workspace path. The version build script +# resolves the absolute workspace path from layout at build time; set CARGO_WORKSPACE_DIR +# explicitly when this repo is built as a submodule. diff --git a/src/common/meta/AGENTS.md b/src/common/meta/AGENTS.md index 3684f8f9db3..6c9710b9e45 100644 --- a/src/common/meta/AGENTS.md +++ b/src/common/meta/AGENTS.md @@ -41,3 +41,12 @@ cargo nextest run -p common-meta ``` Backend-specific coverage may also need `pg_kvbackend` or `mysql_kvbackend`. + +`test_util::etcd_certs_dir` and `test_util::test_certs_dir` resolve fixture +certificates using `CARGO_MANIFEST_DIR` at **run time** (`std::env::var`, not `env!`). +`cargo nextest run` sets it; direct test-binary execution requires setting the +variable to the crate directory. Run this example from the repository root: + +```bash +CARGO_MANIFEST_DIR="$PWD/src/common/meta" ./target/debug/deps/ +``` diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index 1009bc1c99a..d78dc815d09 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -322,8 +322,8 @@ macro_rules! maybe_skip_etcd_tls_integration_test { /// Returns the directory of the etcd TLS certs. pub fn etcd_certs_dir() -> PathBuf { - let project_path = env!("CARGO_MANIFEST_DIR"); - let project_path = PathBuf::from(project_path); + let project_path = + PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set")); let base = project_path.ancestors().nth(3).unwrap(); base.join("tests-integration") .join("fixtures") @@ -332,8 +332,8 @@ pub fn etcd_certs_dir() -> PathBuf { /// Returns the directory of the test certs. pub fn test_certs_dir() -> PathBuf { - let project_path = env!("CARGO_MANIFEST_DIR"); - let project_path = PathBuf::from(project_path); + let project_path = + PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set")); let base = project_path.ancestors().nth(3).unwrap(); base.join("tests-integration") .join("fixtures") diff --git a/src/common/version/build.rs b/src/common/version/build.rs index 81c805484b4..2037f826616 100644 --- a/src/common/version/build.rs +++ b/src/common/version/build.rs @@ -30,12 +30,21 @@ fn main() -> Result<(), Box> { println!("cargo:rerun-if-env-changed=RUSTC"); println!("cargo:rerun-if-env-changed=GREPTIME_PRODUCT_NAME"); + println!("cargo:rerun-if-env-changed=CARGO_WORKSPACE_DIR"); - // The "CARGO_WORKSPACE_DIR" is set manually (not by Rust itself) in Cargo config file, to - // solve the problem where the "CARGO_MANIFEST_DIR" is not what we want when this repo is - // made as a submodule in another repo. - let workspace_dir = - env::var("CARGO_WORKSPACE_DIR").or_else(|_| env::var("CARGO_MANIFEST_DIR"))?; + // Resolve the workspace root from CARGO_MANIFEST_DIR (this crate is + // src/common/version, three levels below the root). The path remains absolute + // and is used at build time. CARGO_WORKSPACE_DIR remains an explicit override + // for submodule builds; see https://github.com/rust-lang/cargo/issues/3946. + let workspace_dir = match env::var("CARGO_WORKSPACE_DIR") { + Ok(dir) if !dir.is_empty() => dir, + _ => env::var("CARGO_MANIFEST_DIR").map(|dir| { + PathBuf::from(dir) + .join("../../..") + .to_string_lossy() + .into_owned() + })?, + }; let workspace_root = PathBuf::from(&workspace_dir); println!( "cargo:rerun-if-changed={}", diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index d20704cf616..c3b712d3924 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -43,3 +43,13 @@ cargo nextest run -p servers Keep protocol translation here and permissions/database behavior in frontend. Async table batching returns storage-completion futures with its early acknowledgements. Keep completion reporting independent of the caller's lifetime. + +`prom_remote_write/decode.rs` and `benches/prom_decode.rs` resolve +`benches/write_request.pb.data` using `CARGO_MANIFEST_DIR` at **run time** +(`std::env::var`, not `env!`). `cargo nextest run` and `cargo bench` set it; +direct test/benchmark-binary execution requires setting it to the crate directory. +Run this example from the repository root: + +```bash +CARGO_MANIFEST_DIR="$PWD/src/servers" ./target/debug/deps/ +``` diff --git a/src/servers/benches/prom_decode.rs b/src/servers/benches/prom_decode.rs index 1cdc7d3782e..bbca7a4269f 100644 --- a/src/servers/benches/prom_decode.rs +++ b/src/servers/benches/prom_decode.rs @@ -27,7 +27,9 @@ use servers::prom_remote_write::validation::{PromValidationMode, validate_label_ use servers::prom_store::to_grpc_row_insert_requests; fn load_fixture_v1_bytes() -> Vec { - let mut d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let mut d = std::path::PathBuf::from( + std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set"), + ); d.push("benches"); d.push("write_request.pb.data"); std::fs::read(d).expect("read write_request.pb.data fixture") diff --git a/src/servers/src/prom_remote_write/decode.rs b/src/servers/src/prom_remote_write/decode.rs index 6f635c7caf1..f7324b4be6e 100644 --- a/src/servers/src/prom_remote_write/decode.rs +++ b/src/servers/src/prom_remote_write/decode.rs @@ -448,7 +448,9 @@ mod tests { #[test] fn test_decode_write_request() { - let mut d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let mut d = std::path::PathBuf::from( + std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set"), + ); d.push("benches"); d.push("write_request.pb.data"); let data = std::fs::read(d).unwrap(); From 2df33380a7375b473be4a3d728cadcb9ed2cfc41 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 12 Oct 2026 08:49:33 +0000 Subject: [PATCH 13/46] feat(catalog): declare information-schema table scopes (#9504) * feat(catalog): declare information-schema table scopes Signed-off-by: shuiyisong * chore: rename cluster to global Signed-off-by: shuiyisong * fix: update sqlness result Signed-off-by: shuiyisong * refactor(catalog): require information-schema table scopes Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- docs/rfcs/2026-05-28-table-semantic-layer.md | 7 + src/catalog/src/kvbackend/manager.rs | 5 + src/catalog/src/system_schema.rs | 62 ++- .../src/system_schema/information_schema.rs | 466 ++++++++++++------ .../information_schema/cluster_info.rs | 11 +- .../information_schema/columns.rs | 11 +- .../information_schema/flow_statistics.rs | 9 +- .../system_schema/information_schema/flows.rs | 9 +- .../information_schema/key_column_usage.rs | 10 +- .../information_schema/partitions.rs | 11 +- .../information_schema/procedure_info.rs | 11 +- .../information_schema/process_list.rs | 9 +- .../information_schema/recycle_bin.rs | 11 +- .../information_schema/region_info.rs | 10 +- .../information_schema/region_peers.rs | 11 +- .../information_schema/region_statistics.rs | 11 +- .../information_schema/schemata.rs | 11 +- .../system_schema/information_schema/ssts.rs | 23 +- .../information_schema/statistics.rs | 10 +- .../information_schema/table_constraints.rs | 11 +- .../information_schema/table_semantics.rs | 22 +- .../information_schema/tables.rs | 11 +- .../system_schema/information_schema/views.rs | 11 +- src/catalog/src/system_schema/memory_table.rs | 7 + src/table/src/requests/semantic.rs | 12 + .../information_schema/table_semantics.result | 17 + .../information_schema/table_semantics.sql | 8 + .../common/show/show_databases_tables.result | 88 ++-- .../common/system/information_schema.result | 84 ++-- .../standalone/common/view/create.result | 84 ++-- 30 files changed, 749 insertions(+), 314 deletions(-) diff --git a/docs/rfcs/2026-05-28-table-semantic-layer.md b/docs/rfcs/2026-05-28-table-semantic-layer.md index cb5f6e84dc4..79764ad80d2 100644 --- a/docs/rfcs/2026-05-28-table-semantic-layer.md +++ b/docs/rfcs/2026-05-28-table-semantic-layer.md @@ -70,6 +70,13 @@ concern, not query-time semantics). | `greptime.semantic.source_version` | Source protocol version, e.g. Prometheus remote write `1.0` / `2.0` | | `greptime.semantic.pipeline` | `greptime_trace_v1` (the signal-agnostic successor to `table_data_model`) | +**Information-schema tables** carry `greptime.semantic.table_scope`: `catalog` +means the provider filters every query to its construction catalog; `global` +means catalog-independent or shared data, including compatibility stubs. The +scope is declared at registration and appears as `table_scope` in `semantic_options`. +Enterprise authorization trusts only registered system-table metadata. Setting +this option on a user table does not change its permissions. + **Trace**: `greptime.semantic.trace.conventions` (the OTel `schema_url` the rows conform to, or `mixed` / `unknown` when not single-valued). **Metric** — v1 assumes one metric type per table, which is how both Prom RW and the post-v0.16 OTel ingestion path land data today; mixed-type tables are a follow-up. These are stamped for OTLP (which declares them on the wire and then discards them); Prometheus carries its type/unit in the name and gets identity only. diff --git a/src/catalog/src/kvbackend/manager.rs b/src/catalog/src/kvbackend/manager.rs index 2afc77c0d23..0beb2c0e1f5 100644 --- a/src/catalog/src/kvbackend/manager.rs +++ b/src/catalog/src/kvbackend/manager.rs @@ -94,6 +94,11 @@ pub struct KvBackendCatalogManager { pub(super) const CATALOG_CACHE_MAX_CAPACITY: u64 = 128; impl KvBackendCatalogManager { + /// Returns the registered system tables, excluding user-created metadata. + pub fn information_schema_provider(&self) -> &InformationSchemaProvider { + &self.system_catalog.information_schema_provider + } + pub fn view_info_cache(&self) -> Result { self.cache_registry.get().context(CacheNotFoundSnafu { name: "view_info_cache", diff --git a/src/catalog/src/system_schema.rs b/src/catalog/src/system_schema.rs index add91cc75f4..099b8832f40 100644 --- a/src/catalog/src/system_schema.rs +++ b/src/catalog/src/system_schema.rs @@ -33,9 +33,8 @@ use snafu::ResultExt; use store_api::data_source::DataSource; use store_api::storage::ScanRequest; use table::error::{SchemaConversionSnafu, TablesRecordBatchSnafu}; -use table::metadata::{ - FilterPushDownType, TableId, TableInfoBuilder, TableInfoRef, TableMetaBuilder, TableType, -}; +use table::metadata::{FilterPushDownType, TableId, TableInfoBuilder, TableMetaBuilder, TableType}; +use table::requests::TableOptions; use table::{Table, TableRef}; use crate::error::Result; @@ -71,33 +70,44 @@ trait SystemSchemaProviderInner { fn schema_name() -> &'static str; fn build_table(&self, name: &str) -> Option { self.system_table(name).map(|table| { - let table_info = Self::table_info(self.catalog_name().to_string(), &table); - let filter_pushdown = FilterPushDownType::Inexact; - let data_source = Arc::new(SystemTableDataSource::new(table)); - let table = Table::new(table_info, filter_pushdown, data_source); - Arc::new(table) + build_system_table( + self.catalog_name().to_string(), + Self::schema_name(), + table, + TableOptions::default(), + ) }) } fn system_table(&self, name: &str) -> Option; +} - fn table_info(catalog_name: String, table: &SystemTableRef) -> TableInfoRef { - let table_meta = TableMetaBuilder::empty() - .schema(table.schema()) - .primary_key_indices(vec![]) - .next_column_id(0) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(table.table_id()) - .name(table.table_name().to_string()) - .catalog_name(catalog_name) - .schema_name(Self::schema_name().to_string()) - .meta(table_meta) - .table_type(table.table_type()) - .build() - .unwrap(); - Arc::new(table_info) - } +fn build_system_table( + catalog_name: String, + schema_name: &str, + table: SystemTableRef, + options: TableOptions, +) -> TableRef { + let table_meta = TableMetaBuilder::empty() + .schema(table.schema()) + .primary_key_indices(vec![]) + .next_column_id(0) + .options(options) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(table.table_id()) + .name(table.table_name().to_string()) + .catalog_name(catalog_name) + .schema_name(schema_name.to_string()) + .meta(table_meta) + .table_type(table.table_type()) + .build() + .unwrap(); + Arc::new(Table::new( + Arc::new(table_info), + FilterPushDownType::Inexact, + Arc::new(SystemTableDataSource::new(table)), + )) } pub trait SystemTable { diff --git a/src/catalog/src/system_schema/information_schema.rs b/src/catalog/src/system_schema/information_schema.rs index 8c3f4bf0f5f..b2145df9cd3 100644 --- a/src/catalog/src/system_schema/information_schema.rs +++ b/src/catalog/src/system_schema/information_schema.rs @@ -38,7 +38,7 @@ mod views; #[cfg(all(test, feature = "enterprise"))] mod recycle_bin_test; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Weak}; use common_catalog::consts::{self, DEFAULT_CATALOG_NAME, INFORMATION_SCHEMA_NAME}; @@ -63,15 +63,15 @@ use store_api::metric_engine_consts::{ }; use store_api::region_info::RegionInfoEntry; use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; -use store_api::storage::{ScanRequest, TableId}; +use store_api::storage::ScanRequest; use table::TableRef; -use table::metadata::TableType; +use table::requests::{SEMANTIC_TABLE_SCOPE, TableOptions}; pub use table_names::*; use views::InformationSchemaViews; use self::columns::InformationSchemaColumns; use crate::CatalogManager; -use crate::error::{Error, Result}; +use crate::error::Error; use crate::process_manager::ProcessManagerRef; use crate::system_schema::information_schema::cluster_info::InformationSchemaClusterInfo; use crate::system_schema::information_schema::flow_statistics::InformationSchemaFlowStatistics; @@ -92,9 +92,7 @@ use crate::system_schema::information_schema::table_semantics::InformationSchema use crate::system_schema::information_schema::tables::InformationSchemaTables; use crate::system_schema::memory_table::MemoryTable; pub(crate) use crate::system_schema::predicate::Predicates; -use crate::system_schema::{ - SystemSchemaProvider, SystemSchemaProviderInner, SystemTable, SystemTableRef, -}; +use crate::system_schema::{SystemSchemaProvider, SystemTable, build_system_table}; const DENSE_PRIMARY_KEY_ENCODING: &str = "dense"; const SPARSE_PRIMARY_KEY_ENCODING: &str = "sparse"; @@ -134,7 +132,6 @@ lazy_static! { TABLE_PRIVILEGES, GLOBAL_STATUS, SESSION_STATUS, - PARTITIONS, PLUGINS, USER_PRIVILEGES, PROCESSLIST, @@ -163,12 +160,31 @@ pub struct MakeInformationTableRequest { pub kv_backend: KvBackendRef, } +/// The visibility guaranteed by an information-schema table provider. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InformationSchemaTableScope { + /// Every query is filtered to the catalog supplied when constructing the provider. + Catalog, + /// The table contains catalog-independent or shared data, including compatibility stubs. + Global, +} + +impl InformationSchemaTableScope { + pub const fn as_str(self) -> &'static str { + match self { + Self::Catalog => "catalog", + Self::Global => "global", + } + } +} + /// A factory trait for making information schema tables. /// /// This trait allows for extensibility of the information schema by providing /// a way to dynamically create custom information schema tables. pub trait InformationSchemaTableFactory { - fn make_information_table(&self, req: MakeInformationTableRequest) -> SystemTableRef; + /// Creates a table whose scope is declared by its [`InformationTable`] implementation. + fn make_information_table(&self, req: MakeInformationTableRequest) -> InformationTableRef; } pub type InformationSchemaTableFactoryRef = Arc; @@ -192,15 +208,8 @@ impl SystemSchemaProvider for InformationSchemaProvider { } } -impl SystemSchemaProviderInner for InformationSchemaProvider { - fn catalog_name(&self) -> &str { - &self.catalog_name - } - fn schema_name() -> &'static str { - INFORMATION_SCHEMA_NAME - } - - fn system_table(&self, name: &str) -> Option { +impl InformationSchemaProvider { + fn information_table(&self, name: &str) -> Option { if let Some(factory) = self.extra_table_factories.get(name) { let req = MakeInformationTableRequest { catalog_name: self.catalog_name.clone(), @@ -357,142 +366,76 @@ impl InformationSchemaProvider { } fn build_tables(&mut self) { - let mut tables = HashMap::new(); + let mut table_names = HashSet::new(); // SECURITY NOTE: // Carefully consider the tables that may expose sensitive cluster configurations, // authentication details, and other critical information. // Only put these tables under `greptime` catalog to prevent info leak. if self.catalog_name == DEFAULT_CATALOG_NAME { - tables.insert( - BUILD_INFO.to_string(), - self.build_table(BUILD_INFO).unwrap(), - ); - tables.insert( - REGION_PEERS.to_string(), - self.build_table(REGION_PEERS).unwrap(), - ); - tables.insert( - CLUSTER_INFO.to_string(), - self.build_table(CLUSTER_INFO).unwrap(), - ); - tables.insert( - PROCEDURE_INFO.to_string(), - self.build_table(PROCEDURE_INFO).unwrap(), - ); - tables.insert( - REGION_STATISTICS.to_string(), - self.build_table(REGION_STATISTICS).unwrap(), - ); - tables.insert( - REGION_INFO.to_string(), - self.build_table(REGION_INFO).unwrap(), - ); - tables.insert( - SSTS_MANIFEST.to_string(), - self.build_table(SSTS_MANIFEST).unwrap(), - ); - tables.insert( - SSTS_STORAGE.to_string(), - self.build_table(SSTS_STORAGE).unwrap(), - ); - tables.insert( - SSTS_INDEX_META.to_string(), - self.build_table(SSTS_INDEX_META).unwrap(), - ); + table_names.extend([ + BUILD_INFO, + REGION_PEERS, + CLUSTER_INFO, + PROCEDURE_INFO, + REGION_STATISTICS, + REGION_INFO, + SSTS_MANIFEST, + SSTS_STORAGE, + SSTS_INDEX_META, + ]); } - tables.insert(TABLES.to_string(), self.build_table(TABLES).unwrap()); - tables.insert(VIEWS.to_string(), self.build_table(VIEWS).unwrap()); - tables.insert(SCHEMATA.to_string(), self.build_table(SCHEMATA).unwrap()); - tables.insert(COLUMNS.to_string(), self.build_table(COLUMNS).unwrap()); - tables.insert( - KEY_COLUMN_USAGE.to_string(), - self.build_table(KEY_COLUMN_USAGE).unwrap(), - ); - tables.insert( - TABLE_CONSTRAINTS.to_string(), - self.build_table(TABLE_CONSTRAINTS).unwrap(), - ); - tables.insert( - STATISTICS.to_string(), - self.build_table(STATISTICS).unwrap(), - ); - tables.insert(FLOWS.to_string(), self.build_table(FLOWS).unwrap()); - tables.insert( - FLOW_STATISTICS.to_string(), - self.build_table(FLOW_STATISTICS).unwrap(), - ); - #[cfg(feature = "enterprise")] - tables.insert( - RECYCLE_BIN.to_string(), - self.build_table(RECYCLE_BIN).unwrap(), - ); - tables.insert( - TABLE_SEMANTICS.to_string(), - self.build_table(TABLE_SEMANTICS).unwrap(), - ); - if let Some(process_list) = self.build_table(PROCESS_LIST) { - tables.insert(PROCESS_LIST.to_string(), process_list); + table_names.extend([ + TABLES, + VIEWS, + SCHEMATA, + COLUMNS, + KEY_COLUMN_USAGE, + TABLE_CONSTRAINTS, + STATISTICS, + FLOWS, + FLOW_STATISTICS, + #[cfg(feature = "enterprise")] + RECYCLE_BIN, + TABLE_SEMANTICS, + PARTITIONS, + ]); + if self.process_manager.is_some() { + table_names.insert(PROCESS_LIST); } - for name in self.extra_table_factories.keys() { - tables.insert(name.clone(), self.build_table(name).expect(name)); - } - // Add memory tables - for name in MEMORY_TABLES.iter() { - tables.insert((*name).to_string(), self.build_table(name).expect(name)); - } - self.tables = tables; + table_names.extend(MEMORY_TABLES.iter().copied()); + table_names.extend(self.extra_table_factories.keys().map(String::as_str)); + self.tables = table_names + .into_iter() + .map(|name| { + let table = self.information_table(name).expect(name); + let options = TableOptions { + extra_options: HashMap::from([( + SEMANTIC_TABLE_SCOPE.to_string(), + table.scope().as_str().to_string(), + )]), + ..Default::default() + }; + let table = build_system_table( + self.catalog_name.clone(), + INFORMATION_SCHEMA_NAME, + table, + options, + ); + (name.to_string(), table) + }) + .collect(); } } -pub trait InformationTable { - fn table_id(&self) -> TableId; - - fn table_name(&self) -> &'static str; - - fn schema(&self) -> SchemaRef; - - fn to_stream(&self, request: ScanRequest) -> Result; - - fn scan_plan(&self, _request: ScanRequest) -> Result>> { - Ok(None) - } - - fn table_type(&self) -> TableType { - TableType::Temporary - } +/// An information-schema table with an explicit visibility guarantee. +pub trait InformationTable: SystemTable { + /// Returns the visibility guaranteed by this table's query implementation. + fn scope(&self) -> InformationSchemaTableScope; } -// Provide compatibility for legacy `information_schema` code. -impl SystemTable for T -where - T: InformationTable, -{ - fn table_id(&self) -> TableId { - InformationTable::table_id(self) - } - - fn table_name(&self) -> &'static str { - InformationTable::table_name(self) - } - - fn schema(&self) -> SchemaRef { - InformationTable::schema(self) - } - - fn table_type(&self) -> TableType { - InformationTable::table_type(self) - } - - fn to_stream(&self, request: ScanRequest) -> Result { - InformationTable::to_stream(self, request) - } - - fn scan_plan(&self, request: ScanRequest) -> Result>> { - InformationTable::scan_plan(self, request) - } -} +pub type InformationTableRef = Arc; pub type InformationExtensionRef = Arc + Send + Sync>; @@ -597,9 +540,252 @@ impl InformationExtension for NoopInformationExtension { #[cfg(test)] mod tests { + use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry}; + use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; + use common_meta::ddl::test_util::create_table::test_create_table_task; + use common_meta::key::schema_name::SchemaNameKey; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::kv_backend::txn::TxnService; + use common_meta::wal_provider::RegionWalOptions; + use datatypes::schema::Schema; + use futures_util::TryStreamExt; use store_api::region_info::RegionInfoEntry; + use store_api::storage::TableId; use super::*; + use crate::error::Result; + use crate::information_schema::NoopInformationExtension; + use crate::kvbackend::KvBackendCatalogManagerBuilder; + use crate::memory::MemoryCatalogManager; + use crate::process_manager::ProcessManager; + + struct ScopedFactory { + name: &'static str, + scope: InformationSchemaTableScope, + } + + impl InformationSchemaTableFactory for ScopedFactory { + fn make_information_table(&self, _req: MakeInformationTableRequest) -> InformationTableRef { + Arc::new(ScopedTable { + name: self.name, + scope: self.scope, + }) + } + } + + struct ScopedTable { + name: &'static str, + scope: InformationSchemaTableScope, + } + + impl InformationTable for ScopedTable { + fn scope(&self) -> InformationSchemaTableScope { + self.scope + } + } + + impl SystemTable for ScopedTable { + fn table_id(&self) -> TableId { + 999 + } + + fn table_name(&self) -> &'static str { + self.name + } + + fn schema(&self) -> SchemaRef { + Arc::new(Schema::new(vec![])) + } + + fn to_stream(&self, _request: ScanRequest) -> Result { + Ok(common_recordbatch::RecordBatches::empty().as_stream()) + } + } + + #[test] + fn every_registered_information_schema_table_has_scope() { + use InformationSchemaTableScope::{Catalog, Global}; + + let manager: Arc = MemoryCatalogManager::new(); + let backend = Arc::new(MemoryKvBackend::default()); + let factories: HashMap<_, InformationSchemaTableFactoryRef> = [ + ("future_catalog_table", Catalog), + ("future_global_table", Global), + (ENGINES, Catalog), + ] + .into_iter() + .map(|(name, scope)| { + ( + name.to_string(), + Arc::new(ScopedFactory { name, scope }) as _, + ) + }) + .collect(); + let catalog_tables = [ + TABLES, + VIEWS, + SCHEMATA, + COLUMNS, + KEY_COLUMN_USAGE, + TABLE_CONSTRAINTS, + STATISTICS, + FLOWS, + FLOW_STATISTICS, + TABLE_SEMANTICS, + PARTITIONS, + REGION_PEERS, + #[cfg(feature = "enterprise")] + RECYCLE_BIN, + "future_catalog_table", + ENGINES, + ]; + + for (catalog, with_process_manager) in [ + (DEFAULT_CATALOG_NAME, true), + (DEFAULT_CATALOG_NAME, false), + ("tenant", true), + ("tenant", false), + ] { + let provider = InformationSchemaProvider::new( + catalog.to_string(), + Arc::downgrade(&manager), + Arc::new(FlowMetadataManager::new(backend.clone())), + with_process_manager.then(|| Arc::new(ProcessManager::new(String::new(), None))), + backend.clone(), + ); + assert_eq!( + provider + .table(ENGINES) + .unwrap() + .table_info() + .meta + .options + .extra_options + .get(SEMANTIC_TABLE_SCOPE) + .map(String::as_str), + Some(Global.as_str()), + ); + // An extension replacing a built-in supplies the replacement table's scope. + let provider = provider.with_extra_table_factories(factories.clone()); + let expected_count = + 43 + usize::from(with_process_manager) + usize::from(cfg!(feature = "enterprise")) + - if catalog == DEFAULT_CATALOG_NAME { + 0 + } else { + 9 + }; + assert_eq!(provider.tables().len(), expected_count); + for (name, table) in provider.tables() { + let scope = if catalog_tables.contains(&name.as_str()) { + Catalog + } else { + Global + }; + assert_eq!( + table + .table_info() + .meta + .options + .extra_options + .get(SEMANTIC_TABLE_SCOPE) + .map(String::as_str), + Some(scope.as_str()), + "{catalog}.{name}" + ); + } + assert_eq!( + provider.table(REGION_PEERS).is_some(), + catalog == DEFAULT_CATALOG_NAME + ); + assert_eq!(provider.table(PROCESS_LIST).is_some(), with_process_manager); + if let Some(table) = provider.table(PROCESS_LIST) { + assert!(table.schema().column_schema_by_name("catalog").is_some()); + } + } + } + + #[tokio::test] + async fn newly_accessible_information_schema_tables_filter_catalogs() { + let backend = Arc::new(MemoryKvBackend::default()); + let caches = LayeredCacheRegistryBuilder::default() + .add_cache_registry(CacheRegistryBuilder::default().build()) + .add_cache_registry(build_fundamental_cache_registry(backend.clone())); + let manager = KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + backend.clone(), + Arc::new( + with_default_composite_cache_registry(caches) + .unwrap() + .build(), + ), + ) + .build(); + let metadata = manager.table_metadata_manager_ref(); + let flow_metadata = FlowMetadataManager::new(backend.clone()); + for (catalog, name, id) in [ + ("tenant_a", "visible_a", 1024), + ("tenant_b", "visible_b", 1025), + ] { + metadata + .schema_manager() + .create(SchemaNameKey::new(catalog, "public"), None, false) + .await + .unwrap(); + let mut info = test_create_table_task(name, id).table_info; + info.catalog_name = catalog.to_string(); + info.meta + .options + .extra_options + .insert(SEMANTIC_TABLE_SCOPE.to_string(), "catalog".to_string()); + metadata + .create_table_metadata( + info, + TableRouteValue::physical(vec![]), + RegionWalOptions::default(), + ) + .await + .unwrap(); + let (txn, _) = flow_metadata + .flow_name_manager() + .build_create_txn(catalog, name, id) + .unwrap(); + backend.txn(txn).await.unwrap(); + } + + for (catalog, visible, hidden) in [ + ("tenant_a", "visible_a", "visible_b"), + ("tenant_b", "visible_b", "visible_a"), + ] { + for name in [TABLE_SEMANTICS, STATISTICS, FLOW_STATISTICS] { + let table = manager + .table(catalog, INFORMATION_SCHEMA_NAME, name, None) + .await + .unwrap() + .unwrap(); + let batches = table + .scan_to_stream(ScanRequest::default()) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let output = batches + .iter() + .map(|batch| batch.pretty_print()) + .collect::>() + .join("\n"); + assert!(output.contains(visible), "{catalog}.{name}: {output}"); + assert!(!output.contains(hidden), "{catalog}.{name}: {output}"); + } + } + assert!( + !manager + .information_schema_provider() + .tables() + .contains_key("visible_a") + ); + } #[test] fn test_datanode_inspect_region_info_build_plan() { diff --git a/src/catalog/src/system_schema/information_schema/cluster_info.rs b/src/catalog/src/system_schema/information_schema/cluster_info.rs index 1ba1a55fb63..20ecba07d87 100644 --- a/src/catalog/src/system_schema/information_schema/cluster_info.rs +++ b/src/catalog/src/system_schema/information_schema/cluster_info.rs @@ -40,7 +40,10 @@ use store_api::storage::{ScanRequest, TableId}; use crate::CatalogManager; use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result}; -use crate::system_schema::information_schema::{CLUSTER_INFO, InformationTable, Predicates}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + CLUSTER_INFO, InformationSchemaTableScope, InformationTable, Predicates, +}; use crate::system_schema::utils; const PEER_TYPE_FRONTEND: &str = "FRONTEND"; @@ -139,6 +142,12 @@ impl InformationSchemaClusterInfo { } impl InformationTable for InformationSchemaClusterInfo { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaClusterInfo { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_CLUSTER_INFO_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/columns.rs b/src/catalog/src/system_schema/information_schema/columns.rs index e92efd3bebf..6e0524390ee 100644 --- a/src/catalog/src/system_schema/information_schema/columns.rs +++ b/src/catalog/src/system_schema/information_schema/columns.rs @@ -43,7 +43,10 @@ use crate::error::{ CreateRecordBatchSnafu, InternalSnafu, Result, UpgradeWeakCatalogManagerRefSnafu, }; use crate::information_schema::Predicates; -use crate::system_schema::information_schema::{COLUMNS, InformationTable}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + COLUMNS, InformationSchemaTableScope, InformationTable, +}; #[derive(Debug)] pub(super) struct InformationSchemaColumns { @@ -160,6 +163,12 @@ impl InformationSchemaColumns { } impl InformationTable for InformationSchemaColumns { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaColumns { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_COLUMNS_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/flow_statistics.rs b/src/catalog/src/system_schema/information_schema/flow_statistics.rs index 6163461eda7..218104765fd 100644 --- a/src/catalog/src/system_schema/information_schema/flow_statistics.rs +++ b/src/catalog/src/system_schema/information_schema/flow_statistics.rs @@ -41,7 +41,8 @@ use store_api::storage::{ScanRequest, TableId}; use crate::CatalogManager; use crate::error::{CreateRecordBatchSnafu, InternalSnafu, ListFlowsSnafu, Result}; use crate::information_schema::{FLOW_STATISTICS, Predicates}; -use crate::system_schema::information_schema::InformationTable; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{InformationSchemaTableScope, InformationTable}; use crate::system_schema::utils; const INIT_CAPACITY: usize = 42; @@ -108,6 +109,12 @@ impl InformationSchemaFlowStatistics { } impl InformationTable for InformationSchemaFlowStatistics { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaFlowStatistics { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/flows.rs b/src/catalog/src/system_schema/information_schema/flows.rs index 12a8783b58e..5530d67bb8e 100644 --- a/src/catalog/src/system_schema/information_schema/flows.rs +++ b/src/catalog/src/system_schema/information_schema/flows.rs @@ -50,7 +50,8 @@ use crate::error::{ Result, UpgradeWeakCatalogManagerRefSnafu, }; use crate::information_schema::{FLOWS, Predicates}; -use crate::system_schema::information_schema::InformationTable; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{InformationSchemaTableScope, InformationTable}; use crate::system_schema::utils; const INIT_CAPACITY: usize = 42; @@ -195,6 +196,12 @@ impl InformationSchemaFlows { } impl InformationTable for InformationSchemaFlows { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaFlows { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_FLOW_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/key_column_usage.rs b/src/catalog/src/system_schema/information_schema/key_column_usage.rs index 4f8f116afaf..46d10829758 100644 --- a/src/catalog/src/system_schema/information_schema/key_column_usage.rs +++ b/src/catalog/src/system_schema/information_schema/key_column_usage.rs @@ -35,8 +35,10 @@ use crate::CatalogManager; use crate::error::{ CreateRecordBatchSnafu, InternalSnafu, Result, UpgradeWeakCatalogManagerRefSnafu, }; +use crate::system_schema::SystemTable; use crate::system_schema::information_schema::{ - InformationTable, KEY_COLUMN_USAGE, Predicates, primary_key_encoding_index_type, + InformationSchemaTableScope, InformationTable, KEY_COLUMN_USAGE, Predicates, + primary_key_encoding_index_type, }; pub const CONSTRAINT_SCHEMA: &str = "constraint_schema"; @@ -148,6 +150,12 @@ impl InformationSchemaKeyColumnUsage { } impl InformationTable for InformationSchemaKeyColumnUsage { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaKeyColumnUsage { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_KEY_COLUMN_USAGE_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/partitions.rs b/src/catalog/src/system_schema/information_schema/partitions.rs index 1bfe3f634fd..3bb3c29b704 100644 --- a/src/catalog/src/system_schema/information_schema/partitions.rs +++ b/src/catalog/src/system_schema/information_schema/partitions.rs @@ -44,7 +44,10 @@ use crate::error::{ Result, UpgradeWeakCatalogManagerRefSnafu, }; use crate::kvbackend::KvBackendCatalogManager; -use crate::system_schema::information_schema::{InformationTable, PARTITIONS, Predicates}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, PARTITIONS, Predicates, +}; const TABLE_CATALOG: &str = "table_catalog"; const TABLE_SCHEMA: &str = "table_schema"; @@ -167,6 +170,12 @@ impl InformationSchemaPartitions { } impl InformationTable for InformationSchemaPartitions { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaPartitions { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_PARTITIONS_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/procedure_info.rs b/src/catalog/src/system_schema/information_schema/procedure_info.rs index b9a892d2ac8..e9f3490a592 100644 --- a/src/catalog/src/system_schema/information_schema/procedure_info.rs +++ b/src/catalog/src/system_schema/information_schema/procedure_info.rs @@ -35,7 +35,10 @@ use store_api::storage::{ScanRequest, TableId}; use crate::CatalogManager; use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result}; -use crate::system_schema::information_schema::{InformationTable, PROCEDURE_INFO, Predicates}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, PROCEDURE_INFO, Predicates, +}; use crate::system_schema::utils; const PROCEDURE_ID: &str = "procedure_id"; @@ -97,6 +100,12 @@ impl InformationSchemaProcedureInfo { } impl InformationTable for InformationSchemaProcedureInfo { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaProcedureInfo { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_PROCEDURE_INFO_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/process_list.rs b/src/catalog/src/system_schema/information_schema/process_list.rs index fb453c8c333..02d3148197b 100644 --- a/src/catalog/src/system_schema/information_schema/process_list.rs +++ b/src/catalog/src/system_schema/information_schema/process_list.rs @@ -36,7 +36,8 @@ use store_api::storage::{ScanRequest, TableId}; use crate::error::{self, InternalSnafu}; use crate::information_schema::Predicates; use crate::process_manager::ProcessManagerRef; -use crate::system_schema::information_schema::InformationTable; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{InformationSchemaTableScope, InformationTable}; /// Column names of `information_schema.process_list` pub const ID: &str = "id"; @@ -82,6 +83,12 @@ impl InformationSchemaProcessList { } impl InformationTable for InformationSchemaProcessList { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaProcessList { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_PROCESS_LIST_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/recycle_bin.rs b/src/catalog/src/system_schema/information_schema/recycle_bin.rs index c2022cf23da..7019b840507 100644 --- a/src/catalog/src/system_schema/information_schema/recycle_bin.rs +++ b/src/catalog/src/system_schema/information_schema/recycle_bin.rs @@ -31,7 +31,10 @@ use store_api::storage::{ScanRequest, TableId}; use crate::CatalogManager; use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result, TableMetadataManagerSnafu}; -use crate::system_schema::information_schema::{InformationTable, Predicates, RECYCLE_BIN}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, Predicates, RECYCLE_BIN, +}; use crate::system_schema::utils; const OBJECT_ID: &str = "object_id"; @@ -118,6 +121,12 @@ impl InformationSchemaRecycleBin { } impl InformationTable for InformationSchemaRecycleBin { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaRecycleBin { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_RECYCLE_BIN_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/region_info.rs b/src/catalog/src/system_schema/information_schema/region_info.rs index 3fad1c84405..defb32df904 100644 --- a/src/catalog/src/system_schema/information_schema/region_info.rs +++ b/src/catalog/src/system_schema/information_schema/region_info.rs @@ -27,8 +27,10 @@ use store_api::storage::{ScanRequest, TableId}; use crate::CatalogManager; use crate::error::{ProjectSchemaSnafu, Result}; use crate::information_schema::{ - DatanodeInspectKind, DatanodeInspectRequest, InformationTable, REGION_INFO, + DatanodeInspectKind, DatanodeInspectRequest, InformationSchemaTableScope, InformationTable, + REGION_INFO, }; +use crate::system_schema::SystemTable; use crate::system_schema::utils; /// Information schema table for region info. @@ -47,6 +49,12 @@ impl InformationSchemaRegionInfo { } impl InformationTable for InformationSchemaRegionInfo { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaRegionInfo { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_REGION_INFO_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/region_peers.rs b/src/catalog/src/system_schema/information_schema/region_peers.rs index 9eddb061f80..40ed16f1dc4 100644 --- a/src/catalog/src/system_schema/information_schema/region_peers.rs +++ b/src/catalog/src/system_schema/information_schema/region_peers.rs @@ -41,7 +41,10 @@ use crate::error::{ UpgradeWeakCatalogManagerRefSnafu, }; use crate::kvbackend::KvBackendCatalogManager; -use crate::system_schema::information_schema::{InformationTable, Predicates, REGION_PEERS}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, Predicates, REGION_PEERS, +}; pub const TABLE_CATALOG: &str = "table_catalog"; pub const TABLE_SCHEMA: &str = "table_schema"; @@ -105,6 +108,12 @@ impl InformationSchemaRegionPeers { } impl InformationTable for InformationSchemaRegionPeers { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaRegionPeers { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_REGION_PEERS_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/region_statistics.rs b/src/catalog/src/system_schema/information_schema/region_statistics.rs index 8d1a15ddcc3..8a7c19adc00 100644 --- a/src/catalog/src/system_schema/information_schema/region_statistics.rs +++ b/src/catalog/src/system_schema/information_schema/region_statistics.rs @@ -38,7 +38,10 @@ use store_api::storage::{ScanRequest, TableId}; use crate::CatalogManager; use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result}; use crate::information_schema::Predicates; -use crate::system_schema::information_schema::{InformationTable, REGION_STATISTICS}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, REGION_STATISTICS, +}; use crate::system_schema::utils; const REGION_ID: &str = "region_id"; @@ -138,6 +141,12 @@ impl InformationSchemaRegionStatistics { } impl InformationTable for InformationSchemaRegionStatistics { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaRegionStatistics { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_REGION_STATISTICS_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/schemata.rs b/src/catalog/src/system_schema/information_schema/schemata.rs index c8756057cdf..d9654a48bd5 100644 --- a/src/catalog/src/system_schema/information_schema/schemata.rs +++ b/src/catalog/src/system_schema/information_schema/schemata.rs @@ -36,7 +36,10 @@ use crate::error::{ CreateRecordBatchSnafu, InternalSnafu, Result, TableMetadataManagerSnafu, UpgradeWeakCatalogManagerRefSnafu, }; -use crate::system_schema::information_schema::{InformationTable, Predicates, SCHEMATA}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, Predicates, SCHEMATA, +}; use crate::system_schema::utils; pub const CATALOG_NAME: &str = "catalog_name"; @@ -93,6 +96,12 @@ impl InformationSchemaSchemata { } impl InformationTable for InformationSchemaSchemata { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaSchemata { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_SCHEMATA_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/ssts.rs b/src/catalog/src/system_schema/information_schema/ssts.rs index e364bddfd3d..9006ab6fed9 100644 --- a/src/catalog/src/system_schema/information_schema/ssts.rs +++ b/src/catalog/src/system_schema/information_schema/ssts.rs @@ -30,9 +30,10 @@ use store_api::storage::{ScanRequest, TableId}; use crate::CatalogManager; use crate::error::{ProjectSchemaSnafu, Result}; use crate::information_schema::{ - DatanodeInspectKind, DatanodeInspectRequest, InformationTable, SSTS_INDEX_META, SSTS_MANIFEST, - SSTS_STORAGE, + DatanodeInspectKind, DatanodeInspectRequest, InformationSchemaTableScope, InformationTable, + SSTS_INDEX_META, SSTS_MANIFEST, SSTS_STORAGE, }; +use crate::system_schema::SystemTable; use crate::system_schema::utils; fn projected_schema(schema: &SchemaRef, request: &ScanRequest) -> Result { @@ -75,6 +76,12 @@ impl InformationSchemaSstsManifest { } impl InformationTable for InformationSchemaSstsManifest { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaSstsManifest { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_SSTS_MANIFEST_TABLE_ID } @@ -134,6 +141,12 @@ impl InformationSchemaSstsStorage { } impl InformationTable for InformationSchemaSstsStorage { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaSstsStorage { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_SSTS_STORAGE_TABLE_ID } @@ -194,6 +207,12 @@ impl InformationSchemaSstsIndexMeta { } impl InformationTable for InformationSchemaSstsIndexMeta { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + +impl SystemTable for InformationSchemaSstsIndexMeta { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_SSTS_INDEX_META_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/statistics.rs b/src/catalog/src/system_schema/information_schema/statistics.rs index 70312350759..4aeacd21ebd 100644 --- a/src/catalog/src/system_schema/information_schema/statistics.rs +++ b/src/catalog/src/system_schema/information_schema/statistics.rs @@ -39,12 +39,14 @@ use crate::CatalogManager; use crate::error::{ CreateRecordBatchSnafu, InternalSnafu, Result, UpgradeWeakCatalogManagerRefSnafu, }; +use crate::system_schema::SystemTable; use crate::system_schema::information_schema::key_column_usage::{ CONSTRAINT_NAME_FULLTEXT_INDEX, CONSTRAINT_NAME_INVERTED_INDEX, CONSTRAINT_NAME_PRI, CONSTRAINT_NAME_SKIPPING_INDEX, CONSTRAINT_NAME_TIME_INDEX, }; use crate::system_schema::information_schema::{ - InformationTable, Predicates, STATISTICS, primary_key_encoding_index_type, + InformationSchemaTableScope, InformationTable, Predicates, STATISTICS, + primary_key_encoding_index_type, }; pub const TABLE_CATALOG: &str = "table_catalog"; @@ -135,6 +137,12 @@ impl InformationSchemaStatistics { } impl InformationTable for InformationSchemaStatistics { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaStatistics { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_STATISTICS_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/table_constraints.rs b/src/catalog/src/system_schema/information_schema/table_constraints.rs index 2208d303c2e..aedbe251935 100644 --- a/src/catalog/src/system_schema/information_schema/table_constraints.rs +++ b/src/catalog/src/system_schema/information_schema/table_constraints.rs @@ -40,7 +40,10 @@ use crate::information_schema::Predicates; use crate::information_schema::key_column_usage::{ CONSTRAINT_NAME_PRI, CONSTRAINT_NAME_TIME_INDEX, }; -use crate::system_schema::information_schema::{InformationTable, TABLE_CONSTRAINTS}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, TABLE_CONSTRAINTS, +}; /// The `TABLE_CONSTRAINTS` table describes which tables have constraints. #[derive(Debug)] @@ -102,6 +105,12 @@ impl InformationSchemaTableConstraints { } impl InformationTable for InformationSchemaTableConstraints { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaTableConstraints { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_TABLE_CONSTRAINTS_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/table_semantics.rs b/src/catalog/src/system_schema/information_schema/table_semantics.rs index 07f1a844f9e..b571b03d40f 100644 --- a/src/catalog/src/system_schema/information_schema/table_semantics.rs +++ b/src/catalog/src/system_schema/information_schema/table_semantics.rs @@ -59,7 +59,10 @@ use crate::CatalogManager; use crate::error::{ CreateRecordBatchSnafu, InternalSnafu, Result, UpgradeWeakCatalogManagerRefSnafu, }; -use crate::system_schema::information_schema::{InformationTable, Predicates, TABLE_SEMANTICS}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, Predicates, TABLE_SEMANTICS, +}; use crate::system_schema::semantic_graph::{EntityGraphProviderRef, TableEntityDeclaration}; use crate::system_schema::utils; @@ -235,6 +238,12 @@ impl InformationSchemaTableSemantics { } impl InformationTable for InformationSchemaTableSemantics { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaTableSemantics { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_TABLE_SEMANTICS_TABLE_ID } @@ -514,6 +523,17 @@ mod tests { ); } + #[test] + fn extract_system_table_scope() { + let info = table_info(&[(table::requests::SEMANTIC_TABLE_SCOPE, "catalog")]); + let row = SemanticRow::extract(&info).unwrap(); + assert_eq!( + row.options_json.as_deref(), + Some(r#"{"table_scope":"catalog"}"#) + ); + assert!(row.signal_type.is_none()); + } + #[test] fn extract_skips_untagged_table() { let info = table_info(&[("ttl", "7d")]); diff --git a/src/catalog/src/system_schema/information_schema/tables.rs b/src/catalog/src/system_schema/information_schema/tables.rs index 6175c17d39b..6c9c0167da0 100644 --- a/src/catalog/src/system_schema/information_schema/tables.rs +++ b/src/catalog/src/system_schema/information_schema/tables.rs @@ -43,7 +43,10 @@ use crate::error::{ UpgradeWeakCatalogManagerRefSnafu, }; use crate::kvbackend::KvBackendCatalogManager; -use crate::system_schema::information_schema::{InformationTable, Predicates, TABLES}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, Predicates, TABLES, +}; use crate::system_schema::utils; pub const TABLE_CATALOG: &str = "table_catalog"; @@ -139,6 +142,12 @@ impl InformationSchemaTables { } impl InformationTable for InformationSchemaTables { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaTables { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_TABLES_TABLE_ID } diff --git a/src/catalog/src/system_schema/information_schema/views.rs b/src/catalog/src/system_schema/information_schema/views.rs index d1eb3b80ac0..65e16dfe120 100644 --- a/src/catalog/src/system_schema/information_schema/views.rs +++ b/src/catalog/src/system_schema/information_schema/views.rs @@ -38,7 +38,10 @@ use crate::error::{ UpgradeWeakCatalogManagerRefSnafu, ViewInfoNotFoundSnafu, }; use crate::kvbackend::KvBackendCatalogManager; -use crate::system_schema::information_schema::{InformationTable, Predicates, VIEWS}; +use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{ + InformationSchemaTableScope, InformationTable, Predicates, VIEWS, +}; const INIT_CAPACITY: usize = 42; pub const TABLE_CATALOG: &str = "table_catalog"; @@ -102,6 +105,12 @@ impl InformationSchemaViews { } impl InformationTable for InformationSchemaViews { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Catalog + } +} + +impl SystemTable for InformationSchemaViews { fn table_id(&self) -> TableId { INFORMATION_SCHEMA_VIEW_TABLE_ID } diff --git a/src/catalog/src/system_schema/memory_table.rs b/src/catalog/src/system_schema/memory_table.rs index de5ff120b2f..c7246939c23 100644 --- a/src/catalog/src/system_schema/memory_table.rs +++ b/src/catalog/src/system_schema/memory_table.rs @@ -31,6 +31,7 @@ use store_api::storage::{ScanRequest, TableId}; use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result}; use crate::system_schema::SystemTable; +use crate::system_schema::information_schema::{InformationSchemaTableScope, InformationTable}; /// A memory table with specified schema and columns. #[derive(Debug)] @@ -104,6 +105,12 @@ impl DfPartitionStream for MemoryTable { } } +impl InformationTable for MemoryTable { + fn scope(&self) -> InformationSchemaTableScope { + InformationSchemaTableScope::Global + } +} + impl SystemTable for MemoryTable { fn table_id(&self) -> TableId { self.table_id diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index dbe28eb240f..05b1edb56f5 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -49,6 +49,11 @@ pub const SEMANTIC_PER_TABLE_INDEX_KEY: &str = "greptime.internal.semantic.per_t // ---- Common keys (all signals) ---- +/// Information-schema visibility: `catalog` for providers that filter rows to +/// their catalog, or `global` for catalog-independent data. Only registered +/// system-table metadata is trusted for authorization. +pub const SEMANTIC_TABLE_SCOPE: &str = "greptime.semantic.table_scope"; + /// Signal kind: one of [`SIGNAL_TYPE_TRACE`] / [`SIGNAL_TYPE_LOG`] / /// [`SIGNAL_TYPE_METRIC`] / [`SIGNAL_TYPE_EVENT`]. pub const SEMANTIC_SIGNAL_TYPE: &str = "greptime.semantic.signal_type"; @@ -154,6 +159,7 @@ pub const SEMANTIC_VALUE_MIXED: &str = "mixed"; /// so an unknown key like `greptime.semantic.unknown_key` does not silently land /// in a table's options. Adding a key to the vocabulary means adding it here. pub const SEMANTIC_OPTION_KEYS: &[&str] = &[ + SEMANTIC_TABLE_SCOPE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, @@ -282,6 +288,7 @@ pub fn validate_semantic_option(key: &str, value: &str) -> bool { return !value.is_empty() && value.split(',').all(|column| !column.trim().is_empty()); } match key { + SEMANTIC_TABLE_SCOPE => matches!(value, "catalog" | "global"), SEMANTIC_PIPELINE | SEMANTIC_SOURCE_VERSION | SEMANTIC_METRIC_UNIT @@ -404,6 +411,11 @@ mod tests { #[test] fn test_validate_semantic_option() { // Enum keys reject out-of-domain values. + assert!(validate_semantic_option(SEMANTIC_TABLE_SCOPE, "catalog")); + assert!(validate_semantic_option(SEMANTIC_TABLE_SCOPE, "global")); + assert!(!validate_semantic_option(SEMANTIC_TABLE_SCOPE, "cluster")); + assert!(!validate_semantic_option(SEMANTIC_TABLE_SCOPE, "unknown")); + assert!(!validate_semantic_option(SEMANTIC_TABLE_SCOPE, "")); assert!(validate_semantic_option(SEMANTIC_SIGNAL_TYPE, "metric")); assert!(!validate_semantic_option(SEMANTIC_SIGNAL_TYPE, "spans")); assert!(validate_semantic_option(SEMANTIC_METRIC_TYPE, "counter")); diff --git a/tests/cases/standalone/common/information_schema/table_semantics.result b/tests/cases/standalone/common/information_schema/table_semantics.result index ee66ea3019b..6aac406f7a0 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.result +++ b/tests/cases/standalone/common/information_schema/table_semantics.result @@ -54,6 +54,7 @@ Affected Rows: 0 SELECT table_schema, table_name, signal_type, source, source_version, pipeline, metadata_quality, semantic_options FROM information_schema.table_semantics +WHERE table_schema = 'public' ORDER BY table_name; +--------------+----------------+-------------+---------------+----------------+--------------------+------------------+-----------------------------------------------------------------+ @@ -63,6 +64,22 @@ ORDER BY table_name; | public | traces_tagged | trace | opentelemetry | | | | {"trace.conventions":"https://opentelemetry.io/schemas/1.27.0"} | +--------------+----------------+-------------+---------------+----------------+--------------------+------------------+-----------------------------------------------------------------+ +-- System table scope is exposed without changing the view's columns. +SELECT table_name, semantic_options +FROM information_schema.table_semantics +WHERE table_schema = 'information_schema' + AND table_name IN ('build_info', 'partitions', 'process_list', 'table_semantics') +ORDER BY table_name; + ++-----------------+---------------------------+ +| table_name | semantic_options | ++-----------------+---------------------------+ +| build_info | {"table_scope":"global"} | +| partitions | {"table_scope":"catalog"} | +| process_list | {"table_scope":"global"} | +| table_semantics | {"table_scope":"catalog"} | ++-----------------+---------------------------+ + -- Predicate pushdown on a promoted column. SELECT table_name, signal_type FROM information_schema.table_semantics diff --git a/tests/cases/standalone/common/information_schema/table_semantics.sql b/tests/cases/standalone/common/information_schema/table_semantics.sql index 577c1edb693..d20d47e64e0 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.sql +++ b/tests/cases/standalone/common/information_schema/table_semantics.sql @@ -32,6 +32,14 @@ CREATE TABLE plain_table ( SELECT table_schema, table_name, signal_type, source, source_version, pipeline, metadata_quality, semantic_options FROM information_schema.table_semantics +WHERE table_schema = 'public' +ORDER BY table_name; + +-- System table scope is exposed without changing the view's columns. +SELECT table_name, semantic_options +FROM information_schema.table_semantics +WHERE table_schema = 'information_schema' + AND table_name IN ('build_info', 'partitions', 'process_list', 'table_semantics') ORDER BY table_name; -- Predicate pushdown on a promoted column. diff --git a/tests/cases/standalone/common/show/show_databases_tables.result b/tests/cases/standalone/common/show/show_databases_tables.result index 1c9e3a32460..29b11ff23f7 100644 --- a/tests/cases/standalone/common/show/show_databases_tables.result +++ b/tests/cases/standalone/common/show/show_databases_tables.result @@ -135,48 +135,48 @@ SHOW TABLE STATUS; +++++++++++++++++++ |Name|Engine|Version|Row_format|Rows|Avg_row_length|Data_length|Max_data_length|Index_length|Data_free|Auto_increment|Create_time|Update_time|Check_time|Collation|Checksum|Create_options|Comment| +++++++++++++++++++ -|build_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|character_sets||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|check_constraints||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|cluster_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|collation_character_set_applicability||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|collations||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|column_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|column_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|columns||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|engines||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|events||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|files||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|flow_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|flows||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|global_status||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|key_column_usage||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|optimizer_trace||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|parameters||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|partitions||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|plugins||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|procedure_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|process_list||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|processlist||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|profiling||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|referential_constraints||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|region_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|region_peers||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|region_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|routines||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|schema_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|schemata||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|session_status||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|ssts_index_meta||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|ssts_manifest||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|ssts_storage||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|table_constraints||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|table_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|table_semantics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|tables||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|user_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| -|views||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| +|build_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|character_sets||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|check_constraints||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|cluster_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|collation_character_set_applicability||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|collations||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|column_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|column_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|columns||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|engines||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|events||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|files||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|flow_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|flows||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|global_status||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|key_column_usage||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|optimizer_trace||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|parameters||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|partitions||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|plugins||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|procedure_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|process_list||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|processlist||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|profiling||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|referential_constraints||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|region_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|region_peers||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|region_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|routines||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|schema_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|schemata||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|session_status||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|ssts_index_meta||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|ssts_manifest||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|ssts_storage||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|table_constraints||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|table_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|table_semantics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|tables||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +|user_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|views||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +++++++++++++++++++ -- SQLNESS REPLACE (\s[\-0-9T:\.]{15,}) DATETIME @@ -186,7 +186,7 @@ SHOW TABLE STATUS LIKE 'tables'; +++++++++++++++++++ |Name|Engine|Version|Row_format|Rows|Avg_row_length|Data_length|Max_data_length|Index_length|Data_free|Auto_increment|Create_time|Update_time|Check_time|Collation|Checksum|Create_options|Comment| +++++++++++++++++++ -|tables||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| +|tables||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +++++++++++++++++++ -- SQLNESS REPLACE (\s[\-0-9T:\.]{15,}) DATETIME @@ -196,7 +196,7 @@ SHOW TABLE STATUS WHERE Name = 'tables'; +++++++++++++++++++ |Name|Engine|Version|Row_format|Rows|Avg_row_length|Data_length|Max_data_length|Index_length|Data_free|Auto_increment|Create_time|Update_time|Check_time|Collation|Checksum|Create_options|Comment| +++++++++++++++++++ -|tables||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0||| +|tables||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| +++++++++++++++++++ -- SQLNESS REPLACE (\s[\-0-9T:\.]{15,}) DATETIME diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result index fa9dd02a782..a83d66f4109 100644 --- a/tests/cases/standalone/common/system/information_schema.result +++ b/tests/cases/standalone/common/system/information_schema.result @@ -15,48 +15,48 @@ order by table_schema, table_name; +++++++++++++++++++++++++ |table_catalog|table_schema|table_name|table_type|table_id|data_length|max_data_length|index_length|max_index_length|avg_row_length|engine|version|row_format|table_rows|data_free|auto_increment|create_time|update_time|check_time|table_collation|checksum|create_options|table_comment|temporary| +++++++++++++++++++++++++ -|greptime|information_schema|build_info|LOCALTEMPORARY|8|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|character_sets|LOCALTEMPORARY|9|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|check_constraints|LOCALTEMPORARY|12|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|cluster_info|LOCALTEMPORARY|31|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|11|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|collations|LOCALTEMPORARY|10|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|column_privileges|LOCALTEMPORARY|6|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|column_statistics|LOCALTEMPORARY|7|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|columns|LOCALTEMPORARY|4|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|engines|LOCALTEMPORARY|5|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|events|LOCALTEMPORARY|13|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|files|LOCALTEMPORARY|14|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|flow_statistics|LOCALTEMPORARY|45|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|flows|LOCALTEMPORARY|33|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|global_status|LOCALTEMPORARY|25|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|key_column_usage|LOCALTEMPORARY|16|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|optimizer_trace|LOCALTEMPORARY|17|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|parameters|LOCALTEMPORARY|18|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|partitions|LOCALTEMPORARY|28|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|plugins|LOCALTEMPORARY|47|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|procedure_info|LOCALTEMPORARY|34|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|process_list|LOCALTEMPORARY|36|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|processlist|LOCALTEMPORARY|49|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|profiling|LOCALTEMPORARY|19|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|referential_constraints|LOCALTEMPORARY|20|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|region_info|LOCALTEMPORARY|41|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|region_peers|LOCALTEMPORARY|29|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|region_statistics|LOCALTEMPORARY|35|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|routines|LOCALTEMPORARY|21|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|schema_privileges|LOCALTEMPORARY|22|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|schemata|LOCALTEMPORARY|15|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|session_status|LOCALTEMPORARY|26|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|ssts_index_meta|LOCALTEMPORARY|39|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|ssts_manifest|LOCALTEMPORARY|37|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|ssts_storage|LOCALTEMPORARY|38|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|statistics|LOCALTEMPORARY|43|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|table_constraints|LOCALTEMPORARY|30|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|table_privileges|LOCALTEMPORARY|23|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|table_semantics|LOCALTEMPORARY|42|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|tables|LOCALTEMPORARY|3|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|user_privileges|LOCALTEMPORARY|48|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| -|greptime|information_schema|views|LOCALTEMPORARY|32|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| +|greptime|information_schema|build_info|LOCALTEMPORARY|8|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|character_sets|LOCALTEMPORARY|9|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|check_constraints|LOCALTEMPORARY|12|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|cluster_info|LOCALTEMPORARY|31|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|11|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|collations|LOCALTEMPORARY|10|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|column_privileges|LOCALTEMPORARY|6|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|column_statistics|LOCALTEMPORARY|7|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|columns|LOCALTEMPORARY|4|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|engines|LOCALTEMPORARY|5|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|events|LOCALTEMPORARY|13|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|files|LOCALTEMPORARY|14|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|flow_statistics|LOCALTEMPORARY|45|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|flows|LOCALTEMPORARY|33|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|global_status|LOCALTEMPORARY|25|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|key_column_usage|LOCALTEMPORARY|16|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|optimizer_trace|LOCALTEMPORARY|17|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|parameters|LOCALTEMPORARY|18|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|partitions|LOCALTEMPORARY|28|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|plugins|LOCALTEMPORARY|47|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|procedure_info|LOCALTEMPORARY|34|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|process_list|LOCALTEMPORARY|36|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|processlist|LOCALTEMPORARY|49|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|profiling|LOCALTEMPORARY|19|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|referential_constraints|LOCALTEMPORARY|20|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|region_info|LOCALTEMPORARY|41|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|region_peers|LOCALTEMPORARY|29|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|region_statistics|LOCALTEMPORARY|35|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|routines|LOCALTEMPORARY|21|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|schema_privileges|LOCALTEMPORARY|22|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|schemata|LOCALTEMPORARY|15|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|session_status|LOCALTEMPORARY|26|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|ssts_index_meta|LOCALTEMPORARY|39|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|ssts_manifest|LOCALTEMPORARY|37|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|ssts_storage|LOCALTEMPORARY|38|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|statistics|LOCALTEMPORARY|43|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|table_constraints|LOCALTEMPORARY|30|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|table_privileges|LOCALTEMPORARY|23|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|table_semantics|LOCALTEMPORARY|42|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|tables|LOCALTEMPORARY|3|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|user_privileges|LOCALTEMPORARY|48|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|views|LOCALTEMPORARY|32|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| |greptime|public|numbers|LOCALTEMPORARY|2|0|0|0|0|0|test_engine|11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| +++++++++++++++++++++++++ diff --git a/tests/cases/standalone/common/view/create.result b/tests/cases/standalone/common/view/create.result index 577583515e9..82d6765ded8 100644 --- a/tests/cases/standalone/common/view/create.result +++ b/tests/cases/standalone/common/view/create.result @@ -96,52 +96,52 @@ ORDER BY TABLE_NAME, TABLE_TYPE; +++++++++++++++++++++++++ |table_catalog|table_schema|table_name|table_type|table_id|data_length|max_data_length|index_length|max_index_length|avg_row_length|engine|version|row_format|table_rows|data_free|auto_increment|create_time|update_time|check_time|table_collation|checksum|create_options|table_comment|temporary| +++++++++++++++++++++++++ -|greptime|information_schema|build_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|character_sets|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|check_constraints|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|cluster_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|collations|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|column_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|column_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|columns|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|engines|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|events|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|files|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|flow_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|flows|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|global_status|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|key_column_usage|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| +|greptime|information_schema|build_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|character_sets|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|check_constraints|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|cluster_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|collations|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|column_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|column_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|columns|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|engines|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|events|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|files|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|flow_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|flows|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|global_status|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|key_column_usage|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| |greptime|public|numbers|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID|test_engine|ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|optimizer_trace|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|parameters|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|partitions|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|plugins|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|procedure_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|process_list|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|processlist|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|profiling|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|referential_constraints|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|region_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|region_peers|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|region_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|routines|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|schema_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|schemata|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|session_status|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|ssts_index_meta|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|ssts_manifest|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|ssts_storage|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|table_constraints|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|table_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|table_semantics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|tables|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| +|greptime|information_schema|optimizer_trace|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|parameters|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|partitions|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|plugins|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|procedure_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|process_list|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|processlist|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|profiling|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|referential_constraints|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|region_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|region_peers|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|region_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|routines|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|schema_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|schemata|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|session_status|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|ssts_index_meta|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|ssts_manifest|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|ssts_storage|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|table_constraints|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|table_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|table_semantics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +|greptime|information_schema|tables|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| |greptime|public|test_table|BASETABLE|ID|ID|ID|ID|ID|ID|mito|ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||N| |greptime|public|test_view|VIEW|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||N| |greptime|public|test_view2|VIEW|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||N| -|greptime|information_schema|user_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| -|greptime|information_schema|views|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|||Y| +|greptime|information_schema|user_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|views|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| +++++++++++++++++++++++++ -- SQLNESS REPLACE (\s\d+\s) ID From fc5831043f28ca88eff96b75b3da36a027b15833 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 12 Oct 2026 10:09:26 +0000 Subject: [PATCH 14/46] feat(cli): simplify packed metric exports with a boolean flag (#9531) * feat(cli): select packed metric export with the layout option Signed-off-by: jeremyhi * test(cli): inspect the packed capability error cause Signed-off-by: jeremyhi * test(cli): follow the standard capability error source Signed-off-by: jeremyhi * refactor(cli): clarify the metric export gate helper name Signed-off-by: jeremyhi * feat(cli): use a boolean flag for packed metric exports Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- docs/rfcs/2026-09-11-metric-export-import.md | 6 + src/cli/src/data/export_v2.rs | 5 +- src/cli/src/data/export_v2/command.rs | 137 ++++++++++-------- src/cli/src/data/export_v2/coordinator.rs | 6 +- src/cli/src/data/export_v2/data.rs | 24 ++- .../tests/export_logical_tables.rs | 104 +++++++++---- 6 files changed, 183 insertions(+), 99 deletions(-) diff --git a/docs/rfcs/2026-09-11-metric-export-import.md b/docs/rfcs/2026-09-11-metric-export-import.md index 57212ea53cb..bf66df932b9 100644 --- a/docs/rfcs/2026-09-11-metric-export-import.md +++ b/docs/rfcs/2026-09-11-metric-export-import.md @@ -193,6 +193,12 @@ Packed Metric exports write manifest version 2 with a required data-layout marke exports can keep version 1. A version-2 snapshot indexes all data in its schema chunks, including ordinary tables stored as standalone Parquet objects. +`export-v2 create --pack-metric-data` selects shared physical scans and +packed output. Omitting the flag keeps the default version-1 export path; +schema-only snapshots also retain version 1. The source frontend or standalone +server must enable `experimental_metric_export = true`. Resume derives the layout +from the manifest, so packed snapshots do not need the flag repeated. + Before creating or replacing a snapshot, the CLI checks the source server's explicit packed-export capability. Before target DDL or data writes, the importer checks the manifest version/layout and the target's packed-import capability. diff --git a/src/cli/src/data/export_v2.rs b/src/cli/src/data/export_v2.rs index 9887e525030..8785bcf3e31 100644 --- a/src/cli/src/data/export_v2.rs +++ b/src/cli/src/data/export_v2.rs @@ -38,10 +38,11 @@ //! --end-time 2025-01-31T23:59:59Z //! ``` //! -//! `--experimental-metric-export` enables shared Metric physical scans for Parquet. +//! `--pack-metric-data` enables shared Metric physical scans for Parquet +//! and writes version-2 packed snapshots. //! Enable `experimental_metric_export = true` on the frontend or standalone server //! and use an endpoint whose frontends all support and enable this option. -//! Add `--metric-data-layout packed` to write version-2 packed snapshots. +//! Resume reads the layout from the existing manifest. //! This requires the server packed-export capability; schema-only exports remain version 1. //! The command checks the server capability before modifying the snapshot. //! Each snapshot path belongs to one export task. Resume requires stable source diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index 48779f3c435..fe38c6b7c7e 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -292,14 +292,11 @@ pub struct ExportCreateCommand { #[clap(long, value_enum, default_value = "parquet")] format: DataFormat, - /// Use shared Metric physical scans (Parquet only). Resume requires that the previous - /// export and storage writes have ended; an HTTP timeout does not establish this. + /// Use shared Metric physical scans and packed objects (Parquet only). + /// Resume uses the manifest layout and requires that previous export and storage + /// writes have ended; an HTTP timeout does not establish this. #[clap(long)] - experimental_metric_export: bool, - - /// Store independent Metric Parquet streams in shared packed objects. - #[clap(long, value_parser = ["packed"], requires = "experimental_metric_export")] - metric_data_layout: Option, + pack_metric_data: bool, /// Delete existing snapshot and recreate. #[clap(long)] @@ -350,10 +347,7 @@ impl ExportCreateCommand { let time_range = TimeRange::parse(self.start_time.as_deref(), self.end_time.as_deref()) .map_err(BoxedError::new)?; - if self.metric_data_layout.is_some() - && time_range.is_bounded() - && time_range.start == time_range.end - { + if self.pack_metric_data && time_range.is_bounded() && time_range.start == time_range.end { return crate::error::InvalidArgumentsSnafu { msg: "Packed export requires --start-time to be earlier than --end-time", } @@ -413,40 +407,18 @@ impl ExportCreateCommand { ); // The filesystem storage constructor can create the snapshot root. - if self.experimental_metric_export { - if self.format != DataFormat::Parquet { - return crate::data::export_v2::error::MetricExportUnavailableSnafu - .fail() - .map_err(BoxedError::new); - } - let capability = database_client - .sql_response( - "SHOW VARIABLES experimental_metric_export", - DEFAULT_SCHEMA_NAME, - ) + if self.pack_metric_data { + require_metric_export_enabled(&database_client, self.format) .await - .context(DatabaseSnafu) .map_err(BoxedError::new)?; - let expected_schema = OutputSchema::new(vec![ColumnSchema::new( - "EXPERIMENTAL_METRIC_EXPORT".to_string(), - "String".to_string(), - )]); - if !matches!(capability.output(), [GreptimeQueryOutput::Records(records)] - if records.schema() == &expected_schema - && records.rows() == &vec![vec![Value::String("true".to_string())]]) - { - return crate::data::export_v2::error::MetricExportUnavailableSnafu - .fail() - .map_err(BoxedError::new); + if !self.schema_only { + database_client + .require_packed_export() + .await + .context(DatabaseSnafu) + .map_err(BoxedError::new)?; } } - if self.metric_data_layout.is_some() && !self.schema_only { - database_client - .require_packed_export() - .await - .context(DatabaseSnafu) - .map_err(BoxedError::new)?; - } let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?; Ok(Box::new(ExportCreate { @@ -455,8 +427,7 @@ impl ExportCreateCommand { schemas, schema_only: self.schema_only, format: self.format, - experimental_metric_export: self.experimental_metric_export, - packed: self.metric_data_layout.is_some(), + packed: self.pack_metric_data, force: self.force, time_range, chunk_time_window: self.chunk_time_window, @@ -484,7 +455,6 @@ struct ExportConfig { schemas: Option>, schema_only: bool, format: DataFormat, - experimental_metric_export: bool, packed: bool, force: bool, time_range: TimeRange, @@ -538,13 +508,11 @@ impl ExportCreate { validate_resume_config(&manifest, &self.config)?; if manifest.is_packed() { + require_metric_export_enabled(&self.database_client, manifest.format).await?; self.database_client .require_packed_export() .await .context(DatabaseSnafu)?; - if !self.config.experimental_metric_export { - return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail(); - } } info!( @@ -573,7 +541,6 @@ impl ExportCreate { storage_config: &self.config.storage_config, parallelism: self.config.parallelism, chunk_parallelism: self.config.chunk_parallelism, - experimental_metric_export: self.config.experimental_metric_export, resume: true, }, progress.as_ref(), @@ -608,7 +575,7 @@ impl ExportCreate { manifest.version = 2; manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); } - if self.config.experimental_metric_export { + if manifest.is_packed() { for chunk in &manifest.chunks { self.storage .prepare_export_chunk(&schema_names, chunk.id, false, manifest.is_packed()) @@ -651,7 +618,6 @@ impl ExportCreate { storage_config: &self.config.storage_config, parallelism: self.config.parallelism, chunk_parallelism: self.config.chunk_parallelism, - experimental_metric_export: self.config.experimental_metric_export, resume: false, }, progress.as_ref(), @@ -761,10 +727,34 @@ impl ExportCreate { } } +async fn require_metric_export_enabled(client: &DatabaseClient, format: DataFormat) -> Result<()> { + if format != DataFormat::Parquet { + return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail(); + } + let capability = client + .sql_response( + "SHOW VARIABLES experimental_metric_export", + DEFAULT_SCHEMA_NAME, + ) + .await + .context(DatabaseSnafu)?; + let expected_schema = OutputSchema::new(vec![ColumnSchema::new( + "EXPERIMENTAL_METRIC_EXPORT".to_string(), + "String".to_string(), + )]); + if !matches!(capability.output(), [GreptimeQueryOutput::Records(records)] + if records.schema() == &expected_schema + && records.rows() == &vec![vec![Value::String("true".to_string())]]) + { + return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail(); + } + Ok(()) +} + fn validate_resume_config(manifest: &Manifest, config: &ExportConfig) -> Result<()> { if config.packed && !manifest.schema_only && !manifest.is_packed() { return ResumeConfigMismatchSnafu { - field: "metric_data_layout", + field: "data_layout", existing: "standalone".to_string(), requested: "packed".to_string(), } @@ -1683,6 +1673,44 @@ mod tests { ); } + #[tokio::test] + async fn test_pack_metric_data_arguments() { + let directory = tempfile::tempdir().unwrap(); + let snapshot = directory.path().join("snapshot"); + let uri = url::Url::from_file_path(&snapshot).unwrap(); + let args = [ + "export-v2-create", + "--addr", + "127.0.0.1:1", + "--to", + uri.as_str(), + ]; + assert!(!ExportCreateCommand::parse_from(args).pack_metric_data); + let packed = + ExportCreateCommand::parse_from(args.into_iter().chain(["--pack-metric-data"])); + assert!(packed.pack_metric_data); + for removed in [ + &["--experimental-metric-export"][..], + &["--metric-data-layout", "packed"][..], + ] { + let error = ExportCreateCommand::try_parse_from( + args.into_iter().chain(removed.iter().copied()), + ) + .unwrap_err(); + assert_eq!(error.kind(), clap::error::ErrorKind::UnknownArgument); + } + for format in ["csv", "json"] { + let command = ExportCreateCommand::parse_from(args.into_iter().chain([ + "--pack-metric-data", + "--format", + format, + ])); + let error = command.build().await.err().unwrap(); + assert!(error.to_string().contains("Metric export requires Parquet")); + assert!(!snapshot.exists()); + } + } + #[test] fn test_chunk_parallelism_parses_valid_value() { let cmd = ExportCreateCommand::parse_from([ @@ -1755,7 +1783,6 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, - experimental_metric_export: false, packed: false, force: false, time_range: TimeRange::unbounded(), @@ -1794,7 +1821,6 @@ mod tests { ]), schema_only: false, format: DataFormat::Parquet, - experimental_metric_export: false, packed: false, force: false, time_range: TimeRange::unbounded(), @@ -1812,7 +1838,7 @@ mod tests { validate_resume_config(&manifest, &config) .unwrap_err() .to_string() - .contains("metric_data_layout") + .contains("data_layout") ); manifest.version = 2; manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); @@ -1840,7 +1866,6 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, - experimental_metric_export: false, packed: false, force: false, time_range, @@ -1875,7 +1900,6 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Csv, - experimental_metric_export: false, packed: false, force: false, time_range: TimeRange::unbounded(), @@ -1912,7 +1936,6 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, - experimental_metric_export: false, packed: false, force: false, time_range: TimeRange::new(Some(start), Some(start)), diff --git a/src/cli/src/data/export_v2/coordinator.rs b/src/cli/src/data/export_v2/coordinator.rs index c6181e51211..e18014af32f 100644 --- a/src/cli/src/data/export_v2/coordinator.rs +++ b/src/cli/src/data/export_v2/coordinator.rs @@ -39,7 +39,6 @@ struct ExportContext<'a> { schemas: Vec, format: DataFormat, parallelism: usize, - experimental_metric_export: bool, packed: bool, resume: bool, } @@ -49,7 +48,6 @@ pub struct ExportDataOptions<'a> { pub storage_config: &'a ObjectStoreConfig, pub parallelism: usize, pub chunk_parallelism: usize, - pub experimental_metric_export: bool, pub resume: bool, } @@ -74,7 +72,6 @@ pub async fn export_data( format: manifest.format, packed: manifest.is_packed(), parallelism: options.parallelism, - experimental_metric_export: options.experimental_metric_export, resume: options.resume, }; @@ -278,7 +275,7 @@ async fn export_chunk( chunk_id: u32, time_range: TimeRange, ) -> Result> { - if context.experimental_metric_export { + if context.packed { context .storage .prepare_export_chunk(&context.schemas, chunk_id, context.resume, context.packed) @@ -290,7 +287,6 @@ async fn export_chunk( format: context.format, time_range, parallelism: context.parallelism, - experimental_metric_export: context.experimental_metric_export, packed: context.packed, }; diff --git a/src/cli/src/data/export_v2/data.rs b/src/cli/src/data/export_v2/data.rs index 53c3899f59e..dac094c0dda 100644 --- a/src/cli/src/data/export_v2/data.rs +++ b/src/cli/src/data/export_v2/data.rs @@ -30,7 +30,6 @@ pub(super) struct CopyOptions { pub(super) format: DataFormat, pub(super) time_range: TimeRange, pub(super) parallelism: usize, - pub(crate) experimental_metric_export: bool, pub(super) packed: bool, } @@ -218,10 +217,8 @@ pub(crate) async fn execute_copy_database_from( fn build_with_options(options: &CopyOptions) -> String { let mut parts = vec![format!("FORMAT='{}'", options.format)]; - if options.experimental_metric_export { - parts.push("experimental_metric_export='true'".to_string()); - } if options.packed { + parts.push("experimental_metric_export='true'".to_string()); parts.push("metric_data_layout='packed'".to_string()); } if let Some(start) = options.time_range.start { @@ -430,6 +427,25 @@ mod tests { use super::*; use crate::common::{PrefixedAzblobConnection, PrefixedGcsConnection, PrefixedOssConnection}; + #[test] + fn test_packed_copy_options() { + let mut options = CopyOptions { + format: DataFormat::Parquet, + time_range: TimeRange::unbounded(), + parallelism: 2, + packed: false, + }; + assert_eq!( + build_with_options(&options), + "FORMAT='parquet', PARALLELISM=2" + ); + options.packed = true; + assert_eq!( + build_with_options(&options), + "FORMAT='parquet', experimental_metric_export='true', metric_data_layout='packed', PARALLELISM=2" + ); + } + #[test] fn test_build_oss_connection_includes_endpoint() { let storage = ObjectStoreConfig { diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 13f9384cb5c..101b2b6f476 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1713,9 +1713,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { uri.as_str(), "--schemas", "public", - "--experimental-metric-export", - "--metric-data-layout", - "packed", + "--pack-metric-data", "--no-proxy", "--start-time", "1970-01-01T00:00:00Z", @@ -1792,7 +1790,6 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { } else { "public,z_later" }, - "--experimental-metric-export", "--no-proxy", "--start-time", "1970-01-01T00:00:00Z", @@ -1805,7 +1802,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { ]; args.extend(storage_args.iter().map(String::as_str)); if packed { - args.extend(["--metric-data-layout", "packed"]); + args.push("--pack-metric-data"); } assert!(run_data_cli(&args).await.is_err()); let before: cli::export_v2::manifest::Manifest = @@ -1838,25 +1835,28 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .await .unwrap() ); - store - .write("data/public/2/unknown.txt", "keep") - .await - .unwrap(); - let copies = faults.copies.load(std::sync::atomic::Ordering::SeqCst); - assert!(run_data_cli(&args).await.is_err()); - assert_eq!( - faults.copies.load(std::sync::atomic::Ordering::SeqCst), - copies - ); - assert_eq!( + if packed { + args.pop(); store - .read("data/public/2/unknown.txt") + .write("data/public/2/unknown.txt", "keep") .await - .unwrap() - .to_vec(), - b"keep" - ); - store.delete("data/public/2/unknown.txt").await.unwrap(); + .unwrap(); + let copies = faults.copies.load(std::sync::atomic::Ordering::SeqCst); + assert!(run_data_cli(&args).await.is_err()); + assert_eq!( + faults.copies.load(std::sync::atomic::Ordering::SeqCst), + copies + ); + assert_eq!( + store + .read("data/public/2/unknown.txt") + .await + .unwrap() + .to_vec(), + b"keep" + ); + store.delete("data/public/2/unknown.txt").await.unwrap(); + } run_data_cli(&args).await.unwrap(); let after: cli::export_v2::manifest::Manifest = serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); @@ -2031,9 +2031,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { "--schemas", "public", "--schema-only", - "--experimental-metric-export", - "--metric-data-layout", - "packed", + "--pack-metric-data", "--no-proxy", "--progress", "never", @@ -2115,7 +2113,7 @@ async fn metric_export_v2_disabled_and_legacy_cli() { "--to", &uri, "--force", - "--experimental-metric-export", + "--pack-metric-data", "--no-proxy", "--progress", "never", @@ -2156,6 +2154,7 @@ async fn metric_export_v2_disabled_and_legacy_cli() { #[tokio::test] async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() { use axum::http::StatusCode; + use common_error::ext::ErrorExt; use serde_json::json; let column = json!({"name": "EXPERIMENTAL_METRIC_EXPORT", "data_type": "String"}); let records = |rows, columns| { @@ -2168,6 +2167,7 @@ async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() let mut responses = [ (StatusCode::BAD_REQUEST, json!([])), (StatusCode::OK, json!([])), + (StatusCode::OK, json!([["false"]])), (StatusCode::OK, json!([[true]])), (StatusCode::OK, json!([["true", "extra"]])), (StatusCode::OK, json!([["true"], ["true"]])), @@ -2205,9 +2205,11 @@ async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() )]), ), ]); + responses.push((StatusCode::OK, json!([valid.clone()]))); for (status, output) in responses { - let app = - axum::Router::new().route( + let gate_enabled = output == json!([valid.clone()]); + let app = axum::Router::new() + .route( "/v1/sql", axum::routing::post( move |axum::Form(form): axum::Form< @@ -2223,6 +2225,10 @@ async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() ) }, ), + ) + .route( + "/v1/capabilities", + axum::routing::get(|| async { axum::Json(json!({})) }), ); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap().to_string(); @@ -2237,14 +2243,17 @@ async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() &addr, "--to", &uri, - "--experimental-metric-export", + "--pack-metric-data", "--force", "--no-proxy", "--progress", "never", ]; let error = run_data_cli(&args).await.unwrap_err(); - if status == StatusCode::OK { + let cause = error.root_cause().unwrap_or(&error).to_string(); + if gate_enabled { + assert!(cause.contains("metric_packed_export"), "{error:?}"); + } else if status == StatusCode::OK { assert!( error.to_string().contains("Metric export requires"), "{error}" @@ -2255,7 +2264,40 @@ async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() let manifest = snapshot.join("manifest.json"); std::fs::write(&manifest, b"preserve").unwrap(); assert!(run_data_cli(&args).await.is_err()); - assert_eq!(std::fs::read(manifest).unwrap(), b"preserve"); + assert_eq!(std::fs::read(&manifest).unwrap(), b"preserve"); + + let mut packed = cli::export_v2::manifest::Manifest::new_for_export( + "greptime".into(), + vec!["public".into()], + false, + cli::export_v2::manifest::TimeRange::unbounded(), + cli::export_v2::manifest::DataFormat::Parquet, + None, + ) + .unwrap(); + packed.version = 2; + packed.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); + packed.validate_layout().unwrap(); + let original = serde_json::to_vec(&packed).unwrap(); + std::fs::write(&manifest, &original).unwrap(); + let chunk = snapshot.join("data/public/1"); + std::fs::create_dir_all(&chunk).unwrap(); + let part = chunk.join("pack-000000.bin"); + std::fs::write(&part, b"unfinished").unwrap(); + let resume = args + .into_iter() + .filter(|arg| !matches!(*arg, "--force" | "--pack-metric-data")) + .collect::>(); + let resume_error = run_data_cli(&resume).await.unwrap_err(); + assert_eq!( + resume_error + .root_cause() + .unwrap_or(&resume_error) + .to_string(), + cause + ); + assert_eq!(std::fs::read(&manifest).unwrap(), original); + assert_eq!(std::fs::read(&part).unwrap(), b"unfinished"); server.abort(); } } From 7c32c38ba33f488cf08196bb7311f464958c4115 Mon Sep 17 00:00:00 2001 From: liyang Date: Mon, 12 Oct 2026 10:17:15 +0000 Subject: [PATCH 15/46] feat(cmd): add `--http-api-server-addr` CLI flags (#9512) * feat(cmd): add --enable-api-server and --api-server-addr CLI flags Signed-off-by: liyang * apply suggestion Signed-off-by: liyang --------- Signed-off-by: liyang --- src/cmd/src/frontend.rs | 9 +++++++++ src/cmd/src/standalone.rs | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index b000f7e1f02..83c16ffc2dc 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -212,6 +212,10 @@ pub struct StartCommand { internal_grpc_server_addr: Option, #[clap(long)] http_addr: Option, + /// The address to bind the dedicated public HTTP API server. Setting this + /// flag also enables the API server. + #[clap(long)] + http_api_server_addr: Option, #[clap(long)] http_timeout: Option, #[clap(long)] @@ -301,6 +305,11 @@ impl StartCommand { opts.http.disable_dashboard = disable_dashboard; } + if let Some(addr) = &self.http_api_server_addr { + opts.http.api_server_addr.clone_from(addr); + opts.http.enable_api_server = true; + } + if let Some(addr) = &self.grpc_bind_addr { opts.grpc.bind_addr.clone_from(addr); opts.grpc.tls = merge_tls_option(&opts.grpc.tls, tls_opts.clone()); diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 76c95c26d32..ca0f3b74057 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -328,6 +328,10 @@ impl App for Instance { pub struct StartCommand { #[clap(long)] http_addr: Option, + /// The address to bind the dedicated public HTTP API server. Setting this + /// flag also enables the API server. + #[clap(long)] + http_api_server_addr: Option, #[clap(long = "grpc-bind-addr", alias = "rpc-bind-addr", alias = "rpc-addr")] grpc_bind_addr: Option, #[clap(long)] @@ -418,6 +422,11 @@ impl StartCommand { opts.http.addr.clone_from(addr); } + if let Some(addr) = &self.http_api_server_addr { + opts.http.api_server_addr.clone_from(addr); + opts.http.enable_api_server = true; + } + if let Some(data_home) = &self.data_home { opts.storage.data_home.clone_from(data_home); } From 70a0c826ca289e72795523816771b5acd24af845 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 12 Oct 2026 11:16:53 +0000 Subject: [PATCH 16/46] docs(config): qualify when the object store WAL batch seals on size (#9538) The `wal.flush_interval` comment claimed that under sustained load the batch always seals on `max_batch_bytes` and the interval no longer matters. Sustained-load runs on S3 sealed on the timer in every case because the admitted bytes did not reach `max_batch_bytes` within one interval. State the actual condition: the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next timer tick; otherwise it seals on the timer, whatever the load. Regenerated `config/config.md` from the example TOMLs. Signed-off-by: jeremyhi --- config/config.md | 4 ++-- config/datanode.example.toml | 2 +- config/standalone.example.toml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/config/config.md b/config/config.md index 612a5cd43c0..5d7555ad4b0 100644 --- a/config/config.md +++ b/config/config.md @@ -127,7 +127,7 @@ | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `wal.storage_provider` | String | `""` | The name of the storage provider that holds the WAL objects, an empty name selects the default object store.
**It's only used when the provider is `experimental_object_store`**. | | `wal.prefix` | String | `wal` | The path prefix of the WAL objects inside the storage provider.
The objects are written under `/datanodes//epochs/`, which is derived from this prefix.
**It's only used when the provider is `experimental_object_store`**. | -| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next timer tick; otherwise it seals on the timer, whatever the load.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | | `wal.ack_mode` | String | `durable` | When an append to the object store WAL returns.
- `durable`: an append returns after the object holding its entries is durable (the default).
- `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog.
**It's only used when the provider is `experimental_object_store`**. | | `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | @@ -606,7 +606,7 @@ | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `wal.storage_provider` | String | `""` | The name of the storage provider that holds the WAL objects, an empty name selects the default object store.
**It's only used when the provider is `experimental_object_store`**. | | `wal.prefix` | String | `wal` | The path prefix of the WAL objects inside the storage provider.
The objects are written under `/datanodes//epochs/`, which is derived from this prefix.
**It's only used when the provider is `experimental_object_store`**. | -| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next timer tick; otherwise it seals on the timer, whatever the load.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | | `wal.ack_mode` | String | `durable` | When an append to the object store WAL returns.
- `durable`: an append returns after the object holding its entries is durable (the default).
- `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog.
**It's only used when the provider is `experimental_object_store`**. | | `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 413afb609c6..a2586908b70 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -240,7 +240,7 @@ overwrite_entry_start_id = false #+ prefix = "wal" ## The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`. -## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters. +## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next timer tick; otherwise it seals on the timer, whatever the load. ## When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission. ## **It's only used when the provider is `experimental_object_store`**. #+ flush_interval = "100ms" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index c656049c5d2..0da3ad9031c 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -427,7 +427,7 @@ overwrite_entry_start_id = false #+ prefix = "wal" ## The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`. -## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters. +## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next timer tick; otherwise it seals on the timer, whatever the load. ## When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission. ## **It's only used when the provider is `experimental_object_store`**. #+ flush_interval = "100ms" From 4ef312ecf204bce3fc7719a0e5631c2c31e6774f Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 12 Oct 2026 12:00:55 +0000 Subject: [PATCH 17/46] docs: add the object store WAL RFC (#9195) * docs: add the object store WAL RFC Propose an object store backed write-ahead log as a WAL provider next to Raft Engine and Kafka: a datanode batches the entries of its regions into immutable objects under one prefix, creates each object conditionally under a monotonically increasing sequence, acknowledges a write once its object is durable by default, and replays each region from its own segments after listing the prefix. A reference implementation exists and follows this RFC in a series of pull requests; the document describes it and marks what is proposed on top of it. Signed-off-by: jeremyhi * docs: link the object store WAL RFC to its tracking issue Signed-off-by: jeremyhi * docs: describe the predecessor chain and instance epoch of the object store WAL Recovery replays only the chain of objects that ends at the latest complete object, and every open starts a new epoch with an empty start object, so an object whose create was reported as failed is never replayed once a later append has been acknowledged. Update the object format, recovery, writing, the failure matrix, the garbage collection constraints, compatibility, evidence, drawbacks and the cluster mode alignment accordingly. Signed-off-by: jeremyhi * docs: identify object store WAL objects by sequence and epoch Signed-off-by: jeremyhi * docs: derive the object store WAL epoch from the start object sequence Signed-off-by: jeremyhi * docs: separate the retryable and the conflicting start object outcomes Signed-off-by: jeremyhi * docs: describe garbage collection of the object store WAL through the reclaim boundary Signed-off-by: jeremyhi * docs: judge vanished objects after chain selection and fix the evidence references Signed-off-by: jeremyhi * docs: record the stop-time conflict and the completed fault matrix in the object store WAL RFC Signed-off-by: jeremyhi * docs: limit the stop-time conflict to objects of the same or a later epoch Signed-off-by: jeremyhi * docs: keep the conflicting-content condition in the stop-time conflict sentence Signed-off-by: jeremyhi * docs: report the S3 write path measurements of the object store WAL Signed-off-by: jeremyhi * docs: qualify the seal timer findings and link the S3 runs Signed-off-by: jeremyhi * docs: state the serial durable cycle behind the one-batch-per-interval observation Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- docs/rfcs/2026-09-06-object-store-wal.md | 352 +++++++++++++++++++++++ 1 file changed, 352 insertions(+) create mode 100644 docs/rfcs/2026-09-06-object-store-wal.md diff --git a/docs/rfcs/2026-09-06-object-store-wal.md b/docs/rfcs/2026-09-06-object-store-wal.md new file mode 100644 index 00000000000..26970d166f8 --- /dev/null +++ b/docs/rfcs/2026-09-06-object-store-wal.md @@ -0,0 +1,352 @@ +--- +Feature Name: Object Store WAL +Tracking Issue: https://github.com/GreptimeTeam/greptimedb/issues/9197 +Date: 2026-09-06 +Author: jeremyhi +--- + +# Summary + +Add an object store backed write-ahead log as a first-class WAL provider, alongside Raft Engine and Kafka. A datanode batches the entries of all its regions into immutable objects under one prefix, creates each object conditionally under a monotonically increasing sequence number, and by default acknowledges a write only after the object holding it is durable; an opt-in acknowledgement mode returns on admission with a stated loss bound. Recovery lists the objects, selects the chain of objects that ends at the latest complete one, rebuilds an in-memory catalog from their footers, and replays each region from its own segments. The provider is standalone-only and marked experimental in this RFC; the persisted metadata and object format are designed so that lifting either restriction later does not require a migration. + +A reference implementation of everything this document describes as implemented exists and was validated as reported under *Evidence*; it is submitted in a series of pull requests that follow this RFC, each covering one part of the design in the order the *Details* sections give. The document keeps two things apart: what that implementation does, and what is proposed on top of it. Every passage that describes proposed behaviour is introduced with the word *Proposed*; everything else describes the reference implementation. The *Future work* and *Unresolved questions* sections are proposed by definition and carry no marker. + +# Motivation + +GreptimeDB already stores data files in object storage. The WAL is the last component that still needs either local disk with durability guarantees (Raft Engine) or an external cluster (Kafka). Both are awkward for a cloud-native standalone deployment: + +- Raft Engine ties the node to a persistent volume. Losing the volume loses unflushed writes; moving the node means moving the volume. +- Kafka is an operational dependency of its own, with topics, retention and a broker fleet to size, and it is disproportionate for a single node. + +Object storage is already provisioned, already durable across zones, and already the place the rest of the data lives. An object store WAL takes the WAL out of the local state that must survive a restart: unflushed writes are no longer tied to a volume, and a node restarted against the same bucket and prefix replays them. It does not yet make a standalone node independent of local disk altogether, because standalone keeps its metadata in a local key-value store under the data home regardless of the WAL provider; moving that metadata off local disk is a separate change. The cost is write latency under the default durable acknowledgement: a write waits for its batch to seal, for every earlier object to be durable, and for one PUT. That is acceptable for the ingestion patterns GreptimeDB targets when writes are batched; *Write path measurements* under *Evidence* shows where the cost lies. + +# Goals + +- A WAL provider that needs nothing but an object store and a prefix. +- Exactly-once replay of acknowledged writes across graceful and ungraceful restarts under the default acknowledgement mode. The `enqueued` mode is an explicit choice whose loss is bounded to the unpersisted backlog and reported. +- Region isolation: many regions share objects, but a region only ever replays its own entries. +- Recovery that never silently loses data. Structural corruption fails recovery. A corrupted segment is skipped at segment granularity by default, and the skip is recorded for its region, counted and logged rather than hidden; `on_corrupted_segment = "fail"` fails the read instead. +- Default behaviour of the existing providers unchanged; the new provider is opt-in through configuration. + +# Non-goals + +- Distributed mode. The provider is rejected when a datanode is controlled by a metasrv; the metasrv side has no allocation logic for it yet. +- Persistence of per-region obsolete watermarks. Mito re-establishes every watermark when it opens a region, which is what garbage collection reasons from; persisting them is a follow-up. +- Matching Raft Engine's write latency. The provider trades latency for the absence of local durable WAL state. +- Making standalone metadata independent of local disk. The metadata store is unchanged by this RFC. + +# Details + +## Identity + +The provider is added to the four parallel WAL enums so that every resolution path handles it explicitly: + +| Enum | Variant | Notes | +| --- | --- | --- | +| `DatanodeWalConfig` | `ObjectStore(ObjectStoreWalConfig)` | serde tag `experimental_object_store`; fields `storage_provider`, `prefix`, `flush_interval`, `max_batch_bytes`, `ack_mode`, `max_unpersisted_bytes`, `max_unpersisted_age`, `on_corrupted_segment` | +| `WalOptions` | `ObjectStore(ObjectStoreWalOptions { prefix })` | persisted per region with the stable tag `object_store`; flat key `wal.object_store.prefix`; the value is the derived node prefix | +| `WalProvider` | `ObjectStore { prefix }` | allocates the per-region options in standalone under the derived node prefix; `start` is a no-op | +| `Provider` | `ObjectStore(Arc)` | `is_remote_wal()` is true | + +Two choices are deliberate: + +- The persisted `WalOptions` carry the derived node prefix `/datanodes//epochs/` described under *Object format*, not the configured root. On reopen the log store compares the persisted prefix with the one it runs under and fails fast on a mismatch, so an operator who changes the prefix or the node id cannot silently point a region at an empty WAL. Changing the prefix of an existing region requires an explicit migration, which is out of scope here. +- `Provider::ObjectStore` is classified as a remote WAL. Mito's remote-WAL branches (initial flushed entry id from `latest_entry_id`, `topic_latest_entry_id` after replay, region-scoped reads, no topic grouping on batch open) apply unchanged. This is a classification of replay semantics, not of physical location. + +A datanode that is configured with a metasrv client rejects the provider before any store is built: `DatanodeBuilder` returns an error stating that the object store WAL is standalone-only. Separately, `DatanodeWalConfig::ObjectStore` cannot be converted into `MetasrvWalConfig`, so a metasrv cannot be configured with it either. The standalone bootstrap matches the datanode WAL configuration before that conversion and constructs `WalProvider::ObjectStore` directly. + +## Object format + +Every object holds one sealed batch: + +```text +header | segment(region 1) | ... | segment(region N) | footer | trailer +``` + +- The header carries the magic `GTWALOBJ`, format version 1, the object sequence number, the epoch of the store instance that wrote the object, a link to the object this one extends, as that object's sequence number and epoch or none for the first object of a chain, the reclaim boundary described under *Entry ids*, and a CRC32 of the preceding header bytes; the header is 54 bytes. A reclaim boundary above the object's own sequence fails the header check. A sequence number and an epoch together identify an object: an instance never writes two objects under one sequence, and instances have distinct epochs, as described under *Log store*. Recovery selects the history it replays by the links and the epochs, as described under *Log store*. +- An object may hold no segment. A store writes such an object when it opens, as described under *Log store*; it is a chain object like any other, and its segments are trivially all flushed. +- Each segment holds the entries of exactly one region, ordered by entry id; segments are ordered by region id. Ordering makes the encoding deterministic, which the identical-retry rule below depends on. +- The footer records, per segment, the region id, the entry id range, the entry count, the byte range and a CRC32 of the segment. +- The fixed-size trailer carries the magic `GTWALTRL`, the footer location, a CRC32 of the footer and a CRC32 of the whole object. A reader locates the footer by reading the trailer at the end of the object and can decode a single region's segment without touching the others. + +Corruption is handled at two levels; nothing else is checked in production. Structural damage fails recovery: recovery reads only the header, the trailer and the footer of every object, as described under *Log store*, and verifies the header and trailer magic, the format version, the header CRC, the header sequence against the key, the footer CRC and that the segment byte ranges tile the object body exactly; the first object that fails any of these checks fails construction of the store, since nothing can be trusted about such an object. A segment is verified by the read that decodes it, and a segment that does not decode, because its CRC does not match or its content disagrees with its footer entry, is a data-level failure confined to one region. The read fetches the segment once more, since the download rather than the object may be damaged, and decodes it again; if it still does not decode, the behaviour follows `on_corrupted_segment`. `skip`, the default, skips the segment, records a WAL hole for the region with the object key, the object sequence and the entry id range the footer holds for the segment, increments the counter `greptime_logstore_object_store_wal_skipped_segments_total` and logs a warning with those fields, and the read continues with the next object; every other region in the object is unaffected. The holes are kept in memory per region and can be listed from the store; nothing is persisted, so a store opened later on the same prefix records a hole again when a read meets the segment. `fail` fails the read with the decode error, so a region whose replay meets the segment does not open. Skipping loses only the unflushed entries of that region in that object: nothing already flushed is affected, and the loss is reported rather than hidden. A failed fetch is not corruption and fails the read in both modes, unless garbage collection deleted the object, as described under *Entry ids*, in which case the object is skipped; an object the catalog still holds once a delete of it has settled fails the read whether it is missing or unreadable, and a segment of such an object that does not decode follows `on_corrupted_segment` as it would without any collection. The whole-object CRC is not a recovery requirement; it stays in the trailer as a fixture-level check. Because recovery never checks it, the header carries its own CRC, so a damaged link fails recovery instead of silently cutting acknowledged objects off the chain. The format is frozen at version 1. + +The operator configures only the prefix root. The store runs under a node prefix derived from it, `/datanodes//epochs/`, and object keys are `/objects/.wal`, where the sequence is zero-padded to 20 decimal digits. Standalone uses its datanode id, zero unless configured, and generation zero, so with prefix `wal` the first object is `wal/datanodes/0/epochs/0/objects/00000000000000000000.wal`. Every place that names the prefix, whether to open the store or to allocate the WAL options of a region, derives it from the same function, so the prefix a region persists matches the one its store runs under. Listing accepts only well-formed keys under `/objects/` and returns them in sequence order. This is the standalone form of the key layout under *Alignment with cluster mode*; the cluster assignment of node id and generation is proposed there. + +## Catalog + +The catalog is an in-memory index rebuilt at recovery from the objects on the replayed chain: objects by sequence number, and per region the objects that hold its entries with their entry id ranges. Its invariants are enforced on insertion: a region's entry id ranges are strictly increasing along the object sequence, an object holds at most one segment per region, and an already indexed sequence number is rejected, whether or not the footer matches; an object without segments is indexed like any other and takes its sequence. The next sequence number is the larger of the largest indexed one plus one and the floor that any region's largest entry id sets, as described under *Entry ids*, checked against the sequence limit; construction raises it above every listed object, as described under *Log store*, and gaps in the sequence are tolerated. + +## Conditional create + +Objects are created with an if-not-exists put. Three outcomes matter: + +| Outcome | Meaning | +| --- | --- | +| created | the batch is durable | +| exists with identical content | a retry of a batch that was durable before the client observed it; accepted | +| exists with different content | another object owns this sequence; its epoch decides, as described under *Log store* | + +Deterministic encoding is what makes the middle row decidable. + +## Entry ids + +Entry ids are object-sequence-major. The high bits are the sequence number of the object that holds the entry, the low bits are the entry's position among that region's entries inside the object: + +```text +entry_id = (object_seq << 20) | position (1 <= position < 2^20, object_seq < 2^44) +``` + +Positions start at one, so entry id zero is never assigned: Mito treats zero as the watermark of a region that has no durable entry and replays from the next id, and the first object has sequence zero. Ids are assigned at admission, from the sequence the open batch takes when it is sealed. A store never hands out a sequence twice, so a batch that fails is never written again under its sequence, and a retry of its entries is a new batch with new ids. A batch that would take a region past the position range is sealed first and the entries open the next batch; an append that alone exceeds the range is rejected. This keeps two properties at once. `entry_id >> 20` names the object that holds an entry without consulting any index, which is what a future cross-node takeover needs to describe a WAL position as a range of object sequences. And every entry has a unique, strictly increasing id within its region, so a flush that lands between two appends of the same region that ended up in the same object records exactly which entries it covered; replay from `flushed_entry_id` can never skip an unflushed entry or repeat a flushed one. A region's ids have gaps wherever other regions or other positions took the sequence, exactly as Kafka offsets do, and Mito already tolerates gaps; the recovered history of a region may also lack whole objects where a batch was never acknowledged, as described under *Writing*. The id values are encoded in two places that must agree: in each record inside its segment and in the footer's entry id range for that segment. + +`latest_entry_id(provider)` is region-scoped: the highest entry id durable under the prefix for the provider's region, or zero when the region has none. A freshly created region therefore starts at zero even if other regions under the prefix have entries. This differs from Kafka, whose high watermark is per topic; it is sufficient here because ids are per region, and Mito consumes the value only for the region that asked. + +Two rules keep every new id of a region above every id it already has, on a prefix that already holds objects and across the watermark Mito hands the store. On open, the next object sequence is the larger of the largest listed sequence plus one, orphans included, and, over every region, the object named by the region's largest existing entry id plus one, that is `(max_entry_id >> 20) + 1`, so that the catalog's strictly increasing ranges hold; sequence gaps are tolerated, and a sequence at or above `2^44` is rejected as exhausted. Objects written under an earlier contiguous scheme are not rewritten: they remain readable, their ids carry no object information, and the rule places them by their high bits just the same. + +The second rule is the sequence floor of `obsolete`. Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, and the actor applies the watermark and a floor derived from that entry id together: the next object sequence is raised to `(entry_id >> 20) + 1` unless it is there already, while zero, the watermark of a region without entries, names no object and needs no floor, so that a watermark assigned under another prefix, or one whose object this prefix no longer holds, is never passed by a new id. The floor moves only the next sequence, never the chain: the next batch still links to its predecessor, so a skipped sequence leaves no gap in the history recovery replays. While the open batch holds ids under the next sequence, the call is refused with a retryable error and records neither the watermark nor the floor, because those ids would fall below the floor. In the `enqueued` mode an id the store itself handed out needs no floor, even before it is durable: it is never handed out again while the store runs, so every later id of the region is greater. In the current single-prefix layout the floor never needs to move, since every watermark Mito hands the store names an object below the next sequence; the rule exists for the cross-prefix cases under *Alignment with cluster mode*. No additional persisted state is introduced. + +Garbage collection deletes the objects below a reclaim boundary that the store persists in every header, so that recovery can tell a collected predecessor from one that never landed. The boundary is computed from the catalog and the watermarks the store already keeps: each region with a recorded watermark contributes the sequence of its first indexed object holding a segment whose largest entry id is above the watermark; a region without a recorded watermark contributes its first indexed object; a region whose watermark is the maximum entry id, after `obsolete_all`, contributes nothing; the boundary is the smallest contribution, or the last indexed sequence plus one when nothing contributes, and it never decreases. Whether a segment is flushed is decided per segment from the footer, on the full entry id, so a watermark that stops in the middle of an object keeps the object. The boundary is recomputed after every successful `obsolete`, after `obsolete_all` and after every object is indexed; Mito calls `obsolete` after a flush, after a truncate, when it drops a region and when it opens one, so a restart re-establishes every watermark and the boundary advances on the same rule. Every sealed batch carries the boundary current at sealing and the start object carries the boundary recovery found, so a boundary is always at or below the sequence of the object carrying it. The boundary that permits deletion is the one carried by the last indexed object; the boundary of a batch whose create failed or is in flight permits nothing. + +Every present object with a sequence below the durable boundary is deletable, on the chain or not: the objects the catalog indexes, the objects recovery listed but did not index, the objects of an earlier epoch that the start object moved past, and the sequences of batches that failed in this run, whose objects may or may not exist. A collection starts after a successful `obsolete`, after `obsolete_all` and after an indexing that raised the durable boundary. Deletes run in the background, at most four at a time and in no particular order, and never hold up admission, sealing, uploads or acknowledgements; an indexed object is unindexed only after its delete succeeded, and the WAL holes recorded for it are dropped with it. Every deleted object is counted in `greptime_logstore_object_store_wal_deleted_objects_total`; a failed delete is logged, counted in `greptime_logstore_object_store_wal_failed_deletes_total` and retried at the next collection, and the object stays indexed meanwhile. An object deleted but not yet unindexed when the process crashes is simply absent at the next recovery, which lists the prefix. Nothing is collected once `stop` has begun or while the store is poisoned; a delete in flight when `stop` begins runs to completion before `stop` returns, and the actor settles every attempt it abandons when it exits. + +A read never needs a deleted object, since entries at or below a region's watermark are hidden anyway, but a read that listed an object before it was collected skips the object when it gets there. A read that meets an error on such an object makes two observations, in the reverse of the order the actor writes them: it looks at the set of objects whose delete is in flight, and then at the catalog. The actor registers an object in that set before its delete starts, unindexes it if the delete succeeded, and only then takes it out, so an object absent from the set was either never collected, in which case the catalog still holds it, or is already unindexed by the time the catalog is read. Observing the catalog first would let a collection complete between the two observations and leave the read with an object that still looked indexed and no longer looked deleting. An object that is still in the set has an undecided delete, so the read waits for that one delete, a single object store request, rather than assume it succeeds: a delete that fails leaves the object present and indexed, and the error the read met is then reported as any other, `on_corrupted_segment` included. What the read waits for is the delete attempt its observation found, taken together with the lookup, not the object: a retry of the same object is a separate attempt, so the attempt the read holds releases it even while that retry is in flight, and a read is never handed from one attempt to the next. A read stream outlives the store it came from, so the actor settles every attempt it abandons when it exits, whether the store was stopped or dropped; the objects of those attempts stay present and indexed, so the reads waiting for them report the error they met. A collected object is skipped silently: it is not recorded as a hole and does not count as a skipped segment. An object the catalog still holds after those observations fails the read like any other fetch failure, `on_corrupted_segment` included. Within one run the catalog keeps the largest entry id of a region after the objects holding it were collected, so `latest_entry_id`, the durability wait and the `enqueued` clamp keep their reference point; a store opened later learns only what is listed, so after a restart `latest_entry_id` of a region whose later objects were collected is the largest id among the objects left, which only the pruning hint of Mito consumes. No state is persisted outside the headers and no configuration option is added. + +## Log store + +`ObjectStoreLogStore` implements the `LogStore` trait for `Provider::ObjectStore`. + +**Construction and recovery.** The prefix, flush interval, batch size and backlog thresholds are validated, then the store lists the prefix and fetches, for up to eight objects at a time, the header and a window at the end of each object with range reads; a short object is read whole, and the footer is read separately only when it starts before the window. Each header and footer is verified as described under *Object format*; the first corrupted object fails construction; a fetch that finds the object gone is judged once the history is selected, as described below, and any other failed fetch abandons the rest. Segments are not read, so recovery costs a few small reads per object however large the objects are. + +The store then selects the history it replays from the headers. An object's link holds when the object has no predecessor, when its predecessor is present and carries the epoch the link records, or when the predecessor's sequence is below the reclaim boundary the object carries, in which case the chain ends there whether or not the predecessor is present; above that boundary a missing predecessor breaks the link. An object is complete when every link on its chain holds, each judged with the boundary of the object carrying the link. The tip is the complete object with the largest epoch and, among those, the largest sequence. Only the objects on the tip's chain are indexed and replayed, in sequence order whatever the order the fetches completed in, so the catalog checks the entry ranges of every object against its predecessors like a sequential replay, and a conflicting object fails construction. Every other present object is an orphan: it is neither indexed nor read, but its sequence counts, so the next sequence is above every listed object, and the objects below the tip's boundary are collected. A listed object whose header or footer is gone by the time it is fetched is ignored when it lies below the tip's boundary and fails construction when it lies at or above it, since collection never deletes one there. Construction also fails when a listed object carries an epoch above the next sequence, which no instance can have produced, since an epoch is a start sequence plus one. When objects are present but none is complete, construction fails rather than start an empty history. Construction then records the largest durable entry id per region. + +Before it admits any append, the store writes a start object: an object without segments at the next sequence, which is above every listed object, orphans included, linked to the tip, or to nothing on an empty prefix. The instance's epoch is the sequence of its start object plus one, and every object the instance writes carries it. The conditional create assigns each sequence to exactly one object, so no two instances share an epoch; and the start sequence is above every object present at open, including the start objects of every earlier instance, so the epoch of every later instance is greater than that of any object an earlier instance ever writes, including one whose create lands after that instance stopped or crashed, and such an object never becomes the tip. Only once the start object is created does the store accept appends, and the first batch links to it. If the sequence turns out to be taken, the store reads the object there: an object of an earlier epoch is stale, and the start object moves to the next sequence, taking the epoch that sequence gives; an object of the same or a later epoch means another writer, and construction fails with a conflict error that is not retried; an object with identical content, which another open of the same prefix has just written, fails construction with a retryable error, so that the next attempt lists it and starts above it. A create whose outcome is unknown fails construction with a retryable error too: an object it may have left is listed by the next attempt, which therefore starts above it. Two assumptions remain. Listing the prefix is strongly consistent, so an open sees every start object that was durable before it listed; this holds for S3, GCS, Azure Blob Storage and MinIO, and an S3-compatible store without it is not supported. And one process at a time opens a prefix: two processes opening it concurrently either contend for the same start sequence, where the conditional create lets one through and fails the other's construction, or, if one of them listed a late object the other did not, start under different sequences and epochs and both proceed until their sequences meet, where the writer that meets the other's object follows the conflict rules under *Writing* and the *Failure matrix*: in the `durable` mode an object of an earlier epoch fails its batch and one of a later epoch poisons it, and in the `enqueued` mode any conflict poisons it, so which writer is poisoned depends on who observes the collision, not on the epochs alone; the store does not fence concurrent writers beyond that, and in cluster mode the generation in the key prefix does. + +**Writing.** One background actor per store admits appended entries into a single open batch and assigns entry ids as described under *Entry ids* above. The batch is sealed when its estimated size reaches `max_batch_bytes` or when `flush_interval` elapses with entries in it; a timer tick on an empty batch creates nothing. Sealing encodes the batch under the next sequence number and hands it to a pipelined uploader: the actor keeps admitting entries into the next batch while uploads are in flight, and up to four conditional creates run at a time, started in sequence order. Every sealed batch is encoded with the store's epoch and a link to its pipeline predecessor: the last batch sealed before it that is not indexed yet, or, when there is none, the last indexed object, which after opening is the start object. A link to the last indexed object while an earlier batch is still uploading would let a later chain skip that batch once it is acknowledged. Batches are indexed in sequence order in both acknowledgement modes: a created object is not indexed until every earlier batch is. The rest of this paragraph is `durable` mode behaviour; the `enqueued` mode, whose appends have already returned, is described under *Acknowledgement*. In the `durable` mode a batch is acknowledged when it is indexed, so a batch whose object is durable is not acknowledged until every earlier batch is durable too, and the acknowledged prefix of a region's history never has a missing predecessor. When a create fails transiently, every batch that is not indexed fails at once, including batches whose objects were already created and batches whose creates are still in flight. Their sequence numbers are never handed out again, the next batch links to the last indexed object, and the store keeps serving. A create still in flight completes later and finds no batch to acknowledge, so whatever it stored is an orphan. The next object acknowledged links past every failed batch, and from then on recovery replays none of them, whether their objects exist, land later or never do. A create that meets a different object reads that object's header: an object of an earlier epoch is stale and can never be the tip, so the batch fails as for a transient error and the store keeps serving; an object of the same or a later epoch belongs to another writer on the prefix, and the store poisons itself. + +The guarantee is therefore about acknowledged history. An acknowledged entry is always replayed; an entry whose append failed is replayed only when nothing was acknowledged after it, the same outcome as a crash between an object's creation and its acknowledgement, which the engine already tolerates. The distinction matters because Mito does not advance a region's committed sequence when an append fails and reuses those row sequences in its next write, so replaying a failed batch next to a later acknowledged one could let the unacknowledged rows win deduplication. + +**Acknowledgement.** `ack_mode` selects when `append_batch` returns. `durable`, the default, returns only after the object holding the entries is durable and indexed, so callers never see an entry id for an entry that is not durable; a write-ahead log is expected to mean durability on return, which is why this stays the default. `enqueued` returns as soon as the entries are admitted, with their entry ids already assigned, and creates the object in the background. In that mode the recovery point objective is the unpersisted backlog. `max_unpersisted_bytes` and `max_unpersisted_age` are admission thresholds: when the open batch plus the sealed batches that are not durable yet reach either, new appends are held back in arrival order until an upload completes, and the store seals the open batch if no create is in flight so that there is an upload to wait for; nothing is dropped and the backlog cannot grow without limit. They are not a hard bound on what a crash can lose: during an outage the writes already acknowledged stay in memory for as long as the outage lasts, however old they get, and a crash in that window loses them. + +The failure contract of `enqueued` differs from the durable one because no caller remains to retry. A transient create failure is repeated by the store after a short delay instead of failing the batch. A permanent failure, meaning a conflicting object or an encoding or catalog error, poisons the store. Once `stop` has begun a transient failure is no longer repeated: the backlog is dropped and `stop` reports the loss as its error, since those entries were acknowledged but never became durable. Otherwise `stop` seals and uploads the remaining backlog before it returns. + +The mode also needs a barrier that the durable mode gets for free. An entry id that Mito has been handed may not yet be durable, and a flush that persisted such an id as the manifest's `flushed_entry_id` would, after a crash before the object was created, leave a watermark that refers to an object which never existed, so that the sequence is reused on restart and later entries land below the watermark and are skipped by replay. The `LogStore` trait therefore has `wait_durable(provider, entry_id)`, with a default that returns at once, so that every existing provider keeps its own acknowledgement and persistence semantics, whatever they are; the durability guarantee this section describes belongs to the object store provider alone. The object store WAL answers it once the region is durable and indexed through `entry_id`; an id the store never handed out has nothing to wait for. Mito waits before it publishes a manifest watermark: a flush waits after its SSTs are written and before the manifest edit, a full truncate waits before it records the truncated entry id as the replay frontier, and discarding unflushed data waits before it records the discarded entry id as flushed. The wait in a flush observes the flush's cancellation, so a cancelled flush does not hold up the drop or truncate queued behind it; a full truncate and a discard of unflushed data still run their own wait before publishing their frontier. After a backlog was lost, a durability wait for an entry that is not durable fails rather than certifying it, whether the entry was in that backlog or not. In `enqueued` mode `obsolete` also clamps the watermark to the region's durable entry id, since an entry that is not durable yet is replayed after a crash and hiding it would skip it. A manifest watermark can then only ever name a durable entry, and the restart allocation rule under *Entry ids* keeps every new id above it. + +**Failure matrix.** + +| Situation | Sequence | Waiters of the batch | Store | +| --- | --- | --- | --- | +| create succeeds, or identical retry | advances | acknowledged in sequence order | healthy | +| transient I/O error, `durable` mode | consumed and never reused, nor are the sequences of the other failed batches | this and every batch that is not indexed fail, including those whose objects were created; a create still in flight completes as an orphan | healthy: the next batch links to the last indexed object | +| an object of an earlier epoch holds the sequence, `durable` mode | consumed and never reused | as for a transient I/O error | healthy | +| transient I/O error, `enqueued` mode | unchanged | already acknowledged | healthy: the create is repeated after a short delay | +| transient I/O error, `enqueued` mode after `stop` began | unchanged | already acknowledged | the backlog is dropped and `stop` reports the loss | +| an object of the same or a later epoch holds the sequence, or any conflicting object in the `enqueued` mode | unchanged | every unacknowledged waiter fails, including those of batches whose objects were created but not yet indexed; an `enqueued` acknowledgement already returned cannot be revoked | poisoned | +| encoding or catalog error | unchanged | every unacknowledged waiter fails, including those of batches whose objects were created but not yet indexed; an `enqueued` acknowledgement already returned cannot be revoked | poisoned | +| create succeeds at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence, every later append fails | +| the start object meets an object of an earlier epoch | moves to the next sequence | none admitted yet | construction continues | +| the start object meets an object of the same or a later epoch, or its create has an unknown outcome | unchanged | none admitted yet | construction fails; an unknown outcome is retryable | + +A store cannot be constructed on a prefix whose largest object already carries the maximum sequence. A poisoned store fails every `LogStore` operation except `stop` with the same terminal error; in particular `obsolete` can no longer move a watermark. In the `durable` mode transient errors are not retried inside the store, but a failed create is reconciled once: the store immediately reads the key back, accepts the batch if the object exists with identical bytes, treats it as a conflict if the bytes differ, and reports the transient error only when the key is absent or that read fails too. So a lost response alone does not fail the append. When the append does fail, neither the store nor Mito keeps the batch: it surfaces through Mito to the caller of the write, and a retry is a new write by the caller under a new sequence. If the create had in fact succeeded and both its response and the read-back were lost, or its request lands only later, the object is an orphan as soon as a later object is acknowledged. Retrying the batch itself inside the store in the `durable` mode would only hide the latency and would hold writes up for as long as an outage lasts; the `enqueued` mode repeats the create with the same bytes, so a create that succeeded without a response is accepted as an identical retry. + +**Stopping.** `stop` is idempotent and awaits the actor. Once stop has begun, queued appends are not admitted, no timer or size triggered seal starts, and appends after stop, including empty ones, fail with the stopped error. In the `durable` mode the open batch and every sealed batch whose create has not started are dropped: their waiters receive the stopped error. A create that is already in flight runs to completion; if it succeeds while its batch is still pending, the batch is indexed in order and acknowledged. A transient failure after stop began fails every batch that is not indexed, as it does while the store runs, and those batches receive the stopped error; a create of such a batch that completes later is an orphan and acknowledges nothing. A create that finds a conflicting object of the same or a later epoch after stop began poisons the store as it would while the store runs, but the waiters of the batches that are not indexed receive the stopped error, the conflict is kept only as the terminal error, and `stop` returns successfully once the creates in flight have completed. In the `enqueued` mode the held-back appends receive the stopped error, the open batch is sealed and the whole backlog is uploaded before `stop` returns; a transient failure in that phase drops the backlog and `stop` returns the error, as described under *Acknowledgement*. A durability wait still pending when the actor exits receives the stopped error, and a delete already in flight runs to completion before `stop` returns. + +**Reading.** `read(provider, entry_id)` returns the region's entries with id at least `entry_id`, in order, fetching only that region's segments with range reads and decoding nothing else, and hides entries at or below the region's obsolete watermark. `obsolete` records the watermark in memory, clamped to the durable entry id in the `enqueued` mode; the watermark is re-established after a restart because Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, so no WAL-owned watermark state is persisted. `obsolete_all` hides the whole region. `latest_entry_id` is not affected by the watermark. Every method taking a provider requires `Provider::ObjectStore` with the store's prefix and the caller's region. + +**Metrics.** The store exports metrics under `greptime_logstore_object_store_wal_`. The write path reports `seal_to_durable_seconds`, the time from sealing a batch until its object is durable, `append_ack_seconds`, the time an append waits for its acknowledgement in the `durable` mode, `object_bytes` and `object_entries`, the size and the entry count of every object sealed, and `created_objects_total`, the creates that succeeded, an identical retry included; a create is counted and timed whether or not its batch is still there to be indexed, since an object a poisoned store gave up on is as durable as any other. `indexed_objects` and `indexed_bytes` are the objects and the bytes the catalog holds, set by a recovery that succeeded and moved by every object indexed or unindexed afterwards, so a collection takes the object it deleted off both, and a recovery that is abandoned publishes neither. A recovery that succeeds reports `recovery_seconds` and the objects it read footers from in `recovered_objects_total`, both of the same recoveries, so one that is abandoned at a fetch or rejected by the catalog contributes to neither; a read stream reports `read_seconds` once it has been read to its end, and one that fails or is dropped part way reports nothing. The failure paths are counted: `create_conflicts_total` for a create that found a conflicting object, `create_failures_total` for one the object store did not confirm, `poisoned_total` for a store that recorded a terminal error, `skipped_segments_total` for a segment a read skipped, and `stalled_appends_total` with `stalled_append_seconds` for an append the `enqueued` backlog thresholds held back, timed until it is admitted or refused, a stop, a terminal error or the teardown of the actor releasing it included; the delete counters of a collection are described under *Entry ids*. The histogram buckets are chosen for the default flush interval of 100 ms and the default `max_batch_bytes` of 8 MiB. + +A `testing` cargo feature exposes hooks to wait for admitted appends, to seal the open batch, to hold and release creates and to make creates fail, so engine tests can flush deterministically and observe entries that are admitted but not durable. + +## Mito integration + +Mito maps `WalOptions::ObjectStore { prefix }` to `Provider::object_store_provider(region_id, prefix)`; resolving object store options against a Raft Engine or Kafka log store reports the existing incompatible-provider error. Reads go through the region-scoped reader used for Kafka, so a namespace shared by many regions never leaks entries across regions even if a log store returned them. Object store regions are opened individually rather than grouped by topic. After replay, `topic_latest_entry_id` follows the existing remote-WAL rule: it is taken from the store's latest entry id only when the region's memtables are empty, otherwise it stays at the replay start, because that value feeds pruning and must never advance past unflushed data. Before a flush, a full truncate or a discard of unflushed data publishes a manifest watermark, Mito waits on the durability barrier described under *Acknowledgement*; with the `durable` mode and with every other provider the wait returns at once. + +## Datanode wiring and lifecycle + +The scalar options are validated before any store or engine exists: the prefix must be non-empty, relative and free of empty or `..` components; `flush_interval` at least 10 ms; `max_batch_bytes`, `max_unpersisted_bytes` and `max_unpersisted_age` positive. `storage_provider` is resolved later, once the object store manager has been built and before the log store and the Mito engine are: empty selects the default store, otherwise the name must match a configured provider. Failures in both stages are `InvalidArguments` naming the field, and both happen before any WAL object is touched. The datanode requires a node id to derive the node prefix it opens the store under, and the standalone bootstrap derives the same prefix for the WAL options it allocates. + +`DatanodeBuilder` keeps the store handle from the moment it is created and stops it if a later build step fails. On success the handle moves into `Datanode`, whose shutdown stops the store after the region server. Shutdown of the datanode and of the standalone instance is best-effort and ordered: every step runs, the first error is returned, so a failing region engine does not leave the WAL actor alive. + +## Configuration + +The options: + +```toml +[wal] +provider = "experimental_object_store" +# Name of a configured storage provider; empty selects the default store. +storage_provider = "" +prefix = "wal" +# How long a non-empty batch may stay open before it is sealed; at +# least 10ms. A shorter interval lowers the acknowledgement latency of +# sparse `durable` writes and raises the request count: timer-triggered +# sealing creates at most one object per interval per node, and in the +# `enqueued` mode the backlog thresholds can seal earlier. The batch +# seals on size only when the admitted bytes reach `max_batch_bytes` +# before the next timer tick. +flush_interval = "100ms" +# Estimated batch size at which the open batch is sealed. It is a +# threshold, not an upper bound: a single append larger than this is +# admitted whole, and object framing adds bytes on top. Object stores +# bill per request, and the cost-effective request size is 8 to 16 MiB. +max_batch_bytes = "8MB" +# "durable": append returns after the object is durable (the default). +# "enqueued": append returns on admission. New appends stall once the +# unpersisted backlog reaches either threshold below; the thresholds do +# not bound what a crash during an outage can lose. +ack_mode = "durable" +max_unpersisted_bytes = "64MB" +max_unpersisted_age = "8s" +# What a read does with a segment that still does not decode after a +# second fetch, because its checksum does not match or its content +# disagrees with its footer entry. "skip": the segment is skipped and +# recorded as a WAL hole of its region, counted and logged; the other +# regions of the object are unaffected (the default). "fail": the read +# fails, so the region does not open. +on_corrupted_segment = "skip" +``` + +`config/config.md` is generated from the example files. + +## Compatibility + +- `WalOptions::ObjectStore` is a new persisted variant. Existing `raft_engine`, `kafka` and `noop` options decode unchanged; a region's options never change provider implicitly. +- The object format is version 1 and frozen once the provider is released; a version bump after that is a new decoder, not an in-place change. The version 1 header layout was revised before release to carry the epoch, the link, the reclaim boundary and the header CRC. No object in an earlier layout exists, since the provider has not been enabled anywhere. +- The configuration tag carries `experimental_`; the persisted tag does not, so promoting the provider to stable is a configuration rename with an alias, not a metadata migration. +- `LogStore::wait_durable` has a default implementation that returns at once, so Raft Engine, Kafka and the noop provider keep their own acknowledgement and persistence semantics; the durability barrier binds only the object store provider. +- The object-sequence-major id scheme needs no migration of existing objects. A prefix written under the contiguous scheme keeps its objects, and the sequence allocation rule under *Entry ids* keeps new ids above the old ones. Garbage collection reasons from footers and from the reclaim boundary, not from the id encoding, so it applies to objects of either scheme. +- `on_corrupted_segment` is a new configuration option with a default, so existing configurations decode unchanged; it changes what a read does with a damaged segment, not what is written. + +## Evidence + +In summary, the evidence supports going ahead with the experimental standalone implementation. Targeted tests and a process-level crash gate cover the recovery paths and the durability barrier; a local MinIO run confirms recovery and garbage collection. In same-session comparisons on a laptop the default `durable` mode reaches about 23% to 28% less throughput than Raft Engine with synchronous writes; on S3 from an instance in the bucket's region, against Raft Engine on a local NVMe store, it reaches 31% of Raft Engine's throughput with one writer and 54% with eight, bound by the seal interval rather than by S3, while the `enqueued` mode matches Raft Engine and never stalls. Under sustained load the default interval costs about 600 object creates per minute per node. The pipelined uploader has shown no consistent throughput benefit; in a multi-region experiment with injected object store latency it lowered write latency at the cost of more object requests. Long-running stability and recovery at production scale remain unverified. + +The reference implementation was validated on the following evidence; the tests are part of the pull requests that submit it. The first table maps each guarantee of this RFC to what was verified, the list after it names the failures that were driven through the store, and the last list states what the evidence does not cover. + +### Guarantees verified + +| Guarantee | Verified by | Result | +| --- | --- | --- | +| Acknowledged writes replay exactly once across graceful and ungraceful restarts under `durable` | Deterministic recovery tests on a real store over an in-memory object store: two regions sharing a prefix with only one flushed, an abrupt drop after an object is durable but before any flush, two consecutive restarts with writes and a flush between them. A process-level crash gate: concurrent writers, SIGKILL at a random point of a window, restart on the same bucket and data home, repeated cycles with replay accumulating. A standalone round trip against MinIO with a restart before and one after a flush. | The concrete flushed, replayed, manifest and latest entry ids and row-level scan equality after every restart. In the crash gate every acknowledged row is present exactly once after the restart, and every present row that was not acknowledged was in flight at the kill. | +| The loss of `enqueued` is bounded to the unpersisted backlog and reported | The same recovery tests and the same crash gate run in `enqueued` mode; unit tests of both admission thresholds stalling until an upload completes, of `stop` uploading the backlog, and of `stop` reporting a backlog lost before it was called. | No row is present twice and no present row was never attempted; the crash gate reports the acknowledged rows lost per cycle. | +| A manifest watermark never names an entry that is not durable | Durability barrier tests: a flush held until its entry is durable; a crash in `enqueued` mode after an acknowledgement and a flush attempt but before the object exists, followed by a restart, new writes and a second crash; a drop cancelling a flush that waits; a truncate that waits; a flush after a backlog was lost. | Every durable entry replays; the flush after the lost backlog publishes no frontier; `obsolete` never passes the durable id. | +| A region replays only its own entries | The two-region recovery test above. | Per-region row-level scan equality with the other region unflushed. | +| Recovery never silently loses data | Unit coverage of every corruption class at recovery and on read; footer-only recovery compared with a full decode; a footer longer than the tail window. | Structural damage fails construction of the store; a corrupted segment is skipped, recorded as a hole of its region and counted, or fails the read under `fail`; the footer-only catalog equals the full decode. | +| A region cannot be silently pointed at another WAL | Reopen on a mismatched prefix; reopen on an empty prefix. | The mismatch is rejected with both prefixes named and the region is not opened; the empty prefix opens cleanly with the expected row counts. | +| Garbage collection deletes everything below the persisted reclaim boundary, and recovery stays correct while objects disappear | Unit tests of the boundary computation (a region without a watermark holds it at its first object, a watermark inside an object keeps the object, `obsolete_all` releases it, last indexed plus one when everything is flushed), of stamping and durability (headers carry the boundary current at sealing, a failed create's boundary permits nothing, the start object carries the recovered boundary), of deletion (chain objects, listed orphans and failed batches below the durable boundary, nothing at or above it, a failed delete retried, `stop` waiting for deletes in flight, holes dropped), of recovery (two restarts with a collection between them, a late object of an earlier epoch below the boundary collected, a missing predecessor above the boundary still failing) and of a read that listed an object before it was collected; the Mito flush-then-collect test in both acknowledgement modes; the standalone integration test's assertion after a flush. | Objects below the durable boundary are gone after a flush, the start object included; a restart replays everything above the boundary; exactly one WAL object remains after a flush and one more insert in the integration test. | +| Every row of the *Failure matrix* holds, and every metric moves as documented | A table-driven fault matrix, `test_store_fault_matrix`, that drives every failure the object store access can produce through the store and asserts, per case, which waiters fail and with what error, whether the store is poisoned, whether the sequence is consumed, the indexed gauges at the failure itself, the exact movement of every metric, and what a read of the store sees afterwards; every timer is observed twice around two forced waits whose intervals cannot overlap. | 33 cases: 17 for the ten rows of the matrix and 16 for the failures listed under *Failures driven through the store*. The test fails if a row of the matrix is left uncovered. | +| Recovery cost does not grow with object size | Footer-only recovery over 2,000 objects; recovery concurrency bounded and ordered by sequence. | Construction of the store fell from 3.6 s to 0.9 s. | + +### Failures driven through the store + +- Conditional create outcomes: success, identical retry, conflicting content from an earlier epoch and from the same or a later one, lost response, transient failure in `durable` mode with later objects already created and with creates still in flight (every batch that is not indexed fails, no sequence is reused, late completions are orphans), transient failure in `enqueued` mode and after `stop` began, completion after a conflict had already poisoned the store, a create at the last representable sequence. +- Object store access: an unwritable store; a listing, a trailer fetch, a footer fetch, a segment fetch and a delete that fail; a delete that fails and is retried by the next collection; an object collected between the listing and its fetch, below and at the tip's reclaim boundary; an object store that answers slowly. +- Store state: an object the catalog rejects, per-region entry id exhaustion, the catalog refusing a next sequence past the maximum, every stop path with several creates in flight that succeed, fail transiently or conflict after stop began, a stop with a delete in flight, and the appends the backlog thresholds were holding back released by admission, by a stop, by a terminal error and by the teardown of the actor. +- Pipelined uploader: concurrent creates acknowledged in sequence order, every sealed batch linked to its pipeline predecessor, the open batch sealing at the threshold while creates are in flight, a transient failure failing every batch that is not indexed, created ones included, a permanent failure poisoning the store. +- Chain recovery: a linear chain, orphans below and above the tip, a broken tip falling back to the latest complete object, a late object of an earlier epoch next to the chain of a later one, a missing predecessor, objects present without a complete chain, orphans raising the next sequence, and the start object on an empty prefix and meeting objects of an earlier and of a later epoch. + +### What the evidence does not cover + +- The cloud measurement covers one instance type, one region and one bucket, with the Raft Engine control on the instance's local NVMe store; no crash gate ran on the cloud host, and the request cost is the list price, not a bill. +- Loss during writes is measured only by the crash gate, in single samples. The Raft Engine comparison below killed the server after every row was counted, so its replay checks confirm completeness and say nothing about loss. +- The store's own work, encoding, indexing and the actor, is not quantified separately from the object store round trip. +- A create whose response and immediate read-back are both lost is covered only at recovery, by placing such objects next to a chain; no test loses a response and its read-back at the object store. +- The benefit of the pipelined uploader was measured only with injected latency on a local object store. + +### Write path measurements + +A log workload benchmark was run against the write path on a laptop with a local MinIO server. The numbers calibrate the defaults; the same workload on S3 is reported under *Cloud write path measurements* below. The profile writes 2,000,000 rows of about 900 bytes in 1000-row batches into one region at client concurrency 1, 8 and 32. Raft Engine with `sync_write = true` on the same build and the same storage is the control. Rows per second: + +| Write path | 1 | 8 | 32 | +| --- | --- | --- | --- | +| Raft Engine, `sync_write = true` | 3,500 to 4,000 | 19,000 to 25,000 | 20,000 to 28,000 | +| one create in flight, `durable`, 1 s | 998 | 4,001 | 19,029 | +| one create in flight, `durable`, 100 ms | 3,287 | 13,662 | 17,204 | +| one create in flight, `durable`, 10 ms | 3,632 | 19,835 | 19,309 | +| pipelined uploader, `durable`, 1 s | 1,000 | 3,999 | 22,043 | +| pipelined uploader, `durable`, 100 ms | 3,304 | 17,207 | 23,185 | +| pipelined uploader, `enqueued`, 1 s | 3,955 | 24,074 | 22,168 | +| pipelined uploader, `enqueued`, 100 ms | 3,999 | 23,595 | 22,662 | + +The rows come from separate runs, and the same build with the same configuration, on an idle host each time, measured 1.4 times apart in runs days apart, so a difference between two rows of the table is not a measurement of the difference between those write paths. The Raft Engine range covers several runs. Both write paths converge on the host's CPU ceiling at concurrency 32, 6 to 7.5 of 10 cores busy with one region, so the last column measures the host rather than the WAL. A serial writer under the `durable` mode is phase-locked to the seal timer: its extra latency over Raft Engine is the interval minus the base insert cost modulo the interval, which is 747 ms at 1 s, 52 ms at 100 ms, and a floor of about 14 ms, one seal plus one PUT, at 10 ms. The pipelined uploader alone was measured by counterbalanced A/B runs of the builds before and after it, run in both orders within one session: at 1 s and concurrency 32 the two builds were within about 1.4% and the two pairs disagreed on which was faster, and at 100 ms and concurrency 8, where both builds sealed about ten objects per second, they were 0.25% apart on the geometric mean of two pairs, while whichever build ran first in a pair was 3.5% to 4% faster, and their object counts and sizes were within 2%. No throughput benefit was measured at either point. Every restart under the `durable` mode replayed exactly the acknowledged rows. In the `enqueued` mode a SIGKILL lost one to two 1000-row batches at 1 s and none at 100 ms, in single samples, always whole batches and never a duplicate. Footer-only recovery cut log store construction from 3.6 s to 0.9 s at 2,000 objects. + +The conditions under which the pipelined uploader does help were then examined, again comparing the builds before and after it within one session in both orders at concurrency 8 and 100 ms, with a proxy in front of the object store that delays every request of the WAL while the data files keep a direct connection. In the build before it the actor awaits each create before it admits more entries, and Mito's region worker in turn waits for one append to become durable before it sends the next, so what the uploader can run side by side depends on how many appends wait at once, not only on how slow a create is. With one region, adding about 40 ms or about 160 ms to each create left the two builds within the order effect of these runs, since one region never has a second sealed batch while the first uploads. With a table split into three regions, each served by a different region worker, and about 160 ms added, the pipelined uploader was faster by 13.6% and 2.1%, 7.7% on the geometric mean, still within that margin, and it alone let the regions fall out of step: 27% to 46% of its sampled objects held fewer than three regions, against 0% to 4% before it, because a region whose append became durable while another create was still in flight started the next batch. The clients of those runs wait for each response, which ties every region back to the seal. With the three regions driven at a fixed request rate of 0.7 times that closed-loop throughput instead, the pipelined uploader lowered the median insert latency by 24.0% and 12.2% in the two pairs, 18.3% on the geometric mean, and the 99th percentile in both pairs; at 1.0 times the rate one valid pair showed 23.8%, while the second pair was not completed because the load generator repeatedly fell behind its send schedule. The lower latency is paid for in requests: with about 160 ms added, the pipelined uploader created 19% to 28% more objects with one region, 51% to 67% more with three regions in the closed loop, and 35% to 63% more at 0.7 times the rate in the open loop. These measurements used a local object store with added latency and do not measure a cloud object store. + +The comparison with Raft Engine was measured again within one session, with one binary and only the `[wal]` section differing, each comparison run in both orders to reduce the bias of running first or second. Runs on this host still varied by 4% to 8% with their position alone, which limits how small a difference these runs can resolve. Both runs of a pair killed the server only after every row was written and counted, so the replay checks confirm a complete replay and say nothing of what a crash during writes loses. Rows per second: + +| Write path | Concurrency 1 | Concurrency 8 | +| --- | --- | --- | +| Raft Engine, `sync_write = true` | 3,817 and 4,503 | 25,124 to 26,880, four runs | +| object store WAL, `durable`, 100 ms | 3,184 and 3,198 | 17,779 to 19,174, four runs | +| object store WAL slower by, geometric mean of the pairs | about 23% | about 28% | + +The gap does not close under concurrency; it is slightly wider at 8 than at 1. At concurrency 1 the object store WAL measured the same in both runs, 301 ms at the median, while Raft Engine moved between a median of 237 ms and 221 ms, both of which round up to the same 300 ms seal, so the single-writer figure depends on how fast Raft Engine happens to be. Raft Engine exceeded the upper end of the range in the table above in every one of the four runs at concurrency 8. + +Two further configurations of the same binary, run in both orders in two rounds at concurrency 8, locate the gap. With the WAL objects written to local disk through a file object store instead of MinIO, still in the `durable` mode, MinIO was 4.5% slower. Switching the MinIO configuration to the `enqueued` mode was 35% faster, and the `enqueued` configuration was 4.4% slower than Raft Engine. The first and third of these differences are within the run-to-run variation of the host, while the second is far larger, which points to the `durable` mode waiting for the seal and the create before it acknowledges as the main contributor to the gap in this workload. These are differences between whole configurations: they do not measure the MinIO round trip or the store's own work, encoding, indexing and the actor, individually, and the runs leave those costs unquantified. The file object store is not equivalent to a remote one: its conditional create checks, writes a temporary file and renames it rather than creating atomically, and it does not sync the directory after the rename. The `enqueued` mode does not provide the durability of Raft Engine with `sync_write = true`, so this comparison separates costs and is not a comparison of equivalent write paths; against Raft Engine with its default `sync_write = false`, also run in both orders within one session, the `enqueued` mode was 5.1% slower, again within the noise. These local runs do not measure the latency or throughput of a cloud object store. + +The default interval is 100 ms because it removes most of the seal-timer latency of a sparse `durable` writer, 52 ms of residual wait against 747 ms at 1 s, at a request cost that stays bounded: timer-triggered sealing creates at most one object per interval per node, ten per second at the default, the `enqueued` mode can seal earlier only when the unpersisted backlog reaches a threshold, and the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next timer tick, which the measured S3 workload did not reach at about 60 MB/s: every object store cell sealed on the timer. Going to 10 ms buys another 40 ms for ten times the request ceiling. On S3 at the default configuration the seal timer, not the PUT, bounded a serial `durable` writer to one batch per interval, see *Cloud write path measurements*; a shorter interval removes part of that wait but not the create's latency, and was not measured on S3. The measured way to Raft Engine's throughput from one writer is the `enqueued` mode, at the cost of its loss bound. The default acknowledgement mode stays `durable` because that is what a write-ahead log means: an acknowledged write survives a crash. `enqueued` is an explicit choice with a documented loss bound, the unpersisted backlog, and the crash gate reports what it loses. + +### Cloud write path measurements + +The same workload was run against S3 from an EC2 instance in the bucket's region (`c6id.2xlarge`, 8 vCPU, `ap-southeast-1`, the nightly image of `main` of 2026-10-12) by the workflow of pull request #9536 (https://github.com/GreptimeTeam/greptimedb/pull/9536), which starts GreptimeDB with its data files on S3 and differs between cells only in the `[wal]` section: Raft Engine with `sync_write = true` on the instance's local NVMe store, and the object store WAL in the `durable` and in the `enqueued` mode, both with the defaults. 2,000,000 rows of about 900 bytes in 1000-row batches at concurrency 1 and 8, two rounds with the order of the configurations reversed between them, every cell on a fresh container and a fresh prefix. Rows per second, geometric mean of the two rounds of run 38232084219 (https://github.com/GreptimeTeam/greptimedb/actions/runs/38232084219): + +| Write path | 1 | 8 | +| --- | --- | --- | +| Raft Engine, `sync_write = true`, local NVMe | 31,700 | 67,800 | +| object store WAL, `durable`, 100 ms | 9,670 | 36,900 | +| object store WAL, `enqueued`, 100 ms | 32,900 | 67,900 | + +A second run of one round at concurrency 8, 38234710697 (https://github.com/GreptimeTeam/greptimedb/actions/runs/38234710697), reproduced the ratios within 1%: the `durable` mode at 0.54 of Raft Engine, the `enqueued` mode at 1.00. The latencies, object counts and histograms below are read from the metrics scraped after every cell, kept in the runs' artifacts. + +A conditional create on S3 took about 50 ms at the median and 110 to 130 ms at the 99th percentile from seal to durable, against single-digit milliseconds on the local server, and still a serial `durable` writer is bound by the seal timer, not by S3: every 1000-row batch took about 103 ms, one object per batch, with its acknowledgement at 91 ms at the median, so at the default configuration its throughput was the batch size over the interval, about ten batches per second: the measured 103 ms per batch is one interval plus the admission of the next batch, since the create, the acknowledgement and the next append all complete before the following tick. No shorter interval was measured on S3. Raft Engine on the instance store took 32 ms per batch, hence 0.31. The local ratio of 0.77 was higher only because the local Raft Engine batch took over 220 ms, which rounds up to the next seal with little loss: the faster the control, the lower the ratio at low concurrency under the default interval. At concurrency 8 the `durable` mode still seals on the timer, about four batches per object, and reaches 0.54. The `enqueued` mode matched Raft Engine at both concurrencies and never stalled: at about 60 MB/s of WAL the backlog thresholds were not approached, since one create per interval of 3.5 to 8 MiB drains it. Under sustained load the timer ceiling sets the request count: 550 to 600 creates per minute per node in every cell, which at the S3 standard list price of $0.005 per thousand PUTs is about $0.18 per node-hour of continuous writing, plus about 30 deletes per cell for garbage collection. No create failed, found a conflicting object or poisoned the store in 7,609 creates, the 6,809 of the first run's eight object store cells and the 800 of the second run's two. + +# Alternatives + +## A generic external provider injected from outside the repository + +The backend could stay outside the main crates and be injected through an opaque provider variant plus a resolution hook on the `LogStore` trait. That avoids touching the four enums, but every resolution path then has to special-case an opaque value, distributed rejection can only be enforced by convention, and nothing in the repository can test the backend. Making the provider an explicit variant keeps every match exhaustive and lets the engine tests run against the real store, at the cost of one more arm in each enum. + +## Reusing the Kafka provider with an object store "topic" + +The Kafka code path assumes a broker assigns offsets and that a topic is shared across datanodes. Neither holds here: entry ids are assigned by the datanode, and one prefix belongs to one datanode. Reusing the Kafka provider would have inherited the topic tickers, the failover preconditions and the metasrv allocation logic, all of which are wrong for this design. + +## One object per region per flush + +Simpler catalog, but every region then produces its own timer-driven object, so the request rate scales with the number of regions even when they write little, which is the cost trap object stores punish. Sharing one batch across regions removes that per-region floor: at low volume timer-triggered sealing creates one object per interval regardless of region count, with the `enqueued` backlog thresholds able to seal earlier, and at high volume the count is governed by aggregate bytes and `max_batch_bytes`, which is the same request volume a single busy region would produce. + +# Drawbacks + +- Write latency under the `durable` mode is the time until the batch seals (up to `flush_interval`, 100 ms by default), plus waiting for every earlier object, plus one PUT. A serial writer waits for the next seal, so its extra latency over Raft Engine is the interval minus the base insert cost modulo the interval, and concurrency does not close the gap: measured within one session on a laptop it was about 23% of throughput at concurrency 1 and about 28% at concurrency 8, and on S3 against Raft Engine on a local NVMe store the `durable` mode reached 31% of Raft Engine's throughput at concurrency 1 and 54% at concurrency 8, a serial writer being bound to one batch per interval, while the `enqueued` mode matched it, see *Write path measurements* and *Cloud write path measurements*. +- Every object costs a request. Timer-triggered sealing creates at most one object per interval per node, ten per second at the default; the `enqueued` mode can seal earlier when the unpersisted backlog reaches a threshold, and the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next tick. Under sustained load on S3 the timer ceiling was reached, about 600 creates per minute per node, about $0.18 per node-hour at the S3 standard list price. +- Garbage collection keeps every object at or above the reclaim boundary, so a region that is written but rarely flushed, or never flushed, holds the boundary at its first object and keeps every later object the prefix shares with other regions; recovery time grows with the number of retained objects. A prefix that receives no appends cannot persist a higher boundary, so an idle prefix is collected only after its next object. +- Every open of the store writes one start object without segments, one extra PUT per open, and every object carries a 54-byte header. Orphans, the objects of failed batches that were stored anyway, hold storage together with the rows Mito wrote again after those failures until garbage collection deletes them. +- Object-sequence-major entry ids consume 20 bits per object for positions, which caps a region at about a million entries per object; the batch seals when the cap is reached. Sequence numbers are 44 bits wide as a result, enough for one object per second for half a million years. + +# Alignment with cluster mode + +Distributed mode is out of scope, but several decisions were taken so that the cluster design can build on this backend without a metadata migration. Except where a paragraph says otherwise, all of them are *Proposed*. + +**Key layout.** The layout is implemented in its standalone form, see *Object format*: the operator configures only the prefix root, the store derives the prefix it writes under as `/datanodes//epochs/`, so objects live at `/datanodes//epochs//objects/.wal`, standalone uses its datanode id and generation zero, and the persisted `WalOptions` of a region carry the full derived prefix, so a region's WAL location is self-describing. *Proposed*: in cluster mode the metasrv assigns both the node id and the generation, and issues a new generation every time it hands a region's write ownership to a node; the self-describing location then lets a node other than the writer open a region's WAL read-only, which today is rejected because a datanode builds its store only under its own derived prefix and the store refuses providers with another prefix. Putting the generation into the key rather than only into the object header means that a writer which lost its lease and keeps writing lands in its old epoch directory, where the new writer never creates objects, so the two can never contend for a sequence number; the generation in the key is the fence. The version 1 header carries no generation; adding one as a read-side check would be a format evolution, a header version bump whose reader keeps decoding version 1 objects, and is not part of this proposal. The header does carry an epoch, which is implemented: the epoch orders the store instances that wrote under one prefix, while the generation separates prefixes, so the two are orthogonal. The generation could later seed the epoch, but the epoch does not depend on it. Despite the name, the `epochs/` component of the key holds the generation, not the header epoch. + +**Planned migration.** A region that moves while its source node is alive follows the existing migration flow, and the cheap path applies only when that flow proves the handoff: the source stops admitting writes, flushes, and the resulting manifest watermark covers the last acknowledged entry. Then the target opens the region from that watermark and starts writing under its own prefix, and no cross-node read is needed. If the final flush fails or times out, the fact that the source process is alive proves nothing about its unflushed tail; the migration then either takes the failover path below, with the source's prefix in the chain, or is aborted. + +**Failover.** When the source is gone, the target must replay the region's unflushed tail from the source's prefix. The region's WAL position becomes a chain of segments whose last element is the target's own prefix. Each earlier segment is `(prefix, first sequence, cutover sequence, holes)`: the metasrv appends it to the region's `WalOptions` when it reassigns the region, after the source's lease has expired, and it selects the history under the source prefix at that moment by the rule the store applies at open: the complete object with the largest epoch and, among those, the largest sequence, and every object it links to. It records that chain as the segment's replay history, the cutover sequence and the sequences of the chain's objects, and the record is sealed: a present object outside the chain, whether the orphan of a create the source reported as failed or a stale object of an earlier epoch, is excluded, an object that becomes visible later, because the old writer's upload was still in flight, stays excluded, and an object above the cutover is ignored, so every open of the region replays the same source history. On open the target replays the earlier segments in order, reading only footers and its own region's segments through range reads, and never creates an object under a foreign prefix, so there is nothing to conflict with. + +Sealing may only exclude objects whose writes were never acknowledged, which takes two fences. Admission is fenced by the lease mechanism Mito already has: when a region's lease expires the datanode demotes it to a follower, which refuses writes, and the metasrv waits for that deadline before it upgrades the candidate. Completion is fenced by the store: every conditional create runs under an upload deadline, a create that has not completed by its deadline is reported failed to its waiters and is never acknowledged afterwards even if the object turns out to exist, and when a region's lease expires the store fails every waiter of that region that is not yet acknowledged, so that a create still in flight at that moment cannot be acknowledged later. The metasrv seals the segment only after the lease deadline plus the upload deadline have both passed, so an object that appears after sealing belongs to a create the old writer had already reported as failed. Under the durable acknowledgement mode this means a sealed set never excludes an acknowledged write; under the `enqueued` mode the entries in the old writer's unpersisted backlog at that moment are lost, which is that mode's stated recovery point objective. The late-completing create is part of the acceptance criteria for takeover, stated in terms of the sealing instant: an object that was absent at sealing, whether it fills a hole later or lies above the cutover, is excluded on every open of the region, and so is a visible object outside the selected chain; an object on the sealed chain replays on every open even if its write was never acknowledged, which is the same outcome as a crash between an object's creation and its acknowledgement; and in the durable mode every caller whose create had not completed when the lease expired must have been failed, so no acknowledged write is among the excluded objects. The generation in the key is an additional guarantee that such a late object lands in a directory the new writer never reads beyond the sealed set. Object-sequence-major entry ids are what make this cheap: `entry_id >> 20` locates the object under the right prefix without the source node's in-memory catalog. + +**Entry id floor across prefixes.** A region that arrives under a new prefix, by migration or by failover, brings a watermark and replayed ids that were assigned under another prefix's sequence. Before the target assigns the region its first new id, the target store raises its next object sequence above `max_id >> 20`, where `max_id` is the largest of the region's inherited watermark and every id it replayed, so that every new id of that region is greater than every id it already has; sequence gaps are tolerated and the sequence space is wide enough for this. The rule is applied on every open from the region's persisted state, in the same place Mito already hands the store the manifest watermark, so it needs no persisted state of its own and survives a crash before the target's first flush: the restart re-derives the floor and the region's acknowledged-but-unflushed entries under the target prefix are above the inherited watermark and replay. + +**Acknowledgement modes.** Both are implemented. Cluster mode is expected to run `enqueued` for latency-sensitive workloads; the object layout, the flush durability barrier and the recovery path are identical in both modes. + +**Corruption.** Segment-level skipping with the hole recorded per region is implemented, see *Object format*. *Proposed*: cluster mode routes the WAL hole event to the metasrv so that a follower or a takeover knows the region needs a fresh replica rather than a replay. + +**Garbage collection across nodes.** Within one prefix, garbage collection is implemented as described under *Entry ids*: watermark from each region's manifest, per-segment footer comparison, a reclaim boundary persisted in every header. A source prefix that still holds segments of regions which have moved away cannot be trimmed by the source alone, because it no longer learns those regions' watermarks. Either the target reports to the metasrv, after each flush, the source sequence below which it no longer needs the region, and the metasrv drives deletion, or a cluster-level janitor collects the watermarks; which of the two is the last of the *Unresolved questions*. + +# Future work + +1. Probe the next sequence with a GET from the last known sequence instead of listing the whole prefix at recovery; on a prefix with many objects the probe is cheaper than the listing. +2. Distributed mode, which needs metasrv-side allocation of node ids and generations, the takeover chain in `WalOptions`, read-only replay of a foreign prefix, and cross-node garbage collection, as sketched under *Alignment with cluster mode*. + +# Unresolved questions + +- Whether concurrent writers on one prefix need a fence beyond the epoch conflict. Conditional creates already prevent overwrites, and a concurrent writer is noticed when its epoch meets another at a sequence, which can be after both acknowledged writes. +- Whether garbage collection of a source prefix after a failover is driven by the target reporting its watermark to the metasrv, or by a cluster-level janitor that collects watermarks; the same question covers regions that are closed on a node but still own objects under its prefix. From 674ba7266949c3c05b71e6e7fb99891a71cc8d2e Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 12 Oct 2026 13:16:52 +0000 Subject: [PATCH 18/46] fix(mito): preserve version ordering in bulk merges and flushes (#9517) * fix(mito): preserve delete barriers in flat last-non-null dedup Signed-off-by: evenyag * fix(mito): respect sequence ranges in bulk merges and flushes Signed-off-by: evenyag * perf(mito): avoid merge selection on the bulk write path Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/src/access_layer.rs | 2 +- src/mito2/src/engine/flush_test.rs | 313 ++++++++++++++++- src/mito2/src/flush.rs | 330 +++++++++--------- src/mito2/src/memtable.rs | 26 ++ src/mito2/src/memtable/bulk.rs | 514 ++++++++++++++++------------ src/mito2/src/memtable/bulk/part.rs | 32 +- src/mito2/src/read/flat_dedup.rs | 163 +++++---- 7 files changed, 924 insertions(+), 456 deletions(-) diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 3959e0c26a4..b1457935c16 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -579,7 +579,7 @@ pub struct SstWriteRequest { pub source: FlatSource, pub cache_manager: CacheManagerRef, /// Optional uniform row sequence for writes that do not preserve sequences. - /// Compaction passes `None` to retain the reader's effective input sequences. + /// Compaction and LastNonNull flushes pass `None` to retain input sequences. pub max_sequence: Option, pub sst_write_format: FormatType, diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index b040ae5ec54..313860c0857 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -27,6 +27,7 @@ use common_base::readable_size::ReadableSize; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use datafusion_expr::{col, lit}; use parquet::basic::{Encoding, Type as PhysicalType}; use rstest::rstest; use rstest_reuse::{self, apply}; @@ -35,9 +36,9 @@ use store_api::metadata::RegionMetadataRef; use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - AlterKind, PathType, RegionAlterRequest, RegionCloseRequest, RegionFlushRequest, - RegionOpenRequest, RegionPutRequest, RegionRequest, ReplayCheckpoint, SetRegionOption, - UnsetRegionOption, + AlterKind, PathType, RegionAlterRequest, RegionCloseRequest, RegionCompactRequest, + RegionFlushRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, ReplayCheckpoint, + SetRegionOption, UnsetRegionOption, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::Notify; @@ -49,9 +50,11 @@ use crate::engine::listener::{EventListener, FlushListener, StallListener}; use crate::engine::region_hook::{RegionGcInfo, RegionHook, RegionHookRef, SstFileInfo}; use crate::error::Error; use crate::manifest::action::RegionMetaActionList; +use crate::memtable::RangesOptions; use crate::test_util::{ - CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, build_rows, - build_rows_for_key, flush_region, kafka_log_store_factory, multiple_log_store_factories, + CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, + build_delete_rows_for_key, build_rows, build_rows_for_key, build_rows_with_fields, delete_rows, + delete_rows_schema, flush_region, kafka_log_store_factory, multiple_log_store_factories, prepare_test_for_kafka_log_store, put_rows, raft_engine_log_store_factory, reopen_region, rows_schema, single_kafka_log_store_factory, }; @@ -1818,3 +1821,303 @@ async fn test_region_hook() { manifest_count, ); } + +/// Explicitly finish a merge instead of sleeping for the background dispatcher. +fn compact_bulk_memtables(engine: &MitoEngine, region_id: RegionId) { + let version = engine.get_region(region_id).unwrap().version(); + for memtable in version.memtables.list_memtables() { + memtable.compact(false).unwrap(); + } +} + +async fn scan_bulk_test_keys(engine: &MitoEngine, region_id: RegionId) -> String { + let mut schema = None; + let mut batches = Vec::new(); + for key in ["deleted", "k"] { + let request = ScanRequest { + filters: vec![col("tag_0").eq(lit(key))], + ..Default::default() + }; + let stream = engine.scan_to_stream(region_id, request).await.unwrap(); + schema = Some(stream.schema()); + batches.extend(RecordBatches::try_collect(stream).await.unwrap().take()); + } + RecordBatches::try_new(schema.unwrap(), batches) + .unwrap() + .pretty_print() + .unwrap() +} + +#[rstest] +#[tokio::test] +async fn test_bulk_flush_keeps_newer_version(#[values(false, true)] delete_newer: bool) { + let mut env = TestEnv::new().await; + let config = MitoConfig { + default_flat_format: true, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env.create_engine(config.clone()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .insert_option("memtable.bulk.merge_threshold", "2") + .insert_option("memtable.bulk.encode_row_threshold", "1") + .build(); + let schema = rows_schema(&request); + let delete_schema = delete_rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("k", 1, 2, 1), + }, + ) + .await; + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("a", 0, 1024, 0), + }, + ) + .await; + if delete_newer { + let mut rows = build_delete_rows_for_key("b", 0, 1023); + rows.extend(build_delete_rows_for_key("k", 1, 2)); + delete_rows( + &engine, + region_id, + Rows { + schema: delete_schema, + rows, + }, + ) + .await; + } else { + let mut rows = build_rows_for_key("b", 0, 1023, 0); + rows.extend(build_rows_for_key("k", 1, 2, 2)); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + } + compact_bulk_memtables(&engine, region_id); + // Previously this advanced the old unordered row above the newer encoded row. + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows_for_key("c", 0, 1, 0), + }, + ) + .await; + let expected = if delete_newer { + "\ ++-------+---------+----+ +| tag_0 | field_0 | ts | ++-------+---------+----+ ++-------+---------+----+" + } else { + "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| k | 2.0 | 1970-01-01T00:00:01 | ++-------+---------+---------------------+" + }; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + flush_region(&engine, region_id, None).await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + let engine = env.reopen_engine(engine, config).await; + reopen_region( + &engine, + region_id, + request.table_dir, + false, + request.options, + ) + .await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); +} + +#[rstest] +#[tokio::test] +async fn test_bulk_last_non_null_merge_and_flush(#[values(false, true)] encoded: bool) { + let mut env = TestEnv::new().await; + let config = MitoConfig { + default_flat_format: true, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env.create_engine(config.clone()).await; + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let request = CreateRequestBuilder::new() + .field_num(2) + .insert_option("memtable.type", "bulk") + .insert_option("merge_mode", "last_non_null") + .insert_option("memtable.bulk.merge_threshold", "4") + .insert_option( + "memtable.bulk.encode_row_threshold", + if encoded { "1" } else { "1000000" }, + ) + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "2") + .build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + // An older SST must remain hidden behind a delete retained in the memtable. + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_with_fields("deleted", &[1], &[(Some(1), Some(1))]), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + + let mut a = build_rows_with_fields( + "a", + &(0..1023).collect::>(), + &[(Some(0), None); 1023], + ); + a.extend(build_rows_with_fields("k", &[1], &[(Some(1), None)])); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: a, + }, + ) + .await; + // Intermediate writes must be bulk parts, not an unordered merge barrier. + let mut b = build_rows_with_fields( + "b", + &(0..1023).collect::>(), + &[(Some(0), None); 1023], + ); + b.extend(build_rows_with_fields("k", &[1], &[(Some(2), None)])); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: b, + }, + ) + .await; + let mut deletes = build_delete_rows_for_key("d", 0, 1023); + deletes.extend(build_delete_rows_for_key("deleted", 1, 2)); + delete_rows( + &engine, + region_id, + Rows { + schema: delete_rows_schema(&request), + rows: deletes, + }, + ) + .await; + let before = scan_bulk_test_keys(&engine, region_id).await; + assert!(before.contains("| k | 2.0")); + assert!(!before.contains("| deleted")); + + let mut c = build_rows_with_fields( + "c", + &(0..1022).collect::>(), + &[(Some(0), None); 1022], + ); + c.extend(build_rows_with_fields("k", &[1], &[(None, Some(3))])); + c.extend(build_rows_with_fields("deleted", &[1], &[(None, Some(3))])); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: c, + }, + ) + .await; + compact_bulk_memtables(&engine, region_id); + // All four sources must have merged, with the requested output format. + let version = engine.get_region(region_id).unwrap().version(); + let mut ranges = Vec::new(); + for memtable in version.memtables.list_memtables() { + ranges.extend( + memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .ranges + .into_values(), + ); + } + assert_eq!(1, ranges.len()); + assert_eq!(encoded, ranges[0].encoded().is_some()); + drop(ranges); + drop(version); + let expected = "\ ++---------+---------+---------+---------------------+ +| tag_0 | field_0 | field_1 | ts | ++---------+---------+---------+---------------------+ +| deleted | | 3.0 | 1970-01-01T00:00:01 | +| k | 2.0 | 3.0 | 1970-01-01T00:00:01 | ++---------+---------+---------+---------------------+"; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + // Leave a raw source beside merged parts so flush also exercises grouping. + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows_with_fields("z", &[0], &[(Some(0), None)]), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + .unwrap(); + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + let engine = env.reopen_engine(engine, config).await; + reopen_region( + &engine, + region_id, + request.table_dir, + false, + request.options, + ) + .await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); +} diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index e910c81f896..d5191bdcadb 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -46,7 +46,10 @@ use crate::error::{ }; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::memtable::bulk::ENCODE_ROW_THRESHOLD; -use crate::memtable::{BoxedRecordBatchIterator, EncodedRange, MemtableRanges, RangesOptions}; +use crate::memtable::{ + BoxedRecordBatchIterator, EncodedRange, MemtableRange, MemtableRanges, RangesOptions, + group_by_sequence, +}; use crate::metrics::{ FLUSH_BYTES_TOTAL, FLUSH_ELAPSED, FLUSH_FAILURE_TOTAL, FLUSH_FILE_TOTAL, FLUSH_REQUESTS_TOTAL, INFLIGHT_FLUSH_COUNT, @@ -850,7 +853,10 @@ impl RegionFlushTask { metadata: version.metadata.clone(), source, cache_manager: self.cache_manager.clone(), - max_sequence: Some(max_sequence), + // Sources have disjoint sequence intervals. LastNonNull can still + // contain Puts and retained delete barriers for the same row. + max_sequence: (version.options.merge_mode() != MergeMode::LastNonNull) + .then_some(max_sequence), sst_write_format: if flat_format { FormatType::Flat } else { @@ -920,7 +926,7 @@ struct FlatSources { encoded: SmallVec<[(EncodedRange, SequenceNumber); 4]>, } -/// Returns the max sequence and [FlatSource] for the given memtable. +/// Builds sequence-disjoint sources and their maximum input sequences. fn memtable_flat_sources( schema: SchemaRef, mem_ranges: MemtableRanges, @@ -928,190 +934,158 @@ fn memtable_flat_sources( options: &RegionOptions, field_column_start: usize, ) -> Result { - let MemtableRanges { ranges } = mem_ranges; + let components = group_by_sequence( + mem_ranges.ranges.into_values().collect(), + |range: &MemtableRange| { + let stats = range.stats(); + (stats.min_sequence, stats.max_sequence()) + }, + ); let mut flat_sources = FlatSources { sources: SmallVec::new(), encoded: SmallVec::new(), }; - - if ranges.len() == 1 { - debug!("Flushing single flat range"); - - let only_range = ranges.into_values().next().unwrap(); - let max_sequence = only_range.stats().max_sequence(); - if let Some(encoded) = only_range.encoded() { - flat_sources.encoded.push((encoded, max_sequence)); - } else { - let schema = only_range.record_batch_schema_hint().unwrap_or(schema); - let iter = only_range.build_record_batch_iter(None, None)?; - // Dedup according to append mode and merge mode. - // Even single range may have duplicate rows. - let iter = maybe_dedup_one( - options.append_mode, - options.merge_mode(), - field_column_start, - iter, - ); - flat_sources - .sources - .push((FlatSource::new_iter(schema, iter), max_sequence)); - }; - } else { - let min_flush_rows = *ENCODE_ROW_THRESHOLD; - // Calculate total rows from non-encoded ranges. - let total_rows: usize = ranges - .values() - .filter(|r| r.encoded().is_none()) - .map(|r| r.num_rows()) - .sum(); - debug!( - "Flushing multiple flat ranges, total_rows: {}, min_flush_rows: {}, num_ranges: {}", - total_rows, - min_flush_rows, - ranges.len() - ); - let mut rows_remaining = total_rows; - let mut last_iter_rows = 0; - let num_ranges = ranges.len(); - let mut input_iters = Vec::with_capacity(num_ranges); - let mut current_ranges = Vec::new(); - - let schemas = ranges - .values() - .filter(|range| range.encoded().is_none()) - .map(|range| { - ( - range - .record_batch_schema_hint() - .unwrap_or_else(|| schema.clone()), - range.num_rows() as u64, - ) - }) - .collect::>(); - let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?); - let schema = rewrite_json2_schema( - schemas.first().map(|(schema, _)| schema).unwrap_or(&schema), - &plans, - ); - - for (_range_id, range) in ranges { - if let Some(encoded) = range.encoded() { - let max_sequence = range.stats().max_sequence(); - flat_sources.encoded.push((encoded, max_sequence)); - continue; - } - - if let Some(actual) = range.record_batch_schema_hint() { - let actual = rewrite_json2_schema(&actual, &plans); - ensure!( - actual == schema, - UnexpectedSnafu { - reason: format!( - "Different schemas found in a MemtableRanges, expected: {}, actual: {}", - schema, actual, - ), - } - ) - } - - let iter = range.build_record_batch_iter(None, None)?; - let iter: BoxedRecordBatchIterator = if plans.is_empty() { - iter - } else { - let plans = plans.clone(); - Box::new(iter.map(move |batch| rewrite_json2_batch(batch?, &plans))) - }; - input_iters.push(iter); - let range_rows = range.num_rows(); - last_iter_rows += range_rows; - rows_remaining -= range_rows; - current_ranges.push(range); - - // Flush if we have enough rows, but don't flush if the remaining rows - // would be less than DEFAULT_ROW_GROUP_SIZE (to avoid small last files). - if last_iter_rows >= min_flush_rows - && (rows_remaining == 0 || rows_remaining >= DEFAULT_ROW_GROUP_SIZE) - { - debug!( - "Flush batch ready, rows: {}, min_rows: {}, num_iters: {}, remaining: {}", - last_iter_rows, - min_flush_rows, - input_iters.len(), - rows_remaining - ); - - // Calculate max_sequence from all merged ranges - let max_sequence = current_ranges - .iter() - .map(|r| r.stats().max_sequence()) - .max() - .unwrap_or(0); - let batch_size = - crate::batch_size::estimate_batch_size(current_ranges.iter().map(|range| { - let stats = range.stats(); - (stats.num_rows() as u64, stats.bytes_allocated() as u64) - })); - - let input_iters = - std::mem::replace(&mut input_iters, Vec::with_capacity(num_ranges)); - - let maybe_dedup = merge_and_dedup_with_batch_size( + let mut rows_remaining = components + .iter() + .filter(|component| component.len() != 1 || component[0].encoded().is_none()) + .flatten() + .map(MemtableRange::num_rows) + .sum::(); + let mut pending = Vec::new(); + let mut pending_rows = 0; + for component in components { + // A standalone encoded component is an ordering barrier. Do not merge + // raw ranges on either side while uploading the intervening rows as-is. + if component.len() == 1 + && let Some(encoded) = component[0].encoded() + { + if !pending.is_empty() { + flat_sources.sources.push(build_flush_source( &schema, - options.append_mode, - options.merge_mode(), + std::mem::take(&mut pending), + metadata, + options, field_column_start, - input_iters, - batch_size, - )?; - - flat_sources.sources.push(( - FlatSource::new_iter(schema.clone(), maybe_dedup), - max_sequence, - )); - last_iter_rows = 0; - current_ranges.clear(); + )?); + pending_rows = 0; } - } - - // Handle remaining iters. - if !input_iters.is_empty() { - debug!( - "Flush remaining batch, rows: {}, min_rows: {}, num_iters: {}, remaining: {}", - last_iter_rows, - min_flush_rows, - input_iters.len(), - rows_remaining - ); - - let max_sequence = current_ranges - .iter() - .map(|r| r.stats().max_sequence()) - .max() - .unwrap_or(0); - let batch_size = - crate::batch_size::estimate_batch_size(current_ranges.iter().map(|range| { - let stats = range.stats(); - (stats.num_rows() as u64, stats.bytes_allocated() as u64) - })); - - let maybe_dedup = merge_and_dedup_with_batch_size( - &schema, - options.append_mode, - options.merge_mode(), - field_column_start, - input_iters, - batch_size, - )?; - flat_sources - .sources - .push((FlatSource::new_iter(schema, maybe_dedup), max_sequence)); + .encoded + .push((encoded, component[0].stats().max_sequence())); + continue; + } + let component_rows = component.iter().map(MemtableRange::num_rows).sum::(); + pending_rows += component_rows; + rows_remaining -= component_rows; + pending.extend(component); + // Only split between disjoint components, retaining the existing + // minimum-size rule for the final source. + if pending_rows >= *ENCODE_ROW_THRESHOLD + && (rows_remaining == 0 || rows_remaining >= DEFAULT_ROW_GROUP_SIZE) + { + flat_sources.sources.push(build_flush_source( + &schema, + std::mem::take(&mut pending), + metadata, + options, + field_column_start, + )?); + pending_rows = 0; } } - + if !pending.is_empty() { + flat_sources.sources.push(build_flush_source( + &schema, + pending, + metadata, + options, + field_column_start, + )?); + } Ok(flat_sources) } +/// Builds one sequence-contiguous source, decoding encoded ranges if their +/// intervals overlap other ranges. JSON2 layouts are reconciled for this group. +fn build_flush_source( + schema: &SchemaRef, + ranges: Vec, + metadata: &RegionMetadataRef, + options: &RegionOptions, + field_column_start: usize, +) -> Result<(FlatSource, SequenceNumber)> { + let max_sequence = ranges + .iter() + .map(|range| range.stats().max_sequence()) + .max() + .unwrap_or(0); + if let [range] = ranges.as_slice() { + // A single range needs no layout reconciliation, including legacy JSON2. + let schema = range + .record_batch_schema_hint() + .unwrap_or_else(|| schema.clone()); + let iter = maybe_dedup_one( + options.append_mode, + options.merge_mode(), + field_column_start, + range.build_record_batch_iter(None, None)?, + ); + return Ok((FlatSource::new_iter(schema, iter), max_sequence)); + } + let batch_size = crate::batch_size::estimate_batch_size(ranges.iter().map(|range| { + let stats = range.stats(); + (stats.num_rows() as u64, stats.bytes_allocated() as u64) + })); + let schemas = ranges + .iter() + .map(|range| { + ( + range + .record_batch_schema_hint() + .unwrap_or_else(|| schema.clone()), + range.num_rows() as u64, + ) + }) + .collect::>(); + let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?); + let schema = rewrite_json2_schema( + schemas.first().map(|(schema, _)| schema).unwrap_or(schema), + &plans, + ); + let mut iters = Vec::with_capacity(ranges.len()); + for range in ranges { + if let Some(actual) = range.record_batch_schema_hint() { + let actual = rewrite_json2_schema(&actual, &plans); + ensure!( + actual == schema, + UnexpectedSnafu { + reason: format!( + "Different schemas found in a MemtableRanges, expected: {}, actual: {}", + schema, actual + ), + } + ); + } + let iter = range.build_record_batch_iter(None, None)?; + let iter: BoxedRecordBatchIterator = if plans.is_empty() { + iter + } else { + let plans = plans.clone(); + Box::new(iter.map(move |batch| rewrite_json2_batch(batch?, &plans))) + }; + iters.push(iter); + } + let iter = merge_and_dedup_with_batch_size( + &schema, + options.append_mode, + options.merge_mode(), + field_column_start, + iters, + batch_size, + )?; + Ok((FlatSource::new_iter(schema, iter), max_sequence)) +} + /// Merges multiple record batch iterators and applies deduplication based on the specified mode. /// /// This function is used during the flush process to combine data from multiple memtable ranges @@ -1143,7 +1117,7 @@ fn memtable_flat_sources( /// 3. If `append_mode` is false, wraps the merge iterator with a `FlatDedupIterator` that /// applies the specified merge mode: /// - `LastRow`: Removes duplicate rows, keeping only the last one -/// - `LastNonNull`: Removes duplicates but preserves the last non-null value for each field +/// - `LastNonNull`: Folds fields between delete markers, retaining every marker /// /// # Examples /// @@ -2373,7 +2347,7 @@ mod tests { 1, std::iter::once(1000), std::iter::once(Some(1.0)), - 1, + 2, ); converter.append_key_values(&kvs)?; let batch = converter.convert()?.batch; @@ -2388,6 +2362,8 @@ mod tests { )), MemtableStats { num_rows: 1, + min_sequence: id as u64 + 1, + max_sequence: id as u64 + 1, ..Default::default() }, ) @@ -2406,7 +2382,7 @@ mod tests { ranges.insert( 1, new_range( - 0, + 1, TestIterBuilder { schema: schema.clone(), batch: Some(batch), diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index 0f0a50875cc..afb546340b1 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -64,6 +64,32 @@ pub use bulk::part::{ #[cfg(any(test, feature = "test"))] pub use time_partition::filter_record_batch; +/// Groups sequence intervals into ordered, disjoint components. A component must +/// not be split during a partial merge: an excluded interval could contain an +/// intermediate row version. Bounds may conservatively include removed rows. +pub(crate) fn group_by_sequence( + mut parts: Vec, + bounds: impl Fn(&T) -> (SequenceNumber, SequenceNumber), +) -> Vec> { + parts.sort_by_key(&bounds); + let mut groups: Vec> = Vec::new(); + let mut group_max = 0; + for part in parts { + let (min, max) = bounds(&part); + match groups.last_mut() { + Some(group) if min <= group_max => { + group.push(part); + group_max = group_max.max(max); + } + _ => { + groups.push(vec![part]); + group_max = max; + } + } + } + groups +} + /// Id for memtables. /// /// Should be unique under the same region. diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index afede8904d8..7402eeb0582 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -21,7 +21,7 @@ pub mod part_reader; mod row_group_reader; use std::collections::{BTreeMap, HashSet}; -use std::sync::atomic::{AtomicI64, AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, Mutex, RwLock}; use std::time::Instant; @@ -60,6 +60,7 @@ use crate::memtable::{ AllocTracker, BoxedBatchIterator, BoxedRecordBatchIterator, EncodedBulkPart, EncodedRange, IterBuilder, KeyValues, MemScanMetrics, Memtable, MemtableBuilder, MemtableId, MemtableRange, MemtableRangeContext, MemtableRanges, MemtableRef, MemtableStats, RangesOptions, + group_by_sequence, }; use crate::read::flat_dedup::{FlatDedupIterator, FlatLastNonNull, FlatLastRow}; use crate::read::flat_merge::FlatMergeIterator; @@ -177,7 +178,7 @@ enum MergedPart { /// Result of collecting parts to merge struct CollectedParts { - /// Groups of parts ready for merging (each group has up to 16 parts) + /// Consecutive groups of parts ready for merging. groups: Vec>, } @@ -202,38 +203,18 @@ impl BulkParts { self.unordered_part.is_empty() && self.parts.is_empty() } - /// Returns true if enough parts of the same type are available to merge. - fn should_merge_parts(&self, merge_threshold: usize, encode_bytes_threshold: usize) -> bool { - let mut bulk_count = 0; - let mut encoded_count = 0; - - for wrapper in &self.parts { - if !Self::is_merge_candidate(wrapper, encode_bytes_threshold) { - continue; - } - - if wrapper.part.is_encoded() { - encoded_count += 1; - } else { - bulk_count += 1; - } - - if bulk_count >= merge_threshold || encoded_count >= merge_threshold { - return true; - } - } - - false - } - - /// Returns whether a part is small enough to benefit from merging. - fn is_merge_candidate(wrapper: &BulkPartWrapper, encode_bytes_threshold: usize) -> bool { - !wrapper.merging - && Self::is_merge_candidate_by_size( - wrapper.part.is_encoded(), - wrapper.part.estimated_size(), - encode_bytes_threshold, - ) + /// Returns a conservative scheduling hint without sorting or allocating. + /// The picker checks sequence barriers and sizes when compaction runs. + fn may_merge_parts(&self, merge_threshold: usize) -> bool { + let merge_threshold = merge_threshold.max(2); + self.parts.len() >= merge_threshold + && self + .parts + .iter() + .filter(|wrapper| !wrapper.merging) + .take(merge_threshold) + .count() + == merge_threshold } fn is_merge_candidate_by_size( @@ -250,79 +231,105 @@ impl BulkParts { || self.unordered_part.estimated_bytes() > encode_bytes_threshold } - /// Collects unmerged parts and marks them as being merged. - /// Only collects parts of types that meet the threshold. - /// Parts are grouped by row count for parallel processing. + /// Converts buffered small writes into a sorted part once a threshold is reached. + fn seal_unordered(&mut self, metadata: &RegionMetadataRef) -> Result<()> { + if let Some(part) = self.unordered_part.to_bulk_part(metadata)? { + self.parts.push(BulkPartWrapper { + part: PartToMerge::Bulk { + part, + file_id: FileId::random(), + }, + merging: false, + }); + self.unordered_part.clear(); + } + Ok(()) + } + + /// Selects consecutive sequence components without crossing unordered, busy, + /// or oversized parts. Overlapping components are indivisible, even above the size limit. + fn pick_merge_groups( + &self, + merge_threshold: usize, + max_merge_groups: usize, + encode_bytes_threshold: usize, + ) -> Vec> { + if self.parts.len() < merge_threshold.max(2) || max_merge_groups == 0 { + return Vec::new(); + } + // The unordered interval participates in overlap grouping as a barrier, + // so neither overlapping parts nor parts on opposite sides can be merged. + let intervals = self + .parts + .iter() + .enumerate() + .map(|(index, wrapper)| { + let part = &wrapper.part; + (Some(index), (part.min_sequence(), part.max_sequence())) + }) + .chain( + self.unordered_part + .sequence_bounds() + .map(|bounds| (None, bounds)), + ) + .collect(); + let components = group_by_sequence(intervals, |(_, bounds)| *bounds); + let mut groups = Vec::new(); + let mut pending = Vec::new(); + for component in components { + if groups.len() >= max_merge_groups { + break; + } + let Some(component) = component + .into_iter() + .map(|(index, _)| index) + .collect::>>() + else { + pending.clear(); + continue; + }; + let busy = component.iter().any(|&i| self.parts[i].merging); + let oversized = component.len() == 1 && { + let part = &self.parts[component[0]].part; + !Self::is_merge_candidate_by_size( + part.is_encoded(), + part.estimated_size(), + encode_bytes_threshold, + ) + }; + if busy || oversized { + pending.clear(); + continue; + } + pending.extend(component); + if pending.len() >= merge_threshold.max(2) { + groups.push(std::mem::take(&mut pending)); + } + } + groups + } + fn collect_parts_to_merge( &mut self, merge_threshold: usize, max_merge_groups: usize, encode_bytes_threshold: usize, ) -> CollectedParts { - let mut bulk_indices = Vec::new(); - let mut encoded_indices = Vec::new(); - - for (idx, wrapper) in self.parts.iter().enumerate() { - if !Self::is_merge_candidate(wrapper, encode_bytes_threshold) { - continue; - } - let num_rows = wrapper.part.num_rows(); - if wrapper.part.is_encoded() { - encoded_indices.push((idx, num_rows)); - } else { - bulk_indices.push((idx, num_rows)); - } - } - - let mut groups = Vec::new(); - - // Process bulk parts if threshold met - if bulk_indices.len() >= merge_threshold { - groups.extend(self.collect_and_group_parts( - bulk_indices, - merge_threshold, - max_merge_groups, - )); - } - - // Process encoded parts if threshold met - if encoded_indices.len() >= merge_threshold { - groups.extend(self.collect_and_group_parts( - encoded_indices, - merge_threshold, - max_merge_groups, - )); - } - - CollectedParts { groups } - } - - /// Sorts indices by row count, groups into chunks, marks as merging, and returns groups. - fn collect_and_group_parts( - &mut self, - mut indices: Vec<(usize, usize)>, - merge_threshold: usize, - max_merge_groups: usize, - ) -> Vec> { - if indices.is_empty() { - return Vec::new(); - } - - indices.sort_unstable_by_key(|(_, num_rows)| *num_rows); - indices - .chunks(merge_threshold) - .take(max_merge_groups) - .map(|chunk| { - chunk - .iter() - .map(|(idx, _)| { - let wrapper = &mut self.parts[*idx]; + let groups = self + .pick_merge_groups(merge_threshold, max_merge_groups, encode_bytes_threshold) + .into_iter() + .map(|group| { + group + .into_iter() + .map(|i| { + let wrapper = &mut self.parts[i]; wrapper.merging = true; wrapper.part.clone() }) .collect() }) - .collect() + .collect(); + CollectedParts { groups } } /// Installs merged parts and removes the original parts by file ids. @@ -431,6 +438,8 @@ pub struct BulkMemtable { compactor: Arc>, /// Dispatcher for scheduling compaction tasks compact_dispatcher: Option>, + /// Whether a write-triggered compaction is queued or running. + compacting: Arc, /// Whether the append mode is enabled append_mode: bool, /// Mode to handle duplicate rows while merging @@ -493,18 +502,8 @@ impl Memtable for BulkMemtable { bulk_parts.unordered_part.push(fragment); // Compacts unordered_part if the row or byte threshold is exceeded. - if bulk_parts.should_compact_unordered_part(self.config.encode_bytes_threshold) - && let Some(bulk_part) = - bulk_parts.unordered_part.to_bulk_part(&self.metadata)? - { - bulk_parts.parts.push(BulkPartWrapper { - part: PartToMerge::Bulk { - part: bulk_part, - file_id: FileId::random(), - }, - merging: false, - }); - bulk_parts.unordered_part.clear(); + if bulk_parts.should_compact_unordered_part(self.config.encode_bytes_threshold) { + bulk_parts.seal_unordered(&self.metadata)?; } } else { bulk_parts.parts.push(BulkPartWrapper { @@ -523,9 +522,7 @@ impl Memtable for BulkMemtable { self.update_stats(local_metrics); } - if self.should_compact() { - self.schedule_compact(); - } + self.schedule_compact(); Ok(()) } @@ -688,21 +685,12 @@ impl Memtable for BulkMemtable { return Ok(()); } - // Unified merge for all parts - let should_merge = self.parts.read().unwrap().should_merge_parts( - self.config.merge_threshold, - self.config.encode_bytes_threshold, - ); - if should_merge { - compactor.merge_parts( - &self.parts, - &self.metadata, - !self.append_mode, - self.merge_mode, - )?; - } - - Ok(()) + compactor.merge_parts( + &self.parts, + &self.metadata, + !self.append_mode, + self.merge_mode, + ) } } @@ -727,6 +715,7 @@ impl BulkMemtable { self.float_field_encoding, ))), compact_dispatcher: self.compact_dispatcher.clone(), + compacting: Arc::new(AtomicBool::new(false)), append_mode: self.append_mode, merge_mode: self.merge_mode, row_group_size: self.row_group_size, @@ -815,6 +804,7 @@ impl BulkMemtable { float_field_encoding, ))), compact_dispatcher, + compacting: Arc::new(AtomicBool::new(false)), append_mode, merge_mode, row_group_size, @@ -870,22 +860,27 @@ impl BulkMemtable { .sum() } - /// Returns whether the memtable should be compacted. - fn should_compact(&self) -> bool { - let parts = self.parts.read().unwrap(); - parts.should_merge_parts( - self.config.merge_threshold, - self.config.encode_bytes_threshold, - ) - } - - /// Schedules a compaction task using the CompactDispatcher. + /// Schedules at most one write-triggered compaction at a time. fn schedule_compact(&self) { + if self.config.max_merge_groups == 0 || self.compacting.load(Ordering::Relaxed) { + return; + } + if !self + .parts + .read() + .unwrap() + .may_merge_parts(self.config.merge_threshold) + { + return; + } + let Some(guard) = CompactTaskGuard::try_new(self.compacting.clone()) else { + return; + }; if let Some(dispatcher) = &self.compact_dispatcher { let task = MemCompactTask { metadata: self.metadata.clone(), parts: self.parts.clone(), - config: self.config.clone(), + _guard: guard, compactor: self.compactor.clone(), append_mode: self.append_mode, merge_mode: self.merge_mode, @@ -1110,6 +1105,15 @@ impl PartToMerge { } } + /// Returns the lower bound of input sequences. + fn min_sequence(&self) -> SequenceNumber { + match self { + PartToMerge::Bulk { part, .. } => part.min_sequence, + PartToMerge::Multi { part, .. } => part.min_sequence(), + PartToMerge::Encoded { part, .. } => part.metadata().min_sequence, + } + } + /// Gets the maximum sequence number of this part. fn max_sequence(&self) -> u64 { match self { @@ -1261,12 +1265,14 @@ impl MemtableCompactor { ) -> Result<()> { let start = Instant::now(); - // Collect pre-grouped parts - let collected = bulk_parts.write().unwrap().collect_parts_to_merge( - self.config.merge_threshold, - self.config.max_merge_groups, - self.config.encode_bytes_threshold, - ); + let collected = { + let mut parts = bulk_parts.write().unwrap(); + parts.collect_parts_to_merge( + self.config.merge_threshold, + self.config.max_merge_groups, + self.config.encode_bytes_threshold, + ) + }; if collected.groups.is_empty() { return Ok(()); @@ -1355,6 +1361,11 @@ impl MemtableCompactor { .map(|p| p.max_timestamp()) .max() .unwrap_or(i64::MIN); + let min_sequence = parts_to_merge + .iter() + .map(PartToMerge::min_sequence) + .min() + .unwrap_or(0); let max_sequence = parts_to_merge .iter() .map(|p| p.max_sequence()) @@ -1450,7 +1461,7 @@ impl MemtableCompactor { schema, min_timestamp, max_timestamp, - max_sequence, + min_sequence..=max_sequence, &mut metrics, )?; @@ -1476,6 +1487,7 @@ impl MemtableCompactor { batches, min_timestamp, max_timestamp, + min_sequence, max_sequence, estimated_series_count, metadata, @@ -1492,12 +1504,31 @@ impl MemtableCompactor { } } +/// Keeps the compaction flag set while queued or running and clears it on drop, +/// including failures and cancellation while waiting for a dispatcher permit. +struct CompactTaskGuard(Arc); + +impl CompactTaskGuard { + fn try_new(compacting: Arc) -> Option { + // The compactor mutex protects merging; this flag only avoids duplicate tasks. + compacting + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .ok() + .map(|_| Self(compacting)) + } +} + +impl Drop for CompactTaskGuard { + fn drop(&mut self) { + self.0.store(false, Ordering::Relaxed); + } +} + /// A memtable compact task to run in background. struct MemCompactTask { metadata: RegionMetadataRef, parts: Arc>, - /// Configuration for the bulk memtable. - config: BulkMemtableConfig, + _guard: CompactTaskGuard, /// Compactor for merging bulk parts compactor: Arc>, /// Whether the append mode is enabled @@ -1508,22 +1539,12 @@ struct MemCompactTask { impl MemCompactTask { fn compact(&self) -> Result<()> { - let mut compactor = self.compactor.lock().unwrap(); - - let should_merge = self.parts.read().unwrap().should_merge_parts( - self.config.merge_threshold, - self.config.encode_bytes_threshold, - ); - if should_merge { - compactor.merge_parts( - &self.parts, - &self.metadata, - !self.append_mode, - self.merge_mode, - )?; - } - - Ok(()) + self.compactor.lock().unwrap().merge_parts( + &self.parts, + &self.metadata, + !self.append_mode, + self.merge_mode, + ) } } @@ -1794,11 +1815,16 @@ mod tests { assert_eq!(expected_min, memtable.stats().min_sequence); assert_eq!(expected_min, memtable.min_sequence()); assert_eq!(expected_ranges, memtable.stats().num_ranges); + assert!(!memtable.compacting.load(Ordering::Relaxed)); } // The third write synchronously merges all parts without a dispatcher. // Dedup can remove the oldest rows, but their lower bound remains safe. assert_eq!(10, memtable.stats().min_sequence); assert_eq!(201, memtable.stats().max_sequence); + let ranges = memtable.ranges(None, RangesOptions::default()).unwrap(); + let merged = ranges.ranges.values().next().unwrap(); + assert!(merged.encoded().is_some()); + assert_eq!(10, merged.stats().min_sequence); let fork = memtable.fork(2, &metadata); assert!(fork.is_empty()); assert_eq!(0, fork.min_sequence()); @@ -2615,7 +2641,7 @@ mod tests { assert_eq!(5, large_part.num_rows()); memtable.write_bulk(large_part).unwrap(); - // Write another small part (2 rows) - should trigger compaction of unordered_part + // A later small write fills the unordered part to its compact threshold. let part = create_bulk_part_with_converter( "small_2", 2, @@ -2780,10 +2806,10 @@ mod tests { } #[test] - fn test_should_merge_parts_below_threshold() { + fn test_may_merge_parts_below_threshold() { let mut bulk_parts = BulkParts::default(); - // Add 7 bulk parts (below DEFAULT_MERGE_THRESHOLD of 8) + // Add one fewer part than the merge threshold. for i in 0..DEFAULT_MERGE_THRESHOLD - 1 { let part = create_bulk_part_with_converter( &format!("key_{}", i), @@ -2794,14 +2820,18 @@ mod tests { ) .unwrap(); bulk_parts.parts.push(create_bulk_part_wrapper(part)); + if i == 0 { + assert!(!bulk_parts.may_merge_parts(0)); + assert!(!bulk_parts.may_merge_parts(1)); + } } - // Should not trigger merge since we have only 7 parts - assert!(!bulk_parts.should_merge_parts(DEFAULT_MERGE_THRESHOLD, usize::MAX)); + // There are not enough parts to merge. + assert!(!bulk_parts.may_merge_parts(DEFAULT_MERGE_THRESHOLD)); } #[test] - fn test_should_merge_parts_at_threshold() { + fn test_may_merge_parts_at_threshold() { let mut bulk_parts = BulkParts::default(); let merge_threshold = 8; @@ -2819,7 +2849,7 @@ mod tests { } // Should trigger merge since we have 8 parts - assert!(bulk_parts.should_merge_parts(merge_threshold, usize::MAX)); + assert!(bulk_parts.may_merge_parts(merge_threshold)); } #[test] @@ -2839,7 +2869,7 @@ mod tests { } let max_size = bulk_parts.parts[0].part.estimated_size(); - assert!(bulk_parts.should_merge_parts(merge_threshold, max_size)); + assert!(bulk_parts.may_merge_parts(merge_threshold)); assert_eq!( 1, bulk_parts @@ -2850,7 +2880,7 @@ mod tests { } #[test] - fn test_should_merge_parts_with_merging_flag() { + fn test_may_merge_parts_with_merging_flag() { let mut bulk_parts = BulkParts::default(); let merge_threshold = 8; @@ -2868,7 +2898,7 @@ mod tests { } // Should trigger merge since we have 10 parts - assert!(bulk_parts.should_merge_parts(merge_threshold, usize::MAX)); + assert!(bulk_parts.may_merge_parts(merge_threshold)); // Mark first 3 parts as merging for wrapper in bulk_parts.parts.iter_mut().take(3) { @@ -2876,53 +2906,115 @@ mod tests { } // Now only 7 parts are available for merging, should not trigger - assert!(!bulk_parts.should_merge_parts(merge_threshold, usize::MAX)); + assert!(!bulk_parts.may_merge_parts(merge_threshold)); } #[test] - fn test_collect_parts_to_merge_grouping() { - let mut bulk_parts = BulkParts::default(); + fn test_compact_task_guard() { + let compacting = Arc::new(AtomicBool::new(false)); + let guard = CompactTaskGuard::try_new(compacting.clone()).unwrap(); + assert!(compacting.load(Ordering::Relaxed)); + assert!(CompactTaskGuard::try_new(compacting.clone()).is_none()); + drop(guard); + assert!(!compacting.load(Ordering::Relaxed)); + assert!(CompactTaskGuard::try_new(compacting.clone()).is_some()); + assert!(!compacting.load(Ordering::Relaxed)); + } - // Add 16 bulk parts with different row counts - for i in 0..16 { - let num_rows = (i % 4) + 1; // 1 to 4 rows - let timestamps: Vec = (0..num_rows) - .map(|j| 1000 + i as i64 * 100 + j as i64) - .collect(); - let values: Vec> = - (0..num_rows).map(|j| Some((i * 10 + j) as f64)).collect(); - let part = create_bulk_part_with_converter( - &format!("key_{}", i), - i as u32, - timestamps, - values, - 100 + i as u64, - ) - .unwrap(); - bulk_parts.parts.push(create_bulk_part_wrapper(part)); - } - - // Should trigger merge since we have 16 parts - assert!(bulk_parts.should_merge_parts(DEFAULT_MERGE_THRESHOLD, usize::MAX)); - - // Collect parts to merge - let collected = bulk_parts.collect_parts_to_merge( - DEFAULT_MERGE_THRESHOLD, - DEFAULT_MAX_MERGE_GROUPS, - usize::MAX, + #[test] + fn test_compact_scheduling_disabled() { + let memtable = BulkMemtable::new( + 0, + BulkMemtableConfig { + merge_threshold: 2, + max_merge_groups: 0, + ..Default::default() + }, + metadata_for_test(), + None, + Some(Arc::new(CompactDispatcher::new(0))), + false, + MergeMode::LastRow, ); + let part = + create_bulk_part_with_converter("key", 0, vec![1000], vec![Some(1.0)], 0).unwrap(); + memtable.parts.write().unwrap().parts = vec![ + create_bulk_part_wrapper(part.clone()), + create_bulk_part_wrapper(part), + ]; + memtable.schedule_compact(); + assert!(!memtable.compacting.load(Ordering::Relaxed)); + } - // Should have groups - assert!(!collected.groups.is_empty()); + #[test] + fn test_collect_parts_to_merge_sequence_components() { + let make_part = |min, max| { + let mut part = + create_bulk_part_with_converter("key", 0, vec![1000], vec![Some(1.0)], max) + .unwrap(); + part.min_sequence = min; + part + }; + let metadata = metadata_for_test(); + let encoder = BulkPartEncoder::new(metadata.clone(), DEFAULT_ROW_GROUP_SIZE).unwrap(); + // Deliberately use installation order rather than sequence order. + let mut parts = BulkParts { + parts: vec![ + create_bulk_part_wrapper(make_part(70, 70)), + create_bulk_part_wrapper(make_part(10, 30)), + BulkPartWrapper { + part: PartToMerge::Multi { + part: MultiBulkPart::from_bulk_part(make_part(35, 50), &metadata), + file_id: FileId::random(), + }, + merging: false, + }, + create_bulk_part_wrapper(make_part(60, 60)), + BulkPartWrapper { + part: PartToMerge::Encoded { + part: encoder.encode_part(&make_part(20, 40)).unwrap().unwrap(), + file_id: FileId::random(), + }, + merging: false, + }, + ], + ..Default::default() + }; + // The oversized encoded part connects three intervals transitively; + // neither the merge threshold nor its size may split that component. + assert!(parts.may_merge_parts(parts.parts.len())); + assert_eq!( + vec![vec![1, 4, 2], vec![3, 0]], + parts.pick_merge_groups(2, 2, 1) + ); + // The hint may schedule a merge even when unordered writes block all parts. + parts.unordered_part.push(make_part(0, 100)); + assert!(parts.may_merge_parts(2)); + assert!(parts.pick_merge_groups(2, 2, 1).is_empty()); + parts.unordered_part.clear(); - // All groups should have parts - for group in &collected.groups { - assert!(!group.is_empty()); - } + // Touching the unordered interval blocks the entire transitive component. + parts.unordered_part.push(make_part(50, 55)); + assert_eq!(vec![vec![3, 0]], parts.pick_merge_groups(2, 2, 1)); + parts.unordered_part.clear(); - // Total parts collected should be 16 - let total_parts: usize = collected.groups.iter().map(|g| g.len()).sum(); - assert_eq!(16, total_parts); + // Even without overlap, a group cannot straddle pending small writes. + parts.unordered_part.push(make_part(65, 65)); + assert_eq!(vec![vec![1, 4, 2]], parts.pick_merge_groups(2, 2, 1)); + parts.unordered_part.clear(); + + parts.parts[4].merging = true; + let selected = parts.collect_parts_to_merge(2, 2, 1); + assert_eq!(1, selected.groups.len()); + assert_eq!( + vec![60, 70], + selected.groups[0] + .iter() + .map(PartToMerge::max_sequence) + .collect::>() + ); + assert!(!parts.parts[1].merging); + assert!(!parts.parts[2].merging); } #[test] diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 444cc59e317..fecfceafeee 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -15,6 +15,7 @@ //! Bulk part encoder/decoder. use std::collections::{HashMap, HashSet}; +use std::ops::RangeInclusive; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -436,6 +437,11 @@ impl UnorderedPart { self.total_bytes } + /// Returns the sequence interval occupied by pending small writes. + pub(super) fn sequence_bounds(&self) -> Option<(SequenceNumber, SequenceNumber)> { + (!self.is_empty()).then_some((self.min_sequence, self.max_sequence)) + } + /// Adds a BulkPart to this unordered collection. pub fn push(&mut self, part: BulkPart) { self.total_rows += part.num_rows(); @@ -1204,7 +1210,7 @@ impl EncodedBulkPart { num_rows: meta.num_rows, num_ranges: 1, max_sequence: meta.max_sequence, - min_sequence: 0, + min_sequence: meta.min_sequence, series_count: meta.num_series as usize, } } @@ -1291,6 +1297,8 @@ pub struct BulkPartMeta { pub region_metadata: RegionMetadataRef, /// Number of series. pub num_series: u64, + /// Lower bound of input sequences, including rows removed by deduplication. + pub min_sequence: SequenceNumber, /// Maximum sequence number in part. pub max_sequence: u64, } @@ -1356,7 +1364,7 @@ impl BulkPartEncoder { arrow_schema: SchemaRef, min_timestamp: i64, max_timestamp: i64, - max_sequence: u64, + sequence_bounds: RangeInclusive, metrics: &mut BulkPartEncodeMetrics, ) -> Result> { let mut buf = Vec::with_capacity(4096); @@ -1408,7 +1416,8 @@ impl BulkPartEncoder { parquet_metadata, region_metadata: self.metadata.clone(), num_series, - max_sequence, + min_sequence: *sequence_bounds.start(), + max_sequence: *sequence_bounds.end(), }, schema: arrow_schema, })) @@ -1444,6 +1453,7 @@ impl BulkPartEncoder { region_metadata: self.metadata.clone(), num_series: part.estimated_series_count() as u64, max_sequence: part.sequence, + min_sequence: part.min_sequence, }, schema: arrow_schema, })) @@ -1653,6 +1663,8 @@ pub struct MultiBulkPart { max_timestamp: i64, /// Min timestamp in part. min_timestamp: i64, + /// Lower bound of input sequences, including rows removed by deduplication. + min_sequence: SequenceNumber, /// Max sequence number in part. max_sequence: SequenceNumber, /// Number of series. @@ -1676,6 +1688,7 @@ impl MultiBulkPart { total_rows: num_rows, max_timestamp: part.max_timestamp, min_timestamp: part.min_timestamp, + min_sequence: part.min_sequence, max_sequence: part.sequence, series_count, batch_stats, @@ -1688,6 +1701,7 @@ impl MultiBulkPart { /// * `batches` - Ordered record batches /// * `min_timestamp` - Minimum timestamp across all batches /// * `max_timestamp` - Maximum timestamp across all batches + /// * `min_sequence` - Lower bound of input sequences /// * `max_sequence` - Maximum sequence number across all batches /// * `series_count` - Number of series in the batches /// * `metadata` - Region metadata for computing batch statistics @@ -1698,6 +1712,7 @@ impl MultiBulkPart { batches: Vec, min_timestamp: i64, max_timestamp: i64, + min_sequence: SequenceNumber, max_sequence: SequenceNumber, series_count: usize, metadata: &RegionMetadata, @@ -1712,6 +1727,7 @@ impl MultiBulkPart { total_rows, max_timestamp, min_timestamp, + min_sequence, max_sequence, series_count, batch_stats, @@ -1737,6 +1753,11 @@ impl MultiBulkPart { self.max_timestamp } + /// Returns the lower bound of input sequences. + pub fn min_sequence(&self) -> SequenceNumber { + self.min_sequence + } + /// Returns the maximum sequence number. pub fn max_sequence(&self) -> SequenceNumber { self.max_sequence @@ -1827,7 +1848,7 @@ impl MultiBulkPart { num_rows: self.num_rows(), num_ranges: 1, max_sequence: self.max_sequence, - min_sequence: 0, + min_sequence: self.min_sequence, series_count: self.series_count, } } @@ -3147,12 +3168,14 @@ mod tests { let mut all_batches = Vec::new(); let mut min_ts = i64::MAX; let mut max_ts = i64::MIN; + let mut min_seq = u64::MAX; let mut max_seq = 0u64; for inputs in groups { let part = build_converted_bulk_part(inputs); min_ts = min_ts.min(part.min_timestamp); max_ts = max_ts.max(part.max_timestamp); + min_seq = min_seq.min(part.min_sequence); max_seq = max_seq.max(part.sequence); all_batches.push(part.batch); } @@ -3161,6 +3184,7 @@ mod tests { all_batches, min_ts, max_ts, + min_seq, max_seq, groups.len(), &metadata, diff --git a/src/mito2/src/read/flat_dedup.rs b/src/mito2/src/read/flat_dedup.rs index 64c78895a99..b166d4c7f32 100644 --- a/src/mito2/src/read/flat_dedup.rs +++ b/src/mito2/src/read/flat_dedup.rs @@ -160,8 +160,8 @@ impl Drop for FlatDedupReader { /// Strategy to remove duplicate rows from sorted record batches. pub trait RecordBatchDedupStrategy: Send { /// Pushes a batch to the dedup strategy. - /// Returns a batch if the strategy ensures there is no duplications based on - /// the input batch. + /// Returns rows that no longer need to be updated by subsequent input. + /// A strategy for partial merges may retain duplicates separated by Deletes. fn push_batch( &mut self, batch: RecordBatch, @@ -297,11 +297,11 @@ impl RecordBatchDedupStrategy for FlatLastRow { pub struct FlatLastNonNull { /// The start index of field columns: field_column_start: usize, - /// Filter deleted rows. + /// Filter deleted rows; otherwise preserve every delete barrier. filter_deleted: bool, /// Buffered batch to check whether the next batch have duplicated rows with this batch. /// Fields in the last row of this batch may be updated by the next batch. - /// The buffered batch should contain no duplication. + /// In preserve mode, duplicate keys may be separated by delete markers. buffer: Option, /// Whether the last row range contains a delete operation. /// If so, we don't need to update null fields. @@ -326,7 +326,7 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { // If the buffer is None, dedup the batch, put the batch into the buffer and return. // There is no previous batch with the same key, we can pass contains_delete as false. let (record_batch, contains_delete) = - Self::dedup_one_batch(batch, self.field_column_start, false)?; + Self::dedup_one_batch(batch, self.field_column_start, false, self.filter_deleted)?; metrics.num_unselected_rows += row_before_dedup - record_batch.num_rows(); self.buffer = BatchLastRow::try_new(record_batch); self.contains_delete = contains_delete; @@ -341,7 +341,7 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { // Dedup the batch. // There is no previous batch with the same key, we can pass contains_delete as false. let (record_batch, contains_delete) = - Self::dedup_one_batch(batch, self.field_column_start, false)?; + Self::dedup_one_batch(batch, self.field_column_start, false, self.filter_deleted)?; metrics.num_unselected_rows += row_before_dedup - record_batch.num_rows(); debug_assert!(record_batch.num_rows() > 0); self.buffer = BatchLastRow::try_new(record_batch); @@ -369,8 +369,12 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { let merged = concat_batches(&schema, &[last_row, batch]).context(ComputeArrowSnafu)?; let merged_row_count = merged.num_rows(); // Dedup the merged batch and update the buffer. - let (record_batch, contains_delete) = - Self::dedup_one_batch(merged, self.field_column_start, self.contains_delete)?; + let (record_batch, contains_delete) = Self::dedup_one_batch( + merged, + self.field_column_start, + self.contains_delete, + self.filter_deleted, + )?; metrics.num_unselected_rows += merged_row_count - record_batch.num_rows(); debug_assert!(record_batch.num_rows() > 0); self.buffer = BatchLastRow::try_new(record_batch); @@ -397,7 +401,11 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { } impl FlatLastNonNull { - /// Creates a new strategy with the given `filter_deleted` flag. + /// Creates a strategy with the given `filter_deleted` flag. When false, + /// every delete barrier is preserved and Puts are folded only between + /// barriers so later reads cannot fill fields from a deleted row. + /// Inputs for partial merges must form a sequence-contiguous group to avoid + /// folding fields across an intermediate version left outside the merge. pub fn new(field_column_start: usize, filter_deleted: bool) -> Self { Self { field_column_start, @@ -413,6 +421,7 @@ impl FlatLastNonNull { batch: RecordBatch, field_column_start: usize, prev_batch_contains_delete: bool, + filter_deleted: bool, ) -> Result<(RecordBatch, bool)> { // Get op type array for checking delete operations let op_type_column = batch @@ -460,6 +469,7 @@ impl FlatLastNonNull { field_column_start, op_types, prev_batch_contains_delete, + filter_deleted, ) } @@ -471,9 +481,29 @@ impl FlatLastNonNull { field_column_start: usize, op_types: &UInt8Array, first_range_contains_delete: bool, + filter_deleted: bool, ) -> Result<(RecordBatch, bool)> { - let ranges = partitions.ranges(); - let contains_delete = Self::last_range_has_delete(&ranges, op_types); + let mut ranges = partitions.ranges(); + let contains_delete = filter_deleted && Self::last_range_has_delete(&ranges, op_types); + if !filter_deleted && op_types.values().contains(&(OpType::Delete as u8)) { + let mut segments = Vec::with_capacity(ranges.len()); + for range in ranges { + let mut start = range.start; + for i in range.clone() { + if op_types.value(i) == OpType::Delete as u8 { + if start < i { + segments.push(start..i); + } + segments.push(i..i + 1); + start = i + 1; + } + } + if start < range.end { + segments.push(start..range.end); + } + } + ranges = segments; + } // Each range at least has 1 row. let num_duplications: usize = ranges.iter().map(|r| r.end - r.start - 1).sum(); @@ -498,7 +528,7 @@ impl FlatLastNonNull { &ranges, column, op_types, - first_range_contains_delete, + first_range_contains_delete && filter_deleted, ); take(column, &field_indices, take_options.clone()).context(ComputeArrowSnafu) } else { @@ -1239,52 +1269,6 @@ mod tests { check_record_batches_equal(&expected_filter_deleted, &result); assert_eq!(6, dedup_iter.metrics.num_unselected_rows); assert_eq!(2, dedup_iter.metrics.num_deleted_rows); - - // Test with filter_deleted = false - let expected_no_filter = vec![ - new_record_batch_multi_fields( - &[b"k1"], - &[1], - &[13], - &[OpType::Put], - &[(Some(11), Some(11))], - ), - new_record_batch_multi_fields( - &[b"k1", b"k1", b"k1"], - &[2, 3, 4], - &[11, 13, 13], - &[OpType::Put, OpType::Put, OpType::Delete], - &[(Some(12), Some(22)), (Some(13), None), (None, Some(14))], - ), - new_record_batch_multi_fields( - &[b"k2"], - &[1], - &[20], - &[OpType::Put], - &[(Some(101), Some(101))], - ), - new_record_batch_multi_fields( - &[b"k2"], - &[2], - &[20], - &[OpType::Delete], - &[(None, None)], - ), - new_record_batch_multi_fields( - &[b"k3"], - &[2], - &[20], - &[OpType::Put], - &[(Some(202), Some(202))], - ), - ]; - - let iter = input.clone().into_iter().map(Ok); - let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastNonNull::new(1, false)); - let result = collect_iterator_results(&mut dedup_iter); - check_record_batches_equal(&expected_no_filter, &result); - assert_eq!(4, dedup_iter.metrics.num_unselected_rows); - assert_eq!(0, dedup_iter.metrics.num_deleted_rows); } #[test] @@ -1588,6 +1572,69 @@ mod tests { ); } + #[test] + fn test_last_non_null_preserves_delete_barriers() { + let input = new_record_batch_multi_fields( + &[b"k1".as_slice(); 9], + &[1; 9], + &[9, 8, 7, 6, 5, 4, 3, 2, 1], + &[ + OpType::Put, + OpType::Put, + OpType::Delete, + OpType::Delete, + OpType::Put, + OpType::Put, + OpType::Delete, + OpType::Put, + OpType::Put, + ], + &[ + (None, Some(9)), + (Some(8), None), + (None, None), + (None, None), + (None, Some(5)), + (Some(4), None), + (None, None), + (None, None), + (Some(1), Some(1)), + ], + ); + let expected = new_record_batch_multi_fields( + &[b"k1".as_slice(); 6], + &[1; 6], + &[9, 7, 6, 5, 3, 2], + &[ + OpType::Put, + OpType::Delete, + OpType::Delete, + OpType::Put, + OpType::Delete, + OpType::Put, + ], + &[ + (Some(8), Some(9)), + (None, None), + (None, None), + (Some(4), Some(5)), + (None, None), + (Some(1), Some(1)), + ], + ); + // Check both intra-batch partitioning and barriers split across batches. + for batch_size in [1, input.num_rows()] { + let batches = (0..input.num_rows()) + .step_by(batch_size) + .map(|start| Ok(input.slice(start, batch_size.min(input.num_rows() - start)))); + let output = FlatDedupIterator::new(batches, FlatLastNonNull::new(1, false)) + .collect::>>() + .unwrap(); + let actual = concat_batches(&input.schema(), &output).unwrap(); + assert_eq!(expected, actual); + } + } + #[test] fn test_flat_last_non_null_strategy_delete_middle() { let input = vec![ From c1a0ca9595c3a26a4273e85944927215383d7b7e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 12 Oct 2026 13:36:55 +0000 Subject: [PATCH 19/46] ci: benchmark the object store WAL against S3 from an EC2 runner (#9536) * ci: benchmark the object store WAL against S3 from an EC2 runner Add a workflow_dispatch workflow that loads the o11ybench agent observability dataset into a fresh GreptimeDB container per cell and compares three WAL configurations that differ only in [wal]: Raft Engine with sync_write on the instance store, and the object store WAL in durable and enqueued mode, with the data files on S3 in the runner's region for all three. Rounds alternate the configuration order. A summary reports rows per second, D/R and E/R as the geometric mean over rounds, append acknowledgement p50 and p99, and the object store WAL counters. A smoke mode runs the same chain on ubuntu-latest against the MinIO fixture. Signed-off-by: jeremyhi * ci: scrape WAL metrics on failed loads and keep logic out of the benchmark YAML Scrape the object store WAL metrics before stopping a cell's server whether the load succeeded or failed, so failed cells keep their stall, poison and create failure counters. Build the summary from the plan so that cells which never ran are listed and leave their ratios undefined. Move budget, region, runtime image, manifest and credential scan logic from the workflow into the helper scripts, and drive the cell lifecycle in the tests through fake docker and curl binaries. Signed-off-by: jeremyhi * ci: drop redundant required flags from benchmark inputs Signed-off-by: jeremyhi * ci: fail a benchmark cell whose metrics scrape fails and zero untouched WAL counters A load that succeeds without its metrics scrape now fails the cell with exit code 4 and records metrics_scraped in cell.json. The summary reports an event counter that never registered as zero when the cell is an object store WAL cell whose scrape succeeded and created objects, and as unknown otherwise. Planned cells that never ran are read through the same path as recorded ones. Signed-off-by: jeremyhi * ci: register the S3 benchmark workflow for dispatch from its branch A workflow_dispatch workflow can only be dispatched once GitHub has registered it, which happens on the default branch or when one of its triggers fires. A push trigger limited to this branch registers it; validate runs only for workflow_dispatch, so a push run skips every job. Signed-off-by: jeremyhi * ci: delete the benchmark S3 prefix one object at a time boto3 now sends a CRC checksum instead of Content-MD5 with DeleteObjects, which the MinIO fixture of the smoke run rejects. DeleteObject needs neither and the prefix of one run holds at most a few thousand objects. Signed-off-by: jeremyhi * ci: default the S3 benchmark to c6id.2xlarge c7id is not offered in the runner region ap-southeast-1, where DescribeInstanceTypes rejects it. c6id is the compute optimized x86 family with a local NVMe instance store that the region offers, with the same 8 vCPU and 16 GiB at 2xlarge. Signed-off-by: jeremyhi * ci: drop the registration trigger of the S3 benchmark workflow The workflow is registered now, so dispatches no longer need a push trigger on its branch. Signed-off-by: jeremyhi * ci: fail an S3 benchmark run on EC2 without a local instance store Without an instance store the Raft Engine WAL of the R cells silently stayed on the root EBS volume, so the run measured a different baseline. Only a smoke run may keep it on the root volume now. Signed-off-by: jeremyhi * ci: let the instance store check fail its workflow step The step piped wal-disk into tee without pipefail, so a rejected host still passed. wal-disk now writes its record itself and the step runs it alone. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- .github/scripts/oswal-s3-bench-run.sh | 179 ++++++++++++ .github/scripts/oswal-s3-bench-summary.py | 162 +++++++++++ .github/scripts/oswal-s3-bench.py | 222 +++++++++++++++ .github/workflows/oswal-s3-bench.yml | 229 ++++++++++++++++ tests/perf/test_oswal_s3_bench_scripts.py | 318 ++++++++++++++++++++++ 5 files changed, 1110 insertions(+) create mode 100644 .github/scripts/oswal-s3-bench-run.sh create mode 100644 .github/scripts/oswal-s3-bench-summary.py create mode 100644 .github/scripts/oswal-s3-bench.py create mode 100644 .github/workflows/oswal-s3-bench.yml create mode 100644 tests/perf/test_oswal_s3_bench_scripts.py diff --git a/.github/scripts/oswal-s3-bench-run.sh b/.github/scripts/oswal-s3-bench-run.sh new file mode 100644 index 00000000000..a9835c1e5f1 --- /dev/null +++ b/.github/scripts/oswal-s3-bench-run.sh @@ -0,0 +1,179 @@ +#!/usr/bin/env bash +# Host side of oswal-s3-bench.yml: one subcommand per step; each cell of $PLAN gets a fresh GreptimeDB. +set -euo pipefail + +here=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) + +# Same isolation as the o11ybench runtime's own docker run. +runtime() { + docker run --rm --user "$(id -u):$(id -g)" --label "o11ybench.owner=$OWNER" \ + --security-opt no-new-privileges --cap-drop ALL "$@" +} + +check_runtime() { + local image=$1 resolved + [[ "$image" == public.ecr.aws/* ]] || { echo 'The runtime image must come from AWS ECR Public' >&2; return 1; } + bash "$here/pull-benchmark-image.sh" "$image" + resolved=$(docker image inspect "$image" | jq -er '.[0].RepoDigests[0]') + docker run --rm --entrypoint sh "$resolved" -c \ + 'test -f scripts/agent-observability/greptimedb.sh && test -f scripts/generate_agent_observations.py' \ + || { echo 'The runtime image lacks the agent observability loader or generator' >&2; return 1; } + printf 'runtime_image=%s\n' "$resolved" >> "${GITHUB_OUTPUT:-/dev/stdout}" +} + +pull_images() { + : "${RESULTS:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_TAG:?}" + mkdir -p "$RESULTS/images" + bash "$here/pull-benchmark-image.sh" "$RUNTIME_IMAGE" + bash "$here/pull-benchmark-image.sh" "greptime/greptimedb:$GREPTIMEDB_TAG" + docker image inspect "$RUNTIME_IMAGE" > "$RESULTS/images/runtime.json" + docker image inspect "greptime/greptimedb:$GREPTIMEDB_TAG" > "$RESULTS/images/greptimedb.json" +} + +wal_disk() { + : "${RESULTS:?}" + local root=$1 device + device=$(lsblk -dnpo NAME,MODEL | awk '/Instance Storage/ {print $1; exit}') + # Only a smoke run may keep the Raft Engine WAL on the root volume. + [[ -n "$device" || "${SMOKE:-}" == true ]] || { echo 'No local NVMe instance store; choose an instance type with one' >&2; return 1; } + sudo -n mkdir -p "$root" + if [[ -n "$device" ]]; then + sudo -n umount "$device" 2>/dev/null || true + sudo -n mkfs.ext4 -q -F "$device" + sudo -n mount -o noatime "$device" "$root" + fi + sudo -n chown "$(id -u):$(id -g)" "$root" + findmnt -no SOURCE,FSTYPE -T "$root" | tee "$RESULTS/wal-disk.txt" +} + +generate() { + : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" + local args=() arg + while IFS= read -r arg; do args+=("$arg"); done < <(jq -r '.generator_args[]' "$PLAN") + mkdir -p "$CORPUS_DIR" + runtime --name "$OWNER-generator" --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus" \ + "$RUNTIME_IMAGE" scripts/generate_agent_observations.py "${args[@]}" \ + --output /work/corpus/agent_observations.jsonl \ + --summary-output /work/corpus/summary.json --evidence-output /work/corpus/evidence.json \ + > "$RESULTS/generator.log" 2>&1 + jq -e --argjson rows "$(jq .rows "$PLAN")" '.line_count == $rows' "$CORPUS_DIR/summary.json" >/dev/null +} + +# Scrape once, after the load succeeded or failed, while the server still runs. +scrape_metrics() { + curl --noproxy '*' -fsS --max-time 30 http://127.0.0.1:4000/metrics > "$dir/metrics.all" || { rm -f "$dir/metrics.all"; return 1; } + grep -E '^(# (HELP|TYPE) )?greptime_logstore_object_store_wal_' "$dir/metrics.all" > "$dir/metrics.prom" || true + rm -f "$dir/metrics.all" + metrics_scraped=true +} + +finish() { + local rc=$? + trap - EXIT + set +e + if [[ -n "$cid" ]]; then + # A load that succeeded without its metrics is not a measurement. + scrape_metrics || { echo "$id: metrics scrape failed" >&2; [[ $rc != 0 ]] || rc=4; } + docker stop -t 120 "$cid" >/dev/null + docker logs "$cid" > "$dir/server.log" 2>&1 + docker rm -fv "$cid" >/dev/null + fi + [[ -z "$wal_dir" ]] || sudo -n rm -rf "$wal_dir" + jq -n --argjson spec "$spec" --arg device "$wal_device" --arg version "$version" \ + --arg load_seconds "$load_seconds" --argjson exit_code "$rc" --argjson scraped "$metrics_scraped" \ + '$spec + {wal_device: $device, version: $version, exit_code: $exit_code, metrics_scraped: $scraped, + load_wall_seconds: ($load_seconds | tonumber? // null)}' > "$dir/cell.json" + exit "$rc" +} + +# One cell in its own process, so errexit holds inside it. The cell's state is +# global because the EXIT trap runs after this function returns. +cell() { + : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGE:?}" + : "${RUN_PREFIX:?}" "${WAL_ROOT:?}" "${S3_BUCKET:?}" "${S3_REGION:?}" "${AWS_ACCESS_KEY_ID:?}" "${AWS_SECRET_ACCESS_KEY:?}" + spec=$1 cid= wal_dir= version= load_seconds= metrics_scraped=false id= + local concurrency name setting started ended + id=$(jq -r .id <<< "$spec"); concurrency=$(jq -r .concurrency <<< "$spec") + dir="$RESULTS/cells/$id"; name="$OWNER-$id" + mkdir -p "$dir" + # The root holds data/ and the WAL's default prefix wal/ side by side. + local settings=(STORAGE__TYPE=S3 "STORAGE__BUCKET=$S3_BUCKET" "STORAGE__REGION=$S3_REGION" + "STORAGE__ROOT=$RUN_PREFIX/$id") + [[ -z "${S3_CONTAINER_ENDPOINT:-}" ]] || settings+=("STORAGE__ENDPOINT=$S3_CONTAINER_ENDPOINT") + while IFS= read -r setting; do settings+=("$setting"); done \ + < <(jq -r '.wal | to_entries[] | "WAL__\(.key)=\(.value)"' <<< "$spec") + local docker_env=() mounts=() + for setting in "${settings[@]}"; do + docker_env+=(-e "GREPTIMEDB_STANDALONE__$setting") + printf 'GREPTIMEDB_STANDALONE__%s\n' "$setting" >> "$dir/config.env" + done + # Pass the two keys by name: their values never reach argv, a file or the artifact. + export GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID + export GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY + docker_env+=(-e GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID -e GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY) + printf 'GREPTIMEDB_STANDALONE__STORAGE__%s=\n' ACCESS_KEY_ID SECRET_ACCESS_KEY >> "$dir/config.env" + if jq -e '.wal.DIR' <<< "$spec" >/dev/null; then + wal_dir="$WAL_ROOT/$id"; mkdir -p "$wal_dir" + mounts=(--mount "type=bind,src=$wal_dir,dst=$(jq -r .wal.DIR <<< "$spec")") + wal_device=$(findmnt -no SOURCE -T "$wal_dir") + else + wal_device="s3://$S3_BUCKET/$RUN_PREFIX/$id/wal" + fi + trap finish EXIT + cid=$(docker run -d --init --name "$name" --label "o11ybench.owner=$OWNER" \ + --add-host host.docker.internal:host-gateway \ + -p 127.0.0.1:4000:4000 -p 127.0.0.1:4001:4001 "${docker_env[@]}" ${mounts[@]+"${mounts[@]}"} \ + "$GREPTIMEDB_IMAGE" standalone start --http-addr=0.0.0.0:4000 --grpc-bind-addr=0.0.0.0:4001) + local deadline=$((SECONDS + 300)) + until curl --noproxy '*' -fsS --max-time 5 http://127.0.0.1:4000/health >/dev/null 2>&1; do + [[ $(docker inspect --format '{{.State.Running}}' "$cid") == true ]] || { echo "$id: server exited" >&2; exit 1; } + ((SECONDS < deadline)) || { echo "$id: server not healthy in 300s" >&2; exit 1; } + sleep 2 + done + version=$(docker exec "$cid" greptime --version | sed -n 's/^version: //p') + [[ -n "$version" ]] + started=$(date +%s.%N) + # Only the load phase, run exactly as run-agent-observability.sh runs it. + runtime --name "$name-load" --network host \ + --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus,readonly" --mount "type=bind,src=$dir,dst=/results" \ + -e GRPC_ENDPOINT=127.0.0.1:4001 -e "DATA_PROFILE=$(jq -r .data_profile "$PLAN")" -e RESULT_NAME=greptimedb-1 \ + -e "ENDPOINT=http://127.0.0.1:4000/v1/sql?db=public&format=greptimedb_v1" \ + -e "TARGET_IMAGE=$GREPTIMEDB_IMAGE" -e "TARGET_VERSION=$version" -e "TARGET_DIGEST=${GREPTIMEDB_IMAGE##*@}" \ + -e "TARGET_ID=$name" -e TARGET_HOST=127.0.0.1 -e PREPARE_WORKERS= -e "MAX_IN_FLIGHT=$concurrency" \ + -e "GREPTIMEDB_BATCH_SIZE=$(jq -r .batch_rows "$PLAN")" --entrypoint /bin/bash \ + "$RUNTIME_IMAGE" scripts/agent-observability/greptimedb.sh load > "$dir/load.log" 2>&1 + ended=$(date +%s.%N) + load_seconds=$(awk -v a="$started" -v b="$ended" 'BEGIN {printf "%.3f", b - a}') + jq -e --argjson rows "$(jq .rows "$PLAN")" '.count_before == 0 and .count_after == $rows' \ + "$dir/greptimedb-1/load-result.json" >/dev/null +} + +scrub() { + : "${RESULTS:?}" + [[ -n "${AWS_ACCESS_KEY_ID:-}" && -n "${AWS_SECRET_ACCESS_KEY:-}" && -d "$RESULTS" ]] || return 0 + if grep -rqF -e "$AWS_ACCESS_KEY_ID" -e "$AWS_SECRET_ACCESS_KEY" "$RESULTS"; then + echo '::error::A credential value appears in the results; removing them instead of uploading' >&2 + rm -rf "$RESULTS" + return 1 + fi +} + +case "${1:-}" in + check-runtime) check_runtime "${2:?}" ;; + pull) pull_images ;; + wal-disk) wal_disk "${2:?}" ;; + generate) generate ;; + cell) cell "${2:?}" ;; + cells) + : "${PLAN:?}" + failed=0 + while IFS= read -r spec; do + id=$(jq -r .id <<< "$spec") + echo "::group::cell $id" + bash "${BASH_SOURCE[0]}" cell "$spec" < /dev/null || { echo "::error::cell $id failed"; failed=1; } + echo '::endgroup::' + done < <(jq -c '.cells[]' "$PLAN") + exit "$failed" ;; + scrub) scrub ;; + *) echo "usage: $0 check-runtime IMAGE | pull | wal-disk DIR | generate | cells | scrub" >&2; exit 2 ;; +esac diff --git a/.github/scripts/oswal-s3-bench-summary.py b/.github/scripts/oswal-s3-bench-summary.py new file mode 100644 index 00000000000..67186d422da --- /dev/null +++ b/.github/scripts/oswal-s3-bench-summary.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +"""Summarize the object store WAL S3 benchmark as Markdown and summary.json.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +from pathlib import Path + +PREFIX = "greptime_logstore_object_store_wal_" +COUNTERS = ("created_objects_total", "stalled_appends_total", "poisoned_total", + "create_failures_total", "create_conflicts_total") +SAMPLE = re.compile(r'([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{([^}]*)\})?\s+(\S+)') + + +def parse_metrics(text: str) -> dict: + """Map each unlabeled sample to its value and each histogram to its buckets.""" + samples, buckets = {}, {} + for line in text.splitlines(): + match = SAMPLE.fullmatch(line.strip()) + if not match or not match.group(1).startswith(PREFIX): + continue + name, labels, value = match.group(1)[len(PREFIX):], match.group(2) or "", float(match.group(3)) + le = re.search(r'le="([^"]+)"', labels) + if name.endswith("_bucket") and le: + buckets.setdefault(name[:-len("_bucket")], []).append((float(le.group(1)), value)) + else: + samples[name] = value + return {"samples": samples, "buckets": {name: sorted(items) for name, items in buckets.items()}} + + +def quantile(q: float, buckets: list) -> float | None: + """Prometheus histogram_quantile over cumulative (upper bound, count) buckets.""" + if not buckets or buckets[-1][1] == 0: + return None + rank = q * buckets[-1][1] + lower, below = 0.0, 0.0 + for upper, count in buckets: + if count >= rank: + if math.isinf(upper): + return lower + return lower + (upper - lower) * (rank - below) / (count - below) if count > below else upper + lower, below = upper, count + return lower + + +def geomean(values: list) -> float | None: + return math.exp(sum(map(math.log, values)) / len(values)) if values and all(v > 0 for v in values) else None + + +def load_cell(path: Path, spec: dict | None = None) -> dict: + """One cell's row; a planned cell that never recorded a result has no values, so no ratio uses its round.""" + recorded = path / "cell.json" + cell = json.loads(recorded.read_text()) if recorded.exists() else spec | {"exit_code": None, "missing": True} + result_path = path / "greptimedb-1" / "load-result.json" + result = json.loads(result_path.read_text()) if result_path.exists() else {} + metrics_path = path / "metrics.prom" + metrics = parse_metrics(metrics_path.read_text() if metrics_path.exists() else "") + ack = metrics["buckets"].get("append_ack_seconds", []) + load_seconds = result.get("load_seconds") + row = { + "id": cell["id"], "round": cell["round"], "concurrency": cell["concurrency"], + "config": cell["config"], "exit_code": cell["exit_code"], "wal_device": cell.get("wal_device"), + "missing": cell.get("missing", False), + "version": cell.get("version"), "load_wall_seconds": cell.get("load_wall_seconds"), + "rows": result.get("count_after"), "load_seconds": load_seconds, + "rows_per_second": result.get("rows_per_second") if cell["exit_code"] == 0 else None, + "append_ack_p50_seconds": quantile(0.5, ack), "append_ack_p99_seconds": quantile(0.99, ack), + } + # The counters register on their first event, so a scrape of an object store WAL + # that created objects omits the ones that never moved: those are zero, not unknown. + active = (cell.get("metrics_scraped") and cell.get("wal", {}).get("PROVIDER") == "experimental_object_store" + and metrics["samples"].get("created_objects_total")) + for counter in COUNTERS: + row[counter] = metrics["samples"].get(counter, 0 if active else None) + created = row["created_objects_total"] + row["creates_per_minute"] = created * 60 / load_seconds if created is not None and load_seconds else None + return row + + +def ratios(rows: list) -> list: + """D/R and E/R per round and their geometric mean over the rounds, per concurrency.""" + speed = {(r["round"], r["concurrency"], r["config"]): r["rows_per_second"] for r in rows} + result = [] + for concurrency in sorted({r["concurrency"] for r in rows}): + rounds = sorted({r["round"] for r in rows if r["concurrency"] == concurrency}) + for config in ("D", "E"): + per_round = [] + for round_number in rounds: + base, value = speed.get((round_number, concurrency, "R")), speed.get((round_number, concurrency, config)) + per_round.append(value / base if base and value else None) + complete = [value for value in per_round if value is not None] + result.append({"concurrency": concurrency, "ratio": f"{config}/R", "per_round": per_round, + "geomean": geomean(complete) if len(complete) == len(per_round) else None}) + return result + + +def fmt(value, digits=0, scale=1.0) -> str: + return "—" if value is None else f"{value * scale:,.{digits}f}" + + +def render(run: dict, rows: list, pairs: list) -> str: + lines = [f"## Object store WAL against S3 ({'smoke, MinIO' if run.get('smoke') else run.get('instance_type')})", ""] + lines += [f"GreptimeDB `{run.get('greptimedb_image')}` ({next((r['version'] for r in rows if r['version']), '—')}); " + f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}.", ""] + lines += ["| Cell | Rows/s | Load s | Wall s | Rows | Ack p50 ms | Ack p99 ms | Creates | Creates/min " + "| Stalled | Poisoned | Create failures | Conflicts | WAL on |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |"] + for r in rows: + cell = f"{r['id']} (not run)" if r["missing"] else r["id"] if r["exit_code"] == 0 \ + else f"{r['id']} (failed, exit {r['exit_code']})" + lines.append(" | ".join([ + f"| {cell}", fmt(r["rows_per_second"]), fmt(r["load_seconds"], 1), fmt(r["load_wall_seconds"], 1), + fmt(r["rows"]), fmt(r["append_ack_p50_seconds"], 1, 1000), fmt(r["append_ack_p99_seconds"], 1, 1000), + *(fmt(r[c]) for c in COUNTERS[:1]), fmt(r["creates_per_minute"], 1), + *(fmt(r[c]) for c in COUNTERS[1:]), f"`{r['wal_device'] or '—'}` |"])) + lines += ["", "| Concurrency | Ratio | Per round | Geometric mean |", "| ---: | --- | --- | ---: |"] + for p in pairs: + lines.append(f"| {p['concurrency']} | {p['ratio']} | {', '.join(fmt(v, 3) for v in p['per_round'])} " + f"| {fmt(p['geomean'], 3)} |") + return "\n".join(lines) + "\n" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, required=True) + args = parser.parse_args() + args.root.mkdir(parents=True, exist_ok=True) + run_path = args.root / "run.json" + run = json.loads(run_path.read_text()) if run_path.exists() else {} + plan_path = args.root / "plan.json" + planned = json.loads(plan_path.read_text())["cells"] if plan_path.exists() else [] + # The plan, not what happened to run, defines the expected cells and rounds. + rows = [load_cell(args.root / "cells" / spec["id"], spec) for spec in planned] + rows += [load_cell(path) for path in sorted((args.root / "cells").glob("*")) + if path.name not in {spec["id"] for spec in planned} and (path / "cell.json").exists()] + pairs = ratios(rows) + missing = [r["id"] for r in rows if r["missing"]] + (args.root / "summary.json").write_text(json.dumps( + {"run": run, "planned_cells": len(planned), "missing_cells": missing, "cells": rows, "ratios": pairs}, + indent=2) + "\n") + print(render(run, rows, pairs) if rows else "No benchmark cell completed.\n") + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/oswal-s3-bench.py b/.github/scripts/oswal-s3-bench.py new file mode 100644 index 00000000000..faad8ace674 --- /dev/null +++ b/.github/scripts/oswal-s3-bench.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +# /// script +# requires-python = ">=3.10" +# dependencies = ["boto3>=1.35,<2"] +# /// + +"""Helpers of the object store WAL S3 benchmark; only `probe` and `cleanup` need boto3 (`uv run`).""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +import time +import urllib.request +from pathlib import Path + +SMOKE_ROWS = 50_000 +# The o11ybench data profiles; a matching row count generates the canonical dataset. +DATA_PROFILES = {5_000: "S", 2_000_000: "D", 10_000_000: "P", 100_000_000: "M"} +# Identical [storage] for every configuration; only [wal] differs. +WAL_CONFIGS = { + "R": {"PROVIDER": "raft_engine", "SYNC_WRITE": "true", "DIR": "/wal"}, + "D": {"PROVIDER": "experimental_object_store", "ACK_MODE": "durable"}, + "E": {"PROVIDER": "experimental_object_store", "ACK_MODE": "enqueued"}, +} +ORDERS = (("R", "D", "E"), ("E", "D", "R")) +NIGHTLY = re.compile(r"v\d+\.\d+\.\d+(?:-[0-9A-Za-z.]+)?-nightly-(\d{8})") +TAG = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}") +DIGEST = re.compile(r".+@sha256:[0-9a-f]{64}") + + +def append_github(variable: str, name: str, value: str) -> None: + if path := os.environ.get(variable): + with open(path, "a") as output: + output.write(f"{name}={value}\n") + + +def positive(name: str, value: str, upper: int) -> int: + if not re.fullmatch(r"[1-9][0-9]*", value) or int(value) > upper: + raise ValueError(f"{name} must be an integer in 1..{upper}, got {value!r}") + return int(value) + + +def validate_budgets(timeout_minutes: str, ttl_hours: str) -> None: + timeout, ttl = positive("timeout", timeout_minutes, 360), positive("ttl", ttl_hours, 168) + # Provisioning, teardown and queueing need two hours beyond the benchmark job. + if ttl * 60 < timeout + 120: + raise ValueError(f"janitor TTL {ttl}h must cover the {timeout}-minute timeout plus 2 hours") + + +def check_regions(bucket: str, runner: str) -> None: + if not bucket or bucket != runner: + raise ValueError(f"bucket region {bucket!r} differs from runner region {runner!r}; " + "a cross-region bucket measures the WAN, not S3") + + +def build_manifest(plan: dict, greptimedb_inspect: list, runtime_image: str, + instance_type: str, region: str, prefix: str) -> dict: + image = next((d for d in greptimedb_inspect[0].get("RepoDigests") or [] if DIGEST.fullmatch(d)), None) + if image is None: + raise ValueError("the GreptimeDB image has no sha256 repository digest") + return {"smoke": plan["smoke"], "rows": plan["rows"], "rounds": plan["rounds"], + "concurrency": plan["concurrency"], "batch_rows": plan["batch_rows"], + "greptimedb_image": image, "runtime_image": runtime_image, + "instance_type": instance_type, "region": region, "prefix": prefix} + + +def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: bool) -> dict: + levels = [positive("concurrency", level.strip(), 256) for level in concurrency.split(",")] + if len(set(levels)) != len(levels): + raise ValueError(f"concurrency repeats a level: {concurrency!r}") + row_count = SMOKE_ROWS if smoke else positive("rows", rows, 100_000_000) + round_count = 1 if smoke else positive("rounds", rounds, 10) + profile = DATA_PROFILES.get(row_count) + cells = [] + for round_number in range(1, round_count + 1): + # Alternate the order between rounds so no configuration always runs first. + for level in levels: + for config in ORDERS[(round_number - 1) % 2]: + cells.append({ + "id": f"{round_number}-{level}-{config}", "round": round_number, + "concurrency": level, "config": config, "wal": WAL_CONFIGS[config], + }) + return { + "smoke": smoke, "rows": row_count, "rounds": round_count, + "batch_rows": positive("batch_rows", batch_rows, 100_000), "concurrency": levels, + "data_profile": profile or "custom", + "generator_args": ["--data-profile", profile] if profile else ["--line-count", str(row_count)], + "cells": cells, + } + + +def newest_nightly(tags: list[str]) -> str: + dated = [(match.group(1), tag) for tag in tags if (match := NIGHTLY.fullmatch(tag))] + if not dated: + raise ValueError("no v*-nightly-* tag found") + return max(dated)[1] + + +def resolve_tag(requested: str) -> str: + if requested != "nightly": + if not TAG.fullmatch(requested): + raise ValueError(f"invalid image tag: {requested!r}") + return requested + url = ("https://hub.docker.com/v2/repositories/greptime/greptimedb/tags" + "?page_size=100&name=nightly&ordering=last_updated") + with urllib.request.urlopen(url, timeout=30) as response: + return newest_nightly([tag["name"] for tag in json.load(response)["results"]]) + + +def s3_client(args): + import boto3 + from botocore.config import Config + + return boto3.client("s3", region_name=args.region, endpoint_url=args.endpoint or None, + config=Config(retries={"mode": "standard", "max_attempts": 5})) + + +def probe(args) -> None: + client, key = s3_client(args), f"{args.prefix}/probe/object" + deadline = time.monotonic() + 60 + while True: + try: + client.put_object(Bucket=args.bucket, Key=key, Body=b"probe") + break + except Exception: + # A fresh MinIO creates its bucket after the container starts. + if time.monotonic() > deadline: + raise + time.sleep(3) + listed = client.list_objects_v2(Bucket=args.bucket, Prefix=f"{args.prefix}/probe/") + if [item["Key"] for item in listed.get("Contents", [])] != [key]: + raise RuntimeError(f"listing {args.prefix}/probe/ did not return exactly {key}") + client.delete_object(Bucket=args.bucket, Key=key) + print(f"put, list and delete succeeded under s3://{args.bucket}/{args.prefix}/probe/") + + +def cleanup(args) -> None: + # One DeleteObject per key: DeleteObjects needs a checksum header the MinIO fixture rejects. + client, deleted = s3_client(args), 0 + for page in client.get_paginator("list_objects_v2").paginate(Bucket=args.bucket, Prefix=f"{args.prefix}/"): + for item in page.get("Contents", []): + client.delete_object(Bucket=args.bucket, Key=item["Key"]) + deleted += 1 + print(f"deleted {deleted} objects under s3://{args.bucket}/{args.prefix}/") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + plan = commands.add_parser("plan") + for name in ("concurrency", "rows", "batch-rows", "rounds"): + plan.add_argument(f"--{name}", required=True) + plan.add_argument("--smoke", choices=("true", "false"), required=True) + for name in ("timeout-minutes", "ttl-hours", "output"): + plan.add_argument(f"--{name}") + commands.add_parser("tag").add_argument("requested") + regions = commands.add_parser("regions") + for name in ("bucket", "runner"): + regions.add_argument(f"--{name}", required=True) + manifest = commands.add_parser("manifest") + for name in ("results", "runtime-image", "instance-type", "region", "prefix"): + manifest.add_argument(f"--{name}", required=True) + for name in ("probe", "cleanup"): + bucket = commands.add_parser(name) + bucket.add_argument("--endpoint", default="") + for option in ("bucket", "region", "prefix"): + bucket.add_argument(f"--{option}", required=True) + args = parser.parse_args() + if args.command == "plan": + if args.timeout_minutes or args.ttl_hours: + validate_budgets(args.timeout_minutes or "", args.ttl_hours or "") + result = build_plan(args.concurrency, args.rows, args.batch_rows, args.rounds, args.smoke == "true") + if args.output: + with open(args.output, "w") as output: + json.dump(result, output, indent=2) + print(f"{len(result['cells'])} cells, {result['rows']} rows, profile {result['data_profile']}") + elif args.command == "tag": + tag = resolve_tag(args.requested) + append_github("GITHUB_OUTPUT", "greptimedb_tag", tag) + print(f"GreptimeDB image tag: {tag}") + elif args.command == "regions": + check_regions(args.bucket, args.runner) + print(f"bucket and runner both in {args.runner}") + elif args.command == "manifest": + results = Path(args.results) + manifest = build_manifest(json.loads((results / "plan.json").read_text()), + json.loads((results / "images" / "greptimedb.json").read_text()), + args.runtime_image, args.instance_type, args.region, args.prefix) + (results / "run.json").write_text(json.dumps(manifest, indent=2) + "\n") + append_github("GITHUB_ENV", "GREPTIMEDB_IMAGE", manifest["greptimedb_image"]) + print(f"GreptimeDB image: {manifest['greptimedb_image']}") + else: + if not re.fullmatch(r"oswal-bench/[0-9]+-[0-9]+", args.prefix): + raise ValueError(f"refusing a prefix outside oswal-bench/: {args.prefix!r}") + (probe if args.command == "probe" else cleanup)(args) + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except ValueError as error: + print(f"::error::{error}", file=sys.stderr) + sys.exit(1) diff --git a/.github/workflows/oswal-s3-bench.yml b/.github/workflows/oswal-s3-bench.yml new file mode 100644 index 00000000000..c31dfe8f1c3 --- /dev/null +++ b/.github/workflows/oswal-s3-bench.yml @@ -0,0 +1,229 @@ +name: Object Store WAL S3 Benchmark +on: + workflow_dispatch: + inputs: + smoke: + description: 'Smoke run: ubuntu-latest and the MinIO fixture instead of EC2 and S3, 50000 rows, one round; its numbers mean nothing' + type: boolean + default: false + greptimedb_tag: + description: 'Docker Hub greptime/greptimedb tag; nightly = the newest v*-nightly-* tag (latest is a release without the object store WAL)' + type: string + default: nightly + instance_type: + description: EC2 instance type in the bucket region; the d of c6id gives a local NVMe instance store for the Raft Engine WAL + type: string + default: c6id.2xlarge + concurrency: + description: Comma-separated loader concurrency levels + type: string + default: '1,8' + rows: + description: Dataset rows (2000000 = o11ybench profile D) + type: number + default: 2000000 + batch_rows: + description: Rows per write request + type: number + default: 1000 + rounds: + description: Rounds; the configuration order alternates between rounds (R, D, E then E, D, R) + type: number + default: 2 + runtime_image: + description: 'o11ybench runtime image: auto = AWS ECR Public; explicit tag/digest overrides' + type: string + default: auto + benchmark_timeout_minutes: + description: Benchmark job timeout in minutes (1..360) + type: number + default: 240 + janitor_ttl_hours: + description: Instance lifetime before janitor eligibility (1..168 hours; at least timeout + 2h) + type: number + default: 8 +permissions: + contents: read +env: + SMOKE: ${{ inputs.smoke }} + GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} + REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} + CONCURRENCY: ${{ inputs.concurrency }} + ROWS: ${{ inputs.rows }} + BATCH_ROWS: ${{ inputs.batch_rows }} + ROUNDS: ${{ inputs.rounds }} + RUNTIME_IMAGE: ${{ inputs.runtime_image != 'auto' && inputs.runtime_image || 'public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime:77a84e1c8136@sha256:c482993ec177b026235d6cd0c61d4fb6d6f1088bc78bdc0e5964a5398de6dba0' }} + BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} + JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} +jobs: + validate: + runs-on: ubuntu-latest + timeout-minutes: 15 + outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} + greptimedb_tag: ${{ steps.plan.outputs.greptimedb_tag }} + runtime_image: ${{ steps.runtime.outputs.runtime_image }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Test benchmark scripts + run: | + python3 greptimedb/tests/perf/test_oswal_s3_bench_scripts.py + python3 greptimedb/tests/perf/test_ci_runner_scripts.py + - name: Validate plan, budgets and image tag + id: plan + run: | + python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" \ + --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" \ + --timeout-minutes "$BENCHMARK_TIMEOUT_MINUTES" --ttl-hours "$JANITOR_TTL_HOURS" + python3 greptimedb/.github/scripts/oswal-s3-bench.py tag "$GREPTIMEDB_TAG" + - name: Resolve CI runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider AWS --instance-type "$REQUESTED_INSTANCE_TYPE" + - name: Check bucket and runner regions + env: + BUCKET_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + RUNNER_REGION: ${{ vars.EC2_RUNNER_REGION }} + run: python3 greptimedb/.github/scripts/oswal-s3-bench.py regions --bucket "$BUCKET_REGION" --runner "$RUNNER_REGION" + - name: Check the runtime image + id: runtime + run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh check-runtime "$RUNTIME_IMAGE" + provision: + needs: validate + if: ${{ !inputs.smoke }} + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.aws.outputs.label }} + instance_id: ${{ steps.aws.outputs.instance_id }} + runner_name: ${{ steps.aws.outputs.runner_name }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + - name: Create ephemeral EC2 runner + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: '60' + AWS_EC2_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py + benchmark: + needs: + - validate + - provision + if: ${{ !cancelled() && needs.validate.result == 'success' && (inputs.smoke || needs.provision.result == 'success') }} + runs-on: ${{ inputs.smoke && 'ubuntu-latest' || needs.provision.outputs.label }} + timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} + env: + INSTANCE_TYPE: ${{ inputs.smoke && 'ubuntu-latest' || needs.validate.outputs.instance_type }} + GREPTIMEDB_TAG: ${{ needs.validate.outputs.greptimedb_tag }} + RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} + RUN_PREFIX: oswal-bench/${{ github.run_id }}-${{ github.run_attempt }} + OWNER: o11ybench-oswal-${{ github.run_id }}-${{ github.run_attempt }} + RESULTS: ${{ github.workspace }}/oswal-results + CORPUS_DIR: ${{ github.workspace }}/oswal-corpus + PLAN: ${{ github.workspace }}/oswal-results/plan.json + WAL_ROOT: /mnt/oswal-wal + S3_BUCKET: ${{ inputs.smoke && 'greptime' || vars.AWS_CI_TEST_BUCKET }} + S3_REGION: ${{ inputs.smoke && 'us-west-2' || vars.AWS_CI_TEST_BUCKET_REGION }} + S3_ENDPOINT: ${{ inputs.smoke && 'http://127.0.0.1:9000' || '' }} + S3_CONTAINER_ENDPOINT: ${{ inputs.smoke && 'http://host.docker.internal:9000' || '' }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark + - uses: astral-sh/setup-uv@v6 + - name: Start the MinIO fixture + if: inputs.smoke + run: | + bash greptimedb/.github/scripts/pull-benchmark-image.sh greptime/minio:2024 + docker compose -f greptimedb/tests-integration/fixtures/docker-compose.yml up -d --wait minio + - name: Pull images and write the plan + run: | + bash greptimedb/.github/scripts/oswal-s3-bench-run.sh pull + python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" --output "$PLAN" \ + --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" + python3 greptimedb/.github/scripts/oswal-s3-bench.py manifest --results "$RESULTS" \ + --runtime-image "$RUNTIME_IMAGE" --instance-type "$INSTANCE_TYPE" --region "$S3_REGION" --prefix "$RUN_PREFIX" + - name: Probe the bucket + env: + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/oswal-s3-bench.py probe --bucket "$S3_BUCKET" --region "$S3_REGION" --endpoint "$S3_ENDPOINT" --prefix "$RUN_PREFIX" + - name: Put the Raft Engine WAL on the instance store + run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh wal-disk "$WAL_ROOT" + - name: Generate the dataset once + run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh generate + - name: Load every cell + env: + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh cells + - name: Summary + if: always() + run: python3 greptimedb/.github/scripts/oswal-s3-bench-summary.py --root "$RESULTS" >> "$GITHUB_STEP_SUMMARY" + - name: Delete the run's S3 prefix + if: always() + env: + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/oswal-s3-bench.py cleanup --bucket "$S3_BUCKET" --region "$S3_REGION" --endpoint "$S3_ENDPOINT" --prefix "$RUN_PREFIX" + - name: Check the results hold no credentials + if: always() + env: + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh scrub + - name: Upload results (no corpus) + if: always() + uses: actions/upload-artifact@v4 + with: + name: oswal-s3-bench-${{ inputs.smoke && 'smoke' || 'ec2' }}-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.RESULTS }} + if-no-files-found: warn + retention-days: 14 + teardown: + needs: + - provision + - benchmark + if: ${{ always() && !inputs.smoke && (needs.provision.outputs.instance_id != '' || needs.provision.result != 'skipped') }} + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + - name: Delete EC2 and unregister runner + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py diff --git a/tests/perf/test_oswal_s3_bench_scripts.py b/tests/perf/test_oswal_s3_bench_scripts.py new file mode 100644 index 00000000000..dc81240f1ae --- /dev/null +++ b/tests/perf/test_oswal_s3_bench_scripts.py @@ -0,0 +1,318 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +"""Offline tests of the object store WAL S3 benchmark plan and summary.""" + +import importlib.util +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +SCRIPTS = Path(__file__).parents[2] / ".github/scripts" + + +def load(name, filename): + spec = importlib.util.spec_from_file_location(name, SCRIPTS / filename) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +bench = load("oswal_s3_bench", "oswal-s3-bench.py") +summary = load("oswal_s3_bench_summary", "oswal-s3-bench-summary.py") + +ACK_BUCKETS = """# TYPE greptime_logstore_object_store_wal_append_ack_seconds histogram +greptime_logstore_object_store_wal_append_ack_seconds_bucket{le="0.064"} 0 +greptime_logstore_object_store_wal_append_ack_seconds_bucket{le="0.128"} 50 +greptime_logstore_object_store_wal_append_ack_seconds_bucket{le="0.256"} 100 +greptime_logstore_object_store_wal_append_ack_seconds_bucket{le="+Inf"} 100 +greptime_logstore_object_store_wal_append_ack_seconds_sum 12 +greptime_logstore_object_store_wal_append_ack_seconds_count 100 +greptime_logstore_object_store_wal_created_objects_total 120 +greptime_logstore_object_store_wal_stalled_appends_total 3 +greptime_logstore_object_store_wal_poisoned_total 0 +greptime_logstore_kafka_client_produce_elapsed_count 9 +""" + + +class PlanTest(unittest.TestCase): + def test_default_plan_alternates_the_order_between_rounds(self): + plan = bench.build_plan("1,8", "2000000", "1000", "2", smoke=False) + self.assertEqual([c["id"] for c in plan["cells"]], [ + "1-1-R", "1-1-D", "1-1-E", "1-8-R", "1-8-D", "1-8-E", + "2-1-E", "2-1-D", "2-1-R", "2-8-E", "2-8-D", "2-8-R"]) + self.assertEqual(plan["generator_args"], ["--data-profile", "D"]) + self.assertEqual((plan["rows"], plan["batch_rows"], plan["data_profile"]), (2_000_000, 1000, "D")) + self.assertEqual({c["config"]: c["wal"] for c in plan["cells"]}, { + "R": {"PROVIDER": "raft_engine", "SYNC_WRITE": "true", "DIR": "/wal"}, + "D": {"PROVIDER": "experimental_object_store", "ACK_MODE": "durable"}, + "E": {"PROVIDER": "experimental_object_store", "ACK_MODE": "enqueued"}}) + + def test_smoke_overrides_rows_and_rounds(self): + plan = bench.build_plan("8", "2000000", "1000", "2", smoke=True) + self.assertEqual((plan["rows"], plan["rounds"], plan["data_profile"]), (50_000, 1, "custom")) + self.assertEqual(plan["generator_args"], ["--line-count", "50000"]) + self.assertEqual([c["id"] for c in plan["cells"]], ["1-8-R", "1-8-D", "1-8-E"]) + + def test_invalid_inputs_are_rejected(self): + for args in (("0", "10", "10", "1"), ("1,1", "10", "10", "1"), ("1", "abc", "10", "1"), + ("1", "10", "10", "11"), ("1,", "10", "10", "1")): + with self.subTest(args=args), self.assertRaises(ValueError): + bench.build_plan(*args, smoke=False) + + def test_nightly_tag_is_the_newest_by_date(self): + tags = ["latest", "v1.3.0-alpha.1-nightly-20260907", "v1.4.0-nightly-20261012", + "v1.4.0-nightly-20261005", "nightly-arm64"] + self.assertEqual(bench.newest_nightly(tags), "v1.4.0-nightly-20261012") + with self.assertRaises(ValueError): + bench.newest_nightly(["latest"]) + self.assertEqual(bench.resolve_tag("v1.4.0-nightly-20261012"), "v1.4.0-nightly-20261012") + with self.assertRaises(ValueError): + bench.resolve_tag("bad tag") + + def test_budgets_and_regions(self): + bench.validate_budgets("240", "8") + for timeout, ttl in (("360", "7"), ("0", "8"), ("361", "168"), ("240", "169")): + with self.subTest(timeout=timeout, ttl=ttl), self.assertRaises(ValueError): + bench.validate_budgets(timeout, ttl) + bench.check_regions("ap-southeast-1", "ap-southeast-1") + for bucket, runner in (("us-east-1", "ap-southeast-1"), ("", "")): + with self.subTest(bucket=bucket), self.assertRaisesRegex(ValueError, "cross-region"): + bench.check_regions(bucket, runner) + + def test_manifest_takes_the_sha256_digest(self): + plan = bench.build_plan("1,8", "2000000", "1000", "2", smoke=False) + inspect = [{"RepoDigests": ["greptime/greptimedb@sha256:" + "a" * 64]}] + manifest = bench.build_manifest(plan, inspect, "runtime@sha256:" + "b" * 64, "c6id.2xlarge", + "ap-southeast-1", "oswal-bench/1-1") + self.assertEqual(manifest["greptimedb_image"], "greptime/greptimedb@sha256:" + "a" * 64) + self.assertEqual((manifest["rows"], manifest["rounds"], manifest["concurrency"]), (2_000_000, 2, [1, 8])) + with self.assertRaises(ValueError): + bench.build_manifest(plan, [{"RepoDigests": []}], "", "", "", "") + + +class SummaryTest(unittest.TestCase): + def test_quantile_interpolates_like_prometheus(self): + buckets = summary.parse_metrics(ACK_BUCKETS)["buckets"]["append_ack_seconds"] + self.assertAlmostEqual(summary.quantile(0.5, buckets), 0.128) + self.assertAlmostEqual(summary.quantile(0.99, buckets), 0.128 + 0.128 * 49 / 50) + self.assertAlmostEqual(summary.quantile(0.25, buckets), 0.064 + 0.064 * 25 / 50) + self.assertIsNone(summary.quantile(0.5, [])) + self.assertEqual(summary.quantile(0.99, [(0.1, 1.0), (float("inf"), 4.0)]), 0.1) + self.assertNotIn("kafka_client_produce_elapsed_count", summary.parse_metrics(ACK_BUCKETS)["samples"]) + + def write_cell(self, root, cell_id, rows_per_second, exit_code=0, metrics=""): + round_number, concurrency, config = cell_id.split("-") + path = root / "cells" / cell_id + (path / "greptimedb-1").mkdir(parents=True) + (path / "cell.json").write_text(json.dumps({ + "id": cell_id, "round": int(round_number), "concurrency": int(concurrency), "config": config, + "wal": bench.WAL_CONFIGS[config], "metrics_scraped": True, + "exit_code": exit_code, "wal_device": "/dev/nvme1n1", "version": "1.4.0", "load_wall_seconds": 61.0})) + if exit_code == 0: + (path / "greptimedb-1" / "load-result.json").write_text(json.dumps({ + "count_after": 6000, "load_seconds": 60.0, "rows_per_second": rows_per_second})) + (path / "metrics.prom").write_text(metrics) + + def test_summary_reports_cells_and_geometric_mean_ratios(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + plan = bench.build_plan("8", "6000", "1000", "2", smoke=False) + (root / "plan.json").write_text(json.dumps(plan)) + (root / "run.json").write_text(json.dumps({"rows": 6000, "batch_rows": 1000, "instance_type": "c6id.2xlarge"})) + speeds = {"1-8-R": 100.0, "1-8-D": 80.0, "1-8-E": 90.0, "2-8-E": 99.0, "2-8-D": 50.0, "2-8-R": 110.0} + for cell_id, speed in speeds.items(): + self.write_cell(root, cell_id, speed, metrics=ACK_BUCKETS if cell_id.endswith("D") else "") + output = subprocess.run([sys.executable, str(SCRIPTS / "oswal-s3-bench-summary.py"), "--root", str(root)], + check=True, capture_output=True, text=True).stdout + result = json.loads((root / "summary.json").read_text()) + self.assertEqual([c["id"] for c in result["cells"]], [c["id"] for c in plan["cells"]]) + d_ratio = next(r for r in result["ratios"] if r["ratio"] == "D/R") + self.assertAlmostEqual(d_ratio["geomean"], (0.8 * (50 / 110)) ** 0.5) + self.assertAlmostEqual(next(r for r in result["ratios"] if r["ratio"] == "E/R")["geomean"], 0.9) + d_cell = next(c for c in result["cells"] if c["id"] == "1-8-D") + self.assertEqual(d_cell["creates_per_minute"], 120.0) + self.assertEqual(d_cell["stalled_appends_total"], 3) + self.assertIn("| 1-8-D | 80 | 60.0 | 61.0 | 6,000 | 128.0 | 253.4 | 120 | 120.0 | 3 | 0 | 0 | 0 |", output) + r_cell = next(c for c in result["cells"] if c["id"] == "1-8-R") + self.assertIsNone(r_cell["stalled_appends_total"]) + + def test_untouched_counters_are_zero_only_for_an_active_object_store_wal(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + created = "greptime_logstore_object_store_wal_created_objects_total 40\n" + self.write_cell(root, "1-8-E", 90.0, metrics=created) + self.write_cell(root, "1-8-D", 80.0, metrics="") + cell_e, cell_d = (summary.load_cell(root / "cells" / i) for i in ("1-8-E", "1-8-D")) + self.assertEqual([cell_e[c] for c in summary.COUNTERS], [40, 0, 0, 0, 0]) + self.assertEqual([cell_d[c] for c in summary.COUNTERS], [None] * 5) + recorded = json.loads((root / "cells" / "1-8-E" / "cell.json").read_text()) + (root / "cells" / "1-8-E" / "cell.json").write_text(json.dumps(recorded | {"metrics_scraped": False})) + unscraped = summary.load_cell(root / "cells" / "1-8-E") + self.assertIsNone(unscraped["stalled_appends_total"]) + + def test_a_planned_round_that_never_ran_leaves_the_ratio_undefined(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + plan = bench.build_plan("8", "6000", "1000", "2", smoke=False) + (root / "plan.json").write_text(json.dumps(plan)) + for cell_id, speed in {"1-8-R": 100.0, "1-8-D": 80.0, "1-8-E": 90.0}.items(): + self.write_cell(root, cell_id, speed) + output = subprocess.run([sys.executable, str(SCRIPTS / "oswal-s3-bench-summary.py"), "--root", str(root)], + check=True, capture_output=True, text=True).stdout + result = json.loads((root / "summary.json").read_text()) + self.assertEqual(result["missing_cells"], ["2-8-E", "2-8-D", "2-8-R"]) + self.assertEqual(result["planned_cells"], 6) + for ratio in result["ratios"]: + self.assertIsNone(ratio["geomean"]) + self.assertIsNone(ratio["per_round"][1]) + self.assertIn("2-8-D (not run)", output) + + def test_a_failed_cell_leaves_its_ratio_undefined(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_cell(root, "1-1-R", 100.0) + self.write_cell(root, "1-1-D", 0.0, exit_code=1) + rows = [summary.load_cell(path) for path in sorted((root / "cells").iterdir())] + ratio = next(r for r in summary.ratios(rows) if r["ratio"] == "D/R") + self.assertEqual((ratio["per_round"], ratio["geomean"]), ([None], None)) + self.assertIn("1-1-D (failed, exit 1)", summary.render({}, rows, summary.ratios(rows))) + + +FAKE_DOCKER = r"""#!/usr/bin/env bash +printf '%s\n' "docker $*" >> "$FAKE_LOG" +case "$1" in + run) + for arg in "$@"; do [[ "$arg" != -d ]] || { echo "keys=${GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID:+set}${GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY:+set}" >> "$FAKE_LOG"; echo fakecid; exit 0; }; done + for arg in "$@"; do + case "$arg" in type=bind,src=*,dst=/results) results=${arg#type=bind,src=}; results=${results%,dst=/results} ;; esac + done + [[ "$FAKE_LOAD_RC" == 0 ]] || exit "$FAKE_LOAD_RC" + mkdir -p "$results/greptimedb-1" + echo "{\"count_before\": 0, \"count_after\": $FAKE_ROWS, \"load_seconds\": 2, \"rows_per_second\": 3000}" \ + > "$results/greptimedb-1/load-result.json" ;; + inspect) echo true ;; + exec) echo 'version: 1.4.0-test' ;; + logs) echo 'server log' ;; +esac +""" +FAKE_CURL = r"""#!/usr/bin/env bash +printf '%s\n' "curl $*" >> "$FAKE_LOG" +for arg in "$@"; do + [[ "$arg" != */metrics ]] || { + [[ "$FAKE_METRICS_RC" == 0 ]] || exit "$FAKE_METRICS_RC" + printf '%s\n' 'greptime_logstore_object_store_wal_poisoned_total 1' 'other_total 2'; exit 0; + } +done +""" + + +class CellLifecycleTest(unittest.TestCase): # the run script against fake docker and curl + + def run_cells(self, root, load_rc, metrics_rc=0): + bin_dir = root / "bin" + bin_dir.mkdir() + for name, body in (("docker", FAKE_DOCKER), ("curl", FAKE_CURL)): + (bin_dir / name).write_text(body) + (bin_dir / name).chmod(0o755) + plan = bench.build_plan("8", "6000", "1000", "1", smoke=False) + plan["cells"] = [cell for cell in plan["cells"] if cell["config"] != "R"] + (root / "results").mkdir() + (root / "results" / "plan.json").write_text(json.dumps(plan)) + env = os.environ | { + "PATH": f"{bin_dir}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), + "FAKE_LOAD_RC": str(load_rc), "FAKE_METRICS_RC": str(metrics_rc), "FAKE_ROWS": "6000", "PLAN": str(root / "results" / "plan.json"), + "RESULTS": str(root / "results"), "CORPUS_DIR": str(root / "corpus"), "OWNER": "o11ybench-oswal-1-1", + "RUNTIME_IMAGE": "runtime@sha256:" + "b" * 64, "GREPTIMEDB_IMAGE": "greptimedb@sha256:" + "a" * 64, + "RUN_PREFIX": "oswal-bench/1-1", "WAL_ROOT": str(root / "wal"), "S3_BUCKET": "bucket", + "S3_REGION": "ap-southeast-1", "AWS_ACCESS_KEY_ID": "AKIAFAKEKEY", "AWS_SECRET_ACCESS_KEY": "fake-secret-value", + } + script = str(SCRIPTS / "oswal-s3-bench-run.sh") + cells = subprocess.run(["bash", script, "cells"], env=env, capture_output=True, text=True, timeout=60) + scrub = subprocess.run(["bash", script, "scrub"], env=env, capture_output=True, text=True, timeout=60) + return cells, scrub, (root / "fake.log").read_text(), env + + def test_a_failed_load_still_scrapes_metrics_and_removes_the_server(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, scrub, log, _ = self.run_cells(root, load_rc=3) + self.assertNotEqual(cells.returncode, 0) + self.assertEqual(scrub.returncode, 0, scrub.stderr) + for cell_id in ("1-8-D", "1-8-E"): + cell = root / "results" / "cells" / cell_id + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 3) + self.assertEqual((cell / "metrics.prom").read_text(), + "greptime_logstore_object_store_wal_poisoned_total 1\n") + self.assertFalse((cell / "metrics.all").exists()) + self.assertEqual((cell / "server.log").read_text(), "server log\n") + self.assertEqual(log.count("docker stop -t 120 fakecid"), 2) + self.assertEqual(log.count("docker rm -fv fakecid"), 2) + + def test_an_ec2_run_without_an_instance_store_fails(self): + with tempfile.TemporaryDirectory() as directory: + (Path(directory) / "lsblk").write_text("#!/usr/bin/env bash\necho '/dev/nvme0n1 Amazon Elastic Block Store'\n") + (Path(directory) / "lsblk").chmod(0o755) + env = os.environ | {"PATH": f"{directory}:{os.environ['PATH']}", "SMOKE": "false", "RESULTS": directory} + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "wal-disk", f"{directory}/wal"], + env=env, capture_output=True, text=True, timeout=60) + self.assertNotEqual(done.returncode, 0) + self.assertIn("No local NVMe instance store", done.stderr) + + def test_a_failed_scrape_fails_a_successful_load(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, log, _ = self.run_cells(root, load_rc=0, metrics_rc=22) + self.assertNotEqual(cells.returncode, 0) + for cell_id in ("1-8-D", "1-8-E"): + cell = root / "results" / "cells" / cell_id + recorded = json.loads((cell / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["metrics_scraped"]), (4, False)) + self.assertFalse((cell / "metrics.prom").exists()) + self.assertFalse((cell / "metrics.all").exists()) + self.assertEqual(log.count("docker rm -fv fakecid"), 2) + + def test_a_cell_passes_the_keys_by_name_and_records_a_redacted_config(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, scrub, log, env = self.run_cells(root, load_rc=0) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + self.assertEqual(scrub.returncode, 0, scrub.stderr) + self.assertNotIn("fake-secret-value", log) + self.assertNotIn("AKIAFAKEKEY", log) + self.assertEqual(log.count("keys=setset"), 2) + self.assertIn("-e GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY ", log) + cell = root / "results" / "cells" / "1-8-D" + config = (cell / "config.env").read_text().splitlines() + self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__ROOT=oswal-bench/1-1/1-8-D", config) + self.assertIn("GREPTIMEDB_STANDALONE__WAL__ACK_MODE=durable", config) + self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY=", config) + recorded = json.loads((cell / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["version"], recorded["metrics_scraped"]), + (0, "1.4.0-test", True)) + self.assertEqual(recorded["wal_device"], "s3://bucket/oswal-bench/1-1/1-8-D/wal") + # A leaked value makes scrub refuse the results. + (cell / "leak.txt").write_text("fake-secret-value") + leaked = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "scrub"], env=env, + capture_output=True, text=True, timeout=60) + self.assertNotEqual(leaked.returncode, 0) + self.assertFalse((root / "results").exists()) + + +if __name__ == "__main__": + unittest.main() From 684c8ff2917c10dcf8103664f5eb2b01ef6720e9 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 02:29:54 +0000 Subject: [PATCH 20/46] feat(client): support region request reuse and route refresh hints (#9540) * feat(client): support shared region requests without deep cloning Signed-off-by: WenyXu * feat: propagate route refresh retry hints Signed-off-by: WenyXu * style: preserve local error variant imports in retry hints Signed-off-by: WenyXu * fix(mito2): distinguish staging rejection from write backpressure Signed-off-by: WenyXu * fix(meta): preserve legacy retry hints in instruction replies Signed-off-by: WenyXu * test(client): inline retry hint compatibility cases Signed-off-by: WenyXu * test(datanode): align batch-open retry hint assertion Signed-off-by: WenyXu * style: use imported types in region request tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 1 + Cargo.toml | 1 + src/client/Cargo.toml | 1 + src/client/src/client.rs | 22 +- src/client/src/client_manager.rs | 25 +- src/client/src/error.rs | 15 +- src/client/src/region.rs | 616 ++++++++++++++++++++++----- src/client/src/region/codec.rs | 63 +++ src/common/error/src/ext.rs | 18 +- src/common/error/tests/ext.rs | 9 + src/common/meta/src/error.rs | 7 +- src/common/meta/src/instruction.rs | 34 +- src/datanode/src/error.rs | 30 +- src/datanode/src/region_server.rs | 2 +- src/metric-engine/src/error.rs | 20 + src/mito2/src/engine/flush_test.rs | 2 + src/mito2/src/engine/staging_test.rs | 73 +++- src/mito2/src/error.rs | 73 +++- src/mito2/src/worker/handle_write.rs | 14 +- 19 files changed, 850 insertions(+), 176 deletions(-) create mode 100644 src/client/src/region/codec.rs diff --git a/Cargo.lock b/Cargo.lock index 61a9c246ed1..f75bde16328 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2122,6 +2122,7 @@ dependencies = [ "tokio", "tokio-stream", "tonic 0.14.2", + "tonic-prost", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index f47907aa797..702498ee5e5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -272,6 +272,7 @@ url = "2.3" # Pin below 0.14.5 to avoid the `connection_timeout_future` "async fn resumed after # completion" panic (greptimedb#8265 / grpc-rust#2522). Remove once tonic ships a fix. tonic = { version = "0.14, <0.14.5", features = ["tls-aws-lc", "gzip", "zstd"] } +tonic-prost = "0.14" tower = "0.5" tracing = "0.1" tracing-appender = "0.2" diff --git a/src/client/Cargo.toml b/src/client/Cargo.toml index 36c3f60a9a4..e5bd479cc7d 100644 --- a/src/client/Cargo.toml +++ b/src/client/Cargo.toml @@ -45,3 +45,4 @@ substrait.workspace = true tokio.workspace = true tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true +tonic-prost.workspace = true diff --git a/src/client/src/client.rs b/src/client/src/client.rs index dfc0e44fad5..1746854cd26 100644 --- a/src/client/src/client.rs +++ b/src/client/src/client.rs @@ -20,7 +20,6 @@ use api::v1::HealthCheckRequest; use api::v1::flow::flow_client::FlowClient as PbFlowClient; use api::v1::health_check_client::HealthCheckClient; use api::v1::prometheus_gateway_client::PrometheusGatewayClient; -use api::v1::region::region_client::RegionClient as PbRegionClient; use arrow_flight::flight_service_client::FlightServiceClient; use common_grpc::channel_manager::{ ChannelConfig, ChannelManager, ClientTlsOption, load_client_tls_config, @@ -526,23 +525,12 @@ impl Client { Ok(FlightClient { addr, client }) } - pub(crate) fn raw_region_client(&self) -> Result<(String, PbRegionClient)> { + /// Creates the region gRPC transport with the configured message limits. + pub(crate) fn raw_region_client(&self) -> Result<(String, tonic::client::Grpc)> { let (addr, channel) = self.find_channel()?; - let client = PbRegionClient::new(channel) - .max_decoding_message_size( - self.inner - .control_channel_manager - .config() - .max_recv_message_size - .as_bytes() as usize, - ) - .max_encoding_message_size( - self.inner - .control_channel_manager - .config() - .max_send_message_size - .as_bytes() as usize, - ); + let client = tonic::client::Grpc::new(channel) + .max_decoding_message_size(self.max_grpc_recv_message_size()) + .max_encoding_message_size(self.max_grpc_send_message_size()); Ok((addr, client)) } diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs index 4d723361e80..ef26c107cd5 100644 --- a/src/client/src/client_manager.rs +++ b/src/client/src/client_manager.rs @@ -52,18 +52,7 @@ impl Debug for NodeClients { #[async_trait::async_trait] impl DatanodeManager for NodeClients { async fn datanode(&self, datanode: &Peer) -> DatanodeRef { - let client = self.get_client(datanode).await; - - let ChannelConfig { - send_compression, - accept_compression, - .. - } = self.control_channel_manager.config(); - Arc::new(RegionRequester::new( - client, - *send_compression, - *accept_compression, - )) + Arc::new(self.make_region_requester(datanode).await) } } @@ -89,6 +78,18 @@ impl NodeClients { } } + /// Creates a [`RegionRequester`] with the cached client and control channel compression settings. + pub async fn make_region_requester(&self, datanode: &Peer) -> RegionRequester { + let client = self.get_client(datanode).await; + + let ChannelConfig { + send_compression, + accept_compression, + .. + } = self.control_channel_manager.config(); + RegionRequester::new(client, *send_compression, *accept_compression) + } + pub async fn get_client(&self, datanode: &Peer) -> Client { self.clients .get_with_by_ref(datanode, async move { diff --git a/src/client/src/error.rs b/src/client/src/error.rs index 168ba422715..5c70e5fbe91 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -18,7 +18,7 @@ use common_error::define_from_tonic_status; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; -use snafu::{Location, Snafu}; +use snafu::{Location, Snafu, location}; use tonic::Code; use tonic::metadata::errors::InvalidMetadataValue; @@ -205,6 +205,19 @@ impl ErrorExt for Error { define_from_tonic_status!(Error, Tonic); +impl From for common_meta::error::Error { + fn from(error: Error) -> Self { + if error.should_retry() { + Self::retry_later(error) + } else { + Self::External { + source: BoxedError::new(error), + location: location!(), + } + } + } +} + impl Error { /// Returns the gRPC status code if this error is caused by a gRPC request failure. pub fn tonic_code(&self) -> Option { diff --git a/src/client/src/region.rs b/src/client/src/region.rs index fb2e4d10ca1..47c477b06b0 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod codec; + +use std::borrow::Borrow; use std::sync::Arc; use std::time::Duration; @@ -49,10 +52,12 @@ use crate::error::{ MissingFieldSnafu, Result, ServerSnafu, }; use crate::flight::{FlightMessageReader, decode_flight_data}; +use crate::region::codec::RegionCodec; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); +/// Sends region operations through a configured gRPC client. #[derive(Debug)] pub struct RegionRequester { client: Client, @@ -63,19 +68,7 @@ pub struct RegionRequester { #[async_trait] impl Datanode for RegionRequester { async fn handle(&self, request: RegionRequest) -> MetaResult { - self.handle_inner(request).await.map_err(|err| { - if err.should_retry() { - meta_error::Error::RetryLater { - source: BoxedError::new(err), - clean_poisons: false, - } - } else { - meta_error::Error::External { - source: BoxedError::new(err), - location: location!(), - } - } - }) + self.handle_inner(request).await.map_err(Into::into) } async fn handle_query(&self, request: QueryRequest) -> MetaResult { @@ -144,11 +137,49 @@ impl RegionRequester { recordbatches_from_flight_message_stream(addr, flight_message_stream).await } - async fn handle_inner(&self, request: RegionRequest) -> Result { + /// Sends an owned or shared request without deep-cloning its protobuf payload. + /// + /// Accepts both [`RegionRequest`] and [`Arc`]. Callers can retain a + /// shared request for another attempt; this method does not retry or determine + /// whether replaying the request is safe. + pub async fn handle(&self, request: R) -> Result + where + R: Borrow + Send + Sync + 'static, + { + self.handle_inner(request).await + } + + /// Sends a remote dynamic filter update to the region service. + pub async fn handle_remote_dyn_filter_update( + &self, + query_id: impl Into, + update: RemoteDynFilterUpdate, + ) -> Result { + self.handle_inner(build_remote_dyn_filter_update_request(query_id, update)) + .await + } + + /// Unregisters a remote dynamic filter in the region service. + pub async fn handle_remote_dyn_filter_unregister( + &self, + query_id: impl Into, + unregister: RemoteDynFilterUnregister, + ) -> Result { + self.handle_inner(build_remote_dyn_filter_unregister_request( + query_id, unregister, + )) + .await + } + + async fn handle_inner(&self, request: R) -> Result + where + R: Borrow + Send + Sync + 'static, + { let request_body = request + .borrow() .body .as_ref() - .with_context(|| MissingFieldSnafu { field: "body" })?; + .context(MissingFieldSnafu { field: "body" })?; let is_insert = matches!( request_body, region_request::Body::Inserts(_) | region_request::Body::BulkInsert(_) @@ -168,49 +199,43 @@ impl RegionRequester { } } - let response = client - .handle(request) - .await - .map_err(|e| { - let code = e.code(); - // Uses `Error::RegionServer` instead of `Error::Server` - error::Error::RegionServer { - addr, - code, - source: BoxedError::new(error::Error::from(e)), - location: location!(), - } - })? - .into_inner(); + let response = async { + client + .ready() + .await + .map_err(|e| tonic::Status::unknown(format!("Service was not ready: {e}")))?; + let mut request = tonic::Request::new(request); + request.extensions_mut().insert(tonic::GrpcMethod::new( + "greptime.v1.region.Region", + "Handle", + )); + client + .unary( + request, + tonic::codegen::http::uri::PathAndQuery::from_static( + "/greptime.v1.region.Region/Handle", + ), + RegionCodec::default(), + ) + .await + } + .await + .map_err(|e| { + let code = e.code(); + // Uses `Error::RegionServer` instead of `Error::Server` + error::Error::RegionServer { + addr, + code, + source: BoxedError::new(error::Error::from(e)), + location: location!(), + } + })? + .into_inner(); check_response_header(&response.header)?; Ok(RegionResponse::from_region_response(response)) } - - pub async fn handle(&self, request: RegionRequest) -> Result { - self.handle_inner(request).await - } - - pub async fn handle_remote_dyn_filter_update( - &self, - query_id: impl Into, - update: RemoteDynFilterUpdate, - ) -> Result { - self.handle_inner(build_remote_dyn_filter_update_request(query_id, update)) - .await - } - - pub async fn handle_remote_dyn_filter_unregister( - &self, - query_id: impl Into, - unregister: RemoteDynFilterUnregister, - ) -> Result { - self.handle_inner(build_remote_dyn_filter_unregister_request( - query_id, unregister, - )) - .await - } } async fn recordbatches_from_flight_message_stream( @@ -385,13 +410,22 @@ pub fn check_response_header(header: &Option) -> Result<()> { #[cfg(test)] #[allow(deprecated)] -mod test { - use api::v1::Status as PbStatus; +mod tests { + use std::collections::VecDeque; + use std::sync::Mutex; + use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - BulkInsertRequest, RegionResponse as PbRegionResponse, RemoteDynFilterUnregister, - RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, + BulkInsertRequest, InsertRequest, InsertRequests, RegionResponse as PbRegionResponse, + RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, + remote_dyn_filter_request, }; + use api::v1::value::ValueData; + use api::v1::{ + ColumnDataType, ColumnSchema as PbColumnSchema, Row, Rows, Status as PbStatus, Value, + }; + use common_error::ext::RetryHint; + use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_recordbatch::adapter::RecordBatchMetrics; use datatypes::arrow::array::Int32Array; use datatypes::prelude::{ConcreteDataType, VectorRef}; @@ -399,74 +433,440 @@ mod test { use datatypes::vectors::Int32Vector; use futures_util::stream; use tokio::net::TcpListener; + use tokio::task::JoinHandle; use tokio_stream::wrappers::TcpListenerStream; use tonic::codec::CompressionEncoding; use tonic::{Request, Response, Status}; use super::*; - use crate::Error::{self, IllegalDatabaseResponse, Server}; + use crate::Error; - #[derive(Clone)] - struct CompressionRecordingRegionService { - zstd_headers: Arc>>, + #[derive(Clone, Default)] + struct RecordingRegionService { + requests: Arc>>, + responses: Arc>>>, } #[tonic::async_trait] - impl Region for CompressionRecordingRegionService { + impl Region for RecordingRegionService { async fn handle( &self, request: Request, ) -> std::result::Result, Status> { let metadata = request.metadata(); - let sends_zstd = metadata - .get("grpc-encoding") - .and_then(|v| v.to_str().ok()) - .is_some_and(|value| value == "zstd"); + let sends_zstd = metadata.get("grpc-encoding").is_some_and(|v| v == "zstd"); let accepts_zstd = metadata .get("grpc-accept-encoding") - .and_then(|value| value.to_str().ok()) - .is_some_and(|value| value.split(',').any(|encoding| encoding == "zstd")); - self.zstd_headers + .and_then(|v| v.to_str().ok()) + .is_some_and(|v| v.split(',').any(|encoding| encoding == "zstd")); + self.requests .lock() .unwrap() - .push((sends_zstd, accepts_zstd)); + .push((request.into_inner(), sends_zstd, accepts_zstd)); + self.responses + .lock() + .unwrap() + .pop_front() + .unwrap_or_else(|| Ok(build_response(StatusCode::Success))) + .map(Response::new) + } + } - Ok(Response::new(PbRegionResponse { - header: Some(ResponseHeader { - status: Some(PbStatus { - status_code: StatusCode::Success as u32, - ..Default::default() - }), + struct RegionTestServer { + addr: String, + service: RecordingRegionService, + task: JoinHandle<()>, + } + + impl RegionTestServer { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let service = RecordingRegionService::default(); + let server_service = service.clone(); + let task = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service( + RegionServer::new(server_service) + .accept_compressed(CompressionEncoding::Zstd) + .send_compressed(CompressionEncoding::Zstd), + ) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + Self { + addr, + service, + task, + } + } + + fn create_requester( + &self, + config: ChannelConfig, + send: bool, + accept: bool, + ) -> RegionRequester { + RegionRequester::new( + Client::with_manager_and_urls( + ChannelManager::with_config(config, None), + [self.addr.clone()], + ), + send, + accept, + ) + } + } + + impl Drop for RegionTestServer { + fn drop(&mut self) { + self.task.abort(); + } + } + + fn build_response(code: StatusCode) -> PbRegionResponse { + PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: code as u32, + ..Default::default() }), + }), + affected_rows: 1, + ..Default::default() + } + } + + fn build_insert_request() -> RegionRequest { + RegionRequest { + header: Some(RegionRequestHeader { + dbname: "shared-request-test".to_string(), ..Default::default() - })) + }), + body: Some(region_request::Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id: 42, + rows: Some(Rows { + schema: vec![ + PbColumnSchema { + column_name: "message".to_string(), + datatype: ColumnDataType::String as i32, + ..Default::default() + }, + PbColumnSchema { + column_name: "payload".to_string(), + datatype: ColumnDataType::Binary as i32, + ..Default::default() + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::StringValue( + "hello 世界".to_string(), + )), + }, + Value { + value_data: Some(ValueData::BinaryValue(vec![0, 1, 127, 255])), + }, + ], + }], + }), + ..Default::default() + }], + })), } } #[tokio::test] - async fn test_inserts_and_bulk_insert_use_transport_compression() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let zstd_headers = Arc::new(std::sync::Mutex::new(Vec::new())); - let service = CompressionRecordingRegionService { - zstd_headers: zstd_headers.clone(), - }; - let server = tokio::spawn(async move { - tonic::transport::Server::builder() - .add_service( - RegionServer::new(service) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd), - ) - .serve_with_incoming(TcpListenerStream::new(listener)) + async fn test_owned_and_shared_requests_preserve_payload_and_compression() { + let server = RegionTestServer::start().await; + let shared = Arc::new(build_insert_request()); + for (send, accept) in [(false, false), (true, false), (false, true), (true, true)] { + let requester = server.create_requester(ChannelConfig::default(), send, accept); + let channel = + tonic::transport::Endpoint::from_shared(format!("http://{}", server.addr)) + .unwrap() + .connect() + .await + .unwrap(); + let mut baseline = api::v1::region::region_client::RegionClient::new(channel); + if send { + baseline = baseline.send_compressed(CompressionEncoding::Zstd); + } + if accept { + baseline = baseline.accept_compressed(CompressionEncoding::Zstd); + } + assert_eq!( + baseline + .handle(build_insert_request()) + .await + .unwrap() + .into_inner() + .affected_rows, + 1 + ); + assert_eq!( + requester + .handle(build_insert_request()) + .await + .unwrap() + .affected_rows, + 1 + ); + assert_eq!( + requester + .handle(Arc::clone(&shared)) + .await + .unwrap() + .affected_rows, + 1 + ); + } + let requests = server.service.requests.lock().unwrap(); + assert_eq!(requests.len(), 12); + for ((request, sends, accepts), (send, accept)) in requests.iter().zip( + [(false, false), (true, false), (false, true), (true, true)] + .into_iter() + .flat_map(|settings| [settings, settings, settings]), + ) { + assert_eq!(request, shared.as_ref()); + assert_eq!((*sends, *accepts), (send, accept)); + } + } + + #[tokio::test] + async fn test_shared_request_can_be_resent_after_transport_error() { + let server = RegionTestServer::start().await; + let mut status = Status::unavailable("migrating"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + RetryHint::Retryable.as_str().parse().unwrap(), + ); + server + .service + .responses + .lock() + .unwrap() + .push_back(Err(status)); + let requester = server.create_requester(ChannelConfig::default(), false, false); + let shared = Arc::new(build_insert_request()); + let error = requester.handle(Arc::clone(&shared)).await.unwrap_err(); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert!( + matches!(error, Error::RegionServer { addr, code: tonic::Code::Unavailable, .. } if addr == server.addr) + ); + requester.handle(Arc::clone(&shared)).await.unwrap(); + let requests = server.service.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(&requests[0].0, shared.as_ref()); + assert_eq!(&requests[1].0, shared.as_ref()); + } + + #[test] + fn test_retry_hint_wire_compatibility() { + for (name, wire_hint, expected_hint) in [ + ("legacy_missing", None, RetryHint::NonRetryable), + ("legacy_retryable", Some("retryable"), RetryHint::Retryable), + ( + "legacy_non_retryable", + Some("non_retryable"), + RetryHint::NonRetryable, + ), + ( + "route_refresh", + Some("retry_after_route_refresh"), + RetryHint::RetryAfterRouteRefresh, + ), + ( + "unknown_future_hint", + Some("future_retry_hint"), + RetryHint::NonRetryable, + ), + ] { + let mut status = Status::failed_precondition("region unavailable"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::RegionNotReady as u32) + .to_string() + .parse() + .unwrap(), + ); + if let Some(hint) = wire_hint { + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.parse().unwrap(), + ); + } + let error: Error = status.into(); + assert_eq!(error.retry_hint(), expected_hint, "{name}"); + assert_eq!(error.status_code(), StatusCode::RegionNotReady, "{name}"); + assert_eq!( + error.tonic_code(), + Some(tonic::Code::FailedPrecondition), + "{name}" + ); + } + } + + #[tokio::test] + async fn test_transport_error_parity_and_datanode_mapping() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + for hint in [ + RetryHint::Retryable, + RetryHint::RetryAfterRouteRefresh, + RetryHint::NonRetryable, + ] { + let mut status = Status::failed_precondition("region unavailable"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.as_str().parse().unwrap(), + ); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::RegionNotFound as u32) + .to_string() + .parse() + .unwrap(), + ); + server + .service + .responses + .lock() + .unwrap() + .extend((0..4).map(|_| Err(status.clone()))); + let mut baseline = api::v1::region::region_client::RegionClient::connect(format!( + "http://{}", + server.addr + )) + .await + .unwrap(); + let baseline: Error = baseline + .handle(build_insert_request()) .await - .unwrap(); - }); - let client = Client::with_urls([addr.to_string()]); - let requester = RegionRequester::new(client.clone(), true, true); - let send_only_requester = RegionRequester::new(client.clone(), true, false); - let accept_only_requester = RegionRequester::new(client.clone(), false, true); - let disabled_requester = RegionRequester::new(client, false, false); + .unwrap_err() + .into(); + assert_eq!(baseline.status_code(), StatusCode::RegionNotFound); + assert_eq!(baseline.retry_hint(), hint); + assert_eq!(baseline.tonic_code(), Some(tonic::Code::FailedPrecondition)); + let owned = requester.handle(build_insert_request()).await.unwrap_err(); + let shared = requester + .handle(Arc::new(build_insert_request())) + .await + .unwrap_err(); + for error in [owned, shared] { + assert_eq!(error.status_code(), baseline.status_code()); + assert_eq!(error.retry_hint(), baseline.retry_hint()); + assert_eq!(error.tonic_code(), baseline.tonic_code()); + } + let error = Datanode::handle(&requester, build_insert_request()) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), hint); + match hint { + RetryHint::Retryable | RetryHint::RetryAfterRouteRefresh => assert!(matches!( + error, + meta_error::Error::RetryLater { + clean_poisons: false, + .. + } + )), + RetryHint::NonRetryable => { + assert!(matches!(error, meta_error::Error::External { .. })) + } + } + } + } + + #[tokio::test] + async fn test_shared_request_preserves_response_header_errors() { + let server = RegionTestServer::start().await; + server.service.responses.lock().unwrap().extend([ + Ok(PbRegionResponse::default()), + Ok(build_response(StatusCode::RegionNotFound)), + ]); + let requester = server.create_requester(ChannelConfig::default(), false, false); + let shared = Arc::new(build_insert_request()); + assert!(matches!( + requester.handle(Arc::clone(&shared)).await.unwrap_err(), + Error::IllegalDatabaseResponse { .. } + )); + assert!(matches!( + requester.handle(shared).await.unwrap_err(), + Error::Server { + code: StatusCode::RegionNotFound, + .. + } + )); + } + + #[tokio::test] + async fn test_shared_request_respects_configured_message_limits() { + let server = RegionTestServer::start().await; + let shared = Arc::new(build_insert_request()); + for send_limit in [true, false] { + let mut config = ChannelConfig::default(); + if send_limit { + config.max_send_message_size = "1B".parse().unwrap(); + } else { + config.max_recv_message_size = "1B".parse().unwrap(); + } + let mut baseline = api::v1::region::region_client::RegionClient::connect(format!( + "http://{}", + server.addr + )) + .await + .unwrap() + .max_encoding_message_size(config.max_send_message_size.as_bytes() as usize) + .max_decoding_message_size(config.max_recv_message_size.as_bytes() as usize); + let baseline: Error = baseline + .handle(build_insert_request()) + .await + .unwrap_err() + .into(); + if !send_limit { + assert_eq!( + baseline.tonic_code(), + Some(tonic::Code::OutOfRange), + "{baseline:?}" + ); + } + let requester = server.create_requester(config, false, false); + let error = requester.handle(Arc::clone(&shared)).await.unwrap_err(); + assert_eq!( + error.tonic_code(), + baseline.tonic_code(), + "shared: {error:?}; generated: {baseline:?}" + ); + assert_eq!( + error.status_code(), + baseline.status_code(), + "shared: {error:?}; generated: {baseline:?}" + ); + assert_eq!( + error.retry_hint(), + baseline.retry_hint(), + "shared: {error:?}; generated: {baseline:?}" + ); + assert_eq!( + server.service.requests.lock().unwrap().len(), + if send_limit { 0 } else { 2 } + ); + } + let requester = server.create_requester(ChannelConfig::default(), false, false); + requester.handle(shared).await.unwrap(); + } + + #[tokio::test] + async fn test_inserts_and_bulk_insert_use_transport_compression() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), true, true); + let send_only_requester = server.create_requester(ChannelConfig::default(), true, false); + let accept_only_requester = server.create_requester(ChannelConfig::default(), false, true); + let disabled_requester = server.create_requester(ChannelConfig::default(), false, false); let inserts = || RegionRequest { body: Some(region_request::Body::Inserts(Default::default())), ..Default::default() @@ -502,9 +902,15 @@ mod test { (false, false), (false, false) ], - *zstd_headers.lock().unwrap() + server + .service + .requests + .lock() + .unwrap() + .iter() + .map(|(_, send, accept)| (*send, *accept)) + .collect::>() ); - server.abort(); } #[test] @@ -566,13 +972,13 @@ mod test { let result = check_response_header(&None); assert!(matches!( result.unwrap_err(), - IllegalDatabaseResponse { .. } + Error::IllegalDatabaseResponse { .. } )); let result = check_response_header(&Some(ResponseHeader { status: None })); assert!(matches!( result.unwrap_err(), - IllegalDatabaseResponse { .. } + Error::IllegalDatabaseResponse { .. } )); let result = check_response_header(&Some(ResponseHeader { @@ -591,7 +997,7 @@ mod test { })); assert!(matches!( result.unwrap_err(), - IllegalDatabaseResponse { .. } + Error::IllegalDatabaseResponse { .. } )); let result = check_response_header(&Some(ResponseHeader { @@ -600,7 +1006,7 @@ mod test { err_msg: "blabla".to_string(), }), })); - let Server { code, msg, .. } = result.unwrap_err() else { + let Error::Server { code, msg, .. } = result.unwrap_err() else { unreachable!() }; assert_eq!(code, StatusCode::Internal); diff --git a/src/client/src/region/codec.rs b/src/client/src/region/codec.rs new file mode 100644 index 00000000000..9476385df36 --- /dev/null +++ b/src/client/src/region/codec.rs @@ -0,0 +1,63 @@ +// Copyright 2023 Greptime Team +// +// 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 std::borrow::Borrow; +use std::marker::PhantomData; + +use api::v1::region::{RegionRequest, RegionResponse}; +use prost::Message; +use tonic::Status; +use tonic::codec::{Codec, EncodeBuf, Encoder}; +use tonic_prost::ProstDecoder; + +/// Encodes owned or shared region requests using the existing protobuf wire format. +pub(crate) struct RegionCodec(PhantomData); + +impl Default for RegionCodec { + fn default() -> Self { + Self(PhantomData) + } +} + +impl Codec for RegionCodec +where + R: Borrow + Send + 'static, +{ + type Encode = R; + type Decode = RegionResponse; + type Encoder = RegionEncoder; + type Decoder = ProstDecoder; + + fn encoder(&mut self) -> Self::Encoder { + RegionEncoder(PhantomData) + } + + fn decoder(&mut self) -> Self::Decoder { + ProstDecoder::default() + } +} + +/// Borrows the protobuf payload while consuming its owned or shared container. +pub(crate) struct RegionEncoder(PhantomData); + +impl> Encoder for RegionEncoder { + type Item = R; + type Error = Status; + + fn encode(&mut self, item: R, buf: &mut EncodeBuf<'_>) -> Result<(), Status> { + item.borrow() + .encode(buf) + .map_err(|error| Status::internal(error.to_string())) + } +} diff --git a/src/common/error/src/ext.rs b/src/common/error/src/ext.rs index 0c267dc15b2..f2d01914818 100644 --- a/src/common/error/src/ext.rs +++ b/src/common/error/src/ext.rs @@ -46,6 +46,10 @@ use crate::status_code::StatusCode; pub enum RetryHint { /// The operation may succeed if retried later. Retryable, + /// The operation may succeed after refreshing routing and partition information. + /// The request may need rebuilding when its partition assignment changes. + /// This hint does not imply that a failed operation had no side effects. + RetryAfterRouteRefresh, /// Retrying the same operation is not expected to help. /// /// This is the default for errors that do not explicitly opt in to retry. @@ -53,16 +57,23 @@ pub enum RetryHint { } const RETRY_HINT_RETRYABLE: &str = "retryable"; +const RETRY_HINT_AFTER_ROUTE_REFRESH: &str = "retry_after_route_refresh"; const RETRY_HINT_NON_RETRYABLE: &str = "non_retryable"; impl RetryHint { + /// Returns whether another attempt may succeed, possibly after refreshing routes. pub fn is_retryable(self) -> bool { - matches!(self, RetryHint::Retryable) + matches!( + self, + RetryHint::Retryable | RetryHint::RetryAfterRouteRefresh + ) } + /// Returns the retry hint encoded in gRPC error metadata. pub fn as_str(self) -> &'static str { match self { RetryHint::Retryable => RETRY_HINT_RETRYABLE, + RetryHint::RetryAfterRouteRefresh => RETRY_HINT_AFTER_ROUTE_REFRESH, RetryHint::NonRetryable => RETRY_HINT_NON_RETRYABLE, } } @@ -90,6 +101,7 @@ impl FromStr for RetryHint { fn from_str(s: &str) -> Result { match s { RETRY_HINT_RETRYABLE => Ok(RetryHint::Retryable), + RETRY_HINT_AFTER_ROUTE_REFRESH => Ok(RetryHint::RetryAfterRouteRefresh), RETRY_HINT_NON_RETRYABLE => Ok(RetryHint::NonRetryable), _ => Err(()), } @@ -133,7 +145,9 @@ pub trait ErrorExt: StackError { /// /// Implementations should return [`RetryHint::Retryable`] only when retrying the /// same operation may succeed without changing the request. The default is - /// [`RetryHint::NonRetryable`] to avoid accidental retry loops. + /// [`RetryHint::NonRetryable`] to avoid accidental retry loops. Use + /// [`RetryHint::RetryAfterRouteRefresh`] when routing or partition information + /// must be refreshed before another attempt. fn retry_hint(&self) -> RetryHint { RetryHint::NonRetryable } diff --git a/src/common/error/tests/ext.rs b/src/common/error/tests/ext.rs index 8be8afa4748..c34a47c400f 100644 --- a/src/common/error/tests/ext.rs +++ b/src/common/error/tests/ext.rs @@ -164,6 +164,15 @@ fn test_transparent_flag() { #[test] fn test_retry_hint_helpers() { assert!(RetryHint::Retryable.is_retryable()); + assert!(RetryHint::RetryAfterRouteRefresh.is_retryable()); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + "retry_after_route_refresh" + ); + assert_eq!( + RetryHint::from_str("retry_after_route_refresh").unwrap(), + RetryHint::RetryAfterRouteRefresh + ); assert!(!RetryHint::NonRetryable.is_retryable()); assert_eq!(RetryHint::Retryable.as_str(), "retryable"); diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 5c1febb9b4d..297af792799 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -1346,8 +1346,11 @@ impl ErrorExt for Error { use Error::*; match self { - RetryLater { .. } - | GetLatestCacheRetryExceeded { .. } + RetryLater { source, .. } => match source.retry_hint() { + RetryHint::RetryAfterRouteRefresh => RetryHint::RetryAfterRouteRefresh, + RetryHint::Retryable | RetryHint::NonRetryable => RetryHint::Retryable, + }, + GetLatestCacheRetryExceeded { .. } | NoLeader { .. } | ElectionNoLeader { .. } | ElectionLeaderLeaseExpired { .. } diff --git a/src/common/meta/src/instruction.rs b/src/common/meta/src/instruction.rs index 6ae074b0ac9..827f06e944d 100644 --- a/src/common/meta/src/instruction.rs +++ b/src/common/meta/src/instruction.rs @@ -49,7 +49,7 @@ pub struct InstructionError { pub message: String, /// Retry hint serialized as a string. #[serde( - serialize_with = "RetryHint::serialize_as_str", + serialize_with = "InstructionError::serialize_retry_hint", deserialize_with = "RetryHint::deserialize_from_str" )] pub retry_hint: RetryHint, @@ -108,6 +108,18 @@ impl InstructionError { retry_hint: error.retry_hint(), } } + + /// Keeps instruction replies readable by peers that only know the legacy hints. + fn serialize_retry_hint(hint: &RetryHint, serializer: S) -> Result + where + S: Serializer, + { + let hint = match hint { + RetryHint::RetryAfterRouteRefresh => RetryHint::Retryable, + hint => *hint, + }; + RetryHint::serialize_as_str(&hint, serializer) + } } impl Display for InstructionError { @@ -1450,6 +1462,26 @@ mod tests { ); } + #[test] + fn test_instruction_error_route_refresh_compatibility() { + let error = InstructionError::new( + StatusCode::RegionNotFound, + "region not found", + RetryHint::RetryAfterRouteRefresh, + ); + let serialized = serde_json::to_string(&error).unwrap(); + assert_eq!( + r#"{"code":4005,"message":"region not found","retry_hint":"retryable"}"#, + serialized + ); + let decoded: InstructionError = serde_json::from_str(&serialized).unwrap(); + assert_eq!(decoded.code, error.code); + assert_eq!(decoded.message, error.message); + assert_eq!(decoded.retry_hint, RetryHint::Retryable); + assert_eq!(error.retry_hint, RetryHint::RetryAfterRouteRefresh); + assert_eq!(error.retry_hint.as_str(), "retry_after_route_refresh"); + } + #[test] fn test_instruction_error_from_error() { let error = MockError::new(StatusCode::RegionNotFound); diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index e9b5d5ec1dc..3f61d4b3d76 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -561,8 +561,8 @@ impl ErrorExt for Error { use Error::*; match self { + RegionNotReady { .. } | RegionNotFound { .. } => RetryHint::RetryAfterRouteRefresh, RegionBusy { .. } - | RegionNotReady { .. } | ConcurrentQueryLimiterClosed { .. } | ConcurrentQueryLimiterTimeout { .. } => RetryHint::Retryable, NewPlanDecoder { source, .. } | ExecuteLogicalPlan { source, .. } => { @@ -602,14 +602,28 @@ mod tests { use super::*; #[test] - fn test_region_state_hints_are_retryable() { + fn test_region_state_retry_hints() { let region_id = RegionId::new(1024, 1); - - let err = RegionBusySnafu { region_id }.build(); - assert_eq!(err.retry_hint(), RetryHint::Retryable); - - let err = RegionNotReadySnafu { region_id }.build(); - assert_eq!(err.retry_hint(), RetryHint::Retryable); + for (error, status, hint) in [ + ( + RegionBusySnafu { region_id }.build(), + StatusCode::RegionBusy, + RetryHint::Retryable, + ), + ( + RegionNotReadySnafu { region_id }.build(), + StatusCode::RegionNotReady, + RetryHint::RetryAfterRouteRefresh, + ), + ( + RegionNotFoundSnafu { region_id }.build(), + StatusCode::RegionNotFound, + RetryHint::RetryAfterRouteRefresh, + ), + ] { + assert_eq!(error.status_code(), status); + assert_eq!(error.retry_hint(), hint); + } } #[test] diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index ef333cff04d..a1ccc382ef1 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -2589,7 +2589,7 @@ mod tests { .unwrap_err(); assert_matches!(&err, error::Error::HandleBatchOpenRequest { .. }); assert_eq!(err.status_code(), StatusCode::RegionNotFound); - assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); } struct CurrentEngineTest { diff --git a/src/metric-engine/src/error.rs b/src/metric-engine/src/error.rs index 0803223ba33..f9d61d6af31 100644 --- a/src/metric-engine/src/error.rs +++ b/src/metric-engine/src/error.rs @@ -486,6 +486,9 @@ impl ErrorExt for Error { use Error::*; match self { + PhysicalRegionNotFound { .. } | LogicalRegionNotFound { .. } => { + RetryHint::RetryAfterRouteRefresh + } CreateMitoRegion { source, .. } | OpenMitoRegion { source, .. } | CleanUpMitoRegion { source, .. } @@ -510,3 +513,20 @@ impl ErrorExt for Error { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_missing_region_retry_hints() { + let region_id = RegionId::new(1024, 1); + for error in [ + PhysicalRegionNotFoundSnafu { region_id }.build(), + LogicalRegionNotFoundSnafu { region_id }.build(), + ] { + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + } +} diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index 313860c0857..c7b4be14587 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -24,6 +24,7 @@ use api::v1::Rows; use async_trait::async_trait; use common_base::Plugins; use common_base::readable_size::ReadableSize; +use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; @@ -1006,6 +1007,7 @@ async fn test_region_write_buffer_rejects_only_full_region_queue() { .expect("the stalled queue should be rejected at the region hard limit") .unwrap() .unwrap_err(); + assert_eq!(err.retry_hint(), RetryHint::Retryable); assert_matches!( err.into_inner().as_any().downcast_ref::().unwrap(), Error::RejectWrite { .. } diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index bd823708509..1b2a7ffee13 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -18,11 +18,14 @@ use std::sync::Arc; use std::time::Duration; use std::{assert_matches, fs}; -use api::v1::Rows; -use common_error::ext::ErrorExt; +use api::v1::{ArrowIpc, Rows}; +use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_function::utils::partition_expr_version; +use common_grpc::flight; use common_recordbatch::RecordBatches; +use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; +use datatypes::arrow::record_batch::RecordBatch; use datatypes::value::Value; use object_store::Buffer; use object_store::layers::mock::{ @@ -34,8 +37,9 @@ use store_api::region_engine::{ RegionEngine, RemapManifestsRequest, SetRegionRoleStateResponse, SettableRegionRoleState, }; use store_api::region_request::{ - ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, RegionFlushRequest, - RegionPutRequest, RegionRequest, RegionTruncateRequest, StagingPartitionDirective, + ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, RegionBulkInsertsRequest, + RegionFlushRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, + StagingPartitionDirective, }; use store_api::storage::{RegionId, ScanRequest}; @@ -254,7 +258,7 @@ fn default_partition_expr() -> String { } #[tokio::test] -async fn test_staging_reject_all_writes_rejects_put() { +async fn test_staging_reject_all_writes_rejects_put_and_bulk() { let mut env = TestEnv::new().await; let engine = env.create_engine(MitoConfig::default()).await; @@ -281,23 +285,50 @@ async fn test_staging_reject_all_writes_rejects_put() { schema: column_schemas, rows: build_rows(0, 2), }; - let err = engine - .handle_request( + let payload = RecordBatch::try_from_iter([ + ( + "tag_0", + Arc::new(StringArray::from(vec!["0", "1"])) as ArrayRef, + ), + ( + "field_0", + Arc::new(Float64Array::from(vec![0.0, 1.0])) as ArrayRef, + ), + ( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![0, 1000])) as ArrayRef, + ), + ]) + .unwrap(); + let (schema, data_header, data) = flight::record_batch_to_ipc(payload.clone()).unwrap(); + for request in [ + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows, + hint: None, + partition_expr_version: None, + }), + RegionRequest::BulkInserts(RegionBulkInsertsRequest { + skip_wal: false, region_id, - RegionRequest::Put(RegionPutRequest { - skip_wal: false, - rows, - hint: None, - partition_expr_version: None, - }), - ) - .await - .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::StorageUnavailable); - assert_matches!( - err.into_inner().as_any().downcast_ref::().unwrap(), - Error::RejectWrite { .. } - ); + payload, + raw_data: ArrowIpc { + schema, + data_header, + payload: data, + }, + partition_expr_version: None, + aligned_schema_version: None, + }), + ] { + let err = engine.handle_request(region_id, request).await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::StorageUnavailable); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RejectWriteInStaging { .. } + ); + } } #[tokio::test] diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index f738044c9ef..5af4b7cec09 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -641,6 +641,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Region {} rejects writes while in staging", region_id))] + RejectWriteInStaging { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to compact region {}", region_id))] CompactRegion { region_id: RegionId, @@ -1541,7 +1548,7 @@ impl ErrorExt for Error { RegionClosed { .. } => StatusCode::Cancelled, StaleCompactionExecution { .. } => StatusCode::Cancelled, RegionTruncated { .. } => StatusCode::Cancelled, - RejectWrite { .. } => StatusCode::StorageUnavailable, + RejectWrite { .. } | RejectWriteInStaging { .. } => StatusCode::StorageUnavailable, CompactRegion { source, .. } => source.status_code(), EditRegion { source, .. } => source.status_code(), CompatReader { .. } => StatusCode::Unexpected, @@ -1639,6 +1646,13 @@ impl ErrorExt for Error { use Error::*; match self { + PartitionExprVersionMismatch { .. } + | RegionNotFound { .. } + | RejectWriteInStaging { .. } + | RegionState { + state: RegionRoleState::Follower, + .. + } => RetryHint::RetryAfterRouteRefresh, ReadParquet { .. } | WriteParquet { .. } | RejectWrite { .. } @@ -1720,6 +1734,63 @@ mod tests { use snafu::IntoError; use super::*; + use crate::region::RegionLeaderState; + + #[test] + fn test_partition_version_mismatch_retry_hint() { + let error = PartitionExprVersionMismatchSnafu { + region_id: RegionId::new(1, 1), + request_version: 1u64, + expected_version: 2u64, + } + .build(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + + #[test] + fn test_region_state_retry_hint() { + let region_id = RegionId::new(1, 1); + let writable = RegionRoleState::Leader(RegionLeaderState::Writable); + let follower_error = RegionStateSnafu { + region_id, + state: RegionRoleState::Follower, + expect: writable, + } + .build(); + assert_eq!( + follower_error.retry_hint(), + RetryHint::RetryAfterRouteRefresh + ); + assert_eq!(follower_error.status_code(), StatusCode::RegionNotReady); + let leader_error = RegionStateSnafu { + region_id, + state: writable, + expect: RegionRoleState::Follower, + } + .build(); + assert_eq!(leader_error.retry_hint(), RetryHint::Retryable); + for (error, status, hint) in [ + ( + RejectWriteSnafu { region_id }.build(), + StatusCode::StorageUnavailable, + RetryHint::Retryable, + ), + ( + RejectWriteInStagingSnafu { region_id }.build(), + StatusCode::StorageUnavailable, + RetryHint::RetryAfterRouteRefresh, + ), + ( + RegionNotFoundSnafu { region_id }.build(), + StatusCode::RegionNotFound, + RetryHint::RetryAfterRouteRefresh, + ), + ] { + assert_eq!(error.status_code(), status); + assert_eq!(error.retry_hint(), hint); + } + } #[test] fn test_manifest_update_persistence() { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 41af569c649..1383a3e9a5d 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -26,7 +26,7 @@ use store_api::storage::RegionId; use crate::error::{ InvalidRequestSnafu, PartitionExprVersionMismatchSnafu, RegionNotFoundSnafu, RegionStateSnafu, - RejectWriteSnafu, Result, + RejectWriteInStagingSnafu, RejectWriteSnafu, Result, }; use crate::metrics; use crate::metrics::{ @@ -335,7 +335,7 @@ impl RegionWorkerLoop { if region.reject_all_writes_in_staging() { sender_req .sender - .send(RejectWriteSnafu { region_id }.fail()); + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } @@ -399,7 +399,7 @@ impl RegionWorkerLoop { if region.reject_all_writes_in_staging() { sender_req .sender - .send(RejectWriteSnafu { region_id }.fail()); + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let expected_version = region.expected_partition_expr_version(); @@ -481,7 +481,9 @@ impl RegionWorkerLoop { RegionRoleState::Leader(RegionLeaderState::Writable) | RegionRoleState::Leader(RegionLeaderState::Staging) => { if region.reject_all_writes_in_staging() { - bulk_req.sender.send(RejectWriteSnafu { region_id }.fail()); + bulk_req + .sender + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let region_ctx = RegionWriteCtx::new( @@ -529,7 +531,9 @@ impl RegionWorkerLoop { continue; }; if region.reject_all_writes_in_staging() { - bulk_req.sender.send(RejectWriteSnafu { region_id }.fail()); + bulk_req + .sender + .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } let expected_version = region.expected_partition_expr_version(); From a2d28e6060df4c1e40de23df995d0f04f0adb193 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 02:30:14 +0000 Subject: [PATCH 21/46] refactor(operator): centralize write partitioning and dispatch (#9541) * refactor(operator): centralize write partitioning and dispatch Signed-off-by: WenyXu * fix(operator): preserve write dispatch behavior during refactor Signed-off-by: WenyXu * style(operator): import write dispatch error paths Signed-off-by: WenyXu * style(operator): align dispatch imports and test utilities Signed-off-by: WenyXu * style(operator): use imported types in partition tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/operator/AGENTS.md | 1 + src/operator/src/bulk_insert.rs | 328 +++++++--------- src/operator/src/delete.rs | 73 ++-- src/operator/src/insert.rs | 93 ++--- src/operator/src/lib.rs | 1 + src/operator/src/req_convert.rs | 3 + .../src/req_convert/common/partitioner.rs | 237 ++++++++++-- src/operator/src/write_dispatcher.rs | 256 +++++++++++++ src/operator/src/write_dispatcher/tests.rs | 350 ++++++++++++++++++ 9 files changed, 1014 insertions(+), 328 deletions(-) create mode 100644 src/operator/src/write_dispatcher.rs create mode 100644 src/operator/src/write_dispatcher/tests.rs diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 49cbd617567..46359eb5c50 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -14,6 +14,7 @@ insert/delete routing. | Statements | `src/operator/src/statement.rs`, `src/operator/src/statement/` | `StatementExecutor` and per-statement handlers | | Inserts | `src/operator/src/insert.rs`, `src/operator/src/bulk_insert.rs` | Insert, auto-create/alter, routing, dispatch | | Deletes | `src/operator/src/delete.rs` | Delete conversion and dispatch | +| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch and batched peer resolution | | Request conversion | `src/operator/src/req_convert.rs`, `src/operator/src/req_convert/` | Table/row/column requests to region requests | | Region requests | `src/operator/src/region_req_factory.rs` | Region-level DDL/DML request construction | | Procedures | `src/operator/src/procedure.rs` | Metasrv procedure administration | diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index f8d58f5d173..d5aa7168b44 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -14,33 +14,33 @@ use std::collections::HashSet; -use ahash::{HashMap, HashMapExt}; +use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, PartitionExprVersion}; -use arrow::compute::filter_record_batch; +use arrow::compute; use arrow::record_batch::RecordBatch; use bytes::Bytes; use common_base::AffectedRows; -use common_error::ext::BoxedError; use common_grpc::FlightData; -use common_grpc::flight::{FlightEncoder, FlightMessage, record_batch_to_ipc}; +use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_telemetry::error; use common_telemetry::tracing_context::TracingContext; -use futures::future::{join_all, try_join_all}; +use futures::future; use meter_core::data::MeterRecord; use meter_macros::write_meter; use session::context::{Channel, QueryContextRef}; -use snafu::{ResultExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfoRef; use crate::error::Result; use crate::insert::Inserter; -use crate::req_convert::insert::extract_timestamps; +use crate::req_convert::TablePartitioner; +use crate::req_convert::insert; use crate::{error, metrics}; impl Inserter { @@ -54,72 +54,28 @@ impl Inserter { batch: RecordBatch, ctx: QueryContextRef, ) -> Result { - let (rule, versions) = self - .partition_manager - .find_table_partition_rule(&table_info) + let requests = TablePartitioner::try_new(&self.partition_manager, &table_info) .await - .context(error::InvalidPartitionSnafu)?; - let masks = rule - .split_record_batch(&batch) - .context(error::SplitInsertSnafu)?; - let mut writes = Vec::with_capacity(masks.len()); - for (region_number, mask) in masks { - if mask.select_none() { - continue; - } - let region_id = RegionId::new(table_info.table_id(), region_number); - let selected = if mask.select_all() { - batch.clone() - } else { - filter_record_batch(&batch, mask.array()).context(error::ComputeArrowSnafu)? - }; - let (schema, data_header, payload) = record_batch_to_ipc(selected) - .map_err(BoxedError::new) - .context(error::ExternalSnafu)?; - let peer = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - let request = RegionRequest { + .context(error::InvalidPartitionSnafu)? + .partition_bulk_insert_requests(batch, None, ctx.skip_wal())?; + let writes = requests + .into_iter() + .map(|request| RegionRequest { header: Some(RegionRequestHeader { dbname: ctx.get_db_string(), tracing_context: TracingContext::from_current_span().to_w3c(), ..Default::default() }), - body: Some(region_request::Body::BulkInsert(BulkInsertRequest { - skip_wal: ctx.skip_wal(), - region_id: region_id.as_u64(), - partition_expr_version: versions - .get(®ion_number) - .copied() - .flatten() - .map(|value| PartitionExprVersion { value }), - // Let the datanode revalidate against its current schema. - aligned_schema_version: None, - body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { - schema, - data_header, - payload, - })), - })), - }; - writes.push((peer, request)); - } - let results = join_all(writes.into_iter().map(|(peer, request)| async move { - self.node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(error::RequestInsertsSnafu) - })) - .await; - let affected_rows = results - .into_iter() - .map(|result| result.map(|response| response.affected_rows)) - .sum::>()?; - Ok(affected_rows) + body: Some(region_request::Body::BulkInsert(request)), + }) + .collect(); + let writes = self.write_dispatcher.prepare(writes).await?; + Ok(self + .write_dispatcher + .dispatch_prepared(writes, None) + .await + .context(error::RequestInsertsSnafu)? + .affected_rows) } /// Handle bulk insert request. @@ -208,12 +164,6 @@ impl Inserter { .get(®ion_number) .copied() .unwrap_or_default(); - let datanode = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - let request = RegionRequest { header: Some(RegionRequestHeader { tracing_context: TracingContext::from_current_span().to_w3c(), @@ -233,12 +183,15 @@ impl Inserter { })), }; - let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["datanode_handle"]) - .start_timer(); - let datanode = self.node_manager.datanode(&datanode).await; - let result = datanode - .handle(request) + let writes = self.write_dispatcher.prepare(vec![request]).await?; + let result = self + .write_dispatcher + .dispatch_prepared( + writes, + Some( + metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), + ), + ) .await .context(error::RequestRegionSnafu) .map(|r| r.affected_rows); @@ -250,121 +203,126 @@ impl Inserter { return result; } - let mut mask_per_datanode = HashMap::with_capacity(region_masks.len()); - for (region_number, mask) in region_masks { - let region_id = RegionId::new(table_id, region_number); - let datanode = self - .partition_manager - .find_region_leader(region_id) - .await - .context(error::FindRegionLeaderSnafu)?; - mask_per_datanode - .entry(datanode) - .or_insert_with(Vec::new) - .push((region_id, mask)); - } + let regions = region_masks + .keys() + .map(|region| RegionId::new(table_id, *region)) + .collect::>(); + let mut peers = self + .write_dispatcher + .resolve_peers(®ions) + .await + .context(error::FindRegionLeaderSnafu)?; let wait_all_datanode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["wait_all_datanode"]) .start_timer(); - let mut handles = Vec::with_capacity(mask_per_datanode.len()); + let mut handles = Vec::with_capacity(region_masks.len()); - for (peer, masks) in mask_per_datanode { - for (region_id, mask) in masks { - if mask.select_none() { - continue; - } - let partition_expr_version = partition_versions - .get(®ion_id.region_number()) - .copied() - .unwrap_or_default(); - let rb = record_batch.clone(); - let schema_bytes = schema_bytes.clone(); - let node_manager = self.node_manager.clone(); - let peer = peer.clone(); - let raw_header_and_data = if mask.select_all() { - Some(( - raw_flight_data.data_header.clone(), - raw_flight_data.data_body.clone(), - )) - } else { - None - }; - let handle: common_runtime::JoinHandle> = - common_runtime::spawn_global(async move { - let (header, payload) = if mask.select_all() { - // SAFETY: raw data must be present, we can avoid re-encoding. - raw_header_and_data.unwrap() - } else { - let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["filter"]) - .start_timer(); - let batch = filter_record_batch(&rb, mask.array()) - .context(error::ComputeArrowSnafu)?; - filter_timer.observe_duration(); - metrics::BULK_REQUEST_ROWS - .with_label_values(&["rows_per_region"]) - .observe(batch.num_rows() as f64); - - let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["encode"]) - .start_timer(); - let mut iter = FlightEncoder::default() - .encode(FlightMessage::RecordBatch(batch)) - .into_iter(); - let Some(flight_data) = iter.next() else { - // Safety: `iter` on a type of `Vec1`, which is guaranteed to have - // at least one element. - unreachable!() - }; - ensure!( - iter.next().is_none(), - error::NotSupportedSnafu { - feat: "bulk insert RecordBatch with dictionary arrays", - } - ); - encode_timer.observe_duration(); - (flight_data.data_header, flight_data.data_body) - }; - let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED - .with_label_values(&["datanode_handle"]) - .start_timer(); - let request = RegionRequest { - header: Some(RegionRequestHeader { - tracing_context: TracingContext::from_current_span().to_w3c(), - ..Default::default() - }), - body: Some(region_request::Body::BulkInsert(BulkInsertRequest { - skip_wal, - region_id: region_id.as_u64(), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - aligned_schema_version: None, - body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { - schema: schema_bytes, - data_header: header, - payload, - })), - })), - }; - - let datanode = node_manager.datanode(&peer).await; - datanode - .handle(request) - .await - .context(error::RequestRegionSnafu) - }); - handles.push(handle); + for (region_number, mask) in region_masks { + let region_id = RegionId::new(table_id, region_number); + if mask.select_none() { + continue; } + let partition_expr_version = partition_versions + .get(®ion_id.region_number()) + .copied() + .unwrap_or_default(); + let peer = peers + .remove(®ion_id) + .context(partition::error::FindLeaderSnafu { + region_id, + table_id, + }) + .context(error::FindRegionLeaderSnafu)?; + let dispatcher = self.write_dispatcher.clone(); + let rb = record_batch.clone(); + let schema_bytes = schema_bytes.clone(); + let raw_header_and_data = if mask.select_all() { + Some(( + raw_flight_data.data_header.clone(), + raw_flight_data.data_body.clone(), + )) + } else { + None + }; + let handle: common_runtime::JoinHandle> = + common_runtime::spawn_global(async move { + let (header, payload) = if mask.select_all() { + // SAFETY: raw data must be present, we can avoid re-encoding. + raw_header_and_data.unwrap() + } else { + let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["filter"]) + .start_timer(); + let batch = compute::filter_record_batch(&rb, mask.array()) + .context(error::ComputeArrowSnafu)?; + filter_timer.observe_duration(); + metrics::BULK_REQUEST_ROWS + .with_label_values(&["rows_per_region"]) + .observe(batch.num_rows() as f64); + + let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["encode"]) + .start_timer(); + let mut iter = FlightEncoder::default() + .encode(FlightMessage::RecordBatch(batch)) + .into_iter(); + let Some(flight_data) = iter.next() else { + // Safety: `iter` on a type of `Vec1`, which is guaranteed to have + // at least one element. + unreachable!() + }; + ensure!( + iter.next().is_none(), + error::NotSupportedSnafu { + feat: "bulk insert RecordBatch with dictionary arrays", + } + ); + encode_timer.observe_duration(); + (flight_data.data_header, flight_data.data_body) + }; + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal, + region_id: region_id.as_u64(), + partition_expr_version: partition_expr_version + .map(|value| PartitionExprVersion { value }), + aligned_schema_version: None, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: schema_bytes, + data_header: header, + payload, + })), + })), + }; + + dispatcher + .dispatch_prepared( + vec![(peer, request)], + Some( + metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["datanode_handle"]), + ), + ) + .await + .context(error::RequestRegionSnafu) + }); + handles.push(handle); } - let region_responses = try_join_all(handles).await.context(error::JoinTaskSnafu)?; + let results = future::try_join_all(handles) + .await + .context(error::JoinTaskSnafu)?; + let rows_inserted = results + .into_iter() + .map(|result| result.map(|response| response.affected_rows)) + .sum::>()?; wait_all_datanode_timer.observe_duration(); - let mut rows_inserted: usize = 0; - for res in region_responses { - rows_inserted += res?.affected_rows; - } crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[db_name.as_str()]) .inc_by(rows_inserted as u64); @@ -393,7 +351,7 @@ impl Inserter { return; } - let Ok(timestamps) = extract_timestamps( + let Ok(timestamps) = insert::extract_timestamps( &record_batch, &table_info .meta diff --git a/src/operator/src/delete.rs b/src/operator/src/delete.rs index 9641e0e27d4..abd8302c965 100644 --- a/src/operator/src/delete.rs +++ b/src/operator/src/delete.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use std::{iter, mem}; @@ -20,10 +20,8 @@ use api::v1::region::{DeleteRequests as RegionDeleteRequests, RegionRequestHeade use api::v1::{DeleteRequests, RowDeleteRequests}; use catalog::CatalogManagerRef; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; -use common_meta::peer::Peer; use common_query::Output; use common_telemetry::tracing_context::TracingContext; -use futures_util::future; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; @@ -31,34 +29,45 @@ use table::TableRef; use table::requests::DeleteRequest as TableDeleteRequest; use crate::error::{ - CatalogSnafu, FindRegionLeaderSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, - MissingTimeIndexColumnSnafu, RequestDeletesSnafu, Result, TableNotFoundSnafu, + CatalogSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, MissingTimeIndexColumnSnafu, + RequestDeletesSnafu, Result, TableNotFoundSnafu, }; use crate::region_req_factory::RegionRequestFactory; use crate::req_convert::common::preprocess_row_delete_requests; use crate::req_convert::delete::{ColumnToRow, RowToRegion, TableToRegion}; +use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef}; pub struct Deleter { catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, } pub type DeleterRef = Arc; impl Deleter { + /// Creates a row deleter with the default write dispatcher. pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, ) -> Self { Self { + write_dispatcher: Arc::new(WriteDispatcher::new( + node_manager, + partition_manager.clone(), + )), catalog_manager, partition_manager, - node_manager, } } + /// Installs the dispatcher shared with the other write paths. + pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { + self.write_dispatcher = dispatcher; + self + } + pub async fn handle_column_deletes( &self, requests: DeleteRequests, @@ -129,52 +138,22 @@ impl Deleter { ..Default::default() }); - let tasks = self - .group_requests_by_peer(requests) - .await? - .into_iter() - .map(|(peer, deletes)| { - let request = request_factory.build_delete(deletes); - let node_manager = self.node_manager.clone(); - common_runtime::spawn_global(async move { - node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(RequestDeletesSnafu) - }) - }); - let results = future::try_join_all(tasks).await.context(JoinTaskSnafu)?; - - let affected_rows = results - .into_iter() - .map(|resp| resp.map(|r| r.affected_rows)) - .sum::>()?; + let dispatcher = self.write_dispatcher.clone(); + let request = request_factory.build_delete(requests); + let writes = dispatcher.prepare(vec![request]).await?; + let affected_rows = common_runtime::spawn_global(async move { + dispatcher.dispatch_prepared(writes, None).await + }) + .await + .context(JoinTaskSnafu)? + .context(RequestDeletesSnafu)? + .affected_rows; crate::metrics::DIST_DELETE_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); Ok(affected_rows) } - async fn group_requests_by_peer( - &self, - requests: RegionDeleteRequests, - ) -> Result> { - let mut deletes: HashMap = HashMap::new(); - - for req in requests.requests { - let peer = self - .partition_manager - .find_region_leader(req.region_id.into()) - .await - .context(FindRegionLeaderSnafu)?; - deletes.entry(peer).or_default().requests.push(req); - } - - Ok(deletes) - } - async fn trim_columns( &self, mut requests: RowDeleteRequests, diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index c7f60c0883c..89c1195cda4 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -42,7 +42,7 @@ use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; -use common_meta::node_manager::{AffectedRows, NodeManagerRef}; +use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; use common_meta::rpc::ddl::TriggerReason; use common_query::Output; @@ -83,9 +83,8 @@ use table::table_reference::TableReference; use crate::batcher::PendingRowsBatcher; use crate::error::{ - CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, FindRegionLeaderSnafu, - InvalidInsertRequestSnafu, JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, - WriteRejectedSnafu, + CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, InvalidInsertRequestSnafu, + JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, WriteRejectedSnafu, }; use crate::expr_helper; use crate::region_req_factory::RegionRequestFactory; @@ -95,11 +94,13 @@ use crate::req_convert::insert::{ fill_reqs_with_impure_default, rows_to_record_batch, }; use crate::statement::StatementExecutor; +use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef}; pub struct Inserter { catalog_manager: CatalogManagerRef, pub(crate) partition_manager: PartitionRuleManagerRef, pub(crate) node_manager: NodeManagerRef, + pub(crate) write_dispatcher: WriteDispatcherRef, pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef, /// Server-side upper bound for auto table creation on write. /// When `false`, missing tables are never auto-created regardless of the @@ -292,6 +293,7 @@ impl Inserter { cost } + /// Creates a row inserter with the default write dispatcher. pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, @@ -300,6 +302,10 @@ impl Inserter { auto_create_table: bool, ) -> Self { Self { + write_dispatcher: Arc::new(WriteDispatcher::new( + node_manager.clone(), + partition_manager.clone(), + )), catalog_manager, partition_manager, node_manager, @@ -826,30 +832,16 @@ impl Inserter { has_instant_rows, )?; - // Write requests to datanode and wait for response - let write_tasks = self - .group_requests_by_peer(normal_requests) - .await? - .into_iter() - .map(|(peer, inserts)| { - let node_manager = self.node_manager.clone(); - let request = request_factory.build_insert(inserts); - common_runtime::spawn_global(async move { - node_manager - .datanode(&peer) - .await - .handle(request) - .await - .context(RequestInsertsSnafu) - }) - }); - let results = future::try_join_all(write_tasks) - .await - .context(JoinTaskSnafu)?; - let affected_rows = results - .into_iter() - .map(|resp| resp.map(|r| r.affected_rows)) - .sum::>()?; + let dispatcher = self.write_dispatcher.clone(); + let request = request_factory.build_insert(normal_requests); + let writes = dispatcher.prepare(vec![request]).await?; + let affected_rows = common_runtime::spawn_global(async move { + dispatcher.dispatch_prepared(writes, None).await + }) + .await + .context(JoinTaskSnafu)? + .context(RequestInsertsSnafu)? + .affected_rows; crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); @@ -859,40 +851,6 @@ impl Inserter { )) } - async fn group_requests_by_peer( - &self, - requests: RegionInsertRequests, - ) -> Result> { - // group by region ids first to reduce repeatedly call `find_region_leader` - // TODO(discord9): determine if a addition clone is worth it - let mut requests_per_region: HashMap = HashMap::new(); - for req in requests.requests { - let region_id = RegionId::from_u64(req.region_id); - requests_per_region - .entry(region_id) - .or_default() - .requests - .push(req); - } - - let mut inserts: HashMap = HashMap::new(); - - for (region_id, reqs) in requests_per_region { - let peer = self - .partition_manager - .find_region_leader(region_id) - .await - .context(FindRegionLeaderSnafu)?; - inserts - .entry(peer) - .or_default() - .requests - .extend(reqs.requests); - } - - Ok(inserts) - } - /// Returns `Some(reason)` if the config or request hint disables automatic /// table creation. Exempt private system tables are handled by /// [`Self::is_auto_create_exempt_private_table`]. @@ -1717,6 +1675,17 @@ impl Inserter { } } + /// Installs the shared dispatcher used by direct writes and batch flushes. + pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { + self.write_dispatcher = dispatcher; + self + } + + /// Returns the dispatcher shared with batching execution paths. + pub fn write_dispatcher(&self) -> &WriteDispatcherRef { + &self.write_dispatcher + } + pub fn node_manager(&self) -> &NodeManagerRef { &self.node_manager } diff --git a/src/operator/src/lib.rs b/src/operator/src/lib.rs index a971ca41f70..901f351e22c 100644 --- a/src/operator/src/lib.rs +++ b/src/operator/src/lib.rs @@ -29,3 +29,4 @@ pub mod table; #[cfg(any(test, feature = "testing"))] pub mod test_util; pub mod utils; +pub mod write_dispatcher; diff --git a/src/operator/src/req_convert.rs b/src/operator/src/req_convert.rs index 456ff0b328d..1161f62795d 100644 --- a/src/operator/src/req_convert.rs +++ b/src/operator/src/req_convert.rs @@ -15,3 +15,6 @@ pub(crate) mod common; pub mod delete; pub mod insert; + +pub use common::partitioner::Partitioner; +pub(crate) use common::partitioner::TablePartitioner; diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 3353e2fcda7..044d3b97db5 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -12,72 +12,241 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::region::{DeleteRequest, InsertRequest}; -use api::v1::{PartitionExprVersion, Rows}; +use std::collections::HashMap; + +use api::v1::region::{BulkInsertRequest, DeleteRequest, InsertRequest, bulk_insert_request}; +use api::v1::{ArrowIpc, PartitionExprVersion, Rows}; +use common_error::ext::BoxedError; +use common_grpc::flight; +use datatypes::arrow::compute; +use datatypes::arrow::record_batch::RecordBatch; +use partition::PartitionRuleRef; use partition::manager::PartitionRuleManager; +use partition::splitter::RowSplitter; use snafu::ResultExt; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, RegionNumber, TableId}; use table::metadata::TableInfo; -use crate::error::{Result, SplitDeleteSnafu, SplitInsertSnafu}; +use crate::error::{self, Result, SplitDeleteSnafu, SplitInsertSnafu}; +/// Loads table routing rules and converts rows to region write requests. pub struct Partitioner<'a> { partition_manager: &'a PartitionRuleManager, } +/// A single snapshot of a table's partition rule and expression versions. +pub(crate) struct TablePartitioner { + table_id: TableId, + rule: PartitionRuleRef, + versions: HashMap>, +} + impl<'a> Partitioner<'a> { + /// Creates a partitioner backed by the shared routing caches. pub fn new(partition_manager: &'a PartitionRuleManager) -> Self { Self { partition_manager } } + /// Routes insert rows using the current table partitions. pub async fn partition_insert_requests( &self, table_info: &TableInfo, rows: Rows, skip_wal: bool, ) -> Result> { - let table_id = table_info.table_id(); - let requests = self - .partition_manager - .split_rows(table_info, rows) + TablePartitioner::try_new(self.partition_manager, table_info) .await .context(SplitInsertSnafu)? - .into_iter() - .map( - |(region_number, (rows, partition_expr_version))| InsertRequest { - skip_wal, - region_id: RegionId::new(table_id, region_number).into(), - rows: Some(rows), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - }, - ) - .collect(); - Ok(requests) + .partition_insert_requests(rows, skip_wal) } + /// Routes delete rows using the current table partitions. pub async fn partition_delete_requests( &self, table_info: &TableInfo, rows: Rows, ) -> Result> { - let table_id = table_info.table_id(); - - let requests = self - .partition_manager - .split_rows(table_info, rows) + TablePartitioner::try_new(self.partition_manager, table_info) .await .context(SplitDeleteSnafu)? + .partition_delete_requests(rows) + } +} + +impl TablePartitioner { + /// Reads a table's partition rule and expression versions into one routing snapshot. + pub async fn try_new( + partition_manager: &PartitionRuleManager, + table_info: &TableInfo, + ) -> partition::error::Result { + let (rule, versions) = partition_manager + .find_table_partition_rule(table_info) + .await?; + Ok(Self { + table_id: table_info.table_id(), + rule, + versions: versions + .into_iter() + .map(|(region, version)| { + (region, version.map(|value| PartitionExprVersion { value })) + }) + .collect(), + }) + } + + /// Moves insert rows into requests for this snapshot's destination regions. + pub fn partition_insert_requests( + &self, + rows: Rows, + skip_wal: bool, + ) -> Result> { + Ok(RowSplitter::new(self.rule.clone()) + .split(rows) + .context(SplitInsertSnafu)? .into_iter() - .map( - |(region_number, (rows, partition_expr_version))| DeleteRequest { - region_id: RegionId::new(table_id, region_number).into(), - rows: Some(rows), - partition_expr_version: partition_expr_version - .map(|value| PartitionExprVersion { value }), - }, - ) - .collect(); + .map(|(region, rows)| InsertRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + rows: Some(rows), + partition_expr_version: self.version(region), + skip_wal, + }) + .collect()) + } + + /// Moves delete rows into requests for this snapshot's destination regions. + pub fn partition_delete_requests(&self, rows: Rows) -> Result> { + Ok(RowSplitter::new(self.rule.clone()) + .split(rows) + .context(SplitDeleteSnafu)? + .into_iter() + .map(|(region, rows)| DeleteRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + rows: Some(rows), + partition_expr_version: self.version(region), + }) + .collect()) + } + + /// Splits an Arrow batch, reusing the original IPC when one region selects all rows. + pub fn partition_bulk_insert_requests( + &self, + batch: RecordBatch, + mut original_ipc: Option, + skip_wal: bool, + ) -> Result> { + let masks = self + .rule + .split_record_batch(&batch) + .context(SplitInsertSnafu)?; + let mut requests = Vec::with_capacity(masks.len()); + for (region, mask) in masks { + if mask.select_none() { + continue; + } + let original = if mask.select_all() { + original_ipc.take() + } else { + None + }; + let ipc = if let Some(ipc) = original { + ipc + } else { + let selected = if mask.select_all() { + batch.clone() + } else { + compute::filter_record_batch(&batch, mask.array()) + .context(error::ComputeArrowSnafu)? + }; + let (schema, data_header, payload) = flight::record_batch_to_ipc(selected) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + ArrowIpc { + schema, + data_header, + payload, + } + }; + requests.push(BulkInsertRequest { + region_id: RegionId::new(self.table_id, region).as_u64(), + partition_expr_version: self.version(region), + // The destination datanode validates against its current schema. + aligned_schema_version: None, + skip_wal, + body: Some(bulk_insert_request::Body::ArrowIpc(ipc)), + }); + } Ok(requests) + } + + fn version(&self, region: RegionNumber) -> Option { + self.versions.get(®ion).copied().flatten() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + use common_meta::cache; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::store::RangeResponse; + use moka::future::Cache; + + use super::*; + use crate::error::Error; + use crate::test_util; + + #[tokio::test] + async fn test_partition_requests_preserve_routing_errors() { + let table_info = test_util::new_test_table_info(1, "table_1", [1].into_iter()); + for (backend_unavailable, expected_status, expected_hint) in [ + (false, StatusCode::TableNotFound, RetryHint::NonRetryable), + (true, StatusCode::Internal, RetryHint::Retryable), + ] { + let backend = Arc::new( + MockKvBackendBuilder::default() + .range_fn(Arc::new(move |_| { + if backend_unavailable { + common_meta::error::NoLeaderSnafu { table_id: 1u32 }.fail() + } else { + Ok(RangeResponse { + kvs: Vec::new(), + more: false, + }) + } + })) + .build() + .unwrap(), + ); + let routes = Arc::new(cache::new_table_route_cache( + "partitioner_test_routes".to_string(), + Cache::new(16), + backend.clone(), + )); + let partitions = Arc::new(partition::cache::new_partition_info_cache( + "partitioner_test_partitions".to_string(), + Cache::new(16), + routes.clone(), + )); + let manager = PartitionRuleManager::new(backend, routes, partitions); + let partitioner = Partitioner::new(&manager); + let insert_error = partitioner + .partition_insert_requests(&table_info, Rows::default(), false) + .await + .unwrap_err(); + let delete_error = partitioner + .partition_delete_requests(&table_info, Rows::default()) + .await + .unwrap_err(); + + assert!(matches!(insert_error, Error::SplitInsert { .. })); + assert!(matches!(delete_error, Error::SplitDelete { .. })); + for error in [insert_error, delete_error] { + assert_eq!(error.status_code(), expected_status); + assert_eq!(error.retry_hint(), expected_hint); + } + } } } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs new file mode 100644 index 00000000000..4af4fbfafdf --- /dev/null +++ b/src/operator/src/write_dispatcher.rs @@ -0,0 +1,256 @@ +// Copyright 2023 Greptime Team +// +// 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. + +#[cfg(test)] +mod tests; + +use std::collections::{HashMap, HashSet, hash_map::Entry}; +use std::sync::Arc; + +use api::region::RegionResponse; +use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body}; +use common_error::ext::BoxedError; +use common_meta::error::{ExternalSnafu, Result, UnexpectedSnafu}; +use common_meta::node_manager::NodeManagerRef; +use common_meta::peer::Peer; +use common_meta::rpc::router; +use futures::future; +use partition::manager::PartitionRuleManagerRef; +use prometheus::Histogram; +use snafu::{OptionExt, ResultExt}; +use store_api::storage::RegionId; + +use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; + +/// Dispatches already partitioned insert, delete, and bulk requests to their peers. +pub struct WriteDispatcher { + node_manager: NodeManagerRef, + partition_manager: PartitionRuleManagerRef, +} + +/// Shared reference to a [`WriteDispatcher`]. +pub type WriteDispatcherRef = Arc; + +impl WriteDispatcher { + /// Creates an owned-request dispatcher without retries, including for standalone. + pub fn new(node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef) -> Self { + Self { + node_manager, + partition_manager, + } + } + + /// Writes planned requests without repeating admission, default values, or Flow side effects. + pub async fn dispatch(&self, requests: Vec) -> Result { + self.dispatch_inner(requests, None).await + } + + /// Records each datanode send in the supplied histogram. + pub async fn dispatch_with_metrics( + &self, + requests: Vec, + histogram: Histogram, + ) -> Result { + self.dispatch_inner(requests, Some(histogram)).await + } + + /// Resolves routes before writes are detached from the caller. + pub(crate) async fn prepare( + &self, + requests: Vec, + ) -> error::Result> { + let requests = requests.into_iter().map(split_request).collect::>(); + let regions = requests + .iter() + .flatten() + .map(write_region_id) + .collect::>>() + .context(RequestRegionSnafu)?; + let peers = self + .resolve_peers(®ions) + .await + .context(FindRegionLeaderSnafu)?; + let mut writes = Vec::new(); + // Preserve each original envelope and peer batching when retries are disabled. + for requests in requests { + let mut by_peer: HashMap = HashMap::new(); + for request in requests { + let peer = get_peer( + &peers, + write_region_id(&request).context(RequestRegionSnafu)?, + ) + .context(RequestRegionSnafu)?; + match by_peer.entry(peer) { + Entry::Vacant(entry) => { + entry.insert(request); + } + Entry::Occupied(mut entry) => match (&mut entry.get_mut().body, request.body) { + (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) => { + existing.requests.extend(next.requests) + } + (Some(Body::Deletes(existing)), Some(Body::Deletes(next))) => { + existing.requests.extend(next.requests) + } + _ => { + return UnexpectedSnafu { + err_msg: "Cannot combine different write kinds", + } + .fail() + .context(RequestRegionSnafu); + } + }, + } + } + writes.extend(by_peer); + } + Ok(writes) + } + + /// Sends requests using routes already resolved by the caller. + pub(crate) async fn dispatch_prepared( + &self, + writes: Vec<(Peer, RegionRequest)>, + histogram: Option, + ) -> Result { + let histogram = histogram.as_ref(); + let results = future::join_all(writes.into_iter().map(|(peer, request)| async move { + let _timer = histogram.map(Histogram::start_timer); + let datanode = self.node_manager.datanode(&peer).await; + datanode.handle(request).await + })) + .await; + let mut response = RegionResponse::new(0); + for result in results { + merge_response(&mut response, result?); + } + Ok(response) + } + + /// Resolves only the requested peers, loading each table route once. + pub(crate) async fn resolve_peers( + &self, + regions: &[RegionId], + ) -> partition::error::Result> { + let table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_tables = HashMap::new(); + let mut routes = HashMap::new(); + for table_id in table_ids { + let (physical_id, route) = self + .partition_manager + .find_physical_table_route_with_id(table_id) + .await?; + physical_tables.insert(table_id, physical_id); + routes.entry(physical_id).or_insert(route); + } + let mut peers = HashMap::with_capacity(regions.len()); + for region in regions { + if peers.contains_key(region) { + continue; + } + let peer = physical_tables + .get(®ion.table_id()) + .and_then(|physical_id| routes.get(physical_id)) + .and_then(|route| { + router::find_region_leader(&route.region_routes, region.region_number()) + }) + .context(partition::error::FindLeaderSnafu { + region_id: *region, + table_id: region.table_id(), + })?; + peers.insert(*region, peer); + } + Ok(peers) + } + + async fn dispatch_inner( + &self, + requests: Vec, + histogram: Option, + ) -> Result { + let writes = self + .prepare(requests) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + self.dispatch_prepared(writes, histogram).await + } +} + +fn get_peer(peers: &HashMap, region: RegionId) -> Result { + peers.get(®ion).cloned().context(UnexpectedSnafu { + err_msg: "Planned write has no resolved peer", + }) +} + +fn merge_response(response: &mut RegionResponse, result: RegionResponse) { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); + response.metadata = result.metadata; +} + +fn split_request(request: RegionRequest) -> Vec { + let RegionRequest { header, body } = request; + let bodies: Vec = match body { + Some(Body::Inserts(inserts)) => inserts + .requests + .into_iter() + .map(|request| { + Body::Inserts(InsertRequests { + requests: vec![request], + }) + }) + .collect(), + Some(Body::Deletes(deletes)) => deletes + .requests + .into_iter() + .map(|request| { + Body::Deletes(DeleteRequests { + requests: vec![request], + }) + }) + .collect(), + body => return vec![RegionRequest { header, body }], + }; + bodies + .into_iter() + .map(|body| RegionRequest { + header: header.clone(), + body: Some(body), + }) + .collect() +} + +fn write_region_id(request: &RegionRequest) -> Result { + let region_id = match &request.body { + Some(Body::Inserts(inserts)) => match inserts.requests.as_slice() { + [request] => Some(request.region_id), + _ => None, + }, + Some(Body::Deletes(deletes)) => match deletes.requests.as_slice() { + [request] => Some(request.region_id), + _ => None, + }, + Some(Body::BulkInsert(bulk)) => Some(bulk.region_id), + _ => None, + }; + region_id.map(RegionId::from_u64).ok_or_else(|| { + UnexpectedSnafu { + err_msg: "Expected a single-region write request", + } + .build() + }) +} diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs new file mode 100644 index 00000000000..1191f73d55a --- /dev/null +++ b/src/operator/src/write_dispatcher/tests.rs @@ -0,0 +1,350 @@ +// Copyright 2023 Greptime Team +// +// 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 std::collections::HashMap; +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::task::Poll; +use std::time::{Duration, Instant}; + +use api::v1::region::{ + BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, bulk_insert_request, +}; +use api::v1::{ + ArrowIpc, ColumnDataType, PartitionExprVersion, Row, Rows, helper, value::ValueData, +}; +use client::client_manager::NodeClients; +use common_error::ext::ErrorExt; +use common_grpc::channel_manager::ChannelConfig; +use common_meta::cache as meta_cache; +use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; +use common_meta::key::TableMetadataManager; +use common_meta::key::table_route::TableRouteValue; +use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; +use common_query::request::QueryRequest; +use common_recordbatch::SendableRecordBatchStream; +use partition::manager::PartitionRuleManager; +use tokio::sync::{Barrier, Notify}; + +use super::*; +use crate::error::Error; +use crate::test_util; + +fn build_insert_request(regions: &[u32]) -> RegionRequest { + RegionRequest { + body: Some(Body::Inserts(InsertRequests { + requests: regions + .iter() + .map(|region| InsertRequest { + region_id: RegionId::new(1, *region).as_u64(), + ..Default::default() + }) + .collect(), + })), + ..Default::default() + } +} + +#[tokio::test] +async fn test_resolve_peers_batches_logical_and_physical_regions() { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend) + .create_table_metadata( + test_util::new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()); + let regions = [ + RegionId::new(1, 1), + RegionId::new(1, 2), + RegionId::new(42, 1), + RegionId::new(42, 2), + RegionId::new(42, 1), + ]; + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + assert_eq!(peers.len(), 4); + for region in regions { + assert_eq!( + peers[®ion], + partition_manager.find_region_leader(region).await.unwrap() + ); + } + assert_eq!(peers[&RegionId::new(1, 1)], peers[&RegionId::new(42, 1)]); + let missing = RegionId::new(42, 999); + let expected = partition_manager + .find_region_leader(missing) + .await + .unwrap_err(); + let actual = dispatcher.resolve_peers(&[missing]).await.unwrap_err(); + assert_eq!(actual.status_code(), expected.status_code()); + assert_eq!(actual.retry_hint(), expected.retry_hint()); + assert!(dispatcher.resolve_peers(&[]).await.unwrap().is_empty()); +} + +#[tokio::test] +async fn test_dispatch_preserves_write_payloads_and_envelopes() { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend) + .create_table_metadata( + test_util::new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let (sender, mut received) = tokio::sync::mpsc::channel(8); + let watcher = DatanodeWatcher::new(sender).with_handler(|_, _| Ok(RegionResponse::new(1))); + let dispatcher = WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(watcher)), + partition_manager, + ); + // Physical and logical regions share a peer but retain distinct region IDs. + let regions = [RegionId::new(1, 1), RegionId::new(42, 1)]; + let rows = Rows { + schema: vec![helper::tag_column_schema("a", ColumnDataType::Int32)], + rows: vec![Row { + values: vec![ValueData::I32Value(7).into()], + }], + }; + let version = Some(PartitionExprVersion { value: 17 }); + let bodies = [ + Body::Inserts(InsertRequests { + requests: regions + .iter() + .map(|region| InsertRequest { + region_id: region.as_u64(), + rows: Some(rows.clone()), + partition_expr_version: version, + skip_wal: true, + }) + .collect(), + }), + Body::Deletes(DeleteRequests { + requests: regions + .iter() + .map(|region| DeleteRequest { + region_id: region.as_u64(), + rows: Some(rows.clone()), + partition_expr_version: version, + }) + .collect(), + }), + ]; + for body in bodies { + let requests = ["first", "second"].map(|dbname| RegionRequest { + header: Some(RegionRequestHeader { + dbname: dbname.to_string(), + ..Default::default() + }), + body: Some(body.clone()), + }); + let writes = dispatcher.prepare(requests.to_vec()).await.unwrap(); + assert!(received.try_recv().is_err()); + let response = dispatcher.dispatch_prepared(writes, None).await.unwrap(); + assert_eq!(response.affected_rows, 2); + for request in requests { + assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); + } + assert!(received.try_recv().is_err()); + } + let requests = regions.map(|region| RegionRequest { + header: Some(RegionRequestHeader { + dbname: "bulk".to_string(), + ..Default::default() + }), + body: Some(Body::BulkInsert(BulkInsertRequest { + region_id: region.as_u64(), + partition_expr_version: version, + aligned_schema_version: None, + skip_wal: true, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: "schema".into(), + data_header: "header".into(), + payload: "payload".into(), + })), + })), + }); + let response = dispatcher.dispatch(requests.to_vec()).await.unwrap(); + assert_eq!(response.affected_rows, 2); + for request in requests { + assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); + } + assert!(received.try_recv().is_err()); + + let error = dispatcher + .prepare(vec![build_insert_request(&[1, 999])]) + .await + .unwrap_err(); + assert!(matches!(error, Error::FindRegionLeader { .. })); + assert!(received.try_recv().is_err()); +} + +#[derive(Clone)] +struct GatedDatanodeHandler { + entered: Arc, + rejected: Arc, + completed: Arc, + release: Arc, +} + +#[async_trait::async_trait] +impl MockDatanodeHandler for GatedDatanodeHandler { + async fn handle(&self, peer: &Peer, _request: RegionRequest) -> Result { + self.entered.wait().await; + if peer.id == 1 { + self.rejected.store(true, Ordering::SeqCst); + return UnexpectedSnafu { + err_msg: "first region rejected", + } + .fail(); + } + self.release.notified().await; + self.completed.store(true, Ordering::SeqCst); + Ok(RegionResponse::new(1)) + } + + async fn handle_query( + &self, + _peer: &Peer, + _request: QueryRequest, + ) -> Result { + unimplemented!() + } +} + +#[tokio::test] +async fn test_dispatch_sends_concurrently_and_waits_for_siblings_after_error() { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend).await; + let handler = GatedDatanodeHandler { + entered: Arc::new(Barrier::new(2)), + rejected: Arc::default(), + completed: Arc::default(), + release: Arc::default(), + }; + let dispatcher = WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(handler.clone())), + partition_manager, + ); + let dispatch = dispatcher.dispatch(vec![build_insert_request(&[1, 2])]); + tokio::pin!(dispatch); + tokio::time::timeout( + Duration::from_secs(1), + futures::future::poll_fn(|cx| { + assert!( + dispatch.as_mut().poll(cx).is_pending(), + "dispatch returned while a sibling remained blocked" + ); + if handler.rejected.load(Ordering::SeqCst) { + Poll::Ready(()) + } else { + Poll::Pending + } + }), + ) + .await + .expect("both region sends must reach the barrier concurrently"); + assert!(!handler.completed.load(Ordering::SeqCst)); + handler.release.notify_one(); + let error = tokio::time::timeout(Duration::from_secs(1), dispatch) + .await + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("first region rejected")); + assert!(handler.completed.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn test_resolve_peers_cache_lookup_ablation() { + common_telemetry::init_default_ut_logging(); + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + test_util::new_test_table_info(42, "resolution_ablation", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + // A dedicated cache label isolates counters from other concurrently running tests. + let cache_name = "write_dispatcher_resolution_ablation_table_route"; + let routes = Arc::new(meta_cache::new_table_route_cache( + cache_name.to_string(), + moka::future::Cache::new(16), + backend.clone(), + )); + let partitions = Arc::new(partition::cache::new_partition_info_cache( + "write_dispatcher_resolution_ablation_partitions".to_string(), + moka::future::Cache::new(16), + routes.clone(), + )); + let partition_manager = Arc::new(PartitionRuleManager::new(backend, routes, partitions)); + let dispatcher = + WriteDispatcher::new(Arc::new(NodeClients::default()), partition_manager.clone()); + let regions = (0..256) + .map(|index| RegionId::new(if index % 2 == 0 { 1 } else { 42 }, index % 3 + 1)) + .collect::>(); + let expected = dispatcher.resolve_peers(®ions).await.unwrap(); + let counter = common_meta::metrics::CACHE_CONTAINER_CACHE_GET.with_label_values(&[cache_name]); + for run in 0..6 { + // Alternate order; compare the same warm cache, routes and deterministic workload. + let modes = if run % 2 == 0 { + ["delete", "insert", "batched"] + } else { + ["batched", "insert", "delete"] + }; + for mode in modes { + let before = counter.get(); + let started = Instant::now(); + let peers = if mode == "batched" { + dispatcher.resolve_peers(®ions).await.unwrap() + } else { + let mut peers = HashMap::new(); + for region in ®ions { + if mode == "insert" && peers.contains_key(region) { + continue; + } + peers.insert( + *region, + partition_manager.find_region_leader(*region).await.unwrap(), + ); + } + peers + }; + let elapsed = started.elapsed(); + let gets = counter.get() - before; + assert_eq!(peers, expected); + let expected_gets = match mode { + "delete" => 384, + "insert" => 9, + "batched" => 3, + _ => unreachable!(), + }; + assert_eq!(gets, expected_gets); + common_telemetry::info!( + "route_resolution_ablation run={run} mode={mode} regions={} cache_gets={gets} elapsed_ns={}", + regions.len(), + elapsed.as_nanos() + ); + } + } +} From 83e62a772590cce7dd9a854f9a11e01d9b413050 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 13 Oct 2026 02:32:30 +0000 Subject: [PATCH 22/46] feat: remove the server-side Metric export gate (#9547) * feat: remove the server-side Metric export gate Signed-off-by: jeremyhi * chore: remove unused export capability import Signed-off-by: jeremyhi * chore: remove unused CLI capability import Signed-off-by: jeremyhi * style: import the packed export error selector Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- config/config.md | 2 - config/frontend.example.toml | 4 - config/standalone.example.toml | 4 - docs/rfcs/2026-09-11-metric-export-import.md | 26 ++-- src/cli/src/data/export_v2.rs | 3 +- src/cli/src/data/export_v2/command.rs | 36 ++---- src/cli/src/data/export_v2/error.rs | 8 +- src/frontend/src/frontend.rs | 6 - src/frontend/src/instance.rs | 14 +-- src/frontend/src/instance/builder.rs | 1 - src/frontend/src/instance/export_database.rs | 31 ----- src/standalone/src/options.rs | 7 -- tests-integration/src/standalone.rs | 10 -- .../tests/export_logical_tables.rs | 112 ++++-------------- tests-integration/tests/http.rs | 1 - 15 files changed, 53 insertions(+), 212 deletions(-) diff --git a/config/config.md b/config/config.md index 5d7555ad4b0..5d76c8f7445 100644 --- a/config/config.md +++ b/config/config.md @@ -15,7 +15,6 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | | `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL:
`static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | 0 | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | `wait` | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | @@ -285,7 +284,6 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | | `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL:
`static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | 0 | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | `wait` | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index e4ef6e36935..9a1c3db63fa 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -11,10 +11,6 @@ default_column_prefix = "greptime" ## When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. #+ auto_create_table = true -## Enables experimental Parquet COPY DATABASE using shared Metric physical scans. -## Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. -#+ experimental_metric_export = false - ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 0da3ad9031c..82fb3d2c489 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -11,10 +11,6 @@ default_column_prefix = "greptime" ## When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. #+ auto_create_table = true -## Enables experimental Parquet COPY DATABASE using shared Metric physical scans. -## Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. -#+ experimental_metric_export = false - ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", diff --git a/docs/rfcs/2026-09-11-metric-export-import.md b/docs/rfcs/2026-09-11-metric-export-import.md index bf66df932b9..4b504065f6c 100644 --- a/docs/rfcs/2026-09-11-metric-export-import.md +++ b/docs/rfcs/2026-09-11-metric-export-import.md @@ -195,15 +195,18 @@ chunks, including ordinary tables stored as standalone Parquet objects. `export-v2 create --pack-metric-data` selects shared physical scans and packed output. Omitting the flag keeps the default version-1 export path; -schema-only snapshots also retain version 1. The source frontend or standalone -server must enable `experimental_metric_export = true`. Resume derives the layout +schema-only snapshots also retain version 1. Default frontend and standalone +configurations accept explicit packed export requests. Resume derives the layout from the manifest, so packed snapshots do not need the flag repeated. -Before creating or replacing a snapshot, the CLI checks the source server's -explicit packed-export capability. Before target DDL or data writes, the importer +Before creating or replacing a packed data snapshot, the CLI checks the source +server's explicit packed-export capability. Schema-only exports do not require +packed-data capability. Before target DDL or data writes, the importer checks the manifest version/layout and the target's packed-import capability. An unsupported packed layout fails; it must not fall through to listing `.parquet` files and silently omit packed tables. Batch-DDL capability is checked separately. +General batch-DDL and schema metadata optimizations apply independently of packed +export selection. | Reader and snapshot | Behavior | | --- | --- | @@ -438,8 +441,8 @@ The CLI retains one `COPY DATABASE FROM` operation per `(chunk, schema)` task. Add an explicit `metric_data_layout = 'packed'` COPY option, paired with `FORMAT = 'parquet'`, for the advertised capability. The server validates the index and selected destination tables, then schedules table streams by object -and offset. Use the same explicit layout option for packed export. The existing -Metric feature gate continues to control the experimental export path. +and offset. Use the same explicit layout option for packed export, with +per-table authorization and storage access checks. A reader owned by that request translates each Parquet-relative read into a bounded object range. Start with two shared 8 MiB read windows per request; @@ -518,16 +521,15 @@ the existing durable state update. - Pin supported old-reader binaries for version-2 rejection and version-1 compatibility tests. The first-release storage targets are local files and S3-compatible storage exercised through MinIO. -- Decide when release acceptance permits removing the existing experimental - gate. Packed export/restore capability is required while mixed versions exist. +- Packed export/restore capability is required while mixed versions exist. # Integration and release acceptance Use existing COPY DATABASE orchestration to form Metric export units. Retain ordinary COPY for non-Metric tables and other formats, and direct COPY TABLE -behavior. Gate the unfinished Metric fast path behind an `experimental_` -configuration option; disabling it selects the existing path. Add packed-import -support before enabling packed snapshot creation. The writer and reader form one +behavior. Explicit packed requests select the Metric fast path; ordinary exports +keep the existing path. Packed snapshot creation requires packed-export capability, +and restore requires packed-import capability. The writer and reader form one first-release feature, even though they land in separate PRs. The [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120) owns @@ -540,7 +542,7 @@ must verify its own correctness and failure behavior. Release acceptance require | Restore | Authenticated remote batch DDL, per-table authorization, fresh target IDs and reconstructed metadata, dependency ordering, older-server behavior and batch failure/retry. | | V2 lifecycle | New-reader version-1/2 coverage, old-reader version-2 rejection, target capability preflight; local/object-store interruption, multipart abort, owned-output cleanup, index/manifest failure and import replay. | | Scale | Repeated release-build measurements varying table count up to approximately 100,000, rows/file, physical union width, regions and backend. Report median/spread, control execution order/cache, and include schema export and metadata discovery in end-to-end timing. | -| Open decisions | Agreed batch-DDL interface, resource limits, compatibility targets and experimental gate. | +| Open decisions | Agreed batch-DDL interface, resource limits, compatibility targets. | ## Snapshot format acceptance diff --git a/src/cli/src/data/export_v2.rs b/src/cli/src/data/export_v2.rs index 8785bcf3e31..08652d28c6d 100644 --- a/src/cli/src/data/export_v2.rs +++ b/src/cli/src/data/export_v2.rs @@ -40,8 +40,7 @@ //! //! `--pack-metric-data` enables shared Metric physical scans for Parquet //! and writes version-2 packed snapshots. -//! Enable `experimental_metric_export = true` on the frontend or standalone server -//! and use an endpoint whose frontends all support and enable this option. +//! Use an endpoint whose frontends all support packed export. //! Resume reads the layout from the existing manifest. //! This requires the server packed-export capability; schema-only exports remain version 1. //! The command checks the server capability before modifying the snapshot. diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index fe38c6b7c7e..ee10091e5ae 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -20,11 +20,9 @@ use std::time::Duration; use async_trait::async_trait; use clap::{Parser, Subcommand}; -use common_catalog::consts::DEFAULT_SCHEMA_NAME; use common_error::ext::BoxedError; use common_telemetry::info; use serde_json::Value; -use servers::http::{ColumnSchema, GreptimeQueryOutput, OutputSchema}; use snafu::ResultExt; use crate::Tool; @@ -32,8 +30,9 @@ use crate::common::ObjectStoreConfig; use crate::data::export_v2::coordinator::{ExportDataOptions, export_data}; use crate::data::export_v2::error::{ ChunkTimeWindowRequiresBoundsSnafu, DatabaseSnafu, IoSnafu, ManifestVersionMismatchSnafu, - Result, ResumeConfigMismatchSnafu, SchemaOnlyArgsNotAllowedSnafu, SchemaOnlyModeMismatchSnafu, - SnapshotVerifyFailedSnafu, UnexpectedValueTypeSnafu, + MetricExportRequiresParquetSnafu, Result, ResumeConfigMismatchSnafu, + SchemaOnlyArgsNotAllowedSnafu, SchemaOnlyModeMismatchSnafu, SnapshotVerifyFailedSnafu, + UnexpectedValueTypeSnafu, }; use crate::data::export_v2::extractor::SchemaExtractor; use crate::data::export_v2::manifest::{ @@ -408,9 +407,7 @@ impl ExportCreateCommand { // The filesystem storage constructor can create the snapshot root. if self.pack_metric_data { - require_metric_export_enabled(&database_client, self.format) - .await - .map_err(BoxedError::new)?; + require_metric_export_format(self.format).map_err(BoxedError::new)?; if !self.schema_only { database_client .require_packed_export() @@ -508,7 +505,7 @@ impl ExportCreate { validate_resume_config(&manifest, &self.config)?; if manifest.is_packed() { - require_metric_export_enabled(&self.database_client, manifest.format).await?; + require_metric_export_format(manifest.format)?; self.database_client .require_packed_export() .await @@ -727,26 +724,9 @@ impl ExportCreate { } } -async fn require_metric_export_enabled(client: &DatabaseClient, format: DataFormat) -> Result<()> { +fn require_metric_export_format(format: DataFormat) -> Result<()> { if format != DataFormat::Parquet { - return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail(); - } - let capability = client - .sql_response( - "SHOW VARIABLES experimental_metric_export", - DEFAULT_SCHEMA_NAME, - ) - .await - .context(DatabaseSnafu)?; - let expected_schema = OutputSchema::new(vec![ColumnSchema::new( - "EXPERIMENTAL_METRIC_EXPORT".to_string(), - "String".to_string(), - )]); - if !matches!(capability.output(), [GreptimeQueryOutput::Records(records)] - if records.schema() == &expected_schema - && records.rows() == &vec![vec![Value::String("true".to_string())]]) - { - return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail(); + return MetricExportRequiresParquetSnafu.fail(); } Ok(()) } @@ -1756,7 +1736,7 @@ mod tests { #[test] fn test_schema_only_mode_mismatch_error_message() { - let error = crate::data::export_v2::error::SchemaOnlyModeMismatchSnafu { + let error = SchemaOnlyModeMismatchSnafu { existing_schema_only: false, requested_schema_only: true, } diff --git a/src/cli/src/data/export_v2/error.rs b/src/cli/src/data/export_v2/error.rs index dbb4a88ae46..a64ff93865b 100644 --- a/src/cli/src/data/export_v2/error.rs +++ b/src/cli/src/data/export_v2/error.rs @@ -24,10 +24,8 @@ use snafu::{Location, Snafu}; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { - #[snafu(display( - "Experimental Metric export requires Parquet and a frontend that explicitly reports experimental_metric_export=true. Use an upgraded, consistently configured frontend entry point." - ))] - MetricExportUnavailable { + #[snafu(display("Packed Metric export requires Parquet"))] + MetricExportRequiresParquet { #[snafu(implicit)] location: Location, }, @@ -219,7 +217,7 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { - Error::MetricExportUnavailable { .. } + Error::MetricExportRequiresParquet { .. } | Error::InvalidUri { .. } | Error::UnsupportedScheme { .. } | Error::SchemaOnlyModeMismatch { .. } diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 90e7abc4517..e594ba9f7b1 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -50,8 +50,6 @@ pub struct FrontendOptions { /// even if a request sets the `auto_create_table` hint to `true`. When `true` /// (default), the per-request hint still applies. Default: `true`. pub auto_create_table: bool, - /// Enables experimental Parquet database exports using shared Metric scans. - pub experimental_metric_export: bool, /// Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). /// Set to 0 to disable the limit. Default: "0" (unlimited) pub max_in_flight_write_bytes: ReadableSize, @@ -96,7 +94,6 @@ impl Default for FrontendOptions { default_timezone: None, default_column_prefix: None, auto_create_table: true, - experimental_metric_export: false, max_in_flight_write_bytes: ReadableSize(0), write_bytes_exhausted_policy: OnExhaustedPolicy::default(), http: HttpOptions::default(), @@ -311,9 +308,6 @@ max_batch_rows = 25 #[test] fn test_toml() { let opts = FrontendOptions::default(); - assert!(!opts.experimental_metric_export); - let enabled: FrontendOptions = toml::from_str("experimental_metric_export = true").unwrap(); - assert!(enabled.experimental_metric_export); let toml_string = toml::to_string(&opts).unwrap(); let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); assert_eq!(parsed.otlp, opts.otlp); diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index df30106c54a..59559950aef 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -127,7 +127,6 @@ lazy_static! { #[derive(Clone)] pub struct Instance { frontend_peer_addr: String, - experimental_metric_export: bool, catalog_manager: CatalogManagerRef, pipeline_operator: Arc, statement_executor: Arc, @@ -409,14 +408,6 @@ impl Instance { .copy_packed_database(arg.clone(), &stmt, query_ctx) .await; } - if let Statement::ShowVariables(show) = &stmt - && show - .variable - .to_string() - .eq_ignore_ascii_case("experimental_metric_export") - { - return self.show_metric_export_capability(); - } if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase(CopyDatabase::To( arg, ))) = &stmt @@ -2921,10 +2912,7 @@ mod tests { let catalog = catalog::memory::MemoryCatalogManager::new_with_table(source); let kv = Arc::new(MemoryKvBackend::new()); let instance = FrontendBuilder::new( - FrontendOptions { - experimental_metric_export: true, - ..Default::default() - }, + FrontendOptions::default(), kv.clone(), test_cache_registry(kv).unwrap(), catalog, diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 032722499f9..9db7b70c221 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -399,7 +399,6 @@ impl FrontendBuilder { Ok(Instance { logical_batcher: Default::default(), frontend_peer_addr, - experimental_metric_export: self.options.experimental_metric_export, catalog_manager: self.catalog_manager, pipeline_operator, statement_executor, diff --git a/src/frontend/src/instance/export_database.rs b/src/frontend/src/instance/export_database.rs index f36bcf562c7..022a50482d1 100644 --- a/src/frontend/src/instance/export_database.rs +++ b/src/frontend/src/instance/export_database.rs @@ -12,19 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; - use auth::{ PermissionAction, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget, PermissionTableTargets, }; -use common_error::ext::BoxedError; use common_query::Output; -use common_recordbatch::RecordBatches; use common_telemetry::error; -use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{ColumnSchema, Schema}; -use datatypes::vectors::StringVector; use operator::statement::export_database::{DatabaseExportSummary, PreparedDatabaseExport}; use session::context::QueryContextRef; use snafu::ResultExt; @@ -51,35 +44,11 @@ pub(crate) fn parse_metric_export_requested(options: &OptionMap) -> Result } impl Instance { - pub(crate) fn show_metric_export_capability(&self) -> Result { - let schema = Arc::new(Schema::new(vec![ColumnSchema::new( - "EXPERIMENTAL_METRIC_EXPORT", - ConcreteDataType::string_datatype(), - false, - )])); - let batches = RecordBatches::try_from_columns( - schema, - vec![Arc::new(StringVector::from(vec![ - self.experimental_metric_export.to_string(), - ])) as datatypes::vectors::VectorRef], - ) - .map_err(BoxedError::new) - .context(crate::error::ExternalSnafu)?; - Ok(Output::new_with_record_batches(batches)) - } - pub(crate) async fn copy_metric_database( &self, arg: CopyDatabaseArgument, ctx: QueryContextRef, ) -> Result { - if !self.experimental_metric_export { - return Err(operator::error::InvalidDatabaseExportSnafu { - reason: "experimental_metric_export is disabled on this frontend", - } - .build() - .into()); - } let req = operator::statement::to_copy_database_request(arg, &ctx)?; let plan = self.prepare_database_export(req, None, &ctx).await?; let cancellation = CancellationToken::new(); diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index fbfcc2ccc1d..079cae4bca6 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -44,8 +44,6 @@ pub struct StandaloneOptions { /// Upper bound: when `false`, missing tables are never auto-created even if a /// request sets the `auto_create_table` hint to `true`. Default: `true`. pub auto_create_table: bool, - /// Enables experimental Parquet database exports using shared Metric scans. - pub experimental_metric_export: bool, /// Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). /// Set to 0 to disable the limit. Default: "0" (unlimited) pub max_in_flight_write_bytes: ReadableSize, @@ -93,7 +91,6 @@ impl Default for StandaloneOptions { default_timezone: None, default_column_prefix: None, auto_create_table: true, - experimental_metric_export: false, max_in_flight_write_bytes: ReadableSize(0), write_bytes_exhausted_policy: OnExhaustedPolicy::default(), http: HttpOptions::default(), @@ -158,7 +155,6 @@ impl StandaloneOptions { FrontendOptions { default_timezone: cloned_opts.default_timezone, auto_create_table: cloned_opts.auto_create_table, - experimental_metric_export: cloned_opts.experimental_metric_export, max_in_flight_write_bytes: cloned_opts.max_in_flight_write_bytes, write_bytes_exhausted_policy: cloned_opts.write_bytes_exhausted_policy, http: cloned_opts.http, @@ -368,9 +364,6 @@ flow_notification_queue_capacity = 17 #[test] fn test_query_options_propagated_to_components() { let mut options = StandaloneOptions::default(); - assert!(!options.frontend_options().experimental_metric_export); - options.experimental_metric_export = true; - assert!(options.frontend_options().experimental_metric_export); options.query.parallelism = 4; assert_eq!(options.frontend_options().query.parallelism, 4); diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 241d4daf5b5..074be282543 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -88,7 +88,6 @@ pub struct GreptimeDbStandaloneBuilder { slow_query_options: SlowQueryOptions, event_recorder_options: EventRecorderOptions, auto_create_table: bool, - experimental_metric_export: bool, logical_batcher: Option, table_batcher: BatcherOptions, mito_config: Option, @@ -110,7 +109,6 @@ impl GreptimeDbStandaloneBuilder { }, event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, - experimental_metric_export: false, logical_batcher: None, table_batcher: BatcherOptions::default(), mito_config: None, @@ -124,13 +122,6 @@ impl GreptimeDbStandaloneBuilder { self } - /// Enables experimental Metric export for the standalone test instance. - #[must_use] - pub fn with_experimental_metric_export(mut self) -> Self { - self.experimental_metric_export = true; - self - } - /// Configures ordinary-table batching for protocol integration tests. #[must_use] pub fn with_table_batcher(mut self, options: BatcherOptions) -> Self { @@ -398,7 +389,6 @@ impl GreptimeDbStandaloneBuilder { slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, - experimental_metric_export: self.experimental_metric_export, pending_rows_batcher: PendingRowsBatcherOptions { logical_table: self.logical_batcher.clone(), table: self.table_batcher.clone(), diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 101b2b6f476..eec32cb6b2d 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1587,7 +1587,6 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { }); plugins.insert::>(faults.clone()); let source = GreptimeDbStandaloneBuilder::new("metric_v2_source") - .with_experimental_metric_export() .with_plugin(plugins) .build() .await; @@ -2094,8 +2093,8 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { } #[tokio::test] -async fn metric_export_v2_disabled_and_legacy_cli() { - let source = GreptimeDbStandaloneBuilder::new("metric_v2_disabled") +async fn metric_export_v2_requires_parquet_and_keeps_default_layout() { + let source = GreptimeDbStandaloneBuilder::new("metric_v2_default") .build() .await; let (addr, server) = export_http(source.fe_instance().clone()).await; @@ -2113,12 +2112,15 @@ async fn metric_export_v2_disabled_and_legacy_cli() { "--to", &uri, "--force", + "--format", + "csv", "--pack-metric-data", "--no-proxy", "--progress", "never", ]; - assert!(run_data_cli(&args).await.is_err()); + let error = run_data_cli(&args).await.unwrap_err(); + assert!(error.to_string().contains("requires Parquet"), "{error}"); assert_eq!( std::fs::read(&manifest).unwrap(), b"preserve-before-capability-check" @@ -2148,6 +2150,10 @@ async fn metric_export_v2_disabled_and_legacy_cli() { run_data_cli(&["export-v2", "verify", "--snapshot", &uri]) .await .unwrap(); + let saved: cli::export_v2::manifest::Manifest = + serde_json::from_slice(&std::fs::read(&manifest).unwrap()).unwrap(); + assert_eq!(saved.version, 1); + assert!(!saved.is_packed()); server.abort(); } @@ -2156,80 +2162,18 @@ async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() use axum::http::StatusCode; use common_error::ext::ErrorExt; use serde_json::json; - let column = json!({"name": "EXPERIMENTAL_METRIC_EXPORT", "data_type": "String"}); - let records = |rows, columns| { - json!({"records": { - "schema": {"column_schemas": columns}, - "rows": rows - }}) - }; - let valid = records(json!([["true"]]), json!([column.clone()])); - let mut responses = [ - (StatusCode::BAD_REQUEST, json!([])), - (StatusCode::OK, json!([])), - (StatusCode::OK, json!([["false"]])), - (StatusCode::OK, json!([[true]])), - (StatusCode::OK, json!([["true", "extra"]])), - (StatusCode::OK, json!([["true"], ["true"]])), - ] - .map(|(status, rows)| (status, json!([records(rows, json!([column.clone()]))]))) - .to_vec(); - responses.extend([ - ( - StatusCode::OK, - json!([ - valid.clone(), - records(json!([["false"]]), json!([column.clone()])) - ]), - ), - ( - StatusCode::OK, - json!([records(json!([["true"]]), json!([]))]), - ), - ( - StatusCode::OK, - json!([records(json!([["true"]]), json!([column.clone(), column]))]), - ), - ( - StatusCode::OK, - json!([records( - json!([["true"]]), - json!([{"name": "EXPERIMENTAL_METRIC_EXPORT", "data_type": "Boolean"}]) - )]), - ), - ( - StatusCode::OK, - json!([records( - json!([["true"]]), - json!([{"name": "OTHER", "data_type": "String"}]) - )]), - ), - ]); - responses.push((StatusCode::OK, json!([valid.clone()]))); - for (status, output) in responses { - let gate_enabled = output == json!([valid.clone()]); - let app = axum::Router::new() - .route( - "/v1/sql", - axum::routing::post( - move |axum::Form(form): axum::Form< - std::collections::HashMap, - >| async move { - assert_eq!(form["sql"], "SHOW VARIABLES experimental_metric_export"); - ( - status, - axum::Json(json!({ - "execution_time_ms": 0, - "output": output - })), - ) - }, - ), - ) - .route( - "/v1/capabilities", - axum::routing::get(|| async { axum::Json(json!({})) }), - ); + for (status, capability) in [ + (StatusCode::NOT_FOUND, json!({})), + (StatusCode::OK, json!({})), + (StatusCode::OK, json!({"metric_packed_export": 0})), + (StatusCode::OK, json!({"metric_packed_export": 2})), + (StatusCode::OK, json!({"metric_packed_export": true})), + (StatusCode::OK, json!({"metric_packed_export": "1"})), + ] { + let app = axum::Router::new().route( + "/v1/capabilities", + axum::routing::get(move || async move { (status, axum::Json(capability)) }), + ); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap().to_string(); let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); @@ -2251,14 +2195,10 @@ async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() ]; let error = run_data_cli(&args).await.unwrap_err(); let cause = error.root_cause().unwrap_or(&error).to_string(); - if gate_enabled { - assert!(cause.contains("metric_packed_export"), "{error:?}"); - } else if status == StatusCode::OK { - assert!( - error.to_string().contains("Metric export requires"), - "{error}" - ); - } + assert!( + cause.contains("metric_packed_export") || cause.contains("invalid capability version"), + "{error:?}" + ); assert!(!snapshot.exists()); std::fs::create_dir(&snapshot).unwrap(); let manifest = snapshot.join("manifest.json"); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index dd2fc51a9b9..f3b439e19bc 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3283,7 +3283,6 @@ providers = []"#, r#" enable_telemetry = true auto_create_table = true -experimental_metric_export = false max_in_flight_write_bytes = "0KiB" write_bytes_exhausted_policy = "wait" init_regions_in_background = false From 287ced2de21b8e986f2becad59169132770f3ed0 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 03:18:37 +0000 Subject: [PATCH 23/46] fix(mito2): ignore missing inverted indexes in old SSTs (#9545) Signed-off-by: WenyXu --- .../src/sst/index/inverted_index/applier.rs | 5 +- .../src/sst/index/inverted_index/creator.rs | 68 +++++++- .../alter/change_col_inverted_index.result | 148 ++++++++++++++++++ .../alter/change_col_inverted_index.sql | 45 ++++++ 4 files changed, 258 insertions(+), 8 deletions(-) diff --git a/src/mito2/src/sst/index/inverted_index/applier.rs b/src/mito2/src/sst/index/inverted_index/applier.rs index 0f665eec549..8f0c309a003 100644 --- a/src/mito2/src/sst/index/inverted_index/applier.rs +++ b/src/mito2/src/sst/index/inverted_index/applier.rs @@ -221,8 +221,9 @@ impl InvertedIndexApplier { let start = Instant::now(); let context = SearchContext { - // Encountering a non-existing column indicates that it doesn't match predicates. - index_not_found_strategy: IndexNotFoundStrategy::ReturnEmpty, + // Older SSTs may lack an index enabled by ALTER, or a column with a default value. + // Missing indexes cannot rule out matching rows. + index_not_found_strategy: IndexNotFoundStrategy::Ignore, }; let mut cache_miss = 0; diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index d1dee226515..9c1ba5cbca0 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -556,6 +556,7 @@ mod tests { async fn build_applier_factory( prefix: &str, rows: BTreeSet<(&'static str, i32, [u64; 2])>, + sst_indexed_column_ids: HashSet, ) -> impl Fn(DfExpr) -> BoxFuture<'static, Vec> { let (d, factory) = PuffinManagerFactory::new_for_test_async(prefix).await; let table_dir = "table0".to_string(); @@ -573,7 +574,7 @@ mod tests { intm_mgr, memory_threshold, NonZeroUsize::new(segment_row_count).unwrap(), - indexed_column_ids.clone(), + sst_indexed_column_ids, ); for (str_tag, i32_tag, u64_field) in &rows { @@ -624,6 +625,42 @@ mod tests { } } + #[tokio::test] + async fn test_create_and_query_missing_index() { + let rows = BTreeSet::from_iter([("aaa", 1, [1, 2]), ("bbb", 2, [3, 4])]); + // The SST predates enabling indexes on tag_i32 and field_u64. + let applier_factory = build_applier_factory( + "test_create_and_query_missing_index_", + rows, + HashSet::from([1]), + ) + .await; + + let cases = [ + (col("tag_i32").eq(lit(1)), vec![0, 1]), + ( + col("tag_i32") + .eq(lit(1)) + .and(col("field_u64").eq(lit(1u64))), + vec![0, 1], + ), + ( + col("tag_str").eq(lit("aaa")).and(col("tag_i32").eq(lit(1))), + vec![0], + ), + ( + col("tag_str") + .eq(lit("missing")) + .and(col("tag_i32").eq(lit(1))), + vec![], + ), + ]; + for (expr, expected) in cases { + let actual = applier_factory(expr.clone()).await; + assert_eq!(actual, expected, "{expr}"); + } + } + #[tokio::test] async fn test_create_and_query_get_key() { let rows = BTreeSet::from_iter([ @@ -638,7 +675,12 @@ mod tests { ("abc", 3, [9, 10]), ]); - let applier_factory = build_applier_factory("test_create_and_query_get_key_", rows).await; + let applier_factory = build_applier_factory( + "test_create_and_query_get_key_", + rows, + HashSet::from([1, 2, 4]), + ) + .await; let expr = col("tag_str").eq(lit("aaa")); let res = applier_factory(expr).await; @@ -681,7 +723,12 @@ mod tests { ("abc", 3, [9, 10]), ]); - let applier_factory = build_applier_factory("test_create_and_query_range_", rows).await; + let applier_factory = build_applier_factory( + "test_create_and_query_range_", + rows, + HashSet::from([1, 2, 4]), + ) + .await; let expr = col("tag_str").between(lit("aaa"), lit("aab")); let res = applier_factory(expr).await; @@ -718,8 +765,12 @@ mod tests { ("abc", 3, [9, 10]), ]); - let applier_factory = - build_applier_factory("test_create_and_query_comparison_", rows).await; + let applier_factory = build_applier_factory( + "test_create_and_query_comparison_", + rows, + HashSet::from([1, 2, 4]), + ) + .await; let expr = col("tag_str").lt(lit("aab")); let res = applier_factory(expr).await; @@ -800,7 +851,12 @@ mod tests { ("abc", 3, [9, 10]), ]); - let applier_factory = build_applier_factory("test_create_and_query_regex_", rows).await; + let applier_factory = build_applier_factory( + "test_create_and_query_regex_", + rows, + HashSet::from([1, 2, 4]), + ) + .await; let expr = binary_expr(col("tag_str"), Operator::RegexMatch, lit(".*")); let res = applier_factory(expr).await; diff --git a/tests/cases/standalone/common/alter/change_col_inverted_index.result b/tests/cases/standalone/common/alter/change_col_inverted_index.result index df684fec436..c1265f2958e 100644 --- a/tests/cases/standalone/common/alter/change_col_inverted_index.result +++ b/tests/cases/standalone/common/alter/change_col_inverted_index.result @@ -180,3 +180,151 @@ DROP TABLE test_pk; Affected Rows: 0 +-- The old SST has an inverted index, but not on host. +CREATE TABLE missing_inverted_index ( + ts TIMESTAMP TIME INDEX, + indexed_col STRING INVERTED INDEX, + host STRING +) WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +INSERT INTO missing_inverted_index VALUES + (1, 'a', 'server1'), + (2, 'b', 'server2'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('missing_inverted_index'); + ++---------------------------------------------+ +| ADMIN FLUSH_TABLE('missing_inverted_index') | ++---------------------------------------------+ +| 0 | ++---------------------------------------------+ + +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; + ++-------------------------+-------------+---------+ +| ts | indexed_col | host | ++-------------------------+-------------+---------+ +| 1970-01-01T00:00:00.001 | a | server1 | ++-------------------------+-------------+---------+ + +ALTER TABLE missing_inverted_index MODIFY COLUMN host SET INVERTED INDEX; + +Affected Rows: 0 + +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; + ++-------------------------+-------------+---------+ +| ts | indexed_col | host | ++-------------------------+-------------+---------+ +| 1970-01-01T00:00:00.001 | a | server1 | ++-------------------------+-------------+---------+ + +SELECT * FROM missing_inverted_index WHERE host = 'missing' ORDER BY ts; + ++----+-------------+------+ +| ts | indexed_col | host | ++----+-------------+------+ ++----+-------------+------+ + +SELECT * FROM missing_inverted_index WHERE indexed_col = 'a' AND host = 'server1' ORDER BY ts; + ++-------------------------+-------------+---------+ +| ts | indexed_col | host | ++-------------------------+-------------+---------+ +| 1970-01-01T00:00:00.001 | a | server1 | ++-------------------------+-------------+---------+ + +SELECT * FROM missing_inverted_index WHERE indexed_col = 'b' AND host = 'server1' ORDER BY ts; + ++----+-------------+------+ +| ts | indexed_col | host | ++----+-------------+------+ ++----+-------------+------+ + +-- New SSTs have both indexes; queries must also retain rows from the old SST. +INSERT INTO missing_inverted_index VALUES + (3, 'a', 'server1'), + (4, 'b', 'server2'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('missing_inverted_index'); + ++---------------------------------------------+ +| ADMIN FLUSH_TABLE('missing_inverted_index') | ++---------------------------------------------+ +| 0 | ++---------------------------------------------+ + +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; + ++-------------------------+-------------+---------+ +| ts | indexed_col | host | ++-------------------------+-------------+---------+ +| 1970-01-01T00:00:00.001 | a | server1 | +| 1970-01-01T00:00:00.003 | a | server1 | ++-------------------------+-------------+---------+ + +SELECT * FROM missing_inverted_index WHERE indexed_col = 'a' AND host = 'server1' ORDER BY ts; + ++-------------------------+-------------+---------+ +| ts | indexed_col | host | ++-------------------------+-------------+---------+ +| 1970-01-01T00:00:00.001 | a | server1 | +| 1970-01-01T00:00:00.003 | a | server1 | ++-------------------------+-------------+---------+ + +-- SQLNESS ARG restart=true +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; + ++-------------------------+-------------+---------+ +| ts | indexed_col | host | ++-------------------------+-------------+---------+ +| 1970-01-01T00:00:00.001 | a | server1 | +| 1970-01-01T00:00:00.003 | a | server1 | ++-------------------------+-------------+---------+ + +-- A column absent from old SSTs can still match its default value. +ALTER TABLE missing_inverted_index ADD COLUMN source STRING DEFAULT 'legacy'; + +Affected Rows: 0 + +ALTER TABLE missing_inverted_index MODIFY COLUMN source SET INVERTED INDEX; + +Affected Rows: 0 + +SELECT * FROM missing_inverted_index WHERE source = 'legacy' ORDER BY ts; + ++-------------------------+-------------+---------+--------+ +| ts | indexed_col | host | source | ++-------------------------+-------------+---------+--------+ +| 1970-01-01T00:00:00.001 | a | server1 | legacy | +| 1970-01-01T00:00:00.002 | b | server2 | legacy | +| 1970-01-01T00:00:00.003 | a | server1 | legacy | +| 1970-01-01T00:00:00.004 | b | server2 | legacy | ++-------------------------+-------------+---------+--------+ + +SELECT * FROM missing_inverted_index WHERE source = 'missing' ORDER BY ts; + ++----+-------------+------+--------+ +| ts | indexed_col | host | source | ++----+-------------+------+--------+ ++----+-------------+------+--------+ + +SELECT * FROM missing_inverted_index WHERE indexed_col = 'a' AND source = 'legacy' ORDER BY ts; + ++-------------------------+-------------+---------+--------+ +| ts | indexed_col | host | source | ++-------------------------+-------------+---------+--------+ +| 1970-01-01T00:00:00.001 | a | server1 | legacy | +| 1970-01-01T00:00:00.003 | a | server1 | legacy | ++-------------------------+-------------+---------+--------+ + +DROP TABLE missing_inverted_index; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/alter/change_col_inverted_index.sql b/tests/cases/standalone/common/alter/change_col_inverted_index.sql index aadc463d6fc..857bbb0aaf2 100644 --- a/tests/cases/standalone/common/alter/change_col_inverted_index.sql +++ b/tests/cases/standalone/common/alter/change_col_inverted_index.sql @@ -52,3 +52,48 @@ ALTER TABLE test_pk MODIFY COLUMN foo SET INVERTED INDEX; SHOW INDEX FROM test_pk; DROP TABLE test_pk; + +-- The old SST has an inverted index, but not on host. +CREATE TABLE missing_inverted_index ( + ts TIMESTAMP TIME INDEX, + indexed_col STRING INVERTED INDEX, + host STRING +) WITH ('append_mode' = 'true'); + +INSERT INTO missing_inverted_index VALUES + (1, 'a', 'server1'), + (2, 'b', 'server2'); + +ADMIN FLUSH_TABLE('missing_inverted_index'); + +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; + +ALTER TABLE missing_inverted_index MODIFY COLUMN host SET INVERTED INDEX; + +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; +SELECT * FROM missing_inverted_index WHERE host = 'missing' ORDER BY ts; +SELECT * FROM missing_inverted_index WHERE indexed_col = 'a' AND host = 'server1' ORDER BY ts; +SELECT * FROM missing_inverted_index WHERE indexed_col = 'b' AND host = 'server1' ORDER BY ts; + +-- New SSTs have both indexes; queries must also retain rows from the old SST. +INSERT INTO missing_inverted_index VALUES + (3, 'a', 'server1'), + (4, 'b', 'server2'); + +ADMIN FLUSH_TABLE('missing_inverted_index'); + +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; +SELECT * FROM missing_inverted_index WHERE indexed_col = 'a' AND host = 'server1' ORDER BY ts; + +-- SQLNESS ARG restart=true +SELECT * FROM missing_inverted_index WHERE host = 'server1' ORDER BY ts; + +-- A column absent from old SSTs can still match its default value. +ALTER TABLE missing_inverted_index ADD COLUMN source STRING DEFAULT 'legacy'; +ALTER TABLE missing_inverted_index MODIFY COLUMN source SET INVERTED INDEX; + +SELECT * FROM missing_inverted_index WHERE source = 'legacy' ORDER BY ts; +SELECT * FROM missing_inverted_index WHERE source = 'missing' ORDER BY ts; +SELECT * FROM missing_inverted_index WHERE indexed_col = 'a' AND source = 'legacy' ORDER BY ts; + +DROP TABLE missing_inverted_index; From 19a128f8d31a5e9264ccd6ebcf10240590873147 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 03:37:15 +0000 Subject: [PATCH 24/46] fix(session): respect supported read preferences in wire tests (#9550) Signed-off-by: WenyXu --- src/session/src/context.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 2f81dfabffb..20462933a39 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -974,7 +974,7 @@ mod test { assert_eq!(ctx.read_preference(), ReadPreference::Leader); assert!(!ctx.extensions().contains_key(READ_PREFERENCE_EXTENSION_KEY)); } - for value in ["", "unknown", "follower", "follower_preferred", "random"] { + for value in ["", "unknown"] { let error = QueryContext::try_from(wire(Some(value))).unwrap_err(); assert!(matches!( error, From 394562aa9ed2ad3cc6fef1e15aa1c09f42f759e1 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 13 Oct 2026 03:48:17 +0000 Subject: [PATCH 25/46] fix: prevent stale leader metadata after concurrent writes (#9549) fix: fence leader cache fills around store mutations Signed-off-by: jeremyhi --- src/meta-srv/src/service/store/cached_kv.rs | 361 +++++++++++++++++++- 1 file changed, 359 insertions(+), 2 deletions(-) diff --git a/src/meta-srv/src/service/store/cached_kv.rs b/src/meta-srv/src/service/store/cached_kv.rs index 6526ebcd41e..1c2990f5446 100644 --- a/src/meta-srv/src/service/store/cached_kv.rs +++ b/src/meta-srv/src/service/store/cached_kv.rs @@ -69,6 +69,9 @@ pub struct LeaderCachedKvBackend { check_leader: CheckLeaderRef, store: KvBackendRef, cache: ResettableKvBackendRef, + /// Incremented before and after each store mutation. The first increment + /// detects overlapping writers; the second invalidates cache fills that + /// read the old value while the store mutation was in flight. version: AtomicUsize, name: String, } @@ -214,9 +217,11 @@ impl KvBackend for LeaderCachedKvBackend { let ver = self.create_new_version(); let res = self.store.put(req.clone()).await?; + self.create_new_version(); let _ = self.cache.put(req.clone()).await?; - if !self.validate_version(ver) { + // Only our own completion increment may follow the starting version. + if !self.validate_version(ver + 1) { self.invalid_key(req.key).await?; } @@ -231,9 +236,11 @@ impl KvBackend for LeaderCachedKvBackend { let ver = self.create_new_version(); let res = self.store.batch_put(req.clone()).await?; + self.create_new_version(); let _ = self.cache.batch_put(req.clone()).await?; - if !self.validate_version(ver) { + // Only our own completion increment may follow the starting version. + if !self.validate_version(ver + 1) { let keys = req.kvs.into_iter().map(|kv| kv.key).collect::>(); self.invalid_keys(keys).await?; } @@ -306,6 +313,7 @@ impl KvBackend for LeaderCachedKvBackend { let key = req.key.clone(); let res = self.store.compare_and_put(req).await?; + self.create_new_version(); // Delete key in the cache. // // Cache can not deal with the CAS operation, because it does @@ -322,6 +330,7 @@ impl KvBackend for LeaderCachedKvBackend { let _ = self.create_new_version(); let res = self.store.delete_range(req.clone()).await?; + self.create_new_version(); let _ = self.cache.delete_range(req).await?; Ok(res) } @@ -334,6 +343,7 @@ impl KvBackend for LeaderCachedKvBackend { let _ = self.create_new_version(); let res = self.store.batch_delete(req.clone()).await?; + self.create_new_version(); let _ = self.cache.batch_delete(req).await?; Ok(res) } @@ -351,6 +361,7 @@ impl TxnService for LeaderCachedKvBackend { let _ = self.create_new_version(); let res = self.store.txn(txn.clone()).await?; + self.create_new_version(); let TxnRequest { success, failure, .. } = txn.into(); @@ -384,6 +395,7 @@ impl TxnService for LeaderCachedKvBackend { impl ResettableKvBackend for LeaderCachedKvBackend { fn reset(&self) { + self.create_new_version(); self.cache.reset() } @@ -394,10 +406,355 @@ impl ResettableKvBackend for LeaderCachedKvBackend { #[cfg(test)] mod tests { + use std::sync::Mutex; + + use common_meta::kv_backend::txn::{Compare, CompareOp}; use common_meta::rpc::KeyValue; + use futures::poll; + use tokio::sync::oneshot; use super::*; + /// Pauses one backend call at a known point, without relying on task scheduling. + #[derive(Default)] + struct Pause { + receiver: Mutex>>, + } + + impl Pause { + fn arm(&self) -> oneshot::Sender<()> { + let (sender, receiver) = oneshot::channel(); + assert!(self.receiver.lock().unwrap().replace(receiver).is_none()); + sender + } + + async fn wait(&self) { + let receiver = self.receiver.lock().unwrap().take(); + if let Some(receiver) = receiver { + receiver.await.unwrap(); + } + } + } + + #[derive(Default)] + struct ControlledBackend { + inner: MemoryKvBackend, + before_write: Pause, + after_write: Pause, + after_read: Pause, + } + + #[async_trait::async_trait] + impl TxnService for ControlledBackend { + type Error = Error; + + async fn txn(&self, txn: Txn) -> Result { + self.before_write.wait().await; + let res = self.inner.txn(txn).await; + self.after_write.wait().await; + res + } + + fn max_txn_ops(&self) -> usize { + self.inner.max_txn_ops() + } + } + + #[async_trait::async_trait] + impl KvBackend for ControlledBackend { + fn name(&self) -> &str { + "ControlledBackend" + } + + fn as_any(&self) -> &dyn Any { + self + } + + async fn range(&self, req: RangeRequest) -> Result { + let res = self.inner.range(req).await; + self.after_read.wait().await; + res + } + + async fn batch_get(&self, req: BatchGetRequest) -> Result { + let res = self.inner.batch_get(req).await; + self.after_read.wait().await; + res + } + + async fn put(&self, req: PutRequest) -> Result { + self.before_write.wait().await; + let res = self.inner.put(req).await; + self.after_write.wait().await; + res + } + + async fn batch_put(&self, req: BatchPutRequest) -> Result { + self.before_write.wait().await; + let res = self.inner.batch_put(req).await; + self.after_write.wait().await; + res + } + + async fn delete_range(&self, req: DeleteRangeRequest) -> Result { + self.before_write.wait().await; + let res = self.inner.delete_range(req).await; + self.after_write.wait().await; + res + } + + async fn batch_delete(&self, req: BatchDeleteRequest) -> Result { + self.before_write.wait().await; + let res = self.inner.batch_delete(req).await; + self.after_write.wait().await; + res + } + } + + #[derive(Clone, Copy, Debug)] + enum Write { + Txn, + TxnFailure, + Put, + BatchPut, + CompareAndPut, + DeleteRange, + BatchDelete, + } + + const WRITES: [Write; 7] = [ + Write::Txn, + Write::TxnFailure, + Write::Put, + Write::BatchPut, + Write::CompareAndPut, + Write::DeleteRange, + Write::BatchDelete, + ]; + const KEY: &[u8] = b"route"; + + impl Write { + async fn execute(self, store: &LeaderCachedKvBackend, value: &[u8]) { + match self { + Self::Txn => { + store + .txn(Txn::new().and_then(vec![TxnOp::Put(KEY.to_vec(), value.to_vec())])) + .await + .unwrap(); + } + Self::TxnFailure => { + let res = store + .txn( + Txn::new() + .when(vec![Compare::with_value_not_exists( + KEY.to_vec(), + CompareOp::Equal, + )]) + .or_else(vec![TxnOp::Put(KEY.to_vec(), value.to_vec())]), + ) + .await + .unwrap(); + assert!(!res.succeeded); + } + Self::Put => { + store + .put(PutRequest::new().with_key(KEY).with_value(value)) + .await + .unwrap(); + } + Self::BatchPut => { + store + .batch_put(BatchPutRequest { + kvs: vec![KeyValue { + key: KEY.to_vec(), + value: value.to_vec(), + }], + ..Default::default() + }) + .await + .unwrap(); + } + Self::CompareAndPut => { + let res = store + .compare_and_put( + CompareAndPutRequest::new() + .with_key(KEY) + .with_expect(b"old") + .with_value(value), + ) + .await + .unwrap(); + assert!(res.success); + } + Self::DeleteRange => { + store + .delete_range(DeleteRangeRequest::new().with_key(KEY)) + .await + .unwrap(); + } + Self::BatchDelete => { + store + .batch_delete(BatchDeleteRequest { + keys: vec![KEY.to_vec()], + prev_kv: false, + }) + .await + .unwrap(); + } + } + } + } + + async fn read(store: &LeaderCachedKvBackend, batch: bool) -> Option { + if batch { + store + .batch_get(BatchGetRequest { + keys: vec![KEY.to_vec()], + }) + .await + .unwrap() + .kvs + .pop() + } else { + store.get(KEY).await.unwrap() + } + } + + /// A cache miss reads the old value after a writer starts, then returns after + /// the writer has committed and updated/invalidated the cache (#9548). + #[tokio::test] + async fn test_cache_miss_overlaps_write() { + for write in WRITES { + for batch in [true, false] { + let backing = Arc::new(ControlledBackend::default()); + backing + .inner + .put(PutRequest::new().with_key(KEY).with_value(b"old")) + .await + .unwrap(); + let store = LeaderCachedKvBackend::with_always_leader(backing.clone()); + + let commit = backing.before_write.arm(); + let writer = write.execute(&store, b"new"); + tokio::pin!(writer); + assert!(poll!(writer.as_mut()).is_pending()); + + let return_read = backing.after_read.arm(); + let reader = read(&store, batch); + tokio::pin!(reader); + assert!(poll!(reader.as_mut()).is_pending()); + + commit.send(()).unwrap(); + writer.await; + return_read.send(()).unwrap(); + assert_eq!(reader.await.unwrap().value(), b"old"); + + let durable = backing.inner.get(KEY).await.unwrap(); + assert_eq!( + read(&store, batch).await, + durable, + "{write:?}, batch={batch}" + ); + } + } + } + + /// Start order, commit order and response order can all differ. An older + /// put response must not overwrite the cache after a newer mutation. + #[tokio::test] + async fn test_overlapping_writes() { + for first in [Write::Put, Write::BatchPut] { + for second in WRITES { + for reverse_commit in [false, true] { + let backing = Arc::new(ControlledBackend::default()); + backing + .inner + .put(PutRequest::new().with_key(KEY).with_value(b"old")) + .await + .unwrap(); + let store = LeaderCachedKvBackend::with_always_leader(backing.clone()); + + let resume = if reverse_commit { + backing.before_write.arm() + } else { + backing.after_write.arm() + }; + // Keep the initial value so the second CAS succeeds. + let writer = first.execute(&store, b"old"); + tokio::pin!(writer); + assert!(poll!(writer.as_mut()).is_pending()); + second.execute(&store, b"new").await; + resume.send(()).unwrap(); + writer.await; + + let durable = backing.inner.get(KEY).await.unwrap(); + assert_eq!( + read(&store, true).await, + durable, + "{first:?}, {second:?}, reverse_commit={reverse_commit}" + ); + } + } + } + } + + /// The later-started put commits first but returns last, after its stored + /// value has already been replaced by the first writer. + #[tokio::test] + async fn test_write_response_reordering() { + for first in WRITES { + for second in [Write::Put, Write::BatchPut] { + let backing = Arc::new(ControlledBackend::default()); + backing + .inner + .put(PutRequest::new().with_key(KEY).with_value(b"old")) + .await + .unwrap(); + let store = LeaderCachedKvBackend::with_always_leader(backing.clone()); + + let commit_first = backing.before_write.arm(); + let first_writer = first.execute(&store, b"new"); + tokio::pin!(first_writer); + assert!(poll!(first_writer.as_mut()).is_pending()); + + let return_second = backing.after_write.arm(); + let second_writer = second.execute(&store, b"old"); + tokio::pin!(second_writer); + assert!(poll!(second_writer.as_mut()).is_pending()); + + commit_first.send(()).unwrap(); + first_writer.await; + return_second.send(()).unwrap(); + second_writer.await; + + let durable = backing.inner.get(KEY).await.unwrap(); + assert_eq!(read(&store, true).await, durable, "{first:?}, {second:?}"); + } + } + } + + #[tokio::test] + async fn test_cache_miss_overlaps_reset() { + for batch in [true, false] { + let backing = Arc::new(ControlledBackend::default()); + backing + .inner + .put(PutRequest::new().with_key(KEY).with_value(b"old")) + .await + .unwrap(); + let store = LeaderCachedKvBackend::with_always_leader(backing.clone()); + + let return_read = backing.after_read.arm(); + let reader = read(&store, batch); + tokio::pin!(reader); + assert!(poll!(reader.as_mut()).is_pending()); + store.reset(); + return_read.send(()).unwrap(); + reader.await; + assert!(store.cache.get(KEY).await.unwrap().is_none()); + } + } + fn create_leader_cached_kv_backend() -> LeaderCachedKvBackend { let store = Arc::new(MemoryKvBackend::new()); LeaderCachedKvBackend::with_always_leader(store) From 18634e17bb1a50c1e0abd2322d4d0cf4ad548c31 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 13 Oct 2026 05:17:41 +0000 Subject: [PATCH 26/46] refactor(mito): retain region compaction and DDL lifecycle state (#9498) * refactor(mito): retain region compaction and DDL lifecycle state Signed-off-by: Lei, HUANG * fix(mito): finalize opening batch in the worker on registration Clean up the `opening_regions` entry and reply to its waiters in the same serialized worker step that registers an opened region, instead of waiting for the background opener to resume after the registration acknowledgement. Previously the entry was only removed by the background opener, so a close followed by a reopen could join the stale opening batch and receive a false success after the region no longer existed. Replying from the worker closes that window because registration and waiter cleanup are now atomic with respect to the worker loop. Add a test-only `EventListener::on_region_open_registered` hook and a regression test that holds the first background opener to verify a reopen does not wait for the previous open task. On failure the region is still stopped and waiters receive `RegionClosed`. Signed-off-by: Lei, HUANG * fix(mito): unregister drop runtime state before WAL cleanup The region leaves the region map before its WAL entries are obsoleted. A failure in `wal.obsolete(...).await?` returned early and skipped `cleanup_dropped_region_runtime_state`, so the resident compaction, flush, and index scheduler entries were never unregistered. Since the region is no longer in the map, a later close cannot clean them up either, leaking the entries for the rest of the worker's lifetime. Move the runtime-state cleanup to immediately after the region leaves the map, before the fallible WAL cleanup. The schedulers only touch their own state and do not depend on the region map or WAL, so the reordering is safe and leaves the success path unchanged. Signed-off-by: Lei, HUANG * refactor(mito): tighten visibility of new compaction state items Keep `ActiveCompaction::idle` private (it is only used inside `state.rs`) and expose the other items added by the resident-region refactor as `pub(crate)` instead of `pub(super)`, per the style guide. Also drop a stale comment in `on_compaction_finished`: DDL handoff now happens in `take_ready_ddls` before this block, which only runs when no DDL is queued. Signed-off-by: Lei, HUANG * fix(mito): skip compaction scheduling while a region is truncating `handle_flush_finished` dispatches queued DDLs before scheduling compaction. A truncate queued behind a committing flush marks the region as `Truncating` synchronously, but `schedule_compaction` only skipped staging states, so it could schedule a new automatic compaction while the truncate was still running. Add `MitoRegion::is_truncating` and skip compaction scheduling for a truncating region. Add a regression test that queues a full truncate behind a committing flush and asserts no compaction is scheduled. Signed-off-by: Lei, HUANG * fix(mito): log DDL-fenced compaction rejection The refactor dropped the log emitted when a region was skipped because of pending DDL requests. Automatic triggers carry no waiter, so the rejection became silent and could not be told apart from a scheduling interval that had not elapsed. Log the queued and executing DDL state before cancelling. Signed-off-by: Lei, HUANG * docs(mito): document DDL reply fence invariants Explain why fence_ddl_replies installs executing_ddl before handing off the batch and why CompactionDdlComplete carries both a scheduler-wide generation and a per-slot request_id. The generation fences replies from superseded batches or region incarnations, while request_id makes completion idempotent so only the last distinct reply releases the fence. Signed-off-by: Lei, HUANG * refactor(mito): rename DDL completion request_id to request_index The value is the DDL's index within the handed-off batch, not a globally stable request identifier; request_id suggested otherwise. Rename it to request_index across the scheduler, worker dispatch, and tests. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/AGENTS.md | 3 + src/mito2/src/compaction/scheduler.rs | 238 ++++++++++---- .../src/compaction/scheduler/planning.rs | 19 +- src/mito2/src/compaction/scheduler/state.rs | 92 ++++-- src/mito2/src/compaction/scheduler_test.rs | 311 +++++++++++++++--- src/mito2/src/engine/compaction_test.rs | 55 ++++ src/mito2/src/engine/listener.rs | 3 + src/mito2/src/engine/open_test.rs | 109 +++++- src/mito2/src/engine/truncate_test.rs | 100 ++++++ src/mito2/src/region.rs | 5 + src/mito2/src/request.rs | 11 + src/mito2/src/worker.rs | 37 +++ src/mito2/src/worker/handle_catchup.rs | 2 +- src/mito2/src/worker/handle_compaction.rs | 8 +- src/mito2/src/worker/handle_create.rs | 2 +- src/mito2/src/worker/handle_drop.rs | 6 +- src/mito2/src/worker/handle_enter_staging.rs | 35 +- src/mito2/src/worker/handle_open.rs | 39 ++- src/mito2/src/worker/handle_truncate.rs | 27 +- 19 files changed, 908 insertions(+), 194 deletions(-) diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index b8fd2f8892c..efee5f6cca1 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -88,6 +88,9 @@ Tests live next to the code as `*_test.rs` (e.g. `src/mito2/src/engine/flush_tes - Lock ordering: take the manifest lock before updating `version_control`; the reverse deadlocks against concurrent flush/compaction. - All region I/O runs on tokio workers; never `block_on` inside a worker. +- Loaded regions have resident compaction state, including `Idle` and handed-off + DDLs. Publish create/open/catchup results through the worker's `register_region`; + only close/drop unregisters them. See `compaction/scheduler/state.rs`. ## Maintenance contract diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index ff92f8a41c3..e9da3002254 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -27,7 +27,7 @@ use common_time::range::TimestampRange; pub(crate) use planning::CompactionPickFinished; pub use planning::CompactionRequest; pub(crate) use state::CompactionExecution; -use state::{CompactionStatus, CompactionTrigger, PendingCompaction}; +use state::{CompactionStatus, CompactionTrigger, DdlExecution, PendingCompaction}; use store_api::storage::RegionId; use tokio::sync::mpsc::Sender; @@ -51,7 +51,7 @@ use crate::worker::WorkerListener; /// Compaction scheduler tracks and manages compaction tasks. pub(crate) struct CompactionScheduler { scheduler: SchedulerRef, - /// Compacting regions. + /// All loaded regions, including idle regions and regions executing DDLs. region_status: HashMap, /// Request sender of the worker that this scheduler belongs to. request_sender: Sender, @@ -139,11 +139,38 @@ impl CompactionScheduler { } } + /// Registers a loaded region; replacing its VersionControl invalidates the old incarnation. + pub(crate) fn register_region( + &mut self, + version_control: &VersionControlRef, + access_layer: &AccessLayerRef, + ) { + let region_id = version_control.region_id(); + if let Some(status) = self.region_status.get_mut(®ion_id) { + if Arc::ptr_eq(&status.version_control, version_control) { + return; + } + status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build())); + } + self.region_status.insert( + region_id, + CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()), + ); + } + + /// Reports active work rather than the existence of a resident region record. + #[cfg(test)] + fn is_compacting(&self, region_id: RegionId) -> bool { + self.region_status + .get(®ion_id) + .is_some_and(CompactionStatus::is_compacting) + } + /// Accepts an automatic compaction trigger for a region. /// /// # Effects /// - /// Starts planning when no scheduler status exists for the region. + /// Starts planning when the registered region is idle. /// Otherwise, an active status coalesces the trigger into one unrestricted /// follow-up cycle. Returns `true` only when this call dispatches planning. /// @@ -178,7 +205,7 @@ impl CompactionScheduler { /// /// # Effects /// - /// Starts planning when no scheduler status exists for the region. During + /// Starts planning when the registered region is idle. During /// an automatic compaction, queues the request and replaces any older /// pending manual request; during a manual compaction, rejects the new /// request through `waiter`. Returns `true` only when this call dispatches @@ -265,7 +292,7 @@ impl CompactionScheduler { /// # Effects /// /// Notifies waiters, schedules a pending manual or automatic follow-up, or - /// removes the region status. A follow-up is scheduled whenever the cycle + /// returns the region to idle. A follow-up is scheduled whenever the cycle /// latched an automatic trigger, or when `made_progress` is true (the /// execution produced output or reduced the file count): successful compaction /// keeps draining the region until the picker returns no plan. The returned @@ -301,7 +328,7 @@ impl CompactionScheduler { /// /// # Effects /// - /// Removes the matching region status, notifies compaction waiters, and + /// Resets the matching region to idle, notifies compaction waiters, and /// returns DDLs that were waiting for cancellation. /// /// # Constraints @@ -323,7 +350,7 @@ impl CompactionScheduler { /// /// # Effects /// - /// Removes the matching region status and fails all compaction waiters and + /// Resets the matching region to idle and fails all compaction waiters and /// dependent DDL requests with the supplied error. /// /// # Constraints @@ -352,7 +379,7 @@ impl CompactionScheduler { /// /// The owning worker must invoke this as part of serialized region teardown. pub(crate) fn on_region_dropped(&mut self, region_id: RegionId) { - self.remove_region_on_failure( + self.remove_region( region_id, Arc::new(RegionDroppedSnafu { region_id }.build()), ); @@ -368,10 +395,10 @@ impl CompactionScheduler { /// /// The owning worker must invoke this as part of serialized region teardown. pub(crate) fn on_region_closed(&mut self, region_id: RegionId) { - self.remove_region_on_failure(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); + self.remove_region(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); } - /// Removes compaction state because the region was truncated. + /// Resets the cycle after truncate, retaining registration and outstanding DDL replies. /// /// # Effects /// @@ -381,7 +408,7 @@ impl CompactionScheduler { /// /// The owning worker must invoke this after truncate completes. pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) { - self.remove_region_on_failure( + self.fail_compaction( region_id, Arc::new(RegionTruncatedSnafu { region_id }.build()), ); @@ -412,6 +439,9 @@ impl CompactionScheduler { let Some(status) = self.region_status.get_mut(®ion_id) else { return Err((sender, request)); }; + if !status.is_compacting() { + return Err((sender, request)); + } status.request_cancel(); let request = SenderDdlRequest { @@ -454,6 +484,19 @@ impl CompactionScheduler { time_range: Option, ) -> Result { let region_id = version_control.region_id(); + let status = self.region_status.entry(region_id).or_insert_with(|| { + CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()) + }); + if status.has_ddl() { + info!( + "Region {} has pending DDL requests, ignoring compaction: queued={}, executing={}", + region_id, + status.pending_ddl_requests.len(), + status.executing_ddl.is_some(), + ); + waiter.send(CompactionCancelledSnafu.fail()); + return Ok(false); + } let current_state = manifest_ctx.current_state(); if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) { info!( @@ -464,18 +507,7 @@ impl CompactionScheduler { return Ok(false); } - if let Some(status) = self.region_status.get_mut(®ion_id) { - // Pending Truncate/EnterStaging requests form a scheduling fence. Any later - // manual request receives CompactionCancelled; automatic triggers are ignored. - if !status.pending_ddl_requests.is_empty() { - waiter.send(CompactionCancelledSnafu.fail()); - info!( - "Region {} has pending DDL requests, ignoring compaction: {:?}", - region_id, compact_options - ); - return Ok(false); - } - + if status.is_compacting() { match trigger { CompactionTrigger::Automatic => status.mark_automatic_trigger(), CompactionTrigger::Manual if status.is_manual_compaction() => { @@ -503,13 +535,7 @@ impl CompactionScheduler { // Publish the picking phase before dispatching background planning. let plan_id = Self::next_plan_id(&mut self.next_plan_id); - let mut status = CompactionStatus::new( - region_id, - version_control.clone(), - access_layer.clone(), - plan_id, - trigger, - ); + status.start_picking_with_trigger(plan_id, trigger); let request = status.new_compaction_request( self.request_sender.clone(), self.engine_config.clone(), @@ -520,7 +546,6 @@ impl CompactionScheduler { max_parallelism, ); status.merge_waiter(waiter); - self.region_status.insert(region_id, status); self.dispatch_compaction_planning(plan_id, request, compact_options, time_range); self.listener.on_compaction_scheduled(region_id); Ok(true) @@ -587,6 +612,10 @@ impl CompactionScheduler { return CompactionTransition::NoAction; } + if let Some(ddls) = self.take_ready_ddls(region_id) { + return CompactionTransition::DdlReady(ddls); + } + if self.handle_pending_compaction_request( region_id, manifest_ctx, @@ -595,8 +624,6 @@ impl CompactionScheduler { return CompactionTransition::NoAction; } - // The region status might be removed by the previous steps. - // So we return empty DDL requests. let Some(status) = self.region_status.get_mut(®ion_id) else { return CompactionTransition::NoAction; }; @@ -604,18 +631,6 @@ impl CompactionScheduler { waiter.send(Ok(0)); } - // A queued DDL was waiting for the current task to terminate; chaining - // another compaction ahead of it would delay the DDL by a whole extra - // plan/execution cycle, so dispatch the DDLs first. - let pending_ddl_requests = std::mem::take(&mut status.pending_ddl_requests); - if !pending_ddl_requests.is_empty() { - // The DDL supersedes any retained automatic follow-up. - self.region_status.remove(®ion_id); - // If there are pending DDL requests, we should return them to the caller. - // And skip try to schedule next compaction task. - return CompactionTransition::DdlReady(pending_ddl_requests); - } - // Keep draining when the cycle latched an automatic trigger, or when the // execution produced output or reduced the file count. An empty edit stops // here; the next flush trigger resumes. @@ -625,7 +640,9 @@ impl CompactionScheduler { { return CompactionTransition::AutomaticFollowupScheduled; } - self.region_status.remove(®ion_id); + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.become_idle(); + } CompactionTransition::NoAction } @@ -668,13 +685,13 @@ impl CompactionScheduler { /// Notifies the scheduler that the compaction job is cancelled cooperatively. async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec { - self.remove_region_on_cancel(region_id) + self.finish_compaction_on_cancel(region_id) } /// Notifies the scheduler that the compaction job is failed. fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc) { error!(err; "Region {} failed to compact, cancel all pending tasks", region_id); - self.remove_region_on_failure(region_id, err); + self.fail_compaction(region_id, err); } #[cfg(test)] @@ -708,28 +725,129 @@ impl CompactionScheduler { .request_cancel() } - fn remove_region_on_failure(&mut self, region_id: RegionId, err: Arc) { - // Remove this region. - let Some(status) = self.region_status.remove(®ion_id) else { - return; - }; - - // Notifies all pending tasks. - status.on_failure(err); + /// Unregisters a closed or dropped region and fails its queued requests. + fn remove_region(&mut self, region_id: RegionId, err: Arc) { + if let Some(mut status) = self.region_status.remove(®ion_id) { + status.on_failure(err); + } } - fn remove_region_on_cancel(&mut self, region_id: RegionId) -> Vec { - let Some(status) = self.region_status.remove(®ion_id) else { + /// Fails a cycle without discarding the loaded region's registration. + fn fail_compaction(&mut self, region_id: RegionId, err: Arc) { + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.on_failure(err); + } + } + + /// Returns the cycle to idle and records any DDL handoff in the same resident entry. + fn finish_compaction_on_cancel(&mut self, region_id: RegionId) -> Vec { + let Some(status) = self.region_status.get_mut(®ion_id) else { return Vec::new(); }; - status.on_cancel() + let ddls = status.on_cancel(); + self.fence_ddl_replies(region_id, ddls) + } + + /// Hands off pending DDLs after compaction ends and cancels remaining compaction waiters. + fn take_ready_ddls(&mut self, region_id: RegionId) -> Option> { + if self + .region_status + .get(®ion_id)? + .pending_ddl_requests + .is_empty() + { + return None; + } + Some(self.finish_compaction_on_cancel(region_id)) + } + + /// Hands off a drained batch of DDLs and fences new compaction until every reply lands. + /// + /// Installs the fence, intercepts each DDL's reply, and returns the same batch so the caller + /// can execute it. The real result of each DDL is still forwarded to its original sender. + /// + /// # Constraints + /// Must only be called from a compaction terminal transition, and only by the worker that owns + /// `region_id`. + fn fence_ddl_replies( + &mut self, + region_id: RegionId, + mut ddls: Vec, + ) -> Vec { + if ddls.is_empty() { + return ddls; + } + let generation = Self::next_plan_id(&mut self.next_plan_id); + if let Some(status) = self.region_status.get_mut(®ion_id) { + // Set region status to "executing DDLs" so future compaction won't be scheduled. + status.executing_ddl = Some(DdlExecution::new(generation, ddls.len())); + } + + for (request_index, ddl) in ddls.iter_mut().enumerate() { + // Swap the DDL's real reply sender for a temporary one so we can observe it. The + // spawned task forwards the observed result to the original caller and, in between, + // notifies the scheduler that this slot of the batch has completed. + let (tx, rx) = tokio::sync::oneshot::channel(); + let reply = std::mem::replace(&mut ddl.sender, tx.into()); + let sender = self.request_sender.clone(); + common_runtime::spawn_global(async move { + use snafu::ResultExt; + let result = rx + .await + .context(crate::error::RecvSnafu) + .and_then(|result| result); + // Route the completion through the worker channel so `on_ddl_complete` sees it + // serialized with other region state transitions. + // + // At most one batch can be outstanding per region: [`CompactionStatus::has_ddl`] blocks + // compaction admission while `executing_ddl` is set, so `fence_ddl_replies` cannot run again + // for the same region before the current batch drains. A single `(generation, request_index)` + // pair therefore identifies every reply. `generation` comes from the scheduler-wide counter, + // which outlives region statuses, so a late reply from a previous incarnation cannot clear the + // current fence. Inside the batch `request_index` makes completion idempotent: only the last + // distinct reply releases the fence. + let _ = sender + .send(WorkerRequestWithTime::new( + crate::request::WorkerRequest::Background { + region_id, + notify: crate::request::BackgroundNotify::CompactionDdlComplete { + generation, + request_index, + }, + }, + )) + .await; + reply.send(result); + }); + } + ddls + } + + /// Releases the DDL fence after its last reply, ignoring replies from an older generation. + pub(crate) fn on_ddl_complete( + &mut self, + region_id: RegionId, + generation: u64, + request_index: usize, + ) { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return; + }; + // `complete` returns false for a stale generation or a duplicate `request_index`, so a late or + // repeated notify cannot end the batch early. It only clears the fence on the last distinct + // reply, which is what lets `has_ddl` unblock compaction admission again. + if let Some(ddl) = &mut status.executing_ddl + && ddl.complete(generation, request_index) + { + status.executing_ddl = None; + } } } impl Drop for CompactionScheduler { fn drop(&mut self) { - for (region_id, status) in self.region_status.drain() { + for (region_id, mut status) in self.region_status.drain() { // We are shutting down so notify all pending tasks. status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build())); } diff --git a/src/mito2/src/compaction/scheduler/planning.rs b/src/mito2/src/compaction/scheduler/planning.rs index 23318ad7669..5553d6435da 100644 --- a/src/mito2/src/compaction/scheduler/planning.rs +++ b/src/mito2/src/compaction/scheduler/planning.rs @@ -298,7 +298,7 @@ impl CompactionScheduler { // execution callback will follow, so return DDLs released by the fence. if !status.accept_plan(finished.plan_id) { return CompactionTransition::from_pending_ddls( - self.remove_region_on_cancel(region_id), + self.finish_compaction_on_cancel(region_id), ); } @@ -354,7 +354,7 @@ impl CompactionScheduler { .await } Err(err) => { - self.remove_region_on_failure(region_id, Arc::new(err)); + self.fail_compaction(region_id, Arc::new(err)); CompactionTransition::NoAction } } @@ -400,6 +400,9 @@ impl CompactionScheduler { } } + if let Some(ddls) = self.take_ready_ddls(region_id) { + return CompactionTransition::DdlReady(ddls); + } if self.handle_pending_compaction_request( region_id, manifest_ctx, @@ -412,21 +415,15 @@ impl CompactionScheduler { return CompactionTransition::NoAction; }; - // A queued DDL supersedes a retained automatic follow-up, matching the - // execution terminal path in `on_compaction_finished`. - let pending_ddls = std::mem::take(&mut status.pending_ddl_requests); - if !pending_ddls.is_empty() { - self.region_status.remove(®ion_id); - return CompactionTransition::DdlReady(pending_ddls); - } - if status.active.reset_automatic_followup() && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager) { return CompactionTransition::AutomaticFollowupScheduled; } - self.region_status.remove(®ion_id); + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.become_idle(); + } CompactionTransition::NoAction } diff --git a/src/mito2/src/compaction/scheduler/state.rs b/src/mito2/src/compaction/scheduler/state.rs index c5aa0816c55..4fe4a686751 100644 --- a/src/mito2/src/compaction/scheduler/state.rs +++ b/src/mito2/src/compaction/scheduler/state.rs @@ -68,6 +68,8 @@ impl CompactionExecution { #[derive(Debug)] pub(super) enum CompactionPhase { + /// The region is registered but has no compaction in progress. + Idle, Picking { plan_id: u64, cancelled: bool, @@ -102,6 +104,17 @@ pub(super) struct ActiveCompaction { } impl ActiveCompaction { + /// Creates an empty cycle while retaining the region's scheduling record. + fn idle() -> Self { + Self { + phase: CompactionPhase::Idle, + trigger: CompactionTrigger::Automatic, + waiters: Vec::new(), + automatic_followup_required: false, + } + } + + #[cfg(test)] pub(super) fn picking( plan_id: u64, waiters: Vec, @@ -165,7 +178,7 @@ impl ActiveCompaction { pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool { match &self.phase { - CompactionPhase::Picking { .. } => None, + CompactionPhase::Idle | CompactionPhase::Picking { .. } => None, CompactionPhase::Local { execution, .. } | CompactionPhase::Remote { execution } => { Some(execution) } @@ -175,6 +188,7 @@ impl ActiveCompaction { pub(super) fn request_cancel(&mut self) -> RequestCancelResult { match &mut self.phase { + CompactionPhase::Idle => RequestCancelResult::TooLateToCancel, CompactionPhase::Picking { cancelled, .. } => { if *cancelled { RequestCancelResult::AlreadyCancelling @@ -250,7 +264,31 @@ impl Drop for CompactingFilesInner { } } -/// Status of running and pending region compaction tasks. +/// A handed-off DDL batch, retained until each reply has been observed once. +#[derive(Debug)] +pub(crate) struct DdlExecution { + generation: u64, + pending_replies: HashSet, +} + +impl DdlExecution { + /// Tracks a batch's individual replies so duplicate notifications are harmless. + pub(crate) fn new(generation: u64, count: usize) -> Self { + Self { + generation, + pending_replies: (0..count).collect(), + } + } + + /// Returns true only for the last distinct reply of this batch. + pub(crate) fn complete(&mut self, generation: u64, request_index: usize) -> bool { + self.generation == generation + && self.pending_replies.remove(&request_index) + && self.pending_replies.is_empty() + } +} + +/// Scheduling state retained for the lifetime of a loaded region. pub(super) struct CompactionStatus { /// Id of the region. pub(super) region_id: RegionId, @@ -271,24 +309,25 @@ pub(super) struct CompactionStatus { /// [`crate::request::DdlRequest::Truncate`] and [`crate::request::DdlRequest::EnterStaging`] here. Both must serialize with /// compaction so they observe the version after compaction terminates. pub(super) pending_ddl_requests: Vec, + /// DDLs already handed to the worker keep new compaction paused until their replies. + pub(crate) executing_ddl: Option, } impl CompactionStatus { - /// Creates a new picking [CompactionStatus]. + /// Registers an idle region without starting a compaction cycle. pub(super) fn new( region_id: RegionId, version_control: VersionControlRef, access_layer: AccessLayerRef, - plan_id: u64, - trigger: CompactionTrigger, ) -> CompactionStatus { CompactionStatus { region_id, version_control, access_layer, - active: ActiveCompaction::picking(plan_id, Vec::new(), trigger), + active: ActiveCompaction::idle(), pending_request: None, pending_ddl_requests: Vec::new(), + executing_ddl: None, } } @@ -298,13 +337,24 @@ impl CompactionStatus { version_control: VersionControlRef, access_layer: AccessLayerRef, ) -> Self { - Self::new( - region_id, - version_control, - access_layer, - 0, - CompactionTrigger::Automatic, - ) + let mut status = Self::new(region_id, version_control, access_layer); + status.active = ActiveCompaction::picking(0, Vec::new(), CompactionTrigger::Automatic); + status + } + + /// Registration alone does not indicate active work. + pub(crate) fn is_compacting(&self) -> bool { + !matches!(self.active.phase, CompactionPhase::Idle) + } + + /// Both queued and handed-off DDLs prevent new compaction admission. + pub(crate) fn has_ddl(&self) -> bool { + !self.pending_ddl_requests.is_empty() || self.executing_ddl.is_some() + } + + /// Ends the current cycle without losing region or DDL execution state. + pub(crate) fn become_idle(&mut self) { + self.active = ActiveCompaction::idle(); } #[cfg(test)] @@ -395,14 +445,15 @@ impl CompactionStatus { } } - pub(super) fn on_failure(mut self, err: Arc) { - for waiter in self.active.waiters.drain(..) { + pub(super) fn on_failure(&mut self, err: Arc) { + let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle()); + for waiter in active.waiters.drain(..) { waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id: self.region_id, })); } - if let Some(pending_compaction) = self.pending_request { + if let Some(pending_compaction) = self.pending_request.take() { pending_compaction .waiter .send(Err(err.clone()).context(CompactRegionSnafu { @@ -410,7 +461,7 @@ impl CompactionStatus { })); } - for pending_ddl in self.pending_ddl_requests { + for pending_ddl in self.pending_ddl_requests.drain(..) { pending_ddl .sender .send(Err(err.clone()).context(CompactRegionSnafu { @@ -420,12 +471,13 @@ impl CompactionStatus { } #[must_use] - pub(super) fn on_cancel(mut self) -> Vec { - for waiter in self.active.waiters.drain(..) { + pub(super) fn on_cancel(&mut self) -> Vec { + let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle()); + for waiter in active.waiters.drain(..) { waiter.send(CompactionCancelledSnafu.fail()); } - if let Some(pending_compaction) = self.pending_request { + if let Some(pending_compaction) = self.pending_request.take() { pending_compaction.waiter.send( Err(Arc::new(CompactionCancelledSnafu.build())).context(CompactRegionSnafu { region_id: self.region_id, diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index 42c1df23f07..fe624316357 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -358,12 +358,214 @@ async fn test_find_compaction_options_priority() { } } +#[tokio::test] +async fn test_resident_region_survives_terminal_cycles_and_truncate() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(8); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version = Arc::new(VersionControlBuilder::new().build()); + let id = version.region_id(); + scheduler.register_region(&version, &env.access_layer); + assert!(!scheduler.is_compacting(id)); + assert!(scheduler.region_status.contains_key(&id)); + for cancelled in [false, true] { + let status = scheduler.region_status.get_mut(&id).unwrap(); + status.start_picking(7); + let (reply, response) = oneshot::channel(); + status.merge_waiter(reply.into()); + scheduler.register_region(&version, &env.access_layer); + assert!( + scheduler.is_compacting(id), + "duplicate registration must retain the active cycle" + ); + if cancelled { + assert!(scheduler.on_compaction_cancelled(id).await.is_empty()); + } else { + scheduler.on_compaction_failed(id, Arc::new(InvalidSchedulerStateSnafu.build())); + } + assert!(response.await.unwrap().is_err()); + assert!(!scheduler.is_compacting(id)); + assert!(Arc::ptr_eq( + &scheduler.region_status[&id].version_control, + &version + )); + } + scheduler.on_region_truncated(id); + assert!(scheduler.region_status.contains_key(&id)); + assert!(!scheduler.is_compacting(id)); + scheduler.on_region_closed(id); + assert!(!scheduler.region_status.contains_key(&id)); + scheduler.register_region(&version, &env.access_layer); + let (reply, response) = oneshot::channel(); + let status = scheduler.region_status.get_mut(&id).unwrap(); + status.start_picking(8); + status.merge_waiter(reply.into()); + let replacement = Arc::new(VersionControlBuilder::new().build()); + scheduler.register_region(&replacement, &env.access_layer); + assert!(response.await.unwrap().is_err()); + assert!(Arc::ptr_eq( + &scheduler.region_status[&id].version_control, + &replacement + )); + assert!(!scheduler.is_compacting(id)); + scheduler.on_region_dropped(id); + assert!(!scheduler.region_status.contains_key(&id)); +} + +/// Reads the real completion event produced by a wrapped DDL reply. +async fn recv_ddl_completion(rx: &mut mpsc::Receiver) -> (u64, usize) { + let request = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + let WorkerRequest::Background { + notify: + BackgroundNotify::CompactionDdlComplete { + generation, + request_index, + }, + .. + } = request.request + else { + panic!("expected DDL completion"); + }; + (generation, request_index) +} + +#[tokio::test] +async fn test_resident_region_ddl_handoff_and_late_replies() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(8); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version = Arc::new(VersionControlBuilder::new().build()); + let id = version.region_id(); + let manifest = env + .mock_manifest_context(version.current().version.metadata.clone()) + .await; + let (schemas, _) = mock_schema_metadata_manager(); + scheduler.register_region(&version, &env.access_layer); + scheduler + .region_status + .get_mut(&id) + .unwrap() + .start_picking(7); + let mut responses = Vec::new(); + for _ in 0..2 { + let (reply, response) = oneshot::channel(); + scheduler + .try_cancel_and_add_ddl(id, reply.into(), (), |_| { + DdlRequest::Truncate(store_api::region_request::RegionTruncateRequest::All) + }) + .unwrap(); + responses.push(response); + } + let mut ddls = scheduler.on_compaction_cancelled(id).await; + assert_eq!(2, ddls.len()); + assert!(!scheduler.is_compacting(id)); + assert!(scheduler.region_status[&id].has_ddl()); + assert!(scheduler.region_status[&id].pending_ddl_requests.is_empty()); + assert!( + !scheduler + .schedule_automatic_compaction( + Options::Regular(Default::default()), + &version, + &env.access_layer, + &manifest, + schemas.clone() + ) + .unwrap() + ); + let (reply, response) = oneshot::channel(); + assert!( + !scheduler + .schedule_manual_compaction( + Options::Regular(Default::default()), + &version, + &env.access_layer, + reply.into(), + &manifest, + schemas.clone(), + 1, + None + ) + .unwrap() + ); + assert_matches!( + response.await.unwrap(), + Err(Error::CompactionCancelled { .. }) + ); + + ddls.remove(0).sender.send(Ok(0)); + let (old_generation, first) = recv_ddl_completion(&mut rx).await; + scheduler.on_ddl_complete(id, old_generation, first); + scheduler.on_ddl_complete(id, old_generation, first); + scheduler.on_region_truncated(id); + scheduler.register_region(&version, &env.access_layer); + assert!( + scheduler.region_status[&id].has_ddl(), + "duplicate replies and truncate must not end the batch" + ); + ddls.pop().unwrap().sender.send(Ok(0)); + let (_, late_reply) = recv_ddl_completion(&mut rx).await; + for response in responses { + assert_eq!(0, response.await.unwrap().unwrap()); + } + + scheduler.on_region_closed(id); + let reopened = Arc::new(VersionControlBuilder::new().build()); + scheduler.register_region(&reopened, &env.access_layer); + scheduler + .region_status + .get_mut(&id) + .unwrap() + .start_picking(8); + let (reply, response) = oneshot::channel(); + scheduler + .try_cancel_and_add_ddl(id, reply.into(), (), |_| { + DdlRequest::Truncate(store_api::region_request::RegionTruncateRequest::All) + }) + .unwrap(); + let mut ddls = scheduler.on_compaction_cancelled(id).await; + scheduler.on_ddl_complete(id, old_generation, late_reply); + assert!( + scheduler.region_status[&id].has_ddl(), + "old incarnation must not clear the new batch" + ); + ddls.pop() + .unwrap() + .sender + .send(InvalidSchedulerStateSnafu.fail()); + let (generation, request_index) = recv_ddl_completion(&mut rx).await; + scheduler.on_ddl_complete(id, generation, request_index); + scheduler.on_ddl_complete(id, generation, request_index); + assert!(response.await.unwrap().is_err()); + assert!(!scheduler.region_status[&id].has_ddl()); + assert!( + scheduler + .schedule_automatic_compaction( + Options::Regular(Default::default()), + &reopened, + &env.access_layer, + &manifest, + schemas.clone() + ) + .unwrap() + ); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest, schemas) + .await; + assert!(!scheduler.is_compacting(id)); + assert!(scheduler.region_status.contains_key(&id)); +} + #[tokio::test] async fn test_schedule_empty() { let env = SchedulerEnv::new().await; let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); schema_metadata_manager .register_region_table_info( @@ -402,10 +604,12 @@ async fn test_schedule_empty() { .await; let output = output_rx.await.unwrap().unwrap(); assert_eq!(output, 0); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); // Only one file, picker won't compact it. let version_control = Arc::new(builder.push_l0_file(0, 1000).build()); + scheduler.register_region(&version_control, &env.access_layer); let (output_tx, output_rx) = oneshot::channel(); let waiter = OptionOutputTx::from(output_tx); let scheduled = scheduler @@ -428,7 +632,8 @@ async fn test_schedule_empty() { .await; let output = output_rx.await.unwrap().unwrap(); assert_eq!(output, 0); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -484,7 +689,7 @@ async fn test_schedule_compaction_returns_true_when_task_scheduled() { .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) .await; assert_eq!(1, job_scheduler.num_jobs()); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); } #[tokio::test] @@ -521,7 +726,7 @@ async fn test_planning_followup_reports_automatic_schedule() { } #[tokio::test] -async fn test_planning_panic_notifies_and_clears_status() { +async fn test_planning_panic_notifies_and_returns_region_to_idle() { let env = SchedulerEnv::new().await; let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx.clone()); @@ -555,7 +760,8 @@ async fn test_planning_panic_notifies_and_clears_status() { assert!(pending_ddls.is_empty()); assert!(waiter_rx.await.unwrap().is_err()); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -613,14 +819,11 @@ async fn test_ddl_fence_prevents_repeated_regular_followups() { schema_metadata_manager.clone(), ) .await; - assert!(pending_ddls.is_empty()); - - let mut followup_finished = tokio::time::timeout( - Duration::from_secs(5), - recv_compaction_pick_finished(&mut rx), - ) - .await - .expect("pre-fence regular follow-up was not planned"); + assert_eq!(pending_ddls.len(), 1); + assert!( + rx.try_recv().is_err(), + "DDL must precede even an earlier pending manual request" + ); let (post_fence_tx, post_fence_rx) = oneshot::channel(); assert!( !scheduler @@ -641,13 +844,9 @@ async fn test_ddl_fence_prevents_repeated_regular_followups() { Error::CompactionCancelled { .. } ); - followup_finished.result = CompactionPlanningResult::NoPlan; - let pending_ddls = scheduler - .handle_compaction_pick_finished(followup_finished, &manifest_ctx, schema_metadata_manager) - .await; - assert_eq!(pending_ddls.len(), 1); - assert_eq!(pre_fence_rx.await.unwrap().unwrap(), 0); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(pre_fence_rx.await.unwrap().is_err()); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert!(rx.try_recv().is_err()); } @@ -776,7 +975,8 @@ async fn test_pick_result_local_submission_failure_releases_and_notifies_once() assert!(waiter_rx.await.unwrap().is_err()); assert!(selected.iter().all(|file| !file.compacting())); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -806,7 +1006,8 @@ async fn test_pick_result_remote_submission_failure_releases_and_notifies_once() assert!(waiter_rx.await.unwrap().is_err()); assert!(selected.iter().all(|file| !file.compacting())); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -882,13 +1083,13 @@ async fn test_stale_plan_execution_does_not_affect_replacement_status() { ) .await; assert!(pending_ddls.is_empty()); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); scheduler.on_execution_failed( region_id, &stale_execution, Arc::new(InvalidSchedulerStateSnafu.build()), ); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); assert_matches!( waiter_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty) @@ -955,7 +1156,8 @@ async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { assert_eq!(output_rx.await.unwrap().unwrap(), 0); assert_eq!(rejected_before + 1, rejected.get()); assert_eq!(0, job_scheduler.num_jobs()); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert!(selected.iter().all(|file| !file.compacting())); } @@ -1039,7 +1241,7 @@ async fn test_execution_finished_drains_until_no_plan() { ) .await; assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); let finished = recv_compaction_pick_finished(&mut rx).await; scheduler .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) @@ -1069,12 +1271,13 @@ async fn test_execution_finished_drains_until_no_plan() { .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) .await; assert_matches!(transition, CompactionTransition::NoAction); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] -async fn test_execution_finished_without_progress_removes_status() { +async fn test_execution_finished_without_progress_keeps_idle_status() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); @@ -1136,7 +1339,8 @@ async fn test_execution_finished_without_progress_removes_status() { ) .await; assert_matches!(transition, CompactionTransition::NoAction); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(1, job_scheduler.num_jobs()); } @@ -1672,7 +1876,8 @@ async fn test_compaction_bypass_in_staging_mode() { let result = rx.await.unwrap(); assert_eq!(result.unwrap(), 0); - assert_eq!(0, scheduler.region_status.len()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(version_control.region_id())); } #[tokio::test] @@ -1851,6 +2056,10 @@ async fn test_try_cancel_and_add_ddl_returns_request_when_not_running() { let (tx, _rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let region_id = RegionId::new(1, 1); + let version = Arc::new(VersionControlBuilder::new().build()); + scheduler.register_region(&version, &env.access_layer); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); let (ddl_tx, ddl_rx) = oneshot::channel(); let result = @@ -1940,7 +2149,8 @@ async fn test_on_compaction_cancelled_returns_pending_ddl_requests() { assert_eq!(pending_ddls.len(), 1); assert!(!scheduler.has_pending_ddls(region_id)); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(job_scheduler.num_jobs(), 0); } @@ -1987,7 +2197,8 @@ async fn test_on_compaction_cancelled_prioritizes_pending_ddls_over_pending_comp let pending_ddls = scheduler.on_compaction_cancelled(region_id).await; assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(job_scheduler.num_jobs(), 0); assert_matches!(manual_rx.await.unwrap(), Err(_)); } @@ -2173,7 +2384,8 @@ async fn test_on_compaction_finished_returns_pending_ddl_requests() { assert_eq!(pending_ddls.len(), 1); assert!(!scheduler.has_pending_ddls(region_id)); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(job_scheduler.num_jobs(), 0); } @@ -2225,14 +2437,11 @@ async fn test_planning_terminal_prioritizes_pending_ddl_over_automatic_followup( ) .await; - assert!(pending_ddls.is_empty()); - let finished = recv_compaction_pick_finished(&mut rx).await; - let pending_ddls = scheduler - .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) - .await; assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert_eq!(manual_rx.await.unwrap().unwrap(), 0); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); + assert!(rx.try_recv().is_err()); + assert!(manual_rx.await.unwrap().is_err()); } #[tokio::test] @@ -2275,7 +2484,8 @@ async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regul .await; assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert!(rx.try_recv().is_err()); } @@ -2300,7 +2510,7 @@ async fn test_on_compaction_finished_returns_empty_when_region_absent() { } #[tokio::test] -async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { +async fn test_on_compaction_finished_ddl_bypasses_failing_manual_scheduler() { let env = SchedulerEnv::new() .await .scheduler(Arc::new(FailingScheduler)); @@ -2356,15 +2566,16 @@ async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { ) .await; - assert!(pending_ddls.is_empty()); - let finished = recv_compaction_pick_finished(&mut rx).await; - let pending_ddls = scheduler - .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) - .await; - assert!(pending_ddls.is_empty()); - assert!(!scheduler.region_status.contains_key(®ion_id)); + let super::CompactionTransition::DdlReady(mut ddls) = pending_ddls else { + panic!("DDL not released") + }; + assert_eq!(ddls.len(), 1); + assert!(rx.try_recv().is_err()); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_matches!(manual_rx.await.unwrap(), Err(_)); - assert_matches!(ddl_rx.await.unwrap(), Err(_)); + ddls.pop().unwrap().sender.send(Ok(0)); + assert_eq!(ddl_rx.await.unwrap().unwrap(), 0); } #[tokio::test] diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 092abc296a6..f20eef75460 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -1097,6 +1097,61 @@ impl EventListener for CompactionPlanningGate { } } +#[tokio::test] +async fn test_idle_region_accepts_ddl_after_create_and_reopen() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(45, 1); + let create = CreateRequestBuilder::new().build(); + let table_dir = create.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + for reopen in [false, true] { + if reopen { + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: table_dir.clone(), + path_type: PathType::Bare, + options: Default::default(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + } + for request in [ + RegionRequest::Compact(RegionCompactRequest::default()), + RegionRequest::Truncate(RegionTruncateRequest::All), + RegionRequest::Compact(RegionCompactRequest::default()), + ] { + tokio::time::timeout( + Duration::from_secs(5), + engine.handle_request(region_id, request), + ) + .await + .expect("idle registration must not stall DDL or the next compaction") + .unwrap(); + } + } +} + #[tokio::test] async fn test_planning_followup_updates_schedule_time() { assert_automatic_followup_updates_schedule_time(0).await; diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index 432e3367603..bbfdc97afd6 100644 --- a/src/mito2/src/engine/listener.rs +++ b/src/mito2/src/engine/listener.rs @@ -28,6 +28,9 @@ use crate::sst::file::RegionFileId; /// Mito engine background event listener. #[async_trait] pub trait EventListener: Send + Sync { + /// Notifies the background opener after the worker acknowledges region registration. + async fn on_region_open_registered(&self, _region_id: RegionId) {} + /// Notifies the listener that a region is flushed successfully. fn on_flush_success(&self, region_id: RegionId) { let _ = region_id; diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index c29c9aa3915..126dde56697 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -31,11 +31,12 @@ use store_api::region_request::{ RegionRequest, }; use store_api::storage::{RegionId, ScanRequest}; -use tokio::sync::oneshot; +use tokio::sync::{Notify, oneshot}; use crate::compaction::compactor::{OpenCompactionRegionRequest, open_compaction_region}; use crate::config::MitoConfig; use crate::engine::flush_test::MockRegionHook; +use crate::engine::listener::EventListener; use crate::engine::region_hook::RegionHookRef; use crate::error; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; @@ -763,6 +764,112 @@ async fn test_open_region_skip_wal_replay_with_format(flat_format: bool) { assert_eq!(expected, batches.pretty_print().unwrap()); } +/// Holds the first background opener after registration without blocking the worker. +#[derive(Default)] +struct OpenRegistrationGate { + registrations: AtomicUsize, + entered: Notify, + resume: Notify, +} + +#[async_trait::async_trait] +impl EventListener for OpenRegistrationGate { + async fn on_region_open_registered(&self, _region_id: RegionId) { + if self.registrations.fetch_add(1, Ordering::Relaxed) == 0 { + self.entered.notify_one(); + self.resume.notified().await; + } + } +} + +/// Releases the opener even if a regression assertion fails. +struct OpenRegistrationGateGuard(Arc); + +impl Drop for OpenRegistrationGateGuard { + fn drop(&mut self) { + self.0.resume.notify_one(); + } +} + +#[tokio::test] +async fn test_reopen_does_not_wait_for_previous_open_task() { + let mut env = TestEnv::with_prefix("reopen-before-open-task-exits").await; + let gate = Arc::new(OpenRegistrationGate::default()); + let engine = env + .create_engine_with( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + None, + Some(gate.clone()), + None, + ) + .await; + let region_id = RegionId::new(1, 1); + let create = CreateRequestBuilder::new().build(); + let open = RegionOpenRequest { + engine: String::new(), + table_dir: create.table_dir.clone(), + path_type: PathType::Bare, + options: Default::default(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }; + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + + let guard = OpenRegistrationGateGuard(gate.clone()); + let first_engine = engine.clone(); + let first_request = open.clone(); + let first_open = tokio::spawn(async move { + first_engine + .handle_request(region_id, RegionRequest::Open(first_request)) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.entered.notified()) + .await + .expect("the first opener did not receive its registration acknowledgement"); + assert!(engine.is_region_exists(region_id)); + + // A later open must load the region again, not join the first opener's waiters. + tokio::time::timeout(Duration::from_secs(5), async { + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + assert!(!engine.is_region_exists(region_id)); + engine + .handle_request(region_id, RegionRequest::Open(open)) + .await + .unwrap(); + }) + .await + .expect("reopening must not wait for the previous background opener to exit"); + assert!(engine.is_region_exists(region_id)); + assert!(!engine.is_region_opening(region_id)); + + drop(guard); + tokio::time::timeout(Duration::from_secs(5), first_open) + .await + .unwrap() + .unwrap() + .unwrap(); +} + #[tokio::test] async fn test_open_region_wait_for_opening_region_ok() { test_open_region_wait_for_opening_region_ok_with_format(false).await; diff --git a/src/mito2/src/engine/truncate_test.rs b/src/mito2/src/engine/truncate_test.rs index 61dbc062e78..856644ca12d 100644 --- a/src/mito2/src/engine/truncate_test.rs +++ b/src/mito2/src/engine/truncate_test.rs @@ -42,6 +42,7 @@ struct FlushCommitListener { commit_started: Notify, cancel_requested: Notify, resume_flush: Notify, + compaction_scheduled: AtomicUsize, } impl FlushCommitListener { @@ -68,6 +69,10 @@ impl EventListener for FlushCommitListener { fn on_flush_cancel_requested(&self, _region_id: RegionId) { self.cancel_requested.notify_one(); } + + fn on_compaction_scheduled(&self, _region_id: RegionId) { + self.compaction_scheduled.fetch_add(1, Ordering::Relaxed); + } } #[derive(Default)] @@ -899,6 +904,101 @@ async fn test_engine_discard_unflushed_after_flush_commit_started_with_format(fl ); } +#[tokio::test] +async fn test_engine_truncate_all_after_flush_commit_skips_compaction() { + test_engine_truncate_all_after_flush_commit_skips_compaction_with_format(false).await; + test_engine_truncate_all_after_flush_commit_skips_compaction_with_format(true).await; +} + +async fn test_engine_truncate_all_after_flush_commit_skips_compaction_with_format( + flat_format: bool, +) { + let mut env = TestEnv::with_prefix("truncate-all-after-flush-commit").await; + let listener = Arc::new(FlushCommitListener::default()); + let engine = env + .create_engine_with( + MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 3), + }, + ) + .await; + + let flush_engine = engine.clone(); + let flush_task = tokio::spawn(async move { + flush_engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), listener.wait_commit_started()) + .await + .expect("flush did not reach the commit gate"); + + // The region is idle and the flush is past cancellation, so the truncate + // queues behind the flush and only runs after the flush commits. + let truncate_engine = engine.clone(); + let truncate_task = tokio::spawn(async move { + truncate_engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::All), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), listener.wait_cancel_requested()) + .await + .expect("truncate did not queue behind the committing flush"); + assert!(!truncate_task.is_finished()); + + listener.resume_flush(); + tokio::time::timeout(Duration::from_secs(5), flush_task) + .await + .expect("flush did not finish") + .expect("flush task panicked") + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), truncate_task) + .await + .expect("truncate did not finish") + .expect("truncate task panicked") + .unwrap(); + + assert_eq!( + 0, + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .num_files() + ); + // The committing flush dispatches the queued truncate, which puts the region + // into `Truncating` before `handle_flush_finished` schedules compaction. No + // compaction may be scheduled while the region is truncating. + assert_eq!(0, listener.compaction_scheduled.load(Ordering::Relaxed)); +} + async fn test_engine_truncate_during_flush_with_format(flat_format: bool, unflushed_only: bool) { init_default_ut_logging(); let mut env = TestEnv::with_prefix("truncate-during-flush").await; diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 521a0e08a21..d4753fe14bf 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -377,6 +377,11 @@ impl MitoRegion { == RegionRoleState::Leader(RegionLeaderState::EnteringStaging) } + /// Returns whether the region is truncating. + pub(crate) fn is_truncating(&self) -> bool { + self.manifest_ctx.state.load() == RegionRoleState::Leader(RegionLeaderState::Truncating) + } + pub fn region_id(&self) -> RegionId { self.region_id } diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 10002570008..cd2902f68dd 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -58,6 +58,7 @@ use crate::manifest::action::{RegionEdit, TruncateKind}; use crate::memtable::MemtableId; use crate::memtable::bulk::part::BulkPart; use crate::metrics::COMPACTION_ELAPSED_TOTAL; +use crate::region::MitoRegionRef; use crate::region::options::RegionOptions; use crate::sst::file::FileMeta; use crate::sst::index::IndexBuildType; @@ -899,6 +900,16 @@ pub(crate) struct SenderDdlRequest { /// Notification from a background job. #[derive(Debug)] pub(crate) enum BackgroundNotify { + /// Registers a loaded region and its scheduling state before acknowledging open. + RegionOpened { + region: MitoRegionRef, + registered: Sender<()>, + }, + /// A DDL released by compaction has actually completed, including failure. + CompactionDdlComplete { + generation: u64, + request_index: usize, + }, /// Compaction planning has finished. CompactionPickFinished(CompactionPickFinished), /// Flush has finished. diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 2eda4fff49f..7fca91cf3a3 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1010,6 +1010,13 @@ struct RegionWorkerLoop { } impl RegionWorkerLoop { + /// Publishes a loaded region together with its resident compaction state. + fn register_region(&mut self, region: MitoRegionRef) { + self.compaction_scheduler + .register_region(®ion.version_control, ®ion.access_layer); + self.regions.insert_region(region); + } + /// Starts the worker loop. async fn run(&mut self) { let init_check_delay = worker_init_check_delay(); @@ -1338,6 +1345,28 @@ impl RegionWorkerLoop { /// Handles region background request async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) { match notify { + BackgroundNotify::RegionOpened { region, registered } => { + self.register_region(region); + self.region_count.inc(); + + // Finish this opening batch before the worker can process a close/reopen, + // without waiting for the background opener to resume after the ACK. + let senders = self.opening_regions.remove_sender(region_id); + if let Some(state) = &self.series_index_task_state { + state.wake(); + } + let _ = registered.send(()); + for sender in senders { + sender.send(Ok(0)); + } + } + BackgroundNotify::CompactionDdlComplete { + generation, + request_index, + } => { + self.compaction_scheduler + .on_ddl_complete(region_id, generation, request_index); + } BackgroundNotify::CompactionPickFinished(req) => { self.handle_compaction_pick_finished(region_id, req).await } @@ -1448,6 +1477,14 @@ impl WorkerListener { WorkerListener { listener } } + /// Notifies the background opener after the worker acknowledges region registration. + pub(crate) async fn on_region_open_registered(&self, _region_id: RegionId) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_region_open_registered(_region_id).await; + } + } + /// Flush is finished successfully. pub(crate) fn on_flush_success(&self, region_id: RegionId) { #[cfg(any(test, feature = "test"))] diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index 576ffed7ba2..81838060385 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -133,7 +133,7 @@ impl RegionWorkerLoop { .open(&self.config, &self.wal) .await?; debug_assert!(!reopened_region.is_writable()); - self.regions.insert_region(reopened_region.clone()); + self.register_region(reopened_region.clone()); if let Some(state) = &self.series_index_task_state { state.wake(); } diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index ad7872b7237..238c6c216da 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -45,8 +45,8 @@ impl RegionWorkerLoop { let Some(region) = self.regions.get_region(region_id) else { return; }; - // A terminal pick (canceled, no plan, or failed) may remove the - // compaction status and release DDLs fenced behind its picking phase. + // A terminal pick (canceled, no plan, or failed) may return the region + // to idle and release DDLs fenced behind its picking phase. // Such a pick never produces an execution callback, so the worker must // execute the returned DDLs from this notification. let transition = self @@ -242,9 +242,9 @@ impl RegionWorkerLoop { /// Schedule compaction for the region if necessary. pub(crate) async fn schedule_compaction(&mut self, region: &MitoRegionRef) { - if region.is_staging() || region.is_enter_staging() { + if region.is_staging() || region.is_enter_staging() || region.is_truncating() { info!( - "Region {} is staging or entering staging, skip compaction", + "Region {} is staging, entering staging, or truncating, skip compaction", region.region_id ); return; diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 2263cfa2ffb..cff77b50f5f 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -95,7 +95,7 @@ impl RegionWorkerLoop { } // Insert the MitoRegion into the RegionMap. - self.regions.insert_region(region); + self.register_region(region); if let Some(state) = &self.series_index_task_state { state.wake(); } diff --git a/src/mito2/src/worker/handle_drop.rs b/src/mito2/src/worker/handle_drop.rs index 8112d6a1592..a62e8f0068e 100644 --- a/src/mito2/src/worker/handle_drop.rs +++ b/src/mito2/src/worker/handle_drop.rs @@ -115,6 +115,11 @@ where // Removes this region from region map to prevent other requests from accessing this region self.regions.remove_region(region_id); self.dropping_regions.insert_region(region.clone()); + // The region is no longer served once it leaves the map, so unregister its + // runtime state before the fallible WAL cleanup below. Otherwise a WAL + // failure would leak the resident compaction/flush/index entries that only + // this cleanup and close/drop remove. + self.cleanup_dropped_region_runtime_state(region_id).await; // Delete region data in WAL. self.wal @@ -124,7 +129,6 @@ where ®ion.provider, ) .await?; - self.cleanup_dropped_region_runtime_state(region_id).await; if let Some(store) = &self.series_index_store { crate::series_index::delete_catalogs(store, region_id).await; diff --git a/src/mito2/src/worker/handle_enter_staging.rs b/src/mito2/src/worker/handle_enter_staging.rs index f9e24541241..e5b0674aaac 100644 --- a/src/mito2/src/worker/handle_enter_staging.rs +++ b/src/mito2/src/worker/handle_enter_staging.rs @@ -64,6 +64,24 @@ impl RegionWorkerLoop { return; } + // Retain the compaction fence while the DDL later waits for flush. + let (sender, partition_directive) = match self.compaction_scheduler.try_cancel_and_add_ddl( + region_id, + sender, + partition_directive, + |partition_directive| { + DdlRequest::EnterStaging(EnterStagingRequest { + partition_directive, + }) + }, + ) { + Ok(()) => { + self.listener.on_compaction_cancel_requested(region_id); + return; + } + Err(request) => request, + }; + let version = region.version(); if !version.memtables.is_empty() { // If memtable is not empty, we can't enter staging directly and need to flush @@ -98,23 +116,6 @@ impl RegionWorkerLoop { return; } - let (sender, partition_directive) = match self.compaction_scheduler.try_cancel_and_add_ddl( - region_id, - sender, - partition_directive, - |partition_directive| { - DdlRequest::EnterStaging(EnterStagingRequest { - partition_directive, - }) - }, - ) { - Ok(()) => { - self.listener.on_compaction_cancel_requested(region_id); - return; - } - Err(request) => request, - }; - self.handle_enter_staging(region, partition_directive, sender); } diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index 52b902ffe92..b734d7e9b86 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -35,7 +35,7 @@ use crate::region::opener::{ RegionOpener, get_object_store, provider_from_wal_options, sanitize_open_request_options, }; use crate::region::options::RegionOptions; -use crate::request::OptionOutputTx; +use crate::request::{BackgroundNotify, OptionOutputTx, WorkerRequest, WorkerRequestWithTime}; use crate::sst::location::region_dir_from_table_dir; use crate::wal::entry_distributor::WalEntryReceiver; use crate::worker::handle_drop::{ @@ -221,13 +221,12 @@ impl RegionWorkerLoop { } let now = Instant::now(); - let regions = self.regions.clone(); + let request_sender = self.sender.clone(); let wal = self.wal.clone(); let config = self.config.clone(); let opening_regions = self.opening_regions.clone(); - let region_count = self.region_count.clone(); let worker_id = self.id; - let series_index_task_state = self.series_index_task_state.clone(); + let listener = self.listener.clone(); opening_regions.insert_sender(region_id, sender); common_runtime::spawn_global(async move { match opener.open(&config, &wal).await { @@ -239,8 +238,6 @@ impl RegionWorkerLoop { worker_id, now.elapsed() ); - region_count.inc(); - // Notify the region hook that the region has been opened. // Fires before registration; allocates nothing when no hook // is registered. @@ -248,15 +245,27 @@ impl RegionWorkerLoop { hook.on_region_opened(region_id, ®ion.metadata()).await; } - // Insert the Region into the RegionMap. - regions.insert_region(region); - if let Some(state) = &series_index_task_state { - state.wake(); - } - - let senders = opening_regions.remove_sender(region_id); - for sender in senders { - sender.send(Ok(0)); + // The worker installs both the region and its resident scheduler state. + let (registered, receiver) = tokio::sync::oneshot::channel(); + let sent = request_sender + .send(WorkerRequestWithTime::new(WorkerRequest::Background { + region_id, + notify: BackgroundNotify::RegionOpened { + region: region.clone(), + registered, + }, + })) + .await + .is_ok(); + let accepted = sent && receiver.await.is_ok(); + if accepted { + listener.on_region_open_registered(region_id).await; + } else { + region.stop().await; + let senders = opening_regions.remove_sender(region_id); + for sender in senders { + sender.send(crate::error::RegionClosedSnafu { region_id }.fail()); + } } } Err(err) => { diff --git a/src/mito2/src/worker/handle_truncate.rs b/src/mito2/src/worker/handle_truncate.rs index d62bf8369a4..409998592d1 100644 --- a/src/mito2/src/worker/handle_truncate.rs +++ b/src/mito2/src/worker/handle_truncate.rs @@ -40,19 +40,7 @@ impl RegionWorkerLoop { } }; - let (sender, req) = match self.flush_scheduler.try_cancel_and_add_ddl( - region_id, - sender, - req, - DdlRequest::Truncate, - ) { - Ok(()) => { - self.listener.on_flush_cancel_requested(region_id); - return; - } - Err(request) => request, - }; - + // Fence compaction before any later flush wait so handed-off DDLs keep admission paused. let (sender, req) = match self.compaction_scheduler.try_cancel_and_add_ddl( region_id, sender, @@ -66,6 +54,19 @@ impl RegionWorkerLoop { Err(request) => request, }; + let (sender, req) = match self.flush_scheduler.try_cancel_and_add_ddl( + region_id, + sender, + req, + DdlRequest::Truncate, + ) { + Ok(()) => { + self.listener.on_flush_cancel_requested(region_id); + return; + } + Err(request) => request, + }; + let version_data = region.version_control.current(); match req { From d86741d0ad92b98cff2cc0cff101734fd02e270a Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 06:02:43 +0000 Subject: [PATCH 27/46] feat(operator): retry failed writes with refreshed partitions (#9542) * feat(operator): retry failed writes with refreshed partitions Signed-off-by: WenyXu * style(operator): simplify retry test type references Signed-off-by: WenyXu * refactor(operator): clarify write request rebuilding names Signed-off-by: WenyXu * refactor(operator): encapsulate write retry attempts Signed-off-by: WenyXu * style(operator): group inserter configuration methods Signed-off-by: WenyXu * refactor(operator): share routed peers and separate retry decisions Signed-off-by: WenyXu * refactor(operator): avoid temporary write routing allocations Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 4 + src/common/meta/src/peer.rs | 3 + src/common/meta/src/rpc/router.rs | 33 +- src/operator/AGENTS.md | 4 +- src/operator/Cargo.toml | 4 + src/operator/src/bulk_insert.rs | 47 +- src/operator/src/delete.rs | 13 +- src/operator/src/insert.rs | 25 +- src/operator/src/req_convert.rs | 1 - .../src/req_convert/common/partitioner.rs | 454 ++++++++- src/operator/src/write_dispatcher.rs | 451 +++++++-- src/operator/src/write_dispatcher/options.rs | 106 ++ src/operator/src/write_dispatcher/tests.rs | 930 +++++++++++++++++- 13 files changed, 1958 insertions(+), 117 deletions(-) create mode 100644 src/operator/src/write_dispatcher/options.rs diff --git a/Cargo.lock b/Cargo.lock index f75bde16328..8678104f5d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10095,7 +10095,9 @@ dependencies = [ "async-stream", "async-trait", "axum 0.8.4", + "backon", "bytes", + "cache", "catalog", "chrono", "client", @@ -10126,6 +10128,7 @@ dependencies = [ "futures", "futures-util", "humantime", + "humantime-serde", "itertools 0.14.0", "jsonb", "lazy_static", @@ -10152,6 +10155,7 @@ dependencies = [ "substrait 1.4.0-alpha.0", "table", "tokio", + "tokio-stream", "tokio-util", "tonic 0.14.2", "tracing", diff --git a/src/common/meta/src/peer.rs b/src/common/meta/src/peer.rs index 924c293bbe0..18ffd92f8bb 100644 --- a/src/common/meta/src/peer.rs +++ b/src/common/meta/src/peer.rs @@ -22,6 +22,9 @@ use crate::cluster::NodeInfo; use crate::error::Error; use crate::{DatanodeId, FlownodeId}; +/// Shared reference to a [`Peer`]. +pub type PeerRef = Arc; + /// [`PeerResolver`] provides methods to look up peer information by node ID. /// /// This trait allows resolving both datanode and flownode peers, regardless of their current activity status. diff --git a/src/common/meta/src/rpc/router.rs b/src/common/meta/src/rpc/router.rs index 20fcd955abb..e089318b272 100644 --- a/src/common/meta/src/rpc/router.rs +++ b/src/common/meta/src/rpc/router.rs @@ -114,15 +114,25 @@ pub fn convert_to_region_leader_map(region_routes: &[RegionRoute]) -> HashMap>() } +/// Returns the leader of the first route matching the region number. pub fn find_region_leader( region_routes: &[RegionRoute], region_number: RegionNumber, ) -> Option { + find_region_leader_ref(region_routes, region_number).cloned() +} + +/// Borrows the leader of the first route matching the region number. +/// +/// Returns `None` if no route matches or the first matching route has no leader. +pub fn find_region_leader_ref( + region_routes: &[RegionRoute], + region_number: RegionNumber, +) -> Option<&Peer> { region_routes .iter() .find(|x| x.region.id.region_number() == region_number) .and_then(|r| r.leader_peer.as_ref()) - .cloned() } /// Returns the region numbers of the leader regions on the target datanode. @@ -442,6 +452,27 @@ mod tests { } } + #[test] + fn test_find_region_leader() { + let mut routes = vec![ + new_test_region_route(RegionId::new(1, 1)), + new_test_region_route(RegionId::new(1, 1)), + ]; + routes[1].leader_peer = Some(Peer::new(2, "a2")); + + assert_eq!( + find_region_leader_ref(&routes, 1), + routes[0].leader_peer.as_ref() + ); + assert_eq!(find_region_leader(&routes, 1), routes[0].leader_peer); + assert_eq!(find_region_leader_ref(&routes, 2), None); + assert_eq!(find_region_leader(&routes, 2), None); + + routes[0].leader_peer = None; + assert_eq!(find_region_leader_ref(&routes, 1), None); + assert_eq!(find_region_leader(&routes, 1), None); + } + #[test] fn test_leader_is_downgraded() { let mut region_route = RegionRoute { diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 46359eb5c50..4d31f44de96 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -14,7 +14,7 @@ insert/delete routing. | Statements | `src/operator/src/statement.rs`, `src/operator/src/statement/` | `StatementExecutor` and per-statement handlers | | Inserts | `src/operator/src/insert.rs`, `src/operator/src/bulk_insert.rs` | Insert, auto-create/alter, routing, dispatch | | Deletes | `src/operator/src/delete.rs` | Delete conversion and dispatch | -| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch and batched peer resolution | +| Write dispatch | `src/operator/src/write_dispatcher.rs` | Shared insert/delete/bulk dispatch, bounded retries, route refresh and write request rebuilding | | Request conversion | `src/operator/src/req_convert.rs`, `src/operator/src/req_convert/` | Table/row/column requests to region requests | | Region requests | `src/operator/src/region_req_factory.rs` | Region-level DDL/DML request construction | | Procedures | `src/operator/src/procedure.rs` | Metasrv procedure administration | @@ -39,6 +39,8 @@ cargo sqlness bare -t ``` Check standalone and distributed routing when changing region dispatch. +Write retries operate on prepared region requests in `write_dispatcher`; keep +admission and Flow side effects outside retries, and preserve successful regions. Write admission uses fallible `write_meter!` calls in `Inserter::do_request` and `handle_bulk_insert`, before flow side effects and datanode dispatch. Bulk diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 72f953e0f89..7f4233e633e 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -24,6 +24,7 @@ api.workspace = true arrow.workspace = true async-stream.workspace = true async-trait.workspace = true +backon.workspace = true bytes.workspace = true catalog.workspace = true chrono.workspace = true @@ -54,6 +55,7 @@ file-engine.workspace = true futures.workspace = true futures-util.workspace = true humantime.workspace = true +humantime-serde.workspace = true itertools.workspace = true jsonb.workspace = true lazy_static.workspace = true @@ -86,8 +88,10 @@ url.workspace = true [dev-dependencies] axum.workspace = true +cache.workspace = true catalog = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true object-store = { workspace = true, features = ["testing"] } path-slash = "0.2" +tokio-stream = { workspace = true, features = ["net"] } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index d5aa7168b44..29537959df5 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -14,7 +14,6 @@ use std::collections::HashSet; -use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, @@ -41,6 +40,7 @@ use crate::error::Result; use crate::insert::Inserter; use crate::req_convert::TablePartitioner; use crate::req_convert::insert; +use crate::write_dispatcher::WriteAttempt; use crate::{error, metrics}; impl Inserter { @@ -72,7 +72,7 @@ impl Inserter { let writes = self.write_dispatcher.prepare(writes).await?; Ok(self .write_dispatcher - .dispatch_prepared(writes, None) + .execute(writes, None) .await .context(error::RequestInsertsSnafu)? .affected_rows) @@ -186,7 +186,7 @@ impl Inserter { let writes = self.write_dispatcher.prepare(vec![request]).await?; let result = self .write_dispatcher - .dispatch_prepared( + .execute( writes, Some( metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), @@ -246,7 +246,7 @@ impl Inserter { } else { None }; - let handle: common_runtime::JoinHandle> = + let handle: common_runtime::JoinHandle> = common_runtime::spawn_global(async move { let (header, payload) = if mask.select_all() { // SAFETY: raw data must be present, we can avoid re-encoding. @@ -301,27 +301,40 @@ impl Inserter { })), }; - dispatcher - .dispatch_prepared( - vec![(peer, request)], + Ok(dispatcher + .send( + peer, + request, Some( - metrics::HANDLE_BULK_INSERT_ELAPSED + &metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["datanode_handle"]), ), ) - .await - .context(error::RequestRegionSnafu) + .await) }); handles.push(handle); } - let results = future::try_join_all(handles) - .await - .context(error::JoinTaskSnafu)?; - let rows_inserted = results - .into_iter() - .map(|result| result.map(|response| response.affected_rows)) - .sum::>()?; + let dispatcher = self.write_dispatcher.clone(); + let rows_inserted = common_runtime::spawn_global(async move { + let results = future::try_join_all(handles) + .await + .context(error::JoinTaskSnafu)?; + let attempts = results.into_iter().collect::>>()?; + dispatcher + .finish( + attempts, + Some( + &metrics::HANDLE_BULK_INSERT_ELAPSED + .with_label_values(&["datanode_handle"]), + ), + ) + .await + .context(error::RequestRegionSnafu) + .map(|response| response.affected_rows) + }) + .await + .context(error::JoinTaskSnafu)??; wait_all_datanode_timer.observe_duration(); crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[db_name.as_str()]) diff --git a/src/operator/src/delete.rs b/src/operator/src/delete.rs index abd8302c965..c389374b7ac 100644 --- a/src/operator/src/delete.rs +++ b/src/operator/src/delete.rs @@ -141,13 +141,12 @@ impl Deleter { let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_delete(requests); let writes = dispatcher.prepare(vec![request]).await?; - let affected_rows = common_runtime::spawn_global(async move { - dispatcher.dispatch_prepared(writes, None).await - }) - .await - .context(JoinTaskSnafu)? - .context(RequestDeletesSnafu)? - .affected_rows; + let affected_rows = + common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await }) + .await + .context(JoinTaskSnafu)? + .context(RequestDeletesSnafu)? + .affected_rows; crate::metrics::DIST_DELETE_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 89c1195cda4..2c01a35b8d9 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -316,6 +316,12 @@ impl Inserter { } } + /// Installs the shared dispatcher used by direct writes and batch flushes. + pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { + self.write_dispatcher = dispatcher; + self + } + /// Installs the shared batcher; callers explicitly select its ingestion entry point. pub fn with_pending_rows_batcher( mut self, @@ -835,13 +841,12 @@ impl Inserter { let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_insert(normal_requests); let writes = dispatcher.prepare(vec![request]).await?; - let affected_rows = common_runtime::spawn_global(async move { - dispatcher.dispatch_prepared(writes, None).await - }) - .await - .context(JoinTaskSnafu)? - .context(RequestInsertsSnafu)? - .affected_rows; + let affected_rows = + common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await }) + .await + .context(JoinTaskSnafu)? + .context(RequestInsertsSnafu)? + .affected_rows; crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); @@ -1675,12 +1680,6 @@ impl Inserter { } } - /// Installs the shared dispatcher used by direct writes and batch flushes. - pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self { - self.write_dispatcher = dispatcher; - self - } - /// Returns the dispatcher shared with batching execution paths. pub fn write_dispatcher(&self) -> &WriteDispatcherRef { &self.write_dispatcher diff --git a/src/operator/src/req_convert.rs b/src/operator/src/req_convert.rs index 1161f62795d..82681e09968 100644 --- a/src/operator/src/req_convert.rs +++ b/src/operator/src/req_convert.rs @@ -16,5 +16,4 @@ pub(crate) mod common; pub mod delete; pub mod insert; -pub use common::partitioner::Partitioner; pub(crate) use common::partitioner::TablePartitioner; diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 044d3b97db5..f1c3eff4fe3 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -14,16 +14,19 @@ use std::collections::HashMap; -use api::v1::region::{BulkInsertRequest, DeleteRequest, InsertRequest, bulk_insert_request}; +use api::v1::region::{ + BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest, + bulk_insert_request, region_request::Body, +}; use api::v1::{ArrowIpc, PartitionExprVersion, Rows}; use common_error::ext::BoxedError; -use common_grpc::flight; +use common_grpc::flight::{self, FlightDecoder}; use datatypes::arrow::compute; use datatypes::arrow::record_batch::RecordBatch; use partition::PartitionRuleRef; use partition::manager::PartitionRuleManager; use partition::splitter::RowSplitter; -use snafu::ResultExt; +use snafu::{OptionExt, ResultExt}; use store_api::storage::{RegionId, RegionNumber, TableId}; use table::metadata::TableInfo; @@ -94,6 +97,37 @@ impl TablePartitioner { }) } + /// Rebuilds write requests using this partition snapshot after a route refresh. + /// + /// Takes ownership so row payloads can move into their new regions. Requests whose + /// region and partition version still match retain their original payloads. + pub(crate) fn rebuild_write_requests( + &self, + request: RegionRequest, + ) -> Result> { + let RegionRequest { header, body } = request; + let bodies = match body.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a request without a body", + })? { + Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?, + Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?, + Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk)?, + _ => { + return error::UnexpectedSnafu { + violated: "Only writes can be rebuilt", + } + .fail(); + } + }; + Ok(bodies + .into_iter() + .map(|body| RegionRequest { + header: header.clone(), + body: Some(body), + }) + .collect()) + } + /// Moves insert rows into requests for this snapshot's destination regions. pub fn partition_insert_requests( &self, @@ -178,6 +212,84 @@ impl TablePartitioner { Ok(requests) } + fn rebuild_insert_requests(&self, inserts: InsertRequests) -> Result> { + let mut bodies = Vec::new(); + for insert in inserts.requests { + let region_id = RegionId::from_u64(insert.region_id); + if self.is_partition_expr_unchanged(region_id, insert.partition_expr_version) { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![insert], + })); + continue; + } + let rows = insert.rows.context(error::UnexpectedSnafu { + violated: "Cannot rebuild an insert without rows", + })?; + for request in self.partition_insert_requests(rows, insert.skip_wal)? { + bodies.push(Body::Inserts(InsertRequests { + requests: vec![request], + })); + } + } + Ok(bodies) + } + + fn rebuild_delete_requests(&self, deletes: DeleteRequests) -> Result> { + let mut bodies = Vec::new(); + for delete in deletes.requests { + let region_id = RegionId::from_u64(delete.region_id); + if self.is_partition_expr_unchanged(region_id, delete.partition_expr_version) { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![delete], + })); + continue; + } + let rows = delete.rows.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a delete without rows", + })?; + for request in self.partition_delete_requests(rows)? { + bodies.push(Body::Deletes(DeleteRequests { + requests: vec![request], + })); + } + } + Ok(bodies) + } + + fn rebuild_bulk_insert_request(&self, bulk: BulkInsertRequest) -> Result> { + let region_id = RegionId::from_u64(bulk.region_id); + if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) { + return Ok(vec![Body::BulkInsert(bulk)]); + } + let bulk_insert_request::Body::ArrowIpc(ipc) = + bulk.body.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a bulk insert without Arrow IPC data", + })?; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .and_then(|mut decoder| decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload)) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + Ok(self + .partition_bulk_insert_requests(batch, Some(ipc), bulk.skip_wal)? + .into_iter() + .map(Body::BulkInsert) + .collect()) + } + + /// Whether the request's region still exists with the same partition expression version. + /// + /// A leader change leaves the partition unchanged, so the original payload can be + /// sent to the new peer. A removed region or changed partition expression version + /// requires rebuilding the request using this snapshot's partition rule. + fn is_partition_expr_unchanged( + &self, + region_id: RegionId, + version: Option, + ) -> bool { + region_id.table_id() == self.table_id + && self.versions.get(®ion_id.region_number()) == Some(&version) + } + fn version(&self, region: RegionNumber) -> Option { self.versions.get(®ion).copied().flatten() } @@ -187,12 +299,26 @@ impl TablePartitioner { mod tests { use std::sync::Arc; + use api::v1::region::{ + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, + }; + use api::v1::value::ValueData; + use api::v1::{ArrowIpc, PartitionExprVersion}; + use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, Value}; use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; - use common_meta::cache; + use common_grpc::flight; + use common_meta::cache::{self, TableInfoCache}; + use common_meta::key::TableMetadataManager; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::test_util::MockKvBackendBuilder; use common_meta::rpc::store::RangeResponse; + use datatypes::arrow::array::Int32Array; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; use moka::future::Cache; + use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use super::*; use crate::error::Error; @@ -249,4 +375,324 @@ mod tests { } } } + + async fn create_fixture() -> (PartitionRuleManagerRef, TableInfoCache) { + let backend = Arc::new(MemoryKvBackend::new()); + let manager = test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + test_util::new_test_table_info(42, "logical_reroute", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + let cache = cache::new_table_info_cache( + "reroute_test_table_info".to_string(), + Cache::new(16), + backend, + ); + (manager, cache) + } + + fn build_request(body: Body) -> RegionRequest { + RegionRequest { + header: Some(RegionRequestHeader { + dbname: "greptime.public".to_string(), + ..Default::default() + }), + body: Some(body), + } + } + + fn build_rows(values: &[i32]) -> Rows { + Rows { + schema: vec![ColumnSchema { + column_name: "a".to_string(), + datatype: ColumnDataType::Int32 as i32, + ..Default::default() + }], + rows: values + .iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(*value)), + }], + }) + .collect(), + } + } + + fn build_row_request( + rows: Rows, + insert: bool, + region: u32, + version: Option, + ) -> RegionRequest { + let region_id = RegionId::new(1, region).as_u64(); + let partition_expr_version = version.map(|value| PartitionExprVersion { value }); + let body = if insert { + Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id, + partition_expr_version, + rows: Some(rows), + skip_wal: true, + }], + }) + } else { + Body::Deletes(DeleteRequests { + requests: vec![DeleteRequest { + region_id, + partition_expr_version, + rows: Some(rows), + }], + }) + }; + build_request(body) + } + + fn row_request_parts(request: &RegionRequest) -> (u64, Option, &Rows) { + match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => { + assert_eq!(inserts.requests.len(), 1); + let insert = &inserts.requests[0]; + assert!(insert.skip_wal); + ( + insert.region_id, + insert.partition_expr_version.map(|v| v.value), + insert.rows.as_ref().unwrap(), + ) + } + Body::Deletes(deletes) => { + assert_eq!(deletes.requests.len(), 1); + let delete = &deletes.requests[0]; + ( + delete.region_id, + delete.partition_expr_version.map(|v| v.value), + delete.rows.as_ref().unwrap(), + ) + } + _ => panic!("expected row write"), + } + } + + #[tokio::test] + async fn test_rebuild_rows_updates_partitions_without_losing_rows() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + for insert in [true, false] { + let request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0)); + let header = request.header.clone(); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut actual = Vec::new(); + for request in rerouted { + assert_eq!(request.header, header); + let (id, version, rows) = row_request_parts(&request); + let region_id = RegionId::from_u64(id); + assert_eq!(region_id.table_id(), 1); + assert_eq!(version, versions[®ion_id.region_number()]); + assert_eq!(rows.schema, build_rows(&[]).schema); + assert_eq!(rows.rows.len(), 1); + actual.push((region_id.region_number(), rows.rows[0].clone())); + } + actual.sort_by_key(|(region, _)| *region); + assert_eq!( + actual, + vec![ + (1, build_rows(&[70]).rows.remove(0)), + (2, build_rows(&[20]).rows.remove(0)), + (3, build_rows(&[1]).rows.remove(0)) + ] + ); + } + } + + #[tokio::test] + async fn test_unchanged_row_routes_preserve_allocations() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + for insert in [true, false] { + let request = build_row_request(build_rows(&[1]), insert, 3, versions[&3]); + let original_rows = row_request_parts(&request).2; + let row_ptr = original_rows.rows.as_ptr(); + let value_ptr = original_rows.rows[0].values.as_ptr(); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 1); + let rows = row_request_parts(&rerouted[0]).2; + assert_eq!(rows.rows.as_ptr(), row_ptr); + assert_eq!(rows.rows[0].values.as_ptr(), value_ptr); + assert_eq!(rows, &build_rows(&[1])); + } + } + + fn build_bulk_request(region: u32, version: Option, values: Vec) -> RegionRequest { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap(); + build_request(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, region).as_u64(), + partition_expr_version: version.map(|value| PartitionExprVersion { value }), + aligned_schema_version: Some(AlignedSchemaVersion { + schema_version: 123, + }), + skip_wal: true, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + })) + } + + #[tokio::test] + async fn test_rebuild_bulk_reencodes_only_changed_partitions() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let (_, versions) = manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let request = build_bulk_request(99, Some(0), vec![1, 20, 70]); + let header = request.header.clone(); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut actual = Vec::new(); + for request in rerouted { + assert_eq!(request.header, header); + let Some(Body::BulkInsert(bulk)) = request.body else { + panic!("expected bulk insert") + }; + let region = RegionId::from_u64(bulk.region_id); + assert_eq!(region.table_id(), 1); + assert!(bulk.skip_wal); + assert_eq!(bulk.aligned_schema_version, None); + assert_eq!( + bulk.partition_expr_version.map(|v| v.value), + versions[®ion.region_number()] + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else { + panic!("expected Arrow IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + actual.push((region.region_number(), values.value(0))); + } + actual.sort_unstable(); + assert_eq!(actual, [(1, 70), (2, 20), (3, 1)]); + + // A leader-only change retains the original IPC buffers and schema hint. + let request = build_bulk_request(3, versions[&3], vec![1]); + let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else { + unreachable!() + }; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + let pointers = ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr(), + ); + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 1); + let Some(Body::BulkInsert(bulk)) = rerouted[0].body.as_ref() else { + unreachable!() + }; + assert_eq!( + bulk.aligned_schema_version, + Some(AlignedSchemaVersion { + schema_version: 123 + }) + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + assert_eq!( + ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr() + ), + pointers + ); + } + + #[tokio::test] + async fn test_rebuild_logical_table_preserves_logical_region_ids() { + let (manager, cache) = create_fixture().await; + for insert in [true, false] { + let mut request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0)); + match request.body.as_mut().unwrap() { + Body::Inserts(inserts) => { + inserts.requests[0].region_id = RegionId::new(42, 99).as_u64() + } + Body::Deletes(deletes) => { + deletes.requests[0].region_id = RegionId::new(42, 99).as_u64() + } + _ => unreachable!(), + } + let rerouted = rebuild_with_cached_table(&manager, &cache, request) + .await + .unwrap(); + assert_eq!(rerouted.len(), 3); + let mut regions = rerouted + .iter() + .map(|request| { + let region = RegionId::from_u64(row_request_parts(request).0); + assert_eq!(region.table_id(), 42); + region.region_number() + }) + .collect::>(); + regions.sort_unstable(); + assert_eq!(regions, [1, 2, 3]); + } + } + + async fn rebuild_with_cached_table( + manager: &PartitionRuleManager, + cache: &TableInfoCache, + request: RegionRequest, + ) -> Result> { + let region_id = match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => inserts.requests[0].region_id, + Body::Deletes(deletes) => deletes.requests[0].region_id, + Body::BulkInsert(bulk) => bulk.region_id, + _ => unreachable!(), + }; + let table_id = RegionId::from_u64(region_id).table_id(); + let table_info = cache.get(table_id).await.unwrap().unwrap(); + let partitioner = TablePartitioner::try_new(manager, &table_info) + .await + .context(error::InvalidPartitionSnafu)?; + partitioner.rebuild_write_requests(request) + } } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 4af4fbfafdf..c2ac6543411 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod options; #[cfg(test)] mod tests; @@ -20,43 +21,85 @@ use std::sync::Arc; use api::region::RegionResponse; use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body}; -use common_error::ext::BoxedError; -use common_meta::error::{ExternalSnafu, Result, UnexpectedSnafu}; +use backon::{BackoffBuilder, ExponentialBuilder}; +use client::client_manager::NodeClients; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; +use common_meta::cache::TableInfoCacheRef; +use common_meta::cache_invalidator::{CacheInvalidatorRef, Context}; +use common_meta::error::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu}; +use common_meta::instruction::CacheIdent; use common_meta::node_manager::NodeManagerRef; -use common_meta::peer::Peer; +use common_meta::peer::{Peer, PeerRef}; use common_meta::rpc::router; use futures::future; -use partition::manager::PartitionRuleManagerRef; +use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use prometheus::Histogram; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; +pub use self::options::WriteRetryOptions; use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; +use crate::req_convert::TablePartitioner; -/// Dispatches already partitioned insert, delete, and bulk requests to their peers. +/// Dispatches already partitioned writes, retaining only failed payloads for retries. pub struct WriteDispatcher { node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef, + retry: Option, } /// Shared reference to a [`WriteDispatcher`]. pub type WriteDispatcherRef = Arc; +/// The result of one write RPC attempt. +pub(crate) enum WriteAttempt { + Success(RegionResponse), + Failed { + error: Error, + /// Retained when retries are enabled and the original client error allows retry. + /// The remaining retry budget is checked separately. + request: Option, + }, +} + impl WriteDispatcher { /// Creates an owned-request dispatcher without retries, including for standalone. pub fn new(node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef) -> Self { Self { node_manager, partition_manager, + retry: None, } } + /// Enables shared remote sends and retries using frontend configuration and caches. + pub fn with_retry( + mut self, + clients: Arc, + options: WriteRetryOptions, + table_info_cache: TableInfoCacheRef, + cache_invalidator: CacheInvalidatorRef, + ) -> Self { + if options.max_retries > 0 { + self.retry = Some(WriteRetry { + clients, + partition_manager: self.partition_manager.clone(), + options, + table_info_cache, + cache_invalidator, + }); + } else { + self.retry = None; + } + self + } + /// Writes planned requests without repeating admission, default values, or Flow side effects. pub async fn dispatch(&self, requests: Vec) -> Result { self.dispatch_inner(requests, None).await } - /// Records each datanode send in the supplied histogram. + /// Records each datanode send attempt, including retries, in the supplied histogram. pub async fn dispatch_with_metrics( &self, requests: Vec, @@ -69,11 +112,17 @@ impl WriteDispatcher { pub(crate) async fn prepare( &self, requests: Vec, - ) -> error::Result> { - let requests = requests.into_iter().map(split_request).collect::>(); + ) -> error::Result> { + let requests = requests + .into_iter() + .map(|mut request| { + let header = request.header.take(); + (header, split_request(request)) + }) + .collect::>(); let regions = requests .iter() - .flatten() + .flat_map(|(_, requests)| requests) .map(write_region_id) .collect::>>() .context(RequestRegionSnafu)?; @@ -83,8 +132,8 @@ impl WriteDispatcher { .context(FindRegionLeaderSnafu)?; let mut writes = Vec::new(); // Preserve each original envelope and peer batching when retries are disabled. - for requests in requests { - let mut by_peer: HashMap = HashMap::new(); + for (header, requests) in requests { + let mut by_peer: HashMap = HashMap::new(); for request in requests { let peer = get_peer( &peers, @@ -112,27 +161,94 @@ impl WriteDispatcher { }, } } - writes.extend(by_peer); + for (peer, mut request) in by_peer { + request.header = header.clone(); + writes.push((peer, request)); + } } Ok(writes) } /// Sends requests using routes already resolved by the caller. - pub(crate) async fn dispatch_prepared( + pub(crate) async fn execute( &self, - writes: Vec<(Peer, RegionRequest)>, + writes: Vec<(PeerRef, RegionRequest)>, histogram: Option, ) -> Result { let histogram = histogram.as_ref(); - let results = future::join_all(writes.into_iter().map(|(peer, request)| async move { - let _timer = histogram.map(Histogram::start_timer); - let datanode = self.node_manager.datanode(&peer).await; - datanode.handle(request).await - })) + let writes = if self.retry.is_some() { + writes + .into_iter() + .flat_map(|(peer, request)| { + split_request(request) + .into_iter() + .map(move |request| (Arc::clone(&peer), request)) + }) + .collect() + } else { + writes + }; + let attempts = future::join_all( + writes + .into_iter() + .map(|(peer, request)| self.send(peer, request, histogram)), + ) .await; + self.finish(attempts, histogram).await + } + + /// Sends the first attempt immediately so bulk encoding can overlap other sends. + pub(crate) async fn send( + &self, + peer: PeerRef, + request: RegionRequest, + histogram: Option<&Histogram>, + ) -> WriteAttempt { + if let Some(retry) = &self.retry { + return retry + .send( + RoutedWrite { + peer, + request: Arc::new(request), + }, + histogram, + ) + .await; + } + let _timer = histogram.map(Histogram::start_timer); + let datanode = self.node_manager.datanode(&peer).await; + match datanode.handle(request).await { + Ok(response) => WriteAttempt::Success(response), + Err(error) => WriteAttempt::Failed { + error, + request: None, + }, + } + } + + /// Collects first attempts and retries their failures together, refreshing routes once per round. + pub(crate) async fn finish( + &self, + attempts: Vec, + histogram: Option<&Histogram>, + ) -> Result { + if let Some(retry) = &self.retry { + if retry.should_retry(&attempts, 0) { + return retry.retry(attempts, histogram).await; + } + let mut response = RegionResponse::new(0); + let failures = merge_attempts(attempts, &mut response)?; + return match failures.into_iter().next() { + Some((_, error)) => Err(error), + None => Ok(response), + }; + } let mut response = RegionResponse::new(0); - for result in results { - merge_response(&mut response, result?); + for attempt in attempts { + match attempt { + WriteAttempt::Success(result) => merge_response(&mut response, result), + WriteAttempt::Failed { error, .. } => return Err(error), + } } Ok(response) } @@ -141,39 +257,8 @@ impl WriteDispatcher { pub(crate) async fn resolve_peers( &self, regions: &[RegionId], - ) -> partition::error::Result> { - let table_ids = regions - .iter() - .map(|region| region.table_id()) - .collect::>(); - let mut physical_tables = HashMap::new(); - let mut routes = HashMap::new(); - for table_id in table_ids { - let (physical_id, route) = self - .partition_manager - .find_physical_table_route_with_id(table_id) - .await?; - physical_tables.insert(table_id, physical_id); - routes.entry(physical_id).or_insert(route); - } - let mut peers = HashMap::with_capacity(regions.len()); - for region in regions { - if peers.contains_key(region) { - continue; - } - let peer = physical_tables - .get(®ion.table_id()) - .and_then(|physical_id| routes.get(physical_id)) - .and_then(|route| { - router::find_region_leader(&route.region_routes, region.region_number()) - }) - .context(partition::error::FindLeaderSnafu { - region_id: *region, - table_id: region.table_id(), - })?; - peers.insert(*region, peer); - } - Ok(peers) + ) -> partition::error::Result> { + resolve_peers(&self.partition_manager, regions).await } async fn dispatch_inner( @@ -186,16 +271,270 @@ impl WriteDispatcher { .await .map_err(BoxedError::new) .context(ExternalSnafu)?; - self.dispatch_prepared(writes, histogram).await + self.execute(writes, histogram).await } } -fn get_peer(peers: &HashMap, region: RegionId) -> Result { +/// Sends shared requests and coordinates retries with refreshed routes. +struct WriteRetry { + clients: Arc, + partition_manager: PartitionRuleManagerRef, + options: WriteRetryOptions, + table_info_cache: TableInfoCacheRef, + cache_invalidator: CacheInvalidatorRef, +} + +/// A single-region write request paired with its destination peer. +pub(crate) struct RoutedWrite { + peer: PeerRef, + request: Arc, +} + +impl WriteRetry { + async fn send(&self, write: RoutedWrite, histogram: Option<&Histogram>) -> WriteAttempt { + let _timer = histogram.map(Histogram::start_timer); + let requester = self.clients.make_region_requester(&write.peer).await; + match requester.handle(Arc::clone(&write.request)).await { + Ok(response) => WriteAttempt::Success(response), + Err(error) => { + let request = error.retry_hint().is_retryable().then_some(write); + WriteAttempt::Failed { + error: error.into(), + request, + } + } + } + } + + fn should_retry(&self, attempts: &[WriteAttempt], retries: usize) -> bool { + if retries >= self.options.max_retries { + return false; + } + let mut retryable = false; + for attempt in attempts { + match attempt { + WriteAttempt::Success(_) => {} + WriteAttempt::Failed { + request: Some(_), .. + } => retryable = true, + WriteAttempt::Failed { request: None, .. } => return false, + } + } + retryable + } + + async fn retry( + &self, + mut attempts: Vec, + histogram: Option<&Histogram>, + ) -> Result { + let mut response = RegionResponse::new(0); + let mut backoff = ExponentialBuilder::default() + .with_min_delay(self.options.initial_backoff) + .with_max_delay(self.options.max_backoff) + .with_factor(self.options.backoff_multiplier) + .with_max_times(self.options.max_retries) + .with_jitter() + .build(); + let mut retries = 0; + while self.should_retry(&attempts, retries) { + let Some(delay) = backoff.next() else { + break; + }; + let failures = merge_attempts(attempts, &mut response)?; + tokio::time::sleep(delay).await; + + let mut refresh = Vec::new(); + let mut pending = Vec::new(); + for (write, error) in failures { + if error.retry_hint() == RetryHint::RetryAfterRouteRefresh { + refresh.push(write); + } else { + pending.push(write); + } + } + if !refresh.is_empty() { + pending.extend(self.reroute_failed_writes(refresh).await?); + } + attempts = + future::join_all(pending.into_iter().map(|write| self.send(write, histogram))) + .await; + retries += 1; + } + match merge_attempts(attempts, &mut response)?.into_iter().next() { + Some((_, error)) => Err(error), + None => Ok(response), + } + } + + async fn reroute_failed_writes(&self, writes: Vec) -> Result> { + let regions = writes + .iter() + .map(|write| write_region_id(&write.request)) + .collect::>>()?; + self.invalidate_routes(®ions).await?; + let mut partitioners = HashMap::new(); + for table_id in regions + .iter() + .map(|region| region.table_id()) + .collect::>() + { + let table_info = self + .table_info_cache + .get(table_id) + .await? + .with_context(|| TableInfoNotFoundSnafu { + table: table_id.to_string(), + })?; + let partitioner = TablePartitioner::try_new(&self.partition_manager, &table_info) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + partitioners.insert(table_id, partitioner); + } + let mut rebuilt_requests = Vec::new(); + for write in writes { + let request = Arc::try_unwrap(write.request).map_err(|_| { + UnexpectedSnafu { + err_msg: "Write request is still shared after its RPC completed", + } + .build() + })?; + let table_id = write_region_id(&request)?.table_id(); + let partitioner = partitioners.get(&table_id).context(UnexpectedSnafu { + err_msg: "Refreshed write has no partition snapshot", + })?; + let requests = partitioner + .rebuild_write_requests(request) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + rebuilt_requests.extend(requests); + } + self.prepare(rebuilt_requests).await + } + + async fn prepare(&self, requests: Vec) -> Result> { + let regions = requests + .iter() + .map(write_region_id) + .collect::>>()?; + let peers = resolve_peers(&self.partition_manager, ®ions) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + requests + .into_iter() + .zip(regions) + .map(|(request, region)| { + Ok(RoutedWrite { + peer: get_peer(&peers, region)?, + request: Arc::new(request), + }) + }) + .collect() + } + + async fn invalidate_routes(&self, regions: &[RegionId]) -> Result<()> { + let mut table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_ids = HashSet::new(); + for table_id in &table_ids { + let (physical_id, _) = self + .partition_manager + .find_physical_table_route_with_id(*table_id) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + physical_ids.insert(physical_id); + } + table_ids.extend(physical_ids); + let caches = table_ids + .into_iter() + .map(CacheIdent::TableId) + .collect::>(); + self.cache_invalidator + .invalidate(&Context::default(), &caches) + .await + } +} + +async fn resolve_peers( + partition_manager: &PartitionRuleManager, + regions: &[RegionId], +) -> partition::error::Result> { + let table_ids = regions + .iter() + .map(|region| region.table_id()) + .collect::>(); + let mut physical_tables = HashMap::new(); + let mut routes = HashMap::new(); + for table_id in table_ids { + let (physical_id, route) = partition_manager + .find_physical_table_route_with_id(table_id) + .await?; + physical_tables.insert(table_id, physical_id); + routes.entry(physical_id).or_insert(route); + } + let mut shared_peers: HashMap<&Peer, PeerRef> = HashMap::new(); + let mut peers = HashMap::with_capacity(regions.len()); + for region in regions { + if peers.contains_key(region) { + continue; + } + let peer = physical_tables + .get(®ion.table_id()) + .and_then(|physical_id| routes.get(physical_id)) + .and_then(|route| { + router::find_region_leader_ref(&route.region_routes, region.region_number()) + }) + .context(partition::error::FindLeaderSnafu { + region_id: *region, + table_id: region.table_id(), + })?; + let peer = shared_peers + .entry(peer) + .or_insert_with(|| Arc::new(peer.clone())); + peers.insert(*region, Arc::clone(peer)); + } + Ok(peers) +} + +fn get_peer(peers: &HashMap, region: RegionId) -> Result { peers.get(®ion).cloned().context(UnexpectedSnafu { err_msg: "Planned write has no resolved peer", }) } +fn merge_attempts( + attempts: Vec, + response: &mut RegionResponse, +) -> Result> { + let mut failures = Vec::new(); + let mut terminal_error = None; + for attempt in attempts { + match attempt { + WriteAttempt::Success(result) => merge_response(response, result), + WriteAttempt::Failed { + error, + request: Some(write), + } => failures.push((write, error)), + WriteAttempt::Failed { + error, + request: None, + } => { + terminal_error.get_or_insert(error); + } + } + } + // A terminal error takes precedence over retained retryable failures. + match terminal_error { + Some(error) => Err(error), + None => Ok(failures), + } +} + fn merge_response(response: &mut RegionResponse, result: RegionResponse) { response.affected_rows += result.affected_rows; response.extensions.extend(result.extensions); diff --git a/src/operator/src/write_dispatcher/options.rs b/src/operator/src/write_dispatcher/options.rs new file mode 100644 index 00000000000..20c8ba535f3 --- /dev/null +++ b/src/operator/src/write_dispatcher/options.rs @@ -0,0 +1,106 @@ +// Copyright 2023 Greptime Team +// +// 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 std::time::Duration; + +use serde::{Deserialize, Serialize}; + +/// Exponential backoff for retryable distributed writes. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +#[serde(default)] +pub struct WriteRetryOptions { + /// Maximum additional attempts after the original write. Zero disables retries. + pub max_retries: usize, + /// Initial backoff before jitter is applied. + #[serde(with = "humantime_serde")] + pub initial_backoff: Duration, + /// Maximum backoff before jitter is applied. + #[serde(with = "humantime_serde")] + pub max_backoff: Duration, + /// Multiplier applied to the backoff after each retry. + pub backoff_multiplier: f32, +} + +impl Default for WriteRetryOptions { + fn default() -> Self { + Self { + max_retries: 3, + initial_backoff: Duration::from_millis(1500), + max_backoff: Duration::from_secs(10), + backoff_multiplier: 2.0, + } + } +} + +impl WriteRetryOptions { + /// Checks that the configured backoff is positive, bounded, and nondecreasing. + pub fn validate(&self) -> Result<(), &'static str> { + if self.initial_backoff.is_zero() { + return Err("write_retry.initial_backoff must be positive"); + } + if self.max_backoff < self.initial_backoff { + return Err("write_retry.max_backoff must be at least initial_backoff"); + } + if !self.backoff_multiplier.is_finite() || self.backoff_multiplier < 1.0 { + return Err("write_retry.backoff_multiplier must be finite and at least 1"); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_write_retry_options() { + for options in [ + WriteRetryOptions::default(), + WriteRetryOptions { + max_retries: 0, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: 1.0, + ..Default::default() + }, + ] { + options.validate().unwrap(); + } + for options in [ + WriteRetryOptions { + initial_backoff: Duration::ZERO, + ..Default::default() + }, + WriteRetryOptions { + max_backoff: Duration::from_secs(1), + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: 0.5, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: f32::NAN, + ..Default::default() + }, + WriteRetryOptions { + backoff_multiplier: f32::INFINITY, + ..Default::default() + }, + ] { + assert!(options.validate().is_err(), "{options:?}"); + } + } +} diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 1191f73d55a..250fe87f353 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -12,36 +12,249 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::future::Future; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; use std::task::Poll; use std::time::{Duration, Instant}; +use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, bulk_insert_request, + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, + RegionResponse as PbRegionResponse, bulk_insert_request, }; +use api::v1::value::ValueData; use api::v1::{ - ArrowIpc, ColumnDataType, PartitionExprVersion, Row, Rows, helper, value::ValueData, + ArrowIpc, ColumnDataType, ColumnSchema, PartitionExprVersion, Row, Rows, Value, helper, }; +use api::v1::{ResponseHeader, Status as PbStatus}; use client::client_manager::NodeClients; -use common_error::ext::ErrorExt; +use common_error::ext::{ErrorExt, RetryHint}; +use common_error::status_code::StatusCode; use common_grpc::channel_manager::ChannelConfig; -use common_meta::cache as meta_cache; +use common_grpc::flight::{self, FlightDecoder}; +use common_meta::cache::{ + self as meta_cache, LayeredCacheRegistryBuilder, TableInfoCacheRef, TableRouteCacheRef, +}; +use common_meta::cache_invalidator::{CacheInvalidator, CacheInvalidatorRef, Context}; use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; +use common_meta::error::{Result, UnexpectedSnafu}; +use common_meta::instruction::CacheIdent; use common_meta::key::TableMetadataManager; -use common_meta::key::table_route::TableRouteValue; +use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; +use common_meta::kv_backend::KvBackendRef; +use common_meta::peer::Peer; use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; -use partition::manager::PartitionRuleManager; +use datatypes::arrow::array::Int32Array; +use datatypes::arrow::datatypes::{DataType, Field, Schema}; +use datatypes::arrow::record_batch::RecordBatch; +use partition::cache::PartitionInfoCacheRef; +use partition::expr::{Operand, PartitionExpr, RestrictedOp}; +use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; +use store_api::storage::RegionId; +use tokio::net::TcpListener; use tokio::sync::{Barrier, Notify}; +use tokio::task::JoinHandle; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::{Request, Response, Status}; use super::*; use crate::error::Error; use crate::test_util; +#[derive(Clone, Default)] +struct RecordingRegionService { + requests: Arc>>, + rejections: Arc>>>, +} + +#[async_trait::async_trait] +impl Region for RecordingRegionService { + async fn handle( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + let region_id = write_region_id(&request).unwrap().as_u64(); + self.requests.lock().unwrap().push(request); + if let Some(hint) = self + .rejections + .lock() + .unwrap() + .get_mut(®ion_id) + .and_then(VecDeque::pop_front) + { + let mut status = Status::unavailable(format!("region {region_id} rejected the write")); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.as_str().parse().unwrap(), + ); + return Err(status); + } + Ok(Response::new(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::Success as u32, + ..Default::default() + }), + }), + affected_rows: 1, + ..Default::default() + })) + } +} + +#[derive(Default)] +struct RecordingCacheInvalidator { + calls: Mutex>>, + inner: Option, + route_update: Mutex>, +} + +#[async_trait::async_trait] +impl CacheInvalidator for RecordingCacheInvalidator { + async fn invalidate(&self, ctx: &Context, caches: &[CacheIdent]) -> Result<()> { + self.calls.lock().unwrap().push(caches.to_vec()); + let route_update = self.route_update.lock().unwrap().take(); + if let Some((backend, updated)) = route_update { + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await?.unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated)?; + assert!(backend.txn(txn).await?.succeeded); + } + if let Some(inner) = &self.inner { + inner.invalidate(ctx, caches).await + } else { + UnexpectedSnafu { + err_msg: "stop after recording route refresh", + } + .fail() + } + } + + fn invalidate_all(&self) -> Result<()> { + // This recorder owns no cache entries. + Ok(()) + } +} + +struct RegionTestServer { + peer: Peer, + service: Arc, + invalidator: Arc, + task: JoinHandle<()>, +} + +impl RegionTestServer { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let peer = Peer::new(1, listener.local_addr().unwrap().to_string()); + let service = Arc::new(RecordingRegionService::default()); + let server_service = service.clone(); + let task = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(RegionServer::new(server_service.as_ref().clone())) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + Self { + peer, + service, + invalidator: Arc::default(), + task, + } + } + + async fn create_dispatcher(&self, max_retries: usize) -> WriteDispatcher { + let backend = test_util::prepare_mocked_backend().await; + let partition_manager = test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &self.peer).await; + let registry = cache::build_fundamental_cache_registry(backend); + let table_info_cache = registry.get::().unwrap(); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + WriteDispatcher::new(clients.clone(), partition_manager).with_retry( + clients, + build_retry_options(max_retries), + table_info_cache, + self.invalidator.clone(), + ) + } + + fn fail_region(&self, region_number: u32, hints: impl IntoIterator) { + self.service.rejections.lock().unwrap().insert( + RegionId::new(1, region_number).as_u64(), + hints.into_iter().collect(), + ); + } + + fn attempt_counts(&self) -> HashMap { + let mut counts = HashMap::new(); + for request in self.service.requests.lock().unwrap().iter() { + *counts + .entry(write_region_id(request).unwrap().as_u64()) + .or_default() += 1; + } + counts + } +} + +impl Drop for RegionTestServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +fn build_retry_options(max_retries: usize) -> WriteRetryOptions { + WriteRetryOptions { + max_retries, + initial_backoff: Duration::from_millis(1), + max_backoff: Duration::from_millis(2), + ..Default::default() + } +} + +fn create_routing_caches( + backend: KvBackendRef, +) -> ( + PartitionRuleManagerRef, + TableInfoCacheRef, + CacheInvalidatorRef, +) { + let registry = Arc::new( + cache::with_default_composite_cache_registry( + LayeredCacheRegistryBuilder::default() + .add_cache_registry(cache::build_fundamental_cache_registry(backend.clone())), + ) + .unwrap() + .build(), + ); + let partition_manager = Arc::new(PartitionRuleManager::new( + backend.clone(), + registry.get::().unwrap(), + registry.get::().unwrap(), + )); + let table_info_cache = registry.get::().unwrap(); + (partition_manager, table_info_cache, registry) +} + +async fn update_table_leader(backend: &KvBackendRef, peer: &Peer) { + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await.unwrap().unwrap(); + let mut routes = current.region_routes().unwrap().clone(); + for route in &mut routes { + route.leader_peer = Some(peer.clone()); + } + let updated = current.update(routes).unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated).unwrap(); + assert!(backend.txn(txn).await.unwrap().succeeded); +} + fn build_insert_request(regions: &[u32]) -> RegionRequest { RegionRequest { body: Some(Body::Inserts(InsertRequests { @@ -57,6 +270,529 @@ fn build_insert_request(regions: &[u32]) -> RegionRequest { } } +#[tokio::test] +async fn test_write_retry_disabled_and_enabled_ablation() { + common_telemetry::init_default_ut_logging(); + for run in 0..3 { + for max_retries in [0, 1] { + let server = RegionTestServer::start().await; + server.fail_region(1, [RetryHint::Retryable]); + let dispatcher = server.create_dispatcher(max_retries).await; + let result = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; + let attempts = server.attempt_counts()[&RegionId::new(1, 1).as_u64()]; + common_telemetry::info!( + "write_retry_ablation run={run} max_retries={max_retries} success={} attempts={attempts}", + result.is_ok() + ); + assert_eq!(result.is_ok(), max_retries > 0); + assert_eq!(attempts, max_retries + 1); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn test_write_retry_does_not_replay_successful_regions() { + let server = RegionTestServer::start().await; + server.fail_region(2, [RetryHint::Retryable]); + let dispatcher = server.create_dispatcher(2).await; + let response = dispatcher + .dispatch(vec![build_insert_request(&[1, 2])]) + .await + .unwrap(); + assert_eq!(response.affected_rows, 2); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 2), + ]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn test_write_retry_attempt_budget_and_nonretryable_errors() { + for (hint, expected_attempts) in [(RetryHint::Retryable, 3), (RetryHint::NonRetryable, 1)] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint; 4]); + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(&[1])]) + .await + .unwrap_err(); + // The dispatcher uses the explicit client hint for its attempt budget. + // The existing common-meta conversion promotes Unavailable to RetryLater. + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert_eq!( + server.attempt_counts()[&RegionId::new(1, 1).as_u64()], + expected_attempts + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } +} + +#[tokio::test] +async fn test_write_retry_stops_when_a_sibling_is_nonretryable() { + for hint in [RetryHint::Retryable, RetryHint::RetryAfterRouteRefresh] { + for regions in [[1, 2], [2, 1]] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint]); + server.fail_region(2, [RetryHint::NonRetryable]); + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await + .unwrap_err(); + let terminal_region = RegionId::new(1, 2).as_u64(); + assert!( + format!("{error:?}") + .contains(&format!("region {terminal_region} rejected the write")), + "{error:?}" + ); + assert_eq!( + server.attempt_counts(), + HashMap::from([(RegionId::new(1, 1).as_u64(), 1), (terminal_region, 1),]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn test_should_retry_requires_a_retained_request_and_budget() { + let backend = test_util::prepare_mocked_backend().await; + let (partition_manager, table_info_cache, cache_invalidator) = create_routing_caches(backend); + let retry = WriteRetry { + clients: Arc::new(NodeClients::default()), + partition_manager, + options: build_retry_options(2), + table_info_cache, + cache_invalidator, + }; + let failure = |request| { + // The converted error alone cannot distinguish an explicitly nonretryable attempt. + let error: client::Error = Status::unavailable("write rejected").into(); + WriteAttempt::Failed { + error: error.into(), + request, + } + }; + let retryable = || { + failure(Some(RoutedWrite { + peer: Arc::new(Peer::new(1, "peer")), + request: Arc::new(build_insert_request(&[1])), + })) + }; + for (name, attempts, retries, expected) in [ + ("empty", vec![], 0, false), + ( + "all succeeded", + vec![WriteAttempt::Success(RegionResponse::new(1))], + 0, + false, + ), + ("retryable", vec![retryable()], 0, true), + ("last retry", vec![retryable()], 1, true), + ("budget exhausted", vec![retryable()], 2, false), + ("terminal first", vec![failure(None), retryable()], 0, false), + ("terminal last", vec![retryable(), failure(None)], 0, false), + ] { + assert_eq!(retry.should_retry(&attempts, retries), expected, "{name}"); + } +} + +#[tokio::test] +async fn test_write_retry_refreshes_routes_once_per_failed_round() { + let server = RegionTestServer::start().await; + for region in [1, 2] { + server.fail_region(region, [RetryHint::RetryAfterRouteRefresh]); + } + let dispatcher = server.create_dispatcher(2).await; + let error = dispatcher + .dispatch(vec![build_insert_request(&[1, 2])]) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("stop after recording route refresh") + ); + let calls = server.invalidator.calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0], vec![CacheIdent::TableId(1)]); + assert_eq!(server.service.requests.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn test_finish_retries_initial_failures_in_one_round() { + let server = RegionTestServer::start().await; + for region in [2, 3] { + server.fail_region(region, [RetryHint::RetryAfterRouteRefresh]); + } + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &server.peer).await; + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let invalidator = Arc::new(RecordingCacheInvalidator { + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry( + clients, + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let mut attempts = Vec::new(); + for (region, value) in [(1, 101), (2, 11), (3, 1)] { + let request = + build_partitioned_request(WriteKind::BulkInsert, region, versions[®ion], &[value]); + attempts.push( + dispatcher + .send(Arc::new(server.peer.clone()), request, None) + .await, + ); + } + assert!(invalidator.calls.lock().unwrap().is_empty()); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + (RegionId::new(1, 3).as_u64(), 1), + ]) + ); + let response = dispatcher.finish(attempts, None).await.unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 2), + (RegionId::new(1, 3).as_u64(), 2), + ]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); +} + +#[tokio::test] +async fn test_write_retry_refreshes_stale_route_and_sends_only_failed_region_to_new_leader() { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = test_util::prepare_mocked_backend().await; + // Reuse the operator fixture to populate table metadata and partition rules. + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; + + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend.clone()); + let failed_region = RegionId::new(1, 2); + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + + // Metadata has migrated while the frontend still holds the old route. + update_table_leader(&backend, &new_server.peer).await; + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); + + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry( + clients, + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let mut request = build_insert_request(&[1, 2]); + if let Some(Body::Inserts(inserts)) = &mut request.body { + for insert in &mut inserts.requests { + insert.partition_expr_version = versions + [&RegionId::from_u64(insert.region_id).region_number()] + .map(|value| PartitionExprVersion { value }); + } + } + let response = dispatcher.dispatch(vec![request]).await.unwrap(); + assert_eq!(response.affected_rows, 2); + assert_eq!( + old_server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (failed_region.as_u64(), 1), + ]) + ); + assert_eq!( + new_server.attempt_counts(), + HashMap::from([(failed_region.as_u64(), 1)]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); +} + +#[derive(Clone, Copy, Debug)] +enum WriteKind { + Insert, + Delete, + BulkInsert, +} + +fn build_partitioned_request( + kind: WriteKind, + region: u32, + version: Option, + values: &[i32], +) -> RegionRequest { + let region_id = RegionId::new(1, region).as_u64(); + let partition_expr_version = version.map(|value| PartitionExprVersion { value }); + let rows = Rows { + schema: vec![ColumnSchema { + column_name: "a".to_string(), + datatype: ColumnDataType::Int32 as i32, + ..Default::default() + }], + rows: values + .iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(*value)), + }], + }) + .collect(), + }; + let body = match kind { + WriteKind::Insert => Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id, + partition_expr_version, + rows: Some(rows), + skip_wal: true, + }], + }), + WriteKind::Delete => Body::Deletes(DeleteRequests { + requests: vec![DeleteRequest { + region_id, + partition_expr_version, + rows: Some(rows), + }], + }), + WriteKind::BulkInsert => { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(values.to_vec()))], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap(); + Body::BulkInsert(BulkInsertRequest { + region_id, + partition_expr_version, + skip_wal: true, + aligned_schema_version: Some(AlignedSchemaVersion { + schema_version: 123, + }), + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + }) + } + }; + RegionRequest { + header: None, + body: Some(body), + } +} + +fn request_parts(request: &RegionRequest) -> (u32, Option, Vec) { + let region = write_region_id(request).unwrap().region_number(); + let (version, rows) = match request.body.as_ref().unwrap() { + Body::Inserts(inserts) => { + assert!(inserts.requests[0].skip_wal); + ( + inserts.requests[0].partition_expr_version, + inserts.requests[0].rows.as_ref().unwrap(), + ) + } + Body::Deletes(deletes) => ( + deletes.requests[0].partition_expr_version, + deletes.requests[0].rows.as_ref().unwrap(), + ), + Body::BulkInsert(bulk) => { + assert!(bulk.skip_wal); + assert_eq!(bulk.aligned_schema_version, None); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + panic!("expected IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + return ( + region, + bulk.partition_expr_version.map(|v| v.value), + values.values().to_vec(), + ); + } + _ => panic!("expected write request"), + }; + ( + region, + version.map(|v| v.value), + rows.rows + .iter() + .map(|row| { + let Some(ValueData::I32Value(value)) = row.values[0].value_data else { + panic!("expected int32") + }; + value + }) + .collect(), + ) +} + +#[tokio::test] +async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() { + for kind in [WriteKind::Insert, WriteKind::Delete, WriteKind::BulkInsert] { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; + let (partition_manager, table_info_cache, registry) = + create_routing_caches(backend.clone()); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, old_versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let manager = TableRouteManager::new(backend.clone()); + let current = manager + .table_route_storage() + .get_with_raw_bytes(1) + .await + .unwrap() + .unwrap(); + let old_routes = current.region_routes().unwrap(); + let replaced = old_routes + .iter() + .find(|route| route.region.id.region_number() == 2) + .unwrap(); + let mut routes = old_routes + .iter() + .filter(|route| route.region.id.region_number() != 2) + .cloned() + .collect::>(); + for (region, lower, upper) in [(4, 10, 30), (5, 30, 50)] { + let mut route = replaced.clone(); + route.region.id = RegionId::new(1, region); + route.leader_peer = Some(new_server.peer.clone()); + route.region.partition_expr = PartitionExpr::new( + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::GtEq, + Operand::Value(datatypes::value::Value::Int32(lower)), + )), + RestrictedOp::And, + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::Lt, + Operand::Value(datatypes::value::Value::Int32(upper)), + )), + ) + .as_json_str() + .unwrap(); + routes.push(route); + } + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + route_update: Mutex::new(Some((backend, current.update(routes).unwrap()))), + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager.clone()) + .with_retry( + clients, + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let response = dispatcher + .dispatch(vec![ + build_partitioned_request(kind, 1, old_versions[&1], &[70]), + build_partitioned_request(kind, 2, old_versions[&2], &[20, 40]), + ]) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3, "{kind:?}"); + assert_eq!( + old_server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]), + "{kind:?}" + ); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let mut writes = new_server + .service + .requests + .lock() + .unwrap() + .iter() + .map(request_parts) + .collect::>(); + writes.sort_by_key(|write| write.0); + assert_eq!( + writes, + vec![(4, versions[&4], vec![20]), (5, versions[&5], vec![40])], + "{kind:?}" + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); + } +} + #[tokio::test] async fn test_resolve_peers_batches_logical_and_physical_regions() { let backend = test_util::prepare_mocked_backend().await; @@ -82,11 +818,14 @@ async fn test_resolve_peers_batches_logical_and_physical_regions() { assert_eq!(peers.len(), 4); for region in regions { assert_eq!( - peers[®ion], - partition_manager.find_region_leader(region).await.unwrap() + peers[®ion].as_ref(), + &partition_manager.find_region_leader(region).await.unwrap() ); } - assert_eq!(peers[&RegionId::new(1, 1)], peers[&RegionId::new(42, 1)]); + assert!(Arc::ptr_eq( + &peers[&RegionId::new(1, 1)], + &peers[&RegionId::new(42, 1)], + )); let missing = RegionId::new(42, 999); let expected = partition_manager .find_region_leader(missing) @@ -98,6 +837,38 @@ async fn test_resolve_peers_batches_logical_and_physical_regions() { assert!(dispatcher.resolve_peers(&[]).await.unwrap().is_empty()); } +#[tokio::test] +async fn test_resolve_peers_shares_only_identical_peers() { + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + let manager = TableRouteManager::new(backend.clone()); + let storage = manager.table_route_storage(); + let current = storage.get_with_raw_bytes(1).await.unwrap().unwrap(); + let mut routes = current.region_routes().unwrap().clone(); + let peer = Peer::new(7, "same-address"); + let other = Peer::new(7, "other-address"); + for route in &mut routes { + route.leader_peer = Some(if route.region.id.region_number() == 3 { + other.clone() + } else { + peer.clone() + }); + } + let updated = current.update(routes).unwrap(); + let (txn, _) = storage.build_update_txn(1, ¤t, &updated).unwrap(); + assert!(backend.txn(txn).await.unwrap().succeeded); + let (partition_manager, _, _) = create_routing_caches(backend); + let dispatcher = WriteDispatcher::new(Arc::new(NodeClients::default()), partition_manager); + let regions = [1, 2, 3].map(|region| RegionId::new(1, region)); + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + assert_eq!(peers.len(), 3); + assert_eq!(peers[®ions[0]].as_ref(), &peer); + assert_eq!(peers[®ions[1]].as_ref(), &peer); + assert_eq!(peers[®ions[2]].as_ref(), &other); + assert!(Arc::ptr_eq(&peers[®ions[0]], &peers[®ions[1]])); + assert!(!Arc::ptr_eq(&peers[®ions[0]], &peers[®ions[2]])); +} + #[tokio::test] async fn test_dispatch_preserves_write_payloads_and_envelopes() { let backend = test_util::prepare_mocked_backend().await; @@ -158,7 +929,7 @@ async fn test_dispatch_preserves_write_payloads_and_envelopes() { }); let writes = dispatcher.prepare(requests.to_vec()).await.unwrap(); assert!(received.try_recv().is_err()); - let response = dispatcher.dispatch_prepared(writes, None).await.unwrap(); + let response = dispatcher.execute(writes, None).await.unwrap(); assert_eq!(response.affected_rows, 2); for request in requests { assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); @@ -315,8 +1086,9 @@ async fn test_resolve_peers_cache_lookup_ablation() { for mode in modes { let before = counter.get(); let started = Instant::now(); - let peers = if mode == "batched" { - dispatcher.resolve_peers(®ions).await.unwrap() + let (elapsed, peers_match) = if mode == "batched" { + let peers = dispatcher.resolve_peers(®ions).await.unwrap(); + (started.elapsed(), peers == expected) } else { let mut peers = HashMap::new(); for region in ®ions { @@ -328,11 +1100,17 @@ async fn test_resolve_peers_cache_lookup_ablation() { partition_manager.find_region_leader(*region).await.unwrap(), ); } - peers + let elapsed = started.elapsed(); + let peers_match = peers.len() == expected.len() + && peers.iter().all(|(region, peer)| { + expected + .get(region) + .is_some_and(|expected| expected.as_ref() == peer) + }); + (elapsed, peers_match) }; - let elapsed = started.elapsed(); let gets = counter.get() - before; - assert_eq!(peers, expected); + assert!(peers_match, "{mode}"); let expected_gets = match mode { "delete" => 384, "insert" => 9, @@ -348,3 +1126,121 @@ async fn test_resolve_peers_cache_lookup_ablation() { } } } + +#[tokio::test] +async fn test_write_retry_rebuilds_logical_rows_in_mixed_table_envelope() { + for kind in [WriteKind::Insert, WriteKind::Delete] { + let server = RegionTestServer::start().await; + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + TableMetadataManager::new(backend.clone()) + .create_table_metadata( + test_util::new_test_table_info(42, "retry_logical", [1, 2, 3].into_iter()), + TableRouteValue::logical(1), + Default::default(), + ) + .await + .unwrap(); + update_table_leader(&backend, &server.peer).await; + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let invalidator = Arc::new(RecordingCacheInvalidator { + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry( + clients, + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); + let physical_id = RegionId::new(1, 1).as_u64(); + let stale_logical_id = RegionId::new(42, 2).as_u64(); + let rebuilt_logical_id = RegionId::new(42, 3).as_u64(); + server.service.rejections.lock().unwrap().insert( + stale_logical_id, + VecDeque::from([RetryHint::RetryAfterRouteRefresh]), + ); + + let header = Some(RegionRequestHeader { + dbname: "greptime.public".to_string(), + ..Default::default() + }); + let mut physical = build_partitioned_request(kind, 1, versions[&1], &[70]); + physical.header = header.clone(); + // The stale logical request must move to region 3 while keeping table 42. + let mut logical = build_partitioned_request(kind, 2, Some(0), &[1]); + let mut rebuilt = build_partitioned_request(kind, 3, versions[&3], &[1]); + for request in [&mut logical, &mut rebuilt] { + request.header = header.clone(); + let region_id = match request.body.as_mut().unwrap() { + Body::Inserts(inserts) => &mut inserts.requests[0].region_id, + Body::Deletes(deletes) => &mut deletes.requests[0].region_id, + _ => unreachable!(), + }; + *region_id = RegionId::new(42, RegionId::from_u64(*region_id).region_number()).as_u64(); + } + let mut envelope = physical.clone(); + match ( + envelope.body.as_mut().unwrap(), + logical.body.clone().unwrap(), + ) { + (Body::Inserts(physical), Body::Inserts(logical)) => { + physical.requests.extend(logical.requests); + } + (Body::Deletes(physical), Body::Deletes(logical)) => { + physical.requests.extend(logical.requests); + } + _ => unreachable!(), + } + let response = dispatcher.dispatch(vec![envelope]).await.unwrap(); + assert_eq!(response.affected_rows, 2, "{kind:?}"); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (physical_id, 1), + (stale_logical_id, 1), + (rebuilt_logical_id, 1) + ]), + "{kind:?}" + ); + let mut requests = server.service.requests.lock().unwrap().clone(); + requests.sort_by_key(|request| write_region_id(request).unwrap().as_u64()); + assert_eq!(requests, vec![physical, logical, rebuilt], "{kind:?}"); + let calls = invalidator.calls.lock().unwrap(); + assert_eq!(calls.len(), 1, "{kind:?}"); + assert_eq!(calls[0].len(), 2, "{kind:?}"); + assert!(calls[0].contains(&CacheIdent::TableId(1))); + assert!(calls[0].contains(&CacheIdent::TableId(42))); + } +} + +#[tokio::test] +async fn test_with_retry_can_disable_previously_enabled_retries() { + let server = RegionTestServer::start().await; + let dispatcher = server.create_dispatcher(1).await; + let retry = dispatcher.retry.as_ref().unwrap(); + let clients = retry.clients.clone(); + let table_info_cache = retry.table_info_cache.clone(); + let cache_invalidator = retry.cache_invalidator.clone(); + let dispatcher = dispatcher.with_retry( + clients, + build_retry_options(0), + table_info_cache, + cache_invalidator, + ); + assert!(dispatcher.retry.is_none()); + server.fail_region(1, [RetryHint::Retryable]); + assert!( + dispatcher + .dispatch(vec![build_insert_request(&[1])]) + .await + .is_err() + ); + assert_eq!(server.attempt_counts()[&RegionId::new(1, 1).as_u64()], 1); +} From 2039f508fcd9b04532593646361b14d73ad9bc9a Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 13 Oct 2026 07:58:31 +0000 Subject: [PATCH 28/46] fix(promql): stop fabricating histogram fold hash partitioning (#9509) * fix(promql): stop fabricating histogram fold hash partitioning Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): cover histogram fraction binary join missing rows Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): normalize dictionary labels in histogram regression Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: cover histogram joins on logical metric tables Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: specify insert columns for logical metric fixture Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: use supported histogram child replacement API Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): keep histogram fold fixtures truthful Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/extension_plan/histogram_fold.rs | 71 ++++++++-- tests-integration/src/tests/promql_test.rs | 134 ++++++++++++++++++ 2 files changed, 192 insertions(+), 13 deletions(-) diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 709117df294..fb691cbc00d 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -385,10 +385,7 @@ impl HistogramFold { let output_schema: SchemaRef = self.output_schema.inner().clone(); let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(output_schema.clone()), - Partitioning::Hash( - partition_exprs.clone(), - exec_input.output_partitioning().partition_count(), - ), + Partitioning::UnknownPartitioning(exec_input.output_partitioning().partition_count()), EmissionType::Incremental, Boundedness::Bounded, )); @@ -604,10 +601,7 @@ impl ExecutionPlan for HistogramFoldExec { let new_input = children[0].clone(); let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(self.output_schema.clone()), - Partitioning::Hash( - self.partition_exprs.clone(), - new_input.output_partitioning().partition_count(), - ), + Partitioning::UnknownPartitioning(new_input.output_partitioning().partition_count()), EmissionType::Incremental, Boundedness::Bounded, )); @@ -1442,6 +1436,7 @@ mod test { use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::EmptyRelation; + use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion::prelude::SessionContext; use datatypes::arrow_array::StringArray; use futures::FutureExt; @@ -1461,7 +1456,9 @@ mod test { RecordBatch::try_new(schema, columns).unwrap() } - fn prepare_test_data() -> DataSourceExec { + /// The batches and schema shared by [`prepare_test_data`] and + /// [`prepare_test_data_with_two_partitions`]. + fn prepare_test_data_batches() -> (Vec, SchemaRef) { let schema = Arc::new(Schema::new(vec![ Field::new("host", DataType::Utf8, true), Field::new("le", DataType::Utf8, true), @@ -1514,8 +1511,24 @@ mod test { ) .unwrap(); + (vec![data_1, data_2, data_3], schema) + } + + fn prepare_test_data() -> DataSourceExec { + let (batches, schema) = prepare_test_data_batches(); DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[vec![data_1, data_2, data_3]], schema, None).unwrap(), + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + )) + } + + /// The batches from [`prepare_test_data_batches`], laid out over two + /// partitions, so that partitioning metadata is checked against a real + /// multi-partition input. + fn prepare_test_data_with_two_partitions() -> DataSourceExec { + let (mut batches, schema) = prepare_test_data_batches(); + let first_partition = vec![batches.remove(0)]; + DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[first_partition, batches], schema, None).unwrap(), )) } @@ -1603,9 +1616,8 @@ mod test { let properties = PlanProperties::new( EquivalenceProperties::new(output_schema.clone()), - Partitioning::Hash( - partition_exprs.clone(), - input.output_partitioning().partition_count(), + Partitioning::UnknownPartitioning( + input.properties().output_partitioning().partition_count(), ), EmissionType::Incremental, Boundedness::Bounded, @@ -1808,6 +1820,39 @@ mod test { assert!(broken_result.is_err()); } + #[test] + fn fold_does_not_claim_hash_partitioning() { + let input: Arc = Arc::new(prepare_test_data_with_two_partitions()); + let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(input.schema().to_dfschema().unwrap()), + }); + let logical_fold = HistogramFold::new( + "le".to_string(), + "val".to_string(), + "host".to_string(), + 0.5, + logical_input, + ) + .unwrap(); + let constructed = logical_fold.to_execution_plan(input.clone()); + assert!(matches!( + constructed.output_partitioning(), + Partitioning::UnknownPartitioning(2) + )); + + let rebuilt = constructed + .replace_children( + vec![input], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + assert!(matches!( + rebuilt.output_partitioning(), + Partitioning::UnknownPartitioning(2) + )); + } + #[test] fn confirm_schema() { let input_schema = Schema::new(vec![ diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 74114bbd6c4..8cf27d165e9 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -26,6 +26,7 @@ use datatypes::arrow::array::{Float64Array, Int64Array, StringArray, TimestampMi use datatypes::arrow::compute::cast; use datatypes::arrow::datatypes::DataType; use frontend::instance::Instance; +use query::datafusion::QUERY_PARALLELISM_HINT; use query::parser::{PromQuery, QueryLanguageParser, QueryStatement}; use rstest::rstest; use rstest_reuse::apply; @@ -193,6 +194,139 @@ async fn promql_query_as_batches( } } +#[tokio::test(flavor = "multi_thread")] +async fn histogram_fraction_binary_join_preserves_partitioned_rows() { + common_telemetry::init_default_ut_logging(); + + let distributed = crate::tests::create_distributed_instance( + "histogram_fraction_binary_join_preserves_partitioned_rows", + ) + .await; + let instance = distributed.frontend(); + + for (table, create) in [ + ( + "pa_promql_m", + r#"CREATE TABLE pa_promql_m ( + ts TIMESTAMP TIME INDEX, + host STRING, + shard STRING, + val DOUBLE, + PRIMARY KEY (host, shard) + ) PARTITION ON COLUMNS (shard) ( + shard < 'b', + shard >= 'b' AND shard < 'c', + shard >= 'c' AND shard < 'd', + shard >= 'd' + );"#, + ), + ( + "lg_hist_m_2", + r#"CREATE TABLE phy_hist_m_2 ( + ts TIMESTAMP TIME INDEX, + host STRING, + shard STRING, + val DOUBLE, + PRIMARY KEY (host, shard) + ) PARTITION ON COLUMNS (shard) ( + shard < 'b', + shard >= 'b' AND shard < 'c', + shard >= 'c' AND shard < 'd', + shard >= 'd' + ) ENGINE=metric WITH ('physical_metric_table' = 'true'); + CREATE TABLE lg_hist_m_2 ( + ts TIMESTAMP TIME INDEX, + host STRING, + shard STRING, + val DOUBLE, + PRIMARY KEY (host, shard) + ) PARTITION ON COLUMNS (shard) ( + shard < 'b', + shard >= 'b' AND shard < 'c', + shard >= 'c' AND shard < 'd', + shard >= 'd' + ) ENGINE=metric WITH ('on_physical_table' = 'phy_hist_m_2');"#, + ), + ] { + execute_all( + &instance, + &format!( + "{create}\nINSERT INTO {table} (ts, host, shard, val) VALUES\n (0, 'h', 'a', 1),\n (0, 'h', 'b', 1),\n (0, 'h', 'c', 1),\n (0, 'h', 'd', 1);" + ), + QueryContext::arc(), + ) + .await; + + let expr = format!( + r#"histogram_fraction(-Inf,+Inf,abs(label_replace(({table} + on(host) group_left() (sum by(host)({table}))),"le","+Inf","missing_label",".*")))"# + ); + for parallelism in [1, 4] { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + let query_ctx = Arc::new(query_ctx); + for (query, expected_value) in [ + (format!("TQL EVAL (0, 0, '1s') {expr}"), 1.0), + ( + format!("TQL EVAL (0, 0, '1s') {expr} + on(host, shard) {table}"), + 2.0, + ), + ] { + let output = instance + .do_query(&query, query_ctx.clone()) + .await + .remove(0) + .unwrap(); + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + other => panic!("expected result batches, got {other:?}"), + }; + let mut rows = batches + .iter() + .flat_map(|batch| { + let ts = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let host_array = + cast(batch.column_by_name("host").unwrap(), &DataType::Utf8).unwrap(); + let host = host_array.as_any().downcast_ref::().unwrap(); + let shard_array = + cast(batch.column_by_name("shard").unwrap(), &DataType::Utf8).unwrap(); + let shard = shard_array.as_any().downcast_ref::().unwrap(); + let value = batch + .columns() + .iter() + .find_map(|column| column.as_any().downcast_ref::()) + .unwrap(); + (0..batch.num_rows()) + .map(move |row| { + ( + ts.value(row), + host.value(row).to_owned(), + shard.value(row).to_owned(), + value.value(row), + ) + }) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| left.2.cmp(&right.2)); + assert_eq!( + rows, + ["a", "b", "c", "d"] + .into_iter() + .map(|shard| (0, "h".to_owned(), shard.to_owned(), expected_value)) + .collect::>(), + "query_parallelism={parallelism}, query={query}" + ); + } + } + } +} + const ANON_PROMQL_RATIO_REPRO_DB: &str = "repro_db"; const ANON_PROMQL_RATIO_REPRO_CREATE: &str = r#" From 2a701949301bd6e30ae1bbab04c2a955d3e77839 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 13 Oct 2026 08:03:49 +0000 Subject: [PATCH 29/46] feat: add per-column COMPRESSION option for SST columns (#9426) * feat: add per-column COMPRESSION option for SST columns CREATE TABLE accepts COMPRESSION WITH (type = 'zstd', level = N, page_size = '') on a column. The options are stored in the column metadata, carried through gRPC column options, shown by SHOW CREATE TABLE, and applied by the mito SST writer and bulk part encoder as parquet column properties. Unset fields keep the SST defaults. Nested types are rejected. A page size below the dictionary page limit also lowers the column's dictionary page limit: parquet sizes write batches by the dictionary budget and keeps that size after falling back to plain encoding, which would otherwise ignore the smaller data page limit. Signed-off-by: Dennis Zhuang * fix: validate column compression options before persistence Signed-off-by: Dennis Zhuang * fix: reject compression overrides on metric logical tables Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/api/src/error.rs | 11 +- src/api/src/v1/column_def.rs | 80 ++++- src/datatypes/src/error.rs | 10 +- src/datatypes/src/schema.rs | 14 +- src/datatypes/src/schema/column_schema.rs | 275 ++++++++++++++++++ .../src/engine/alter/validate.rs | 54 +++- src/metric-engine/src/engine/create.rs | 13 +- src/mito2/src/engine/create_test.rs | 68 +++++ src/mito2/src/error.rs | 10 + src/mito2/src/memtable/bulk/part.rs | 4 +- src/mito2/src/region/opener.rs | 2 + src/mito2/src/region/options.rs | 29 ++ src/mito2/src/sst/parquet.rs | 177 ++++++++++- src/mito2/src/sst/parquet/writer.rs | 3 +- src/mito2/src/worker/handle_alter.rs | 13 + src/operator/src/expr_helper.rs | 10 +- src/query/src/sql/show_create_table.rs | 4 + src/sql/src/error.rs | 8 + src/sql/src/parsers/create_parser.rs | 135 ++++++++- src/sql/src/parsers/utils.rs | 19 +- src/sql/src/statements.rs | 10 +- src/sql/src/statements/create.rs | 22 +- src/store-api/src/metadata.rs | 4 + .../create/create_with_compression.result | 218 ++++++++++++++ .../common/create/create_with_compression.sql | 96 ++++++ .../cases/column_compression/case.toml | 8 + .../cases/column_compression/setup.sql | 3 + .../cases/column_compression/verify.result | 63 ++++ .../cases/column_compression/verify.sql | 15 + 29 files changed, 1340 insertions(+), 38 deletions(-) create mode 100644 tests/cases/standalone/common/create/create_with_compression.result create mode 100644 tests/cases/standalone/common/create/create_with_compression.sql create mode 100644 tests/compatibility/cases/column_compression/case.toml create mode 100644 tests/compatibility/cases/column_compression/setup.sql create mode 100644 tests/compatibility/cases/column_compression/verify.result create mode 100644 tests/compatibility/cases/column_compression/verify.sql diff --git a/src/api/src/error.rs b/src/api/src/error.rs index 0b35bbc934f..f6694c68679 100644 --- a/src/api/src/error.rs +++ b/src/api/src/error.rs @@ -60,6 +60,14 @@ pub enum Error { source: datatypes::error::Error, }, + #[snafu(display("Invalid compression options for column {column}"))] + InvalidColumnCompression { + column: String, + #[snafu(implicit)] + location: Location, + source: datatypes::error::Error, + }, + #[snafu(display("Failed to serialize JSON"))] SerializeJson { #[snafu(source)] @@ -93,7 +101,8 @@ impl ErrorExt for Error { StatusCode::Unexpected } Error::ConvertColumnDefaultConstraint { source, .. } - | Error::InvalidColumnDefaultConstraint { source, .. } => source.status_code(), + | Error::InvalidColumnDefaultConstraint { source, .. } + | Error::InvalidColumnCompression { source, .. } => source.status_code(), } } diff --git a/src/api/src/v1/column_def.rs b/src/api/src/v1/column_def.rs index b02b06feafd..af62ed716e5 100644 --- a/src/api/src/v1/column_def.rs +++ b/src/api/src/v1/column_def.rs @@ -16,9 +16,9 @@ use std::collections::HashMap; use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datatypes::schema::{ - COMMENT_KEY, ColumnDefaultConstraint, ColumnSchema, FULLTEXT_KEY, FulltextAnalyzer, - FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, - SkippingIndexOptions, SkippingIndexType, + COMMENT_KEY, COMPRESSION_KEY, ColumnDefaultConstraint, ColumnSchema, FULLTEXT_KEY, + FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, + SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, }; use greptime_proto::v1::{ Analyzer, FulltextBackend as PbFulltextBackend, SkippingIndexType as PbSkippingIndexType, @@ -35,11 +35,14 @@ const FULLTEXT_GRPC_KEY: &str = "fulltext"; const INVERTED_INDEX_GRPC_KEY: &str = "inverted_index"; /// Key used to store skip index options in gRPC column options. const SKIPPING_INDEX_GRPC_KEY: &str = "skipping_index"; +/// Key used to store SST compression options in gRPC column options. +const COMPRESSION_GRPC_KEY: &str = "compression"; -const COLUMN_OPTION_MAPPINGS: [(&str, &str); 5] = [ +const COLUMN_OPTION_MAPPINGS: [(&str, &str); 6] = [ (FULLTEXT_GRPC_KEY, FULLTEXT_KEY), (INVERTED_INDEX_GRPC_KEY, INVERTED_INDEX_KEY), (SKIPPING_INDEX_GRPC_KEY, SKIPPING_INDEX_KEY), + (COMPRESSION_GRPC_KEY, COMPRESSION_KEY), (EXTENSION_TYPE_NAME_KEY, EXTENSION_TYPE_NAME_KEY), (EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_METADATA_KEY), ]; @@ -77,6 +80,9 @@ pub fn try_as_column_schema(column_def: &ColumnDef) -> Result { if let Some(skipping_index) = options.options.get(SKIPPING_INDEX_GRPC_KEY) { metadata.insert(SKIPPING_INDEX_KEY.to_string(), skipping_index.to_owned()); } + if let Some(compression) = options.options.get(COMPRESSION_GRPC_KEY) { + metadata.insert(COMPRESSION_KEY.to_string(), compression.to_owned()); + } if let Some(extension_name) = options.options.get(EXTENSION_TYPE_NAME_KEY) { metadata.insert(EXTENSION_TYPE_NAME_KEY.to_string(), extension_name.clone()); } @@ -88,8 +94,14 @@ pub fn try_as_column_schema(column_def: &ColumnDef) -> Result { } } - ColumnSchema::new(&column_def.name, data_type.into(), column_def.is_nullable) - .with_metadata(metadata) + let schema = ColumnSchema::new(&column_def.name, data_type.into(), column_def.is_nullable) + .with_metadata(metadata); + schema + .compression_options() + .context(error::InvalidColumnCompressionSnafu { + column: &column_def.name, + })?; + schema .with_time_index(column_def.semantic_type() == SemanticType::Timestamp) .with_default_constraint(constraint) .context(error::InvalidColumnDefaultConstraintSnafu { @@ -172,6 +184,11 @@ pub fn options_from_column_schema(column_schema: &ColumnSchema) -> Option Skipp #[cfg(test)] mod tests { + use common_base::readable_size::ReadableSize; use datatypes::data_type::ConcreteDataType; - use datatypes::schema::{FulltextAnalyzer, FulltextBackend}; + use datatypes::schema::{CompressionOptions, FulltextAnalyzer, FulltextBackend}; use super::*; use crate::v1::ColumnDataType; @@ -307,6 +325,54 @@ mod tests { assert!(schema.is_inverted_indexed()); } + #[test] + fn test_compression_options_round_trip() { + let options = CompressionOptions { + level: Some(3), + page_size: Some(ReadableSize::mb(8)), + ..Default::default() + }; + let schema = ColumnSchema::new("test", ConcreteDataType::string_datatype(), true) + .with_compression_options(options) + .unwrap(); + + let column_def = try_as_column_def(&schema, false).unwrap(); + let schema = try_as_column_schema(&column_def).unwrap(); + assert_eq!(Some(options), schema.compression_options().unwrap()); + + let metadata = collect_column_options(column_def.options.as_ref()); + assert!(metadata.contains_key(COMPRESSION_KEY)); + } + + #[test] + fn test_reject_invalid_compression_metadata() { + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + + for (data_type, json) in [ + (ConcreteDataType::string_datatype(), r#"{"level":23}"#), + (ConcreteDataType::string_datatype(), r#"{"level":0}"#), + (ConcreteDataType::string_datatype(), r#"{"page_size":"0B"}"#), + (ConcreteDataType::string_datatype(), r#"{"type":"lz4"}"#), + (ConcreteDataType::string_datatype(), "not json"), + ( + ConcreteDataType::list_datatype(std::sync::Arc::new( + ConcreteDataType::string_datatype(), + )), + r#"{"level":3}"#, + ), + ] { + let schema = + ColumnSchema::new("test", data_type, true).with_metadata(HashMap::from([( + COMPRESSION_KEY.to_string(), + json.to_string(), + )])); + let column_def = try_as_column_def(&schema, false).unwrap(); + let err = try_as_column_schema(&column_def).unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code(), "{json}"); + } + } + #[test] fn test_options_from_column_schema() { let schema = ColumnSchema::new("test", ConcreteDataType::string_datatype(), true); diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index 32780270874..745268f6e35 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -285,6 +285,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid compression option: {}", msg))] + InvalidCompressionOption { + msg: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Inconsistent struct field count {field_len} and item count {item_len}"))] InconsistentStructFieldsAndItems { field_len: usize, @@ -360,7 +367,8 @@ impl ErrorExt for Error { | InvalidJsonb { .. } | InvalidVector { .. } | InvalidFulltextOption { .. } - | InvalidSkippingIndexOption { .. } => StatusCode::InvalidArguments, + | InvalidSkippingIndexOption { .. } + | InvalidCompressionOption { .. } => StatusCode::InvalidArguments, ValueExceedsPrecision { .. } | CastType { .. } diff --git a/src/datatypes/src/schema.rs b/src/datatypes/src/schema.rs index 0f1a832ebd2..66867f6141d 100644 --- a/src/datatypes/src/schema.rs +++ b/src/datatypes/src/schema.rs @@ -28,12 +28,14 @@ use snafu::{ResultExt, ensure}; use crate::error::{self, DuplicateColumnSnafu, Error, ProjectArrowSchemaSnafu, Result}; use crate::prelude::ConcreteDataType; pub use crate::schema::column_schema::{ - COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE, COLUMN_FULLTEXT_OPT_KEY_ANALYZER, - COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, - COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, - COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, - COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, ColumnExtType, ColumnSchema, FULLTEXT_KEY, - FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, + COLUMN_COMPRESSION_OPT_KEY_LEVEL, COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, + COLUMN_COMPRESSION_OPT_KEY_TYPE, COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE, + COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, + COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, + COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, + COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, + COMPRESSION_KEY, ColumnExtType, ColumnSchema, CompressionOptions, CompressionType, + FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, TIME_INDEX_KEY, }; pub use crate::schema::constraint::ColumnDefaultConstraint; diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index b3ce9cafedb..2d48ca72ad5 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -20,6 +20,7 @@ use arrow::datatypes::Field; use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; +use common_base::readable_size::ReadableSize; use serde::{Deserialize, Serialize}; use snafu::{ResultExt, ensure}; use sqlparser_derive::{Visit, VisitMut}; @@ -48,6 +49,8 @@ pub const FULLTEXT_KEY: &str = "greptime:fulltext"; pub const INVERTED_INDEX_KEY: &str = "greptime:inverted_index"; /// Key used to store skip options in arrow field's metadata. pub const SKIPPING_INDEX_KEY: &str = "greptime:skipping_index"; +/// Key used to store SST compression options in arrow field's metadata. +pub const COMPRESSION_KEY: &str = "greptime:compression"; /// Keys used in fulltext options pub const COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE: &str = "enable"; @@ -62,6 +65,11 @@ pub const COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY: &str = "granularity"; pub const COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE: &str = "false_positive_rate"; pub const COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE: &str = "type"; +/// Keys used in compression options +pub const COLUMN_COMPRESSION_OPT_KEY_TYPE: &str = "type"; +pub const COLUMN_COMPRESSION_OPT_KEY_LEVEL: &str = "level"; +pub const COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE: &str = "page_size"; + pub const DEFAULT_GRANULARITY: u32 = 10240; pub const DEFAULT_FALSE_POSITIVE_RATE: f64 = 0.01; @@ -427,6 +435,48 @@ impl ColumnSchema { Ok(()) } + /// Retrieves the SST compression options for the column. + pub fn compression_options(&self) -> Result> { + match self.metadata.get(COMPRESSION_KEY) { + None => Ok(None), + Some(json) => { + let options = serde_json::from_str(json).map_err(|e| { + error::InvalidCompressionOptionSnafu { + msg: format!("invalid compression metadata for column {}: {e}", self.name), + } + .build() + })?; + self.validate_compression_options(&options)?; + Ok(Some(options)) + } + } + } + + /// Sets the SST compression options. Nested types are rejected: their values are + /// stored in child columns that the options would not reach. + pub fn with_compression_options(mut self, options: CompressionOptions) -> Result { + self.validate_compression_options(&options)?; + self.metadata.insert( + COMPRESSION_KEY.to_string(), + serde_json::to_string(&options).context(error::SerializeSnafu)?, + ); + Ok(self) + } + + fn validate_compression_options(&self, options: &CompressionOptions) -> Result<()> { + options.validate()?; + ensure!( + !self.data_type.as_arrow_type().is_nested(), + error::InvalidCompressionOptionSnafu { + msg: format!( + "compression options are not supported on column {} of type {}", + self.name, self.data_type + ), + } + ); + Ok(()) + } + pub fn extension_type(&self) -> Result> where E: ExtensionType, @@ -952,6 +1002,141 @@ impl fmt::Display for SkippingIndexType { } } +/// Compression algorithms for SST columns. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CompressionType { + #[default] + Zstd, +} + +impl fmt::Display for CompressionType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + CompressionType::Zstd => write!(f, "zstd"), + } + } +} + +/// Parquet compression settings for a column in SSTs. Unset fields keep the SST defaults. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct CompressionOptions { + #[serde(rename = "type", default)] + pub compression_type: CompressionType, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub level: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub page_size: Option, +} + +impl CompressionOptions { + fn validate(&self) -> Result<()> { + if let Some(level) = self.level { + ensure!( + ZSTD_LEVELS.contains(&level), + error::InvalidCompressionOptionSnafu { + msg: format!("invalid level: {level}, expected: an integer in 1..=22"), + } + ); + } + if let Some(page_size) = self.page_size { + ensure!( + page_size.as_bytes() > 0, + error::InvalidCompressionOptionSnafu { + msg: "page_size must be positive", + } + ); + } + Ok(()) + } + + /// Returns the options as `COMPRESSION WITH (...)` key-value pairs. + pub fn to_options_map(&self) -> HashMap { + let mut map = HashMap::from([( + COLUMN_COMPRESSION_OPT_KEY_TYPE.to_string(), + self.compression_type.to_string(), + )]); + if let Some(level) = self.level { + map.insert( + COLUMN_COMPRESSION_OPT_KEY_LEVEL.to_string(), + level.to_string(), + ); + } + if let Some(page_size) = self.page_size { + // `Display` rounds to one decimal; the serde form is exact and parses back. + let page_size = match serde_json::to_value(page_size) { + Ok(serde_json::Value::String(size)) => size, + _ => page_size.as_bytes().to_string(), + }; + map.insert(COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE.to_string(), page_size); + } + map + } +} + +/// Valid zstd levels, as accepted by the parquet writer. +const ZSTD_LEVELS: std::ops::RangeInclusive = 1..=22; + +impl TryFrom> for CompressionOptions { + type Error = Error; + + fn try_from(options: HashMap) -> Result { + let compression_type = match options.get(COLUMN_COMPRESSION_OPT_KEY_TYPE) { + Some(value) if value.eq_ignore_ascii_case("zstd") => CompressionType::Zstd, + Some(value) => { + return error::InvalidCompressionOptionSnafu { + msg: format!("invalid type: {value}, expected: 'zstd'"), + } + .fail(); + } + None => CompressionType::default(), + }; + + let level = options + .get(COLUMN_COMPRESSION_OPT_KEY_LEVEL) + .map(|value| { + value + .parse::() + .ok() + .filter(|level| ZSTD_LEVELS.contains(level)) + .ok_or_else(|| { + error::InvalidCompressionOptionSnafu { + msg: format!( + "invalid level: {value}, expected: an integer in {}..={}", + ZSTD_LEVELS.start(), + ZSTD_LEVELS.end() + ), + } + .build() + }) + }) + .transpose()?; + + let page_size = options + .get(COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE) + .map(|value| { + ReadableSize::from_str(value) + .ok() + .filter(|size| size.as_bytes() > 0) + .ok_or_else(|| { + error::InvalidCompressionOptionSnafu { + msg: format!( + "invalid page_size: {value}, expected: a positive size such as '8MiB'" + ), + } + .build() + }) + }) + .transpose()?; + + Ok(CompressionOptions { + compression_type, + level, + page_size, + }) + } +} + impl TryFrom> for SkippingIndexOptions { type Error = Error; @@ -1323,6 +1508,96 @@ mod tests { assert!(column_schema.metadata.is_empty()); } + #[test] + fn test_compression_options_from_map() { + let parse = |pairs: &[(&str, &str)]| { + CompressionOptions::try_from( + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect::>(), + ) + }; + + let options = parse(&[("type", "ZSTD"), ("level", "3"), ("page_size", "8MiB")]).unwrap(); + assert_eq!( + CompressionOptions { + compression_type: CompressionType::Zstd, + level: Some(3), + page_size: Some(ReadableSize::mb(8)), + }, + options + ); + // Stored in column metadata; this format must stay readable. + assert_eq!( + r#"{"type":"zstd","level":3,"page_size":"8MiB"}"#, + serde_json::to_string(&options).unwrap() + ); + assert_eq!( + options, + parse( + &options + .to_options_map() + .iter() + .map(|(k, v)| (k.as_str(), v.as_str())) + .collect::>() + ) + .unwrap() + ); + + let odd_size = parse(&[("page_size", "1.25MiB")]).unwrap(); + assert_eq!( + odd_size, + parse(&[("page_size", odd_size.to_options_map()["page_size"].as_str())]).unwrap() + ); + + let page_only = parse(&[("page_size", "1MiB")]).unwrap(); + assert_eq!(None, page_only.level); + assert_eq!( + r#"{"type":"zstd","page_size":"1MiB"}"#, + serde_json::to_string(&page_only).unwrap() + ); + + for invalid in [ + &[("type", "lz4")][..], + &[("level", "0")], + &[("level", "23")], + &[("level", "high")], + &[("page_size", "0")], + &[("page_size", "large")], + ] { + assert!(parse(invalid).is_err(), "{invalid:?}"); + } + } + + #[test] + fn test_compression_options_reject_nested_types() { + let options = CompressionOptions { + level: Some(3), + ..Default::default() + }; + let string_column = ColumnSchema::new("s", ConcreteDataType::string_datatype(), true) + .with_compression_options(options) + .unwrap(); + assert_eq!(Some(options), string_column.compression_options().unwrap()); + + let struct_type = + ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new( + "a".to_string(), + ConcreteDataType::string_datatype(), + true, + )]))); + let list_type = + ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::string_datatype())); + for data_type in [struct_type, list_type] { + assert!( + ColumnSchema::new("n", data_type, true) + .with_compression_options(options) + .is_err() + ); + } + } + #[test] fn test_skipping_index_options_deserialization() { let original_options = "{\"granularity\":1024,\"false-positive-rate-in-10000\":10,\"index-type\":\"BloomFilter\"}"; diff --git a/src/metric-engine/src/engine/alter/validate.rs b/src/metric-engine/src/engine/alter/validate.rs index 2e2d91eccfc..cd7d8f2f172 100644 --- a/src/metric-engine/src/engine/alter/validate.rs +++ b/src/metric-engine/src/engine/alter/validate.rs @@ -12,11 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +use datatypes::schema::COMPRESSION_KEY; use snafu::ensure; use store_api::region_request::{AlterKind, RegionAlterRequest}; use store_api::storage::RegionId; -use crate::error::{Result, UnsupportedAlterKindSnafu}; +use crate::error::{Result, UnexpectedRequestSnafu, UnsupportedAlterKindSnafu}; /// Validate the alter region requests. pub fn validate_alter_region_requests(requests: &[(RegionId, RegionAlterRequest)]) -> Result<()> { @@ -27,7 +28,58 @@ pub fn validate_alter_region_requests(requests: &[(RegionId, RegionAlterRequest) kind: request.kind.as_ref() } ); + if let AlterKind::AddColumns { columns } = &request.kind { + for column in columns { + ensure!( + !column.column_metadata.column_schema.metadata().contains_key(COMPRESSION_KEY), + UnexpectedRequestSnafu { + reason: "COMPRESSION is not supported on logical tables; configure it on the physical table".to_string(), + } + ); + } + } } Ok(()) } + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, CompressionOptions}; + use store_api::metadata::ColumnMetadata; + use store_api::region_request::AddColumn; + + use super::*; + + #[test] + fn test_reject_logical_column_compression() { + let column_schema = ColumnSchema::new("host", ConcreteDataType::string_datatype(), true) + .with_compression_options(CompressionOptions { + level: Some(3), + ..Default::default() + }) + .unwrap(); + let request = RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema, + semantic_type: SemanticType::Tag, + column_id: 3, + }, + location: None, + }], + }, + }; + let err = validate_alter_region_requests(&[(RegionId::new(1, 1), request)]).unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!( + err.to_string() + .contains("configure it on the physical table") + ); + } +} diff --git a/src/metric-engine/src/engine/create.rs b/src/metric-engine/src/engine/create.rs index 952643cd7a4..ec35c0653c4 100644 --- a/src/metric-engine/src/engine/create.rs +++ b/src/metric-engine/src/engine/create.rs @@ -21,7 +21,7 @@ use common_query::native_histogram::is_native_histogram_value_type; use common_telemetry::info; use common_time::{FOREVER, Timestamp}; use datatypes::data_type::ConcreteDataType; -use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; +use datatypes::schema::{COMPRESSION_KEY, ColumnSchema, SkippingIndexOptions}; use datatypes::value::Value; use mito2::engine::MITO_ENGINE_NAME; use snafu::{OptionExt, ResultExt, ensure}; @@ -391,6 +391,17 @@ impl MetricEngineInner { ConflictRegionOptionSnafu {} ); + if request.options.contains_key(LOGICAL_TABLE_METADATA_KEY) { + for column in &request.column_metadatas { + ensure!( + !column.column_schema.metadata().contains_key(COMPRESSION_KEY), + UnexpectedRequestSnafu { + reason: "COMPRESSION is not supported on logical tables; configure it on the physical table".to_string(), + } + ); + } + } + // check if field columns are either a normal metric value or a native histogram. let mut field_cols = Vec::new(); for col in &request.column_metadatas { diff --git a/src/mito2/src/engine/create_test.rs b/src/mito2/src/engine/create_test.rs index 2c04e2e24f1..00c3686c764 100644 --- a/src/mito2/src/engine/create_test.rs +++ b/src/mito2/src/engine/create_test.rs @@ -30,6 +30,74 @@ use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, }; +#[tokio::test] +async fn test_engine_column_compression_validation() { + use datatypes::schema::COMPRESSION_KEY; + use store_api::region_request::{AlterKind, RegionAlterRequest, SetRegionOption}; + + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + for (column_index, compression, options, accepted) in [ + (1, r#"{"level":23}"#, vec![], false), + ( + 0, + r#"{"level":3}"#, + vec![ + ("sst_format", "primary_key"), + ("memtable.type", "time_series"), + ], + false, + ), + ( + 0, + r#"{"level":3}"#, + vec![("primary_key_encoding", "sparse")], + false, + ), + (0, r#"{"level":3}"#, vec![], true), + ] { + let mut request = CreateRequestBuilder::new().build(); + request.options.extend( + options + .into_iter() + .map(|(k, v)| (k.to_string(), v.to_string())), + ); + let column = &mut request.column_metadatas[column_index].column_schema; + *column = column + .clone() + .with_metadata(std::collections::HashMap::from([( + COMPRESSION_KEY.to_string(), + compression.to_string(), + )])); + let result = engine + .handle_request(region_id, RegionRequest::Create(request)) + .await; + if accepted { + result.unwrap(); + } else { + assert_eq!( + StatusCode::InvalidArguments, + result.unwrap_err().status_code() + ); + assert!(engine.get_region(region_id).is_none()); + } + } + + let err = engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::Format("primary_key".to_string())], + }, + }), + ) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); +} + #[tokio::test] async fn test_engine_create_new_region() { test_engine_create_new_region_with_format(false).await; diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 5af4b7cec09..54bf1ba6b6e 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -1094,6 +1094,14 @@ pub enum Error { column_name: String, }, + #[snafu(display("Failed to retrieve compression options of column {column_name}"))] + CompressionOptions { + #[snafu(implicit)] + location: Location, + source: datatypes::error::Error, + column_name: String, + }, + #[snafu(display("Failed to create fulltext index creator"))] CreateFulltextCreator { source: index::fulltext_index::error::Error, @@ -1586,6 +1594,7 @@ impl ErrorExt for Error { RemoteCompaction { .. } => StatusCode::Unexpected, IndexOptions { source, .. } => source.status_code(), + CompressionOptions { source, .. } => source.status_code(), CreateFulltextCreator { source, .. } => source.status_code(), CastVector { source, .. } => source.status_code(), FulltextPushText { source, .. } @@ -1691,6 +1700,7 @@ impl ErrorExt for Error { | ConvertVector { source, .. } | ConvertValue { source, .. } | IndexOptions { source, .. } + | CompressionOptions { source, .. } | CastVector { source, .. } => source.retry_hint(), BuildIndexApplier { source, .. } diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index fecfceafeee..27fc7983513 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -73,7 +73,8 @@ use crate::sst::index::IndexOutput; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyArrayBuilder}; use crate::sst::parquet::{ - COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding, + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, apply_column_compression, + apply_float_field_encoding, }; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -1347,6 +1348,7 @@ impl BulkPartEncoder { .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); props = apply_float_field_encoding(props, &metadata, float_field_encoding); + props = apply_column_compression(props, &metadata)?; let writer_props = Some(props.build()); Ok(Self { diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 71c1192cbf2..7a93799ae5f 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -384,6 +384,7 @@ impl RegionOpener { FormatType::PrimaryKey }; ensure_json2_not_use_time_series_memtable(&metadata, &options)?; + options.validate_column_compression(&metadata)?; // Create a manifest manager for this region and writes regions to the manifest file. let mut region_manifest_options = RegionManifestOptions::new(config, ®ion_dir, &object_store); @@ -539,6 +540,7 @@ impl RegionOpener { // Updates the region options with the manifest. sanitize_region_options(&manifest, &mut region_options); ensure_json2_not_use_time_series_memtable(&metadata, ®ion_options)?; + region_options.validate_column_compression(&metadata)?; let region_id = self.region_id; let provider = self.provider::(®ion_options.wal_options)?; diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 89b413a5d8f..7e050fb4a76 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -19,16 +19,19 @@ use std::collections::HashMap; use std::time::Duration; +use api::v1::SemanticType; use common_base::readable_size::ReadableSize; use common_telemetry::info; use common_time::TimeToLive; use common_wal::options::{WAL_OPTIONS_KEY, WalOptions}; +use datatypes::schema::COMPRESSION_KEY; use serde::de::Error as _; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::Value; use serde_with::{DisplayFromStr, NoneAsEmptyString, serde_as, with_prefix}; use snafu::{ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::RegionMetadata; use store_api::metric_engine_consts::{ MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING, PRIMARY_KEY_ENCODING, }; @@ -194,6 +197,32 @@ impl RegionOptions { Ok(()) } + /// Rejects tag compression when the SST layout packs tags into `__primary_key`. + /// The caller must resolve the SST format before validation. + pub(crate) fn validate_column_compression(&self, metadata: &RegionMetadata) -> Result<()> { + if self.sst_format == Some(FormatType::Flat) + && metadata.primary_key_encoding == PrimaryKeyEncoding::Dense + { + return Ok(()); + } + for column in &metadata.column_metadatas { + ensure!( + column.semantic_type != SemanticType::Tag + || !column + .column_schema + .metadata() + .contains_key(COMPRESSION_KEY), + InvalidRegionOptionsSnafu { + reason: format!( + "COMPRESSION on tag column {} requires flat SST format and dense primary key encoding", + column.column_schema.name + ), + } + ); + } + Ok(()) + } + /// Returns the configured row group size, falling back to [DEFAULT_ROW_GROUP_SIZE]. pub fn row_group_size(&self) -> usize { self.max_row_group_row_count diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 9404d4484fc..e6968447dce 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -20,13 +20,16 @@ use std::sync::Arc; use api::v1::SemanticType; use common_base::readable_size::ReadableSize; use datatypes::json::JsonSettings; +use parquet::basic::{Compression, ZstdLevel}; use parquet::file::metadata::ParquetMetaData; -use parquet::file::properties::WriterPropertiesBuilder; +use parquet::file::properties::{DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT, WriterPropertiesBuilder}; use parquet::schema::types::ColumnPath; +use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; use store_api::mito_engine_options::FloatFieldEncoding; use store_api::storage::{ColumnId, FileId}; +use crate::error::{CompressionOptionsSnafu, Result, WriteParquetSnafu}; use crate::sst::DEFAULT_WRITE_BUFFER_SIZE; use crate::sst::file::FileTimeRange; use crate::sst::index::IndexOutput; @@ -89,6 +92,41 @@ pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option = parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; +/// Applies the per-column options declared with `COMPRESSION WITH (...)`. +pub(crate) fn apply_column_compression( + mut builder: WriterPropertiesBuilder, + metadata: &RegionMetadataRef, +) -> Result { + for column in &metadata.column_metadatas { + let Some(options) = + column + .column_schema + .compression_options() + .context(CompressionOptionsSnafu { + column_name: &column.column_schema.name, + })? + else { + continue; + }; + let path = ColumnPath::new(vec![column.column_schema.name.clone()]); + if let Some(level) = options.level { + let level = ZstdLevel::try_new(level).context(WriteParquetSnafu)?; + builder = builder.set_column_compression(path.clone(), Compression::ZSTD(level)); + } + if let Some(page_size) = options.page_size { + let page_size = page_size.as_bytes() as usize; + builder = builder.set_column_data_page_size_limit(path.clone(), page_size); + // While dictionary encoding, parquet sizes write batches by the dictionary page + // budget and keeps that size for the rest of the batch after falling back to plain + // encoding, so a data page limit below the dictionary limit would be ignored. + if page_size < DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT { + builder = builder.set_column_dictionary_page_size_limit(path, page_size); + } + } + } + Ok(builder) +} + /// Applies the configured encoding to direct floating-point field columns. pub(crate) fn apply_float_field_encoding( mut builder: WriterPropertiesBuilder, @@ -167,6 +205,7 @@ mod tests { use api::v1::{OpType, SemanticType}; use bytes::Bytes; + use common_base::readable_size::ReadableSize; use common_function::function::FunctionRef; use common_function::function_factory::ScalarFunctionFactory; use common_function::scalars::matches::MatchesFunction; @@ -184,7 +223,9 @@ mod tests { use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; use datatypes::arrow::util::pretty::pretty_format_batches; use datatypes::prelude::ConcreteDataType; - use datatypes::schema::{FulltextAnalyzer, FulltextBackend, FulltextOptions}; + use datatypes::schema::{ + CompressionOptions, FulltextAnalyzer, FulltextBackend, FulltextOptions, + }; use object_store::ObjectStore; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::arrow::{ArrowWriter, AsyncArrowWriter}; @@ -820,6 +861,138 @@ mod tests { .await; } + #[tokio::test] + async fn test_column_compression_options() { + // Text-like values: pseudo-random words from a small vocabulary, so zstd levels + // compress differently and the dictionary falls back to plain encoding. + let words: Vec = (0..256).map(|i| format!("word{i:03}")).collect(); + let mut seed = 42u64; + let values: Vec> = (0..2048) + .map(|_| { + (0..250) + .map(|_| { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + words[(seed >> 33) as usize % words.len()].as_str() + }) + .collect::>() + .join(" ") + .into_bytes() + }) + .collect(); + + async fn write( + values: &[Vec], + compression: Option, + ) -> (usize, usize) { + let mut metadata = (*build_test_binary_test_region_metadata()).clone(); + if let Some(options) = compression { + let column = &mut metadata.column_metadatas[1]; + assert_eq!("field_0", column.column_schema.name); + column.column_schema = column + .column_schema + .clone() + .with_compression_options(options) + .unwrap(); + } + let metadata = Arc::new(metadata); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + let num_rows = values.len(); + let pk = new_primary_key(&["a"]); + let mut tags = StringDictionaryBuilder::::new(); + let mut pks = BinaryDictionaryBuilder::::new(); + for _ in 0..num_rows { + tags.append_value("a"); + pks.append(&pk).unwrap(); + } + let batch = RecordBatch::try_new( + flat_schema, + vec![ + Arc::new(tags.finish()), + Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values( + values.iter().map(|v| v.as_slice()), + )), + Arc::new(TimestampMillisecondArray::from_iter_values( + 0..num_rows as i64, + )), + Arc::new(pks.finish()), + Arc::new(UInt64Array::from_value(1000, num_rows)), + Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)), + ], + ) + .unwrap(); + + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let handle = sst_file_handle(0, 1000); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store.clone(), + metadata, + IndexConfig::default(), + NoopIndexBuilder, + FixedPathProvider { + region_file_id: handle.file_id(), + }, + &mut metrics, + ) + .await; + writer + .write_all_flat_as_primary_key( + new_flat_source_from_record_batches(vec![batch]), + None, + &WriteOptions::default(), + ) + .await + .unwrap(); + + let bytes = object_store + .read(&handle.file_path(FILE_DIR, PathType::Bare)) + .await + .unwrap() + .to_bytes(); + let options = parquet::arrow::arrow_reader::ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options).unwrap(); + let schema = builder.metadata().file_metadata().schema_descr(); + let field = (0..schema.num_columns()) + .find(|i| schema.column(*i).name() == "field_0") + .unwrap(); + let pages = builder.metadata().offset_index().unwrap()[0][field] + .page_locations() + .len(); + let size = builder + .metadata() + .row_group(0) + .column(field) + .compressed_size() as usize; + (size, pages) + } + + let (default_size, default_pages) = write(&values, None).await; + let (level_size, level_pages) = write( + &values, + Some(CompressionOptions { + level: Some(19), + ..Default::default() + }), + ) + .await; + let (_, pages) = write( + &values, + Some(CompressionOptions { + page_size: Some(ReadableSize::kb(64)), + ..Default::default() + }), + ) + .await; + + assert!(default_pages > 1, "{default_pages}"); + assert_eq!(default_pages, level_pages); + assert!(level_size < default_size, "{level_size} >= {default_size}"); + assert!(pages > default_pages, "{pages} <= {default_pages}"); + } + #[tokio::test] async fn test_read_large_binary() { let mut env = TestEnv::new().await; diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index 932654dfdd8..9d2defcceb3 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -63,7 +63,7 @@ use crate::sst::parquet::flat_format::{ use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{ COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, - apply_float_field_encoding, + apply_column_compression, apply_float_field_encoding, }; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, @@ -563,6 +563,7 @@ where &self.metadata, opts.float_field_encoding, ); + let props_builder = apply_column_compression(props_builder, &self.metadata)?; let writer_props = props_builder.build(); let sst_file_path = self.path_provider.build_sst_file_path(RegionFileId::new( diff --git a/src/mito2/src/worker/handle_alter.rs b/src/mito2/src/worker/handle_alter.rs index 66cb2dd95bb..2bb1fadc7dc 100644 --- a/src/mito2/src/worker/handle_alter.rs +++ b/src/mito2/src/worker/handle_alter.rs @@ -197,6 +197,10 @@ impl RegionWorkerLoop { }; // Persist the metadata to region's manifest. let options = new_options.as_ref().unwrap_or(&version.options); + if let Err(e) = options.validate_column_compression(&new_meta) { + sender.send(Err(e)); + return; + } let change = RegionChange { metadata: new_meta, sst_format: options.sst_format.unwrap_or_default(), @@ -356,6 +360,15 @@ impl RegionWorkerLoop { } .build() })?; + candidate + .validate_column_compression(&version.metadata) + .map_err(|e| { + store_api::metadata::InvalidRegionRequestSnafu { + region_id: region.region_id, + err: e.to_string(), + } + .build() + })?; if all_options_altered { region.version_control.alter_options(candidate); Ok(None) diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index d69277a1c90..935ce149414 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -375,6 +375,14 @@ pub fn expr_to_create(expr: &CreateTableExpr, quote_style: Option) -> Resu extensions.inverted_index_options = Some(HashMap::new().into()); } + if let Some(opt) = column_schema + .compression_options() + .map_err(BoxedError::new) + .context(ExternalSnafu)? + { + extensions.compression_options = Some(opt.to_options_map().into()); + } + let sql_column = SqlColumn { column_def: ColumnDef { name: Ident::with_quote(quote_style, &column_def.name), @@ -2117,7 +2125,7 @@ ALTER TABLE metrics PARTITION ON COLUMNS (device_id, area) ( `protocol` STRING NULL, `status_code` INT NULL INVERTED INDEX, `response_size` BIGINT NULL, - `message` STRING NULL, + `message` STRING NULL COMPRESSION WITH(level = '3', page_size = '64KiB', type = 'zstd'), TIME INDEX (`timestamp`), PRIMARY KEY (`username`, `status_code`) ) diff --git a/src/query/src/sql/show_create_table.rs b/src/query/src/sql/show_create_table.rs index c56e83d385a..b9b902a0468 100644 --- a/src/query/src/sql/show_create_table.rs +++ b/src/query/src/sql/show_create_table.rs @@ -169,6 +169,10 @@ fn create_column(column_schema: &ColumnSchema, quote_style: char) -> Result StatusCode::PermissionDenied, SetFulltextOption { .. } | SetSkippingIndexOption { .. } => StatusCode::Unexpected, + SetCompressionOption { source, .. } => source.status_code(), } } diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index c87d0274382..20eba9f2087 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -47,8 +47,8 @@ use crate::error::{ use crate::parser::{FLOW, ParserContext}; use crate::parsers::tql_parser; use crate::parsers::utils::{ - self, parse_with_options, validate_column_fulltext_create_option, - validate_column_skipping_index_create_option, + self, parse_with_options, validate_column_compression_create_option, + validate_column_fulltext_create_option, validate_column_skipping_index_create_option, }; use crate::statements::create::{ Column, ColumnExtensions, CreateDatabase, CreateExternalTable, CreateFlow, CreateTable, @@ -891,6 +891,7 @@ impl<'a> ParserContext<'a> { /// This function will handle: /// - Vector type /// - Indexes + /// - SST compression options fn parse_column_extensions( parser: &mut Parser<'_>, column_name: &Ident, @@ -923,7 +924,7 @@ impl<'a> ParserContext<'a> { } // parse index options in column definition - let mut is_index_declared = false; + let mut is_extension_declared = false; // skipping index if let Token::Word(word) = parser.peek_token().token @@ -965,7 +966,7 @@ impl<'a> ParserContext<'a> { let options = OptionMap::new(options); column_extensions.skipping_index_options = Some(options); - is_index_declared |= true; + is_extension_declared |= true; } // fulltext index @@ -1016,7 +1017,7 @@ impl<'a> ParserContext<'a> { let options = OptionMap::new(options); column_extensions.fulltext_index_options = Some(options); - is_index_declared |= true; + is_extension_declared |= true; } // inverted index @@ -1058,10 +1059,54 @@ impl<'a> ParserContext<'a> { ); column_extensions.inverted_index_options = Some(OptionMap::default()); - is_index_declared |= true; + is_extension_declared |= true; } - Ok(is_index_declared) + // SST compression + if parser.parse_keyword(Keyword::COMPRESSION) { + ensure!( + column_extensions.compression_options.is_none(), + InvalidColumnOptionSnafu { + name: column_name.to_string(), + msg: "duplicated COMPRESSION option", + } + ); + + let options = parser + .parse_options(Keyword::WITH) + .context(error::SyntaxSnafu)? + .into_iter() + .map(parse_option_string) + .collect::>>()?; + ensure!( + !options.is_empty(), + InvalidColumnOptionSnafu { + name: column_name.to_string(), + msg: "expect WITH (...) after COMPRESSION keyword", + } + ); + for (key, value) in options.iter() { + ensure!( + value.as_string().is_some(), + InvalidColumnOptionSnafu { + name: column_name.to_string(), + msg: format!("COMPRESSION option {key} requires a scalar value"), + } + ); + ensure!( + validate_column_compression_create_option(key), + InvalidColumnOptionSnafu { + name: column_name.to_string(), + msg: format!("invalid COMPRESSION option: {key}"), + } + ); + } + + column_extensions.compression_options = Some(OptionMap::new(options)); + is_extension_declared |= true; + } + + Ok(is_extension_declared) } fn parse_optional_table_constraint(&mut self) -> Result> { @@ -3194,6 +3239,82 @@ CREATE TABLE log ( } } + #[test] + fn test_parse_create_table_compression_options() { + let parse = |sql: &str| { + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + }; + + for value in ["[3]", "NULL", "['zstd']"] { + for key in ["type", "level", "page_size"] { + let sql = format!( + "CREATE TABLE t (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH ({key}={value}))" + ); + let err = parse(&sql).unwrap_err(); + assert!( + err.to_string().contains("requires a scalar value"), + "{sql}: {err}" + ); + } + } + + let sql = r" +CREATE TABLE spans ( + ts TIMESTAMP TIME INDEX, + prompt STRING COMPRESSION WITH (type='zstd', level='3', page_size='8MiB') SKIPPING INDEX, + reply STRING NULL COMPRESSION WITH (level='9'), +)"; + let result = parse(sql).unwrap(); + let Statement::CreateTable(c) = &result[0] else { + panic!("should be create_table statement"); + }; + let prompt = &c.columns[1]; + assert_eq!( + vec!["level=3", "page_size=8MiB", "type=zstd"], + prompt + .extensions + .compression_options + .as_ref() + .unwrap() + .kv_pairs() + .into_iter() + .map(|kv| kv.replace(['\'', ' '], "")) + .sorted() + .collect::>() + ); + assert!(prompt.extensions.skipping_index_options.is_some()); + // Display must produce SQL that parses back to the same column. + let reparsed = parse(&format!( + "CREATE TABLE t (ts TIMESTAMP TIME INDEX, {prompt})" + )) + .unwrap(); + let Statement::CreateTable(r) = &reparsed[0] else { + panic!("should be create_table statement"); + }; + assert_eq!(prompt.extensions, r.columns[1].extensions); + + for (sql, err) in [ + ( + "CREATE TABLE t (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION)", + "expect WITH (...) after COMPRESSION keyword", + ), + ( + "CREATE TABLE t (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (codec='zstd'))", + "invalid COMPRESSION option: codec", + ), + ( + "CREATE TABLE t (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (level='3') COMPRESSION WITH (level='5'))", + "duplicated COMPRESSION option", + ), + ] { + let result = parse(sql); + assert!( + result.as_ref().is_err_and(|e| e.to_string().contains(err)), + "{sql}: {result:?}" + ); + } + } + #[test] fn test_parse_create_view_with_columns() { let sql = "CREATE VIEW test () AS SELECT * FROM NUMBERS"; diff --git a/src/sql/src/parsers/utils.rs b/src/sql/src/parsers/utils.rs index afba437045c..65744708232 100644 --- a/src/sql/src/parsers/utils.rs +++ b/src/sql/src/parsers/utils.rs @@ -28,10 +28,12 @@ use datafusion_expr::{AggregateUDF, Expr, HigherOrderUDF, ScalarUDF, TableSource use datafusion_sql::planner::{ContextProvider, SqlToRel}; use datatypes::arrow::datatypes::DataType; use datatypes::schema::{ - COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, - COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, - COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, - COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, + COLUMN_COMPRESSION_OPT_KEY_LEVEL, COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, + COLUMN_COMPRESSION_OPT_KEY_TYPE, COLUMN_FULLTEXT_OPT_KEY_ANALYZER, + COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, + COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, + COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, + COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, }; use snafu::{ResultExt, ensure}; use sqlparser::dialect::Dialect; @@ -375,6 +377,15 @@ pub fn validate_column_skipping_index_create_option(key: &str) -> bool { .contains(&key) } +pub fn validate_column_compression_create_option(key: &str) -> bool { + [ + COLUMN_COMPRESSION_OPT_KEY_TYPE, + COLUMN_COMPRESSION_OPT_KEY_LEVEL, + COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, + ] + .contains(&key) +} + /// Convert an [`IntervalMonthDayNano`] to a [`Duration`]. #[cfg(feature = "enterprise")] pub fn convert_month_day_nano_to_duration( diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index c2c97c34567..34a4c483cae 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -58,8 +58,8 @@ use crate::ast::{ }; use crate::error::{ self, ConvertToGrpcDataTypeSnafu, ConvertValueSnafu, Result, - SerializeColumnDefaultConstraintSnafu, SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, - SqlCommonSnafu, + SerializeColumnDefaultConstraintSnafu, SetCompressionOptionSnafu, SetFulltextOptionSnafu, + SetSkippingIndexOptionSnafu, SqlCommonSnafu, }; use crate::statements::create::Column; pub use crate::statements::option_map::OptionMap; @@ -149,6 +149,12 @@ pub fn column_to_schema( .context(SetSkippingIndexOptionSnafu)?; } + if let Some(options) = column.extensions.build_compression_options()? { + column_schema = column_schema + .with_compression_options(options) + .context(SetCompressionOptionSnafu)?; + } + column_schema.set_inverted_index(column.extensions.inverted_index_options.is_some()); set_json2_extension(&mut column_schema, column)?; diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 919c1cd31f9..13fb4cb0a3a 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -18,7 +18,7 @@ use std::fmt::{Display, Formatter, Write}; use common_catalog::consts::FILE_ENGINE; use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings}; use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; +use datatypes::schema::{CompressionOptions, FulltextOptions, SkippingIndexOptions}; use itertools::Itertools; use serde::Serialize; use snafu::ResultExt; @@ -27,8 +27,8 @@ use sqlparser_derive::{Visit, VisitMut}; use crate::ast::{ColumnDef, Ident, ObjectName}; use crate::error::{ - InvalidFlowQuerySnafu, InvalidSqlSnafu, Result, SetFulltextOptionSnafu, - SetSkippingIndexOptionSnafu, + InvalidFlowQuerySnafu, InvalidSqlSnafu, Result, SetCompressionOptionSnafu, + SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, }; use crate::statements::query::Query as GtQuery; use crate::statements::statement::Statement; @@ -160,6 +160,8 @@ pub struct ColumnExtensions { /// /// Inverted index doesn't have options at present. There won't be any options in that map. pub inverted_index_options: Option, + /// SST compression options. + pub compression_options: Option, /// JSON2-specific column options. pub json2_options: Option, } @@ -294,6 +296,11 @@ impl Column { } } + if let Some(compression_options) = &self.extensions.compression_options { + let options = compression_options.kv_pairs(); + write!(f, " COMPRESSION WITH({})", format_list_comma!(options))?; + } + Ok(()) } } @@ -319,6 +326,15 @@ impl ColumnExtensions { )) } + pub fn build_compression_options(&self) -> Result> { + let Some(options) = self.compression_options.as_ref() else { + return Ok(None); + }; + + let options: HashMap = options.clone().into_map(); + Ok(Some(options.try_into().context(SetCompressionOptionSnafu)?)) + } + pub fn build_json_settings(&self) -> Result> { let Some(options) = &self.json2_options else { return Ok(None); diff --git a/src/store-api/src/metadata.rs b/src/store-api/src/metadata.rs index 6983bafb504..55ff240ee4a 100644 --- a/src/store-api/src/metadata.rs +++ b/src/store-api/src/metadata.rs @@ -549,6 +549,10 @@ impl RegionMetadata { /// Checks whether it is a valid column. fn validate_column_metadata(column_metadata: &ColumnMetadata) -> Result<()> { + column_metadata + .column_schema + .compression_options() + .context(InvalidSchemaSnafu)?; if column_metadata.semantic_type == SemanticType::Timestamp { ensure!( column_metadata.column_schema.data_type.is_timestamp(), diff --git a/tests/cases/standalone/common/create/create_with_compression.result b/tests/cases/standalone/common/create/create_with_compression.result new file mode 100644 index 00000000000..8b0e25a7c0f --- /dev/null +++ b/tests/cases/standalone/common/create/create_with_compression.result @@ -0,0 +1,218 @@ +CREATE TABLE spans ( + ts TIMESTAMP TIME INDEX, + trace_id STRING SKIPPING INDEX, + prompt STRING COMPRESSION WITH (type = 'zstd', level = 3, page_size = '8MiB'), + reply STRING COMPRESSION WITH (level = 9), + tokens BIGINT COMPRESSION WITH (page_size = '64KiB'), +) WITH (append_mode = 'true'); + +Affected Rows: 0 + +SHOW CREATE TABLE spans; + ++-------+--------------------------------------------------------------------------------------------------------------------+ +| Table | Create Table | ++-------+--------------------------------------------------------------------------------------------------------------------+ +| spans | CREATE TABLE IF NOT EXISTS "spans" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "trace_id" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'), | +| | "prompt" STRING NULL COMPRESSION WITH(level = '3', page_size = '8MiB', type = 'zstd'), | +| | "reply" STRING NULL COMPRESSION WITH(level = '9', type = 'zstd'), | +| | "tokens" BIGINT NULL COMPRESSION WITH(page_size = '64KiB', type = 'zstd'), | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++-------+--------------------------------------------------------------------------------------------------------------------+ + +INSERT INTO spans VALUES + (1, 't1', 'fix the flaky test', 'running it', 120), + (2, 't1', 'fix the flaky test\nrunning it', 'patched', 340), + (3, 't2', 'another prompt', NULL, NULL); + +Affected Rows: 3 + +ADMIN flush_table('spans'); + ++----------------------------+ +| ADMIN flush_table('spans') | ++----------------------------+ +| 0 | ++----------------------------+ + +SELECT ts, trace_id, prompt, reply, tokens FROM spans ORDER BY ts; + ++-------------------------+----------+--------------------------------+------------+--------+ +| ts | trace_id | prompt | reply | tokens | ++-------------------------+----------+--------------------------------+------------+--------+ +| 1970-01-01T00:00:00.001 | t1 | fix the flaky test | running it | 120 | +| 1970-01-01T00:00:00.002 | t1 | fix the flaky test\nrunning it | patched | 340 | +| 1970-01-01T00:00:00.003 | t2 | another prompt | | | ++-------------------------+----------+--------------------------------+------------+--------+ + +DROP TABLE spans; + +Affected Rows: 0 + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (level = 0)); + +Error: 1004(InvalidArguments), Invalid compression option: invalid level: 0, expected: an integer in 1..=22 + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (type = 'lz4')); + +Error: 1004(InvalidArguments), Invalid compression option: invalid type: lz4, expected: 'zstd' + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (page_size = 'large')); + +Error: 1004(InvalidArguments), Invalid compression option: invalid page_size: large, expected: a positive size such as '8MiB' + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (codec = 'zstd')); + +Error: 1004(InvalidArguments), Invalid column option, column name: s, error: invalid COMPRESSION option: codec + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, j JSON2 COMPRESSION WITH (level = 3)); + +Error: 1004(InvalidArguments), Invalid compression option: compression options are not supported on column j of type Json2"" + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (level = [3])); + +Error: 1004(InvalidArguments), Invalid column option, column name: s, error: COMPRESSION option level requires a scalar value + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (page_size = NULL)); + +Error: 1004(InvalidArguments), Invalid column option, column name: s, error: COMPRESSION option page_size requires a scalar value + +CREATE TABLE bad ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3) +) WITH ('sst_format' = 'primary_key', 'memtable.type' = 'time_series'); + +Error: 1004(InvalidArguments), Invalid region options, COMPRESSION on tag column host requires flat SST format and dense primary key encoding + +CREATE TABLE bad ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3) +) WITH ('primary_key_encoding' = 'sparse'); + +Error: 1004(InvalidArguments), Invalid region options, COMPRESSION on tag column host requires flat SST format and dense primary key encoding + +CREATE TABLE compressed_tags ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3, page_size = '64KiB'), + val DOUBLE +) WITH ('sst_format' = 'flat', 'primary_key_encoding' = 'dense'); + +Affected Rows: 0 + +SHOW CREATE TABLE compressed_tags; + ++-----------------+-----------------------------------------------------------------------------------------+ +| Table | Create Table | ++-----------------+-----------------------------------------------------------------------------------------+ +| compressed_tags | CREATE TABLE IF NOT EXISTS "compressed_tags" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "host" STRING NULL COMPRESSION WITH(level = '3', page_size = '64KiB', type = 'zstd'), | +| | "val" DOUBLE NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | primary_key_encoding = 'dense', | +| | sst_format = 'flat' | +| | ) | ++-----------------+-----------------------------------------------------------------------------------------+ + +INSERT INTO compressed_tags VALUES (1, 'host-a', 1), (2, 'host-b', 2); + +Affected Rows: 2 + +ADMIN flush_table('compressed_tags'); + ++--------------------------------------+ +| ADMIN flush_table('compressed_tags') | ++--------------------------------------+ +| 0 | ++--------------------------------------+ + +SELECT * FROM compressed_tags ORDER BY ts; + ++-------------------------+--------+-----+ +| ts | host | val | ++-------------------------+--------+-----+ +| 1970-01-01T00:00:00.001 | host-a | 1.0 | +| 1970-01-01T00:00:00.002 | host-b | 2.0 | ++-------------------------+--------+-----+ + +DROP TABLE compressed_tags; + +Affected Rows: 0 + +CREATE TABLE compression_phy ( + ts TIMESTAMP TIME INDEX, + val DOUBLE COMPRESSION WITH (level = 3) +) ENGINE = metric WITH (physical_metric_table = ''); + +Affected Rows: 0 + +CREATE TABLE compression_conflict ( + ts TIMESTAMP TIME INDEX, + val DOUBLE COMPRESSION WITH (level = 9) +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +Error: 1004(InvalidArguments), Unexpected request: COMPRESSION is not supported on logical tables; configure it on the physical table + +CREATE TABLE compression_same ( + ts TIMESTAMP TIME INDEX, + val DOUBLE COMPRESSION WITH (level = 3) +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +Error: 1004(InvalidArguments), Unexpected request: COMPRESSION is not supported on logical tables; configure it on the physical table + +CREATE TABLE compression_new_tag ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3), + val DOUBLE +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +Error: 1004(InvalidArguments), Unexpected request: COMPRESSION is not supported on logical tables; configure it on the physical table + +CREATE TABLE compression_logical ( + ts TIMESTAMP TIME INDEX, + val DOUBLE +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +Affected Rows: 0 + +INSERT INTO compression_logical VALUES (1, 42); + +Affected Rows: 1 + +ADMIN flush_table('compression_phy'); + ++--------------------------------------+ +| ADMIN flush_table('compression_phy') | ++--------------------------------------+ +| 0 | ++--------------------------------------+ + +SELECT * FROM compression_logical; + ++-------------------------+------+ +| ts | val | ++-------------------------+------+ +| 1970-01-01T00:00:00.001 | 42.0 | ++-------------------------+------+ + +DROP TABLE compression_logical; + +Affected Rows: 0 + +DROP TABLE compression_phy; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/create/create_with_compression.sql b/tests/cases/standalone/common/create/create_with_compression.sql new file mode 100644 index 00000000000..166c92be952 --- /dev/null +++ b/tests/cases/standalone/common/create/create_with_compression.sql @@ -0,0 +1,96 @@ +CREATE TABLE spans ( + ts TIMESTAMP TIME INDEX, + trace_id STRING SKIPPING INDEX, + prompt STRING COMPRESSION WITH (type = 'zstd', level = 3, page_size = '8MiB'), + reply STRING COMPRESSION WITH (level = 9), + tokens BIGINT COMPRESSION WITH (page_size = '64KiB'), +) WITH (append_mode = 'true'); + +SHOW CREATE TABLE spans; + +INSERT INTO spans VALUES + (1, 't1', 'fix the flaky test', 'running it', 120), + (2, 't1', 'fix the flaky test\nrunning it', 'patched', 340), + (3, 't2', 'another prompt', NULL, NULL); + +ADMIN flush_table('spans'); + +SELECT ts, trace_id, prompt, reply, tokens FROM spans ORDER BY ts; + +DROP TABLE spans; + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (level = 0)); + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (type = 'lz4')); + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (page_size = 'large')); + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (codec = 'zstd')); + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, j JSON2 COMPRESSION WITH (level = 3)); + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (level = [3])); + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (page_size = NULL)); + +CREATE TABLE bad ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3) +) WITH ('sst_format' = 'primary_key', 'memtable.type' = 'time_series'); + +CREATE TABLE bad ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3) +) WITH ('primary_key_encoding' = 'sparse'); + +CREATE TABLE compressed_tags ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3, page_size = '64KiB'), + val DOUBLE +) WITH ('sst_format' = 'flat', 'primary_key_encoding' = 'dense'); + +SHOW CREATE TABLE compressed_tags; + +INSERT INTO compressed_tags VALUES (1, 'host-a', 1), (2, 'host-b', 2); + +ADMIN flush_table('compressed_tags'); + +SELECT * FROM compressed_tags ORDER BY ts; + +DROP TABLE compressed_tags; + +CREATE TABLE compression_phy ( + ts TIMESTAMP TIME INDEX, + val DOUBLE COMPRESSION WITH (level = 3) +) ENGINE = metric WITH (physical_metric_table = ''); + +CREATE TABLE compression_conflict ( + ts TIMESTAMP TIME INDEX, + val DOUBLE COMPRESSION WITH (level = 9) +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +CREATE TABLE compression_same ( + ts TIMESTAMP TIME INDEX, + val DOUBLE COMPRESSION WITH (level = 3) +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +CREATE TABLE compression_new_tag ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY COMPRESSION WITH (level = 3), + val DOUBLE +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +CREATE TABLE compression_logical ( + ts TIMESTAMP TIME INDEX, + val DOUBLE +) ENGINE = metric WITH (on_physical_table = 'compression_phy'); + +INSERT INTO compression_logical VALUES (1, 42); + +ADMIN flush_table('compression_phy'); + +SELECT * FROM compression_logical; + +DROP TABLE compression_logical; + +DROP TABLE compression_phy; diff --git a/tests/compatibility/cases/column_compression/case.toml b/tests/compatibility/cases/column_compression/case.toml new file mode 100644 index 00000000000..6369410dc08 --- /dev/null +++ b/tests/compatibility/cases/column_compression/case.toml @@ -0,0 +1,8 @@ +name = "column_compression" +reason = "Read legacy column metadata after upgrade, then persist compression options and read compressed SSTs across restart." +introduced_by = "PR #9426" +topologies = ["distributed", "standalone"] +from_range = ["*"] +to_range = [">=v1.3.0-beta.1"] +features = ["table", "compression", "write", "read"] +owner = "storage" diff --git a/tests/compatibility/cases/column_compression/setup.sql b/tests/compatibility/cases/column_compression/setup.sql new file mode 100644 index 00000000000..ec5aa23bd9a --- /dev/null +++ b/tests/compatibility/cases/column_compression/setup.sql @@ -0,0 +1,3 @@ +CREATE TABLE legacy_payload (ts TIMESTAMP TIME INDEX, payload STRING); +INSERT INTO legacy_payload VALUES (1, 'before upgrade'); +ADMIN FLUSH_TABLE('legacy_payload'); diff --git a/tests/compatibility/cases/column_compression/verify.result b/tests/compatibility/cases/column_compression/verify.result new file mode 100644 index 00000000000..d4d9809582f --- /dev/null +++ b/tests/compatibility/cases/column_compression/verify.result @@ -0,0 +1,63 @@ +SELECT * FROM legacy_payload ORDER BY ts; + ++-------------------------+----------------+ +| ts | payload | ++-------------------------+----------------+ +| 1970-01-01T00:00:00.001 | before upgrade | ++-------------------------+----------------+ + +CREATE TABLE compressed_payload ( + ts TIMESTAMP TIME INDEX, + payload STRING COMPRESSION WITH (level = 3, page_size = '64KiB') +); + +Affected Rows: 0 + +INSERT INTO compressed_payload SELECT * FROM legacy_payload; + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('compressed_payload'); + ++-----------------------------------------+ +| ADMIN FLUSH_TABLE('compressed_payload') | ++-----------------------------------------+ +| 0 | ++-----------------------------------------+ + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compressed_payload; + ++--------------------+--------------------------------------------------------------------------------------------+ +| Table | Create Table | ++--------------------+--------------------------------------------------------------------------------------------+ +| compressed_payload | CREATE TABLE IF NOT EXISTS "compressed_payload" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "payload" STRING NULL COMPRESSION WITH(level = '3', page_size = '64KiB', type = 'zstd'), | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++--------------------+--------------------------------------------------------------------------------------------+ + +INSERT INTO compressed_payload VALUES (2, 'after restart'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('compressed_payload'); + ++-----------------------------------------+ +| ADMIN FLUSH_TABLE('compressed_payload') | ++-----------------------------------------+ +| 0 | ++-----------------------------------------+ + +SELECT * FROM compressed_payload ORDER BY ts; + ++-------------------------+----------------+ +| ts | payload | ++-------------------------+----------------+ +| 1970-01-01T00:00:00.001 | before upgrade | +| 1970-01-01T00:00:00.002 | after restart | ++-------------------------+----------------+ diff --git a/tests/compatibility/cases/column_compression/verify.sql b/tests/compatibility/cases/column_compression/verify.sql new file mode 100644 index 00000000000..4ddf8082317 --- /dev/null +++ b/tests/compatibility/cases/column_compression/verify.sql @@ -0,0 +1,15 @@ +SELECT * FROM legacy_payload ORDER BY ts; + +CREATE TABLE compressed_payload ( + ts TIMESTAMP TIME INDEX, + payload STRING COMPRESSION WITH (level = 3, page_size = '64KiB') +); +INSERT INTO compressed_payload SELECT * FROM legacy_payload; +ADMIN FLUSH_TABLE('compressed_payload'); + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compressed_payload; + +INSERT INTO compressed_payload VALUES (2, 'after restart'); +ADMIN FLUSH_TABLE('compressed_payload'); +SELECT * FROM compressed_payload ORDER BY ts; From 55d753ca58ad85d888fae3b686d943f808a1f87a Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 13 Oct 2026 12:26:16 +0000 Subject: [PATCH 30/46] perf(log-store): fetch object store WAL replay segments concurrently (#9555) * perf(log-store): fetch object store WAL replay segments concurrently A region read fetched each segment only after the previous object's entries were yielded, so a replay paid the full request latency once per object. Fetch up to READ_PREFETCH_DEPTH (8) segments ahead, each on its own task, and handle their outcomes in object sequence order: entries, collected-object skips, corrupted-segment holes and errors surface exactly as before, and dropping the stream aborts the fetches in flight. A read holds at most eight decoded segments besides the one it yields. Signed-off-by: jeremyhi * perf(log-store): raise object store WAL recovery concurrency to 64 On S3, recovery reads the header and tail of every object, about 57 to 91 ms per object per lane on 4, 8 and 16 vCPUs alike, and its time scales with objects divided by concurrency up to 128 lanes without throttling: about 9,500 objects recovered in 99 s at 8 and 10.5 s at 64. At 64 recovery issues about 1,800 reads per second, a third of S3's request rate per prefix, and holds one 64 KiB tail window per lane, 4 MiB in all. The recovery concurrency tests now scale with the constant. Signed-off-by: jeremyhi * perf(log-store): share one prefetch byte budget among object store WAL reads Each read fetched up to eight segments ahead on its own, so regions opened in parallel (16 by default) could together hold 128 segments of up to max_batch_bytes each. A read now always fetches the segment it yields next, as before, and fetches a later one only while its length fits in a 64 MiB budget shared by every read of the store. Regions replayed in parallel thus hold one segment each plus at most 64 MiB, however many open at once. Nothing waits on the budget: a read that cannot reserve stops fetching ahead until its next poll. Signed-off-by: jeremyhi * fix(log-store): fail an object store WAL read whose fetch task panics A fetch task that panicked was re-raised in the reader with resume_unwind, which unwinds through the caller, such as Mito's region opener task, past its cleanup. Return it as a WalSegmentFetch error instead, so the read ends at that object like any other failed fetch. Signed-off-by: jeremyhi * fix(log-store): bound replay prefetch by decoded size held by fetch tasks The prefetch reservation sat in the read's queue, so dropping a read released it at once while an aborted fetch could still hold its buffers until it next yielded. The fetch task now holds the reservation and returns it with its outcome, so it is released with the buffers. A reservation also counted only the encoded segment length, while a fetched segment is held decoded. It now counts the length plus one Record per entry, and the comment states the whole bound. Signed-off-by: jeremyhi * docs(log-store): round the measured figures in the recovery concurrency comment Keep the magnitudes in the comment; the exact S3 figures belong to the measurement run rather than the code. Signed-off-by: jeremyhi * refactor(log-store): abort replay fetch tasks with tokio-util's AbortOnDropHandle Signed-off-by: jeremyhi * refactor(log-store): rename the replay prefetch reservation to try_reserve_prefetch_bytes Signed-off-by: jeremyhi * test(log-store): share the bounded polling helper across the object store WAL tests Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 2 + src/log-store/Cargo.toml | 1 + src/log-store/src/error.rs | 12 +- src/log-store/src/object_store_wal/store.rs | 580 ++++++++++++++++++-- 4 files changed, 563 insertions(+), 32 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8678104f5d4..d84446150d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8138,6 +8138,7 @@ dependencies = [ "snafu 0.8.6", "store-api", "tokio", + "tokio-util", "uuid", ] @@ -15489,6 +15490,7 @@ dependencies = [ "futures-core", "futures-io", "futures-sink", + "futures-util", "pin-project-lite", "slab", "tokio", diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 7e4da1557e1..5870110be67 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -46,6 +46,7 @@ serde_json.workspace = true snafu.workspace = true store-api.workspace = true tokio.workspace = true +tokio-util = { workspace = true, features = ["rt"] } [dev-dependencies] common-meta = { workspace = true, features = ["testing"] } diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index 85000dbea3c..64d918d5d78 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -427,6 +427,15 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to fetch a segment of WAL object, path: {}", path))] + WalSegmentFetch { + path: String, + #[snafu(source)] + error: tokio::task::JoinError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid WAL object, path: {}", path))] InvalidWalObject { path: String, @@ -509,7 +518,8 @@ impl ErrorExt for Error { | WaitProduceResultReceiver { .. } | WaitDumpIndex { .. } | MetaLengthExceededLimit { .. } - | ObjectStoreWalStopped { .. } => StatusCode::Internal, + | ObjectStoreWalStopped { .. } + | WalSegmentFetch { .. } => StatusCode::Internal, CorruptedWalObject { .. } | WalObjectConflict { .. } diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 98b0c8cad6b..1d75739350a 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -35,15 +35,16 @@ use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::{ObjectStoreProvider, Provider}; use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; use store_api::storage::RegionId; -use tokio::sync::{mpsc, oneshot, watch}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch}; use tokio::time::MissedTickBehavior; +use tokio_util::task::AbortOnDropHandle; use crate::error::{ CorruptedWalObjectSnafu, Error, IncompleteWalEntrySnafu, InvalidProviderSnafu, InvalidWalObjectSnafu, InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, MismatchedWalRegionSnafu, ObjectStoreWalSnafu, ObjectStoreWalStoppedSnafu, Result, StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, - WalObjectSequenceUnsettledSnafu, + WalObjectSequenceUnsettledSnafu, WalSegmentFetchSnafu, }; use crate::metrics::{ METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS, METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL, @@ -81,8 +82,31 @@ const CREATE_RETRY_DELAY: Duration = Duration::from_millis(100); /// append from its channel only while two more fit: a stalled object store /// holds callers back instead of growing the backlog. const MAX_SEALED_BATCHES: usize = 2 * MAX_IN_FLIGHT_CREATES; -/// Number of objects whose footers recovery fetches at a time. -const RECOVERY_CONCURRENCY: usize = 8; +/// Number of objects whose footers recovery fetches at a time. On S3 an +/// object costs two reads, its header and its tail, tens of milliseconds per +/// lane however many vCPUs the node has, and recovery time scales with objects +/// divided by concurrency up to 128 lanes without throttling: at 64 about +/// 10,000 objects recover in about 10 s instead of about 100 s at 8, while the +/// reads use about a third of S3's request rate per prefix, leaving room for +/// the data file reads of start-up. Each lane holds one +/// [`RECOVERY_TAIL_WINDOW`], 4 MiB in all. +const RECOVERY_CONCURRENCY: usize = 64; +/// Number of segments a read fetches ahead of the one whose entries it +/// yields. On S3 a replay that fetched each segment after the previous one +/// spent about 30 ms per object; with eight fetches in flight a segment +/// arrives about every 4 ms. The segments fetched ahead also take their bytes +/// from [`READ_PREFETCH_BYTES`]. +const READ_PREFETCH_DEPTH: usize = 8; +/// Bytes of decoded segments that all reads of a store together fetch ahead; a +/// segment counts as its length, which bounds its payloads, plus a [`Record`] +/// per entry. A read always fetches the segment it yields next, as a read that +/// fetched one segment at a time did, and a later segment only while it fits +/// in what is left, so regions replayed in parallel hold one segment each, up +/// to about `max_batch_bytes` or more for one oversized append, plus at most +/// these bytes, however many regions open at once. While a fetch copies the +/// payloads out of the downloaded segment it briefly holds the segment's +/// length once more. +const READ_PREFETCH_BYTES: usize = 64 * 1024 * 1024; /// Bytes recovery reads from the end of an object in one request. The window /// holds the trailer and the footer of an object with up to 1364 regions of /// 48 bytes each, so a second request for the footer is rare. @@ -113,6 +137,8 @@ pub struct ObjectStoreLogStore { wal_holes: WalHoles, /// Objects with a delete in flight, which a read consults. deleting: Arc, + /// Bytes the reads may fetch ahead, see [`READ_PREFETCH_BYTES`]. + read_prefetch_budget: Arc, terminal_error: TerminalError, stopped: Arc, command_tx: mpsc::Sender, @@ -393,6 +419,7 @@ impl ObjectStoreLogStore { obsolete_entry_ids, wal_holes, deleting, + read_prefetch_budget: Arc::new(Semaphore::new(READ_PREFETCH_BYTES)), terminal_error, stopped, command_tx, @@ -721,7 +748,9 @@ impl LogStore for ObjectStoreLogStore { /// An object collected after the read listed it is skipped. Otherwise a /// segment that does not decode fails the read or is skipped and recorded /// as a hole, as `on_corrupted_segment` says; a failed fetch always fails - /// the read. + /// the read. Up to [`READ_PREFETCH_DEPTH`] segments are fetched ahead + /// while [`READ_PREFETCH_BYTES`] allows, and each outcome is handled in + /// sequence order. async fn read( &self, provider: &Provider, @@ -760,10 +789,39 @@ impl LogStore for ObjectStoreLogStore { let provider = provider.clone(); let on_corrupted_segment = self.on_corrupted_segment; let wal_holes = self.wal_holes.clone(); + let prefetch_budget = self.read_prefetch_budget.clone(); Ok(Box::pin(try_stream! { let started_at = Instant::now(); - for (object_seq, footer_entry) in objects { - let records = match fetch_segment(io.as_ref(), object_seq, &footer_entry).await { + let mut objects = objects.into_iter().peekable(); + let mut fetches = VecDeque::new(); + loop { + // The next segment is always fetched, a later one only with + // its bytes reserved from the budget. + while fetches.len() < READ_PREFETCH_DEPTH { + let Some((_, entry)) = objects.peek() else { + break; + }; + let permit = if fetches.is_empty() { + None + } else { + match try_reserve_prefetch_bytes(&prefetch_budget, entry) { + Some(permit) => Some(permit), + None => break, + } + }; + let Some((object_seq, footer_entry)) = objects.next() else { + break; + }; + let fetch = + spawn_fetch_segment(io.clone(), object_seq, footer_entry.clone(), permit); + fetches.push_back((object_seq, footer_entry, fetch)); + } + let Some((object_seq, footer_entry, fetch)) = fetches.pop_front() else { + break; + }; + // The permit is held until the entries of the segment are yielded. + let (fetched, _permit) = fetch.await; + let records = match fetched { Ok(records) => records, Err(error) if is_collected(&catalog, &deleting, object_seq).await => { debug!( @@ -2135,6 +2193,47 @@ async fn fetch_segment( }) } +/// Takes from `budget` the bytes the segment `entry` describes holds once +/// decoded, if they are free: its length plus a [`Record`] per entry. +fn try_reserve_prefetch_bytes( + budget: &Arc, + entry: &FooterEntry, +) -> Option { + let records = u64::from(entry.entry_count) * std::mem::size_of::() as u64; + let bytes = u32::try_from(entry.segment_len.checked_add(records)?).ok()?; + budget.clone().try_acquire_many_owned(bytes).ok() +} + +/// Runs [`fetch_segment`] on a task of its own, so the fetch proceeds while +/// the read is busy with earlier segments, and returns the outcome with +/// `permit`. The task holds the permit, so the reserved bytes are released +/// only with the buffers of the fetch. Dropping the returned future, as a +/// dropped read does, aborts the task; a task that panics fails the fetch with +/// [`Error::WalSegmentFetch`]. +fn spawn_fetch_segment( + io: Arc, + object_seq: u64, + entry: FooterEntry, + permit: Option, +) -> impl Future>, Option)> + Send + 'static { + let task_io = io.clone(); + let task = AbortOnDropHandle::new(common_runtime::spawn_global(async move { + let fetched = fetch_segment(task_io.as_ref(), object_seq, &entry).await; + (fetched, permit) + })); + async move { + match task.await { + Ok(fetched) => fetched, + Err(error) => ( + Err(error).context(WalSegmentFetchSnafu { + path: io.object_path(object_seq), + }), + None, + ), + } + } +} + /// Records the segment `entry` describes in the object `object_seq` as a hole /// of its region, once per object, counts the skip and warns about it. An /// object collected meanwhile is skipped without a hole: the catalog is held @@ -3547,6 +3646,7 @@ mod tests { obsolete_entry_ids: ObsoleteEntryIds::default(), wal_holes: WalHoles::default(), deleting: Arc::default(), + read_prefetch_budget: Arc::new(Semaphore::new(READ_PREFETCH_BYTES)), terminal_error: Arc::default(), stopped: Arc::new(AtomicBool::new(false)), command_tx, @@ -3629,7 +3729,8 @@ mod tests { async fn test_store_recovery_bounds_concurrency_and_orders_by_sequence() { let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); - populate(&object_store, 16, 3).await; + let depth = RECOVERY_CONCURRENCY as u64; + populate(&object_store, 2 * RECOVERY_CONCURRENCY, 3).await; let (io, mut parked) = ParkedIo::over(object_store); let objects = io.list().await.unwrap(); @@ -3650,12 +3751,12 @@ mod tests { wave.push(timeout(WAIT, parked.recv()).await.unwrap().unwrap()); } assert_eq!( - (0..8).collect::>(), + (0..depth).collect::>(), wave.iter() .map(|(object_seq, _)| *object_seq) .collect::>() ); - // No ninth fetch can start while eight are parked. + // No further fetch can start while a full wave is parked. for _ in 0..16 { tokio::task::yield_now().await; } @@ -3663,7 +3764,7 @@ mod tests { // Complete the initial wave in reverse order, parking each replacement. let mut replacements = Vec::new(); - for (expected_next, (_, release)) in (8..16).zip(wave.into_iter().rev()) { + for (expected_next, (_, release)) in (depth..2 * depth).zip(wave.into_iter().rev()) { release.send(true).unwrap(); let (object_seq, release) = timeout(WAIT, parked.recv()).await.unwrap().unwrap(); assert_eq!(expected_next, object_seq); @@ -3674,17 +3775,21 @@ mod tests { release.send(true).unwrap(); } assert_eq!( - (0..16).collect::>(), + (0..2 * depth).collect::>(), timeout(WAIT, &mut fetch).await.unwrap().unwrap() ); - assert_eq!(8, io.max_in_flight.load(Ordering::SeqCst)); + assert_eq!( + RECOVERY_CONCURRENCY, + io.max_in_flight.load(Ordering::SeqCst) + ); } #[tokio::test] async fn test_store_recovery_abandons_pending_fetches_on_failure() { let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); - populate(&object_store, 16, 3).await; + let objects = 2 * RECOVERY_CONCURRENCY as u64; + populate(&object_store, objects as usize, 3).await; let (io, mut parked) = ParkedIo::over(object_store.clone()); let fetch = { let io = io.clone(); @@ -3709,9 +3814,13 @@ mod tests { } assert!(parked.try_recv().is_err()); let store = open(object_store, &eager()).await; - assert_eq!(id(15, 1), latest(&store, region(1))); - // Sequence 16 holds the object that started the epoch of the store. - assert_eq!(17, store.catalog.read().unwrap().next_object_seq().unwrap()); + assert_eq!(id(objects - 1, 1), latest(&store, region(1))); + // The sequence after the populated objects holds the object that + // started the epoch of the store. + assert_eq!( + objects + 1, + store.catalog.read().unwrap().next_object_seq().unwrap() + ); store.stop().await.unwrap(); } @@ -3846,7 +3955,10 @@ mod tests { }) .collect::>(); assert_eq!(expected, entries); - assert_eq!(expected_ranges[®ion(number)], *reads.lock().unwrap()); + // Segments are fetched concurrently, so in no particular order. + let mut reads = reads.lock().unwrap().clone(); + reads.sort_unstable(); + assert_eq!(expected_ranges[®ion(number)], reads); } reads.lock().unwrap().clear(); assert_eq!( @@ -3950,6 +4062,18 @@ mod tests { store.stop().await.unwrap(); } + /// The range reads of object `object_seq` among `reads`, in the order + /// they were issued. + fn range_reads_of(reads: &RangeReads, object_seq: u64) -> Vec<(u64, u64, u64)> { + reads + .lock() + .unwrap() + .iter() + .filter(|read| read.0 == object_seq) + .copied() + .collect() + } + /// Writes object 0 with one entry of region 1 and one of region 2 and /// object 1 with an entry of region 2, then flips a byte in the segment of /// region 2 in object 0. Returns the key of object 0 and its footer. @@ -4018,8 +4142,7 @@ mod tests { read_entries(&store, region(2), 0).await ); let corrupt = (0, footer[1].segment_offset, footer[1].segment_len); - assert_eq!(corrupt, reads.lock().unwrap()[0]); - assert_eq!(corrupt, reads.lock().unwrap()[1]); + assert_eq!(vec![corrupt, corrupt], range_reads_of(&reads, 0)); let hole = WalHole { path, object_seq: 0, @@ -4087,7 +4210,7 @@ mod tests { &format!("segment of region {} checksum mismatch", region(2)), ); let corrupt = (0, footer[1].segment_offset, footer[1].segment_len); - assert_eq!(vec![corrupt, corrupt], *reads.lock().unwrap()); + assert_eq!(vec![corrupt, corrupt], range_reads_of(&reads, 0)); assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); assert_eq!( skipped, @@ -4158,6 +4281,379 @@ mod tests { } } + /// The entries of region 1 that [`put_object_per_entry`] wrote to the + /// objects `object_seqs`. + fn object_entries(object_seqs: impl IntoIterator) -> Vec { + object_seqs + .into_iter() + .map(|object_seq| { + let entry_id = object_seq + 1; + Entry::Naive(NaiveEntry { + provider: provider(region(1)), + region_id: region(1), + entry_id, + data: format!("e{entry_id}").into_bytes(), + }) + }) + .collect() + } + + /// Writes objects 0 to `objects - 1`, object `n` with entry `n + 1` of + /// region 1. + async fn put_object_per_entry(object_store: &ObjectStore, objects: u64) { + for object_seq in 0..objects { + put_object(object_store, object_seq, region(1), &[object_seq + 1]).await; + } + } + + /// Opens a store over [`put_object_per_entry`] objects, releasing every + /// read of recovery; range reads park from then on. + async fn open_parking_range_reads( + objects: u64, + ) -> ( + Arc, + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + let object_store = memory_store(); + put_object_per_entry(&object_store, objects).await; + let (io, mut parked) = ParkedIo::over(object_store); + let mut opening = { + let io = io.clone(); + tokio::spawn(async move { + ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await + }) + }; + let store = loop { + tokio::select! { + opened = &mut opening => break opened.unwrap().unwrap(), + Some((_, release)) = parked.recv() => release.send(true).unwrap(), + } + }; + io.range_reads_parked.store(true, Ordering::SeqCst); + io.max_in_flight.store(0, Ordering::SeqCst); + (store, io, parked) + } + + /// Polls until `settled` holds. + async fn wait_until(settled: impl Fn() -> bool) { + let deadline = Instant::now() + WAIT; + while !settled() { + assert!(Instant::now() < deadline, "the condition never held"); + tokio::time::sleep(Duration::from_millis(1)).await; + } + } + + async fn assert_nothing_parks( + parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + assert!( + timeout(Duration::from_millis(100), parked.recv()) + .await + .is_err() + ); + } + + async fn next_batch(stream: &mut SendableEntryStream<'static, Entry, Error>) -> Vec { + timeout(WAIT, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap() + } + + #[tokio::test] + async fn test_store_read_prefetches_up_to_the_depth_in_order() { + let _serialized = SERIALIZED.lock().await; + let depth = READ_PREFETCH_DEPTH as u64; + let objects = depth + 4; + let (store, io, mut parked) = open_parking_range_reads(objects).await; + let reads = METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(); + let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + + // The first poll starts the fetches of the first `depth` objects. + assert!(stream.next().now_or_never().is_none()); + let mut wave = parked_operations(&mut parked, READ_PREFETCH_DEPTH).await; + assert_eq!( + (0..depth).collect::>(), + wave.keys().copied().collect() + ); + assert_nothing_parks(&mut parked).await; + + // Later objects complete first: they wait for object 0 and still + // count toward the depth. + for object_seq in (1..depth).rev() { + wave.remove(&object_seq).unwrap().send(true).unwrap(); + } + wait_until(|| io.in_flight.load(Ordering::SeqCst) == 1).await; + assert!(stream.next().now_or_never().is_none()); + assert_nothing_parks(&mut parked).await; + wave.remove(&0).unwrap().send(true).unwrap(); + assert_eq!(object_entries([0]), next_batch(&mut stream).await); + + // No fetch starts while the caller holds the entries of object 0; + // each later poll starts the fetch of the next object. + assert_nothing_parks(&mut parked).await; + let rest = tokio::spawn(stream.try_collect::>()); + let refills = parked_operations(&mut parked, (objects - depth) as usize).await; + assert_eq!( + (depth..objects).collect::>(), + refills.keys().copied().collect() + ); + for release in refills.into_values() { + release.send(true).unwrap(); + } + let rest = timeout(WAIT, rest).await.unwrap().unwrap().unwrap(); + assert_eq!( + object_entries(1..objects), + rest.into_iter().flatten().collect::>() + ); + assert_eq!(READ_PREFETCH_DEPTH, io.max_in_flight.load(Ordering::SeqCst)); + assert_eq!( + reads + 1, + METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_read_fails_at_a_prefetched_object_in_order() { + let _serialized = SERIALIZED.lock().await; + let depth = READ_PREFETCH_DEPTH as u64; + let objects = depth + 4; + let failed = 5; + let (store, io, mut parked) = open_parking_range_reads(objects).await; + let reads = METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(); + let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + assert!(stream.next().now_or_never().is_none()); + let mut wave = parked_operations(&mut parked, READ_PREFETCH_DEPTH).await; + + // The fetch of object 5 fails before any earlier fetch completes. + wave.remove(&failed).unwrap().send(false).unwrap(); + wait_until(|| io.in_flight.load(Ordering::SeqCst) == READ_PREFETCH_DEPTH - 1).await; + for release in wave.into_values() { + release.send(true).unwrap(); + } + let mut refills = Vec::new(); + for object_seq in 0..failed { + assert_eq!(object_entries([object_seq]), next_batch(&mut stream).await); + if object_seq > 0 { + refills.push(next_parked_operation(&mut parked).await); + } + } + let error = timeout(WAIT, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { operation: "read", path, .. } + if *path == io.object_path(failed)), + "unexpected error: {error:?}" + ); + assert!(timeout(WAIT, stream.next()).await.unwrap().is_none()); + + // The fetches started meanwhile are abandoned. + assert_eq!( + (depth..objects).collect::>(), + refills + .iter() + .map(|(object_seq, _)| *object_seq) + .collect::>() + ); + for (_, mut release) in refills { + timeout(WAIT, release.closed()).await.unwrap(); + } + assert_nothing_parks(&mut parked).await; + assert_eq!( + reads, + METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_read_skips_prefetched_collected_and_corrupted_objects() { + let _serialized = SERIALIZED.lock().await; + let objects = READ_PREFETCH_DEPTH as u64 + 4; + let (corrupted, collected) = (3, 9); + let object_store = memory_store(); + put_object_per_entry(&object_store, objects).await; + let path = object_path(&object_store, corrupted); + let bytes = object_store.read(&path).await.unwrap().to_vec(); + let segment = footer_of(&bytes).1.remove(0); + corrupt_object(&object_store, &path, |bytes| { + bytes[(segment.segment_offset + segment.segment_len - 1) as usize] ^= 1; + }) + .await; + let store = open(object_store.clone(), &eager()).await; + assert_eq!(CorruptedSegmentAction::Skip, store.on_corrupted_segment); + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + // A collection removes object 9 after the read listed it. + let stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + store.catalog.write().unwrap().remove_object(collected); + object_store + .delete(&object_path(&object_store, collected)) + .await + .unwrap(); + let entries = timeout(WAIT, stream.try_collect::>()) + .await + .unwrap() + .unwrap() + .into_iter() + .flatten() + .collect::>(); + assert_eq!( + object_entries((0..objects).filter(|seq| ![corrupted, collected].contains(seq))), + entries + ); + assert_eq!( + vec![WalHole { + path, + object_seq: corrupted, + min_entry_id: corrupted + 1, + max_entry_id: corrupted + 1, + }], + store.wal_holes(&provider(region(1))).unwrap() + ); + assert_eq!( + skipped + 1, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_dropped_read_abandons_its_prefetches() { + let _serialized = SERIALIZED.lock().await; + let objects = READ_PREFETCH_DEPTH as u64 + 4; + let (store, _io, mut parked) = open_parking_range_reads(objects).await; + let reads = METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(); + let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + assert!(stream.next().now_or_never().is_none()); + let mut wave = parked_operations(&mut parked, READ_PREFETCH_DEPTH).await; + wave.remove(&0).unwrap().send(true).unwrap(); + assert_eq!(object_entries([0]), next_batch(&mut stream).await); + + drop(stream); + for (_, mut release) in wave { + timeout(WAIT, release.closed()).await.unwrap(); + } + assert_nothing_parks(&mut parked).await; + assert_eq!( + reads, + METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_read_fails_at_a_panicking_fetch() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + put_object_per_entry(&object_store, 3).await; + let (io, _) = RecordingIo::over(object_store); + let store = open_over(io.clone(), &eager()).await; + *io.panic_range_read_of.lock().unwrap() = Some(1); + let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + assert_eq!(object_entries([0]), next_batch(&mut stream).await); + let error = timeout(WAIT, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!( + matches!(&error, Error::WalSegmentFetch { path, .. } if *path == io.object_path(1)), + "unexpected error: {error:?}" + ); + assert!(timeout(WAIT, stream.next()).await.unwrap().is_none()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_reads_fetch_ahead_within_a_shared_budget() { + let _serialized = SERIALIZED.lock().await; + let depth = READ_PREFETCH_DEPTH as u64; + let objects = depth + 4; + let (store, _io, mut parked) = open_parking_range_reads(objects).await; + + // While other reads hold the whole budget, a read fetches only the + // segment it yields next. + let held = store + .read_prefetch_budget + .clone() + .try_acquire_many_owned(READ_PREFETCH_BYTES as u32) + .unwrap(); + let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + assert!(stream.next().now_or_never().is_none()); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + assert_nothing_parks(&mut parked).await; + release.send(true).unwrap(); + assert_eq!(object_entries([0]), next_batch(&mut stream).await); + + // Once the budget is free, the next poll fetches ahead up to the depth. + drop(held); + assert!(stream.next().now_or_never().is_none()); + let wave = parked_operations(&mut parked, READ_PREFETCH_DEPTH).await; + assert_eq!( + (1..=depth).collect::>(), + wave.keys().copied().collect() + ); + assert_nothing_parks(&mut parked).await; + for release in wave.into_values() { + release.send(true).unwrap(); + } + let rest = tokio::spawn(stream.try_collect::>()); + for release in parked_operations(&mut parked, (objects - depth - 1) as usize) + .await + .into_values() + { + release.send(true).unwrap(); + } + let rest = timeout(WAIT, rest).await.unwrap().unwrap().unwrap(); + assert_eq!( + object_entries(1..objects), + rest.into_iter().flatten().collect::>() + ); + assert_eq!( + READ_PREFETCH_BYTES, + store.read_prefetch_budget.available_permits() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_dropped_read_keeps_reservations_until_its_fetches_end() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + put_object_per_entry(&object_store, 2).await; + let (io, _) = RecordingIo::over(object_store); + let store = open_over(io.clone(), &eager()).await; + let budget = store.read_prefetch_budget.clone(); + let (began_tx, began_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + *io.block_range_read_of.lock().unwrap() = Some((1, began_tx, release_rx)); + + // The first poll fetches object 0, and object 1 with its bytes + // reserved; the read of object 1 blocks its thread. + let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let _ = stream.next().now_or_never(); + began_rx.recv_timeout(WAIT).unwrap(); + assert!(budget.available_permits() < READ_PREFETCH_BYTES); + + // The aborted fetch keeps its bytes until it ends. + drop(stream); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(budget.available_permits() < READ_PREFETCH_BYTES); + release_tx.send(()).unwrap(); + wait_until(|| budget.available_permits() == READ_PREFETCH_BYTES).await; + store.stop().await.unwrap(); + } + #[tokio::test] async fn test_store_assigns_ids_from_the_object_sequence() { let _serialized = SERIALIZED.lock().await; @@ -7009,7 +7505,9 @@ mod tests { /// after it wrote the object, fails it with a given error before it /// writes, stores the same content before it writes, or damages the bytes /// the next range read returns; it also fails the next listing, or the - /// next range read of an object at an offset, with a transient error. + /// next range read of an object at an offset, with a transient error, + /// panics in the next range read of a given object, and blocks its thread + /// in the next range read of a given object until the test releases it. struct RecordingIo { inner: ObjectStoreIo, reads: RangeReads, @@ -7019,8 +7517,18 @@ mod tests { duplicate_next_put: AtomicBool, fail_next_list: AtomicBool, fail_range_read_at: Mutex>, + panic_range_read_of: Mutex>, + block_range_read_of: Mutex>, } + /// The object whose range read blocks, the sender that reports the read + /// began and the receiver that releases it. + type BlockedRangeRead = ( + u64, + std::sync::mpsc::Sender<()>, + std::sync::mpsc::Receiver<()>, + ); + impl RecordingIo { fn over(object_store: ObjectStore) -> (Arc, RangeReads) { let reads = RangeReads::default(); @@ -7033,6 +7541,8 @@ mod tests { duplicate_next_put: AtomicBool::new(false), fail_next_list: AtomicBool::new(false), fail_range_read_at: Mutex::new(None), + panic_range_read_of: Mutex::new(None), + block_range_read_of: Mutex::new(None), }; (Arc::new(io), reads) } @@ -7091,6 +7601,23 @@ mod tests { { return injected_failure("read", self.object_path(object_seq)); } + let panics = self + .panic_range_read_of + .lock() + .unwrap() + .take_if(|of| *of == object_seq) + .is_some(); + assert!(!panics, "injected panic"); + let blocked = self + .block_range_read_of + .lock() + .unwrap() + .take_if(|(of, ..)| *of == object_seq); + if let Some((_, began, release)) = blocked { + began.send(()).unwrap(); + // Blocking the thread keeps an abort from cancelling the read. + let _ = release.recv(); + } let bytes = self.inner.get_range(object_seq, offset, len).await?; if self.damage_next_range_read.swap(false, Ordering::Relaxed) { let mut damaged = bytes.to_vec(); @@ -7389,15 +7916,6 @@ mod tests { ); } - /// Polls until `settled` holds. - async fn wait_until(settled: impl Fn() -> bool) { - let deadline = Instant::now() + WAIT; - while !settled() { - assert!(Instant::now() < deadline, "the actor never settled"); - tokio::time::sleep(Duration::from_millis(1)).await; - } - } - #[tokio::test] async fn test_store_fault_matrix() { let _serialized = SERIALIZED.lock().await; From 4bd36cfd20c987b19247e6fe9ef0cad4438047cb Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:28:36 +0000 Subject: [PATCH 31/46] feat: support compaction options in HTTP hints (#9560) * feat: support compaction options in HTTP hints Signed-off-by: Lei, HUANG * refactor: share TWCS compaction option prefix Use the store-api constant for both HTTP hint inference and Mito region option validation so the two paths cannot drift. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/src/region/options.rs | 3 +- src/operator/src/insert.rs | 70 ++++++++++++++++++--- src/servers/src/hint_headers.rs | 8 ++- src/store-api/src/mito_engine_options.rs | 78 +++++++++++++----------- 4 files changed, 109 insertions(+), 50 deletions(-) diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 7e050fb4a76..42f4adcd04a 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -36,7 +36,7 @@ use store_api::metric_engine_consts::{ MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING, PRIMARY_KEY_ENCODING, }; use store_api::mito_engine_options::{ - COMPACTION_OVERRIDE, FloatFieldEncoding, MAX_ROW_GROUP_ROW_COUNT_LIMIT, + COMPACTION_OVERRIDE, COMPACTION_TWCS_PREFIX, FloatFieldEncoding, MAX_ROW_GROUP_ROW_COUNT_LIMIT, }; use store_api::storage::{ColumnId, RegionId}; use strum::EnumString; @@ -47,7 +47,6 @@ use crate::sst::FormatType; use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; const DEFAULT_INDEX_SEGMENT_ROW_COUNT: usize = 1024; -const COMPACTION_TWCS_PREFIX: &str = "compaction.twcs."; const MEMTABLE_PARTITION_TREE_PREFIX: &str = "memtable.partition_tree."; const MEMTABLE_BULK_PREFIX: &str = "memtable.bulk."; const MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD: &str = "memtable.bulk.encode_bytes_threshold"; diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 2c01a35b8d9..2b382dbdb5f 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -67,8 +67,8 @@ use store_api::metric_engine_consts::{ LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, }; use store_api::mito_engine_options::{ - APPEND_MODE_KEY, COMPACTION_TYPE, COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY, - TWCS_TIME_WINDOW, + APPEND_MODE_KEY, COMPACTION_OPTION_KEYS, COMPACTION_TWCS_PREFIX, COMPACTION_TYPE, + COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY, }; use store_api::storage::{RegionId, TableId}; use table::TableRef; @@ -1856,6 +1856,21 @@ pub fn fill_table_options_for_create( } } + if !matches!(create_type, AutoCreateTableType::Logical(_)) { + let mut has_twcs_option = false; + for key in COMPACTION_OPTION_KEYS { + if let Some(value) = ctx.extension(key) { + has_twcs_option |= key.starts_with(COMPACTION_TWCS_PREFIX); + table_options.insert(key.to_string(), value.to_string()); + } + } + if has_twcs_option { + table_options + .entry(COMPACTION_TYPE.to_string()) + .or_insert_with(|| COMPACTION_TYPE_TWCS.to_string()); + } + } + // Semantic keys use their own vocabulary instead of the fixed option list. for (key, value) in ctx.extensions() { if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) { @@ -1877,14 +1892,6 @@ pub fn fill_table_options_for_create( if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) { table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string()); } - if let Some(time_window) = ctx.extension(TWCS_TIME_WINDOW) { - table_options.insert(TWCS_TIME_WINDOW.to_string(), time_window.to_string()); - // We need to set the compaction type explicitly. - table_options.insert( - COMPACTION_TYPE.to_string(), - COMPACTION_TYPE_TWCS.to_string(), - ); - } } // Set append_mode to true for log table. // because log tables should keep rows with the same ts and tags. @@ -3051,6 +3058,49 @@ mod tests { assert_eq!(requests, original); } + #[test] + fn test_trace_create_options_copy_compaction_hints() { + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + ctx.set_extension("compaction.twcs.time_window", "1h"); + ctx.set_extension("compaction.twcs.max_output_file_size", "256MB"); + ctx.set_extension("compaction.twcs.remote_compaction", "true"); + ctx.set_extension("compaction.unknown", "ignored"); + let ctx = Arc::new(ctx); + let mut table_options = Default::default(); + + fill_table_options_for_create( + &mut table_options, + &AutoCreateTableType::Trace { + alter_existing: true, + }, + &ctx, + ); + + assert_eq!( + Some("twcs"), + table_options.get(COMPACTION_TYPE).map(String::as_str) + ); + assert_eq!( + Some("1h"), + table_options + .get("compaction.twcs.time_window") + .map(String::as_str) + ); + assert_eq!( + Some("256MB"), + table_options + .get("compaction.twcs.max_output_file_size") + .map(String::as_str) + ); + assert_eq!( + Some("true"), + table_options + .get("compaction.twcs.remote_compaction") + .map(String::as_str) + ); + assert!(!table_options.contains_key("compaction.unknown")); + } + #[test] fn test_skip_wal_does_not_change_table_options() { check_skip_wal_does_not_change_table_options(false); diff --git a/src/servers/src/hint_headers.rs b/src/servers/src/hint_headers.rs index 3d13dd9bf5e..4d37447e01d 100644 --- a/src/servers/src/hint_headers.rs +++ b/src/servers/src/hint_headers.rs @@ -145,12 +145,12 @@ mod tests { headers.insert( "x-greptime-hints", HeaderValue::from_static(" auto_create_table=true, ttl =3600d, append_mode=true , merge_mode=false , physical_table= table1,\ - read_preference=leader"), + read_preference=leader, compaction.twcs.time_window=1h"), ); let hints = extract_hints(&headers); - assert_eq!(hints.len(), 6); + assert_eq!(hints.len(), 7); assert_eq!( hints[0], ("auto_create_table".to_string(), "true".to_string()) @@ -166,6 +166,10 @@ mod tests { hints[5], ("read_preference".to_string(), "leader".to_string()) ); + assert_eq!( + hints[6], + ("compaction.twcs.time_window".to_string(), "1h".to_string()) + ); } #[test] diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index 07cfe6a6516..457b33665c3 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -35,6 +35,8 @@ pub const COMPACTION_TYPE: &str = "compaction.type"; pub const COMPACTION_OVERRIDE: &str = "compaction.override"; /// TWCS compaction strategy. pub const COMPACTION_TYPE_TWCS: &str = "twcs"; +/// Prefix for TWCS compaction options. +pub const COMPACTION_TWCS_PREFIX: &str = "compaction.twcs."; /// Option key for twcs min file num to trigger a compaction. pub const TWCS_TRIGGER_FILE_NUM: &str = "compaction.twcs.trigger_file_num"; /// Option key for twcs min file num to trigger compaction in the active window. @@ -57,6 +59,20 @@ pub const TWCS_TIME_WINDOW: &str = "compaction.twcs.time_window"; pub const TWCS_REMOTE_COMPACTION: &str = "compaction.twcs.remote_compaction"; /// Option key for twcs fallback to local. pub const TWCS_FALLBACK_TO_LOCAL: &str = "compaction.twcs.fallback_to_local"; +/// Table options that configure compaction. +pub const COMPACTION_OPTION_KEYS: [&str; 11] = [ + COMPACTION_TYPE, + COMPACTION_OVERRIDE, + TWCS_TRIGGER_FILE_NUM, + TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, + TWCS_ACTIVE_WINDOW_L1_MERGE_TRIGGER, + TWCS_INACTIVE_WINDOW_TRIGGER_FILE_NUM, + TWCS_INACTIVE_WINDOW_L1_MERGE_TRIGGER, + TWCS_MAX_OUTPUT_FILE_SIZE, + TWCS_TIME_WINDOW, + TWCS_REMOTE_COMPACTION, + TWCS_FALLBACK_TO_LOCAL, +]; /// Option key for memtable type. pub const MEMTABLE_TYPE: &str = "memtable.type"; /// Option key for bulk memtable merge threshold. @@ -145,42 +161,32 @@ pub fn normalize_twcs_trigger_options( /// Returns true if the `key` is a valid option key for the mito engine. pub fn is_mito_engine_option_key(key: &str) -> bool { - [ - "ttl", - WRITE_BUFFER_SIZE_KEY, - AUTO_FLUSH_INTERVAL_KEY, - COMPACTION_TYPE, - COMPACTION_OVERRIDE, - TWCS_TRIGGER_FILE_NUM, - TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, - TWCS_ACTIVE_WINDOW_L1_MERGE_TRIGGER, - TWCS_INACTIVE_WINDOW_TRIGGER_FILE_NUM, - TWCS_INACTIVE_WINDOW_L1_MERGE_TRIGGER, - TWCS_MAX_OUTPUT_FILE_SIZE, - TWCS_TIME_WINDOW, - TWCS_REMOTE_COMPACTION, - TWCS_FALLBACK_TO_LOCAL, - "storage", - "index.inverted_index.ignore_column_ids", - "index.inverted_index.segment_row_count", - WAL_OPTIONS_KEY, - MEMTABLE_TYPE, - MEMTABLE_BULK_MERGE_THRESHOLD, - MEMTABLE_BULK_ENCODE_ROW_THRESHOLD, - MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD, - MEMTABLE_BULK_MAX_MERGE_GROUPS, - MEMTABLE_PARTITION_TREE_INDEX_MAX_KEYS_PER_SHARD, - MEMTABLE_PARTITION_TREE_DATA_FREEZE_THRESHOLD, - MEMTABLE_PARTITION_TREE_FORK_DICTIONARY_BYTES, - // We don't allow to create a mito table with sparse primary key encoding directly. - APPEND_MODE_KEY, - MERGE_MODE_KEY, - SST_FORMAT_KEY, - MAX_ROW_GROUP_ROW_COUNT, - PRESERVE_ROW_SEQUENCE, - EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, - ] - .contains(&key) + COMPACTION_OPTION_KEYS.contains(&key) + || [ + "ttl", + WRITE_BUFFER_SIZE_KEY, + AUTO_FLUSH_INTERVAL_KEY, + "storage", + "index.inverted_index.ignore_column_ids", + "index.inverted_index.segment_row_count", + WAL_OPTIONS_KEY, + MEMTABLE_TYPE, + MEMTABLE_BULK_MERGE_THRESHOLD, + MEMTABLE_BULK_ENCODE_ROW_THRESHOLD, + MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD, + MEMTABLE_BULK_MAX_MERGE_GROUPS, + MEMTABLE_PARTITION_TREE_INDEX_MAX_KEYS_PER_SHARD, + MEMTABLE_PARTITION_TREE_DATA_FREEZE_THRESHOLD, + MEMTABLE_PARTITION_TREE_FORK_DICTIONARY_BYTES, + // We don't allow to create a mito table with sparse primary key encoding directly. + APPEND_MODE_KEY, + MERGE_MODE_KEY, + SST_FORMAT_KEY, + MAX_ROW_GROUP_ROW_COUNT, + PRESERVE_ROW_SEQUENCE, + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, + ] + .contains(&key) } #[cfg(test)] From 863eebdecbb9c11e7e3249b6effb43a0c53568f3 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 14:51:35 +0000 Subject: [PATCH 32/46] feat(frontend): retry transient distributed write failures (#9543) * refactor(operator): route shared requests through NodeManager Remove the concrete NodeClients requester API and use the existing node manager for owned and shared requests. Preserve the original write hint while the legacy error adapter remains in place. Signed-off-by: WenyXu * fix(meta): preserve retry policy through error hints Propagate client errors through External and let DDL read retry hints while preserving cleanup and terminal-child policies. Region transport failures retain the legacy Retryable classification, which now also participates in bounded write retries. Signed-off-by: WenyXu * fix(operator): restore sparse partition tags when rebuilding writes Decode retained primary keys when refreshed partition expressions need different tags. Keep the partitioner and multi-region sparse write regression tests with the fix. Signed-off-by: WenyXu * fix(operator): stop detached write retries when callers cancel Signed-off-by: WenyXu * feat(frontend): configure retries across distributed write paths Wire the shared dispatcher through frontend and logical batching, preserving partition columns and request headers. Keep standalone retries disabled and include configuration, batching, and storage regression coverage. Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 1 + config/config.md | 5 + config/frontend.example.toml | 12 + src/client/src/client_manager.rs | 25 +- src/client/src/error.rs | 25 +- src/client/src/region.rs | 134 ++++- src/common/meta/src/ddl/drop_table.rs | 4 +- src/common/meta/src/ddl/undrop_table.rs | 60 ++- src/common/meta/src/ddl/utils.rs | 74 ++- src/common/meta/src/error.rs | 43 +- src/common/meta/src/node_manager.rs | 12 +- .../src/reconciliation/reconcile_catalog.rs | 3 +- .../src/reconciliation/reconcile_database.rs | 3 +- .../reconcile_logical_tables.rs | 3 +- .../src/reconciliation/reconcile_table.rs | 3 +- src/common/meta/src/test_util.rs | 3 +- src/common/procedure/src/error.rs | 21 +- src/frontend/AGENTS.md | 7 +- src/frontend/src/error.rs | 10 +- src/frontend/src/frontend.rs | 72 ++- src/frontend/src/instance/builder.rs | 129 ++++- src/frontend/src/instance/logical_batcher.rs | 1 + src/frontend/src/service_config.rs | 1 + src/metric-engine/src/engine/bulk_insert.rs | 36 +- src/operator/Cargo.toml | 1 + src/operator/src/bulk_insert.rs | 8 +- src/operator/src/delete.rs | 17 +- src/operator/src/insert.rs | 21 +- .../src/req_convert/common/partitioner.rs | 404 +++++++++++++++- src/operator/src/test_util.rs | 5 + .../src/test_util/partition_manager.rs | 6 +- src/operator/src/write_dispatcher.rs | 56 ++- src/operator/src/write_dispatcher/tests.rs | 456 ++++++++++++++++-- src/servers/AGENTS.md | 3 + src/servers/benches/flush_batch_physical.rs | 32 +- src/servers/src/batcher/logical_table.rs | 11 +- .../src/batcher/logical_table/batch.rs | 188 +++----- .../batcher/logical_table/batch_convert.rs | 84 +--- .../batcher/logical_table/pending_worker.rs | 35 +- .../src/batcher/logical_table/region_write.rs | 397 +++++---------- .../src/batcher/logical_table/test_util.rs | 17 +- src/standalone/src/datanode_manager.rs | 12 +- src/standalone/src/options.rs | 19 +- tests-integration/src/otlp.rs | 4 +- tests-integration/src/test_util.rs | 1 + 45 files changed, 1773 insertions(+), 691 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d84446150d6..e97dd203ee3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10136,6 +10136,7 @@ dependencies = [ "meta-client", "meter-core", "meter-macros", + "mito-codec", "moka", "object-store", "partition", diff --git a/config/config.md b/config/config.md index 5d76c8f7445..381abde24c4 100644 --- a/config/config.md +++ b/config/config.md @@ -391,6 +391,11 @@ | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | | `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". | | `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | +| `write_retry` | -- | -- | Retry retryable distributed write failures.
Backoff grows exponentially, with jitter to spread concurrent retries. | +| `write_retry.max_retries` | Integer | `3` | Maximum retries after the initial write attempt. Set to 0 to disable retries. | +| `write_retry.initial_backoff` | String | `1500ms` | Initial backoff before jitter is applied. Must be positive. | +| `write_retry.max_backoff` | String | `10s` | Maximum backoff before jitter is applied. Must be at least initial_backoff. | +| `write_retry.backoff_multiplier` | Float | `2.0` | Backoff multiplier. Must be finite and at least 1. | | `datanode` | -- | -- | Datanode options. | | `datanode.client` | -- | -- | Datanode client options. | | `datanode.client.connect_timeout` | String | 1s | Connect server timeout. Defaults to `1s`; this example sets it to `10s`. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 9a1c3db63fa..c42d254671c 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -398,6 +398,18 @@ memory_pool_size = "50%" #+ experimental_spill_compression = "uncompressed" # --- End spill-to-disk --- +## Retry retryable distributed write failures. +## Backoff grows exponentially, with jitter to spread concurrent retries. +[write_retry] +## Maximum retries after the initial write attempt. Set to 0 to disable retries. +max_retries = 3 +## Initial backoff before jitter is applied. Must be positive. +initial_backoff = "1500ms" +## Maximum backoff before jitter is applied. Must be at least initial_backoff. +max_backoff = "10s" +## Backoff multiplier. Must be finite and at least 1. +backoff_multiplier = 2.0 + ## Datanode options. [datanode] ## Datanode client options. diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs index ef26c107cd5..4d723361e80 100644 --- a/src/client/src/client_manager.rs +++ b/src/client/src/client_manager.rs @@ -52,7 +52,18 @@ impl Debug for NodeClients { #[async_trait::async_trait] impl DatanodeManager for NodeClients { async fn datanode(&self, datanode: &Peer) -> DatanodeRef { - Arc::new(self.make_region_requester(datanode).await) + let client = self.get_client(datanode).await; + + let ChannelConfig { + send_compression, + accept_compression, + .. + } = self.control_channel_manager.config(); + Arc::new(RegionRequester::new( + client, + *send_compression, + *accept_compression, + )) } } @@ -78,18 +89,6 @@ impl NodeClients { } } - /// Creates a [`RegionRequester`] with the cached client and control channel compression settings. - pub async fn make_region_requester(&self, datanode: &Peer) -> RegionRequester { - let client = self.get_client(datanode).await; - - let ChannelConfig { - send_compression, - accept_compression, - .. - } = self.control_channel_manager.config(); - RegionRequester::new(client, *send_compression, *accept_compression) - } - pub async fn get_client(&self, datanode: &Peer) -> Client { self.clients .get_with_by_ref(datanode, async move { diff --git a/src/client/src/error.rs b/src/client/src/error.rs index 5c70e5fbe91..fdf0df9582f 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -191,9 +191,22 @@ impl ErrorExt for Error { Error::Tonic { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } | Error::FlightStream { source, .. } - | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } | Error::External { source, .. } => source.retry_hint(), + Error::RegionServer { code, source, .. } => { + let hint = source.retry_hint(); + // Preserve the transport retry policy previously applied by should_retry(). + if !hint.is_retryable() + && matches!( + code, + Code::Unavailable | Code::Cancelled | Code::DeadlineExceeded + ) + { + RetryHint::Retryable + } else { + hint + } + } Error::ConvertFlightData { source, .. } | Error::CreateChannel { source, .. } | Error::CreateTlsChannel { source, .. } => source.retry_hint(), @@ -207,13 +220,9 @@ define_from_tonic_status!(Error, Tonic); impl From for common_meta::error::Error { fn from(error: Error) -> Self { - if error.should_retry() { - Self::retry_later(error) - } else { - Self::External { - source: BoxedError::new(error), - location: location!(), - } + Self::External { + source: BoxedError::new(error), + location: location!(), } } } diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 47c477b06b0..f9ebbb0d0bf 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -67,7 +67,7 @@ pub struct RegionRequester { #[async_trait] impl Datanode for RegionRequester { - async fn handle(&self, request: RegionRequest) -> MetaResult { + async fn handle_request(&self, request: Arc) -> MetaResult { self.handle_inner(request).await.map_err(Into::into) } @@ -220,16 +220,7 @@ impl RegionRequester { .await } .await - .map_err(|e| { - let code = e.code(); - // Uses `Error::RegionServer` instead of `Error::Server` - error::Error::RegionServer { - addr, - code, - source: BoxedError::new(error::Error::from(e)), - location: location!(), - } - })? + .map_err(|status| map_region_error(status, addr))? .into_inner(); check_response_header(&response.header)?; @@ -238,6 +229,15 @@ impl RegionRequester { } } +fn map_region_error(status: tonic::Status, addr: String) -> error::Error { + error::Error::RegionServer { + addr, + code: status.code(), + source: BoxedError::new(error::Error::from(status)), + location: location!(), + } +} + async fn recordbatches_from_flight_message_stream( addr: String, flight_message_stream: S, @@ -761,21 +761,104 @@ mod tests { assert_eq!(error.retry_hint(), baseline.retry_hint()); assert_eq!(error.tonic_code(), baseline.tonic_code()); } - let error = Datanode::handle(&requester, build_insert_request()) - .await - .unwrap_err(); + let datanode: &dyn Datanode = &requester; + let error = datanode.handle(build_insert_request()).await.unwrap_err(); assert_eq!(error.status_code(), StatusCode::RegionNotFound); assert_eq!(error.retry_hint(), hint); - match hint { - RetryHint::Retryable | RetryHint::RetryAfterRouteRefresh => assert!(matches!( - error, - meta_error::Error::RetryLater { - clean_poisons: false, + assert!(matches!(error, meta_error::Error::External { .. })); + } + } + + #[tokio::test] + async fn test_datanode_transport_retry_hints() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + let datanode: &dyn Datanode = &requester; + let statuses = [ + tonic::Code::Unavailable, + tonic::Code::Cancelled, + tonic::Code::DeadlineExceeded, + tonic::Code::FailedPrecondition, + ] + .into_iter() + .flat_map(|code| { + [ + None, + Some(RetryHint::NonRetryable), + Some(RetryHint::Retryable), + Some(RetryHint::RetryAfterRouteRefresh), + ] + .into_iter() + .map(move |hint| (code, hint)) + }); + for (code, hint) in statuses { + for body in [ + region_request::Body::Inserts(Default::default()), + region_request::Body::Deletes(Default::default()), + region_request::Body::BulkInsert(Default::default()), + region_request::Body::Create(Default::default()), + region_request::Body::RemoteDynFilter(Default::default()), + ] { + let request_type = body.as_ref().to_string(); + let mut status = Status::new(code, "request failed"); + // Metadata status is deliberately different from the transport status. + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::RegionNotReady as u32) + .to_string() + .parse() + .unwrap(), + ); + if let Some(hint) = hint { + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.as_str().parse().unwrap(), + ); + } + let original = Error::from(status.clone()); + let expected = if original.should_retry() && !original.retry_hint().is_retryable() { + RetryHint::Retryable + } else { + original.retry_hint() + }; + server + .service + .responses + .lock() + .unwrap() + .extend([Err(status.clone()), Err(status)]); + let request = RegionRequest { + header: None, + body: Some(body), + }; + let shared = Arc::new(request); + let shared_error = datanode.handle(Arc::clone(&shared)).await.unwrap_err(); + assert_eq!(Arc::strong_count(&shared), 1); + let request = Arc::try_unwrap(shared).unwrap(); + let owned_error = datanode.handle(request).await.unwrap_err(); + for error in [owned_error, shared_error] { + assert_eq!( + error.retry_hint(), + expected, + "{code:?}, {hint:?}, {request_type}" + ); + assert_eq!(error.retry_hint().is_retryable(), original.should_retry()); + assert_eq!(error.status_code(), StatusCode::RegionNotReady); + let meta_error::Error::External { source, .. } = error else { + panic!("expected an external error"); + }; + assert_eq!(source.retry_hint(), expected); + let Error::RegionServer { + source, + code: actual_code, .. - } - )), - RetryHint::NonRetryable => { - assert!(matches!(error, meta_error::Error::External { .. })) + } = source.as_any().downcast_ref::().unwrap() + else { + panic!("expected a region server error"); + }; + assert_eq!(*actual_code, code); + assert_eq!(source.retry_hint(), original.retry_hint()); + assert_eq!(source.to_string(), original.to_string()); } } } @@ -846,8 +929,11 @@ mod tests { baseline.status_code(), "shared: {error:?}; generated: {baseline:?}" ); + let Error::RegionServer { source, .. } = &error else { + panic!("expected a region server error: {error:?}"); + }; assert_eq!( - error.retry_hint(), + source.retry_hint(), baseline.retry_hint(), "shared: {error:?}; generated: {baseline:?}" ); diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index b73bd95e1bc..622a27d804b 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -19,6 +19,8 @@ use std::collections::HashMap; use async_trait::async_trait; use common_error::ext::BoxedError; +#[cfg(feature = "enterprise")] +use common_error::ext::ErrorExt; use common_event_recorder::Event; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ @@ -53,7 +55,7 @@ use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; #[cfg(feature = "enterprise")] fn ensure_retry_later(err: error::Error) -> error::Error { - if err.is_retry_later() { + if err.retry_hint().is_retryable() { err } else { error::Error::retry_later(err) diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 73a361ffd0e..efb1d4824d7 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -17,8 +17,8 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, - Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use common_telemetry::tracing_context::TracingContext; use common_telemetry::warn; @@ -156,6 +156,7 @@ impl UndropTableProcedure { } .await; if let Err(err) = restore_result { + // Keep this classification aligned with map_restore_error(). let should_cleanup_opened_regions = !err.is_retry_later(); let err = self.map_restore_metadata_error(err); if should_cleanup_opened_regions @@ -274,6 +275,19 @@ impl UndropTableProcedure { Ok(()) } + fn map_restore_error(err: error::Error) -> ProcedureError { + // Restore retries keep opened regions, while other failures close them. Preserve + // this policy even when the underlying metadata error carries a retry hint: + // retrying after cleanup would restore metadata without reopening the regions. + if err.is_retry_later() { + map_to_procedure_error(err) + } else if err.need_clean_poisons() { + ProcedureError::external_and_clean_poisons(err) + } else { + ProcedureError::external(err) + } + } + fn map_restore_metadata_error(&self, err: error::Error) -> error::Error { match err { error::Error::TombstoneTargetAlreadyExists { .. } => { @@ -413,7 +427,12 @@ impl Procedure for UndropTableProcedure { async fn execute(&mut self, _: &ProcedureContext) -> ProcedureResult { match self.data.state { UndropTableState::Prepare => self.on_prepare().await, - UndropTableState::RestoreMetadata => self.on_restore_metadata().await, + UndropTableState::RestoreMetadata => { + return self + .on_restore_metadata() + .await + .map_err(Self::map_restore_error); + } UndropTableState::OpenRegions => self.on_open_regions().await, UndropTableState::InvalidateTableCache => self.on_broadcast().await, } @@ -606,7 +625,8 @@ mod tests { use api::v1::region::region_request; use async_trait::async_trait; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; - use common_error::ext::ErrorExt; + use common_error::ext::{BoxedError, ErrorExt}; + use common_error::mock::MockError; use common_error::status_code::StatusCode; use store_api::storage::RegionId; use table::table_name::TableName; @@ -635,6 +655,38 @@ mod tests { } } + #[test] + fn test_restore_metadata_retry_policy() { + let source = error::Error::retry_later(MockError::new(StatusCode::Internal)); + let err = Err::<(), _>(BoxedError::new(source)) + .context(error::ExternalSnafu) + .unwrap_err(); + assert!(err.retry_hint().is_retryable()); + assert!(!err.is_retry_later()); + let err = UndropTableProcedure::map_restore_error(err); + assert!(!err.is_retry_later()); + assert!(!err.need_clean_poisons()); + + for retryable in [false, true] { + for clean_poisons in [false, true] { + let source = BoxedError::new(MockError::new(StatusCode::Internal)); + let err = if retryable { + error::Error::RetryLater { + source, + clean_poisons, + } + } else { + Err::<(), _>(source) + .context(error::AbortProcedureSnafu { clean_poisons }) + .unwrap_err() + }; + let err = UndropTableProcedure::map_restore_error(err); + assert_eq!(err.is_retry_later(), retryable); + assert_eq!(err.need_clean_poisons(), clean_poisons); + } + } + } + #[test] fn test_map_restore_metadata_error_without_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); diff --git a/src/common/meta/src/ddl/utils.rs b/src/common/meta/src/ddl/utils.rs index ebfbd6fd522..a8058d89b94 100644 --- a/src/common/meta/src/ddl/utils.rs +++ b/src/common/meta/src/ddl/utils.rs @@ -28,7 +28,7 @@ use api::v1::region::{ region_request, }; use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_procedure::error::Error as ProcedureError; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{error, info, warn}; @@ -61,7 +61,7 @@ use crate::wal_provider::RegionWalOptions; pub fn add_peer_context_if_needed(datanode: Peer) -> impl FnOnce(Error) -> Error { move |err| { error!(err; "Failed to operate datanode, peer: {}", datanode); - if !err.is_retry_later() { + if !err.retry_hint().is_retryable() { return Err::<(), BoxedError>(BoxedError::new(err)) .context(OperateDatanodeSnafu { peer: datanode }) .unwrap_err(); @@ -84,7 +84,7 @@ pub(crate) fn is_metric_engine_logical_table( /// This function determines whether the error should be retried and if poison cleanup is needed, /// then maps it to the appropriate procedure error variant. pub fn map_to_procedure_error(e: Error) -> ProcedureError { - match (e.is_retry_later(), e.need_clean_poisons()) { + match (e.retry_hint().is_retryable(), e.need_clean_poisons()) { (true, true) => ProcedureError::retry_later_and_clean_poisons(e), (true, false) => ProcedureError::retry_later(e), (false, true) => ProcedureError::external_and_clean_poisons(e), @@ -256,7 +256,7 @@ pub fn handle_multiple_results(results: Vec>) -> MultipleRes match result { Ok(value) => ok_results.push(value), Err(err) => { - if err.is_retry_later() { + if err.retry_hint().is_retryable() { retryable_results.push(err); } else { non_retryable_results.push(err); @@ -475,8 +475,74 @@ pub fn extract_column_metadatas( #[cfg(test)] mod tests { + use common_error::mock::MockError; + use common_error::status_code::StatusCode; + use super::*; + fn external_error(retryable: bool) -> Error { + let source = MockError::new(StatusCode::Internal); + let source = if retryable { + BoxedError::new(Error::retry_later(source)) + } else { + BoxedError::new(source) + }; + Err::<(), _>(source) + .context(error::ExternalSnafu) + .unwrap_err() + } + + #[test] + fn test_map_to_procedure_error() { + for retryable in [false, true] { + for clean_poisons in [false, true] { + let err = if retryable { + Error::RetryLater { + source: BoxedError::new(external_error(false)), + clean_poisons, + } + } else { + Err::<(), _>(BoxedError::new(external_error(true))) + .context(error::AbortProcedureSnafu { clean_poisons }) + .unwrap_err() + }; + let err = map_to_procedure_error(err); + assert_eq!(err.is_retry_later(), retryable); + assert_eq!(err.need_clean_poisons(), clean_poisons); + } + let err = map_to_procedure_error(external_error(retryable)); + assert_eq!(err.is_retry_later(), retryable); + assert!(!err.need_clean_poisons()); + } + } + + #[test] + fn test_handle_multiple_results_retry_hint() { + assert!(matches!( + handle_multiple_results(vec![Ok(()), Err(external_error(true))]), + MultipleResults::PartialRetryable(_) + )); + assert!(matches!( + handle_multiple_results::<()>(vec![ + Err(external_error(true)), + Err(external_error(true)) + ]), + MultipleResults::AllRetryable(_) + )); + assert!(matches!( + handle_multiple_results(vec![ + Ok(()), + Err(external_error(true)), + Err(external_error(false)) + ]), + MultipleResults::PartialNonRetryable(_) + )); + assert!(matches!( + handle_multiple_results::<()>(vec![Err(external_error(false))]), + MultipleResults::AllNonRetryable(_) + )); + } + #[test] fn test_get_catalog_and_schema() { let test_catalog = "my_catalog"; diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 297af792799..adfa158adf0 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -1373,15 +1373,15 @@ impl ErrorExt for Error { | StartProcedureManager { source, .. } | StopProcedureManager { source, .. } | RegisterProcedureLoader { source, .. } - | PutPoison { source, .. } - | ProcedureStateReceiver { source, .. } => source.retry_hint(), + | PutPoison { source, .. } => source.retry_hint(), External { source, .. } | ResponseExceededSizeLimit { source, .. } | OperateDatanode { source, .. } - | AbortProcedure { source, .. } | RegisterRepartitionProcedureLoader { source, .. } | CreateRepartitionProcedure { source, .. } | PersistRepartitionGcRequirement { source, .. } => source.retry_hint(), + // These errors end an operation; their sources cannot make it retryable. + AbortProcedure { .. } | ProcedureStateReceiver { .. } => RetryHint::NonRetryable, Table { source, .. } => source.retry_hint(), ConvertAlterTableRequest { source, .. } => source.retry_hint(), ConvertColumnDef { source, .. } => source.retry_hint(), @@ -1466,6 +1466,7 @@ mod retry_hint_tests { use std::sync::Arc; use common_error::mock::MockError; + use snafu::ResultExt; use super::*; @@ -1483,6 +1484,42 @@ mod retry_hint_tests { assert_eq!(err.retry_hint(), RetryHint::Retryable); } + #[test] + fn test_leader_retry_hints() { + for (error, expected) in [ + ( + NoLeaderSnafu { table_id: 1_u32 }.build(), + RetryHint::Retryable, + ), + (ElectionNoLeaderSnafu.build(), RetryHint::Retryable), + ( + ElectionLeaderLeaseExpiredSnafu.build(), + RetryHint::Retryable, + ), + ( + ElectionLeaderLeaseChangedSnafu.build(), + RetryHint::Retryable, + ), + ] { + assert_eq!(error.retry_hint(), expected, "{error:?}"); + } + } + + #[test] + fn test_procedure_error_retry_hints() { + let source = common_procedure::Error::retry_later(MockError::new(StatusCode::Internal)); + let err = Err::<(), _>(source).context(PutPoisonSnafu).unwrap_err(); + assert_eq!(err.retry_hint(), RetryHint::Retryable); + + let source = common_procedure::Error::retry_later(MockError::new(StatusCode::Internal)); + let err = Err::<(), _>(source) + .context(ProcedureStateReceiverSnafu { + procedure_id: ProcedureId::random(), + }) + .unwrap_err(); + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + } + #[test] fn test_get_cache_forwards_retry_hint() { let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal))); diff --git a/src/common/meta/src/node_manager.rs b/src/common/meta/src/node_manager.rs index 6cc6c3e1539..d18e2cc8471 100644 --- a/src/common/meta/src/node_manager.rs +++ b/src/common/meta/src/node_manager.rs @@ -28,12 +28,22 @@ use crate::peer::Peer; #[async_trait::async_trait] pub trait Datanode: Send + Sync { /// Handles DML, and DDL requests. - async fn handle(&self, request: RegionRequest) -> Result; + async fn handle_request(&self, request: Arc) -> Result; /// Handles query requests async fn handle_query(&self, request: QueryRequest) -> Result; } +impl dyn Datanode + '_ { + /// Handles an owned or shared request without cloning its payload. + pub async fn handle(&self, request: R) -> Result + where + R: Into> + Send, + { + self.handle_request(request.into()).await + } +} + pub type DatanodeRef = Arc; /// The trait for handling requests to flownode diff --git a/src/common/meta/src/reconciliation/reconcile_catalog.rs b/src/common/meta/src/reconciliation/reconcile_catalog.rs index e28e81d323d..f015a483b6e 100644 --- a/src/common/meta/src/reconciliation/reconcile_catalog.rs +++ b/src/common/meta/src/reconciliation/reconcile_catalog.rs @@ -16,6 +16,7 @@ use std::any::Any; use std::fmt::Debug; use std::time::Instant; +use common_error::ext::ErrorExt; use common_procedure::error::FromJsonSnafu; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, @@ -191,7 +192,7 @@ impl Procedure for ReconcileCatalogProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/reconciliation/reconcile_database.rs b/src/common/meta/src/reconciliation/reconcile_database.rs index f5e2fe9f9a0..49443aca9db 100644 --- a/src/common/meta/src/reconciliation/reconcile_database.rs +++ b/src/common/meta/src/reconciliation/reconcile_database.rs @@ -23,6 +23,7 @@ use std::fmt::Debug; use std::time::Instant; use async_trait::async_trait; +use common_error::ext::ErrorExt; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, @@ -231,7 +232,7 @@ impl Procedure for ReconcileDatabaseProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/reconciliation/reconcile_logical_tables.rs b/src/common/meta/src/reconciliation/reconcile_logical_tables.rs index 8c16e5a450f..5c5c654f6df 100644 --- a/src/common/meta/src/reconciliation/reconcile_logical_tables.rs +++ b/src/common/meta/src/reconciliation/reconcile_logical_tables.rs @@ -22,6 +22,7 @@ use std::any::Any; use std::fmt::Debug; use async_trait::async_trait; +use common_error::ext::ErrorExt; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, @@ -207,7 +208,7 @@ impl Procedure for ReconcileLogicalTablesProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/reconciliation/reconcile_table.rs b/src/common/meta/src/reconciliation/reconcile_table.rs index d48999b895f..a14fe13a3db 100644 --- a/src/common/meta/src/reconciliation/reconcile_table.rs +++ b/src/common/meta/src/reconciliation/reconcile_table.rs @@ -20,6 +20,7 @@ pub(crate) mod update_table_info; use std::fmt::Debug; +use common_error::ext::ErrorExt; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, @@ -318,7 +319,7 @@ impl Procedure for ReconcileTableProcedure { Ok(status) } Err(e) => { - if e.is_retry_later() { + if e.retry_hint().is_retryable() { metrics::METRIC_META_RECONCILIATION_PROCEDURE_ERROR .with_label_values(&[procedure_name, step, metrics::ERROR_TYPE_RETRYABLE]) .inc(); diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index d78dc815d09..618603c13e5 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -111,7 +111,8 @@ struct MockNode { #[async_trait::async_trait] impl Datanode for MockNode { - async fn handle(&self, request: RegionRequest) -> Result { + async fn handle_request(&self, request: Arc) -> Result { + let request = Arc::try_unwrap(request).expect("mock request must be exclusively owned"); self.handler.handle(&self.peer, request).await } diff --git a/src/common/procedure/src/error.rs b/src/common/procedure/src/error.rs index 67fab54b156..b3f4a83d5f4 100644 --- a/src/common/procedure/src/error.rs +++ b/src/common/procedure/src/error.rs @@ -306,7 +306,8 @@ impl ErrorExt for Error { | Error::DeletePoison { source, .. } | Error::GetPoison { source, .. } | Error::CheckStatus { source, .. } => source.retry_hint(), - Error::ProcedureExec { source, .. } => source.retry_hint(), + // The watcher reports this only after the child has failed or been poisoned. + Error::ProcedureExec { .. } => RetryHint::NonRetryable, Error::StartRemoveOutdatedMetaTask { source, .. } | Error::StopRemoveOutdatedMetaTask { source, .. } => source.retry_hint(), Error::DeleteState { error, .. } => retry_hint_from_opendal_error(error), @@ -379,8 +380,11 @@ mod tests { use std::sync::Arc; use common_error::mock::MockError; + use tokio::sync::watch; use super::*; + use crate::procedure::ProcedureState; + use crate::watcher; #[test] fn test_retry_later_hint_is_retryable() { @@ -397,6 +401,21 @@ mod tests { assert_eq!(err.retry_hint(), RetryHint::Retryable); } + #[tokio::test] + async fn test_finished_procedure_hint_is_non_retryable() { + let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal))); + for state in [ + ProcedureState::failed(source.clone()), + ProcedureState::poisoned(Default::default(), source), + ] { + let (sender, mut receiver) = watch::channel(ProcedureState::Running); + sender.send(state).unwrap(); + let err = watcher::wait(&mut receiver).await.unwrap_err(); + + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + } + } + #[test] fn test_retry_exceeded_hint_is_non_retryable() { let source = Arc::new(Error::retry_later(MockError::new(StatusCode::Internal))); diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index fdaa5dbcfe4..8d858242e8c 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -48,7 +48,9 @@ including caches owned by pipeline and protocol handlers. - **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` → `check_permission` → `operator`'s `Inserter` (schema validation, optional auto-create, partition routing, meter admission) → local `RegionServer` - (standalone) or RPC to datanodes (distributed). Arrow bulk inserts pass the + (standalone) or RPC to datanodes (distributed). `FrontendBuilder` configures + the operator write dispatcher with frontend retry options for distributed + inserts, deletes, and bulk writes. Arrow bulk inserts pass the request channel to `Inserter` and check meter admission for each nonempty batch. - Finite ingestion requests split internally admit their total rows per database before dispatch (`operator::insert::admit_write` / `admit_row_insert_batches`). @@ -113,7 +115,8 @@ cargo nextest run -p frontend `servers`; permissions, planning, and routing live here. - Standalone vs distributed diverge in datanode access (the `standalone` crate's local `RegionServer` adapter vs `NodeClients` RPC), MetaClient usage, and - whether heartbeat matters. In standalone, the cache invalidator is a no-op. + whether heartbeat matters. The builder defaults local invalidation to its + layered cache registry; callers may supply a custom invalidator. ## Maintenance contract diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index 9fd2312e736..451040c884a 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -323,6 +323,13 @@ pub enum Error { name: String, }, + #[snafu(display("Invalid write retry config: {}", reason))] + InvalidWriteRetryConfig { + reason: &'static str, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid tls config"))] InvalidTlsConfig { #[snafu(source)] @@ -391,7 +398,8 @@ impl ErrorExt for Error { | Error::IllegalAuthConfig { .. } | Error::ColumnNoneDefaultValue { .. } | Error::IncompleteGrpcRequest { .. } - | Error::InvalidTlsConfig { .. } => StatusCode::InvalidArguments, + | Error::InvalidTlsConfig { .. } + | Error::InvalidWriteRetryConfig { .. } => StatusCode::InvalidArguments, Error::NotSupported { .. } => StatusCode::Unsupported, diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index e594ba9f7b1..2ccbc3915e5 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -36,9 +36,10 @@ use crate::heartbeat::HeartbeatTask; use crate::instance::Instance; use crate::service_config::{ InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, - PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, + PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, WriteRetryOptions, }; +/// Configures frontend protocols, query execution, and distributed write retries. #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[serde(default)] pub struct FrontendOptions { @@ -74,6 +75,8 @@ pub struct FrontendOptions { pub meta_client: Option, pub logging: LoggingOptions, pub datanode: DatanodeClientOptions, + /// Retry policy for retryable distributed write failures. + pub write_retry: WriteRetryOptions, pub user_provider: Option, pub tracing: TracingOptions, pub query: QueryOptions, @@ -110,6 +113,7 @@ impl Default for FrontendOptions { meta_client: None, logging: LoggingOptions::default(), datanode: DatanodeClientOptions::default(), + write_retry: WriteRetryOptions::default(), user_provider: None, tracing: TracingOptions::default(), query: QueryOptions::default(), @@ -225,6 +229,72 @@ mod tests { type GrpcStream = Pin> + Send + Sync + 'static>>; + #[test] + fn test_write_retry_config() { + let defaults: FrontendOptions = toml::from_str("").unwrap(); + assert_eq!(defaults.write_retry, WriteRetryOptions::default()); + let options: FrontendOptions = toml::from_str( + r#" +[write_retry] +max_retries = 5 +initial_backoff = "250ms" +max_backoff = "5s" +backoff_multiplier = 3.0 +"#, + ) + .unwrap(); + assert_eq!( + options.write_retry, + WriteRetryOptions { + max_retries: 5, + initial_backoff: Duration::from_millis(250), + max_backoff: Duration::from_secs(5), + backoff_multiplier: 3.0, + } + ); + options.write_retry.validate().unwrap(); + let serialized = toml::to_string(&options).unwrap(); + let parsed: FrontendOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(options.write_retry, parsed.write_retry); + let disabled: FrontendOptions = toml::from_str("[write_retry]\nmax_retries = 0").unwrap(); + assert_eq!(disabled.write_retry.max_retries, 0); + assert_eq!( + disabled.write_retry.initial_backoff, + defaults.write_retry.initial_backoff + ); + } + + #[test] + fn test_write_retry_config_from_env() { + temp_env::with_vars( + [ + ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_RETRIES", Some("4")), + ( + "FRONTEND_RETRY_TEST__WRITE_RETRY__INITIAL_BACKOFF", + Some("500ms"), + ), + ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_BACKOFF", Some("8s")), + ( + "FRONTEND_RETRY_TEST__WRITE_RETRY__BACKOFF_MULTIPLIER", + Some("1.5"), + ), + ], + || { + let options = + FrontendOptions::load_layered_options(None, "FRONTEND_RETRY_TEST").unwrap(); + assert_eq!( + options.write_retry, + WriteRetryOptions { + max_retries: 4, + initial_backoff: Duration::from_millis(500), + max_backoff: Duration::from_secs(8), + backoff_multiplier: 1.5, + } + ); + }, + ); + } + #[test] fn test_logical_batcher_protocols_from_env() { temp_env::with_vars( diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 9db7b70c221..f93113d3426 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -25,8 +25,10 @@ use common_base::Plugins; use common_batcher::request_limiter::RequestLimiter; use common_datasource::object_store::LocalFileAccess; use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; -use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef}; -use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; +use common_meta::cache::{ + LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableInfoCacheRef, TableRouteCacheRef, +}; +use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; use common_meta::kv_backend::KvBackendRef; @@ -46,6 +48,7 @@ use operator::statement::{ StatementExecutor, StatementExecutorConfiguratorRef, StatementExecutorRef, }; use operator::table::TableMutationOperator; +use operator::write_dispatcher::WriteDispatcher; use partition::cache::PartitionInfoCacheRef; use partition::manager::PartitionRuleManager; use pipeline::pipeline_operator::PipelineOperator; @@ -70,7 +73,7 @@ pub struct FrontendBuilder { options: FrontendOptions, kv_backend: KvBackendRef, layered_cache_registry: LayeredCacheRegistryRef, - local_cache_invalidator: Option, + local_cache_invalidator: CacheInvalidatorRef, catalog_manager: CatalogManagerRef, node_manager: NodeManagerRef, plugins: Option, @@ -80,6 +83,9 @@ pub struct FrontendBuilder { } impl FrontendBuilder { + /// Creates a frontend builder from its configuration and shared services. + /// + /// Uses the layered cache registry for local invalidation unless overridden. #[allow(clippy::too_many_arguments)] pub fn new( options: FrontendOptions, @@ -93,8 +99,8 @@ impl FrontendBuilder { Self { options, kv_backend, + local_cache_invalidator: layered_cache_registry.clone(), layered_cache_registry, - local_cache_invalidator: None, catalog_manager, node_manager, plugins: None, @@ -140,9 +146,10 @@ impl FrontendBuilder { ) } + /// Overrides the default local cache invalidator. pub fn with_local_cache_invalidator(self, cache_invalidator: CacheInvalidatorRef) -> Self { Self { - local_cache_invalidator: Some(cache_invalidator), + local_cache_invalidator: cache_invalidator, ..self } } @@ -187,7 +194,11 @@ impl FrontendBuilder { } } + /// Validates configuration and assembles the frontend [`Instance`]. pub async fn try_build(self) -> Result { + if let Err(reason) = self.options.write_retry.validate() { + return error::InvalidWriteRetryConfigSnafu { reason }.fail(); + } let kv_backend = self.kv_backend; let node_manager = self.node_manager; let plugins = self.plugins.unwrap_or_default(); @@ -210,9 +221,24 @@ impl FrontendBuilder { partition_info_cache.clone(), )); - let local_cache_invalidator = self - .local_cache_invalidator - .unwrap_or_else(|| Arc::new(DummyCacheInvalidator)); + let local_cache_invalidator = self.local_cache_invalidator; + + let mut write_dispatcher = + WriteDispatcher::new(node_manager.clone(), partition_manager.clone()); + if self.options.write_retry.max_retries > 0 { + let table_info_cache: TableInfoCacheRef = + self.layered_cache_registry + .get() + .context(error::CacheRequiredSnafu { + name: cache::TABLE_INFO_CACHE_NAME, + })?; + write_dispatcher = write_dispatcher.with_retry( + self.options.write_retry.clone(), + table_info_cache, + local_cache_invalidator.clone(), + ); + } + let write_dispatcher = Arc::new(write_dispatcher); let region_query_handler = if let Some(factory) = plugins.get::() { @@ -236,6 +262,7 @@ impl FrontendBuilder { table_flownode_cache.clone(), self.options.auto_create_table, ) + .with_write_dispatcher(write_dispatcher.clone()) }; // The execution-only inserter owns no batchers, avoiding an Arc cycle. let bulk_inserter = Arc::new(create_inserter()); @@ -264,11 +291,14 @@ impl FrontendBuilder { create_inserter() .with_pending_rows_batcher(build_batcher(self.options.table_batcher_options())), ); - let deleter = Arc::new(Deleter::new( - self.catalog_manager.clone(), - partition_manager.clone(), - node_manager.clone(), - )); + let deleter = Arc::new( + Deleter::new( + self.catalog_manager.clone(), + partition_manager.clone(), + node_manager.clone(), + ) + .with_write_dispatcher(write_dispatcher), + ); let requester = Arc::new(Requester::new( self.catalog_manager.clone(), partition_manager.clone(), @@ -421,3 +451,76 @@ impl FrontendBuilder { }) } } + +#[cfg(test)] +mod tests { + use api::v1::meta::Role; + use common_meta::cache_invalidator::{Context, DummyCacheInvalidator}; + use common_meta::instruction::CacheIdent; + use common_meta::key::table_route::{TableRouteKey, TableRouteValue}; + use common_meta::key::{MetadataKey, MetadataValue}; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::rpc::store::PutRequest; + use meta_client::client::MetaClientBuilder; + use store_api::storage::RegionId; + + use super::*; + + #[tokio::test] + async fn test_local_cache_invalidator() { + for custom_invalidator in [false, true] { + let mut builder = FrontendBuilder::new_test( + &FrontendOptions::default(), + Arc::new(MetaClientBuilder::new(0, Role::Frontend).build()), + ); + if custom_invalidator { + builder = builder.with_local_cache_invalidator(Arc::new(DummyCacheInvalidator)); + } + let backend = builder.kv_backend.clone(); + let routes: TableRouteCacheRef = builder.layered_cache_registry.get().unwrap(); + let partitions: PartitionInfoCacheRef = builder.layered_cache_registry.get().unwrap(); + let instance = builder.try_build().await.unwrap(); + let table_id = 1024; + for region_number in [1, 2] { + let route = TableRouteValue::physical(vec![RegionRoute { + region: Region { + id: RegionId::new(table_id, region_number), + ..Default::default() + }, + ..Default::default() + }]); + backend + .put( + PutRequest::new() + .with_key(TableRouteKey::new(table_id).to_bytes()) + .with_value(route.try_as_raw_value().unwrap()), + ) + .await + .unwrap(); + if region_number == 1 { + routes.get(table_id).await.unwrap().unwrap(); + partitions.get(table_id).await.unwrap().unwrap(); + } + } + assert!(routes.contains_key(&table_id)); + assert!(partitions.contains_key(&table_id)); + + instance + .cache_invalidator() + .invalidate(&Context::default(), &[CacheIdent::TableId(table_id)]) + .await + .unwrap(); + + assert_eq!(custom_invalidator, routes.contains_key(&table_id)); + assert_eq!(custom_invalidator, partitions.contains_key(&table_id)); + let route = routes.get(table_id).await.unwrap().unwrap(); + let region = route.as_physical_table_route_ref().unwrap().region_routes[0] + .region + .id; + assert_eq!( + if custom_invalidator { 1 } else { 2 }, + region.region_number() + ); + } + } +} diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs index 255766ac14d..423f2996cae 100644 --- a/src/frontend/src/instance/logical_batcher.rs +++ b/src/frontend/src/instance/logical_batcher.rs @@ -48,6 +48,7 @@ impl Instance { LogicalTablePendingRowsBatcher::try_new( self.partition_manager().clone(), self.node_manager().clone(), + self.inserter.write_dispatcher().clone(), self.catalog_manager().clone(), self.table_flownode_set_cache().clone(), true, diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs index 2c761c4242a..535ae9b655b 100644 --- a/src/frontend/src/service_config.rs +++ b/src/frontend/src/service_config.rs @@ -25,6 +25,7 @@ pub use influxdb::{InfluxdbMergeMode, InfluxdbOptions}; pub use jaeger::JaegerOptions; pub use mysql::MysqlOptions; pub use opentsdb::OpentsdbOptions; +pub use operator::write_dispatcher::WriteRetryOptions; pub use otlp::OtlpOptions; pub use pending_rows_batcher::{BatcherOptions, PendingRowsBatcherOptions}; pub use postgres::PostgresOptions; diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index dd20f34adf8..b7ad7646e3c 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -354,17 +354,32 @@ mod tests { #[tokio::test] async fn test_bulk_insert_physical_region_passthrough() { - check_bulk_insert_physical_region_passthrough(false, false).await; - check_bulk_insert_physical_region_passthrough(true, false).await; + check_bulk_insert_physical_region_passthrough(false, false, false).await; + check_bulk_insert_physical_region_passthrough(true, false, false).await; } #[tokio::test] async fn test_bulk_insert_physical_region_close_recovery() { - check_bulk_insert_physical_region_passthrough(false, true).await; - check_bulk_insert_physical_region_passthrough(true, true).await; + check_bulk_insert_physical_region_passthrough(false, true, false).await; + check_bulk_insert_physical_region_passthrough(true, true, false).await; } - async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool, close: bool) { + #[tokio::test] + async fn test_physical_bulk_partition_columns_preserve_stored_rows() { + for skip_wal in [false, true] { + let stripped = + check_bulk_insert_physical_region_passthrough(skip_wal, true, false).await; + let retained = + check_bulk_insert_physical_region_passthrough(skip_wal, true, true).await; + assert_eq!(retained, stripped, "skip_wal={skip_wal}"); + } + } + + async fn check_bulk_insert_physical_region_passthrough( + skip_wal: bool, + close: bool, + retain_partition_columns: bool, + ) -> String { // Use flat format so that BulkMemtable is used (supports write_bulk). let mito_config = MitoConfig { default_flat_format: true, @@ -394,7 +409,11 @@ mod tests { index: 2, column_id: 2, // column_id for "job" in the physical table }]; - let non_tag_indices = vec![0, 1]; // timestamp, value + let non_tag_indices = if retain_partition_columns { + vec![0, 1, 2] // timestamp, value, and the job partition column + } else { + vec![0, 1] // timestamp, value + }; let second_batch = build_logical_batch(3, 3); let physical_batch = modify_batch_sparse( second_batch, @@ -403,6 +422,10 @@ mod tests { &non_tag_indices, ) .unwrap(); + assert_eq!( + physical_batch.column_by_name("job").is_some(), + retain_partition_columns + ); let request = build_bulk_request(physical_region_id, physical_batch, skip_wal); let response = env .metric() @@ -469,6 +492,7 @@ mod tests { batches.iter().map(|b| b.num_rows()).sum::(), if skip_wal && !close { 0 } else { 6 }, ); + batches.pretty_print().unwrap() } #[tokio::test] diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 7f4233e633e..b89bf4b13d3 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -62,6 +62,7 @@ lazy_static.workspace = true meta-client.workspace = true meter-core.workspace = true meter-macros.workspace = true +mito-codec.workspace = true moka = { workspace = true, features = ["future"] } object-store.workspace = true partition.workspace = true diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 29537959df5..d4a1ed0497c 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -35,6 +35,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfoRef; +use tokio_util::sync::CancellationToken; use crate::error::Result; use crate::insert::Inserter; @@ -72,7 +73,7 @@ impl Inserter { let writes = self.write_dispatcher.prepare(writes).await?; Ok(self .write_dispatcher - .execute(writes, None) + .execute(writes, None, None) .await .context(error::RequestInsertsSnafu)? .affected_rows) @@ -191,6 +192,7 @@ impl Inserter { Some( metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]), ), + None, ) .await .context(error::RequestRegionSnafu) @@ -217,6 +219,9 @@ impl Inserter { .with_label_values(&["wait_all_datanode"]) .start_timer(); + // Dropping the caller stops retries without abandoning in-flight writes. + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); let mut handles = Vec::with_capacity(region_masks.len()); for (region_number, mask) in region_masks { @@ -328,6 +333,7 @@ impl Inserter { &metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["datanode_handle"]), ), + Some(&cancellation), ) .await .context(error::RequestRegionSnafu) diff --git a/src/operator/src/delete.rs b/src/operator/src/delete.rs index c389374b7ac..9bf14471696 100644 --- a/src/operator/src/delete.rs +++ b/src/operator/src/delete.rs @@ -27,6 +27,7 @@ use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use table::TableRef; use table::requests::DeleteRequest as TableDeleteRequest; +use tokio_util::sync::CancellationToken; use crate::error::{ CatalogSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, MissingTimeIndexColumnSnafu, @@ -141,12 +142,16 @@ impl Deleter { let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_delete(requests); let writes = dispatcher.prepare(vec![request]).await?; - let affected_rows = - common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await }) - .await - .context(JoinTaskSnafu)? - .context(RequestDeletesSnafu)? - .affected_rows; + // Dropping the caller stops retries without abandoning in-flight writes. + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + let affected_rows = common_runtime::spawn_global(async move { + dispatcher.execute(writes, None, Some(&cancellation)).await + }) + .await + .context(JoinTaskSnafu)? + .context(RequestDeletesSnafu)? + .affected_rows; crate::metrics::DIST_DELETE_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 2b382dbdb5f..3f67d4fc576 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -80,6 +80,7 @@ use table::requests::{ validate_semantic_option, }; use table::table_reference::TableReference; +use tokio_util::sync::CancellationToken; use crate::batcher::PendingRowsBatcher; use crate::error::{ @@ -841,12 +842,16 @@ impl Inserter { let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_insert(normal_requests); let writes = dispatcher.prepare(vec![request]).await?; - let affected_rows = - common_runtime::spawn_global(async move { dispatcher.execute(writes, None).await }) - .await - .context(JoinTaskSnafu)? - .context(RequestInsertsSnafu)? - .affected_rows; + // Dropping the caller stops retries without abandoning in-flight writes. + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + let affected_rows = common_runtime::spawn_global(async move { + dispatcher.execute(writes, None, Some(&cancellation)).await + }) + .await + .context(JoinTaskSnafu)? + .context(RequestInsertsSnafu)? + .affected_rows; crate::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[ctx.get_db_string().as_str()]) .inc_by(affected_rows as u64); @@ -3599,9 +3604,9 @@ mod tests { #[async_trait::async_trait] impl common_meta::node_manager::Datanode for RecordingNode { - async fn handle( + async fn handle_request( &self, - _request: api::v1::region::RegionRequest, + _request: Arc, ) -> common_meta::error::Result { let rows = match _request.body.as_ref() { Some(api::v1::region::region_request::Body::Inserts(requests)) => requests diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index f1c3eff4fe3..6df73d7d1b9 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::sync::Arc; use api::v1::region::{ BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest, @@ -21,12 +22,17 @@ use api::v1::region::{ use api::v1::{ArrowIpc, PartitionExprVersion, Rows}; use common_error::ext::BoxedError; use common_grpc::flight::{self, FlightDecoder}; +use datatypes::arrow::array::{Array, BinaryArray, StringBuilder}; use datatypes::arrow::compute; +use datatypes::arrow::datatypes::{DataType, Field, Schema}; use datatypes::arrow::record_batch::RecordBatch; +use mito_codec::row_converter::sparse::{SparseOffsetsCache, SparsePrimaryKeyView}; use partition::PartitionRuleRef; use partition::manager::PartitionRuleManager; +use partition::partition::RegionMask; use partition::splitter::RowSplitter; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::storage::consts; use store_api::storage::{RegionId, RegionNumber, TableId}; use table::metadata::TableInfo; @@ -101,9 +107,11 @@ impl TablePartitioner { /// /// Takes ownership so row payloads can move into their new regions. Requests whose /// region and partition version still match retain their original payloads. + /// The table metadata comes from the same refresh that created this partitioner. pub(crate) fn rebuild_write_requests( &self, request: RegionRequest, + table_info: &TableInfo, ) -> Result> { let RegionRequest { header, body } = request; let bodies = match body.context(error::UnexpectedSnafu { @@ -111,7 +119,7 @@ impl TablePartitioner { })? { Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?, Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?, - Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk)?, + Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk, table_info)?, _ => { return error::UnexpectedSnafu { violated: "Only writes can be rebuilt", @@ -165,13 +173,23 @@ impl TablePartitioner { pub fn partition_bulk_insert_requests( &self, batch: RecordBatch, - mut original_ipc: Option, + original_ipc: Option, skip_wal: bool, ) -> Result> { let masks = self .rule .split_record_batch(&batch) .context(SplitInsertSnafu)?; + self.build_bulk_insert_requests(batch, masks, original_ipc, skip_wal) + } + + fn build_bulk_insert_requests( + &self, + batch: RecordBatch, + masks: HashMap, + mut original_ipc: Option, + skip_wal: bool, + ) -> Result> { let mut requests = Vec::with_capacity(masks.len()); for (region, mask) in masks { if mask.select_none() { @@ -256,7 +274,11 @@ impl TablePartitioner { Ok(bodies) } - fn rebuild_bulk_insert_request(&self, bulk: BulkInsertRequest) -> Result> { + fn rebuild_bulk_insert_request( + &self, + bulk: BulkInsertRequest, + table_info: &TableInfo, + ) -> Result> { let region_id = RegionId::from_u64(bulk.region_id); if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) { return Ok(vec![Body::BulkInsert(bulk)]); @@ -269,13 +291,92 @@ impl TablePartitioner { .and_then(|mut decoder| decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload)) .map_err(BoxedError::new) .context(error::ExternalSnafu)?; + let routing_batch = self.restore_partition_tags(&batch, table_info)?; + let masks = self + .rule + .split_record_batch(routing_batch.as_ref().unwrap_or(&batch)) + .context(SplitInsertSnafu)?; Ok(self - .partition_bulk_insert_requests(batch, Some(ipc), bulk.skip_wal)? + .build_bulk_insert_requests(batch, masks, Some(ipc), bulk.skip_wal)? .into_iter() .map(Body::BulkInsert) .collect()) } + /// Restores omitted physical-table tags for routing without changing the write payload. + fn restore_partition_tags( + &self, + batch: &RecordBatch, + table_info: &TableInfo, + ) -> Result> { + let meta = &table_info.meta; + if !table_info.is_physical_table() || meta.column_ids.len() != meta.schema.num_columns() { + return Ok(None); + } + let missing = self + .rule + .partition_columns() + .iter() + .filter(|name| batch.column_by_name(name).is_none()) + .filter_map(|name| { + let index = meta.schema.column_index_by_name(name)?; + (meta.primary_key_indices.contains(&index) + && meta.schema.column_schemas()[index].data_type.is_string()) + .then_some((name, meta.column_ids[index])) + }) + .collect::>(); + if missing.is_empty() { + return Ok(None); + } + let Some(primary_key) = batch.column_by_name(consts::PRIMARY_KEY_COLUMN_NAME) else { + return Ok(None); + }; + let primary_key = + primary_key + .as_any() + .downcast_ref::() + .context(error::UnexpectedSnafu { + violated: "Sparse primary key column must be Binary", + })?; + let mut builders = missing + .iter() + .map(|_| StringBuilder::new()) + .collect::>(); + let mut cache = SparseOffsetsCache::new(); + let mut buffer = Vec::new(); + for row in 0..batch.num_rows() { + ensure!( + !primary_key.is_null(row), + error::UnexpectedSnafu { + violated: "Sparse primary key must not be null", + } + ); + let mut key = SparsePrimaryKeyView::new(primary_key.value(row), &mut cache) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + for ((_, column_id), builder) in missing.iter().zip(&mut builders) { + let label = key + .label(*column_id, &mut buffer) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + builder.append_option(label); + } + } + let schema = batch.schema(); + let mut fields = schema.fields().to_vec(); + let mut columns = batch.columns().to_vec(); + for ((name, _), mut builder) in missing.into_iter().zip(builders) { + fields.push(Arc::new(Field::new(name, DataType::Utf8, true))); + columns.push(Arc::new(builder.finish())); + } + RecordBatch::try_new( + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())), + columns, + ) + .map(Some) + .context(error::ComputeArrowSnafu) + } + /// Whether the request's region still exists with the same partition expression version. /// /// A leader change leaves the partition unchanged, so the original payload can be @@ -313,16 +414,305 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::rpc::store::RangeResponse; - use datatypes::arrow::array::Int32Array; + use datatypes::arrow::array::{Int32Array, StringArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema}; use datatypes::arrow::record_batch::RecordBatch; + use datatypes::data_type::ConcreteDataType; + use datatypes::schema; + use datatypes::value::{Value as DataValue, ValueRef}; + use mito_codec::row_converter::sparse; + use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec}; use moka::future::Cache; + use partition::expr; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; + use partition::multi_dim::MultiDimPartitionRule; + use store_api::metric_engine_consts; use super::*; use crate::error::Error; use crate::test_util; + use crate::test_util::{test_new_partition_info_cache, test_new_table_route_cache}; + + fn build_sparse_partitioner() -> TablePartitioner { + TablePartitioner { + table_id: 1, + rule: Arc::new( + MultiDimPartitionRule::try_new( + vec!["zone".to_string()], + vec![1, 2], + vec![ + expr::col("zone").lt(DataValue::String("m".into())), + expr::col("zone").gt_eq(DataValue::String("m".into())), + ], + true, + ) + .unwrap(), + ), + versions: [ + (1, Some(PartitionExprVersion { value: 2 })), + (2, Some(PartitionExprVersion { value: 2 })), + ] + .into_iter() + .collect(), + } + } + + fn build_sparse_table_info() -> TableInfo { + let mut table_info = test_util::new_test_table_info(1, "physical_tags", [1, 2].into_iter()); + table_info.meta.schema = Arc::new(schema::Schema::new(vec![ + schema::ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + schema::ColumnSchema::new("zone", ConcreteDataType::string_datatype(), true), + ])); + table_info.meta.primary_key_indices = vec![0, 1]; + table_info.meta.partition_key_indices = vec![1]; + table_info.meta.column_ids = vec![10, 20]; + table_info.meta.options.extra_options.insert( + metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(), + String::new(), + ); + table_info + } + + fn build_sparse_batch() -> RecordBatch { + let codec = SparsePrimaryKeyCodec::schemaless(); + let keys = [Some("east"), Some("west"), Some(""), None, None] + .into_iter() + .enumerate() + .map(|(index, zone)| { + let mut key = Vec::new(); + let values = [ + ( + sparse::RESERVED_COLUMN_ID_TABLE_ID, + ValueRef::UInt32(42 + index as u32), + ), + ( + sparse::RESERVED_COLUMN_ID_TSID, + ValueRef::UInt64(index as u64), + ), + (10, ValueRef::String("node")), + ]; + codec.encode_value_refs(&values, &mut key).unwrap(); + if let Some(zone) = zone { + codec + .encode_raw_tag_value([(20, zone.as_bytes())].into_iter(), &mut key) + .unwrap(); + } + if index == 4 { + key.extend_from_slice(&20_u32.to_be_bytes()); + key.push(0); + } + key + }) + .collect::>(); + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["node"; keys.len()])), + Arc::new(BinaryArray::from_iter_values(keys.iter())), + ], + ) + .unwrap() + } + + #[tokio::test] + async fn test_partition_tags_use_physical_table_metadata() { + let backend = Arc::new(MemoryKvBackend::new()); + let table_route_cache = test_new_table_route_cache(backend.clone()); + let partition_info_cache = test_new_partition_info_cache(table_route_cache.clone()); + let manager = + PartitionRuleManager::new(backend.clone(), table_route_cache, partition_info_cache); + let table_info = build_sparse_table_info(); + let routes = [ + expr::col("zone").lt(DataValue::String("m".into())), + expr::col("zone").gt_eq(DataValue::String("m".into())), + ] + .into_iter() + .enumerate() + .map(|(index, expr)| RegionRoute { + region: Region { + id: RegionId::new(1, index as u32 + 1), + partition_expr: expr.as_json_str().unwrap(), + ..Default::default() + }, + ..Default::default() + }) + .collect(); + TableMetadataManager::new(backend) + .create_table_metadata( + table_info.clone(), + TableRouteValue::physical(routes), + Default::default(), + ) + .await + .unwrap(); + let batch = build_sparse_batch(); + let partitioner = TablePartitioner::try_new(&manager, &table_info) + .await + .unwrap(); + let restored = partitioner + .restore_partition_tags(&batch, &table_info) + .unwrap() + .unwrap(); + let zones = restored + .column_by_name("zone") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + zones.iter().collect::>(), + vec![Some("east"), Some("west"), Some(""), None, None] + ); + + for (physical, tag, valid_ids) in [ + (false, true, true), + (true, false, true), + (true, true, false), + ] { + let mut info = table_info.clone(); + if !physical { + info.meta + .options + .extra_options + .remove(metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY); + } + if !tag { + info.meta.primary_key_indices = vec![0]; + } + if !valid_ids { + info.meta.column_ids.clear(); + } + let partitioner = TablePartitioner::try_new(&manager, &info).await.unwrap(); + assert!( + partitioner + .restore_partition_tags(&batch, &info) + .unwrap() + .is_none() + ); + assert!( + partitioner + .partition_bulk_insert_requests(batch.clone(), None, false) + .is_err() + ); + } + } + + #[test] + fn test_rebuild_sparse_bulk_with_changed_partition_column() { + let partitioner = build_sparse_partitioner(); + let table_info = build_sparse_table_info(); + let batch = build_sparse_batch(); + let restored = partitioner + .restore_partition_tags(&batch, &table_info) + .unwrap() + .unwrap(); + let zones = restored + .column_by_name("zone") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + zones.iter().collect::>(), + vec![Some("east"), Some("west"), Some(""), None, None] + ); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap(); + let rebuilt = partitioner + .rebuild_write_requests( + build_request(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, 1).as_u64(), + partition_expr_version: Some(PartitionExprVersion { value: 1 }), + skip_wal: true, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + ..Default::default() + })), + &table_info, + ) + .unwrap(); + assert_eq!(rebuilt.len(), 2); + for request in rebuilt { + let Some(Body::BulkInsert(bulk)) = request.body else { + unreachable!() + }; + assert!(bulk.skip_wal); + assert_eq!( + bulk.partition_expr_version, + Some(PartitionExprVersion { value: 2 }) + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else { + unreachable!() + }; + let actual = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + let mask = if RegionId::from_u64(bulk.region_id).region_number() == 1 { + vec![true, false, true, true, true] + } else { + vec![false, true, false, false, false] + }; + let expected = compute::filter_record_batch(&batch, &mask.into()).unwrap(); + assert_eq!(actual, expected); + } + } + + #[test] + fn test_restore_partition_tags_checks_only_missing_tags() { + let partitioner = build_sparse_partitioner(); + let table_info = build_sparse_table_info(); + let batch = build_sparse_batch(); + let invalid = RecordBatch::try_new( + batch.schema(), + vec![ + batch.column(0).clone(), + Arc::new(BinaryArray::from_iter_values([b"invalid"; 5])), + ], + ) + .unwrap(); + assert!( + partitioner + .restore_partition_tags(&invalid, &table_info) + .is_err() + ); + let mut fields = invalid.schema().fields().to_vec(); + fields.push(Arc::new(Field::new("zone", DataType::Utf8, true))); + let mut columns = invalid.columns().to_vec(); + columns.push(Arc::new(StringArray::from(vec![Some("east"); 5]))); + let complete = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + assert!( + partitioner + .restore_partition_tags(&complete, &table_info) + .unwrap() + .is_none() + ); + let (schema, data_header, payload) = flight::record_batch_to_ipc(invalid).unwrap(); + let request = build_request(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, 1).as_u64(), + partition_expr_version: Some(PartitionExprVersion { value: 2 }), + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + ..Default::default() + })); + let expected = request.clone(); + assert_eq!( + partitioner + .rebuild_write_requests(request, &table_info) + .unwrap(), + vec![expected] + ); + } #[tokio::test] async fn test_partition_requests_preserve_routing_errors() { @@ -693,6 +1083,6 @@ mod tests { let partitioner = TablePartitioner::try_new(manager, &table_info) .await .context(error::InvalidPartitionSnafu)?; - partitioner.rebuild_write_requests(request) + partitioner.rebuild_write_requests(request, &table_info) } } diff --git a/src/operator/src/test_util.rs b/src/operator/src/test_util.rs index aa4a2f30ab8..436c1bbd7b5 100644 --- a/src/operator/src/test_util.rs +++ b/src/operator/src/test_util.rs @@ -17,3 +17,8 @@ mod partition_manager; pub use crate::test_util::kv_backend::prepare_mocked_backend; pub use crate::test_util::partition_manager::{create_partition_rule_manager, new_test_table_info}; + +#[cfg(test)] +pub(crate) use self::partition_manager::{ + test_new_partition_info_cache, test_new_table_route_cache, +}; diff --git a/src/operator/src/test_util/partition_manager.rs b/src/operator/src/test_util/partition_manager.rs index f80e2cc4fe2..260252d2da5 100644 --- a/src/operator/src/test_util/partition_manager.rs +++ b/src/operator/src/test_util/partition_manager.rs @@ -113,7 +113,7 @@ fn new_test_region_wal_options(regions: Vec) -> RegionWalOptions { HashMap::default() } -fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef { +pub(crate) fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef { let cache = CacheBuilder::new(128).build(); Arc::new(new_table_route_cache( "table_route_cache".to_string(), @@ -122,7 +122,9 @@ fn test_new_table_route_cache(kv_backend: KvBackendRef) -> TableRouteCacheRef { )) } -fn test_new_partition_info_cache(table_route_cache: TableRouteCacheRef) -> PartitionInfoCacheRef { +pub(crate) fn test_new_partition_info_cache( + table_route_cache: TableRouteCacheRef, +) -> PartitionInfoCacheRef { let cache = CacheBuilder::new(128).build(); Arc::new(new_partition_info_cache( "partition_info_cache".to_string(), diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index c2ac6543411..da8a02eddcc 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -22,7 +22,6 @@ use std::sync::Arc; use api::region::RegionResponse; use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body}; use backon::{BackoffBuilder, ExponentialBuilder}; -use client::client_manager::NodeClients; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_meta::cache::TableInfoCacheRef; use common_meta::cache_invalidator::{CacheInvalidatorRef, Context}; @@ -36,6 +35,7 @@ use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use prometheus::Histogram; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; +use tokio_util::sync::CancellationToken; pub use self::options::WriteRetryOptions; use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; @@ -75,14 +75,13 @@ impl WriteDispatcher { /// Enables shared remote sends and retries using frontend configuration and caches. pub fn with_retry( mut self, - clients: Arc, options: WriteRetryOptions, table_info_cache: TableInfoCacheRef, cache_invalidator: CacheInvalidatorRef, ) -> Self { if options.max_retries > 0 { self.retry = Some(WriteRetry { - clients, + node_manager: self.node_manager.clone(), partition_manager: self.partition_manager.clone(), options, table_info_cache, @@ -170,10 +169,12 @@ impl WriteDispatcher { } /// Sends requests using routes already resolved by the caller. + /// Cancellation stops future retries while allowing in-flight requests to complete. pub(crate) async fn execute( &self, writes: Vec<(PeerRef, RegionRequest)>, histogram: Option, + cancellation: Option<&CancellationToken>, ) -> Result { let histogram = histogram.as_ref(); let writes = if self.retry.is_some() { @@ -194,7 +195,7 @@ impl WriteDispatcher { .map(|(peer, request)| self.send(peer, request, histogram)), ) .await; - self.finish(attempts, histogram).await + self.finish(attempts, histogram, cancellation).await } /// Sends the first attempt immediately so bulk encoding can overlap other sends. @@ -231,10 +232,11 @@ impl WriteDispatcher { &self, attempts: Vec, histogram: Option<&Histogram>, + cancellation: Option<&CancellationToken>, ) -> Result { if let Some(retry) = &self.retry { if retry.should_retry(&attempts, 0) { - return retry.retry(attempts, histogram).await; + return retry.retry(attempts, histogram, cancellation).await; } let mut response = RegionResponse::new(0); let failures = merge_attempts(attempts, &mut response)?; @@ -271,13 +273,13 @@ impl WriteDispatcher { .await .map_err(BoxedError::new) .context(ExternalSnafu)?; - self.execute(writes, histogram).await + self.execute(writes, histogram, None).await } } /// Sends shared requests and coordinates retries with refreshed routes. struct WriteRetry { - clients: Arc, + node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef, options: WriteRetryOptions, table_info_cache: TableInfoCacheRef, @@ -293,15 +295,12 @@ pub(crate) struct RoutedWrite { impl WriteRetry { async fn send(&self, write: RoutedWrite, histogram: Option<&Histogram>) -> WriteAttempt { let _timer = histogram.map(Histogram::start_timer); - let requester = self.clients.make_region_requester(&write.peer).await; + let requester = self.node_manager.datanode(&write.peer).await; match requester.handle(Arc::clone(&write.request)).await { Ok(response) => WriteAttempt::Success(response), Err(error) => { let request = error.retry_hint().is_retryable().then_some(write); - WriteAttempt::Failed { - error: error.into(), - request, - } + WriteAttempt::Failed { error, request } } } } @@ -327,6 +326,7 @@ impl WriteRetry { &self, mut attempts: Vec, histogram: Option<&Histogram>, + cancellation: Option<&CancellationToken>, ) -> Result { let mut response = RegionResponse::new(0); let mut backoff = ExponentialBuilder::default() @@ -341,9 +341,20 @@ impl WriteRetry { let Some(delay) = backoff.next() else { break; }; + if let Some(cancellation) = cancellation { + tokio::select! { + biased; + _ = cancellation.cancelled() => break, + _ = tokio::time::sleep(delay) => {} + } + } else { + tokio::time::sleep(delay).await; + } + if cancellation.is_some_and(CancellationToken::is_cancelled) { + break; + } let failures = merge_attempts(attempts, &mut response)?; - tokio::time::sleep(delay).await; - + let mut first_error = None; let mut refresh = Vec::new(); let mut pending = Vec::new(); for (write, error) in failures { @@ -352,10 +363,16 @@ impl WriteRetry { } else { pending.push(write); } + first_error.get_or_insert(error); } if !refresh.is_empty() { pending.extend(self.reroute_failed_writes(refresh).await?); } + if let Some(error) = + first_error.filter(|_| cancellation.is_some_and(CancellationToken::is_cancelled)) + { + return Err(error); + } attempts = future::join_all(pending.into_iter().map(|write| self.send(write, histogram))) .await; @@ -390,7 +407,7 @@ impl WriteRetry { .await .map_err(BoxedError::new) .context(ExternalSnafu)?; - partitioners.insert(table_id, partitioner); + partitioners.insert(table_id, (table_info, partitioner)); } let mut rebuilt_requests = Vec::new(); for write in writes { @@ -401,11 +418,12 @@ impl WriteRetry { .build() })?; let table_id = write_region_id(&request)?.table_id(); - let partitioner = partitioners.get(&table_id).context(UnexpectedSnafu { - err_msg: "Refreshed write has no partition snapshot", - })?; + let (table_info, partitioner) = + partitioners.get(&table_id).context(UnexpectedSnafu { + err_msg: "Refreshed write has no partition snapshot", + })?; let requests = partitioner - .rebuild_write_requests(request) + .rebuild_write_requests(request, table_info) .map_err(BoxedError::new) .context(ExternalSnafu)?; rebuilt_requests.extend(requests); diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 250fe87f353..c8daf73f089 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -44,22 +44,30 @@ use common_meta::instruction::CacheIdent; use common_meta::key::TableMetadataManager; use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; use common_meta::kv_backend::KvBackendRef; +use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::peer::Peer; +use common_meta::rpc::router::{self, RegionRoute}; use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; -use datatypes::arrow::array::Int32Array; +use datatypes::arrow::array::{BinaryArray, Int32Array, StringArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema; +use datatypes::value::{self, ValueRef}; +use mito_codec::row_converter::sparse; +use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec}; use partition::cache::PartitionInfoCacheRef; -use partition::expr::{Operand, PartitionExpr, RestrictedOp}; +use partition::expr::{self, Operand, PartitionExpr, RestrictedOp}; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; -use store_api::storage::RegionId; +use store_api::metric_engine_consts; +use store_api::storage::{RegionId, consts}; use tokio::net::TcpListener; -use tokio::sync::{Barrier, Notify}; +use tokio::sync::{Barrier, Notify, Semaphore}; use tokio::task::JoinHandle; use tokio_stream::wrappers::TcpListenerStream; -use tonic::{Request, Response, Status}; +use tonic::{Code, Request, Response, Status}; use super::*; use crate::error::Error; @@ -68,7 +76,8 @@ use crate::test_util; #[derive(Clone, Default)] struct RecordingRegionService { requests: Arc>>, - rejections: Arc>>>, + rejections: Arc>>>, + release: Arc>>>, } #[async_trait::async_trait] @@ -80,18 +89,17 @@ impl Region for RecordingRegionService { let request = request.into_inner(); let region_id = write_region_id(&request).unwrap().as_u64(); self.requests.lock().unwrap().push(request); - if let Some(hint) = self + let release = self.release.lock().unwrap().clone(); + if let Some(release) = release { + release.acquire().await.unwrap().forget(); + } + if let Some(status) = self .rejections .lock() .unwrap() .get_mut(®ion_id) .and_then(VecDeque::pop_front) { - let mut status = Status::unavailable(format!("region {region_id} rejected the write")); - status.metadata_mut().insert( - common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, - hint.as_str().parse().unwrap(), - ); return Err(status); } Ok(Response::new(PbRegionResponse { @@ -177,8 +185,7 @@ impl RegionTestServer { let registry = cache::build_fundamental_cache_registry(backend); let table_info_cache = registry.get::().unwrap(); let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - WriteDispatcher::new(clients.clone(), partition_manager).with_retry( - clients, + WriteDispatcher::new(clients, partition_manager).with_retry( build_retry_options(max_retries), table_info_cache, self.invalidator.clone(), @@ -186,9 +193,22 @@ impl RegionTestServer { } fn fail_region(&self, region_number: u32, hints: impl IntoIterator) { + let region_id = RegionId::new(1, region_number).as_u64(); self.service.rejections.lock().unwrap().insert( - RegionId::new(1, region_number).as_u64(), - hints.into_iter().collect(), + region_id, + hints + .into_iter() + .map(|hint| { + let mut status = Status::failed_precondition(format!( + "region {region_id} rejected the write" + )); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.as_str().parse().unwrap(), + ); + status + }) + .collect(), ); } @@ -321,9 +341,7 @@ async fn test_write_retry_attempt_budget_and_nonretryable_errors() { .dispatch(vec![build_insert_request(&[1])]) .await .unwrap_err(); - // The dispatcher uses the explicit client hint for its attempt budget. - // The existing common-meta conversion promotes Unavailable to RetryLater. - assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert_eq!(error.retry_hint(), hint); assert_eq!( server.attempt_counts()[&RegionId::new(1, 1).as_u64()], expected_attempts @@ -332,6 +350,84 @@ async fn test_write_retry_attempt_budget_and_nonretryable_errors() { } } +#[tokio::test] +async fn test_write_retry_transport_errors() { + for code in [Code::Unavailable, Code::Cancelled, Code::DeadlineExceeded] { + for hint in [ + None, + Some(RetryHint::NonRetryable), + Some(RetryHint::Retryable), + Some(RetryHint::RetryAfterRouteRefresh), + ] { + let server = RegionTestServer::start().await; + let region_id = RegionId::new(1, 1).as_u64(); + let mut status = Status::new(code, "write failed"); + if let Some(hint) = hint { + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + hint.as_str().parse().unwrap(), + ); + } + server + .service + .rejections + .lock() + .unwrap() + .insert(region_id, VecDeque::from([status])); + let dispatcher = server.create_dispatcher(2).await; + let attempt = dispatcher + .send( + Arc::new(server.peer.clone()), + build_insert_request(&[1]), + None, + ) + .await; + let WriteAttempt::Failed { error, request } = attempt else { + panic!("expected a failed write"); + }; + assert!(request.is_some(), "{code:?}, {hint:?}"); + assert_eq!( + error.retry_hint(), + hint.filter(|hint| hint.is_retryable()) + .unwrap_or(RetryHint::Retryable) + ); + assert_eq!(server.attempt_counts(), HashMap::from([(region_id, 1)])); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn test_write_retry_transport_attempt_budget() { + for code in [Code::Unavailable, Code::Cancelled, Code::DeadlineExceeded] { + for failures in [1, 3] { + let server = RegionTestServer::start().await; + let region_id = RegionId::new(1, 1).as_u64(); + server.service.rejections.lock().unwrap().insert( + region_id, + (0..failures) + .map(|_| Status::new(code, "write failed")) + .collect(), + ); + let dispatcher = server.create_dispatcher(2).await; + let response = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; + let expected_attempts = if failures == 1 { + assert_eq!(response.unwrap().affected_rows, 1); + 2 + } else { + assert_eq!(response.unwrap_err().retry_hint(), RetryHint::Retryable); + 3 + }; + assert_eq!( + server.attempt_counts(), + HashMap::from([(region_id, expected_attempts)]), + "{code:?}, failures: {failures}" + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } + } +} + #[tokio::test] async fn test_write_retry_stops_when_a_sibling_is_nonretryable() { for hint in [RetryHint::Retryable, RetryHint::RetryAfterRouteRefresh] { @@ -364,14 +460,13 @@ async fn test_should_retry_requires_a_retained_request_and_budget() { let backend = test_util::prepare_mocked_backend().await; let (partition_manager, table_info_cache, cache_invalidator) = create_routing_caches(backend); let retry = WriteRetry { - clients: Arc::new(NodeClients::default()), + node_manager: Arc::new(NodeClients::default()), partition_manager, options: build_retry_options(2), table_info_cache, cache_invalidator, }; let failure = |request| { - // The converted error alone cannot distinguish an explicitly nonretryable attempt. let error: client::Error = Status::unavailable("write rejected").into(); WriteAttempt::Failed { error: error.into(), @@ -402,6 +497,96 @@ async fn test_should_retry_requires_a_retained_request_and_budget() { } } +#[tokio::test] +async fn test_write_retry_cancellation_stops_backoff() { + for hint in [RetryHint::Retryable, RetryHint::RetryAfterRouteRefresh] { + let server = RegionTestServer::start().await; + server.fail_region(1, [hint]); + let mut dispatcher = server.create_dispatcher(2).await; + let options = &mut dispatcher.retry.as_mut().unwrap().options; + options.initial_backoff = Duration::from_secs(60); + options.max_backoff = Duration::from_secs(60); + let attempt = dispatcher + .send( + Arc::new(server.peer.clone()), + build_insert_request(&[1]), + None, + ) + .await; + let cancellation = CancellationToken::new(); + let guard = cancellation.clone().drop_guard(); + let finish = dispatcher.finish(vec![attempt], None, Some(&cancellation)); + tokio::pin!(finish); + futures::future::poll_fn(|cx| { + assert!(finish.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(guard); + let error = tokio::time::timeout(Duration::from_secs(1), finish) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.retry_hint(), hint); + assert_eq!( + server.attempt_counts(), + HashMap::from([(RegionId::new(1, 1).as_u64(), 1)]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } +} + +#[tokio::test] +async fn test_write_retry_cancellation_drains_inflight_requests() { + let server = RegionTestServer::start().await; + server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let release = Arc::new(Semaphore::new(0)); + *server.service.release.lock().unwrap() = Some(release.clone()); + let dispatcher = server.create_dispatcher(2).await; + let writes = dispatcher + .prepare(vec![build_insert_request(&[1, 2])]) + .await + .unwrap(); + let cancellation = CancellationToken::new(); + let guard = cancellation.clone().drop_guard(); + let execute = dispatcher.execute(writes, None, Some(&cancellation)); + tokio::pin!(execute); + tokio::time::timeout( + Duration::from_secs(1), + futures::future::poll_fn(|cx| { + assert!(execute.as_mut().poll(cx).is_pending()); + if server.service.requests.lock().unwrap().len() == 2 { + Poll::Ready(()) + } else { + cx.waker().wake_by_ref(); + Poll::Pending + } + }), + ) + .await + .unwrap(); + drop(guard); + futures::future::poll_fn(|cx| { + assert!(execute.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + release.add_permits(2); + let error = tokio::time::timeout(Duration::from_secs(1), execute) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]) + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); +} + #[tokio::test] async fn test_write_retry_refreshes_routes_once_per_failed_round() { let server = RegionTestServer::start().await; @@ -444,8 +629,7 @@ async fn test_finish_retries_initial_failures_in_one_round() { ..Default::default() }); let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry( - clients, + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( build_retry_options(2), table_info_cache, invalidator.clone(), @@ -469,7 +653,7 @@ async fn test_finish_retries_initial_failures_in_one_round() { (RegionId::new(1, 3).as_u64(), 1), ]) ); - let response = dispatcher.finish(attempts, None).await.unwrap(); + let response = dispatcher.finish(attempts, None, None).await.unwrap(); assert_eq!(response.affected_rows, 3); assert_eq!( server.attempt_counts(), @@ -526,8 +710,7 @@ async fn test_write_retry_refreshes_stale_route_and_sends_only_failed_region_to_ ..Default::default() }); let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry( - clients, + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( build_retry_options(2), table_info_cache, invalidator.clone(), @@ -745,13 +928,11 @@ async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() { route_update: Mutex::new(Some((backend, current.update(routes).unwrap()))), }); let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager.clone()) - .with_retry( - clients, - build_retry_options(2), - table_info_cache, - invalidator.clone(), - ); + let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()).with_retry( + build_retry_options(2), + table_info_cache, + invalidator.clone(), + ); let response = dispatcher .dispatch(vec![ build_partitioned_request(kind, 1, old_versions[&1], &[70]), @@ -793,6 +974,188 @@ async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() { } } +fn build_sparse_bulk_request( + region: u32, + version: Option, + tags: &[(&str, &str)], +) -> (RegionRequest, RecordBatch) { + let codec = SparsePrimaryKeyCodec::schemaless(); + let keys = tags + .iter() + .enumerate() + .map(|(index, (host, zone))| { + let mut key = Vec::new(); + codec + .encode_value_refs( + &[ + (sparse::RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(42)), + ( + sparse::RESERVED_COLUMN_ID_TSID, + ValueRef::UInt64(index as u64), + ), + (10, ValueRef::String(host)), + (20, ValueRef::String(zone)), + ], + &mut key, + ) + .unwrap(); + key + }) + .collect::>(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false), + Field::new("host", DataType::Utf8, false), + ])), + vec![ + Arc::new(BinaryArray::from_iter_values(keys)), + Arc::new(StringArray::from_iter_values( + tags.iter().map(|(host, _)| *host), + )), + ], + ) + .unwrap(); + let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap(); + let request = RegionRequest { + body: Some(Body::BulkInsert(BulkInsertRequest { + region_id: RegionId::new(1, region).as_u64(), + partition_expr_version: version.map(|value| PartitionExprVersion { value }), + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + ..Default::default() + })), + ..Default::default() + }; + (request, batch) +} + +#[tokio::test] +async fn test_write_retry_restores_changed_partition_tags_for_failed_regions() { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend: KvBackendRef = Arc::new(MemoryKvBackend::new()); + let metadata = TableMetadataManager::new(backend.clone()); + let mut table_info = test_util::new_test_table_info(1, "physical", [1, 2].into_iter()); + table_info.meta.schema = Arc::new(schema::Schema::new(vec![ + schema::ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + schema::ColumnSchema::new("zone", ConcreteDataType::string_datatype(), true), + ])); + table_info.meta.primary_key_indices = vec![0, 1]; + table_info.meta.partition_key_indices = vec![0]; + table_info.meta.column_ids = vec![10, 20]; + table_info.meta.options.extra_options.insert( + metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(), + String::new(), + ); + let routes = [ + expr::col("host").lt(value::Value::String("m".into())), + expr::col("host").gt_eq(value::Value::String("m".into())), + ] + .into_iter() + .enumerate() + .map(|(index, expr)| RegionRoute { + region: router::Region { + id: RegionId::new(1, index as u32 + 1), + partition_expr: expr.as_json_str().unwrap(), + ..Default::default() + }, + leader_peer: Some(old_server.peer.clone()), + ..Default::default() + }) + .collect(); + metadata + .create_table_metadata( + table_info, + TableRouteValue::physical(routes), + Default::default(), + ) + .await + .unwrap(); + let (partition_manager, table_info_cache, registry) = create_routing_caches(backend.clone()); + let cached_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, old_versions) = partition_manager + .find_table_partition_rule(&cached_info) + .await + .unwrap(); + let current_info = metadata.table_info_manager().get(1).await.unwrap().unwrap(); + let mut updated_info = current_info.table_info.clone(); + updated_info.meta.partition_key_indices = vec![1]; + metadata + .update_table_info(¤t_info, None, updated_info) + .await + .unwrap(); + let current_route = metadata + .table_route_manager() + .table_route_storage() + .get_with_raw_bytes(1) + .await + .unwrap() + .unwrap(); + let mut routes = current_route.region_routes().unwrap().clone(); + for (route, expr) in routes.iter_mut().zip([ + expr::col("zone").lt(value::Value::String("m".into())), + expr::col("zone").gt_eq(value::Value::String("m".into())), + ]) { + route.region.partition_expr = expr.as_json_str().unwrap(); + route.leader_peer = Some(new_server.peer.clone()); + } + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + route_update: Mutex::new(Some((backend, current_route.update(routes).unwrap()))), + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()).with_retry( + build_retry_options(1), + table_info_cache.clone(), + invalidator.clone(), + ); + // The successful row would move under the new rule, but must not be replayed. + let (successful, _) = build_sparse_bulk_request(1, old_versions[&1], &[("a", "west")]); + let (failed, batch) = + build_sparse_bulk_request(2, old_versions[&2], &[("z", "east"), ("z", "west")]); + dispatcher.dispatch(vec![successful, failed]).await.unwrap(); + let counts = HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]); + assert_eq!(old_server.attempt_counts(), counts); + assert_eq!(new_server.attempt_counts(), counts); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![vec![CacheIdent::TableId(1)]] + ); + let refreshed_info = table_info_cache.get(1).await.unwrap().unwrap(); + assert_eq!(refreshed_info.meta.partition_key_indices, vec![1]); + let (_, versions) = partition_manager + .find_table_partition_rule(&refreshed_info) + .await + .unwrap(); + for request in new_server.service.requests.lock().unwrap().iter() { + let Some(Body::BulkInsert(bulk)) = &request.body else { + unreachable!() + }; + let region = RegionId::from_u64(bulk.region_id).region_number(); + assert_eq!( + bulk.partition_expr_version.map(|version| version.value), + versions[®ion] + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + let actual = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + assert_eq!(actual, batch.slice(region as usize - 1, 1)); + assert!(actual.column_by_name("zone").is_none()); + } +} + #[tokio::test] async fn test_resolve_peers_batches_logical_and_physical_regions() { let backend = test_util::prepare_mocked_backend().await; @@ -929,7 +1292,7 @@ async fn test_dispatch_preserves_write_payloads_and_envelopes() { }); let writes = dispatcher.prepare(requests.to_vec()).await.unwrap(); assert!(received.try_recv().is_err()); - let response = dispatcher.execute(writes, None).await.unwrap(); + let response = dispatcher.execute(writes, None, None).await.unwrap(); assert_eq!(response.affected_rows, 2); for request in requests { assert_eq!(received.try_recv().unwrap(), (Peer::new(1, ""), request)); @@ -1153,8 +1516,7 @@ async fn test_write_retry_rebuilds_logical_rows_in_mixed_table_envelope() { ..Default::default() }); let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - let dispatcher = WriteDispatcher::new(clients.clone(), partition_manager).with_retry( - clients, + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( build_retry_options(2), table_info_cache, invalidator.clone(), @@ -1162,10 +1524,17 @@ async fn test_write_retry_rebuilds_logical_rows_in_mixed_table_envelope() { let physical_id = RegionId::new(1, 1).as_u64(); let stale_logical_id = RegionId::new(42, 2).as_u64(); let rebuilt_logical_id = RegionId::new(42, 3).as_u64(); - server.service.rejections.lock().unwrap().insert( - stale_logical_id, - VecDeque::from([RetryHint::RetryAfterRouteRefresh]), + let mut status = Status::unavailable("partition changed"); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, + RetryHint::RetryAfterRouteRefresh.as_str().parse().unwrap(), ); + server + .service + .rejections + .lock() + .unwrap() + .insert(stale_logical_id, VecDeque::from([status])); let header = Some(RegionRequestHeader { dbname: "greptime.public".to_string(), @@ -1225,15 +1594,10 @@ async fn test_with_retry_can_disable_previously_enabled_retries() { let server = RegionTestServer::start().await; let dispatcher = server.create_dispatcher(1).await; let retry = dispatcher.retry.as_ref().unwrap(); - let clients = retry.clients.clone(); let table_info_cache = retry.table_info_cache.clone(); let cache_invalidator = retry.cache_invalidator.clone(); - let dispatcher = dispatcher.with_retry( - clients, - build_retry_options(0), - table_info_cache, - cache_invalidator, - ); + let dispatcher = + dispatcher.with_retry(build_retry_options(0), table_info_cache, cache_invalidator); assert!(dispatcher.retry.is_none()); server.fail_region(1, [RetryHint::Retryable]); assert!( diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index c3b712d3924..6ee7675313c 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -43,6 +43,9 @@ cargo nextest run -p servers Keep protocol translation here and permissions/database behavior in frontend. Async table batching returns storage-completion futures with its early acknowledgements. Keep completion reporting independent of the caller's lifetime. +Logical-table physical flushes retain partition columns in Arrow RPC payloads and +pass all region writes to the operator write dispatcher together; the dispatcher +owns peer resolution and retries. See `batcher/logical_table/region_write.rs`. `prom_remote_write/decode.rs` and `benches/prom_decode.rs` resolve `benches/write_request.pb.data` using `CARGO_MANIFEST_DIR` at **run time** diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs index be992692d13..d879486c2d3 100644 --- a/src/servers/benches/flush_batch_physical.rs +++ b/src/servers/benches/flush_batch_physical.rs @@ -15,8 +15,6 @@ use std::collections::HashMap; use std::sync::Arc; -use api::region::RegionResponse; -use api::v1::meta::Peer; use api::v1::region::RegionRequest; use arrow::array::{Float64Array, StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit}; @@ -29,11 +27,10 @@ use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use partition::error::Result as PartitionResult; use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask}; use servers::batcher::logical_table::{ - PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, + PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher, PhysicalTableMetadata, RecordBatchWithTsIdx, TableBatch, flush_batch_physical, }; -use servers::error::{self, Result}; -use store_api::storage::RegionId; +use servers::error; use table::test_util::table_info::test_table_info; use tokio::runtime::Runtime; @@ -95,24 +92,17 @@ impl PhysicalFlushPartitionProvider for BenchPartitionProvider { async fn find_table_partition_rule( &self, _table_info: &table::metadata::TableInfo, - ) -> PartitionResult { - Ok(Arc::new(SingleRegionPartitionRule)) - } - - async fn find_region_leader(&self, _region_id: RegionId) -> Result { - Ok(Peer { - id: 1, - addr: "bench-node".to_string(), - }) + ) -> PartitionResult<(PartitionRuleRef, HashMap>)> { + Ok((Arc::new(SingleRegionPartitionRule), HashMap::new())) } } -struct BenchNodeRequester; +struct BenchWriteDispatcher; #[async_trait] -impl PhysicalFlushNodeRequester for BenchNodeRequester { - async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> error::Result { - Ok(RegionResponse::new(0)) +impl PhysicalFlushWriteDispatcher for BenchWriteDispatcher { + async fn dispatch(&self, _requests: Vec) -> error::Result { + Ok(0) } } @@ -224,7 +214,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { table: make_physical_table_metadata(num_tags), }; let partition_provider = BenchPartitionProvider; - let node_requester = BenchNodeRequester; + let write_dispatcher = BenchWriteDispatcher; let mut group = c.benchmark_group("flush_batch_physical"); @@ -244,7 +234,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { "phy", &ctx, &partition_provider, - &node_requester, + &write_dispatcher, &catalog_provider, ) .await @@ -271,7 +261,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { "phy", &ctx, &partition_provider, - &node_requester, + &write_dispatcher, &catalog_provider, ) .await diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index 7e0cea44b6a..f34b0ae406e 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -39,6 +39,7 @@ use common_query::prelude::GREPTIME_PHYSICAL_TABLE; use common_time::timestamp::TimeUnit; use meter_core::data::MeterRecord; use meter_macros::write_meter; +use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::ResultExt; @@ -52,7 +53,7 @@ use crate::batcher::logical_table::pending_worker::{ PendingWorker, WorkerCommand, remove_worker_if_same_channel, start_worker, }; pub use crate::batcher::logical_table::region_write::{ - PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, + PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher, PhysicalTableMetadata, }; pub use crate::batcher::logical_table::tables::{ @@ -97,7 +98,7 @@ pub struct LogicalTablePendingRowsBatcher { flush_interval: Duration, flush_policy: TimingFlushPolicy, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: FlowNotifier, flush_limiter: FlushLimiter, @@ -110,10 +111,12 @@ pub struct LogicalTablePendingRowsBatcher { } impl LogicalTablePendingRowsBatcher { + /// Creates the batcher when metric-engine batching is enabled. #[allow(clippy::too_many_arguments)] pub fn try_new( partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, table_flownode_set_cache: TableFlownodeSetCacheRef, prom_store_with_metric_engine: bool, @@ -153,7 +156,7 @@ impl LogicalTablePendingRowsBatcher { flush_interval, flush_policy, partition_manager, - node_manager, + write_dispatcher, catalog_manager, flow_notification_tx, prom_store_with_metric_engine, @@ -404,7 +407,7 @@ impl LogicalTablePendingRowsBatcher { rx, self.shutdown.clone(), self.partition_manager.clone(), - self.node_manager.clone(), + self.write_dispatcher.clone(), self.catalog_manager.clone(), self.flow_notification_tx.clone(), worker_idle_timeout, diff --git a/src/servers/src/batcher/logical_table/batch.rs b/src/servers/src/batcher/logical_table/batch.rs index ac2db137de1..1ce90f39d01 100644 --- a/src/servers/src/batcher/logical_table/batch.rs +++ b/src/servers/src/batcher/logical_table/batch.rs @@ -18,45 +18,41 @@ use std::time::Instant; use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; -use common_meta::node_manager::NodeManagerRef; -use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_query::prelude; use common_telemetry::{debug, warn}; +use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::OptionExt; use crate::batcher::flow_notifier::FlowNotifier; -use crate::batcher::logical_table::PHYSICAL_TABLE_KEY; -use crate::batcher::logical_table::batch_convert::{ - TableBatch, concat_modified_batches, transform_logical_batches_to_physical, -}; -use crate::batcher::logical_table::flow_notifier::enqueue_flow_notifications; +use crate::batcher::logical_table; +use crate::batcher::logical_table::batch_convert::{self, TableBatch}; +use crate::batcher::logical_table::flow_notifier; use crate::batcher::logical_table::pending_worker::FlushWaiter; use crate::batcher::logical_table::region_write::{ - CatalogManagerPhysicalFlushAdapter, NodeManagerPhysicalFlushAdapter, - PartitionManagerPhysicalFlushAdapter, PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, - PhysicalFlushPartitionProvider, encode_region_write_requests, flush_region_writes_concurrently, - plan_region_batches, resolve_region_targets, + self, CatalogManagerPhysicalFlushAdapter, PartitionManagerPhysicalFlushAdapter, + PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher, + WriteDispatcherPhysicalFlushAdapter, }; use crate::error; use crate::error::Result; -use crate::metrics::{ - FLUSH_DROPPED_ROWS, FLUSH_ELAPSED, FLUSH_FAILURES, FLUSH_ROWS, FLUSH_TOTAL, - PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED, -}; +use crate::metrics; -pub(in crate::batcher::logical_table) struct Batch { - pub(in crate::batcher::logical_table) table_batches: Vec, - pub(in crate::batcher::logical_table) total_row_count: usize, - pub(in crate::batcher::logical_table) db_string: String, - pub(in crate::batcher::logical_table) ctx: QueryContextRef, - pub(in crate::batcher::logical_table) waiters: Vec, +/// Rows and completion waiters owned by one physical flush. +pub(crate) struct Batch { + pub(crate) table_batches: Vec, + pub(crate) total_row_count: usize, + pub(crate) db_string: String, + pub(crate) ctx: QueryContextRef, + pub(crate) waiters: Vec, } -pub(in crate::batcher::logical_table) async fn spawn_flush( +/// Spawns a physical flush with the configured concurrency limit. +pub(crate) async fn spawn_flush( flush: Batch, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: FlowNotifier, flush_limiter: FlushLimiter, @@ -68,7 +64,7 @@ pub(in crate::batcher::logical_table) async fn spawn_flush( flush_batch_with_managers( flush, partition_manager, - node_manager, + write_dispatcher, catalog_manager, flow_notification_tx, ) @@ -80,7 +76,7 @@ pub(in crate::batcher::logical_table) async fn spawn_flush( flush_batch_with_managers( flush, partition_manager, - node_manager, + write_dispatcher, catalog_manager, flow_notification_tx, ) @@ -89,32 +85,32 @@ pub(in crate::batcher::logical_table) async fn spawn_flush( } } -pub(in crate::batcher::logical_table) async fn flush_batch_with_managers( +/// Connects shared routing and dispatch dependencies to a physical flush. +pub(crate) async fn flush_batch_with_managers( flush: Batch, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: FlowNotifier, ) { let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager }; - let node_requester = NodeManagerPhysicalFlushAdapter { - node_manager: node_manager.clone(), - }; + let write_dispatcher = WriteDispatcherPhysicalFlushAdapter { write_dispatcher }; let catalog_provider = CatalogManagerPhysicalFlushAdapter { catalog_manager }; flush_batch( flush, &partition_provider, - &node_requester, + &write_dispatcher, &catalog_provider, flow_notification_tx, ) .await; } -pub(in crate::batcher::logical_table) async fn flush_batch( +/// Flushes a batch and reports completion to its waiters. +pub(crate) async fn flush_batch( flush: Batch, partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), - node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), + write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized), catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), flow_notification_tx: FlowNotifier, ) { @@ -130,21 +126,21 @@ pub(in crate::batcher::logical_table) async fn flush_batch( // Physical-table-level flush: transform all logical table batches // into physical format and write them together. let physical_table_name = ctx - .extension(PHYSICAL_TABLE_KEY) - .unwrap_or(GREPTIME_PHYSICAL_TABLE) + .extension(logical_table::PHYSICAL_TABLE_KEY) + .unwrap_or(prelude::GREPTIME_PHYSICAL_TABLE) .to_string(); let result = flush_batch_physical( &table_batches, &physical_table_name, &ctx, partition_manager, - node_manager, + write_dispatcher, catalog_manager, ) .await; let elapsed = start.elapsed().as_secs_f64(); - FLUSH_ELAPSED.observe(elapsed); + metrics::FLUSH_ELAPSED.observe(elapsed); debug!( "Pending rows batch flushed, total rows: {}, elapsed time: {}s", @@ -153,18 +149,18 @@ pub(in crate::batcher::logical_table) async fn flush_batch( match result { Ok(affected_rows) => { - FLUSH_TOTAL.inc(); - FLUSH_ROWS.observe(total_row_count as f64); + metrics::FLUSH_TOTAL.inc(); + metrics::FLUSH_ROWS.observe(total_row_count as f64); operator::metrics::DIST_INGEST_ROW_COUNT .with_label_values(&[db_string.as_str()]) .inc_by(affected_rows as u64); notify_waiters(waiters, Ok(())); - enqueue_flow_notifications(table_batches, &flow_notification_tx); + flow_notifier::enqueue_flow_notifications(table_batches, &flow_notification_tx); } Err(err) => { - FLUSH_FAILURES.inc(); - FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64); + metrics::FLUSH_FAILURES.inc(); + metrics::FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64); notify_waiters(waiters, Err(err)); } } @@ -185,12 +181,12 @@ pub async fn flush_batch_physical( physical_table_name: &str, ctx: &QueryContextRef, partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), - node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), + write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized), catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), ) -> Result { // 1. Resolve the physical table and get column ID mapping let physical_table = { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_resolve_table"]) .start_timer(); catalog_manager @@ -220,8 +216,8 @@ pub async fn flush_batch_physical( })?; // 2. Get the physical table's partition rule (one lookup instead of N) - let partition_rule = { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + let (partition_rule, partition_versions) = { + let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_fetch_partition_rule"]) .start_timer(); partition_manager @@ -233,24 +229,29 @@ pub async fn flush_batch_physical( partition_columns.iter().map(String::as_str).collect(); // 3. Transform each logical table batch into physical format - let modified_batches = - transform_logical_batches_to_physical(table_batches, &name_to_ids, &partition_columns_set)?; + let modified_batches = batch_convert::transform_logical_batches_to_physical( + table_batches, + &name_to_ids, + &partition_columns_set, + )?; // 4. Concatenate all modified batches (all share the same physical schema) - let combined_batch = concat_modified_batches(&modified_batches)?; + let combined_batch = batch_convert::concat_modified_batches(&modified_batches)?; // 5. Split by physical partition rule and send to regions let physical_table_id = physical_table_info.table_id(); - let planned_batches = plan_region_batches( + let planned_batches = region_write::plan_region_batches( combined_batch, physical_table_id, partition_rule.as_ref(), - partition_columns, )?; - let resolved_batches = resolve_region_targets(planned_batches, partition_manager).await?; - let region_writes = encode_region_write_requests(resolved_batches, ctx.skip_wal())?; - flush_region_writes_concurrently(node_manager, region_writes).await + let region_writes = region_write::encode_region_write_requests( + planned_batches, + ctx.skip_wal(), + &partition_versions, + )?; + write_dispatcher.dispatch(region_writes).await } pub(in crate::batcher::logical_table) fn notify_waiters( @@ -276,7 +277,6 @@ mod tests { use std::task::Poll; use std::time::Duration; - use api::region::RegionResponse; use api::v1::meta::Peer; use api::v1::region::RegionRequest; use arrow::record_batch::RecordBatch; @@ -288,7 +288,6 @@ mod tests { use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use partition::error::Result as PartitionResult; use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask}; - use store_api::storage::RegionId; use table::metadata::TableId; use table::test_util::table_info::test_table_info; use tokio::sync::{mpsc, oneshot}; @@ -300,7 +299,7 @@ mod tests { use crate::batcher::logical_table::batch_convert::TableBatch; use crate::batcher::logical_table::pending_worker::FlushWaiter; use crate::batcher::logical_table::region_write::{ - PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, + PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher, PhysicalTableMetadata, }; use crate::batcher::logical_table::test_util::{ @@ -344,9 +343,8 @@ mod tests { }, &MockFlushPartitionProvider { partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), }, - &MockFlushNodeRequester { + &MockFlushWriteDispatcher { writes: writes.clone(), fail: false, }, @@ -405,9 +403,8 @@ mod tests { }, &MockFlushPartitionProvider { partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), }, - &MockFlushNodeRequester { + &MockFlushWriteDispatcher { writes: writes.clone(), fail: true, }, @@ -451,8 +448,7 @@ mod tests { row_count: 1, }]; let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); + let dispatcher = MockFlushWriteDispatcher::default(); let ctx = session::context::QueryContext::arc(); flush_batch_physical( @@ -461,9 +457,8 @@ mod tests { &ctx, &MockFlushPartitionProvider { partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), }, - &node, + &dispatcher, &MockFlushCatalogProvider { table: Some(mock_physical_table_metadata(1024)), }, @@ -472,8 +467,7 @@ mod tests { .unwrap(); assert_eq!(1, partition_calls.load(Ordering::SeqCst)); - assert_eq!(1, leader_calls.load(Ordering::SeqCst)); - assert_eq!(1, node.writes.load(Ordering::SeqCst)); + assert_eq!(1, dispatcher.writes.load(Ordering::SeqCst)); } #[tokio::test] @@ -492,9 +486,8 @@ mod tests { &ctx, &MockFlushPartitionProvider { partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), }, - &AffectedRowsFlushNodeRequester { affected_rows: 7 }, + &AffectedRowsFlushWriteDispatcher { affected_rows: 7 }, &MockFlushCatalogProvider { table: Some(mock_physical_table_metadata(1024)), }, @@ -506,7 +499,7 @@ mod tests { } #[tokio::test] - async fn test_flush_batch_physical_stops_before_partition_and_node_when_table_missing() { + async fn test_flush_batch_physical_stops_before_partition_and_dispatch_when_table_missing() { let table_batches = vec![TableBatch { table_name: "t1".to_string(), table_id: 11, @@ -514,8 +507,7 @@ mod tests { row_count: 1, }]; let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); + let dispatcher = MockFlushWriteDispatcher::default(); let ctx = session::context::QueryContext::arc(); let err = flush_batch_physical( @@ -524,9 +516,8 @@ mod tests { &ctx, &MockFlushPartitionProvider { partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), }, - &node, + &dispatcher, &MockFlushCatalogProvider { table: None }, ) .await @@ -537,8 +528,7 @@ mod tests { .contains("Physical table 'missing_phy' not found") ); assert_eq!(0, partition_calls.load(Ordering::SeqCst)); - assert_eq!(0, leader_calls.load(Ordering::SeqCst)); - assert_eq!(0, node.writes.load(Ordering::SeqCst)); + assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst)); } #[tokio::test] @@ -558,8 +548,7 @@ mod tests { }, ]; let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); + let dispatcher = MockFlushWriteDispatcher::default(); let ctx = session::context::QueryContext::arc(); let err = flush_batch_physical( @@ -568,9 +557,8 @@ mod tests { &ctx, &MockFlushPartitionProvider { partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), }, - &node, + &dispatcher, &MockFlushCatalogProvider { table: Some(mock_physical_table_metadata(1024)), }, @@ -580,8 +568,7 @@ mod tests { assert!(err.to_string().contains("unknown_tag")); assert_eq!(1, partition_calls.load(Ordering::SeqCst)); - assert_eq!(0, leader_calls.load(Ordering::SeqCst)); - assert_eq!(0, node.writes.load(Ordering::SeqCst)); + assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst)); } fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata { @@ -671,7 +658,6 @@ mod tests { struct MockFlushPartitionProvider { partition_rule_calls: Arc, - region_leader_calls: Arc, } #[async_trait] @@ -679,40 +665,28 @@ mod tests { async fn find_table_partition_rule( &self, _table_info: &table::metadata::TableInfo, - ) -> PartitionResult { + ) -> PartitionResult<(PartitionRuleRef, HashMap>)> { self.partition_rule_calls.fetch_add(1, Ordering::SeqCst); - Ok(Arc::new(SingleRegionPartitionRule)) - } - - async fn find_region_leader(&self, _region_id: RegionId) -> error::Result { - self.region_leader_calls.fetch_add(1, Ordering::SeqCst); - Ok(Peer { - id: 1, - addr: "node-1".to_string(), - }) + Ok((Arc::new(SingleRegionPartitionRule), HashMap::new())) } } #[derive(Default)] - struct MockFlushNodeRequester { + struct MockFlushWriteDispatcher { writes: Arc, fail: bool, } #[async_trait] - impl PhysicalFlushNodeRequester for MockFlushNodeRequester { - async fn handle( - &self, - _peer: &Peer, - _request: RegionRequest, - ) -> error::Result { + impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher { + async fn dispatch(&self, _requests: Vec) -> error::Result { self.writes.fetch_add(1, Ordering::SeqCst); if self.fail { return Err(Error::Internal { err_msg: "physical write failed".to_string(), }); } - Ok(RegionResponse::new(0)) + Ok(0) } } @@ -726,18 +700,14 @@ mod tests { } #[derive(Default)] - struct AffectedRowsFlushNodeRequester { + struct AffectedRowsFlushWriteDispatcher { affected_rows: usize, } #[async_trait] - impl PhysicalFlushNodeRequester for AffectedRowsFlushNodeRequester { - async fn handle( - &self, - _peer: &Peer, - _request: RegionRequest, - ) -> error::Result { - Ok(RegionResponse::new(self.affected_rows)) + impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher { + async fn dispatch(&self, _requests: Vec) -> error::Result { + Ok(self.affected_rows) } } } diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs index 37080dbe258..65204ec038c 100644 --- a/src/servers/src/batcher/logical_table/batch_convert.rs +++ b/src/servers/src/batcher/logical_table/batch_convert.rs @@ -29,8 +29,6 @@ use crate::error; use crate::error::Result; use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED; -pub(in crate::batcher::logical_table) const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3; - /// An aligned logical record batch and its timestamp column index. #[derive(Debug, Clone)] pub struct RecordBatchWithTsIdx { @@ -188,23 +186,6 @@ pub(in crate::batcher::logical_table) fn columns_taxonomy( Ok((tag_columns, essential_column_indices)) } -pub(in crate::batcher::logical_table) fn strip_partition_columns_from_batch( - batch: RecordBatch, -) -> Result { - ensure!( - batch.num_columns() >= PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT, - error::InternalSnafu { - err_msg: format!( - "Expected at least {} columns in physical batch, got {}", - PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT, - batch.num_columns() - ), - } - ); - let essential_indices: Vec = (0..PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT).collect(); - batch.project(&essential_indices).context(error::ArrowSnafu) -} - /// Transforms logical table batches into physical format (sparse primary key encoding). /// /// It identifies tag columns and essential columns (timestamp, value) for each logical batch @@ -290,79 +271,18 @@ mod tests { use std::collections::{HashMap, HashSet}; use std::sync::Arc; - use arrow::array::{Array, BinaryArray, StringArray, TimestampMillisecondArray}; + use arrow::array::{Array, StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use metric_engine::batch_modifier::modify_batch_sparse; use smallvec::SmallVec; use crate::batcher::logical_table::batch_convert::{ - TableBatch, columns_taxonomy, strip_partition_columns_from_batch, - transform_logical_batches_to_physical, + TableBatch, columns_taxonomy, transform_logical_batches_to_physical, }; use crate::batcher::logical_table::test_util::mock_aligned_tag_batch; use crate::error::Error; - #[test] - fn test_strip_partition_columns_from_batch_removes_partition_tags() { - let batch = RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - Field::new("__primary_key", ArrowDataType::Binary, false), - Field::new( - "greptime_timestamp", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new("host", ArrowDataType::Utf8, true), - ])), - vec![ - Arc::new(BinaryArray::from(vec![b"k1".as_slice()])), - Arc::new(TimestampMillisecondArray::from(vec![1000_i64])), - Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])), - Arc::new(StringArray::from(vec!["node-1"])), - ], - ) - .unwrap(); - - let stripped = strip_partition_columns_from_batch(batch).unwrap(); - - assert_eq!(3, stripped.num_columns()); - assert_eq!("__primary_key", stripped.schema().field(0).name()); - assert_eq!("greptime_timestamp", stripped.schema().field(1).name()); - assert_eq!("greptime_value", stripped.schema().field(2).name()); - } - - #[test] - fn test_strip_partition_columns_from_batch_projects_essential_columns_without_lookup() { - let batch = RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - Field::new("__primary_key", ArrowDataType::Binary, false), - Field::new( - "greptime_timestamp", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new("host", ArrowDataType::Utf8, true), - ])), - vec![ - Arc::new(BinaryArray::from(vec![b"k1".as_slice()])), - Arc::new(TimestampMillisecondArray::from(vec![1000_i64])), - Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])), - Arc::new(StringArray::from(vec!["node-1"])), - ], - ) - .unwrap(); - - let stripped = strip_partition_columns_from_batch(batch).unwrap(); - - assert_eq!(3, stripped.num_columns()); - assert_eq!("__primary_key", stripped.schema().field(0).name()); - assert_eq!("greptime_timestamp", stripped.schema().field(1).name()); - assert_eq!("greptime_value", stripped.schema().field(2).name()); - } - #[test] fn test_collect_tag_columns_and_non_tag_indices_keeps_partition_tag_column() { let schema = Arc::new(ArrowSchema::new(vec![ diff --git a/src/servers/src/batcher/logical_table/pending_worker.rs b/src/servers/src/batcher/logical_table/pending_worker.rs index 4a83d8a87f4..ff363c7637c 100644 --- a/src/servers/src/batcher/logical_table/pending_worker.rs +++ b/src/servers/src/batcher/logical_table/pending_worker.rs @@ -22,8 +22,8 @@ use common_batcher::flush_policy::FlushTrigger; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::pending_worker::PendingWorker as PendingCore; use common_batcher::worker_registry::WorkerRegistry; -use common_meta::node_manager::NodeManagerRef; use common_telemetry::debug; +use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use table::metadata::TableId; @@ -31,7 +31,7 @@ use tokio::sync::{OwnedSemaphorePermit, broadcast, mpsc, oneshot}; use crate::batcher::flow_notifier::FlowNotifier; use crate::batcher::logical_table::BatchKey; -use crate::batcher::logical_table::batch::{Batch, flush_batch_with_managers, spawn_flush}; +use crate::batcher::logical_table::batch::{self, Batch}; use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; use crate::error::Error; use crate::metrics::{PENDING_BATCHES, PENDING_ROWS, PENDING_WORKERS}; @@ -53,7 +53,8 @@ pub(in crate::batcher::logical_table) struct PendingWorker { pub(in crate::batcher::logical_table) tx: mpsc::Sender, } -pub(in crate::batcher::logical_table) enum WorkerCommand { +/// Commands accepted by a logical-table batching worker. +pub(crate) enum WorkerCommand { Submit { table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>, total_rows: usize, @@ -92,15 +93,16 @@ impl PendingBatch { } } +/// Starts a worker that accumulates logical rows and dispatches physical flushes. #[allow(clippy::too_many_arguments)] -pub(in crate::batcher::logical_table) fn start_worker( +pub(crate) fn start_worker( key: BatchKey, worker_tx: mpsc::Sender, workers: Arc>, mut rx: mpsc::Receiver, shutdown: broadcast::Sender<()>, partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, + write_dispatcher: WriteDispatcherRef, catalog_manager: CatalogManagerRef, flow_notification_tx: FlowNotifier, worker_idle_timeout: Duration, @@ -140,10 +142,10 @@ pub(in crate::batcher::logical_table) fn start_worker( PENDING_ROWS.add(total_rows as i64); if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) { - spawn_flush( + batch::spawn_flush( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), flush_limiter.clone(), @@ -152,10 +154,10 @@ pub(in crate::batcher::logical_table) fn start_worker( } None => { if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { - flush_batch_with_managers( + batch::flush_batch_with_managers( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), ).await; @@ -189,10 +191,10 @@ pub(in crate::batcher::logical_table) fn start_worker( } _ = pending_flush.wait_flush() => { if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Deadline)) { - spawn_flush( + batch::spawn_flush( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), flush_limiter.clone(), @@ -201,10 +203,10 @@ pub(in crate::batcher::logical_table) fn start_worker( } _ = shutdown_rx.recv() => { if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { - flush_batch_with_managers( + batch::flush_batch_with_managers( flush, partition_manager.clone(), - node_manager.clone(), + write_dispatcher.clone(), catalog_manager.clone(), flow_notification_tx.clone(), ).await; @@ -286,6 +288,7 @@ mod tests { use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::node_manager::NodeManagerRef; use moka::future::CacheBuilder; + use operator::write_dispatcher::WriteDispatcher; use partition::cache::new_partition_info_cache; use partition::manager::PartitionRuleManager; use tokio::sync::{Semaphore, broadcast, mpsc, oneshot}; @@ -502,6 +505,10 @@ mod tests { let node_manager: NodeManagerRef = Arc::new(ConcurrentMockNodeManager { datanodes: Arc::new(HashMap::new()), }); + let write_dispatcher = Arc::new(WriteDispatcher::new( + node_manager, + partition_manager.clone(), + )); let catalog_manager = MemoryCatalogManager::with_default_setup(); let (flow_notification_tx, _flow_notification_rx) = FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap(); @@ -515,7 +522,7 @@ mod tests { worker_rx, shutdown.clone(), partition_manager, - node_manager, + write_dispatcher, catalog_manager, flow_notification_tx, worker_idle_timeout, diff --git a/src/servers/src/batcher/logical_table/region_write.rs b/src/servers/src/batcher/logical_table/region_write.rs index 71b558d1428..9511d21a998 100644 --- a/src/servers/src/batcher/logical_table/region_write.rs +++ b/src/servers/src/batcher/logical_table/region_write.rs @@ -14,29 +14,27 @@ use std::collections::HashMap; -use api::v1::ArrowIpc; -use api::v1::meta::Peer; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; -use arrow::compute::filter_record_batch; +use api::v1::{ArrowIpc, PartitionExprVersion}; +use arrow::compute; use arrow::record_batch::RecordBatch; use async_trait::async_trait; use catalog::CatalogManagerRef; use common_grpc::error::Error as GrpcError; -use common_grpc::flight::record_batch_to_ipc; -use common_meta::node_manager::NodeManagerRef; +use common_grpc::flight; use common_telemetry::tracing_context::TracingContext; +use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; use partition::partition::PartitionRuleRef; use snafu::ResultExt; use store_api::storage::RegionId; use table::metadata::{TableId, TableInfo, TableInfoRef}; -use crate::batcher::logical_table::batch_convert::strip_partition_columns_from_batch; use crate::error; use crate::error::{Error, Result}; -use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED; +use crate::metrics; #[derive(Clone)] pub struct PhysicalTableMetadata { @@ -56,23 +54,21 @@ pub trait PhysicalFlushCatalogProvider: Send + Sync { ) -> catalog::error::Result>; } +/// Loads the physical partition rule and its expression versions. #[async_trait] pub trait PhysicalFlushPartitionProvider: Send + Sync { + /// Loads a routing snapshot for the physical table. async fn find_table_partition_rule( &self, table_info: &TableInfo, - ) -> partition::error::Result; - - async fn find_region_leader(&self, region_id: RegionId) -> Result; + ) -> partition::error::Result<(PartitionRuleRef, HashMap>)>; } +/// Dispatches the region requests produced by a physical flush. #[async_trait] -pub trait PhysicalFlushNodeRequester: Send + Sync { - async fn handle( - &self, - peer: &Peer, - request: RegionRequest, - ) -> Result; +pub trait PhysicalFlushWriteDispatcher: Send + Sync { + /// Dispatches one physical flush together, sharing routing refreshes across its writes. + async fn dispatch(&self, requests: Vec) -> Result; } #[derive(Clone)] @@ -115,125 +111,61 @@ impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter { async fn find_table_partition_rule( &self, table_info: &TableInfo, - ) -> partition::error::Result { + ) -> partition::error::Result<(PartitionRuleRef, HashMap>)> { self.partition_manager .find_table_partition_rule(table_info) .await - .map(|(rule, _)| rule) - } - - async fn find_region_leader(&self, region_id: RegionId) -> Result { - let peer = self.partition_manager.find_region_leader(region_id).await?; - Ok(peer) } } +/// Sends a physical flush through the shared write dispatcher. #[derive(Clone)] -pub(in crate::batcher::logical_table) struct NodeManagerPhysicalFlushAdapter { - pub(in crate::batcher::logical_table) node_manager: NodeManagerRef, +pub(crate) struct WriteDispatcherPhysicalFlushAdapter { + pub(crate) write_dispatcher: WriteDispatcherRef, } #[async_trait] -impl PhysicalFlushNodeRequester for NodeManagerPhysicalFlushAdapter { - async fn handle( - &self, - peer: &Peer, - request: RegionRequest, - ) -> error::Result { - let datanode = self.node_manager.datanode(peer).await; - datanode - .handle(request) +impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter { + async fn dispatch(&self, requests: Vec) -> Result { + self.write_dispatcher + .dispatch_with_metrics( + requests, + metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_write_region"]), + ) .await .context(error::CommonMetaSnafu) + .map(|response| response.affected_rows) } } -pub(in crate::batcher::logical_table) struct FlushRegionWrite { - pub(in crate::batcher::logical_table) datanode: Peer, - pub(in crate::batcher::logical_table) request: RegionRequest, -} - -pub(in crate::batcher::logical_table) struct PlannedRegionBatch { - pub(in crate::batcher::logical_table) region_id: RegionId, - pub(in crate::batcher::logical_table) batch: RecordBatch, +/// Arrow rows selected for one physical region. +pub(crate) struct PlannedRegionBatch { + pub(crate) region_id: RegionId, + pub(crate) batch: RecordBatch, } #[cfg(test)] impl PlannedRegionBatch { - pub(in crate::batcher::logical_table) fn num_rows(&self) -> usize { + pub(crate) fn num_rows(&self) -> usize { self.batch.num_rows() } } -pub(in crate::batcher::logical_table) struct ResolvedRegionBatch { - pub(in crate::batcher::logical_table) planned: PlannedRegionBatch, - pub(in crate::batcher::logical_table) datanode: Peer, -} - -pub(in crate::batcher::logical_table) fn should_dispatch_concurrently( - region_write_count: usize, -) -> bool { - region_write_count > 1 -} - -pub(in crate::batcher::logical_table) async fn flush_region_writes_concurrently( - node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), - writes: Vec, -) -> Result { - let mut affected_rows = 0; - if !should_dispatch_concurrently(writes.len()) { - for write in writes { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_write_region"]) - .start_timer(); - affected_rows += node_manager - .handle(&write.datanode, write.request) - .await? - .affected_rows; - } - return Ok(affected_rows); - } - - let write_futures = writes.into_iter().map(|write| async move { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_write_region"]) - .start_timer(); - - let response = node_manager.handle(&write.datanode, write.request).await?; - Ok::<_, Error>(response.affected_rows) - }); - - // todo(hl): should be bounded. - let affected_rows = futures::future::try_join_all(write_futures) - .await? - .into_iter() - .sum(); - Ok(affected_rows) -} - pub(in crate::batcher::logical_table) fn split_combined_batch_by_region( combined_batch: &RecordBatch, partition_rule: &dyn partition::partition::PartitionRule, ) -> Result> { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_split_record_batch"]) .start_timer(); let map = partition_rule.split_record_batch(combined_batch)?; Ok(map) } -pub(in crate::batcher::logical_table) fn prepare_physical_region_routing_batch( - combined_batch: RecordBatch, - partition_columns: &[String], -) -> Result { - if partition_columns.is_empty() { - return Ok(combined_batch); - } - strip_partition_columns_from_batch(combined_batch) -} - -pub(in crate::batcher::logical_table) fn plan_region_batch( - stripped_batch: &RecordBatch, +/// Selects the rows assigned to one physical region. +pub(crate) fn plan_region_batch( + combined_batch: &RecordBatch, physical_table_id: TableId, region_number: u32, mask: &partition::partition::RegionMask, @@ -243,12 +175,12 @@ pub(in crate::batcher::logical_table) fn plan_region_batch( } let region_batch = if mask.select_all() { - stripped_batch.clone() + combined_batch.clone() } else { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_filter_record_batch"]) .start_timer(); - filter_record_batch(stripped_batch, mask.array()).context(error::ArrowSnafu)? + compute::filter_record_batch(combined_batch, mask.array()).context(error::ArrowSnafu)? }; let row_count = region_batch.num_rows(); @@ -262,19 +194,20 @@ pub(in crate::batcher::logical_table) fn plan_region_batch( })) } -pub(in crate::batcher::logical_table) fn plan_region_batches( +/// Partitions a physical flush while retaining columns needed for rerouting. +pub(crate) fn plan_region_batches( combined_batch: RecordBatch, physical_table_id: TableId, partition_rule: &dyn partition::partition::PartitionRule, - partition_columns: &[String], ) -> Result> { let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?; - let stripped_batch = prepare_physical_region_routing_batch(combined_batch, partition_columns)?; + // Retain partition columns in the RPC payload so routing can be rebuilt after repartition. + // Mito's sparse bulk conversion projects the stored fields and encoded primary key. let mut planned_batches = Vec::new(); for (region_number, mask) in region_masks { if let Some(planned_batch) = - plan_region_batch(&stripped_batch, physical_table_id, region_number, &mask)? + plan_region_batch(&combined_batch, physical_table_id, region_number, &mask)? { planned_batches.push(planned_batch); } @@ -283,39 +216,20 @@ pub(in crate::batcher::logical_table) fn plan_region_batches( Ok(planned_batches) } -pub(in crate::batcher::logical_table) async fn resolve_region_targets( +/// Encodes each region batch with the partition version used to select its rows. +pub(crate) fn encode_region_write_requests( planned_batches: Vec, - partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), -) -> Result> { - let mut resolved_batches = Vec::with_capacity(planned_batches.len()); - for planned in planned_batches { - let datanode = { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_physical_resolve_region_leader"]) - .start_timer(); - partition_manager - .find_region_leader(planned.region_id) - .await? - }; - - resolved_batches.push(ResolvedRegionBatch { planned, datanode }); - } - - Ok(resolved_batches) -} - -pub(in crate::batcher::logical_table) fn encode_region_write_requests( - resolved_batches: Vec, skip_wal: bool, -) -> Result> { - let mut region_writes = Vec::with_capacity(resolved_batches.len()); - for resolved in resolved_batches { - let region_id = resolved.planned.region_id; + partition_versions: &HashMap>, +) -> Result> { + let mut region_writes = Vec::with_capacity(planned_batches.len()); + for planned in planned_batches { + let region_id = planned.region_id; let (schema_bytes, data_header, payload) = { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_encode_ipc"]) .start_timer(); - record_batch_to_ipc(resolved.planned.batch).map_err(map_ipc_error)? + flight::record_batch_to_ipc(planned.batch).map_err(map_ipc_error)? }; let request = RegionRequest { @@ -326,7 +240,11 @@ pub(in crate::batcher::logical_table) fn encode_region_write_requests( body: Some(region_request::Body::BulkInsert(BulkInsertRequest { skip_wal, region_id: region_id.as_u64(), - partition_expr_version: None, + partition_expr_version: partition_versions + .get(®ion_id.region_number()) + .copied() + .flatten() + .map(|value| PartitionExprVersion { value }), // Set aligned_schema_version to None so that datanode will check the batch schema again to see if any // column is missing. aligned_schema_version: None, @@ -338,10 +256,7 @@ pub(in crate::batcher::logical_table) fn encode_region_write_requests( })), }; - region_writes.push(FlushRegionWrite { - datanode: resolved.datanode, - request, - }); + region_writes.push(request); } Ok(region_writes) @@ -359,89 +274,27 @@ pub(in crate::batcher::logical_table) fn map_ipc_error(error: GrpcError) -> Erro #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::collections::{HashMap, HashSet}; use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; - use api::region::RegionResponse; - use api::v1::meta::Peer; - use api::v1::region::{RegionRequest, region_request}; - use arrow::array::{BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray}; + use api::v1::PartitionExprVersion; + use api::v1::region::{bulk_insert_request, region_request}; + use arrow::array::{Array, BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; - use async_trait::async_trait; - use common_meta::error::Result as MetaResult; - use common_meta::node_manager::{Datanode, DatanodeRef}; - use common_query::request::QueryRequest; - use common_recordbatch::SendableRecordBatchStream; + use common_grpc::flight::FlightDecoder; use partition::partition::{PartitionRule, RegionMask}; use store_api::storage::RegionId; - use tokio::time::sleep; + use crate::batcher::logical_table::batch_convert::{self, RecordBatchWithTsIdx, TableBatch}; use crate::batcher::logical_table::region_write::{ - FlushRegionWrite, PlannedRegionBatch, ResolvedRegionBatch, encode_region_write_requests, - flush_region_writes_concurrently, plan_region_batches, should_dispatch_concurrently, + PlannedRegionBatch, encode_region_write_requests, plan_region_batches, }; - use crate::batcher::logical_table::test_util::ConcurrentMockNodeManager; - - #[tokio::test] - async fn test_flush_region_writes_concurrently_dispatches_multiple_datanodes() { - let inflight = Arc::new(AtomicUsize::new(0)); - let max_inflight = Arc::new(AtomicUsize::new(0)); - let datanode1: DatanodeRef = Arc::new(ConcurrentMockDatanode { - delay: Duration::from_millis(100), - inflight: inflight.clone(), - max_inflight: max_inflight.clone(), - }); - let datanode2: DatanodeRef = Arc::new(ConcurrentMockDatanode { - delay: Duration::from_millis(100), - inflight, - max_inflight: max_inflight.clone(), - }); - - let mut datanodes = HashMap::new(); - datanodes.insert(1, datanode1); - datanodes.insert(2, datanode2); - let node_manager = Arc::new(ConcurrentMockNodeManager { - datanodes: Arc::new(datanodes), - }); - - let writes = vec![ - FlushRegionWrite { - datanode: Peer { - id: 1, - addr: "node1".to_string(), - }, - request: RegionRequest::default(), - }, - FlushRegionWrite { - datanode: Peer { - id: 2, - addr: "node2".to_string(), - }, - request: RegionRequest::default(), - }, - ]; - - flush_region_writes_concurrently(node_manager.as_ref(), writes) - .await - .unwrap(); - assert!(max_inflight.load(Ordering::SeqCst) >= 2); - } #[test] - fn test_should_dispatch_concurrently_by_region_count() { - assert!(!should_dispatch_concurrently(0)); - assert!(!should_dispatch_concurrently(1)); - assert!(should_dispatch_concurrently(2)); - } - - #[test] - fn test_plan_region_batches_splits_and_strips_partition_columns() { - let combined_batch = RecordBatch::try_new( + fn test_plan_region_batches_preserves_partition_columns_through_encoding() { + let logical_batch = RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![ - Field::new("__primary_key", ArrowDataType::Binary, false), Field::new( "greptime_timestamp", ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), @@ -449,22 +302,38 @@ mod tests { ), Field::new("greptime_value", ArrowDataType::Float64, true), Field::new("host", ArrowDataType::Utf8, true), + Field::new("zone", ArrowDataType::Utf8, true), ])), vec![ - Arc::new(BinaryArray::from(vec![b"k1".as_slice(), b"k2".as_slice()])), Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000_i64])), Arc::new(arrow::array::Float64Array::from(vec![1.0_f64, 2.0_f64])), Arc::new(StringArray::from(vec!["node-1", "node-2"])), + Arc::new(StringArray::from(vec!["east", "west"])), ], ) .unwrap(); + let table_batches = [TableBatch { + table_name: "metric".to_string(), + table_id: 2048, + batches: vec![RecordBatchWithTsIdx::try_new(logical_batch, 0).unwrap()], + row_count: 2, + }]; + let modified = batch_convert::transform_logical_batches_to_physical( + &table_batches, + &HashMap::from([("host".to_string(), 1), ("zone".to_string(), 2)]), + &HashSet::from(["host"]), + ) + .unwrap(); + let combined_batch = batch_convert::concat_modified_batches(&modified).unwrap(); + assert!(combined_batch.column_by_name("host").is_some()); + assert!(combined_batch.column_by_name("zone").is_none()); + let expected = combined_batch.clone(); let mut planned_batches = plan_region_batches( combined_batch, 1024, &TwoRegionPartitionRule { partition_columns: vec!["host".to_string()], }, - &["host".to_string()], ) .unwrap(); planned_batches.sort_by_key(|planned| planned.region_id.region_number()); @@ -472,10 +341,45 @@ mod tests { assert_eq!(2, planned_batches.len()); assert_eq!(RegionId::new(1024, 1), planned_batches[0].region_id); assert_eq!(1, planned_batches[0].num_rows()); - assert_eq!(3, planned_batches[0].batch.num_columns()); + assert_eq!(4, planned_batches[0].batch.num_columns()); assert_eq!(RegionId::new(1024, 2), planned_batches[1].region_id); assert_eq!(1, planned_batches[1].num_rows()); - assert_eq!(3, planned_batches[1].batch.num_columns()); + assert_eq!(4, planned_batches[1].batch.num_columns()); + + let writes = encode_region_write_requests( + planned_batches, + false, + &HashMap::from([(1, Some(7)), (2, Some(8))]), + ) + .unwrap(); + for (index, write) in writes.into_iter().enumerate() { + let Some(region_request::Body::BulkInsert(request)) = write.body else { + panic!("expected bulk") + }; + assert_eq!( + request.partition_expr_version.unwrap().value, + 7 + index as u64 + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else { + panic!("expected IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + // The partition-only Arrow projection must retain the complete sparse key, + // including the zone tag that a later partition rule can require. + assert!(batch.column_by_name("zone").is_none()); + assert_eq!(batch, expected.slice(index, 1)); + let tags = batch + .column_by_name("host") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(tags.len(), 1); + assert_eq!(tags.value(0), if index == 0 { "node-1" } else { "node-2" }); + } } #[test] @@ -520,44 +424,6 @@ mod tests { } } - #[derive(Clone)] - struct ConcurrentMockDatanode { - delay: Duration, - inflight: Arc, - max_inflight: Arc, - } - - #[async_trait] - impl Datanode for ConcurrentMockDatanode { - async fn handle(&self, _request: RegionRequest) -> MetaResult { - let now = self.inflight.fetch_add(1, Ordering::SeqCst) + 1; - loop { - let max = self.max_inflight.load(Ordering::SeqCst); - if now <= max { - break; - } - if self - .max_inflight - .compare_exchange(max, now, Ordering::SeqCst, Ordering::SeqCst) - .is_ok() - { - break; - } - } - - sleep(self.delay).await; - self.inflight.fetch_sub(1, Ordering::SeqCst); - Ok(RegionResponse::new(0)) - } - - async fn handle_query( - &self, - _request: QueryRequest, - ) -> MetaResult { - unimplemented!() - } - } - fn check_encode_region_write_requests(skip_wal: bool) { let planned_batch = PlannedRegionBatch { region_id: RegionId::new(1024, 1), @@ -579,21 +445,22 @@ mod tests { ) .unwrap(), }; - let resolved_batch = ResolvedRegionBatch { - planned: planned_batch, - datanode: Peer { - id: 1, - addr: "node-1".to_string(), - }, - }; - let writes = encode_region_write_requests(vec![resolved_batch], skip_wal).unwrap(); + let writes = encode_region_write_requests( + vec![planned_batch], + skip_wal, + &HashMap::from([(1, Some(7))]), + ) + .unwrap(); assert_eq!(1, writes.len()); - assert_eq!(1, writes[0].datanode.id); - let Some(region_request::Body::BulkInsert(request)) = &writes[0].request.body else { + let Some(region_request::Body::BulkInsert(request)) = &writes[0].body else { panic!("expected bulk insert request"); }; assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id); assert_eq!(skip_wal, request.skip_wal); + assert_eq!( + Some(PartitionExprVersion { value: 7 }), + request.partition_expr_version + ); } } diff --git a/src/servers/src/batcher/logical_table/test_util.rs b/src/servers/src/batcher/logical_table/test_util.rs index 1cc81f9b615..491bd6ff2b3 100644 --- a/src/servers/src/batcher/logical_table/test_util.rs +++ b/src/servers/src/batcher/logical_table/test_util.rs @@ -15,10 +15,9 @@ use std::collections::HashMap; use std::sync::Arc; -use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse}; use api::v1::meta::Peer; -use api::v1::region::{InsertRequests, RegionRequest}; +use api::v1::region::InsertRequests; use arrow::array::{StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; @@ -27,12 +26,9 @@ use common_meta::error::Result as MetaResult; use common_meta::node_manager::{ DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, }; -use snafu::ResultExt; use tokio::sync::mpsc; use crate::batcher::logical_table::batch_convert::RecordBatchWithTsIdx; -use crate::batcher::logical_table::region_write::PhysicalFlushNodeRequester; -use crate::error; pub(in crate::batcher::logical_table) fn mock_tag_batch( tag_name: &str, @@ -149,14 +145,3 @@ impl FlownodeManager for FlowNotificationMockNodeManager { self.flownode.clone() } } - -#[async_trait] -impl PhysicalFlushNodeRequester for ConcurrentMockNodeManager { - async fn handle(&self, peer: &Peer, request: RegionRequest) -> error::Result { - let datanode = self.datanode(peer).await; - datanode - .handle(request) - .await - .context(error::CommonMetaSnafu) - } -} diff --git a/src/standalone/src/datanode_manager.rs b/src/standalone/src/datanode_manager.rs index 13a17ea2b43..a27a94b3c8d 100644 --- a/src/standalone/src/datanode_manager.rs +++ b/src/standalone/src/datanode_manager.rs @@ -77,7 +77,17 @@ impl RegionInvoker { #[async_trait] impl Datanode for RegionInvoker { - async fn handle(&self, request: RegionRequest) -> MetaResult { + async fn handle_request(&self, request: Arc) -> MetaResult { + // Local execution consumes the request; standalone does not retain it for retries. + let request = Arc::try_unwrap(request) + .map_err(|_| { + InvalidRegionRequestSnafu { + reason: "local region request has multiple owners", + } + .build() + }) + .map_err(BoxedError::new) + .context(meta_error::ExternalSnafu)?; let span = request .header .as_ref() diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 079cae4bca6..05d929cc7da 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -25,7 +25,7 @@ use flow::FlowConfig; use frontend::frontend::FrontendOptions; use frontend::service_config::{ InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, - PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, + PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, WriteRetryOptions, }; use mito2::config::MitoConfig; use pipeline::PipelineOptions; @@ -168,6 +168,11 @@ impl StandaloneOptions { otlp: cloned_opts.otlp, prom_store: cloned_opts.prom_store, meta_client: None, + // Local region calls consume requests directly and do not use remote write retries. + write_retry: WriteRetryOptions { + max_retries: 0, + ..Default::default() + }, logging: cloned_opts.logging, user_provider: cloned_opts.user_provider, query: cloned_opts.query, @@ -217,6 +222,18 @@ mod tests { use crate::options::*; + #[test] + fn test_frontend_write_retries_disabled() { + assert!(FrontendOptions::default().write_retry.max_retries > 0); + assert_eq!( + StandaloneOptions::default() + .frontend_options() + .write_retry + .max_retries, + 0 + ); + } + #[test] fn test_logical_batcher_config_forwarding() { let opts: StandaloneOptions = toml::from_str("[pending_rows_batcher]\nprotocols = ['influxdb']\n[pending_rows_batcher.logical_table]\nprotocols = ['otlp', 'prom']\npending_rows_flush_interval = '10ms'").unwrap(); diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 3e447bb0194..2349d580c5f 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -134,9 +134,9 @@ mod test { #[async_trait::async_trait] impl Datanode for GatedTraceDatanode { - async fn handle( + async fn handle_request( &self, - request: RegionRequest, + request: Arc, ) -> common_meta::error::Result { let gate = { let mut gate = self.gate.lock().unwrap(); diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 64f9f4c6af1..0729e6ea96b 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -892,6 +892,7 @@ pub fn build_test_prom_server( LogicalTablePendingRowsBatcher::try_new( frontend_ref.partition_manager().clone(), frontend_ref.node_manager().clone(), + frontend_ref.inserter().write_dispatcher().clone(), frontend_ref.catalog_manager().clone(), frontend_ref.table_flownode_set_cache().clone(), true, From f690bbb2040865c0d22b0b57d51612b60ae27cd0 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 13 Oct 2026 15:31:04 +0000 Subject: [PATCH 33/46] refactor!: remove datanode remote WAL index (#9563) * refactor: remove datanode remote WAL index Signed-off-by: WenyXu * fix(log-store): update new WAL replay tests after index removal Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 10 - config/config.md | 2 - config/datanode.example.toml | 8 - src/cmd/src/datanode/scanbench.rs | 2 +- src/common/meta/src/instruction.rs | 2 +- src/common/wal/src/config.rs | 24 + src/common/wal/src/config/kafka/datanode.rs | 6 - src/datanode/src/datanode.rs | 43 +- .../src/heartbeat/handler/upgrade_region.rs | 45 +- src/datanode/src/store.rs | 19 - src/log-store/Cargo.toml | 1 - src/log-store/src/error.rs | 44 +- src/log-store/src/kafka.rs | 2 - src/log-store/src/kafka/client_manager.rs | 16 - src/log-store/src/kafka/consumer.rs | 115 ++--- src/log-store/src/kafka/consumer/range.rs | 118 +++++ src/log-store/src/kafka/index.rs | 30 -- src/log-store/src/kafka/index/collector.rs | 414 --------------- src/log-store/src/kafka/index/encoder.rs | 181 ------- src/log-store/src/kafka/index/iterator.rs | 483 ------------------ src/log-store/src/kafka/log_store.rs | 67 +-- .../src/kafka/periodic_offset_fetcher.rs | 4 +- src/log-store/src/kafka/producer.rs | 59 +-- src/log-store/src/kafka/test_util.rs | 2 +- src/log-store/src/kafka/util.rs | 1 - src/log-store/src/kafka/util/range.rs | 146 ------ src/log-store/src/kafka/worker.rs | 58 +-- src/log-store/src/kafka/worker/dump_index.rs | 39 -- src/log-store/src/kafka/worker/flush.rs | 5 - src/log-store/src/kafka/worker/produce.rs | 11 +- src/log-store/src/noop/log_store.rs | 3 +- src/log-store/src/object_store_wal/store.rs | 41 +- src/log-store/src/raft_engine/log_store.rs | 11 +- src/log-store/src/test_util/log_store_util.rs | 13 +- src/metric-engine/src/engine/catchup.rs | 2 - src/mito2/src/engine/basic_test.rs | 2 +- src/mito2/src/engine/batch_catchup_test.rs | 2 - src/mito2/src/engine/object_store_wal_test.rs | 2 +- src/mito2/src/region/catchup.rs | 13 +- src/mito2/src/region/opener.rs | 2 +- src/mito2/src/wal.rs | 15 +- src/mito2/src/wal/raw_entry_reader.rs | 21 +- src/mito2/src/worker/handle_catchup.rs | 1 - src/mito2/src/worker/handle_write.rs | 3 +- src/store-api/src/logstore.rs | 17 - src/store-api/src/logstore/provider.rs | 3 +- src/store-api/src/region_request.rs | 2 - 47 files changed, 305 insertions(+), 1805 deletions(-) create mode 100644 src/log-store/src/kafka/consumer/range.rs delete mode 100644 src/log-store/src/kafka/index.rs delete mode 100644 src/log-store/src/kafka/index/collector.rs delete mode 100644 src/log-store/src/kafka/index/encoder.rs delete mode 100644 src/log-store/src/kafka/index/iterator.rs delete mode 100644 src/log-store/src/kafka/util/range.rs delete mode 100644 src/log-store/src/kafka/worker/dump_index.rs diff --git a/Cargo.lock b/Cargo.lock index e97dd203ee3..77e60390cb0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4624,15 +4624,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "delta-encoding" -version = "0.4.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cabe72c7de5725dcd4597e7d2bc8a4b7a4698879481607385fbc50a23203eeb9" -dependencies = [ - "num-traits", -] - [[package]] name = "der" version = "0.7.10" @@ -8119,7 +8110,6 @@ dependencies = [ "common-wal", "crc32fast", "dashmap", - "delta-encoding", "derive_builder 0.20.2", "futures", "futures-util", diff --git a/config/config.md b/config/config.md index 381abde24c4..6faae09c347 100644 --- a/config/config.md +++ b/config/config.md @@ -604,8 +604,6 @@ | `wal.timeout` | String | `5s` | The total request timeout for kafka client.
**It's only used when the provider is `kafka`**. | | `wal.max_batch_bytes` | String | `1MB` | The max size of a single producer batch.
Warning: Kafka has a default limit of 1MB per message in a topic.
Defaults to `8MB` when the provider is `experimental_object_store`.
**It's only used when the provider is `kafka` or `experimental_object_store`**. | | `wal.consumer_wait_timeout` | String | `100ms` | The consumer wait timeout.
**It's only used when the provider is `kafka`**. | -| `wal.create_index` | Bool | `false` | Whether to enable WAL index creation.
**It's only used when the provider is `kafka`**. | -| `wal.dump_index_interval` | String | `60s` | The interval for dumping WAL indexes.
**It's only used when the provider is `kafka`**. | | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `wal.storage_provider` | String | `""` | The name of the storage provider that holds the WAL objects, an empty name selects the default object store.
**It's only used when the provider is `experimental_object_store`**. | | `wal.prefix` | String | `wal` | The path prefix of the WAL objects inside the storage provider.
The objects are written under `/datanodes//epochs/`, which is derived from this prefix.
**It's only used when the provider is `experimental_object_store`**. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index a2586908b70..2d72f795189 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -211,14 +211,6 @@ max_batch_bytes = "1MB" ## **It's only used when the provider is `kafka`**. consumer_wait_timeout = "100ms" -## Whether to enable WAL index creation. -## **It's only used when the provider is `kafka`**. -create_index = false - -## The interval for dumping WAL indexes. -## **It's only used when the provider is `kafka`**. -dump_index_interval = "60s" - ## Ignore missing entries during read WAL. ## **It's only used when the provider is `kafka`**. ## diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 42eac50ec45..9fbde16fb7b 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -836,7 +836,7 @@ impl ScanbenchCommand { } ScanbenchLogStore::Kafka(kafka_config) => { let log_store = Arc::new( - KafkaLogStore::try_new(kafka_config, None) + KafkaLogStore::try_new(kafka_config) .await .map_err(BoxedError::new) .context(error::BuildCliSnafu)?, diff --git a/src/common/meta/src/instruction.rs b/src/common/meta/src/instruction.rs index 827f06e944d..0291f3a0946 100644 --- a/src/common/meta/src/instruction.rs +++ b/src/common/meta/src/instruction.rs @@ -381,7 +381,7 @@ pub struct UpgradeRegion { /// it's helpful to verify whether the leader region is ready. #[serde(with = "humantime_serde")] pub replay_timeout: Duration, - /// The hint for replaying memtable. + /// Legacy WAL index location hint, retained for older datanodes. #[serde(default)] pub location_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/src/common/wal/src/config.rs b/src/common/wal/src/config.rs index 0cbebfb20ab..b92c3d13d36 100644 --- a/src/common/wal/src/config.rs +++ b/src/common/wal/src/config.rs @@ -346,6 +346,30 @@ mod tests { assert!(toml::from_str::(toml_str).is_err()); } + #[test] + fn test_removed_datanode_wal_index_config() { + let expected = DatanodeWalConfig::Kafka(DatanodeKafkaConfig::default()); + for create_index in [false, true] { + let config = format!( + r#" + provider = "kafka" + create_index = {create_index} + dump_index_interval = "10s" + "# + ); + let parsed: DatanodeWalConfig = toml::from_str(&config).unwrap(); + assert_eq!(parsed, expected); + let serialized = toml::to_string(&parsed).unwrap(); + let table: toml::Table = toml::from_str(&serialized).unwrap(); + assert!(!table.contains_key("create_index")); + assert!(!table.contains_key("dump_index_interval")); + assert_eq!( + toml::from_str::(&serialized).unwrap(), + expected + ); + } + } + #[test] fn test_kafka_wal_config_debug_redacts_password() { let config = MetasrvWalConfig::Kafka(MetasrvKafkaConfig { diff --git a/src/common/wal/src/config/kafka/datanode.rs b/src/common/wal/src/config/kafka/datanode.rs index 4003b52e4bc..38836509f0a 100644 --- a/src/common/wal/src/config/kafka/datanode.rs +++ b/src/common/wal/src/config/kafka/datanode.rs @@ -40,10 +40,6 @@ pub struct DatanodeKafkaConfig { pub kafka_topic: KafkaTopicConfig, // Automatically create topics for WAL. pub auto_create_topics: bool, - // Create index for WAL. - pub create_index: bool, - #[serde(with = "humantime_serde")] - pub dump_index_interval: Duration, /// Internal interval for fetching latest Kafka topic offsets. #[serde(with = "humantime_serde")] pub topic_latest_offset_fetch_interval: Duration, @@ -60,8 +56,6 @@ impl Default for DatanodeKafkaConfig { consumer_wait_timeout: Duration::from_millis(100), kafka_topic: KafkaTopicConfig::default(), auto_create_topics: true, - create_index: false, - dump_index_interval: Duration::from_secs(60), topic_latest_offset_fetch_interval: DEFAULT_TOPIC_LATEST_OFFSET_FETCH_INTERVAL, overwrite_entry_start_id: false, } diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index cfbcf948b96..8edcd2ce4a1 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -46,7 +46,6 @@ use common_wal::config::raft_engine::RaftEngineConfig; use file_engine::engine::FileRegionEngine; use log_store::ObjectStoreLogStore; use log_store::kafka::log_store::KafkaLogStore; -use log_store::kafka::{GlobalIndexCollector, default_index_file}; use log_store::noop::log_store::NoopLogStore; use log_store::raft_engine::log_store::RaftEngineLogStore; use meta_client::MetaClientRef; @@ -88,7 +87,7 @@ use crate::heartbeat::HeartbeatTask; use crate::partition_expr_fetcher::MetaPartitionExprFetcher; use crate::region_query::DatanodeRegionQueryHandler; use crate::region_server::{DummyTableProviderFactory, RegionServer}; -use crate::store::{self, new_object_store_without_cache}; +use crate::store; use crate::utils::{RegionOpenRequests, build_region_open_requests}; /// Datanode service. @@ -683,29 +682,7 @@ impl DatanodeBuilder { builder.try_build().await.context(BuildMitoEngineSnafu)? } DatanodeWalConfig::Kafka(kafka_config) => { - if kafka_config.create_index && opts.node_id.is_none() { - warn!("The WAL index creation only available in distributed mode.") - } - let global_index_collector = if kafka_config.create_index - && let Some(node_id) = opts.node_id - { - let operator = new_object_store_without_cache( - &opts.storage.store, - &opts.storage.data_home, - ) - .await?; - let path = default_index_file(node_id); - Some(Self::build_global_index_collector( - kafka_config.dump_index_interval, - operator, - path, - )) - } else { - None - }; - - let log_store = - Self::build_kafka_log_store(kafka_config, global_index_collector).await?; + let log_store = Self::build_kafka_log_store(kafka_config).await?; self.topic_stats_reporter = Some(log_store.topic_stats_reporter()); let builder = MitoEngineBuilder::new( &opts.storage.data_home, @@ -825,25 +802,13 @@ impl DatanodeBuilder { } /// Builds [`KafkaLogStore`]. - async fn build_kafka_log_store( - config: &DatanodeKafkaConfig, - global_index_collector: Option, - ) -> Result> { - KafkaLogStore::try_new(config, global_index_collector) + async fn build_kafka_log_store(config: &DatanodeKafkaConfig) -> Result> { + KafkaLogStore::try_new(config) .await .map_err(Box::new) .context(OpenLogStoreSnafu) .map(Arc::new) } - - /// Builds [`GlobalIndexCollector`] - fn build_global_index_collector( - dump_index_interval: Duration, - operator: object_store::ObjectStore, - path: String, - ) -> GlobalIndexCollector { - GlobalIndexCollector::new(dump_index_interval, operator, path) - } } /// Keeps the first shutdown error and logs the later ones. diff --git a/src/datanode/src/heartbeat/handler/upgrade_region.rs b/src/datanode/src/heartbeat/handler/upgrade_region.rs index e2a5677fcad..08f96753ae9 100644 --- a/src/datanode/src/heartbeat/handler/upgrade_region.rs +++ b/src/datanode/src/heartbeat/handler/upgrade_region.rs @@ -130,7 +130,6 @@ impl UpgradeRegionsHandler { let UpgradeRegion { last_entry_id, metadata_last_entry_id, - location_id, replay_entry_id, metadata_replay_entry_id, replay_timeout, @@ -161,7 +160,6 @@ impl UpgradeRegionsHandler { set_writable: true, entry_id: last_entry_id, metadata_entry_id: metadata_last_entry_id, - location_id, checkpoint, }, )); @@ -320,6 +318,49 @@ mod tests { } } + #[tokio::test] + async fn test_legacy_wal_index_hint() { + for location_id in [None, Some(42)] { + let mock_region_server = mock_region_server(); + let region_id = RegionId::new(1024, 1); + let (mock_engine, _) = + MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |engine| { + engine.mock_role = Some(Some(RegionRole::Follower)); + engine.handle_request_mock_fn = Some(Box::new(|_, request| { + let store_api::region_request::RegionRequest::Catchup(request) = request + else { + panic!("expected catchup request"); + }; + assert!(request.set_writable); + assert_eq!(request.entry_id, Some(10)); + assert_eq!(request.metadata_entry_id, Some(5)); + Ok(0) + })); + }); + mock_region_server.register_test_region(region_id, mock_engine); + let ctx = + HandlerContext::new_for_test(mock_region_server, Arc::new(MemoryKvBackend::new())); + let reply = UpgradeRegionsHandler::new_test() + .handle( + &ctx, + vec![UpgradeRegion { + region_id, + last_entry_id: Some(10), + metadata_last_entry_id: Some(5), + location_id, + replay_timeout: Duration::from_secs(1), + ..Default::default() + }], + ) + .await + .unwrap(); + let reply = &reply.expect_upgrade_regions_reply()[0]; + assert!(reply.ready); + assert!(reply.exists); + assert!(reply.error.is_none()); + } + } + #[tokio::test] async fn test_region_not_ready() { let mock_region_server = mock_region_server(); diff --git a/src/datanode/src/store.rs b/src/datanode/src/store.rs index 78db3d61033..3369de5aedc 100644 --- a/src/datanode/src/store.rs +++ b/src/datanode/src/store.rs @@ -23,25 +23,6 @@ use snafu::prelude::*; use crate::config::ObjectStoreConfig; use crate::error::{self, Result}; -pub(crate) async fn new_object_store_without_cache( - store: &ObjectStoreConfig, - data_home: &str, -) -> Result { - let object_store = new_raw_object_store(store, data_home) - .await - .context(error::ObjectStoreSnafu)?; - // Enable retry layer and cache layer for non-fs object storages - let object_store = if store.is_object_storage() { - // Adds retry layer - with_retry_layers(object_store) - } else { - object_store - }; - - let object_store = with_instrument_layers(object_store, true); - Ok(object_store) -} - /// Cleans up old LRU read cache directories that were removed. fn clean_old_read_cache(store: &ObjectStoreConfig, data_home: &str) { if !store.is_object_storage() { diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 5870110be67..54e0b48c41d 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -30,7 +30,6 @@ common-time.workspace = true common-wal.workspace = true crc32fast.workspace = true dashmap.workspace = true -delta-encoding = "0.4" derive_builder.workspace = true futures.workspace = true futures-util.workspace = true diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index 64d918d5d78..836d1fbfe7d 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -259,39 +259,6 @@ pub enum Error { error: tokio::sync::oneshot::error::RecvError, }, - #[snafu(display("Failed to wait for result of DumpIndex"))] - WaitDumpIndex { - #[snafu(implicit)] - location: Location, - #[snafu(source)] - error: tokio::sync::oneshot::error::RecvError, - }, - - #[snafu(display("Failed to create writer"))] - CreateWriter { - #[snafu(implicit)] - location: Location, - #[snafu(source)] - error: object_store::Error, - }, - - #[snafu(display("Failed to write index"))] - WriteIndex { - #[snafu(implicit)] - location: Location, - #[snafu(source)] - error: object_store::Error, - }, - - #[snafu(display("Failed to read index, path: {path}"))] - ReadIndex { - #[snafu(implicit)] - location: Location, - #[snafu(source)] - error: object_store::Error, - path: String, - }, - #[snafu(display( "The length of meta if exceeded the limit: {}, actual: {}", limit, @@ -516,7 +483,6 @@ impl ErrorExt for Error { | DiscontinuousLogIndex { .. } | OrderedBatchProducerStopped { .. } | WaitProduceResultReceiver { .. } - | WaitDumpIndex { .. } | MetaLengthExceededLimit { .. } | ObjectStoreWalStopped { .. } | WalSegmentFetch { .. } => StatusCode::Internal, @@ -531,10 +497,7 @@ impl ErrorExt for Error { ObjectStoreWal { source, .. } => source.status_code(), // Object store related errors - CreateWriter { .. } - | WriteIndex { .. } - | ReadIndex { .. } - | WalObjectStore { .. } + WalObjectStore { .. } | StaleWalObject { .. } | UnconfirmedWalEpochStart { .. } | Io { .. } => StatusCode::StorageUnavailable, @@ -563,10 +526,7 @@ impl ErrorExt for Error { use Error::*; match self { - CreateWriter { error, .. } - | WriteIndex { error, .. } - | ReadIndex { error, .. } - | WalObjectStore { error, .. } => retry_hint_from_opendal_error(error), + WalObjectStore { error, .. } => retry_hint_from_opendal_error(error), ObjectStoreWal { source, .. } => source.retry_hint(), Io { error, .. } => retry_hint_from_io_error(error), FetchEntry { .. } diff --git a/src/log-store/src/kafka.rs b/src/log-store/src/kafka.rs index 2e123cc059c..810e06bb103 100644 --- a/src/log-store/src/kafka.rs +++ b/src/log-store/src/kafka.rs @@ -14,7 +14,6 @@ pub(crate) mod client_manager; pub(crate) mod consumer; -pub(crate) mod index; pub mod log_store; mod periodic_offset_fetcher; pub(crate) mod producer; @@ -24,7 +23,6 @@ pub(crate) mod util; pub(crate) mod worker; pub use client_manager::DEFAULT_PARTITION; -pub use index::{GlobalIndexCollector, default_index_file}; use serde::{Deserialize, Serialize}; use store_api::logstore::entry::Id as EntryId; diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index e594c9148eb..d1cfa2dba41 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -25,7 +25,6 @@ use store_api::logstore::provider::KafkaProvider; use tokio::sync::{Mutex, RwLock}; use crate::error::{BuildClientSnafu, BuildPartitionClientSnafu, Result, TlsConfigSnafu}; -use crate::kafka::index::{GlobalIndexCollector, NoopCollector}; use crate::kafka::log_store::TopicStat; use crate::kafka::producer::{OrderedBatchProducer, OrderedBatchProducerRef}; @@ -60,7 +59,6 @@ pub(crate) struct ClientManager { /// Used to initialize a new [Client]. mutex: Mutex<()>, instances: RwLock, Client>>, - global_index_collector: Option, flush_batch_size: usize, compression: Compression, @@ -73,7 +71,6 @@ impl ClientManager { /// Tries to create a ClientManager. pub(crate) async fn try_new( config: &DatanodeKafkaConfig, - global_index_collector: Option, topic_stats: Arc, TopicStat>>, ) -> Result { // Sets backoff config for the top-level kafka client and all clients constructed by it. @@ -98,7 +95,6 @@ impl ClientManager { instances: RwLock::new(HashMap::new()), flush_batch_size: config.max_batch_bytes.as_bytes() as usize, compression: Compression::Lz4, - global_index_collector, topic_stats, }) } @@ -151,30 +147,18 @@ impl ClientManager { .map(Arc::new)?; let (tx, rx) = OrderedBatchProducer::channel(); - let index_collector = if let Some(global_collector) = self.global_index_collector.as_ref() { - global_collector - .provider_level_index_collector(provider.clone(), tx.clone()) - .await - } else { - Box::new(NoopCollector) - }; let producer = Arc::new(OrderedBatchProducer::new( (tx, rx), provider.clone(), client.clone(), self.compression, self.flush_batch_size, - index_collector, self.topic_stats.clone(), )); Ok(Client { client, producer }) } - pub(crate) fn global_index_collector(&self) -> Option<&GlobalIndexCollector> { - self.global_index_collector.as_ref() - } - /// Lists all topics. pub(crate) async fn list_topics(&self) -> Vec> { self.instances.read().await.keys().cloned().collect() diff --git a/src/log-store/src/kafka/consumer.rs b/src/log-store/src/kafka/consumer.rs index d064a9bc9ac..d8c22109974 100644 --- a/src/log-store/src/kafka/consumer.rs +++ b/src/log-store/src/kafka/consumer.rs @@ -12,6 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod range; + +pub(crate) use range::RegionWalRange; + use std::collections::VecDeque; use std::ops::Range; use std::pin::Pin; @@ -26,7 +30,7 @@ use pin_project::pin_project; use rskafka::client::partition::PartitionClient; use rskafka::record::RecordAndOffset; -use crate::kafka::index::{NextBatchHint, RegionWalIndexIterator}; +use crate::kafka::consumer::range::NextBatchHint; pub struct FetchResult { /// The offsets of the fetched records. @@ -88,7 +92,7 @@ const MAX_BATCH_SIZE: usize = 52428800; const AVG_RECORD_SIZE: usize = 256 * 1024; /// The [`Consumer`] struct represents a Kafka consumer that fetches messages from -/// a Kafka cluster. Yielding records respecting the [`RegionWalIndexIterator`]. +/// a Kafka cluster. Yielding records respecting the [`RegionWalRange`]. #[pin_project] #[derive(Builder)] #[builder(pattern = "owned")] @@ -142,12 +146,12 @@ impl Consumer { pub(crate) struct RecordsBuffer { buffer: VecDeque, - index: Box, + index: RegionWalRange, } impl RecordsBuffer { /// Creates an empty [`RecordsBuffer`] - pub fn new(index: Box) -> Self { + pub fn new(index: RegionWalRange) -> Self { RecordsBuffer { buffer: VecDeque::new(), index, @@ -297,7 +301,6 @@ mod tests { use super::*; use crate::kafka::consumer::{Consumer, RecordsBuffer}; - use crate::kafka::index::{MultipleRegionWalIndexIterator, RegionWalRange, RegionWalVecIndex}; #[derive(Debug)] struct MockFetchClient { @@ -345,40 +348,6 @@ mod tests { } } - #[tokio::test] - async fn test_consumer_with_index() { - common_telemetry::init_default_ut_logging(); - let record = test_record(); - let record_size = record.approximate_size(); - let mock_client = MockFetchClient { - record: record.clone(), - }; - let index = RegionWalVecIndex::new([1, 3, 4, 8, 10, 12], record_size * 3); - let consumer = Consumer { - last_high_watermark: -1, - client: Arc::new(mock_client), - max_batch_size: usize::MAX, - max_wait_ms: 500, - avg_record_size: record_size, - terminated: false, - buffer: RecordsBuffer { - buffer: VecDeque::new(), - index: Box::new(index), - }, - fetch_fut: Fuse::terminated(), - total_fetched_bytes: 0, - }; - - let records = consumer.try_collect::>().await.unwrap(); - assert_eq!( - records - .into_iter() - .map(|(x, _)| x.offset) - .collect::>(), - vec![1, 3, 4, 8, 10, 12] - ) - } - #[tokio::test] async fn test_consumer_without_index() { common_telemetry::init_default_ut_logging(); @@ -396,7 +365,7 @@ mod tests { terminated: false, buffer: RecordsBuffer { buffer: VecDeque::new(), - index: Box::new(index), + index, }, fetch_fut: Fuse::terminated(), total_fetched_bytes: 0, @@ -413,44 +382,34 @@ mod tests { } #[tokio::test] - async fn test_consumer_with_multiple_index() { - common_telemetry::init_default_ut_logging(); - let record = test_record(); - let mock_client = MockFetchClient { - record: record.clone(), - }; - - let iter0 = Box::new(RegionWalRange::new(0..0, 1024)) as _; - let iter1 = Box::new(RegionWalVecIndex::new( - [0, 1, 2, 7, 8, 11], - record.approximate_size() * 4, - )) as _; - let iter2 = Box::new(RegionWalRange::new(12..12, 1024)) as _; - let iter3 = Box::new(RegionWalRange::new(1024..1028, 1024)) as _; - let iter = MultipleRegionWalIndexIterator::new([iter0, iter1, iter2, iter3]); - - let consumer = Consumer { - last_high_watermark: -1, - client: Arc::new(mock_client), - max_batch_size: usize::MAX, - max_wait_ms: 500, - avg_record_size: record.approximate_size(), - terminated: false, - buffer: RecordsBuffer { - buffer: VecDeque::new(), - index: Box::new(iter), - }, - fetch_fut: Fuse::terminated(), - total_fetched_bytes: 0, - }; - - let records = consumer.try_collect::>().await.unwrap(); - assert_eq!( - records + async fn test_consumer_range_boundaries() { + for (start, end, max_batch_size) in [ + (0, 0, 1024), + (5, 5, 1024), + (6, 5, 1024), + (5, 6, 1024), + (5, 35, 128), + ] { + let record = test_record(); + let consumer = ConsumerBuilder::default() + .client(Arc::new(MockFetchClient { + record: record.clone(), + })) + .buffer(RecordsBuffer::new(RegionWalRange::new( + start..end, + max_batch_size, + ))) + .avg_record_size(record.approximate_size()) + .build() + .unwrap(); + let offsets = consumer + .try_collect::>() + .await + .unwrap() .into_iter() - .map(|(x, _)| x.offset) - .collect::>(), - [0, 1, 2, 7, 8, 11, 1024, 1025, 1026, 1027] - ) + .map(|(record, _)| record.offset as u64) + .collect::>(); + assert_eq!(offsets, (start..end).collect::>()); + } } } diff --git a/src/log-store/src/kafka/consumer/range.rs b/src/log-store/src/kafka/consumer/range.rs new file mode 100644 index 00000000000..d4593b6f20c --- /dev/null +++ b/src/log-store/src/kafka/consumer/range.rs @@ -0,0 +1,118 @@ +// Copyright 2023 Greptime Team +// +// 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 std::cmp::min; +use std::ops::Range; + +use store_api::logstore::EntryId; + +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct NextBatchHint { + pub(crate) bytes: usize, + pub(crate) len: usize, +} + +/// Represents a range [next_entry_id, end_entry_id) of WAL entries for a region. +#[derive(Debug)] +pub struct RegionWalRange { + current_entry_id: EntryId, + end_entry_id: EntryId, + max_batch_size: usize, +} + +impl RegionWalRange { + pub fn new(range: Range, max_batch_size: usize) -> Self { + Self { + current_entry_id: range.start, + end_entry_id: range.end, + max_batch_size, + } + } + + fn next_batch_size(&self) -> Option { + if self.current_entry_id < self.end_entry_id { + Some(self.end_entry_id.saturating_sub(self.current_entry_id)) + } else { + None + } + } + + pub(crate) fn next_batch_hint(&self, avg_size: usize) -> Option { + if let Some(size) = self.next_batch_size() { + let bytes = min(size as usize * avg_size, self.max_batch_size); + let len = bytes / avg_size; + + return Some(NextBatchHint { bytes, len }); + } + + None + } + + pub(crate) fn peek(&self) -> Option { + if self.current_entry_id < self.end_entry_id { + Some(self.current_entry_id) + } else { + None + } + } + + pub(crate) fn next(&mut self) -> Option { + if self.current_entry_id < self.end_entry_id { + let next = self.current_entry_id; + self.current_entry_id += 1; + Some(next) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_region_wal_range() { + let range = RegionWalRange::new(0..1024, 1024); + assert_eq!( + range.next_batch_hint(10), + Some(NextBatchHint { + bytes: 1024, + len: 102 + }) + ); + + let mut range = RegionWalRange::new(0..1, 1024); + + assert_eq!(range.next_batch_size(), Some(1)); + assert_eq!(range.peek(), Some(0)); + + // Advance 1 step + assert_eq!(range.next(), Some(0)); + assert_eq!(range.next_batch_size(), None); + + // Advance 1 step + assert_eq!(range.next(), None); + assert_eq!(range.next_batch_size(), None); + // No effect + assert_eq!(range.next(), None); + assert_eq!(range.next_batch_size(), None); + + let mut range = RegionWalRange::new(0..0, 1024); + assert_eq!(range.next_batch_size(), None); + // No effect + assert_eq!(range.next(), None); + assert_eq!(range.next_batch_size(), None); + } +} diff --git a/src/log-store/src/kafka/index.rs b/src/log-store/src/kafka/index.rs deleted file mode 100644 index 23a9378ccb2..00000000000 --- a/src/log-store/src/kafka/index.rs +++ /dev/null @@ -1,30 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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. - -mod collector; -mod encoder; -mod iterator; - -pub use collector::GlobalIndexCollector; -pub(crate) use collector::{IndexCollector, NoopCollector}; -pub(crate) use encoder::{IndexEncoder, JsonIndexEncoder}; -pub(crate) use iterator::{ - MIN_BATCH_WINDOW_SIZE, NextBatchHint, RegionWalIndexIterator, build_region_wal_index_iterator, -}; -#[cfg(test)] -pub(crate) use iterator::{MultipleRegionWalIndexIterator, RegionWalRange, RegionWalVecIndex}; - -pub fn default_index_file(location_id: u64) -> String { - format!("__wal/{location_id}/index.json") -} diff --git a/src/log-store/src/kafka/index/collector.rs b/src/log-store/src/kafka/index/collector.rs deleted file mode 100644 index 54b9b21553d..00000000000 --- a/src/log-store/src/kafka/index/collector.rs +++ /dev/null @@ -1,414 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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 std::collections::{BTreeSet, HashMap}; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; - -use common_telemetry::{error, info}; -use futures::future::try_join_all; -use object_store::ErrorKind; -use serde::{Deserialize, Serialize}; -use snafu::ResultExt; -use store_api::logstore::EntryId; -use store_api::logstore::provider::KafkaProvider; -use store_api::storage::RegionId; -use tokio::select; -use tokio::sync::Mutex as TokioMutex; -use tokio::sync::mpsc::Sender; - -use crate::error::{self, Result}; -use crate::kafka::index::encoder::{DatanodeWalIndexes, IndexEncoder}; -use crate::kafka::index::{JsonIndexEncoder, default_index_file}; -use crate::kafka::worker::{DumpIndexRequest, TruncateIndexRequest, WorkerRequest}; - -/// The [`IndexCollector`] trait defines the operations for managing and collecting index entries. -pub trait IndexCollector: Send + Sync { - /// Appends an [`EntryId`] for a specific region. - fn append(&mut self, region_id: RegionId, entry_id: EntryId); - - /// Truncates the index for a specific region up to a given [`EntryId`]. - /// - /// It removes all [`EntryId`]s smaller than `entry_id`. - fn truncate(&mut self, region_id: RegionId, entry_id: EntryId); - - /// Sets the latest [`EntryId`]. - fn set_latest_entry_id(&mut self, entry_id: EntryId); - - /// Dumps the index. - fn dump(&mut self, encoder: &dyn IndexEncoder); -} - -/// The [`GlobalIndexCollector`] struct is responsible for managing index entries -/// across multiple providers. -#[derive(Debug)] -pub struct GlobalIndexCollector { - providers: Arc, Sender>>>, - operator: object_store::ObjectStore, - _handle: CollectionTaskHandle, -} - -#[derive(Debug, Clone)] -pub struct CollectionTask { - providers: Arc, Sender>>>, - dump_index_interval: Duration, - operator: object_store::ObjectStore, - path: String, - running: Arc, -} - -impl CollectionTask { - async fn dump_index(&self) -> Result<()> { - let encoder = Arc::new(JsonIndexEncoder::default()); - let receivers = { - let providers = self.providers.lock().await; - let mut receivers = Vec::with_capacity(providers.len()); - for (provider, sender) in providers.iter() { - let (req, rx) = DumpIndexRequest::new(encoder.clone()); - receivers.push(rx); - if sender.send(WorkerRequest::DumpIndex(req)).await.is_err() { - error!( - "BackgroundProducerWorker is stopped, topic: {}", - provider.topic - ) - } - } - receivers - }; - try_join_all(receivers) - .await - .context(error::WaitDumpIndexSnafu)?; - let bytes = encoder.finish()?; - let mut writer = self - .operator - .writer(&self.path) - .await - .context(error::CreateWriterSnafu)?; - writer.write(bytes).await.context(error::WriteIndexSnafu)?; - writer.close().await.context(error::WriteIndexSnafu)?; - - Ok(()) - } - - /// The background task performs two main operations: - /// - Persists the WAL index to the specified `path` at every `dump_index_interval`. - /// - Updates the latest index ID for each WAL provider at every `checkpoint_interval`. - fn run(self) -> CollectionTaskHandle { - let mut dump_index_interval = tokio::time::interval(self.dump_index_interval); - let running = self.running.clone(); - let moved_self = self.clone(); - common_runtime::spawn_global(async move { - loop { - if !running.load(Ordering::Relaxed) { - info!("shutdown the index collection task"); - break; - } - select! { - _ = dump_index_interval.tick() => { - if let Err(err) = moved_self.dump_index().await { - error!(err; "Failed to persist the WAL index"); - } - }, - } - } - }); - CollectionTaskHandle { - running: self.running.clone(), - } - } -} - -impl Drop for CollectionTaskHandle { - fn drop(&mut self) { - self.running.store(false, Ordering::Relaxed); - } -} - -#[derive(Debug, Default)] -struct CollectionTaskHandle { - running: Arc, -} - -impl GlobalIndexCollector { - /// Constructs a [`GlobalIndexCollector`]. - /// - /// This method initializes a `GlobalIndexCollector` instance and starts a background task - /// for managing WAL (Write-Ahead Logging) indexes. - /// - /// The background task persists the WAL index to the specified `path` at every `dump_index_interval`. - pub fn new( - dump_index_interval: Duration, - operator: object_store::ObjectStore, - path: String, - ) -> Self { - let providers: Arc, Sender>>> = - Arc::new(Default::default()); - let task = CollectionTask { - providers: providers.clone(), - dump_index_interval, - operator: operator.clone(), - path, - running: Arc::new(AtomicBool::new(true)), - }; - let handle = task.run(); - Self { - providers, - operator, - _handle: handle, - } - } - - #[cfg(test)] - pub fn new_for_test(operator: object_store::ObjectStore) -> Self { - Self { - providers: Default::default(), - operator, - _handle: Default::default(), - } - } -} - -impl GlobalIndexCollector { - /// Retrieve [`EntryId`]s for a specified `region_id` in `datanode_id` - /// that are greater than or equal to a given `entry_id`. - pub(crate) async fn read_remote_region_index( - &self, - location_id: u64, - provider: &KafkaProvider, - region_id: RegionId, - entry_id: EntryId, - ) -> Result, EntryId)>> { - let path = default_index_file(location_id); - - let bytes = match self.operator.read(&path).await { - Ok(bytes) => bytes.to_vec(), - Err(err) => { - if err.kind() == ErrorKind::NotFound { - return Ok(None); - } else { - return Err(err).context(error::ReadIndexSnafu { path }); - } - } - }; - - match DatanodeWalIndexes::decode(&bytes)?.provider(provider) { - Some(indexes) => { - let last_index = indexes.last_index(); - let indexes = indexes - .region(region_id) - .unwrap_or_default() - .split_off(&entry_id); - - Ok(Some((indexes, last_index))) - } - None => Ok(None), - } - } - - /// Creates a new [`ProviderLevelIndexCollector`] for a specified provider. - pub(crate) async fn provider_level_index_collector( - &self, - provider: Arc, - sender: Sender, - ) -> Box { - self.providers.lock().await.insert(provider.clone(), sender); - Box::new(ProviderLevelIndexCollector { - indexes: Default::default(), - provider, - }) - } - - /// Truncates the index for a specific region up to a given [`EntryId`]. - /// - /// It removes all [`EntryId`]s smaller than `entry_id`. - pub(crate) async fn truncate( - &self, - provider: &Arc, - region_id: RegionId, - entry_id: EntryId, - ) -> Result<()> { - if let Some(sender) = self.providers.lock().await.get(provider).cloned() - && sender - .send(WorkerRequest::TruncateIndex(TruncateIndexRequest::new( - region_id, entry_id, - ))) - .await - .is_err() - { - return error::OrderedBatchProducerStoppedSnafu {}.fail(); - } - - Ok(()) - } -} - -/// The [`RegionIndexes`] struct maintains indexes for a collection of regions. -/// Each region is identified by a `RegionId` and maps to a set of [`EntryId`]s, -/// representing the entries within that region. It also keeps track of the -/// latest [`EntryId`] across all regions. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct RegionIndexes { - pub(crate) regions: HashMap>, - pub(crate) latest_entry_id: EntryId, -} - -impl RegionIndexes { - fn append(&mut self, region_id: RegionId, entry_id: EntryId) { - self.regions.entry(region_id).or_default().insert(entry_id); - self.latest_entry_id = self.latest_entry_id.max(entry_id); - } - - fn truncate(&mut self, region_id: RegionId, entry_id: EntryId) { - if let Some(entry_ids) = self.regions.get_mut(®ion_id) { - *entry_ids = entry_ids.split_off(&entry_id); - // The `RegionIndexes` can be empty, keeps to track the latest entry id. - self.latest_entry_id = self.latest_entry_id.max(entry_id); - } - } - - fn set_latest_entry_id(&mut self, entry_id: EntryId) { - self.latest_entry_id = entry_id; - } -} - -/// The [`ProviderLevelIndexCollector`] struct is responsible for managing index entries -/// specific to a particular provider. -#[derive(Debug, Clone)] -pub struct ProviderLevelIndexCollector { - indexes: RegionIndexes, - provider: Arc, -} - -impl IndexCollector for ProviderLevelIndexCollector { - fn append(&mut self, region_id: RegionId, entry_id: EntryId) { - self.indexes.append(region_id, entry_id) - } - - fn truncate(&mut self, region_id: RegionId, entry_id: EntryId) { - self.indexes.truncate(region_id, entry_id) - } - - fn set_latest_entry_id(&mut self, entry_id: EntryId) { - self.indexes.set_latest_entry_id(entry_id); - } - - fn dump(&mut self, encoder: &dyn IndexEncoder) { - encoder.encode(&self.provider, &self.indexes) - } -} - -/// The [`NoopCollector`] struct implements the [`IndexCollector`] trait with no-op methods. -/// -/// This collector effectively ignores all operations, making it suitable for cases -/// where index collection is not required or should be disabled. -pub struct NoopCollector; - -impl IndexCollector for NoopCollector { - fn append(&mut self, _region_id: RegionId, _entry_id: EntryId) {} - - fn truncate(&mut self, _region_id: RegionId, _entry_id: EntryId) {} - - fn set_latest_entry_id(&mut self, _entry_id: EntryId) {} - - fn dump(&mut self, _encoder: &dyn IndexEncoder) {} -} - -#[cfg(test)] -mod tests { - use std::collections::{BTreeSet, HashMap}; - - use store_api::logstore::provider::KafkaProvider; - use store_api::storage::RegionId; - - use crate::kafka::index::JsonIndexEncoder; - use crate::kafka::index::collector::RegionIndexes; - use crate::kafka::index::encoder::IndexEncoder; - use crate::kafka::{GlobalIndexCollector, default_index_file}; - - #[tokio::test] - async fn test_read_remote_region_index() { - let operator = - object_store::ObjectStore::new(object_store::services::Memory::default()).unwrap(); - - let path = default_index_file(0); - let encoder = JsonIndexEncoder::default(); - encoder.encode( - &KafkaProvider::new("my_topic_0".to_string()), - &RegionIndexes { - regions: HashMap::from([(RegionId::new(1, 1), BTreeSet::from([1, 5, 15]))]), - latest_entry_id: 20, - }, - ); - let bytes = encoder.finish().unwrap(); - let mut writer = operator.writer(&path).await.unwrap(); - writer.write(bytes).await.unwrap(); - writer.close().await.unwrap(); - - let collector = GlobalIndexCollector::new_for_test(operator.clone()); - // Index file doesn't exist - let result = collector - .read_remote_region_index( - 1, - &KafkaProvider::new("my_topic_0".to_string()), - RegionId::new(1, 1), - 1, - ) - .await - .unwrap(); - assert!(result.is_none()); - - // RegionId doesn't exist - let (indexes, last_index) = collector - .read_remote_region_index( - 0, - &KafkaProvider::new("my_topic_0".to_string()), - RegionId::new(1, 2), - 5, - ) - .await - .unwrap() - .unwrap(); - assert_eq!(indexes, BTreeSet::new()); - assert_eq!(last_index, 20); - - // RegionId(1, 1), Start EntryId: 5 - let (indexes, last_index) = collector - .read_remote_region_index( - 0, - &KafkaProvider::new("my_topic_0".to_string()), - RegionId::new(1, 1), - 5, - ) - .await - .unwrap() - .unwrap(); - assert_eq!(indexes, BTreeSet::from([5, 15])); - assert_eq!(last_index, 20); - - // RegionId(1, 1), Start EntryId: 20 - let (indexes, last_index) = collector - .read_remote_region_index( - 0, - &KafkaProvider::new("my_topic_0".to_string()), - RegionId::new(1, 1), - 20, - ) - .await - .unwrap() - .unwrap(); - assert_eq!(indexes, BTreeSet::new()); - assert_eq!(last_index, 20); - } -} diff --git a/src/log-store/src/kafka/index/encoder.rs b/src/log-store/src/kafka/index/encoder.rs deleted file mode 100644 index b8d9e09bfed..00000000000 --- a/src/log-store/src/kafka/index/encoder.rs +++ /dev/null @@ -1,181 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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 std::collections::{BTreeSet, HashMap}; -use std::sync::Mutex; - -use delta_encoding::{DeltaDecoderExt, DeltaEncoderExt}; -use serde::{Deserialize, Serialize}; -use snafu::ResultExt; -use store_api::logstore::provider::KafkaProvider; -use store_api::storage::RegionId; - -use crate::error::{self, Result}; -use crate::kafka::index::collector::RegionIndexes; - -/// Converts a [`RegionIndexes`] instance into a [`DeltaEncodedRegionIndexes`]. -/// -/// This conversion encodes the index values using delta encoding to reduce storage space. -impl From<&RegionIndexes> for DeltaEncodedRegionIndexes { - fn from(value: &RegionIndexes) -> Self { - let mut regions = HashMap::with_capacity(value.regions.len()); - for (region_id, indexes) in value.regions.iter() { - let indexes = indexes.iter().copied().deltas().collect(); - regions.insert(*region_id, indexes); - } - Self { - regions, - last_index: value.latest_entry_id, - } - } -} - -/// Represents the delta-encoded version of region indexes for efficient storage. -#[derive(Debug, Default, Serialize, Deserialize)] -pub struct DeltaEncodedRegionIndexes { - regions: HashMap>, - last_index: u64, -} - -impl DeltaEncodedRegionIndexes { - /// Retrieves the original (decoded) index values for a given region. - pub(crate) fn region(&self, region_id: RegionId) -> Option> { - self.regions - .get(®ion_id) - .map(|delta| delta.iter().copied().original().collect::>()) - } - - /// Retrieves the last index. - pub(crate) fn last_index(&self) -> u64 { - self.last_index - } -} - -pub trait IndexEncoder: Send + Sync { - fn encode(&self, provider: &KafkaProvider, region_index: &RegionIndexes); - - fn finish(&self) -> Result>; -} - -/// [`DatanodeWalIndexes`] structure holds the WAL indexes for a datanode. -#[derive(Debug, Default, Serialize, Deserialize)] -pub(crate) struct DatanodeWalIndexes(HashMap); - -impl DatanodeWalIndexes { - fn insert(&mut self, topic: String, region_index: &RegionIndexes) { - self.0.insert(topic, region_index.into()); - } - - fn encode(&mut self) -> Result> { - let value = serde_json::to_vec(&self.0).context(error::EncodeJsonSnafu); - self.0.clear(); - value - } - - pub(crate) fn decode(byte: &[u8]) -> Result { - serde_json::from_slice(byte).context(error::DecodeJsonSnafu) - } - - /// Retrieves the delta encoded region indexes for a given `provider`. - pub(crate) fn provider(&self, provider: &KafkaProvider) -> Option<&DeltaEncodedRegionIndexes> { - self.0.get(&provider.topic) - } -} - -/// [`JsonIndexEncoder`] encodes the [`RegionIndexes`]s into JSON format. -#[derive(Debug, Default)] -pub(crate) struct JsonIndexEncoder { - buf: Mutex, -} - -impl IndexEncoder for JsonIndexEncoder { - fn encode(&self, provider: &KafkaProvider, region_index: &RegionIndexes) { - self.buf - .lock() - .unwrap() - .insert(provider.topic.clone(), region_index); - } - - fn finish(&self) -> Result> { - let mut buf = self.buf.lock().unwrap(); - buf.encode() - } -} - -#[cfg(test)] -mod tests { - use std::collections::{BTreeSet, HashMap}; - - use store_api::logstore::provider::KafkaProvider; - use store_api::storage::RegionId; - - use super::{DatanodeWalIndexes, IndexEncoder, JsonIndexEncoder}; - use crate::kafka::index::collector::RegionIndexes; - - #[test] - fn test_json_index_encoder() { - let encoder = JsonIndexEncoder::default(); - let topic_1 = KafkaProvider::new("my_topic_1".to_string()); - let region_1_indexes = BTreeSet::from([1u64, 2, 4, 5, 20]); - let region_2_indexes = BTreeSet::from([4u64, 12, 43, 54, 75]); - encoder.encode( - &topic_1, - &RegionIndexes { - regions: HashMap::from([ - (RegionId::new(1, 1), region_1_indexes.clone()), - (RegionId::new(1, 2), region_2_indexes.clone()), - ]), - latest_entry_id: 1024, - }, - ); - let topic_2 = KafkaProvider::new("my_topic_2".to_string()); - encoder.encode( - &topic_2, - &RegionIndexes { - regions: HashMap::from([ - ( - RegionId::new(1, 1), - BTreeSet::from([1024u64, 1025, 1026, 1028, 2048]), - ), - (RegionId::new(1, 2), BTreeSet::from([1512])), - ]), - latest_entry_id: 2048, - }, - ); - - let bytes = encoder.finish().unwrap(); - let datanode_index = DatanodeWalIndexes::decode(&bytes).unwrap(); - assert_eq!( - datanode_index - .provider(&topic_1) - .unwrap() - .region(RegionId::new(1, 1)) - .unwrap(), - region_1_indexes, - ); - assert_eq!( - datanode_index - .provider(&topic_1) - .unwrap() - .region(RegionId::new(1, 2)) - .unwrap(), - region_2_indexes, - ); - assert!( - datanode_index - .provider(&KafkaProvider::new("my_topic_3".to_string())) - .is_none() - ); - } -} diff --git a/src/log-store/src/kafka/index/iterator.rs b/src/log-store/src/kafka/index/iterator.rs deleted file mode 100644 index ba8e7273d78..00000000000 --- a/src/log-store/src/kafka/index/iterator.rs +++ /dev/null @@ -1,483 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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 std::cmp::{max, min}; -use std::collections::{BTreeSet, VecDeque}; -use std::fmt::Debug; -use std::ops::Range; - -use store_api::logstore::EntryId; - -use crate::kafka::util::range::{ConvertIndexToRange, MergeRange}; - -#[derive(Debug, PartialEq, Eq)] -pub(crate) struct NextBatchHint { - pub(crate) bytes: usize, - pub(crate) len: usize, -} - -/// An iterator over WAL (Write-Ahead Log) entries index for a region. -pub trait RegionWalIndexIterator: Send + Sync + Debug { - /// Returns next batch hint. - fn next_batch_hint(&self, avg_size: usize) -> Option; - - // Peeks the next EntryId without advancing the iterator. - fn peek(&self) -> Option; - - // Advances the iterator and returns the next EntryId. - fn next(&mut self) -> Option; - - #[cfg(test)] - fn as_any(&self) -> &dyn std::any::Any; -} - -/// Represents a range [next_entry_id, end_entry_id) of WAL entries for a region. -#[derive(Debug)] -pub struct RegionWalRange { - current_entry_id: EntryId, - end_entry_id: EntryId, - max_batch_size: usize, -} - -impl RegionWalRange { - pub fn new(range: Range, max_batch_size: usize) -> Self { - Self { - current_entry_id: range.start, - end_entry_id: range.end, - max_batch_size, - } - } - - fn next_batch_size(&self) -> Option { - if self.current_entry_id < self.end_entry_id { - Some(self.end_entry_id.saturating_sub(self.current_entry_id)) - } else { - None - } - } -} - -impl RegionWalIndexIterator for RegionWalRange { - fn next_batch_hint(&self, avg_size: usize) -> Option { - if let Some(size) = self.next_batch_size() { - let bytes = min(size as usize * avg_size, self.max_batch_size); - let len = bytes / avg_size; - - return Some(NextBatchHint { bytes, len }); - } - - None - } - - fn peek(&self) -> Option { - if self.current_entry_id < self.end_entry_id { - Some(self.current_entry_id) - } else { - None - } - } - - fn next(&mut self) -> Option { - if self.current_entry_id < self.end_entry_id { - let next = self.current_entry_id; - self.current_entry_id += 1; - Some(next) - } else { - None - } - } - - #[cfg(test)] - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -pub const MIN_BATCH_WINDOW_SIZE: usize = 4 * 1024 * 1024; - -/// Represents an index of Write-Ahead Log entries for a region, -/// stored as a vector of [EntryId]s. -#[derive(Debug)] -pub struct RegionWalVecIndex { - index: VecDeque, - min_batch_window_size: usize, -} - -impl RegionWalVecIndex { - pub fn new>(index: I, min_batch_window_size: usize) -> Self { - Self { - index: index.into_iter().collect::>(), - min_batch_window_size, - } - } -} - -impl RegionWalIndexIterator for RegionWalVecIndex { - fn next_batch_hint(&self, avg_size: usize) -> Option { - let merger = MergeRange::new( - ConvertIndexToRange::new(self.index.iter().peekable(), avg_size), - self.min_batch_window_size, - ); - - merger.merge().map(|(range, size)| NextBatchHint { - bytes: range.end - range.start - 1, - len: size, - }) - } - - fn peek(&self) -> Option { - self.index.front().cloned() - } - - fn next(&mut self) -> Option { - self.index.pop_front() - } - - #[cfg(test)] - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -/// Represents an iterator over multiple region WAL indexes. -/// -/// Allowing iteration through multiple WAL indexes. -#[derive(Debug)] -pub struct MultipleRegionWalIndexIterator { - iterator: VecDeque>, -} - -impl MultipleRegionWalIndexIterator { - pub fn new>>(iterator: I) -> Self { - Self { - iterator: iterator.into_iter().collect::>(), - } - } -} - -impl RegionWalIndexIterator for MultipleRegionWalIndexIterator { - fn next_batch_hint(&self, avg_size: usize) -> Option { - for iter in &self.iterator { - if let Some(batch) = iter.next_batch_hint(avg_size) { - return Some(batch); - } - } - - None - } - - fn peek(&self) -> Option { - for iter in &self.iterator { - let peek = iter.peek(); - if peek.is_some() { - return peek; - } - } - - None - } - - fn next(&mut self) -> Option { - while !self.iterator.is_empty() { - let remove = self.iterator.front().and_then(|iter| iter.peek()).is_none(); - if remove { - self.iterator.pop_front(); - } else { - break; - } - } - - self.iterator.front_mut().and_then(|iter| iter.next()) - } - - #[cfg(test)] - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -/// Builds [`RegionWalIndexIterator`]. -/// -/// Returns None means there are no entries to replay. -pub fn build_region_wal_index_iterator( - start_entry_id: EntryId, - end_entry_id: EntryId, - region_indexes: Option<(BTreeSet, EntryId)>, - max_batch_bytes: usize, - min_window_size: usize, -) -> Option> { - if (start_entry_id..end_entry_id).is_empty() { - return None; - } - - match region_indexes { - Some((region_indexes, last_index)) => { - if region_indexes.is_empty() && last_index >= end_entry_id { - return None; - } - - let mut iterator: Vec> = Vec::with_capacity(2); - if !region_indexes.is_empty() { - let index = RegionWalVecIndex::new(region_indexes, min_window_size); - iterator.push(Box::new(index)); - } - let known_last_index = max(last_index, start_entry_id); - if known_last_index < end_entry_id { - let range = known_last_index..end_entry_id; - let index = RegionWalRange::new(range, max_batch_bytes); - iterator.push(Box::new(index)); - } - - Some(Box::new(MultipleRegionWalIndexIterator::new(iterator))) - } - None => { - let range = start_entry_id..end_entry_id; - - Some(Box::new(RegionWalRange::new(range, max_batch_bytes))) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_region_wal_range() { - let range = RegionWalRange::new(0..1024, 1024); - assert_eq!( - range.next_batch_hint(10), - Some(NextBatchHint { - bytes: 1024, - len: 102 - }) - ); - - let mut range = RegionWalRange::new(0..1, 1024); - - assert_eq!(range.next_batch_size(), Some(1)); - assert_eq!(range.peek(), Some(0)); - - // Advance 1 step - assert_eq!(range.next(), Some(0)); - assert_eq!(range.next_batch_size(), None); - - // Advance 1 step - assert_eq!(range.next(), None); - assert_eq!(range.next_batch_size(), None); - // No effect - assert_eq!(range.next(), None); - assert_eq!(range.next_batch_size(), None); - - let mut range = RegionWalRange::new(0..0, 1024); - assert_eq!(range.next_batch_size(), None); - // No effect - assert_eq!(range.next(), None); - assert_eq!(range.next_batch_size(), None); - } - - #[test] - fn test_region_wal_vec_index() { - let mut index = RegionWalVecIndex::new([0, 1, 2, 7, 8, 11], 30); - assert_eq!( - index.next_batch_hint(10), - Some(NextBatchHint { bytes: 30, len: 3 }) - ); - assert_eq!(index.peek(), Some(0)); - // Advance 1 step - assert_eq!(index.next(), Some(0)); - assert_eq!( - index.next_batch_hint(10), - Some(NextBatchHint { bytes: 20, len: 2 }) - ); - // Advance 1 step - assert_eq!(index.next(), Some(1)); - assert_eq!( - index.next_batch_hint(10), - Some(NextBatchHint { bytes: 10, len: 1 }) - ); - // Advance 1 step - assert_eq!(index.next(), Some(2)); - assert_eq!( - index.next_batch_hint(10), - Some(NextBatchHint { bytes: 20, len: 2 }) - ); - // Advance 1 step - assert_eq!(index.next(), Some(7)); - assert_eq!( - index.next_batch_hint(10), - Some(NextBatchHint { bytes: 40, len: 2 }) - ); - // Advance 1 step - assert_eq!(index.next(), Some(8)); - assert_eq!( - index.next_batch_hint(10), - Some(NextBatchHint { bytes: 10, len: 1 }) - ); - // Advance 1 step - assert_eq!(index.next(), Some(11)); - assert_eq!(index.next_batch_hint(10), None); - - // No effect - assert_eq!(index.next(), None); - assert_eq!(index.next_batch_hint(10), None); - - let mut index = RegionWalVecIndex::new([], 1024); - assert_eq!(index.next_batch_hint(10), None); - assert_eq!(index.peek(), None); - // No effect - assert_eq!(index.peek(), None); - assert_eq!(index.next(), None); - assert_eq!(index.next_batch_hint(10), None); - } - - #[test] - fn test_multiple_region_wal_iterator() { - let iter0 = Box::new(RegionWalRange::new(0..0, 1024)) as _; - let iter1 = Box::new(RegionWalVecIndex::new([0, 1, 2, 7, 8, 11], 40)) as _; - let iter2 = Box::new(RegionWalRange::new(1024..1024, 1024)) as _; - let mut iter = MultipleRegionWalIndexIterator::new([iter0, iter1, iter2]); - - // The next batch is 0, 1, 2 - assert_eq!( - iter.next_batch_hint(10), - Some(NextBatchHint { bytes: 30, len: 3 }) - ); - assert_eq!(iter.peek(), Some(0)); - // Advance 1 step - assert_eq!(iter.next(), Some(0)); - - // The next batch is 1, 2 - assert_eq!( - iter.next_batch_hint(10), - Some(NextBatchHint { bytes: 20, len: 2 }) - ); - assert_eq!(iter.peek(), Some(1)); - // Advance 1 step - assert_eq!(iter.next(), Some(1)); - - // The next batch is 2 - assert_eq!( - iter.next_batch_hint(10), - Some(NextBatchHint { bytes: 10, len: 1 }) - ); - assert_eq!(iter.peek(), Some(2)); - - // Advance 1 step - assert_eq!(iter.next(), Some(2)); - // The next batch is 7, 8, 11 - assert_eq!( - iter.next_batch_hint(10), - Some(NextBatchHint { bytes: 50, len: 3 }) - ); - assert_eq!(iter.peek(), Some(7)); - - // Advance 1 step - assert_eq!(iter.next(), Some(7)); - // The next batch is 8, 11 - assert_eq!( - iter.next_batch_hint(10), - Some(NextBatchHint { bytes: 40, len: 2 }) - ); - assert_eq!(iter.peek(), Some(8)); - - // Advance 1 step - assert_eq!(iter.next(), Some(8)); - // The next batch is 11 - assert_eq!( - iter.next_batch_hint(10), - Some(NextBatchHint { bytes: 10, len: 1 }) - ); - assert_eq!(iter.peek(), Some(11)); - // Advance 1 step - assert_eq!(iter.next(), Some(11)); - - assert_eq!(iter.next_batch_hint(10), None,); - assert_eq!(iter.peek(), None); - assert!(!iter.iterator.is_empty()); - assert_eq!(iter.next(), None); - assert!(iter.iterator.is_empty()); - - // No effect - assert_eq!(iter.next(), None); - assert_eq!(iter.next_batch_hint(10), None,); - assert_eq!(iter.peek(), None); - assert_eq!(iter.next(), None); - } - - #[test] - fn test_build_region_wal_index_iterator() { - let iterator = build_region_wal_index_iterator(1024, 1024, None, 5, 5); - assert!(iterator.is_none()); - - let iterator = build_region_wal_index_iterator(1024, 1023, None, 5, 5); - assert!(iterator.is_none()); - - let iterator = - build_region_wal_index_iterator(1024, 1024, Some((BTreeSet::new(), 1024)), 5, 5); - assert!(iterator.is_none()); - - let iterator = - build_region_wal_index_iterator(1, 1024, Some((BTreeSet::new(), 1024)), 5, 5); - assert!(iterator.is_none()); - - let iterator = - build_region_wal_index_iterator(1, 1024, Some((BTreeSet::new(), 1025)), 5, 5); - assert!(iterator.is_none()); - - let iterator = build_region_wal_index_iterator( - 1, - 1024, - Some((BTreeSet::from([512, 756]), 1024)), - 5, - 5, - ) - .unwrap(); - let iter = iterator - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(iter.iterator.len(), 1); - let vec_index = iter.iterator[0] - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(vec_index.index, VecDeque::from([512, 756])); - - let iterator = build_region_wal_index_iterator( - 1, - 1024, - Some((BTreeSet::from([512, 756]), 1023)), - 5, - 5, - ) - .unwrap(); - let iter = iterator - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(iter.iterator.len(), 2); - let vec_index = iter.iterator[0] - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(vec_index.index, VecDeque::from([512, 756])); - let wal_range = iter.iterator[1] - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(wal_range.current_entry_id, 1023); - assert_eq!(wal_range.end_entry_id, 1024); - } -} diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 591a3c88ea5..daa6dfbcf21 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -31,15 +31,12 @@ use store_api::logstore::entry::{ Entry, Id as EntryId, MultiplePartEntry, MultiplePartHeader, NaiveEntry, }; use store_api::logstore::provider::{KafkaProvider, Provider}; -use store_api::logstore::{AppendBatchResponse, LogStore, SendableEntryStream, WalIndex}; +use store_api::logstore::{AppendBatchResponse, LogStore, SendableEntryStream}; use store_api::storage::RegionId; use crate::error::{self, ConsumeRecordSnafu, Error, GetOffsetSnafu, InvalidProviderSnafu, Result}; use crate::kafka::client_manager::{ClientManager, ClientManagerRef}; -use crate::kafka::consumer::{ConsumerBuilder, RecordsBuffer}; -use crate::kafka::index::{ - GlobalIndexCollector, MIN_BATCH_WINDOW_SIZE, build_region_wal_index_iterator, -}; +use crate::kafka::consumer::{ConsumerBuilder, RecordsBuffer, RegionWalRange}; use crate::kafka::periodic_offset_fetcher::PeriodicOffsetFetcher; use crate::kafka::producer::OrderedBatchProducerRef; use crate::kafka::util::record::{ @@ -142,14 +139,9 @@ impl TopicStatsReporter for PeriodicTopicStatsReporter { impl KafkaLogStore { /// Tries to create a Kafka log store. - pub async fn try_new( - config: &DatanodeKafkaConfig, - global_index_collector: Option, - ) -> Result { + pub async fn try_new(config: &DatanodeKafkaConfig) -> Result { let topic_stats = Arc::new(DashMap::new()); - let client_manager = Arc::new( - ClientManager::try_new(config, global_index_collector, topic_stats.clone()).await?, - ); + let client_manager = Arc::new(ClientManager::try_new(config, topic_stats.clone()).await?); let fetcher = PeriodicOffsetFetcher::new( config.topic_latest_offset_fetch_interval, client_manager.clone(), @@ -291,8 +283,7 @@ impl LogStore for KafkaLogStore { for (region_id, (producer, records)) in region_grouped_records { // Safety: `KafkaLogStore::entry` will ensure that the // `Record`'s `approximate_size` must be less or equal to `max_batch_bytes`. - region_grouped_result_receivers - .push((region_id, producer.produce(region_id, records).await?)) + region_grouped_result_receivers.push((region_id, producer.produce(records).await?)) } let region_grouped_max_offset = @@ -318,7 +309,6 @@ impl LogStore for KafkaLogStore { &self, provider: &Provider, mut entry_id: EntryId, - index: Option, ) -> Result> { let provider = provider .as_kafka_provider() @@ -378,28 +368,13 @@ impl LogStore for KafkaLogStore { record_num: 0, }); - let region_indexes = if let (Some(index), Some(collector)) = - (index, self.client_manager.global_index_collector()) - { - collector - .read_remote_region_index(index.location_id, provider, index.region_id, entry_id) - .await? - } else { - None - }; - - let Some(iterator) = build_region_wal_index_iterator( - entry_id, - end_offset as u64, - region_indexes, - self.max_batch_bytes, - MIN_BATCH_WINDOW_SIZE, - ) else { - let range = entry_id..end_offset as u64; + let range = entry_id..end_offset as u64; + if range.is_empty() { warn!("No new entries in range {:?} of ns {}", range, provider); return Ok(futures_util::stream::empty().boxed()); - }; + } + let iterator = RegionWalRange::new(range, self.max_batch_bytes); debug!("Reading entries with {:?} of ns {}", iterator, provider); // Safety: Must be ok. @@ -495,19 +470,11 @@ impl LogStore for KafkaLogStore { /// that the obsolete entries are deleted immediately. async fn obsolete( &self, - provider: &Provider, - region_id: RegionId, - entry_id: EntryId, + _provider: &Provider, + _region_id: RegionId, + _entry_id: EntryId, ) -> Result<()> { - if let Some(collector) = self.client_manager.global_index_collector() { - let provider = provider - .as_kafka_provider() - .with_context(|| InvalidProviderSnafu { - expected: KafkaProvider::type_name(), - actual: provider.type_name(), - })?; - collector.truncate(provider, region_id, entry_id).await?; - } + // Kafka WAL pruning is coordinated by metasrv. Ok(()) } @@ -682,7 +649,7 @@ mod tests { max_batch_bytes: ReadableSize::kb(32), ..Default::default() }; - let logstore = KafkaLogStore::try_new(&config, None).await.unwrap(); + let logstore = KafkaLogStore::try_new(&config).await.unwrap(); let topic_name = uuid::Uuid::new_v4().to_string(); prepare_topic(&logstore, &topic_name).await; let provider = Provider::kafka_provider(topic_name); @@ -708,7 +675,7 @@ mod tests { // 5 region assert_eq!(response.last_entry_ids.len(), 5); let got_entries = logstore - .read(&provider, 0, None) + .read(&provider, 0) .await .unwrap() .try_collect::>() @@ -766,7 +733,7 @@ mod tests { max_batch_bytes: ReadableSize::kb(8), ..Default::default() }; - let logstore = KafkaLogStore::try_new(&config, None).await.unwrap(); + let logstore = KafkaLogStore::try_new(&config).await.unwrap(); let topic_name = uuid::Uuid::new_v4().to_string(); prepare_topic(&logstore, &topic_name).await; let provider = Provider::kafka_provider(topic_name); @@ -792,7 +759,7 @@ mod tests { // 5 region assert_eq!(response.last_entry_ids.len(), 5); let got_entries = logstore - .read(&provider, 0, None) + .read(&provider, 0) .await .unwrap() .try_collect::>() diff --git a/src/log-store/src/kafka/periodic_offset_fetcher.rs b/src/log-store/src/kafka/periodic_offset_fetcher.rs index 4063fb75107..b3f8a7f06df 100644 --- a/src/log-store/src/kafka/periodic_offset_fetcher.rs +++ b/src/log-store/src/kafka/periodic_offset_fetcher.rs @@ -82,7 +82,6 @@ mod tests { use common_wal::maybe_skip_kafka_integration_test; use common_wal::test_util::get_kafka_endpoints; use store_api::logstore::provider::KafkaProvider; - use store_api::storage::RegionId; use super::*; use crate::kafka::test_util::{prepare, record}; @@ -113,8 +112,7 @@ mod tests { assert_eq!(current_latest_offset, 0); let record = vec![record(), record()]; - let region = RegionId::new(1, 1); - producer.produce(region, record.clone()).await.unwrap(); + producer.produce(record.clone()).await.unwrap(); tokio::time::sleep(Duration::from_millis(150)).await; let current_latest_offset = topic_stats.get(&provider).unwrap().latest_offset; assert_eq!(current_latest_offset, record.len() as u64 - 1); diff --git a/src/log-store/src/kafka/producer.rs b/src/log-store/src/kafka/producer.rs index 121bc5ef9cc..d4094a2ba00 100644 --- a/src/log-store/src/kafka/producer.rs +++ b/src/log-store/src/kafka/producer.rs @@ -19,11 +19,9 @@ use dashmap::DashMap; use rskafka::client::partition::{Compression, OffsetAt, PartitionClient}; use rskafka::record::Record; use store_api::logstore::provider::KafkaProvider; -use store_api::storage::RegionId; use tokio::sync::mpsc::{self, Receiver, Sender}; use crate::error::{self, Result}; -use crate::kafka::index::IndexCollector; use crate::kafka::log_store::TopicStat; use crate::kafka::worker::{BackgroundProducerWorker, ProduceResultHandle, WorkerRequest}; use crate::metrics::{ @@ -57,7 +55,6 @@ impl OrderedBatchProducer { client: Arc, compression: Compression, max_batch_bytes: usize, - index_collector: Box, topic_stats: Arc, TopicStat>>, ) -> Self { let mut worker = BackgroundProducerWorker { @@ -67,7 +64,6 @@ impl OrderedBatchProducer { receiver: rx, request_batch_size: REQUEST_BATCH_SIZE, max_batch_bytes, - index_collector, topic_stats, }; tokio::spawn(async move { worker.run().await }); @@ -81,12 +77,8 @@ impl OrderedBatchProducer { /// /// ## Panic /// Panic if any [Record]'s `approximate_size` > `max_batch_bytes`. - pub(crate) async fn produce( - &self, - region_id: RegionId, - batch: Vec, - ) -> Result { - let (req, handle) = WorkerRequest::new_produce_request(region_id, batch); + pub(crate) async fn produce(&self, batch: Vec) -> Result { + let (req, handle) = WorkerRequest::new_produce_request(batch); if self.sender.send(req).await.is_err() { warn!("OrderedBatchProducer is already exited"); return error::OrderedBatchProducerStoppedSnafu {}.fail(); @@ -172,10 +164,8 @@ mod tests { use rskafka::client::partition::Compression; use rskafka::protocol::error::Error as ProtocolError; use rskafka::record::Record; - use store_api::storage::RegionId; use super::*; - use crate::kafka::index::NoopCollector; use crate::kafka::producer::OrderedBatchProducer; use crate::kafka::test_util::record; @@ -246,17 +236,12 @@ mod tests { client.clone(), Compression::NoCompression, ReadableSize((record.approximate_size() * 2) as u64).as_bytes() as usize, - Box::new(NoopCollector), Arc::new(DashMap::new()), ); - let region_id = RegionId::new(1, 1); // Produces 3 records let handle = producer - .produce( - region_id, - vec![record.clone(), record.clone(), record.clone()], - ) + .produce(vec![record.clone(), record.clone(), record.clone()]) .await .unwrap(); assert_eq!(handle.wait().await.unwrap(), 2); @@ -264,17 +249,14 @@ mod tests { // Produces 2 records let handle = producer - .produce(region_id, vec![record.clone(), record.clone()]) + .produce(vec![record.clone(), record.clone()]) .await .unwrap(); assert_eq!(handle.wait().await.unwrap(), 4); assert_eq!(client.batch_sizes.lock().unwrap().as_slice(), &[2, 1, 2]); // Produces 1 records - let handle = producer - .produce(region_id, vec![record.clone()]) - .await - .unwrap(); + let handle = producer.produce(vec![record.clone()]).await.unwrap(); assert_eq!(handle.wait().await.unwrap(), 5); assert_eq!(client.batch_sizes.lock().unwrap().as_slice(), &[2, 1, 2, 1]); } @@ -295,36 +277,25 @@ mod tests { client.clone(), Compression::NoCompression, ReadableSize((record.approximate_size() * 2) as u64).as_bytes() as usize, - Box::new(NoopCollector), Arc::new(DashMap::new()), ); - let region_id = RegionId::new(1, 1); let mut futures = FuturesUnordered::new(); futures.push( producer - .produce( - region_id, - vec![record.clone(), record.clone(), record.clone()], - ) + .produce(vec![record.clone(), record.clone(), record.clone()]) .await .unwrap() .wait(), ); futures.push( producer - .produce(region_id, vec![record.clone(), record.clone()]) - .await - .unwrap() - .wait(), - ); - futures.push( - producer - .produce(region_id, vec![record.clone()]) + .produce(vec![record.clone(), record.clone()]) .await .unwrap() .wait(), ); + futures.push(producer.produce(vec![record.clone()]).await.unwrap().wait()); futures.next().await.unwrap().unwrap_err(); futures.next().await.unwrap().unwrap_err(); @@ -348,27 +319,17 @@ mod tests { client.clone(), Compression::NoCompression, ReadableSize((record.approximate_size() * 2) as u64).as_bytes() as usize, - Box::new(NoopCollector), Arc::new(DashMap::new()), ); - let region_id = RegionId::new(1, 1); let a = producer - .produce( - region_id, - vec![record.clone(), record.clone(), record.clone()], - ) + .produce(vec![record.clone(), record.clone(), record.clone()]) .await .unwrap() .wait() .fuse(); - let b = producer - .produce(region_id, vec![record]) - .await - .unwrap() - .wait() - .fuse(); + let b = producer.produce(vec![record]).await.unwrap().wait().fuse(); let mut b = Box::pin(b); diff --git a/src/log-store/src/kafka/test_util.rs b/src/log-store/src/kafka/test_util.rs index cb4855e467b..99e824a82c8 100644 --- a/src/log-store/src/kafka/test_util.rs +++ b/src/log-store/src/kafka/test_util.rs @@ -70,7 +70,7 @@ pub(crate) async fn prepare( ..Default::default() }; let high_watermark = Arc::new(DashMap::new()); - let manager = ClientManager::try_new(&config, None, high_watermark) + let manager = ClientManager::try_new(&config, high_watermark) .await .unwrap(); diff --git a/src/log-store/src/kafka/util.rs b/src/log-store/src/kafka/util.rs index e871feca429..562d107b429 100644 --- a/src/log-store/src/kafka/util.rs +++ b/src/log-store/src/kafka/util.rs @@ -12,5 +12,4 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub(crate) mod range; pub(crate) mod record; diff --git a/src/log-store/src/kafka/util/range.rs b/src/log-store/src/kafka/util/range.rs deleted file mode 100644 index 282d0aba185..00000000000 --- a/src/log-store/src/kafka/util/range.rs +++ /dev/null @@ -1,146 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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 std::cmp::{max, min}; -use std::iter::Peekable; -use std::ops::Range; - -use store_api::logstore::EntryId; - -/// Convert a sequence of [`EntryId`]s into size ranges. -pub(crate) struct ConvertIndexToRange<'a, I: Iterator> { - base: Option, - iter: Peekable, - avg_size: usize, -} - -impl<'a, I: Iterator> ConvertIndexToRange<'a, I> { - pub fn new(mut iter: Peekable, avg_size: usize) -> Self { - let base = iter.peek().cloned().cloned(); - - Self { - base, - iter, - avg_size, - } - } -} - -impl<'a, I: Iterator> Iterator for ConvertIndexToRange<'a, I> { - type Item = Range; - - fn next(&mut self) -> Option { - let (base, val) = (&self.base?, self.iter.next()?); - let start = (*val - *base) as usize * self.avg_size; - let end = start + self.avg_size + 1; - Some(start..end) - } -} - -/// Merge all ranges smaller than the `window_size`. -/// -/// e.g., -/// -/// Case 1 -/// - input: range: [(0..3), (5..6), (5..8)], window_size: 6 -/// - output: range: (0..6) -/// -/// Case 2 -/// - input: range: [(0..3)], window_size: 6 -/// - output: range: (0..3) -pub(crate) struct MergeRange>> { - iter: I, - window_size: usize, -} - -impl>> MergeRange { - pub fn new(iter: I, window_size: usize) -> Self { - Self { iter, window_size } - } -} - -/// Merges ranges. -fn merge(this: &mut Range, other: &Range) { - this.start = min(this.start, other.start); - this.end = max(this.end, other.end); -} - -impl>> MergeRange { - /// Calculates the size of the next merged range. - pub(crate) fn merge(mut self) -> Option<(Range, usize)> { - let mut merged_range = self.iter.next(); - let this = merged_range.as_mut()?; - let mut merged = 1; - for next in self.iter { - let window = next.start - this.start; - if window > self.window_size { - break; - } else { - merge(this, &next); - merged += 1; - } - } - merged_range.map(|range| (range, merged)) - } -} - -#[cfg(test)] -#[allow(clippy::single_range_in_vec_init)] -mod tests { - use super::*; - - #[test] - fn test_convert_index_to_range() { - let avg_size = 1024; - let index = [1u64, 4, 10, 15]; - let mut converter = ConvertIndexToRange::new(index.iter().peekable(), avg_size); - - assert_eq!(converter.next(), Some(0..avg_size + 1)); - assert_eq!(converter.next(), Some(3 * avg_size..4 * avg_size + 1)); - assert_eq!(converter.next(), Some(9 * avg_size..10 * avg_size + 1)); - assert_eq!(converter.next(), Some(14 * avg_size..15 * avg_size + 1)); - assert_eq!(converter.next(), None); - - let index = []; - let mut converter = ConvertIndexToRange::new(index.iter().peekable(), avg_size); - assert_eq!(converter.next(), None); - } - - #[test] - fn test_merge_range() { - let size_range = [(10usize..13), (12..14), (16..18), (19..29)]; - let merger = MergeRange::new(size_range.into_iter(), 9); - assert_eq!(merger.merge(), Some((10..29, 4))); - - let size_range = [(10usize..13), (12..14), (16..18)]; - let merger = MergeRange::new(size_range.into_iter(), 5); - assert_eq!(merger.merge(), Some((10..14, 2))); - - let size_range = [(10usize..13), (15..17), (16..18)]; - let merger = MergeRange::new(size_range.into_iter(), 5); - assert_eq!(merger.merge(), Some((10..17, 2))); - - let size_range = [(10usize..13)]; - let merger = MergeRange::new(size_range.into_iter(), 4); - assert_eq!(merger.merge(), Some((10..13, 1))); - - let size_range = [(10usize..13)]; - let merger = MergeRange::new(size_range.into_iter(), 2); - assert_eq!(merger.merge(), Some((10..13, 1))); - - let size_range = []; - let merger = MergeRange::new(size_range.into_iter(), 2); - assert_eq!(merger.merge(), None); - } -} diff --git a/src/log-store/src/kafka/worker.rs b/src/log-store/src/kafka/worker.rs index 7e3d00672c8..d73ab076607 100644 --- a/src/log-store/src/kafka/worker.rs +++ b/src/log-store/src/kafka/worker.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub(crate) mod dump_index; pub(crate) mod fetch_latest_offset; pub(crate) mod flush; pub(crate) mod produce; @@ -25,76 +24,31 @@ use futures::future::try_join_all; use rskafka::client::partition::Compression; use rskafka::record::Record; use snafu::{OptionExt, ResultExt}; -use store_api::logstore::EntryId; use store_api::logstore::provider::KafkaProvider; -use store_api::storage::RegionId; use tokio::sync::mpsc::Receiver; use tokio::sync::oneshot::{self}; use crate::error::{self, NoMaxValueSnafu, Result}; -use crate::kafka::index::{IndexCollector, IndexEncoder}; use crate::kafka::log_store::TopicStat; use crate::kafka::producer::ProducerClient; pub(crate) enum WorkerRequest { Produce(ProduceRequest), - TruncateIndex(TruncateIndexRequest), - DumpIndex(DumpIndexRequest), FetchLatestOffset, } impl WorkerRequest { - pub(crate) fn new_produce_request( - region_id: RegionId, - batch: Vec, - ) -> (WorkerRequest, ProduceResultHandle) { + pub(crate) fn new_produce_request(batch: Vec) -> (WorkerRequest, ProduceResultHandle) { let (tx, rx) = oneshot::channel(); ( - WorkerRequest::Produce(ProduceRequest { - region_id, - batch, - sender: tx, - }), + WorkerRequest::Produce(ProduceRequest { batch, sender: tx }), ProduceResultHandle { receiver: rx }, ) } } -pub(crate) struct DumpIndexRequest { - encoder: Arc, - sender: oneshot::Sender<()>, -} - -impl DumpIndexRequest { - pub fn new(encoder: Arc) -> (DumpIndexRequest, oneshot::Receiver<()>) { - let (tx, rx) = oneshot::channel(); - ( - DumpIndexRequest { - encoder, - sender: tx, - }, - rx, - ) - } -} - -pub(crate) struct TruncateIndexRequest { - region_id: RegionId, - entry_id: EntryId, -} - -impl TruncateIndexRequest { - pub fn new(region_id: RegionId, entry_id: EntryId) -> Self { - Self { - region_id, - entry_id, - } - } -} - pub(crate) struct ProduceRequest { - region_id: RegionId, batch: Vec, sender: oneshot::Sender, } @@ -142,7 +96,6 @@ impl ProduceResultReceiver { pub(crate) struct PendingRequest { batch: Vec, - region_ids: Vec, size: usize, sender: oneshot::Sender>>, } @@ -159,8 +112,6 @@ pub(crate) struct BackgroundProducerWorker { pub(crate) request_batch_size: usize, /// Max bytes size for a single flush. pub(crate) max_batch_bytes: usize, - /// Collecting ids of WAL entries. - pub(crate) index_collector: Box, /// The stats of each topic. pub(crate) topic_stats: Arc, TopicStat>>, } @@ -195,11 +146,6 @@ impl BackgroundProducerWorker { for req in buffer.drain(..) { match req { WorkerRequest::Produce(req) => produce_requests.push(req), - WorkerRequest::TruncateIndex(TruncateIndexRequest { - region_id, - entry_id, - }) => self.index_collector.truncate(region_id, entry_id), - WorkerRequest::DumpIndex(req) => self.dump_index(req).await, WorkerRequest::FetchLatestOffset => { self.fetch_latest_offset().await; } diff --git a/src/log-store/src/kafka/worker/dump_index.rs b/src/log-store/src/kafka/worker/dump_index.rs deleted file mode 100644 index 0de02c5c9e7..00000000000 --- a/src/log-store/src/kafka/worker/dump_index.rs +++ /dev/null @@ -1,39 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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 common_telemetry::error; -use rskafka::client::partition::OffsetAt; -use snafu::ResultExt; - -use crate::error; -use crate::kafka::worker::{BackgroundProducerWorker, DumpIndexRequest}; - -impl BackgroundProducerWorker { - pub(crate) async fn dump_index(&mut self, req: DumpIndexRequest) { - match self - .client - .get_offset(OffsetAt::Latest) - .await - .context(error::GetOffsetSnafu { - topic: &self.provider.topic, - }) { - Ok(offset) => { - self.index_collector.set_latest_entry_id(offset as u64); - self.index_collector.dump(req.encoder.as_ref()); - let _ = req.sender.send(()); - } - Err(err) => error!(err; "Failed to do checkpoint"), - } - } -} diff --git a/src/log-store/src/kafka/worker/flush.rs b/src/log-store/src/kafka/worker/flush.rs index 7348c4ec6ff..4c8e6eab034 100644 --- a/src/log-store/src/kafka/worker/flush.rs +++ b/src/log-store/src/kafka/worker/flush.rs @@ -24,7 +24,6 @@ impl BackgroundProducerWorker { &mut self, PendingRequest { batch, - region_ids, sender, size: _size, }: PendingRequest, @@ -38,10 +37,6 @@ impl BackgroundProducerWorker { if let Ok(result) = &result { let total_record_size = result.encoded_request_size as u64; - for (idx, region_id) in result.offsets.iter().zip(region_ids) { - self.index_collector.append(region_id, *idx as u64); - } - let max_offset = result.offsets.iter().max().cloned().unwrap_or_default() as u64; self.topic_stats .entry(self.provider.clone()) diff --git a/src/log-store/src/kafka/worker/produce.rs b/src/log-store/src/kafka/worker/produce.rs index 465d19a57d8..d9e7ba25654 100644 --- a/src/log-store/src/kafka/worker/produce.rs +++ b/src/log-store/src/kafka/worker/produce.rs @@ -30,16 +30,10 @@ impl BackgroundProducerWorker { max_batch_bytes: usize, ) -> Vec { let mut records_buffer = vec![]; - let mut region_ids = vec![]; let mut batch_size = 0; let mut pending_requests = Vec::with_capacity(requests.len()); - for ProduceRequest { - batch, - sender, - region_id, - } in std::mem::take(requests) - { + for ProduceRequest { batch, sender } in std::mem::take(requests) { let mut receiver = ProduceResultReceiver::default(); for record in batch { assert!(record.approximate_size() <= max_batch_bytes); @@ -48,7 +42,6 @@ impl BackgroundProducerWorker { let (tx, rx) = oneshot::channel(); pending_requests.push(PendingRequest { batch: std::mem::take(&mut records_buffer), - region_ids: std::mem::take(&mut region_ids), size: batch_size, sender: tx, }); @@ -58,7 +51,6 @@ impl BackgroundProducerWorker { batch_size += record.approximate_size(); records_buffer.push(record); - region_ids.push(region_id); } // The remaining records. if batch_size > 0 { @@ -66,7 +58,6 @@ impl BackgroundProducerWorker { let (tx, rx) = oneshot::channel(); pending_requests.push(PendingRequest { batch: std::mem::take(&mut records_buffer), - region_ids: std::mem::take(&mut region_ids), size: batch_size, sender: tx, }); diff --git a/src/log-store/src/noop/log_store.rs b/src/log-store/src/noop/log_store.rs index 11b77e79108..8e9d22d8f1d 100644 --- a/src/log-store/src/noop/log_store.rs +++ b/src/log-store/src/noop/log_store.rs @@ -17,7 +17,7 @@ use std::collections::HashMap; use futures::stream; use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::Provider; -use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; +use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream}; use store_api::storage::RegionId; use crate::error::{Error, Result}; @@ -45,7 +45,6 @@ impl LogStore for NoopLogStore { &self, _provider: &Provider, _entry_id: EntryId, - _index: Option, ) -> Result> { Ok(Box::pin(stream::empty())) } diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 1d75739350a..530f76b9ff6 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -33,7 +33,7 @@ use object_store::ObjectStore; use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::{ObjectStoreProvider, Provider}; -use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; +use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream}; use store_api::storage::RegionId; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch}; use tokio::time::MissedTickBehavior; @@ -755,7 +755,6 @@ impl LogStore for ObjectStoreLogStore { &self, provider: &Provider, entry_id: EntryId, - _index: Option, ) -> Result> { self.check_terminal()?; let region_id = self.region_of(provider)?; @@ -3596,7 +3595,7 @@ mod tests { let errors = [ store.append_batch(vec![entry]).await.unwrap_err(), store.latest_entry_id(foreign).unwrap_err(), - store.read(foreign, 0, None).await.err().unwrap(), + store.read(foreign, 0).await.err().unwrap(), store.create_namespace(foreign).await.unwrap_err(), store.delete_namespace(foreign).await.unwrap_err(), store.obsolete(foreign, region_id, 1).await.unwrap_err(), @@ -3705,7 +3704,7 @@ mod tests { let errors = [ store.append_batch(Vec::new()).await.unwrap_err(), store.latest_entry_id(&provider).unwrap_err(), - store.read(&provider, 0, None).await.err().unwrap(), + store.read(&provider, 0).await.err().unwrap(), store.create_namespace(&provider).await.unwrap_err(), store.delete_namespace(&provider).await.unwrap_err(), store.list_namespaces().await.unwrap_err(), @@ -3880,7 +3879,7 @@ mod tests { start: EntryId, ) -> Vec { store - .read(&provider(region_id), start, None) + .read(&provider(region_id), start) .await .unwrap() .try_collect::>() @@ -4198,7 +4197,7 @@ mod tests { ); reads.lock().unwrap().clear(); let error = store - .read(&provider(region(2)), 0, None) + .read(&provider(region(2)), 0) .await .unwrap() .try_collect::>() @@ -4262,7 +4261,7 @@ mod tests { object_store.delete(&path).await.unwrap(); let error = store - .read(&provider(region(1)), 0, None) + .read(&provider(region(1)), 0) .await .unwrap() .try_collect::>() @@ -4369,7 +4368,7 @@ mod tests { let objects = depth + 4; let (store, io, mut parked) = open_parking_range_reads(objects).await; let reads = METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(); - let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let mut stream = store.read(&provider(region(1)), 0).await.unwrap(); // The first poll starts the fetches of the first `depth` objects. assert!(stream.next().now_or_never().is_none()); @@ -4424,7 +4423,7 @@ mod tests { let failed = 5; let (store, io, mut parked) = open_parking_range_reads(objects).await; let reads = METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(); - let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let mut stream = store.read(&provider(region(1)), 0).await.unwrap(); assert!(stream.next().now_or_never().is_none()); let mut wave = parked_operations(&mut parked, READ_PREFETCH_DEPTH).await; @@ -4491,7 +4490,7 @@ mod tests { let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); // A collection removes object 9 after the read listed it. - let stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let stream = store.read(&provider(region(1)), 0).await.unwrap(); store.catalog.write().unwrap().remove_object(collected); object_store .delete(&object_path(&object_store, collected)) @@ -4530,7 +4529,7 @@ mod tests { let objects = READ_PREFETCH_DEPTH as u64 + 4; let (store, _io, mut parked) = open_parking_range_reads(objects).await; let reads = METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(); - let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let mut stream = store.read(&provider(region(1)), 0).await.unwrap(); assert!(stream.next().now_or_never().is_none()); let mut wave = parked_operations(&mut parked, READ_PREFETCH_DEPTH).await; wave.remove(&0).unwrap().send(true).unwrap(); @@ -4556,7 +4555,7 @@ mod tests { let (io, _) = RecordingIo::over(object_store); let store = open_over(io.clone(), &eager()).await; *io.panic_range_read_of.lock().unwrap() = Some(1); - let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let mut stream = store.read(&provider(region(1)), 0).await.unwrap(); assert_eq!(object_entries([0]), next_batch(&mut stream).await); let error = timeout(WAIT, stream.next()) .await @@ -4585,7 +4584,7 @@ mod tests { .clone() .try_acquire_many_owned(READ_PREFETCH_BYTES as u32) .unwrap(); - let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let mut stream = store.read(&provider(region(1)), 0).await.unwrap(); assert!(stream.next().now_or_never().is_none()); let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(0, object_seq); @@ -4638,7 +4637,7 @@ mod tests { // The first poll fetches object 0, and object 1 with its bytes // reserved; the read of object 1 blocks its thread. - let mut stream = store.read(&provider(region(1)), 0, None).await.unwrap(); + let mut stream = store.read(&provider(region(1)), 0).await.unwrap(); let _ = stream.next().now_or_never(); began_rx.recv_timeout(WAIT).unwrap(); assert!(budget.available_permits() < READ_PREFETCH_BYTES); @@ -5771,11 +5770,7 @@ mod tests { assert_stopped(&timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err()); for error in [ store.latest_entry_id(&provider(region_id)).unwrap_err(), - store - .read(&provider(region_id), 1, None) - .await - .err() - .unwrap(), + store.read(&provider(region_id), 1).await.err().unwrap(), ] { assert!( matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), @@ -7108,7 +7103,7 @@ mod tests { append(&store, region_id, "a2").await.unwrap(); // The read lists objects 1 and 2 before a collection deletes object 1. - let stream = store.read(&provider(region_id), 0, None).await.unwrap(); + let stream = store.read(&provider(region_id), 0).await.unwrap(); store .obsolete(&provider(region_id), region_id, id(1, 1)) .await @@ -7152,7 +7147,7 @@ mod tests { // The read lists objects 1 and 2. The delete of object 1 removes // it, but its outcome is not reported yet. - let stream = store.read(&provider(region_id), 0, None).await.unwrap(); + let stream = store.read(&provider(region_id), 0).await.unwrap(); store .obsolete(&provider(region_id), region_id, id(1, 1)) .await @@ -7905,7 +7900,7 @@ mod tests { /// Asserts that a read of the poisoned store fails with the terminal error. async fn assert_read_is_terminal(store: &ObjectStoreLogStore, region_id: RegionId) { let error = store - .read(&provider(region_id), 1, None) + .read(&provider(region_id), 1) .await .err() .expect("a read of a poisoned store must fail"); @@ -9183,7 +9178,7 @@ mod tests { .unwrap(); *io.fail_range_read_at.lock().unwrap() = Some((1, segment.start)); let error = store - .read(&provider(region_id), 1, None) + .read(&provider(region_id), 1) .await .unwrap() .try_collect::>() diff --git a/src/log-store/src/raft_engine/log_store.rs b/src/log-store/src/raft_engine/log_store.rs index 923110858b9..32a028f1691 100644 --- a/src/log-store/src/raft_engine/log_store.rs +++ b/src/log-store/src/raft_engine/log_store.rs @@ -25,7 +25,7 @@ use raft_engine::{Config, Engine, LogBatch, MessageExt, ReadableSize, RecoveryMo use snafu::{OptionExt, ResultExt, ensure}; use store_api::logstore::entry::{Entry, Id as EntryId, NaiveEntry}; use store_api::logstore::provider::{Provider, RaftEngineProvider}; -use store_api::logstore::{AppendBatchResponse, LogStore, SendableEntryStream, WalIndex}; +use store_api::logstore::{AppendBatchResponse, LogStore, SendableEntryStream}; use store_api::storage::RegionId; use crate::error::{ @@ -293,7 +293,6 @@ impl LogStore for RaftEngineLogStore { &self, provider: &Provider, entry_id: EntryId, - _index: Option, ) -> Result> { let ns = provider .as_raft_engine_provider() @@ -637,7 +636,7 @@ mod tests { } let mut entries = HashSet::with_capacity(1024); let mut s = logstore - .read(&Provider::raft_engine_provider(1), 0, None) + .read(&Provider::raft_engine_provider(1), 0) .await .unwrap(); while let Some(r) = s.next().await { @@ -672,7 +671,7 @@ mod tests { .is_ok() ); let entries = logstore - .read(&Provider::raft_engine_provider(1), 1, None) + .read(&Provider::raft_engine_provider(1), 1) .await .unwrap() .collect::>() @@ -690,7 +689,7 @@ mod tests { let entries = collect_entries( logstore - .read(&Provider::raft_engine_provider(1), 1, None) + .read(&Provider::raft_engine_provider(1), 1) .await .unwrap(), ) @@ -776,7 +775,7 @@ mod tests { logstore.obsolete(&namespace, region_id, 100).await.unwrap(); assert_eq!(101, logstore.engine.first_index(namespace_id).unwrap()); - let res = logstore.read(&namespace, 100, None).await.unwrap(); + let res = logstore.read(&namespace, 100).await.unwrap(); let mut vec = collect_entries(res).await; vec.sort_by(|a, b| a.entry_id().partial_cmp(&b.entry_id()).unwrap()); assert_eq!(101, vec.first().unwrap().entry_id()); diff --git a/src/log-store/src/test_util/log_store_util.rs b/src/log-store/src/test_util/log_store_util.rs index 240006595c6..74b63ae080c 100644 --- a/src/log-store/src/test_util/log_store_util.rs +++ b/src/log-store/src/test_util/log_store_util.rs @@ -34,16 +34,13 @@ pub async fn create_tmp_local_file_log_store>(path: P) -> RaftEng /// Create a [KafkaLogStore]. pub async fn create_kafka_log_store(broker_endpoints: Vec) -> KafkaLogStore { - KafkaLogStore::try_new( - &DatanodeKafkaConfig { - connection: KafkaConnectionConfig { - broker_endpoints, - ..Default::default() - }, + KafkaLogStore::try_new(&DatanodeKafkaConfig { + connection: KafkaConnectionConfig { + broker_endpoints, ..Default::default() }, - None, - ) + ..Default::default() + }) .await .unwrap() } diff --git a/src/metric-engine/src/engine/catchup.rs b/src/metric-engine/src/engine/catchup.rs index 61e7591893c..357404545da 100644 --- a/src/metric-engine/src/engine/catchup.rs +++ b/src/metric-engine/src/engine/catchup.rs @@ -54,7 +54,6 @@ impl MetricEngineInner { set_writable: req.set_writable, entry_id: req.metadata_entry_id, metadata_entry_id: None, - location_id: req.location_id, checkpoint: req.checkpoint.map(|c| ReplayCheckpoint { entry_id: c.metadata_entry_id.unwrap_or_default(), metadata_entry_id: None, @@ -67,7 +66,6 @@ impl MetricEngineInner { set_writable: req.set_writable, entry_id: req.entry_id, metadata_entry_id: None, - location_id: req.location_id, checkpoint: req.checkpoint.map(|c| ReplayCheckpoint { entry_id: c.entry_id, metadata_entry_id: None, diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 35f15f2393c..47b34870ccd 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -1714,7 +1714,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( ); // Inspect persisted WAL mutations, not just an encoder or counter. - let mut reader = wal.wal_entry_reader(®ion.provider, region_id, None); + let mut reader = wal.wal_entry_reader(®ion.provider, region_id); let entries = reader .read(®ion.provider, 1) .unwrap() diff --git a/src/mito2/src/engine/batch_catchup_test.rs b/src/mito2/src/engine/batch_catchup_test.rs index 5af92280261..06f2a093523 100644 --- a/src/mito2/src/engine/batch_catchup_test.rs +++ b/src/mito2/src/engine/batch_catchup_test.rs @@ -184,7 +184,6 @@ async fn test_batch_catchup_with_format(factory: Option, flat_f set_writable: true, entry_id: None, metadata_entry_id: None, - location_id: None, checkpoint: None, }, ) @@ -231,7 +230,6 @@ async fn test_batch_catchup_err_with_format(factory: Option, fl set_writable: true, entry_id: None, metadata_entry_id: None, - location_id: None, checkpoint: None, }, ) diff --git a/src/mito2/src/engine/object_store_wal_test.rs b/src/mito2/src/engine/object_store_wal_test.rs index d4225b6fcf4..cd27c63b48d 100644 --- a/src/mito2/src/engine/object_store_wal_test.rs +++ b/src/mito2/src/engine/object_store_wal_test.rs @@ -147,7 +147,7 @@ async fn count_rows(engine: &MitoEngine, region_id: RegionId) -> usize { /// Returns the ids of the entries the store still hands out for `region_id`. async fn stored_entry_ids(log_store: &ObjectStoreLogStore, region_id: RegionId) -> Vec { log_store - .read(&object_store_provider(region_id), 0, None) + .read(&object_store_provider(region_id), 0) .await .unwrap() .try_collect::>() diff --git a/src/mito2/src/region/catchup.rs b/src/mito2/src/region/catchup.rs index 0d8cfa8ed8c..e9f7f7ef43c 100644 --- a/src/mito2/src/region/catchup.rs +++ b/src/mito2/src/region/catchup.rs @@ -31,7 +31,6 @@ pub struct RegionCatchupTask { replay_checkpoint_entry_id: Option, expected_last_entry_id: Option, allow_stale_entries: bool, - location_id: Option, wal: Wal, } @@ -43,17 +42,10 @@ impl RegionCatchupTask { replay_checkpoint_entry_id: None, expected_last_entry_id: None, allow_stale_entries, - location_id: None, wal, } } - /// Sets the location id. - pub(crate) fn with_location_id(mut self, location_id: Option) -> Self { - self.location_id = location_id; - self - } - /// Sets the expected last entry id. pub(crate) fn with_expected_last_entry_id( mut self, @@ -102,10 +94,7 @@ impl RegionCatchupTask { .entry_receiver .take() .map(|r| Box::new(r) as _) - .unwrap_or_else(|| { - self.wal - .wal_entry_reader(&self.region.provider, region_id, self.location_id) - }); + .unwrap_or_else(|| self.wal.wal_entry_reader(&self.region.provider, region_id)); let on_region_opened = self.wal.on_region_opened(); let last_entry_id = replay_memtable( &self.region.provider, diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 7a93799ae5f..37c91946db3 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -547,7 +547,7 @@ impl RegionOpener { let wal_entry_reader = self .wal_entry_reader .take() - .unwrap_or_else(|| wal.wal_entry_reader(&provider, region_id, None)); + .unwrap_or_else(|| wal.wal_entry_reader(&provider, region_id)); let on_region_opened = wal.on_region_opened(); let object_store = get_object_store(®ion_options.storage, &self.object_store_manager)?; diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index 772aac336cf..76b3e3a2cdf 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -33,7 +33,7 @@ use futures::stream::BoxStream; use snafu::ResultExt; use store_api::logstore::entry::Entry; use store_api::logstore::provider::Provider; -use store_api::logstore::{AppendBatchResponse, LogStore, WalIndex}; +use store_api::logstore::{AppendBatchResponse, LogStore}; use store_api::storage::RegionId; use crate::error::{BuildEntrySnafu, DeleteWalSnafu, Result, WaitWalDurableSnafu, WriteWalSnafu}; @@ -163,7 +163,6 @@ impl Wal { &self, provider: &Provider, region_id: RegionId, - location_id: Option, ) -> Box { match provider { Provider::RaftEngine(_) => Box::new(LogStoreEntryReader::new( @@ -171,15 +170,9 @@ impl Wal { )), // Entries of many regions share one namespace, so the reader filters by region. Provider::Kafka(_) | Provider::ObjectStore(_) => { - let reader = if let Some(location_id) = location_id { - LogStoreRawEntryReader::new(self.store.clone()) - .with_wal_index(WalIndex::new(region_id, location_id)) - } else { - LogStoreRawEntryReader::new(self.store.clone()) - }; - Box::new(LogStoreEntryReader::new(RegionRawEntryReader::new( - reader, region_id, + LogStoreRawEntryReader::new(self.store.clone()), + region_id, ))) } Provider::Noop => Box::new(NoopEntryReader), @@ -194,7 +187,7 @@ impl Wal { start_id: EntryId, provider: &'a Provider, ) -> Result> { - let mut reader = self.wal_entry_reader(provider, region_id, None); + let mut reader = self.wal_entry_reader(provider, region_id); reader.read(provider, start_id) } diff --git a/src/mito2/src/wal/raw_entry_reader.rs b/src/mito2/src/wal/raw_entry_reader.rs index 8a79d6c3b16..0841d21f669 100644 --- a/src/mito2/src/wal/raw_entry_reader.rs +++ b/src/mito2/src/wal/raw_entry_reader.rs @@ -18,9 +18,9 @@ use async_stream::try_stream; use common_error::ext::BoxedError; use futures::stream::BoxStream; use snafu::ResultExt; +use store_api::logstore::LogStore; use store_api::logstore::entry::Entry; use store_api::logstore::provider::Provider; -use store_api::logstore::{LogStore, WalIndex}; use store_api::storage::RegionId; use tokio_stream::StreamExt; @@ -38,20 +38,11 @@ pub(crate) trait RawEntryReader: Send + Sync { /// Implement the [RawEntryReader] for the [LogStore]. pub struct LogStoreRawEntryReader { store: Arc, - wal_index: Option, } impl LogStoreRawEntryReader { pub fn new(store: Arc) -> Self { - Self { - store, - wal_index: None, - } - } - - pub fn with_wal_index(mut self, wal_index: WalIndex) -> Self { - self.wal_index = Some(wal_index); - self + Self { store } } } @@ -59,10 +50,9 @@ impl RawEntryReader for LogStoreRawEntryReader { fn read(&self, provider: &Provider, start_id: EntryId) -> Result> { let store = self.store.clone(); let provider = provider.clone(); - let wal_index = self.wal_index; let stream = try_stream!({ let mut stream = store - .read(&provider, start_id, wal_index) + .read(&provider, start_id) .await .map_err(BoxedError::new) .with_context(|_| error::ReadWalSnafu { @@ -129,9 +119,7 @@ mod tests { use futures::{TryStreamExt, stream}; use store_api::logstore::entry::{Entry, NaiveEntry}; - use store_api::logstore::{ - AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex, - }; + use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream}; use store_api::storage::RegionId; use super::*; @@ -161,7 +149,6 @@ mod tests { &self, _provider: &Provider, _id: EntryId, - _index: Option, ) -> Result, Self::Error> { Ok(Box::pin(stream::iter(vec![Ok(self.entries.clone())]))) } diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index 81838060385..c526aab8c0c 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -73,7 +73,6 @@ impl RegionWorkerLoop { let mut task = RegionCatchupTask::new(region.clone(), wal, allow_stale_entries) .with_entry_receiver(entry_receiver) .with_expected_last_entry_id(request.entry_id) - .with_location_id(request.location_id) .with_replay_checkpoint_entry_id(request.checkpoint.map(|c| c.entry_id)); match task.run().await { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 1383a3e9a5d..721235e71b5 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -758,7 +758,7 @@ mod tests { }; use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::Provider; - use store_api::logstore::{AppendBatchResponse, EntryId, SendableEntryStream, WalIndex}; + use store_api::logstore::{AppendBatchResponse, EntryId, SendableEntryStream}; use store_api::region_request::AffectedRows; use tokio::sync::oneshot; @@ -828,7 +828,6 @@ mod tests { &self, _provider: &Provider, _entry_id: EntryId, - _index: Option, ) -> LogStoreResult> { Ok(Box::pin(stream::empty())) } diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index 6084f207641..dd71142def5 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -30,22 +30,6 @@ pub use crate::logstore::entry::Id as EntryId; use crate::logstore::provider::Provider; use crate::storage::RegionId; -// The information used to locate WAL index for the specified region. -#[derive(Debug, Clone, Copy)] -pub struct WalIndex { - pub region_id: RegionId, - pub location_id: u64, -} - -impl WalIndex { - pub fn new(region_id: RegionId, location_id: u64) -> Self { - Self { - region_id, - location_id, - } - } -} - /// `LogStore` serves as a Write-Ahead-Log for storage engine. #[async_trait::async_trait] pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { @@ -64,7 +48,6 @@ pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { &self, provider: &Provider, id: EntryId, - index: Option, ) -> Result, Self::Error>; /// Creates a new `Namespace` from the given ref. diff --git a/src/store-api/src/logstore/provider.rs b/src/store-api/src/logstore/provider.rs index 4fe7e4dfe7b..4af4caaaa60 100644 --- a/src/store-api/src/logstore/provider.rs +++ b/src/store-api/src/logstore/provider.rs @@ -187,7 +187,7 @@ mod tests { use super::*; use crate::logstore::entry::Entry; - use crate::logstore::{AppendBatchResponse, EntryId, SendableEntryStream, WalIndex}; + use crate::logstore::{AppendBatchResponse, EntryId, SendableEntryStream}; /// A log store whose latest entry id is fixed. #[derive(Debug)] @@ -212,7 +212,6 @@ mod tests { &self, _provider: &Provider, _id: EntryId, - _index: Option, ) -> Result, Self::Error> { unreachable!() } diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index e1f0eb22497..ebb2f7c3684 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -1929,8 +1929,6 @@ pub struct RegionCatchupRequest { /// The `entry_id` that was expected to reply to. /// `None` stands replaying to latest. pub metadata_entry_id: Option, - /// The hint for replaying memtable. - pub location_id: Option, /// Replay checkpoint. pub checkpoint: Option, } From caefcfd01244107b6aa4e106529af7b9f2f57c04 Mon Sep 17 00:00:00 2001 From: NightG1owZ <3112428897@qq.com> Date: Wed, 14 Oct 2026 02:00:39 +0000 Subject: [PATCH 34/46] test: add inverted index compatibility case (#9519) * test: add inverted index compatibility case (#5188) Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> * test(compat): assert the newly indexed column stays queryable (#9535) Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> * test(compat): narrow to_range to versions carrying #9545 The filtered read added in the previous commit only returns rows once #9545 is present. That fix is on main only (no backport-* labels), so it ships in the v1.4.0 line. Narrow to_range to >=v1.4.0-alpha.0 so the documented downgrade invocation (--to-version v1.1.4) is not broken. Also add the missing trailing newline to the three case files. Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> --------- Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> Co-authored-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> --- .../cases/inverted_index_survival/case.toml | 8 + .../cases/inverted_index_survival/setup.sql | 14 ++ .../inverted_index_survival/verify.result | 145 ++++++++++++++++++ .../cases/inverted_index_survival/verify.sql | 50 ++++++ 4 files changed, 217 insertions(+) create mode 100644 tests/compatibility/cases/inverted_index_survival/case.toml create mode 100644 tests/compatibility/cases/inverted_index_survival/setup.sql create mode 100644 tests/compatibility/cases/inverted_index_survival/verify.result create mode 100644 tests/compatibility/cases/inverted_index_survival/verify.sql diff --git a/tests/compatibility/cases/inverted_index_survival/case.toml b/tests/compatibility/cases/inverted_index_survival/case.toml new file mode 100644 index 00000000000..e5ce3f93788 --- /dev/null +++ b/tests/compatibility/cases/inverted_index_survival/case.toml @@ -0,0 +1,8 @@ +name = "inverted_index_survival" +reason = "Verify column-level INVERTED INDEX written by old binaries stays readable, writable and alterable after upgrade, and that a column indexed after the fact stays queryable." +introduced_by = "#5188; regression guard for #9535" +topologies = ["distributed", "standalone"] +from_range = ["*"] +to_range = [">=v1.4.0-alpha.0"] +features = ["table", "index", "inverted_index", "metadata", "alter", "read", "write"] +owner = "storage" diff --git a/tests/compatibility/cases/inverted_index_survival/setup.sql b/tests/compatibility/cases/inverted_index_survival/setup.sql new file mode 100644 index 00000000000..ca11fa0c353 --- /dev/null +++ b/tests/compatibility/cases/inverted_index_survival/setup.sql @@ -0,0 +1,14 @@ +CREATE TABLE t_inverted_index_survival ( + host STRING INVERTED INDEX, + msg STRING INVERTED INDEX, + msg2 STRING, + ts TIMESTAMP TIME INDEX, + cpu DOUBLE, + PRIMARY KEY(host) +) ENGINE=mito; + +INSERT INTO t_inverted_index_survival VALUES + ('h1', 'first event', 'first note', '2026-07-07 00:00:00+0000', 1.0), + ('h2', 'second event', 'second note', '2026-07-07 00:01:00+0000', 2.0); + +ADMIN FLUSH_TABLE('t_inverted_index_survival'); diff --git a/tests/compatibility/cases/inverted_index_survival/verify.result b/tests/compatibility/cases/inverted_index_survival/verify.result new file mode 100644 index 00000000000..1678f07a72f --- /dev/null +++ b/tests/compatibility/cases/inverted_index_survival/verify.result @@ -0,0 +1,145 @@ +-- 1. Index metadata written by the old binary must survive. +SHOW CREATE TABLE t_inverted_index_survival; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| t_inverted_index_survival | CREATE TABLE IF NOT EXISTS "t_inverted_index_survival" ( | +| | "host" STRING NULL INVERTED INDEX, | +| | "msg" STRING NULL INVERTED INDEX, | +| | "msg2" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "cpu" DOUBLE NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------------+----------------------------------------------------------+ + +-- 2. Rows written by the old binary must still be readable. +SELECT host, msg, cpu +FROM t_inverted_index_survival +ORDER BY host, ts; + ++------+--------------+-----+ +| host | msg | cpu | ++------+--------------+-----+ +| h1 | first event | 1.0 | +| h2 | second event | 2.0 | ++------+--------------+-----+ + +-- 3. Filtered read on the column indexed by the old binary. +SELECT host, msg +FROM t_inverted_index_survival +WHERE msg = 'first event' +ORDER BY host, ts; + ++------+-------------+ +| host | msg | ++------+-------------+ +| h1 | first event | ++------+-------------+ + +-- 4. The new binary must still be able to write into the old table. +INSERT INTO t_inverted_index_survival VALUES + ('h3', 'third event', 'third note', '2026-07-07 00:02:00+0000', 3.0); + +Affected Rows: 1 + +SELECT host, msg, cpu +FROM t_inverted_index_survival +ORDER BY host, ts; + ++------+--------------+-----+ +| host | msg | cpu | ++------+--------------+-----+ +| h1 | first event | 1.0 | +| h2 | second event | 2.0 | +| h3 | third event | 3.0 | ++------+--------------+-----+ + +-- 5. Enabling an index on an existing column must keep SSTs written before the +-- index existed readable (regression guard for #9535). 'first note' and +-- 'second note' live in the SST written by the old binary, which carries +-- inverted indexes for host/msg but not for msg2. +ALTER TABLE t_inverted_index_survival MODIFY COLUMN msg2 SET INVERTED INDEX; + +Affected Rows: 0 + +SHOW CREATE TABLE t_inverted_index_survival; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| t_inverted_index_survival | CREATE TABLE IF NOT EXISTS "t_inverted_index_survival" ( | +| | "host" STRING NULL INVERTED INDEX, | +| | "msg" STRING NULL INVERTED INDEX, | +| | "msg2" STRING NULL INVERTED INDEX, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "cpu" DOUBLE NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------------+----------------------------------------------------------+ + +SELECT host, msg2 +FROM t_inverted_index_survival +WHERE msg2 = 'first note' +ORDER BY host, ts; + ++------+------------+ +| host | msg2 | ++------+------------+ +| h1 | first note | ++------+------------+ + +SELECT host, msg2 +FROM t_inverted_index_survival +WHERE msg2 = 'second note' +ORDER BY host, ts; + ++------+-------------+ +| host | msg2 | ++------+-------------+ +| h2 | second note | ++------+-------------+ + +-- A non-matching predicate must still return nothing. +SELECT host, msg2 +FROM t_inverted_index_survival +WHERE msg2 = 'no such note' +ORDER BY host, ts; + ++------+------+ +| host | msg2 | ++------+------+ ++------+------+ + +-- 6. The index can still be removed afterwards. +ALTER TABLE t_inverted_index_survival MODIFY COLUMN msg2 UNSET INVERTED INDEX; + +Affected Rows: 0 + +SHOW CREATE TABLE t_inverted_index_survival; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| t_inverted_index_survival | CREATE TABLE IF NOT EXISTS "t_inverted_index_survival" ( | +| | "host" STRING NULL INVERTED INDEX, | +| | "msg" STRING NULL INVERTED INDEX, | +| | "msg2" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "cpu" DOUBLE NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------------+----------------------------------------------------------+ diff --git a/tests/compatibility/cases/inverted_index_survival/verify.sql b/tests/compatibility/cases/inverted_index_survival/verify.sql new file mode 100644 index 00000000000..1cd5f396852 --- /dev/null +++ b/tests/compatibility/cases/inverted_index_survival/verify.sql @@ -0,0 +1,50 @@ +-- 1. Index metadata written by the old binary must survive. +SHOW CREATE TABLE t_inverted_index_survival; + +-- 2. Rows written by the old binary must still be readable. +SELECT host, msg, cpu +FROM t_inverted_index_survival +ORDER BY host, ts; + +-- 3. Filtered read on the column indexed by the old binary. +SELECT host, msg +FROM t_inverted_index_survival +WHERE msg = 'first event' +ORDER BY host, ts; + +-- 4. The new binary must still be able to write into the old table. +INSERT INTO t_inverted_index_survival VALUES + ('h3', 'third event', 'third note', '2026-07-07 00:02:00+0000', 3.0); + +SELECT host, msg, cpu +FROM t_inverted_index_survival +ORDER BY host, ts; + +-- 5. Enabling an index on an existing column must keep SSTs written before the +-- index existed readable (regression guard for #9535). 'first note' and +-- 'second note' live in the SST written by the old binary, which carries +-- inverted indexes for host/msg but not for msg2. +ALTER TABLE t_inverted_index_survival MODIFY COLUMN msg2 SET INVERTED INDEX; + +SHOW CREATE TABLE t_inverted_index_survival; + +SELECT host, msg2 +FROM t_inverted_index_survival +WHERE msg2 = 'first note' +ORDER BY host, ts; + +SELECT host, msg2 +FROM t_inverted_index_survival +WHERE msg2 = 'second note' +ORDER BY host, ts; + +-- A non-matching predicate must still return nothing. +SELECT host, msg2 +FROM t_inverted_index_survival +WHERE msg2 = 'no such note' +ORDER BY host, ts; + +-- 6. The index can still be removed afterwards. +ALTER TABLE t_inverted_index_survival MODIFY COLUMN msg2 UNSET INVERTED INDEX; + +SHOW CREATE TABLE t_inverted_index_survival; From 75bc4ce46d0e7aedca53dc2b0375ca0658456ff6 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 14 Oct 2026 03:35:58 +0000 Subject: [PATCH 35/46] fix(table): stop claiming tag order for logical metric scans (#9522) * test: cover metric logical scan tag ordering Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(table): only claim projected primary key prefix ordering Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(table): keep physical key ordering for metric projections Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(table): compare unknown partition counts explicitly Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(table): model scoped internal TSID scan properties Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/table/src/table/scan.rs | 272 ++++++++++++++++- tests-integration/src/tests/instance_test.rs | 289 +++++++++++++++++++ 2 files changed, 554 insertions(+), 7 deletions(-) diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 3e0eeab84a9..1bfaea0d1df 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -118,21 +118,34 @@ impl RegionScanExec { } let metadata = scanner.metadata(); - let mut pk_names = metadata + // Keep the primary key columns in the order reported by the metadata (for a metric + // region that is the physical row key order). Reordering them would let a projected + // logical tag look like a prefix of the row key, which yields ordering and + // partitioning claims the scanner cannot honour. + let pk_names = metadata .primary_key_columns() .map(|col| col.column_schema.name.clone()) .collect::>(); - // workaround for logical table - if scanner.properties().is_logical_region() { - pk_names.sort_unstable(); - } + // A region emits rows ordered by its primary key columns and then by the time + // index column. A primary key column that is not projected is absent from the + // arrow schema, and it also breaks the ordering of every primary key column that + // follows it (e.g. the physical `__table_id`/`__tsid` prefix of a metric region + // when reading a logical table). Thus only the longest fully projected prefix of + // the primary key can describe the output ordering, and that prefix can be empty. + let projected_pk_prefix_len = pk_names + .iter() + .position(|name| Column::new_with_schema(name, &arrow_schema).is_err()) + .unwrap_or(pk_names.len()); + // Whether the projected prefix covers the whole row key. + let all_pk_projected = projected_pk_prefix_len == pk_names.len(); + let projected_pk_names = &pk_names[..projected_pk_prefix_len]; let pk_columns = pk_names .iter() .filter_map( |col| Some(Arc::new(Column::new_with_schema(col, &arrow_schema).ok()?) as _), ) .collect::>(); - let pk_sort_columns: Vec = pk_names + let pk_sort_columns: Vec = projected_pk_names .iter() .filter_map(|col| { Some(PhysicalSortExpr::new( @@ -165,7 +178,9 @@ impl RegionScanExec { let mut orderings = Vec::with_capacity(2); let mut pk_time_ordering = pk_sort_columns.clone(); - if let Some(ts) = ts_col.clone() { + // The time index column extends the ordering only when the whole row key + // precedes it; otherwise rows of different series can interleave. + if all_pk_projected && let Some(ts) = ts_col.clone() { pk_time_ordering.push(ts); } orderings.push(pk_time_ordering); @@ -694,9 +709,11 @@ mod test { use datatypes::vectors::{Int32Vector, TimestampMillisecondVector, VectorRef}; use futures::TryStreamExt; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME; use store_api::region_engine::{ PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties, SinglePartitionScanner, }; + use store_api::storage::consts::ReservedColumnId; use store_api::storage::{RegionId, TimeSeriesDistribution}; use super::*; @@ -953,6 +970,247 @@ mod test { } } + /// Returns whether the scan claims the given output ordering. + fn claims_ordering(plan: &RegionScanExec, names: &[&str]) -> bool { + plan.properties() + .equivalence_properties() + .ordering_satisfy(names.iter().map(|name| { + PhysicalSortExpr::new( + Arc::new(Column::new_with_schema(name, &plan.arrow_schema).unwrap()), + SortOptions { + descending: false, + nulls_first: true, + }, + ) + })) + .unwrap() + } + + /// Column schema used by the ordering fixtures below. + fn column_schema_of(name: &str) -> ColumnSchema { + let data_type = match name { + "ts" => ConcreteDataType::timestamp_millisecond_datatype(), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME => ConcreteDataType::uint32_datatype(), + DATA_SCHEMA_TSID_COLUMN_NAME => ConcreteDataType::uint64_datatype(), + _ => ConcreteDataType::string_datatype(), + }; + ColumnSchema::new(name, data_type, false) + } + + /// The scan must describe its output ordering with the primary key order reported by + /// the metadata. Sorting the names (the removed logical-table workaround) would let a + /// projected column look like a row key prefix and claim an ordering the scanner does + /// not produce; the time index only extends a fully projected row key. + #[test] + fn test_per_series_ordering_follows_metadata_pk_order() { + // Deliberately not alphabetical: sorting these names would give `[a, b, z]`. + let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5690)); + for (index, name) in ["z", "a", "b", "ts"].into_iter().enumerate() { + builder.push_column_metadata(ColumnMetadata { + column_schema: column_schema_of(name), + semantic_type: if name == "ts" { + SemanticType::Timestamp + } else { + SemanticType::Tag + }, + column_id: index as u32 + 1, + }); + } + builder.primary_key(vec![1, 2, 3]); + let metadata = Arc::new(builder.build().unwrap()); + + // The claims must not depend on the scanner reporting a logical region. + for logical_region in [false, true] { + for (projection, claimed, unclaimed) in [ + // Missing the first key: neither the remaining keys nor the time index are + // ordered. + ( + vec!["a", "b", "ts"], + vec![], + vec![vec!["a"], vec!["b"], vec!["ts"], vec!["a", "b"]], + ), + // Missing a key in the middle: only the prefix before it is ordered. + ( + vec!["z", "b", "ts"], + vec![vec!["z"]], + vec![vec!["b"], vec!["z", "b"], vec!["z", "ts"]], + ), + // A valid but incomplete prefix: the time index is not ordered after it. + ( + vec!["z", "a", "ts"], + vec![vec!["z"], vec!["z", "a"]], + vec![vec!["a", "z"], vec!["z", "a", "ts"]], + ), + // The whole row key followed by the time index. + ( + vec!["z", "a", "b", "ts"], + vec![vec!["z", "a", "b"], vec!["z", "a", "b", "ts"]], + vec![vec!["a"], vec!["a", "b", "z"]], + ), + ] { + let schema = Arc::new(Schema::new( + projection + .iter() + .map(|name| column_schema_of(name)) + .collect(), + )); + let mut properties = ScannerProperties::new(vec![vec![], vec![]], false, 0); + properties.set_logical_region(logical_region); + let scanner = Box::new(RepeatableScanner { + batches: RecordBatches::try_new(schema, vec![]).unwrap(), + properties, + metadata: metadata.clone(), + dynamic_filters: Arc::new(Mutex::new(Vec::new())), + }); + let plan = RegionScanExec::new( + scanner, + ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + None, + ) + .unwrap(); + + for ordering in claimed { + assert!( + claims_ordering(&plan, &ordering), + "logical_region={logical_region} projection {projection:?} must claim \ + ordering {ordering:?}" + ); + } + for ordering in unclaimed { + assert!( + !claims_ordering(&plan, &ordering), + "logical_region={logical_region} projection {projection:?} must not \ + claim ordering {ordering:?}" + ); + } + } + } + } + + /// Metric TSID properties require projecting `__tsid` and restricting the scan to one + /// `__table_id`; hiding the physical row-key prefix must not claim a visible tag or time as + /// order or hash partition columns. + #[test] + fn test_metric_projection_claims_only_projected_row_key_prefix() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5691)); + for (name, column_id) in [ + ( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + ReservedColumnId::table_id(), + ), + (DATA_SCHEMA_TSID_COLUMN_NAME, ReservedColumnId::tsid()), + ("A", 1), + ("marker", 2), + ("ts", 3), + ] { + builder.push_column_metadata(ColumnMetadata { + column_schema: column_schema_of(name), + semantic_type: if name == "ts" { + SemanticType::Timestamp + } else { + SemanticType::Tag + }, + column_id, + }); + } + // The physical row key order: `__table_id`, `__tsid`, then the tag columns. + builder.primary_key(vec![ + ReservedColumnId::table_id(), + ReservedColumnId::tsid(), + 1, + 2, + ]); + let metadata = Arc::new(builder.build().unwrap()); + + for (projection, partitioned_by_tsid, claimed, unclaimed) in [ + // The logical projection drops the row key prefix, so `A` is not a prefix of the + // row key and neither `A` nor `ts` is ordered. + ( + vec!["A", "marker", "ts"], + false, + vec![], + vec![vec!["A"], vec!["marker"], vec!["ts"], vec!["A", "ts"]], + ), + // Projecting internal `__tsid` plus the single-table filter preserves metric TSID properties. + ( + vec![DATA_SCHEMA_TSID_COLUMN_NAME, "A", "marker", "ts"], + true, + vec![ + vec![DATA_SCHEMA_TSID_COLUMN_NAME], + vec![DATA_SCHEMA_TSID_COLUMN_NAME, "ts"], + ], + vec![vec!["A"], vec!["ts"], vec!["A", "ts"]], + ), + ] { + let schema = Arc::new(Schema::new( + projection + .iter() + .map(|name| column_schema_of(name)) + .collect(), + )); + let mut properties = ScannerProperties::new(vec![vec![], vec![]], false, 0); + properties.set_logical_region(true); + let scanner = Box::new(RepeatableScanner { + batches: RecordBatches::try_new(schema, vec![]).unwrap(), + properties, + metadata: metadata.clone(), + dynamic_filters: Arc::new(Mutex::new(Vec::new())), + }); + let plan = RegionScanExec::new( + scanner, + ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + filters: if partitioned_by_tsid { + vec![ + datafusion_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + .eq(datafusion_expr::lit(42u32)), + ] + } else { + vec![] + }, + ..Default::default() + }, + None, + ) + .unwrap(); + + let expected_partitioning = if partitioned_by_tsid { + Partitioning::Hash( + vec![Arc::new( + Column::new_with_schema(DATA_SCHEMA_TSID_COLUMN_NAME, &plan.arrow_schema) + .unwrap(), + ) as _], + 2, + ) + } else { + Partitioning::UnknownPartitioning(2) + }; + match (&plan.properties.partitioning, &expected_partitioning) { + ( + Partitioning::UnknownPartitioning(actual_count), + Partitioning::UnknownPartitioning(expected_count), + ) => assert_eq!(actual_count, expected_count, "projection {projection:?}"), + (actual, expected) => assert_eq!(actual, expected, "projection {projection:?}"), + } + + for ordering in claimed { + assert!( + claims_ordering(&plan, &ordering), + "projection {projection:?} must claim ordering {ordering:?}" + ); + } + for ordering in unclaimed { + assert!( + !claims_ordering(&plan, &ordering), + "projection {projection:?} must not claim ordering {ordering:?}" + ); + } + } + } + #[test] fn test_count_statistics_require_exact_source_rows() { for (append_mode, exact, expected) in [ diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 1a9cc1790f7..22a0b993d17 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -258,6 +258,295 @@ WHERE greptime_timestamp = ( ); } +#[apply(both_instances_cases)] +async fn test_metric_logical_scan_tag_order(instance: Arc) { + let frontend = instance.frontend(); + + execute_sql( + &frontend, + r#"CREATE TABLE f04_scan_phy ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX + ) ENGINE = metric WITH (physical_metric_table = 'true')"#, + ) + .await; + execute_sql( + &frontend, + r#"CREATE TABLE f04_scan_logical ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX + ) ENGINE = metric WITH (on_physical_table = 'f04_scan_phy')"#, + ) + .await; + execute_sql( + &frontend, + "INSERT INTO f04_scan_logical (host, val, ts) VALUES \ + ('a', 1, 0), ('b', 2, 0), ('c', 3, 0), ('d', 4, 0), \ + ('e', 5, 0), ('f', 6, 0), ('g', 7, 0), ('h', 8, 0)", + ) + .await; + + // The ordered rows are compared exactly as printed. Sorting them before + // comparing (or comparing them unordered) would accept a scan that returns + // the physical table's tag order instead of the requested `host` order. + async fn ordered_rows(instance: &Arc, sql: &str, parallelism: usize) -> String { + execute_sql_with_query_parallelism(instance, sql, parallelism) + .await + .data + .pretty_print() + .await + } + + let all_hosts = "\ ++------+---------------------+-----+ +| host | ts | val | ++------+---------------------+-----+ +| a | 1970-01-01T00:00:00 | 1.0 | +| b | 1970-01-01T00:00:00 | 2.0 | +| c | 1970-01-01T00:00:00 | 3.0 | +| d | 1970-01-01T00:00:00 | 4.0 | +| e | 1970-01-01T00:00:00 | 5.0 | +| f | 1970-01-01T00:00:00 | 6.0 | +| g | 1970-01-01T00:00:00 | 7.0 | +| h | 1970-01-01T00:00:00 | 8.0 | ++------+---------------------+-----+"; + let first_three_hosts = "\ ++------+---------------------+-----+ +| host | ts | val | ++------+---------------------+-----+ +| a | 1970-01-01T00:00:00 | 1.0 | +| b | 1970-01-01T00:00:00 | 2.0 | +| c | 1970-01-01T00:00:00 | 3.0 | ++------+---------------------+-----+"; + + // `host` is NOT NULL, so `NULLS FIRST` and the SQL default `NULLS LAST` return the same + // rows. The `NULLS FIRST` queries spell out the ordering the scan used to claim, so the + // optimizer may drop the sort and expose a wrong tag order; the `NULLS LAST` queries + // keep the sort. Both are replayed in memory and after flushing the physical table. + let mut mismatches = Vec::new(); + for flushed in [false, true] { + if flushed { + execute_sql(&frontend, "ADMIN FLUSH_TABLE('f04_scan_phy')").await; + } + for parallelism in [1, 4] { + for (query, expected) in [ + ( + "SELECT host, ts, val FROM f04_scan_logical ORDER BY host NULLS FIRST", + all_hosts, + ), + ( + "SELECT host, ts, val FROM f04_scan_logical ORDER BY host NULLS FIRST LIMIT 3", + first_three_hosts, + ), + ( + "SELECT host, ts, val FROM f04_scan_logical ORDER BY host NULLS LAST", + all_hosts, + ), + ( + "SELECT host, ts, val FROM f04_scan_logical ORDER BY host NULLS LAST LIMIT 3", + first_three_hosts, + ), + ] { + let actual = ordered_rows(&frontend, query, parallelism).await; + if actual.trim() != expected.trim() { + mismatches.push(format!( + "query `{query}` with query_parallelism={parallelism} flushed={flushed} \ + returned:\n{actual}" + )); + } + } + } + } + assert!( + mismatches.is_empty(), + "logical metric table rows are not ordered by host, expected `a`..`h`:\n{}", + mismatches.join("\n---\n") + ); +} + +/// The uppercase tag `"A"` sorts before the physical `__table_id`/`__tsid` row key prefix +/// by name, but it is not a prefix of the row key. A scan that claims it as ordering or +/// partitioning lets the optimizer skip the sort or the repartition the query needs, so the +/// per-host count comes back split across scan partitions in physical tsid order. +#[apply(both_instances_cases)] +async fn test_metric_logical_scan_uppercase_tag_order(instance: Arc) { + let frontend = instance.frontend(); + + execute_sql( + &frontend, + r#"CREATE TABLE f04_upper_phy ( + "A" STRING, + marker STRING, + val DOUBLE, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY ("A", marker) + ) ENGINE = metric WITH (physical_metric_table = 'true')"#, + ) + .await; + execute_sql( + &frontend, + r#"CREATE TABLE f04_upper_logical ( + "A" STRING, + marker STRING, + val DOUBLE, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY ("A", marker) + ) ENGINE = metric WITH (on_physical_table = 'f04_upper_phy')"#, + ) + .await; + + // 8 hosts x 4 markers x 2 timestamps = 64 rows. Every host spans four series, so a count + // computed per scan partition would return several rows per host with partial counts. + let mut rows = Vec::with_capacity(64); + for host in 0..8 { + for marker in ["a", "b", "c", "d"] { + for ts in [1000, 2000] { + rows.push(format!("('h{host}', '{marker}', {host}, {ts})")); + } + } + } + execute_sql( + &frontend, + &format!( + "INSERT INTO f04_upper_logical (\"A\", marker, val, ts) VALUES {}", + rows.join(", ") + ), + ) + .await; + + let expected = "\ ++----+---+ +| A | n | ++----+---+ +| h0 | 8 | +| h1 | 8 | +| h2 | 8 | +| h3 | 8 | +| h4 | 8 | +| h5 | 8 | +| h6 | 8 | +| h7 | 8 | ++----+---+"; + + // The rows are compared exactly as printed: sorting them before comparing would accept + // a wrong count or a wrong row order. `NULLS FIRST` spells out the ordering the scan + // used to claim, so the optimizer may drop the sort; the default `NULLS LAST` keeps it + // and a failure there can only come from a wrong partitioning claim. Both orders are + // replayed in memory and after flushing the physical table. + let queries = [ + "SELECT \"A\", count(*) AS n FROM f04_upper_logical GROUP BY \"A\" ORDER BY \"A\"", + "SELECT \"A\", count(*) AS n FROM f04_upper_logical GROUP BY \"A\" ORDER BY \"A\" NULLS FIRST", + ]; + + let mut mismatches = Vec::new(); + for flushed in [false, true] { + if flushed { + execute_sql(&frontend, "ADMIN FLUSH_TABLE('f04_upper_phy')").await; + } + for parallelism in [1, 4] { + for query in queries { + let actual = execute_sql_with_query_parallelism(&frontend, query, parallelism) + .await + .data + .pretty_print() + .await; + if actual.trim() != expected.trim() { + mismatches.push(format!( + "query `{query}` with query_parallelism={parallelism} flushed={flushed} \ + returned:\n{actual}" + )); + } + } + } + } + assert!( + mismatches.is_empty(), + "logical metric table must return one row of eight per host in `A` order:\n{}", + mismatches.join("\n---\n") + ); +} + +/// The small shape of the same problem: four rows and two groups. Each host has two series, +/// so the count must not be computed per scan partition, and `ORDER BY "A"` pins the exact +/// row order at query_parallelism 4 as well as at 1. +#[apply(both_instances_cases)] +async fn test_metric_logical_scan_small_group_count(instance: Arc) { + let frontend = instance.frontend(); + + execute_sql( + &frontend, + r#"CREATE TABLE f04_upper_min_phy ( + "A" STRING, + marker STRING, + val DOUBLE, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY ("A", marker) + ) ENGINE = metric WITH (physical_metric_table = 'true')"#, + ) + .await; + execute_sql( + &frontend, + r#"CREATE TABLE f04_upper_min_logical ( + "A" STRING, + marker STRING, + val DOUBLE, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY ("A", marker) + ) ENGINE = metric WITH (on_physical_table = 'f04_upper_min_phy')"#, + ) + .await; + execute_sql( + &frontend, + "INSERT INTO f04_upper_min_logical (\"A\", marker, val, ts) VALUES \ + ('h0', 'a', 1, 1000), ('h0', 'b', 2, 1000), \ + ('h1', 'a', 3, 1000), ('h1', 'b', 4, 1000)", + ) + .await; + + let expected = "\ ++----+---+ +| A | n | ++----+---+ +| h0 | 2 | +| h1 | 2 | ++----+---+"; + // `NULLS FIRST` spells out the ordering the scan used to claim and is replayed next to + // the default `NULLS LAST` order, in memory and after flushing the physical table. + let queries = [ + "SELECT \"A\", count(*) AS n FROM f04_upper_min_logical GROUP BY \"A\" ORDER BY \"A\"", + "SELECT \"A\", count(*) AS n FROM f04_upper_min_logical GROUP BY \"A\" ORDER BY \"A\" NULLS FIRST", + ]; + + let mut mismatches = Vec::new(); + for flushed in [false, true] { + if flushed { + execute_sql(&frontend, "ADMIN FLUSH_TABLE('f04_upper_min_phy')").await; + } + for parallelism in [1, 4] { + for query in queries { + let actual = execute_sql_with_query_parallelism(&frontend, query, parallelism) + .await + .data + .pretty_print() + .await; + if actual.trim() != expected.trim() { + mismatches.push(format!( + "query `{query}` with query_parallelism={parallelism} flushed={flushed} \ + returned:\n{actual}" + )); + } + } + } + } + assert!( + mismatches.is_empty(), + "small logical metric grouping must return [h0, 2] and [h1, 2]:\n{}", + mismatches.join("\n---\n") + ); +} + #[apply(both_instances_cases)] async fn test_partitioned_merge_scan_left_join_repartitions(instance: Arc) { let frontend = instance.frontend(); From f41e86f43a7c3f78d12782df0d65a0a23e04ce4c Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 14 Oct 2026 04:11:53 +0000 Subject: [PATCH 36/46] feat(meta-client): forward actor and event context for region follower RPCs (#9568) The manage-region-follower chain submitted its procedures with an empty executor context: `ProcedureServiceHandler::manage_region_follower` accepted no query context, the operator built an `ExecutorContext::default()`, and `MetaClient` dropped the context before delegating to `RegionFollowerClient`. The submitting user (actor) and the procedure event metadata therefore never reached the submitted procedures. - `ProcedureServiceHandler::manage_region_follower` now takes a `QueryContextRef`, consistent with `migrate_region`, `gc_regions` and `gc_table`. - `ProcedureServiceOperator` derives the `ExecutorContext` (actor and `TriggerReason::Manual` event input) from the query context. - The `RegionFollowerClient` trait methods now accept `&ExecutorContext` and `MetaClient` forwards it, so client implementations can transport the actor and event context on the wire. - `procedure_event_context`/`procedure_actor` are now public so client implementations can reuse the trusted protocol derivation. --- src/common/function/src/admin/gc.rs | 3 +- src/common/function/src/admin/purge_table.rs | 6 ++- src/common/function/src/handlers.rs | 6 ++- src/common/function/src/state.rs | 1 + src/meta-client/src/client.rs | 44 +++++++++++----- src/meta-client/src/client/procedure.rs | 4 +- src/operator/src/procedure.rs | 53 ++++++++++++++++++-- 7 files changed, 97 insertions(+), 20 deletions(-) diff --git a/src/common/function/src/admin/gc.rs b/src/common/function/src/admin/gc.rs index 0090140a460..63b249cb54d 100644 --- a/src/common/function/src/admin/gc.rs +++ b/src/common/function/src/admin/gc.rs @@ -191,7 +191,7 @@ mod tests { use common_meta::rpc::procedure::{ GcResponse, ManageRegionFollowerRequest, MigrateRegionRequest, ProcedureStateResponse, }; - use session::context::QueryContext; + use session::context::{QueryContext, QueryContextRef}; use super::*; use crate::handlers::ProcedureServiceHandler; @@ -298,6 +298,7 @@ mod tests { async fn manage_region_follower( &self, + _query_ctx: QueryContextRef, _request: ManageRegionFollowerRequest, ) -> Result<()> { unreachable!() diff --git a/src/common/function/src/admin/purge_table.rs b/src/common/function/src/admin/purge_table.rs index 7e5fcd84371..88ed26ae9f6 100644 --- a/src/common/function/src/admin/purge_table.rs +++ b/src/common/function/src/admin/purge_table.rs @@ -157,7 +157,11 @@ mod tests { async fn query_procedure_state(&self, _: &str) -> Result { unreachable!() } - async fn manage_region_follower(&self, _: ManageRegionFollowerRequest) -> Result<()> { + async fn manage_region_follower( + &self, + _query_ctx: QueryContextRef, + _request: ManageRegionFollowerRequest, + ) -> Result<()> { unreachable!() } fn catalog_manager(&self) -> &CatalogManagerRef { diff --git a/src/common/function/src/handlers.rs b/src/common/function/src/handlers.rs index 6ee8064d6c7..9dafd853a84 100644 --- a/src/common/function/src/handlers.rs +++ b/src/common/function/src/handlers.rs @@ -105,7 +105,11 @@ pub trait ProcedureServiceHandler: Send + Sync { async fn query_procedure_state(&self, pid: &str) -> Result; /// Manage a region follower to a region. - async fn manage_region_follower(&self, request: ManageRegionFollowerRequest) -> Result<()>; + async fn manage_region_follower( + &self, + query_ctx: QueryContextRef, + request: ManageRegionFollowerRequest, + ) -> Result<()>; /// Get the catalog manager fn catalog_manager(&self) -> &CatalogManagerRef; diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index f60b37dd36d..bf2405930c3 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -92,6 +92,7 @@ impl FunctionState { async fn manage_region_follower( &self, + _query_ctx: QueryContextRef, _request: ManageRegionFollowerRequest, ) -> Result<()> { Ok(()) diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index e0932af5a69..44b6b21dd23 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -13,12 +13,11 @@ // limitations under the License. mod ask_leader; +mod cluster; mod config; pub mod heartbeat; mod load_balance; mod procedure; - -mod cluster; mod store; mod util; @@ -69,7 +68,8 @@ use common_time::util::DefaultSystemTimer; use config::Client as ConfigClient; use futures::TryStreamExt; use heartbeat::{Client as HeartbeatClient, HeartbeatConfig}; -use procedure::{Client as ProcedureClient, procedure_actor, procedure_event_context}; +use procedure::Client as ProcedureClient; +pub use procedure::{procedure_actor, procedure_event_context}; use serde::de::DeserializeOwned; use snafu::{OptionExt, ResultExt}; use store::Client as StoreClient; @@ -318,15 +318,35 @@ impl MetaClient { pub type RegionFollowerClientRef = Arc; /// A trait for clients that can manage region followers. +/// +/// The executor context carries the actor and the event input of the +/// submission, so implementations can forward them to the metasrv and the +/// region follower procedures record who submitted the operation. #[async_trait::async_trait] pub trait RegionFollowerClient: Sync + Send + Debug { - async fn add_region_follower(&self, request: AddRegionFollowerRequest) -> Result<()>; + async fn add_region_follower( + &self, + ctx: &ExecutorContext, + request: AddRegionFollowerRequest, + ) -> Result<()>; - async fn remove_region_follower(&self, request: RemoveRegionFollowerRequest) -> Result<()>; + async fn remove_region_follower( + &self, + ctx: &ExecutorContext, + request: RemoveRegionFollowerRequest, + ) -> Result<()>; - async fn add_table_follower(&self, request: AddTableFollowerRequest) -> Result<()>; + async fn add_table_follower( + &self, + ctx: &ExecutorContext, + request: AddTableFollowerRequest, + ) -> Result<()>; - async fn remove_table_follower(&self, request: RemoveTableFollowerRequest) -> Result<()>; + async fn remove_table_follower( + &self, + ctx: &ExecutorContext, + request: RemoveTableFollowerRequest, + ) -> Result<()>; async fn start(&self, urls: &[&str]) -> Result<()>; @@ -370,31 +390,31 @@ impl ProcedureExecutor for MetaClient { async fn manage_region_follower( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: ManageRegionFollowerRequest, ) -> MetaResult<()> { if let Some(region_follower) = &self.region_follower { match request { ManageRegionFollowerRequest::AddRegionFollower(add_region_follower_request) => { region_follower - .add_region_follower(add_region_follower_request) + .add_region_follower(ctx, add_region_follower_request) .await } ManageRegionFollowerRequest::RemoveRegionFollower( remove_region_follower_request, ) => { region_follower - .remove_region_follower(remove_region_follower_request) + .remove_region_follower(ctx, remove_region_follower_request) .await } ManageRegionFollowerRequest::AddTableFollower(add_table_follower_request) => { region_follower - .add_table_follower(add_table_follower_request) + .add_table_follower(ctx, add_table_follower_request) .await } ManageRegionFollowerRequest::RemoveTableFollower(remove_table_follower_request) => { region_follower - .remove_table_follower(remove_table_follower_request) + .remove_table_follower(ctx, remove_table_follower_request) .await } } diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index 9c32b4fc880..222a525dcbb 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -49,7 +49,7 @@ use crate::error::Result; /// /// The caller can only supply reason/extensions. Protocol is derived here from /// the trusted, typed query channel held locally in the executor context. -pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option { +pub fn procedure_event_context(context: &ExecutorContext) -> Option { context.event_input.as_ref().map(|input| { let mut event_context = ProcedureEventContext::from(input); event_context.protocol = context @@ -62,7 +62,7 @@ pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option Option { +pub fn procedure_actor(context: &ExecutorContext) -> Option { context.actor.as_ref().map(|username| ProcedureActor { username: username.clone(), }) diff --git a/src/operator/src/procedure.rs b/src/operator/src/procedure.rs index 03fb4f909c9..bf9ae106080 100644 --- a/src/operator/src/procedure.rs +++ b/src/operator/src/procedure.rs @@ -124,10 +124,12 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { async fn manage_region_follower( &self, + query_ctx: QueryContextRef, request: ManageRegionFollowerRequest, ) -> QueryResult<()> { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); self.procedure_executor - .manage_region_follower(&ExecutorContext::default(), request) + .manage_region_follower(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu) @@ -178,8 +180,11 @@ mod tests { ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; - use common_meta::rpc::procedure::{MigrateRegionResponse, ProcedureStateResponse}; - use session::context::QueryContextBuilder; + use common_meta::rpc::procedure::{ + AddRegionFollowerRequest, ManageRegionFollowerRequest, MigrateRegionResponse, + ProcedureStateResponse, + }; + use session::context::{QueryContextBuilder, QueryContextRef}; use table::table_name::TableName; use super::*; @@ -229,6 +234,14 @@ mod tests { ) -> common_meta::error::Result { unimplemented!() } + async fn manage_region_follower( + &self, + context: &ExecutorContext, + _: ManageRegionFollowerRequest, + ) -> common_meta::error::Result<()> { + self.contexts.lock().unwrap().push(context.clone()); + Ok(()) + } async fn list_procedures( &self, _: &ExecutorContext, @@ -303,6 +316,40 @@ mod tests { ); } + #[tokio::test] + async fn test_manage_region_follower_forwards_query_context() { + let manager = Arc::new(TableMetadataManager::new(Arc::new( + MemoryKvBackend::default(), + ))); + let executor = Arc::new(RecordingProcedureExecutor::default()); + let operator = ProcedureServiceOperator::new( + executor.clone(), + catalog::memory::MemoryCatalogManager::new(), + manager, + ); + let query_ctx: QueryContextRef = QueryContextBuilder::default().build().into(); + + operator + .manage_region_follower( + query_ctx, + ManageRegionFollowerRequest::AddRegionFollower(AddRegionFollowerRequest { + region_id: 1, + peer_id: 2, + }), + ) + .await + .unwrap(); + + let contexts = executor.contexts.lock().unwrap(); + assert_eq!(contexts.len(), 1); + // The default query context runs as the default user. + assert_eq!(contexts[0].actor.as_deref(), Some("greptime")); + assert_eq!( + contexts[0].event_input.as_ref().map(|input| input.reason), + Some(TriggerReason::Manual) + ); + } + #[tokio::test] async fn test_purge_table_missing_tombstone_is_table_not_found() { let manager = Arc::new(TableMetadataManager::new(Arc::new( From 0e40281262c243ddbe908456e892850aeb435bc0 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 14 Oct 2026 04:20:39 +0000 Subject: [PATCH 37/46] fix(promql): stop claiming scalar calculate preserves order (#9523) * test: cover scalar calculate order contract Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): stop claiming scalar calculate preserves order Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(test): describe scalar ordering regression as pre-fix behavior Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/extension_plan/scalar_calculate.rs | 2 +- tests-integration/src/tests/promql_test.rs | 444 +++++++++++++++++- 2 files changed, 444 insertions(+), 2 deletions(-) diff --git a/src/promql/src/extension_plan/scalar_calculate.rs b/src/promql/src/extension_plan/scalar_calculate.rs index 4ad188a1d4f..0b330dfc00f 100644 --- a/src/promql/src/extension_plan/scalar_calculate.rs +++ b/src/promql/src/extension_plan/scalar_calculate.rs @@ -413,7 +413,7 @@ impl ExecutionPlan for ScalarCalculateExec { } fn maintains_input_order(&self) -> Vec { - vec![true; self.children().len()] + vec![false] } fn input_distribution_requirements(&self) -> InputDistributionRequirements { diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 8cf27d165e9..7796be4d15b 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -22,7 +22,9 @@ use api::v1::{ use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_query::{Output, OutputData}; use common_recordbatch::util::collect_batches; -use datatypes::arrow::array::{Float64Array, Int64Array, StringArray, TimestampMillisecondArray}; +use datatypes::arrow::array::{ + Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, +}; use datatypes::arrow::compute::cast; use datatypes::arrow::datatypes::DataType; use frontend::instance::Instance; @@ -1310,3 +1312,443 @@ async fn promql_stale_marker_excludes_series_across_flushes(instance: Arc)] = &[ + (3000, 2.0, None), + (2000, 3.0, Some(2.0)), + (1000, 1.0, Some(3.0)), + (0, 4.0, Some(1.0)), +]; + +/// Output column of a `scalar()` TQL query; it has to be quoted in SQL. +const F06_SCALAR_VALUE_COLUMN: &str = "scalar(val)"; + +/// Ordinary Mito fixture. `host` stays the first column on purpose: the wrong +/// push-down reuses output index `0` (`ts`) as input column `0` (`host`). +const F06_ORDINARY_TABLE: &str = "f06_scalar_plain"; + +const F06_ORDINARY_CREATE: &str = r#" +CREATE TABLE f06_scalar_plain ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX, + val DOUBLE +); +"#; + +/// Metric fixture; the physical table only carries the logical table, all +/// `INSERT`s, reads and TQL queries go to the logical `f06_scalar_metric`. +/// `ScalarCalculate` sees its input columns as `(val, host, ts)`, so the wrong +/// push-down turns `ts@0` into `val@0`. +const F06_METRIC_TABLE: &str = "f06_scalar_metric"; + +const F06_METRIC_CREATE: &str = r#" +CREATE TABLE f06_scalar_phy ( + ts TIMESTAMP TIME INDEX, + val DOUBLE +) ENGINE=metric WITH ("physical_metric_table" = ""); + +CREATE TABLE f06_scalar_metric ( + host STRING, + ts TIMESTAMP TIME INDEX, + val DOUBLE, + PRIMARY KEY (host) +) ENGINE=metric WITH (on_physical_table = 'f06_scalar_phy'); +"#; + +fn f06_insert_sql(table: &str) -> String { + format!( + "INSERT INTO {table} (host, ts, val) VALUES \ + ('h', 0, 4), ('h', 1000, 1), ('h', 2000, 3), ('h', 3000, 2)" + ) +} + +fn f06_fixture_sql(table: &str) -> String { + format!("SELECT ts, val FROM {table} ORDER BY ts") +} + +/// `ORDER BY ts DESC` over `scalar()`; this is the repro of the wrong sort +/// push-down. +fn f06_scalar_order_by_ts_desc_sql(table: &str) -> String { + format!("WITH q AS (TQL EVAL (0, 3, '1s') scalar({table})) SELECT * FROM q ORDER BY ts DESC") +} + +/// Same as [`f06_scalar_order_by_ts_desc_sql`] but with a fetch, which keeps the +/// correct `TopK` above `ScalarCalculate`. +fn f06_scalar_order_by_ts_desc_limit_3_sql(table: &str) -> String { + format!( + "WITH q AS (TQL EVAL (0, 3, '1s') scalar({table})) SELECT * FROM q \ + ORDER BY ts DESC LIMIT 3" + ) +} + +/// `ORDER BY` on the scalar value instead of on the time index. +fn f06_scalar_order_by_value_asc_sql(table: &str) -> String { + format!( + r#"WITH q AS (TQL EVAL (0, 3, '1s') scalar({table})) SELECT * FROM q ORDER BY "scalar(val)" ASC"# + ) +} + +/// A SQL window over `scalar()`; the window requires the `ts` order that the +/// wrong push-down removes. +fn f06_scalar_lag_order_by_ts_desc_sql(table: &str) -> String { + format!( + r#"WITH q AS (TQL EVAL (0, 3, '1s') scalar({table})) SELECT ts, "scalar(val)" AS val, LAG("scalar(val)") OVER (ORDER BY ts DESC) AS prev FROM q ORDER BY ts DESC"# + ) +} + +/// Selector control: the `ts` order of a bare selector must stay correct. +fn f06_selector_order_by_ts_desc_sql(table: &str) -> String { + format!("WITH q AS (TQL EVAL (0, 3, '1s') {table}) SELECT ts, val FROM q ORDER BY ts DESC") +} + +/// Creates and fills one F06 fixture and also guards the fixture itself: an +/// ordering regression is only meaningful if the stored rows read back in `ts` +/// order. +async fn f06_setup(instance: &Arc, create: &str, table: &str) { + let insert = f06_insert_sql(table); + for sql in [create, insert.as_str()] { + for result in instance.do_query(sql, QueryContext::arc()).await { + if let Err(err) = result { + panic!("F06: setup statement failed: {err}\n{sql}"); + } + } + } + + let sql = f06_fixture_sql(table); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(instance, &sql, parallelism, "val").await; + assert_eq!( + rows.as_slice(), + F06_FIXTURE_ROWS, + "F06: fixture `{table}` must read back in `ts` order (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Runs one F06 query with a query parallelism hint and returns its rows in +/// result order as `(ts_millis, value_column)`; the rows are never sorted here. +async fn f06_ordered_rows( + instance: &Arc, + sql: &str, + parallelism: usize, + value_column: &str, +) -> Vec<(i64, f64)> { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + let output = instance + .do_query(sql, Arc::new(query_ctx)) + .await + .remove(0) + .unwrap_or_else(|err| panic!("F06: query failed: {err}\n{sql}")); + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("F06: expected a query result"), + }; + + let mut rows = Vec::new(); + for batch in batches.iter() { + let ts = cast( + batch.column_by_name("ts").unwrap().as_ref(), + &DataType::Int64, + ) + .unwrap(); + let ts = ts.as_any().downcast_ref::().unwrap(); + let values = batch.column_by_name(value_column).unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + for row in 0..batch.num_rows() { + rows.push((ts.value(row), values.value(row))); + } + } + rows +} + +/// Runs the F06 `LAG` query and returns its rows in result order as +/// `(ts_millis, val, prev)`; `prev` is `None` for the first row of the window. +/// The rows are never sorted here. +async fn f06_ordered_lag_rows( + instance: &Arc, + sql: &str, + parallelism: usize, +) -> Vec<(i64, f64, Option)> { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + let output = instance + .do_query(sql, Arc::new(query_ctx)) + .await + .remove(0) + .unwrap_or_else(|err| panic!("F06: query failed: {err}\n{sql}")); + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("F06: expected a query result"), + }; + + let mut rows = Vec::new(); + for batch in batches.iter() { + let ts = cast( + batch.column_by_name("ts").unwrap().as_ref(), + &DataType::Int64, + ) + .unwrap(); + let ts = ts.as_any().downcast_ref::().unwrap(); + let values = batch.column_by_name("val").unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + let prevs = batch.column_by_name("prev").unwrap(); + let prevs = prevs.as_any().downcast_ref::().unwrap(); + for row in 0..batch.num_rows() { + let prev = if prevs.is_null(row) { + None + } else { + Some(prevs.value(row)) + }; + rows.push((ts.value(row), values.value(row), prev)); + } + } + rows +} + +/// Repro of the wrong sort push-down: `ORDER BY ts DESC` over `scalar()` is +/// pushed into the TQL input as `host@0 DESC` and silently returns the wrong row +/// order (the fixture order, not `ts`-descending). +#[apply(both_instances_cases)] +async fn f06_scalar_ordinary_table_order_by_ts_desc(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_ORDINARY_CREATE, F06_ORDINARY_TABLE).await; + + let sql = f06_scalar_order_by_ts_desc_sql(F06_ORDINARY_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, F06_SCALAR_VALUE_COLUMN).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_TS_DESC, + "F06: `ORDER BY ts DESC` over scalar() must return `ts`-descending rows \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Same repro as [`f06_scalar_ordinary_table_order_by_ts_desc`] for a metric +/// logical table, where the wrong push-down uses `val@0 DESC`. +#[apply(both_instances_cases)] +async fn f06_scalar_metric_table_order_by_ts_desc(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_METRIC_CREATE, F06_METRIC_TABLE).await; + + let sql = f06_scalar_order_by_ts_desc_sql(F06_METRIC_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, F06_SCALAR_VALUE_COLUMN).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_TS_DESC, + "F06: `ORDER BY ts DESC` over scalar() must return `ts`-descending rows \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Repro: ordering by the scalar value is broken as well, the rows come back in +/// the fixture order instead of `1.0, 2.0, 3.0, 4.0`. +#[apply(both_instances_cases)] +async fn f06_scalar_ordinary_table_order_by_value_asc(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_ORDINARY_CREATE, F06_ORDINARY_TABLE).await; + + let sql = f06_scalar_order_by_value_asc_sql(F06_ORDINARY_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, F06_SCALAR_VALUE_COLUMN).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_VALUE_ASC, + "F06: `ORDER BY \"scalar(val)\" ASC` must order by the scalar value \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Same as [`f06_scalar_ordinary_table_order_by_value_asc`] for a metric logical +/// table. +#[apply(both_instances_cases)] +async fn f06_scalar_metric_table_order_by_value_asc(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_METRIC_CREATE, F06_METRIC_TABLE).await; + + let sql = f06_scalar_order_by_value_asc_sql(F06_METRIC_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, F06_SCALAR_VALUE_COLUMN).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_VALUE_ASC, + "F06: `ORDER BY \"scalar(val)\" ASC` must order by the scalar value \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Control that must stay correct: the fetch keeps the `TopK` above +/// `ScalarCalculate`, so the `ts` order is right although the plain +/// `ORDER BY ts DESC` repro is wrong. +#[apply(both_instances_cases)] +async fn f06_scalar_ordinary_table_order_by_ts_desc_limit_3_control( + instance: Arc, +) { + let instance = instance.frontend(); + f06_setup(&instance, F06_ORDINARY_CREATE, F06_ORDINARY_TABLE).await; + + let sql = f06_scalar_order_by_ts_desc_limit_3_sql(F06_ORDINARY_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, F06_SCALAR_VALUE_COLUMN).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_TS_DESC_LIMIT_3, + "F06: `ORDER BY ts DESC LIMIT 3` over scalar() must keep returning the newest rows \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Metric logical table variant of +/// [`f06_scalar_ordinary_table_order_by_ts_desc_limit_3_control`]. +#[apply(both_instances_cases)] +async fn f06_scalar_metric_table_order_by_ts_desc_limit_3_control(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_METRIC_CREATE, F06_METRIC_TABLE).await; + + let sql = f06_scalar_order_by_ts_desc_limit_3_sql(F06_METRIC_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, F06_SCALAR_VALUE_COLUMN).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_TS_DESC_LIMIT_3, + "F06: `ORDER BY ts DESC LIMIT 3` over scalar() must keep returning the newest rows \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Selector control that must stay correct: a bare selector does not go through +/// `ScalarCalculate`, so `ORDER BY ts DESC` over it is right. +#[apply(both_instances_cases)] +async fn f06_selector_ordinary_table_order_by_ts_desc_control(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_ORDINARY_CREATE, F06_ORDINARY_TABLE).await; + + let sql = f06_selector_order_by_ts_desc_sql(F06_ORDINARY_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, "val").await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_TS_DESC, + "F06: a bare selector must keep its `ts` order under `ORDER BY ts DESC` \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Metric logical table variant of +/// [`f06_selector_ordinary_table_order_by_ts_desc_control`]. +#[apply(both_instances_cases)] +async fn f06_selector_metric_table_order_by_ts_desc_control(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_METRIC_CREATE, F06_METRIC_TABLE).await; + + let sql = f06_selector_order_by_ts_desc_sql(F06_METRIC_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_rows(&instance, &sql, parallelism, "val").await; + assert_eq!( + rows.as_slice(), + F06_ROWS_BY_TS_DESC, + "F06: a bare selector must keep its `ts` order under `ORDER BY ts DESC` \ + (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Repro of the planning failure that the wrong sort push-down used to cause: +/// the window needs `ts@0 DESC`, the pushed-down sort only provided +/// `host@0 DESC`, so this query failed with code 3000 while `ScalarCalculate` +/// claimed `maintains_input_order`. +/// +/// Now that the node stops claiming the input order, the window is planned +/// above the correct sort and the query must return exactly +/// `(3000, 2.0, NULL)`, `(2000, 3.0, 2.0)`, `(1000, 1.0, 3.0)`, +/// `(0, 4.0, 1.0)`. +#[apply(both_instances_cases)] +async fn f06_scalar_ordinary_table_lag_order_by_ts_desc(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_ORDINARY_CREATE, F06_ORDINARY_TABLE).await; + + let sql = f06_scalar_lag_order_by_ts_desc_sql(F06_ORDINARY_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_lag_rows(&instance, &sql, parallelism).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_LAG_BY_TS_DESC, + "F06: `LAG(...) OVER (ORDER BY ts DESC)` over scalar() must return `ts`-descending \ + rows with the lagged scalar values (parallelism={parallelism})\n{sql}" + ); + } +} + +/// Metric logical table variant of +/// [`f06_scalar_ordinary_table_lag_order_by_ts_desc`]. +#[apply(both_instances_cases)] +async fn f06_scalar_metric_table_lag_order_by_ts_desc(instance: Arc) { + let instance = instance.frontend(); + f06_setup(&instance, F06_METRIC_CREATE, F06_METRIC_TABLE).await; + + let sql = f06_scalar_lag_order_by_ts_desc_sql(F06_METRIC_TABLE); + for parallelism in [1, 4] { + let rows = f06_ordered_lag_rows(&instance, &sql, parallelism).await; + assert_eq!( + rows.as_slice(), + F06_ROWS_LAG_BY_TS_DESC, + "F06: `LAG(...) OVER (ORDER BY ts DESC)` over scalar() must return `ts`-descending \ + rows with the lagged scalar values (parallelism={parallelism})\n{sql}" + ); + } +} From a08e5bccebea274dfe6254912ad75be9e3b09de5 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 14 Oct 2026 06:10:26 +0000 Subject: [PATCH 38/46] fix(log-store): keep the enqueued backlog on permanent object store errors (#9558) * fix(log-store): keep the enqueued backlog on permanent object store errors On object storage the WAL's object store sits behind OpenDAL's retry layer, which marks every error it returns as persistent. The enqueued mode classified create failures by that status, so a 403 or a missing bucket was repeated every 100 ms on every create in flight, logged nothing, and made appends hang once the sealed batches were full. Object store errors are now classified by kind: denied access, an invalid configuration (how OpenDAL reports a missing S3 bucket) and a write to a missing location are non-retriable. In the enqueued mode such an error makes the store failing instead of poisoning it: the acknowledged backlog is kept under its sequences and bytes, its lowest pending batch is attempted alone with a delay that doubles from 100 ms to 30 s, appends are rejected at once with a retryable error that carries the last object store error, and the first create that succeeds ends the state. Stop uploads the backlog at once and reports a failure as before. The state is logged on entry, at most once a minute while it lasts and on exit, and exported as the `failing` gauge and the `rejected_appends_total` counter. Conflicts, encoding and catalog errors still poison the store; the durable mode is unchanged. An append held back because the sealed batches are full is now counted and timed by the stall metrics in both modes. Signed-off-by: jeremyhi * docs(rfc): describe how the enqueued object store WAL handles refused creates Describe the enqueued failure contract as the store now implements it: non-retriable object store errors are told apart by kind because OpenDAL's retry layer under the store marks every error persistent, a failing store keeps its backlog and attempts it with backoff while it rejects appends, and stop uploads the backlog at once. Add the failure matrix row, the failing gauge and the rejected appends counter, the sealed-batch admission limit and the stall metrics that now count it, and the fault matrix case counts. Signed-off-by: jeremyhi * fix(log-store): time the failing attempts where they start and fix two statements The test of the attempts of a failing store now records when each refused write starts and checks the gaps between them, instead of sampling a counter after a sleep of the test task. The delay comment and the RFC describe 30 seconds as the longest delay between two attempts, not as a bound on when appends resume, and the RFC states that the next append is held once seven sealed batches wait. Signed-off-by: jeremyhi * test(log-store): drop the wait_until helper that main now defines Signed-off-by: jeremyhi * refactor(log-store): name errors that need an operator and trim the RFC contract Rename is_non_retriable to requires_intervention: such a create is still repeated, only slowly, until someone fixes the credentials, configuration or bucket. Name the backoff test after what it asserts, cut the test of the lowest pending batch to two failed rounds since the cadence is covered by the backoff test, and shorten the RFC failure contract to the contract and its rationale, leaving error mappings, request counts and logging cadence to the code comments. Signed-off-by: jeremyhi * docs(log-store): tie the slow repeat to errors requiring intervention Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- docs/rfcs/2026-09-06-object-store-wal.md | 21 +- src/log-store/src/error.rs | 13 +- src/log-store/src/metrics.rs | 32 +- src/log-store/src/object_store_wal/store.rs | 1232 ++++++++++++++--- .../engine/object_store_wal_recovery_test.rs | 110 +- 5 files changed, 1232 insertions(+), 176 deletions(-) diff --git a/docs/rfcs/2026-09-06-object-store-wal.md b/docs/rfcs/2026-09-06-object-store-wal.md index 26970d166f8..b49f97c330c 100644 --- a/docs/rfcs/2026-09-06-object-store-wal.md +++ b/docs/rfcs/2026-09-06-object-store-wal.md @@ -121,13 +121,17 @@ The store then selects the history it replays from the headers. An object's link Before it admits any append, the store writes a start object: an object without segments at the next sequence, which is above every listed object, orphans included, linked to the tip, or to nothing on an empty prefix. The instance's epoch is the sequence of its start object plus one, and every object the instance writes carries it. The conditional create assigns each sequence to exactly one object, so no two instances share an epoch; and the start sequence is above every object present at open, including the start objects of every earlier instance, so the epoch of every later instance is greater than that of any object an earlier instance ever writes, including one whose create lands after that instance stopped or crashed, and such an object never becomes the tip. Only once the start object is created does the store accept appends, and the first batch links to it. If the sequence turns out to be taken, the store reads the object there: an object of an earlier epoch is stale, and the start object moves to the next sequence, taking the epoch that sequence gives; an object of the same or a later epoch means another writer, and construction fails with a conflict error that is not retried; an object with identical content, which another open of the same prefix has just written, fails construction with a retryable error, so that the next attempt lists it and starts above it. A create whose outcome is unknown fails construction with a retryable error too: an object it may have left is listed by the next attempt, which therefore starts above it. Two assumptions remain. Listing the prefix is strongly consistent, so an open sees every start object that was durable before it listed; this holds for S3, GCS, Azure Blob Storage and MinIO, and an S3-compatible store without it is not supported. And one process at a time opens a prefix: two processes opening it concurrently either contend for the same start sequence, where the conditional create lets one through and fails the other's construction, or, if one of them listed a late object the other did not, start under different sequences and epochs and both proceed until their sequences meet, where the writer that meets the other's object follows the conflict rules under *Writing* and the *Failure matrix*: in the `durable` mode an object of an earlier epoch fails its batch and one of a later epoch poisons it, and in the `enqueued` mode any conflict poisons it, so which writer is poisoned depends on who observes the collision, not on the epochs alone; the store does not fence concurrent writers beyond that, and in cluster mode the generation in the key prefix does. -**Writing.** One background actor per store admits appended entries into a single open batch and assigns entry ids as described under *Entry ids* above. The batch is sealed when its estimated size reaches `max_batch_bytes` or when `flush_interval` elapses with entries in it; a timer tick on an empty batch creates nothing. Sealing encodes the batch under the next sequence number and hands it to a pipelined uploader: the actor keeps admitting entries into the next batch while uploads are in flight, and up to four conditional creates run at a time, started in sequence order. Every sealed batch is encoded with the store's epoch and a link to its pipeline predecessor: the last batch sealed before it that is not indexed yet, or, when there is none, the last indexed object, which after opening is the start object. A link to the last indexed object while an earlier batch is still uploading would let a later chain skip that batch once it is acknowledged. Batches are indexed in sequence order in both acknowledgement modes: a created object is not indexed until every earlier batch is. The rest of this paragraph is `durable` mode behaviour; the `enqueued` mode, whose appends have already returned, is described under *Acknowledgement*. In the `durable` mode a batch is acknowledged when it is indexed, so a batch whose object is durable is not acknowledged until every earlier batch is durable too, and the acknowledged prefix of a region's history never has a missing predecessor. When a create fails transiently, every batch that is not indexed fails at once, including batches whose objects were already created and batches whose creates are still in flight. Their sequence numbers are never handed out again, the next batch links to the last indexed object, and the store keeps serving. A create still in flight completes later and finds no batch to acknowledge, so whatever it stored is an orphan. The next object acknowledged links past every failed batch, and from then on recovery replays none of them, whether their objects exist, land later or never do. A create that meets a different object reads that object's header: an object of an earlier epoch is stale and can never be the tip, so the batch fails as for a transient error and the store keeps serving; an object of the same or a later epoch belongs to another writer on the prefix, and the store poisons itself. +**Writing.** One background actor per store admits appended entries into a single open batch and assigns entry ids as described under *Entry ids* above. The batch is sealed when its estimated size reaches `max_batch_bytes` or when `flush_interval` elapses with entries in it; a timer tick on an empty batch creates nothing. Sealing encodes the batch under the next sequence number and hands it to a pipelined uploader: the actor keeps admitting entries into the next batch while uploads are in flight, and up to four conditional creates run at a time, started in sequence order. At most eight sealed batches wait to be indexed: once seven do, which leaves no room for the two one admission can seal, the next append is held back in both acknowledgement modes until a batch is indexed, and later appends wait behind it. Every sealed batch is encoded with the store's epoch and a link to its pipeline predecessor: the last batch sealed before it that is not indexed yet, or, when there is none, the last indexed object, which after opening is the start object. A link to the last indexed object while an earlier batch is still uploading would let a later chain skip that batch once it is acknowledged. Batches are indexed in sequence order in both acknowledgement modes: a created object is not indexed until every earlier batch is. The rest of this paragraph is `durable` mode behaviour; the `enqueued` mode, whose appends have already returned, is described under *Acknowledgement*. In the `durable` mode a batch is acknowledged when it is indexed, so a batch whose object is durable is not acknowledged until every earlier batch is durable too, and the acknowledged prefix of a region's history never has a missing predecessor. When a create fails transiently, every batch that is not indexed fails at once, including batches whose objects were already created and batches whose creates are still in flight. Their sequence numbers are never handed out again, the next batch links to the last indexed object, and the store keeps serving. A create still in flight completes later and finds no batch to acknowledge, so whatever it stored is an orphan. The next object acknowledged links past every failed batch, and from then on recovery replays none of them, whether their objects exist, land later or never do. A create that meets a different object reads that object's header: an object of an earlier epoch is stale and can never be the tip, so the batch fails as for a transient error and the store keeps serving; an object of the same or a later epoch belongs to another writer on the prefix, and the store poisons itself. The guarantee is therefore about acknowledged history. An acknowledged entry is always replayed; an entry whose append failed is replayed only when nothing was acknowledged after it, the same outcome as a crash between an object's creation and its acknowledgement, which the engine already tolerates. The distinction matters because Mito does not advance a region's committed sequence when an append fails and reuses those row sequences in its next write, so replaying a failed batch next to a later acknowledged one could let the unacknowledged rows win deduplication. **Acknowledgement.** `ack_mode` selects when `append_batch` returns. `durable`, the default, returns only after the object holding the entries is durable and indexed, so callers never see an entry id for an entry that is not durable; a write-ahead log is expected to mean durability on return, which is why this stays the default. `enqueued` returns as soon as the entries are admitted, with their entry ids already assigned, and creates the object in the background. In that mode the recovery point objective is the unpersisted backlog. `max_unpersisted_bytes` and `max_unpersisted_age` are admission thresholds: when the open batch plus the sealed batches that are not durable yet reach either, new appends are held back in arrival order until an upload completes, and the store seals the open batch if no create is in flight so that there is an upload to wait for; nothing is dropped and the backlog cannot grow without limit. They are not a hard bound on what a crash can lose: during an outage the writes already acknowledged stay in memory for as long as the outage lasts, however old they get, and a crash in that window loses them. -The failure contract of `enqueued` differs from the durable one because no caller remains to retry. A transient create failure is repeated by the store after a short delay instead of failing the batch. A permanent failure, meaning a conflicting object or an encoding or catalog error, poisons the store. Once `stop` has begun a transient failure is no longer repeated: the backlog is dropped and `stop` reports the loss as its error, since those entries were acknowledged but never became durable. Otherwise `stop` seals and uploads the remaining backlog before it returns. +The failure contract of `enqueued` differs from the durable one because no caller remains to retry: a create the object store did not confirm is repeated under the same sequence with the same bytes instead of failing the batch. A transient failure is repeated after 100 ms while appends are still admitted. A failure that requires intervention, such as denied access or a missing bucket, makes the store failing. The store tells the two apart by the kind of the object store error, not by its status, because OpenDAL's retry layer under the store marks every error it returns as persistent. + +A failing store keeps its acknowledged backlog instead of poisoning itself, since poisoning would drop those entries and need a restart for a cause that can be fixed in place. It attempts the lowest pending batch alone, with exponential backoff, and rejects every new append at once with a retryable error that carries the last object store error, so no append is admitted or waits behind the backlog. A durability wait for an entry in the backlog keeps waiting. The first create that succeeds ends the state: the backlog is uploaded and appends are admitted again. + +A conflicting object or an encoding or catalog error poisons the store. Once `stop` has begun a failed create is no longer repeated: the store attempts its whole backlog at once, failing or not, and any object store failure drops the backlog and `stop` reports the loss as its error, since those entries were acknowledged but never became durable. Otherwise `stop` seals and uploads the remaining backlog before it returns. The mode also needs a barrier that the durable mode gets for free. An entry id that Mito has been handed may not yet be durable, and a flush that persisted such an id as the manifest's `flushed_entry_id` would, after a crash before the object was created, leave a watermark that refers to an object which never existed, so that the sequence is reused on restart and later entries land below the watermark and are skipped by replay. The `LogStore` trait therefore has `wait_durable(provider, entry_id)`, with a default that returns at once, so that every existing provider keeps its own acknowledgement and persistence semantics, whatever they are; the durability guarantee this section describes belongs to the object store provider alone. The object store WAL answers it once the region is durable and indexed through `entry_id`; an id the store never handed out has nothing to wait for. Mito waits before it publishes a manifest watermark: a flush waits after its SSTs are written and before the manifest edit, a full truncate waits before it records the truncated entry id as the replay frontier, and discarding unflushed data waits before it records the discarded entry id as flushed. The wait in a flush observes the flush's cancellation, so a cancelled flush does not hold up the drop or truncate queued behind it; a full truncate and a discard of unflushed data still run their own wait before publishing their frontier. After a backlog was lost, a durability wait for an entry that is not durable fails rather than certifying it, whether the entry was in that backlog or not. In `enqueued` mode `obsolete` also clamps the watermark to the region's durable entry id, since an entry that is not durable yet is replayed after a crash and hiding it would skip it. A manifest watermark can then only ever name a durable entry, and the restart allocation rule under *Entry ids* keeps every new id above it. @@ -139,7 +143,8 @@ The mode also needs a barrier that the durable mode gets for free. An entry id t | transient I/O error, `durable` mode | consumed and never reused, nor are the sequences of the other failed batches | this and every batch that is not indexed fail, including those whose objects were created; a create still in flight completes as an orphan | healthy: the next batch links to the last indexed object | | an object of an earlier epoch holds the sequence, `durable` mode | consumed and never reused | as for a transient I/O error | healthy | | transient I/O error, `enqueued` mode | unchanged | already acknowledged | healthy: the create is repeated after a short delay | -| transient I/O error, `enqueued` mode after `stop` began | unchanged | already acknowledged | the backlog is dropped and `stop` reports the loss | +| I/O error that requires intervention, `enqueued` mode | unchanged | already acknowledged; later appends are rejected with a retryable error until a create succeeds | failing: the lowest pending create is repeated alone with backoff; the first create that succeeds ends it | +| any I/O error, `enqueued` mode after `stop` began | unchanged | already acknowledged | the backlog is dropped and `stop` reports the loss | | an object of the same or a later epoch holds the sequence, or any conflicting object in the `enqueued` mode | unchanged | every unacknowledged waiter fails, including those of batches whose objects were created but not yet indexed; an `enqueued` acknowledgement already returned cannot be revoked | poisoned | | encoding or catalog error | unchanged | every unacknowledged waiter fails, including those of batches whose objects were created but not yet indexed; an `enqueued` acknowledgement already returned cannot be revoked | poisoned | | create succeeds at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence, every later append fails | @@ -148,11 +153,11 @@ The mode also needs a barrier that the durable mode gets for free. An entry id t A store cannot be constructed on a prefix whose largest object already carries the maximum sequence. A poisoned store fails every `LogStore` operation except `stop` with the same terminal error; in particular `obsolete` can no longer move a watermark. In the `durable` mode transient errors are not retried inside the store, but a failed create is reconciled once: the store immediately reads the key back, accepts the batch if the object exists with identical bytes, treats it as a conflict if the bytes differ, and reports the transient error only when the key is absent or that read fails too. So a lost response alone does not fail the append. When the append does fail, neither the store nor Mito keeps the batch: it surfaces through Mito to the caller of the write, and a retry is a new write by the caller under a new sequence. If the create had in fact succeeded and both its response and the read-back were lost, or its request lands only later, the object is an orphan as soon as a later object is acknowledged. Retrying the batch itself inside the store in the `durable` mode would only hide the latency and would hold writes up for as long as an outage lasts; the `enqueued` mode repeats the create with the same bytes, so a create that succeeded without a response is accepted as an identical retry. -**Stopping.** `stop` is idempotent and awaits the actor. Once stop has begun, queued appends are not admitted, no timer or size triggered seal starts, and appends after stop, including empty ones, fail with the stopped error. In the `durable` mode the open batch and every sealed batch whose create has not started are dropped: their waiters receive the stopped error. A create that is already in flight runs to completion; if it succeeds while its batch is still pending, the batch is indexed in order and acknowledged. A transient failure after stop began fails every batch that is not indexed, as it does while the store runs, and those batches receive the stopped error; a create of such a batch that completes later is an orphan and acknowledges nothing. A create that finds a conflicting object of the same or a later epoch after stop began poisons the store as it would while the store runs, but the waiters of the batches that are not indexed receive the stopped error, the conflict is kept only as the terminal error, and `stop` returns successfully once the creates in flight have completed. In the `enqueued` mode the held-back appends receive the stopped error, the open batch is sealed and the whole backlog is uploaded before `stop` returns; a transient failure in that phase drops the backlog and `stop` returns the error, as described under *Acknowledgement*. A durability wait still pending when the actor exits receives the stopped error, and a delete already in flight runs to completion before `stop` returns. +**Stopping.** `stop` is idempotent and awaits the actor. Once stop has begun, queued appends are not admitted, no timer or size triggered seal starts, and appends after stop, including empty ones, fail with the stopped error. In the `durable` mode the open batch and every sealed batch whose create has not started are dropped: their waiters receive the stopped error. A create that is already in flight runs to completion; if it succeeds while its batch is still pending, the batch is indexed in order and acknowledged. A transient failure after stop began fails every batch that is not indexed, as it does while the store runs, and those batches receive the stopped error; a create of such a batch that completes later is an orphan and acknowledges nothing. A create that finds a conflicting object of the same or a later epoch after stop began poisons the store as it would while the store runs, but the waiters of the batches that are not indexed receive the stopped error, the conflict is kept only as the terminal error, and `stop` returns successfully once the creates in flight have completed. In the `enqueued` mode the held-back appends receive the stopped error, the open batch is sealed and the whole backlog is uploaded before `stop` returns, by a failing store too; any object store failure in that phase drops the backlog and `stop` returns the error, as described under *Acknowledgement*. A durability wait still pending when the actor exits receives the stopped error, and a delete already in flight runs to completion before `stop` returns. **Reading.** `read(provider, entry_id)` returns the region's entries with id at least `entry_id`, in order, fetching only that region's segments with range reads and decoding nothing else, and hides entries at or below the region's obsolete watermark. `obsolete` records the watermark in memory, clamped to the durable entry id in the `enqueued` mode; the watermark is re-established after a restart because Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, so no WAL-owned watermark state is persisted. `obsolete_all` hides the whole region. `latest_entry_id` is not affected by the watermark. Every method taking a provider requires `Provider::ObjectStore` with the store's prefix and the caller's region. -**Metrics.** The store exports metrics under `greptime_logstore_object_store_wal_`. The write path reports `seal_to_durable_seconds`, the time from sealing a batch until its object is durable, `append_ack_seconds`, the time an append waits for its acknowledgement in the `durable` mode, `object_bytes` and `object_entries`, the size and the entry count of every object sealed, and `created_objects_total`, the creates that succeeded, an identical retry included; a create is counted and timed whether or not its batch is still there to be indexed, since an object a poisoned store gave up on is as durable as any other. `indexed_objects` and `indexed_bytes` are the objects and the bytes the catalog holds, set by a recovery that succeeded and moved by every object indexed or unindexed afterwards, so a collection takes the object it deleted off both, and a recovery that is abandoned publishes neither. A recovery that succeeds reports `recovery_seconds` and the objects it read footers from in `recovered_objects_total`, both of the same recoveries, so one that is abandoned at a fetch or rejected by the catalog contributes to neither; a read stream reports `read_seconds` once it has been read to its end, and one that fails or is dropped part way reports nothing. The failure paths are counted: `create_conflicts_total` for a create that found a conflicting object, `create_failures_total` for one the object store did not confirm, `poisoned_total` for a store that recorded a terminal error, `skipped_segments_total` for a segment a read skipped, and `stalled_appends_total` with `stalled_append_seconds` for an append the `enqueued` backlog thresholds held back, timed until it is admitted or refused, a stop, a terminal error or the teardown of the actor releasing it included; the delete counters of a collection are described under *Entry ids*. The histogram buckets are chosen for the default flush interval of 100 ms and the default `max_batch_bytes` of 8 MiB. +**Metrics.** The store exports metrics under `greptime_logstore_object_store_wal_`. The write path reports `seal_to_durable_seconds`, the time from sealing a batch until its object is durable, `append_ack_seconds`, the time an append waits for its acknowledgement in the `durable` mode, `object_bytes` and `object_entries`, the size and the entry count of every object sealed, and `created_objects_total`, the creates that succeeded, an identical retry included; a create is counted and timed whether or not its batch is still there to be indexed, since an object a poisoned store gave up on is as durable as any other. `indexed_objects` and `indexed_bytes` are the objects and the bytes the catalog holds, set by a recovery that succeeded and moved by every object indexed or unindexed afterwards, so a collection takes the object it deleted off both, and a recovery that is abandoned publishes neither. A recovery that succeeds reports `recovery_seconds` and the objects it read footers from in `recovered_objects_total`, both of the same recoveries, so one that is abandoned at a fetch or rejected by the catalog contributes to neither; a read stream reports `read_seconds` once it has been read to its end, and one that fails or is dropped part way reports nothing. The failure paths are counted: `create_conflicts_total` for a create that found a conflicting object, `create_failures_total` for one the object store did not confirm, `poisoned_total` for a store that recorded a terminal error, `skipped_segments_total` for a segment a read skipped, and `stalled_appends_total` with `stalled_append_seconds` for an append held back because the sealed batches were full, in either mode, or by the `enqueued` backlog thresholds, timed until it is admitted or refused, a stop, a terminal error, a failing store or the teardown of the actor releasing it included. The gauge `failing` counts the stores whose backlog the object store refuses with an error that requires intervention, and `rejected_appends_total` the appends such a store rejected; the attempts of a failing store are counted in `create_failures_total`. The delete counters of a collection are described under *Entry ids*. The histogram buckets are chosen for the default flush interval of 100 ms and the default `max_batch_bytes` of 8 MiB. A `testing` cargo feature exposes hooks to wait for admitted appends, to seal the open batch, to hold and release creates and to make creates fail, so engine tests can flush deterministically and observe entries that are admitted but not durable. @@ -233,14 +238,14 @@ The reference implementation was validated on the following evidence; the tests | Recovery never silently loses data | Unit coverage of every corruption class at recovery and on read; footer-only recovery compared with a full decode; a footer longer than the tail window. | Structural damage fails construction of the store; a corrupted segment is skipped, recorded as a hole of its region and counted, or fails the read under `fail`; the footer-only catalog equals the full decode. | | A region cannot be silently pointed at another WAL | Reopen on a mismatched prefix; reopen on an empty prefix. | The mismatch is rejected with both prefixes named and the region is not opened; the empty prefix opens cleanly with the expected row counts. | | Garbage collection deletes everything below the persisted reclaim boundary, and recovery stays correct while objects disappear | Unit tests of the boundary computation (a region without a watermark holds it at its first object, a watermark inside an object keeps the object, `obsolete_all` releases it, last indexed plus one when everything is flushed), of stamping and durability (headers carry the boundary current at sealing, a failed create's boundary permits nothing, the start object carries the recovered boundary), of deletion (chain objects, listed orphans and failed batches below the durable boundary, nothing at or above it, a failed delete retried, `stop` waiting for deletes in flight, holes dropped), of recovery (two restarts with a collection between them, a late object of an earlier epoch below the boundary collected, a missing predecessor above the boundary still failing) and of a read that listed an object before it was collected; the Mito flush-then-collect test in both acknowledgement modes; the standalone integration test's assertion after a flush. | Objects below the durable boundary are gone after a flush, the start object included; a restart replays everything above the boundary; exactly one WAL object remains after a flush and one more insert in the integration test. | -| Every row of the *Failure matrix* holds, and every metric moves as documented | A table-driven fault matrix, `test_store_fault_matrix`, that drives every failure the object store access can produce through the store and asserts, per case, which waiters fail and with what error, whether the store is poisoned, whether the sequence is consumed, the indexed gauges at the failure itself, the exact movement of every metric, and what a read of the store sees afterwards; every timer is observed twice around two forced waits whose intervals cannot overlap. | 33 cases: 17 for the ten rows of the matrix and 16 for the failures listed under *Failures driven through the store*. The test fails if a row of the matrix is left uncovered. | +| Every row of the *Failure matrix* holds, and every metric moves as documented | A table-driven fault matrix, `test_store_fault_matrix`, that drives every failure the object store access can produce through the store and asserts, per case, which waiters fail and with what error, whether the store is poisoned, whether the sequence is consumed, the indexed gauges at the failure itself, the exact movement of every metric, and what a read of the store sees afterwards; every timer is observed twice around two forced waits whose intervals cannot overlap. | 36 cases: 19 for the eleven rows of the matrix and 17 for the failures listed under *Failures driven through the store*. The test fails if a row of the matrix is left uncovered. | | Recovery cost does not grow with object size | Footer-only recovery over 2,000 objects; recovery concurrency bounded and ordered by sequence. | Construction of the store fell from 3.6 s to 0.9 s. | ### Failures driven through the store -- Conditional create outcomes: success, identical retry, conflicting content from an earlier epoch and from the same or a later one, lost response, transient failure in `durable` mode with later objects already created and with creates still in flight (every batch that is not indexed fails, no sequence is reused, late completions are orphans), transient failure in `enqueued` mode and after `stop` began, completion after a conflict had already poisoned the store, a create at the last representable sequence. +- Conditional create outcomes: success, identical retry, conflicting content from an earlier epoch and from the same or a later one, lost response, transient failure in `durable` mode with later objects already created and with creates still in flight (every batch that is not indexed fails, no sequence is reused, late completions are orphans), transient failure in `enqueued` mode and after `stop` began, a failure that requires intervention in `enqueued` mode and a `stop` while it lasts, completion after a conflict had already poisoned the store, a create at the last representable sequence. - Object store access: an unwritable store; a listing, a trailer fetch, a footer fetch, a segment fetch and a delete that fail; a delete that fails and is retried by the next collection; an object collected between the listing and its fetch, below and at the tip's reclaim boundary; an object store that answers slowly. -- Store state: an object the catalog rejects, per-region entry id exhaustion, the catalog refusing a next sequence past the maximum, every stop path with several creates in flight that succeed, fail transiently or conflict after stop began, a stop with a delete in flight, and the appends the backlog thresholds were holding back released by admission, by a stop, by a terminal error and by the teardown of the actor. +- Store state: an object the catalog rejects, per-region entry id exhaustion, the catalog refusing a next sequence past the maximum, every stop path with several creates in flight that succeed, fail transiently or conflict after stop began, a stop with a delete in flight, an append held back by full sealed batches, and the appends the backlog thresholds were holding back released by admission, by a stop, by a terminal error and by the teardown of the actor. - Pipelined uploader: concurrent creates acknowledged in sequence order, every sealed batch linked to its pipeline predecessor, the open batch sealing at the threshold while creates are in flight, a transient failure failing every batch that is not indexed, created ones included, a permanent failure poisoning the store. - Chain recovery: a linear chain, orphans below and above the tip, a broken tip falling back to the latest complete object, a late object of an earlier epoch next to the chain of a later one, a missing predecessor, objects present without a complete chain, orphans raising the next sequence, and the start object on an empty prefix and meeting objects of an earlier and of a later epoch. diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index 836d1fbfe7d..c7b9c981c44 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -436,6 +436,15 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(display( + "Object store WAL rejects appends until the object store accepts its acknowledged backlog" + ))] + ObjectStoreWalFailing { + source: Arc, + #[snafu(implicit)] + location: Location, + }, } pub type Result = std::result::Result; @@ -495,6 +504,7 @@ impl ErrorExt for Error { InvalidWalObject { source, .. } => source.status_code(), ObjectStoreWal { source, .. } => source.status_code(), + ObjectStoreWalFailing { .. } => StatusCode::StorageUnavailable, // Object store related errors WalObjectStore { .. } @@ -534,7 +544,8 @@ impl ErrorExt for Error { | AddEntryLogBatch { .. } | WalObjectSequenceUnsettled { .. } | StaleWalObject { .. } - | UnconfirmedWalEpochStart { .. } => RetryHint::Retryable, + | UnconfirmedWalEpochStart { .. } + | ObjectStoreWalFailing { .. } => RetryHint::Retryable, ProduceRecord { error, .. } => match error { rskafka::client::producer::Error::Client(error) => { rskafka_client_error_to_retry_hint(error) diff --git a/src/log-store/src/metrics.rs b/src/log-store/src/metrics.rs index 03a3dc1f174..c787820aaeb 100644 --- a/src/log-store/src/metrics.rs +++ b/src/log-store/src/metrics.rs @@ -244,7 +244,7 @@ lazy_static! { /// `durable` mode and are repeated in the `enqueued` mode. pub static ref METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL: IntCounter = register_int_counter!( "greptime_logstore_object_store_wal_create_failures_total", - "object store logstore transient conditional create failures total", + "object store logstore conditional creates the object store did not confirm total", ) .unwrap(); @@ -256,12 +256,14 @@ lazy_static! { ) .unwrap(); - /// Counter of appends the object store logstore held back because the - /// unpersisted backlog reached a threshold in the `enqueued` mode, which - /// says the write rate is above what the uploads sustain. + /// Counter of appends the object store logstore held back until an upload + /// completed: in either mode because the sealed batches waiting to become + /// durable left no room for another, and in the `enqueued` mode also + /// because the unpersisted backlog reached a threshold. Either says the + /// write rate is above what the uploads sustain. pub static ref METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL: IntCounter = register_int_counter!( "greptime_logstore_object_store_wal_stalled_appends_total", - "object store logstore appends stalled by the backlog thresholds total", + "object store logstore appends stalled by full sealed batches or the backlog thresholds total", ) .unwrap(); @@ -278,4 +280,24 @@ lazy_static! { exponential_buckets(0.005, 3.0, 9).unwrap(), ) .unwrap(); + + /// Number of object store logstores in the `enqueued` mode whose backlog + /// the object store refuses with an error that repeating the create + /// cannot settle, such as denied access or a missing bucket. Such a store + /// keeps its backlog, attempts it again with backoff and rejects appends + /// until a create succeeds. + pub static ref METRIC_OBJECT_STORE_WAL_FAILING: IntGauge = register_int_gauge!( + "greptime_logstore_object_store_wal_failing", + "object store logstores rejecting appends because the object store refuses their backlog with an error that requires intervention", + ) + .unwrap(); + + /// Counter of appends the object store logstore rejected while its + /// backlog is refused, see the gauge above. They take no entry id and + /// fail with a retryable error. + pub static ref METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_rejected_appends_total", + "object store logstore appends rejected while the object store refuses the backlog total", + ) + .unwrap(); } diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 530f76b9ff6..f233f73bdea 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -24,12 +24,12 @@ use std::time::{Duration, Instant}; use async_stream::try_stream; use bytes::Bytes; -use common_telemetry::{debug, info, warn}; +use common_telemetry::{debug, error, info, warn}; use common_wal::config::object_store::{AckMode, CorruptedSegmentAction, ObjectStoreWalConfig}; -use futures::future::BoxFuture; +use futures::future::{BoxFuture, OptionFuture}; use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt}; -use object_store::ObjectStore; +use object_store::{ErrorKind, ObjectStore}; use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::{ObjectStoreProvider, Provider}; @@ -42,18 +42,19 @@ use tokio_util::task::AbortOnDropHandle; use crate::error::{ CorruptedWalObjectSnafu, Error, IncompleteWalEntrySnafu, InvalidProviderSnafu, InvalidWalObjectSnafu, InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, - MismatchedWalRegionSnafu, ObjectStoreWalSnafu, ObjectStoreWalStoppedSnafu, Result, - StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, - WalObjectSequenceUnsettledSnafu, WalSegmentFetchSnafu, + MismatchedWalRegionSnafu, ObjectStoreWalFailingSnafu, ObjectStoreWalSnafu, + ObjectStoreWalStoppedSnafu, Result, StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, + WalObjectSequenceExhaustedSnafu, WalObjectSequenceUnsettledSnafu, WalSegmentFetchSnafu, }; use crate::metrics::{ METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS, METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL, METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL, METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL, - METRIC_OBJECT_STORE_WAL_INDEXED_BYTES, METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS, - METRIC_OBJECT_STORE_WAL_OBJECT_BYTES, METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES, - METRIC_OBJECT_STORE_WAL_POISONED_TOTAL, METRIC_OBJECT_STORE_WAL_READ_SECONDS, - METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS, + METRIC_OBJECT_STORE_WAL_FAILING, METRIC_OBJECT_STORE_WAL_INDEXED_BYTES, + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS, METRIC_OBJECT_STORE_WAL_OBJECT_BYTES, + METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES, METRIC_OBJECT_STORE_WAL_POISONED_TOTAL, + METRIC_OBJECT_STORE_WAL_READ_SECONDS, METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL, + METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS, METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL, METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS, METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL, METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS, METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL, @@ -76,11 +77,21 @@ const MAX_IN_FLIGHT_DELETES: usize = 4; /// Delay before a create that failed transiently is attempted again in the /// `enqueued` acknowledgement mode, where no caller is left to retry it. const CREATE_RETRY_DELAY: Duration = Duration::from_millis(100); +/// Longest delay from a failed attempt of a backlog that the object store +/// refuses with an error that requires intervention, see [`Failing`], to the +/// next attempt. +/// The delay doubles from [`CREATE_RETRY_DELAY`] after every attempt. At the +/// cap a dead bucket costs four requests a minute, a conditional create and +/// the read that reconciles it every 30 seconds; a longer cap would keep +/// appends rejected longer after the cause is fixed. +const MAX_FAILING_RETRY_DELAY: Duration = Duration::from_secs(30); +/// How often a store logs that the object store still refuses its backlog. +const FAILING_LOG_INTERVAL: Duration = Duration::from_secs(60); /// Number of sealed batches that may wait to be created or indexed. One /// admission can seal two batches, the open batch before an append that would -/// exhaust its positions and the append's own batch, so the actor takes an -/// append from its channel only while two more fit: a stalled object store -/// holds callers back instead of growing the backlog. +/// exhaust its positions and the append's own batch, so the actor admits an +/// append only while two more fit and holds it back otherwise: a stalled +/// object store holds callers back instead of growing the backlog. const MAX_SEALED_BATCHES: usize = 2 * MAX_IN_FLIGHT_CREATES; /// Number of objects whose footers recovery fetches at a time. On S3 an /// object costs two reads, its header and its tail, tens of milliseconds per @@ -123,7 +134,9 @@ const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// indexed, so an acknowledged entry never has a missing predecessor; in the /// `enqueued` mode it returns on admission and the object is created in the /// background. While [`MAX_SEALED_BATCHES`] batches -/// wait to become durable no append is admitted. A segment that a read cannot +/// wait to become durable no append is admitted, and in the `enqueued` mode +/// appends are rejected while the object store refuses the backlog with an +/// error that requires intervention. A segment that a read cannot /// decode is handled as `on_corrupted_segment` says. Objects below the reclaim /// boundary of the last indexed object are deleted in the background. pub struct ObjectStoreLogStore { @@ -377,6 +390,7 @@ impl ObjectStoreLogStore { sealed: VecDeque::new(), creates: FuturesUnordered::new(), stalled: None, + failing: None, durable_waiters: Vec::new(), stop: Vec::new(), stop_error: None, @@ -1046,7 +1060,8 @@ struct DurableWaiter { /// Where the conditional create of a sealed batch stands. enum CreateState { /// The create has not started: no slot was free, or in the `enqueued` - /// mode an earlier attempt failed transiently and is repeated. + /// mode an earlier attempt failed with an object store error and is + /// repeated. Pending, InFlight, Created, @@ -1087,6 +1102,48 @@ impl SealedBatch { type CreateOutcome = (u64, Result); +/// The state of an `enqueued` store whose backlog the object store refuses +/// with an error that requires intervention, see [`requires_intervention`], +/// such as denied +/// access or a missing bucket. +/// +/// Poisoning would drop the acknowledged backlog and need a restart, while +/// such a cause can be fixed in place. So the backlog keeps its sequences and +/// bytes, and its lowest pending batch is attempted alone, once nothing is in +/// flight, after a delay that doubles from [`CREATE_RETRY_DELAY`] up to +/// [`MAX_FAILING_RETRY_DELAY`]. Appends are rejected at once with a retryable +/// error that carries the last error: they take no entry id and never wait +/// behind the backlog. The first create that succeeds ends the state. +struct Failing { + /// The last error a create met, which rejected appends carry. + error: Arc, + since: Instant, + /// The delay before the next attempt is due once nothing is in flight. + delay: Duration, + /// When the next attempt starts, `None` until it is scheduled. + probe_at: Option, + failed_creates: u64, + rejected_appends: u64, + logged_at: Instant, +} + +impl Failing { + /// Schedules the next attempt after the current delay and doubles the + /// delay for the one after it. + fn schedule(&mut self) { + self.probe_at = Some(Instant::now() + self.delay); + self.delay = (self.delay * 2).min(MAX_FAILING_RETRY_DELAY); + } + + /// Counts an admission the store refuses and returns the error it + /// fails with. + fn reject(&mut self) -> Error { + self.rejected_appends += 1; + METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.inc(); + ObjectStoreWalFailingSnafu.into_error(self.error.clone()) + } +} + /// The actor that owns the open batch and the sealed batches until they are /// durable. /// @@ -1103,7 +1160,8 @@ type CreateOutcome = (u64, Result); /// | created, or identical retry | advances | acknowledged in order | healthy | /// | transient error, `durable` mode | never reused: a create that reported an error may have written its object | every batch that is not indexed fails, created ones included; retries are assigned new ids | healthy: the next batch links to the last indexed object, so no chain reaches a failed batch | /// | transient error, `enqueued` mode | unchanged | already acknowledged | healthy: the create is repeated with the same bytes after [`CREATE_RETRY_DELAY`], which an identical retry accepts | -/// | transient error, `enqueued` mode after `stop` began | never reused | already acknowledged | the backlog is dropped and `stop` reports the error | +/// | error that requires intervention, `enqueued` mode | unchanged | already acknowledged; later appends are rejected with a retryable error | failing, see [`Failing`]: the lowest pending create is repeated alone with backoff, and the first create that succeeds ends it | +/// | object store error, `enqueued` mode after `stop` began | never reused | already acknowledged | the backlog is dropped and `stop` reports the error | /// | an object of an earlier epoch holds the sequence, `durable` mode | never reused | as for a transient error | healthy: that object can never be on the chain | /// | an object of the same or a later epoch holds the sequence, any conflicting object in the `enqueued` mode, encoding or catalog error | unchanged | every batch that is not indexed fails | poisoned | /// | created at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence | @@ -1114,10 +1172,11 @@ type CreateOutcome = (u64, Result); /// uploads its backlog; creates in flight run to completion and acknowledge /// if they succeed. /// -/// Appends arrive on their own channel, which the actor does not read while -/// [`MAX_SEALED_BATCHES`] batches wait or an append is held back at a backlog -/// threshold of the `enqueued` mode, so commands such as `stop` are handled -/// while admission is held back. +/// Appends arrive on their own channel. The actor holds one back while the +/// sealed batches leave no room for the two an admission can seal, or at a +/// backlog threshold of the `enqueued` mode, and does not read the channel +/// meanwhile, so commands such as `stop` are handled while admission is held +/// back. A failing store reads the channel to reject every append. /// /// The actor also collects garbage. Every batch carries the reclaim boundary /// current when it was sealed, see [`ObjectCatalog::reclaim_boundary`], which @@ -1152,10 +1211,13 @@ struct Actor { sealed: VecDeque, /// Creates in flight, including those of batches that already failed. creates: FuturesUnordered>, - /// The append held back in the `enqueued` mode while the unpersisted - /// backlog is at a threshold, and when it was held back. Later appends - /// wait in their channel. + /// The append held back while the sealed batches are full or, in the + /// `enqueued` mode, the unpersisted backlog is at a threshold, and when + /// it was held back. Later appends wait in their channel. stalled: Option<(QueuedAppend, Instant)>, + /// Set while the object store refuses the `enqueued` backlog with an + /// error that requires intervention. + failing: Option, durable_waiters: Vec, /// Callers of `stop`, answered once nothing is in flight. stop: Vec>>, @@ -1207,6 +1269,7 @@ impl Actor { interval.tick().await; loop { + let probe_at = self.failing.as_ref().and_then(|failing| failing.probe_at); tokio::select! { _ = interval.tick() => { // Nothing starts after stop began; the `enqueued` mode @@ -1221,7 +1284,10 @@ impl Actor { Some((object_seq, result)) = self.deletes.next(), if !self.deletes.is_empty() => { self.on_delete_completed(object_seq, result); } - Some((entries, response)) = self.append_rx.recv(), if self.stalled.is_none() && self.sealed.len() + 2 <= MAX_SEALED_BATCHES => { + Some(()) = OptionFuture::from(probe_at.map(|at| tokio::time::sleep_until(at.into()))) => { + self.probe(); + } + Some((entries, response)) = self.append_rx.recv(), if self.stalled.is_none() => { self.handle_append(entries, response); } command = self.command_rx.recv() => match command { @@ -1271,7 +1337,11 @@ impl Actor { let _ = response.send(Err(shared(&error))); return; } - if self.ack_mode == AckMode::Enqueued && self.backlog_at_threshold() { + if let Some(failing) = &mut self.failing { + let _ = response.send(Err(failing.reject())); + return; + } + if self.should_hold_back() { METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.inc(); self.stalled = Some(((entries, response), Instant::now())); self.ensure_create_in_flight(); @@ -1336,6 +1406,14 @@ impl Actor { } } + /// Returns true while an append has to wait for an upload: the sealed + /// batches leave no room for the two an admission can seal, or in the + /// `enqueued` mode the unpersisted backlog is at a threshold. + fn should_hold_back(&self) -> bool { + self.sealed.len() + 2 > MAX_SEALED_BATCHES + || (self.ack_mode == AckMode::Enqueued && self.backlog_at_threshold()) + } + /// Returns true once the unpersisted backlog, the open batch and every /// sealed batch that is not durable, reaches the size or the age threshold. fn backlog_at_threshold(&self) -> bool { @@ -1364,8 +1442,7 @@ impl Actor { } } - /// Admits the stalled append once the backlog is below the thresholds and - /// the sealed batches leave room for the two an admission can seal. + /// Admits the stalled append once it no longer has to be held back. fn release_stalled(&mut self) { if self.stalled.is_none() { return; @@ -1376,14 +1453,11 @@ impl Actor { } return; } - if self.backlog_at_threshold() { + if self.should_hold_back() { // The append needs an upload to complete. self.ensure_create_in_flight(); return; } - if self.sealed.len() + 2 > MAX_SEALED_BATCHES { - return; - } if let Some((entries, response)) = self.take_stalled() { self.admit(entries, response); } @@ -1475,88 +1549,122 @@ impl Actor { /// Starts the creates of pending batches in sequence order while fewer /// than [`MAX_IN_FLIGHT_CREATES`] are in flight, counting the creates of /// batches that already failed. Nothing starts once stop began, except - /// the backlog of the `enqueued` mode. + /// the backlog of the `enqueued` mode. A failing store only schedules the + /// next attempt once nothing is in flight, see [`Failing`]; after stop + /// began it uploads its backlog at once. fn start_creates(&mut self) { if self.is_stopped() && self.ack_mode == AckMode::Durable { return; } + if !self.is_stopped() + && let Some(failing) = &mut self.failing + { + if self.creates.is_empty() && failing.probe_at.is_none() { + failing.schedule(); + } + return; + } let mut in_flight = self.creates.len(); - for batch in self.sealed.iter_mut() { + for index in 0..self.sealed.len() { if in_flight >= MAX_IN_FLIGHT_CREATES { break; } + let batch = &self.sealed[index]; if !matches!(batch.state, CreateState::Pending) { continue; } - batch.state = CreateState::InFlight; - in_flight += 1; let delay = if batch.attempts == 0 { Duration::ZERO } else { CREATE_RETRY_DELAY }; - batch.attempts += 1; - let io = self.io.clone(); - let object_seq = batch.object_seq; - let sealed_at = batch.sealed_at; - self.creating.insert(object_seq); - let bytes = batch.bytes.clone(); - let epoch = self.epoch; - let ack_mode = self.ack_mode; - #[cfg(any(test, feature = "testing"))] - let mut creates_held = self.creates_held.clone(); - #[cfg(any(test, feature = "testing"))] - let creates_fail = self.creates_fail.clone(); - #[cfg(any(test, feature = "testing"))] - let next_create_fails_after_write = self.next_create_fails_after_write.clone(); - #[cfg(any(test, feature = "testing"))] - let parked_creates = self.parked_creates.clone(); - self.creates.push(Box::pin(async move { - if !delay.is_zero() { - tokio::time::sleep(delay).await; - } - // The store was dropped while the create was parked: it never - // runs. - #[cfg(any(test, feature = "testing"))] - if *creates_held.borrow() { - parked_creates.send_modify(|count| *count += 1); - } - #[cfg(any(test, feature = "testing"))] - if creates_held.wait_for(|held| !*held).await.is_err() { - return (object_seq, Err(ObjectStoreWalStoppedSnafu.build())); - } - #[cfg(any(test, feature = "testing"))] - if creates_fail.load(Ordering::Acquire) { - let result = injected_create_failure(io.as_ref(), object_seq); - record_create(&result, Some(sealed_at)); - return (object_seq, result); - } - let created = io.put_if_absent(object_seq, bytes).await; - #[cfg(any(test, feature = "testing"))] - let created = match created { - Ok(_) if next_create_fails_after_write.swap(false, Ordering::AcqRel) => { - injected_create_failure(io.as_ref(), object_seq) - } - created => created, - }; - // Counted as the object store answered, before the epoch of a - // conflicting object is read. - record_create(&created, Some(sealed_at)); - // Any conflict poisons an `enqueued` store, so only the - // `durable` mode reads the epoch of the existing object. - let result = match created { - Err(error @ Error::WalObjectConflict { .. }) - if ack_mode == AckMode::Durable => - { - stale_conflict(io.as_ref(), object_seq, epoch, error).await - } - result => result, - }; - (object_seq, result) - })); + self.start_create(index, delay); + in_flight += 1; } } + /// Attempts the lowest pending batch of a failing store, whose delay + /// elapsed. After stop began `start_creates` uploads the backlog instead. + fn probe(&mut self) { + let Some(failing) = &mut self.failing else { + return; + }; + failing.probe_at = None; + if self.is_stopped() { + return; + } + if let Some(index) = self + .sealed + .iter() + .position(|batch| matches!(batch.state, CreateState::Pending)) + { + self.start_create(index, Duration::ZERO); + } + } + + /// Starts the create of the sealed batch at `index` after `delay`. + fn start_create(&mut self, index: usize, delay: Duration) { + let batch = &mut self.sealed[index]; + batch.state = CreateState::InFlight; + batch.attempts += 1; + let object_seq = batch.object_seq; + let sealed_at = batch.sealed_at; + let bytes = batch.bytes.clone(); + let io = self.io.clone(); + self.creating.insert(object_seq); + let epoch = self.epoch; + let ack_mode = self.ack_mode; + #[cfg(any(test, feature = "testing"))] + let mut creates_held = self.creates_held.clone(); + #[cfg(any(test, feature = "testing"))] + let creates_fail = self.creates_fail.clone(); + #[cfg(any(test, feature = "testing"))] + let next_create_fails_after_write = self.next_create_fails_after_write.clone(); + #[cfg(any(test, feature = "testing"))] + let parked_creates = self.parked_creates.clone(); + self.creates.push(Box::pin(async move { + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } + // The store was dropped while the create was parked: it never + // runs. + #[cfg(any(test, feature = "testing"))] + if *creates_held.borrow() { + parked_creates.send_modify(|count| *count += 1); + } + #[cfg(any(test, feature = "testing"))] + if creates_held.wait_for(|held| !*held).await.is_err() { + return (object_seq, Err(ObjectStoreWalStoppedSnafu.build())); + } + #[cfg(any(test, feature = "testing"))] + if creates_fail.load(Ordering::Acquire) { + let result = injected_create_failure(io.as_ref(), object_seq); + record_create(&result, Some(sealed_at)); + return (object_seq, result); + } + let created = io.put_if_absent(object_seq, bytes).await; + #[cfg(any(test, feature = "testing"))] + let created = match created { + Ok(_) if next_create_fails_after_write.swap(false, Ordering::AcqRel) => { + injected_create_failure(io.as_ref(), object_seq) + } + created => created, + }; + // Counted as the object store answered, before the epoch of a + // conflicting object is read. + record_create(&created, Some(sealed_at)); + // Any conflict poisons an `enqueued` store, so only the + // `durable` mode reads the epoch of the existing object. + let result = match created { + Err(error @ Error::WalObjectConflict { .. }) if ack_mode == AckMode::Durable => { + stale_conflict(io.as_ref(), object_seq, epoch, error).await + } + result => result, + }; + (object_seq, result) + })); + } + fn on_create_completed(&mut self, object_seq: u64, result: Result) { self.creating.remove(&object_seq); // A batch that already failed, or that the store gave up on when it @@ -1572,28 +1680,27 @@ impl Actor { return; }; match result { - Ok(_) => self.sealed[index].state = CreateState::Created, + Ok(_) => { + self.sealed[index].state = CreateState::Created; + self.leave_failing("a create succeeded"); + } // Nobody is left to retry in the `enqueued` mode, so the store - // repeats a create that failed transiently under the same sequence - // with the same bytes, which an identical retry accepts. - Err(ref error) - if self.ack_mode == AckMode::Enqueued - && !self.is_stopped() - && is_transient(error) => + // repeats a create that the object store did not confirm under + // the same sequence with the same bytes, which an identical retry + // accepts: after `CREATE_RETRY_DELAY`, or with backoff while the + // store is failing. + Err(error @ Error::WalObjectStore { .. }) + if self.ack_mode == AckMode::Enqueued && !self.is_stopped() => { self.sealed[index].state = CreateState::Pending; + self.on_create_failed(error); } // The object store did not confirm the object, which may still // exist or land later, or an earlier epoch holds the sequence and // can never be on the chain. A caller of the `durable` mode - // retries the append itself; after stop began a transient failure - // drops the `enqueued` backlog. - Err(error @ Error::WalObjectStore { .. }) - if self.ack_mode == AckMode::Durable - || (self.is_stopped() && is_transient(&error)) => - { - self.roll_back(Arc::new(error)) - } + // retries the append itself; after stop began a failure drops the + // `enqueued` backlog. + Err(error @ Error::WalObjectStore { .. }) => self.roll_back(Arc::new(error)), Err(error @ Error::StaleWalObject { .. }) if self.ack_mode == AckMode::Durable => { self.roll_back(Arc::new(error)) } @@ -1605,6 +1712,70 @@ impl Actor { self.settle(); } + /// Records the failure of a create the `enqueued` mode repeats. An error + /// that requires intervention makes the store failing; while it is, every + /// error becomes the one rejected appends carry. + fn on_create_failed(&mut self, error: Error) { + let Some(failing) = &mut self.failing else { + if requires_intervention(&error) { + self.enter_failing(error); + } + return; + }; + failing.failed_creates += 1; + failing.error = Arc::new(error); + if failing.logged_at.elapsed() >= FAILING_LOG_INTERVAL { + failing.logged_at = Instant::now(); + warn!( + failing.error; + "Object store WAL still cannot upload its backlog of {} objects after {:?}: {} creates failed, {} appends rejected, next attempt in {:?}", + self.sealed.len(), + failing.since.elapsed(), + failing.failed_creates, + failing.rejected_appends, + failing.delay + ); + } + } + + fn enter_failing(&mut self, error: Error) { + error!( + error; + "Object store WAL cannot upload its backlog of {} objects: the object store refused a create with an error that repeating it cannot settle. Appends are rejected until a create succeeds; the backlog is kept and attempted again with backoff", + self.sealed.len() + ); + METRIC_OBJECT_STORE_WAL_FAILING.inc(); + // The append held back would wait for an upload that is not coming. + let stalled = self.take_stalled(); + let now = Instant::now(); + let failing = self.failing.insert(Failing { + error: Arc::new(error), + since: now, + delay: CREATE_RETRY_DELAY, + probe_at: None, + failed_creates: 1, + rejected_appends: 0, + logged_at: now, + }); + if let Some((_, response)) = stalled { + let _ = response.send(Err(failing.reject())); + } + } + + /// Ends the failing state, if the store is in it, because of `reason`. + fn leave_failing(&mut self, reason: &str) { + let Some(failing) = self.failing.take() else { + return; + }; + METRIC_OBJECT_STORE_WAL_FAILING.dec(); + info!( + "Object store WAL is no longer failing, {reason}, after {:?}: {} creates failed, {} appends rejected", + failing.since.elapsed(), + failing.failed_creates, + failing.rejected_appends + ); + } + /// Indexes and acknowledges the sealed batches from the front as far as /// they are created, and starts the creates that a free slot allows. fn settle(&mut self) { @@ -1718,12 +1889,15 @@ impl Actor { batch.fail(failure); } self.reset_open_batch(); - self.fail_unacknowledged(failure); + // The stalled append was never admitted; the sealed batches are gone, + // so `settle` admits it, or refuses it once stop began. + self.fail_waiters(failure); // An acknowledged backlog was dropped: `stop` reports it, whether // its caller has arrived yet or not. if self.ack_mode == AckMode::Enqueued { self.stop_error.get_or_insert(error); } + self.leave_failing("its backlog was dropped"); } /// Records `error` as terminal and fails every waiter that is not @@ -1751,6 +1925,7 @@ impl Actor { if self.ack_mode == AckMode::Enqueued { self.stop_error.get_or_insert(error.clone()); } + self.leave_failing("the store is poisoned"); error } @@ -1763,12 +1938,17 @@ impl Actor { /// Fails the waiters of the open batch, the stalled append and the /// durability waiters. fn fail_unacknowledged(&mut self, error: impl Fn() -> Error) { - for pending in self.pending.drain(..) { - let _ = pending.response.send(Err(error())); - } if let Some((_, response)) = self.take_stalled() { let _ = response.send(Err(error())); } + self.fail_waiters(error); + } + + /// Fails the waiters of the open batch and the durability waiters. + fn fail_waiters(&mut self, error: impl Fn() -> Error) { + for pending in self.pending.drain(..) { + let _ = pending.response.send(Err(error())); + } for waiter in self.durable_waiters.drain(..) { let _ = waiter.response.send(Err(error())); } @@ -1943,10 +2123,10 @@ impl Actor { return Err(shared(&error)); } // In the `enqueued` mode an id the store handed out is never handed - // out again while it runs: a create that fails transiently is - // repeated under its sequence and a permanent failure poisons the - // store. Every later id of the region is greater, so the id needs no - // floor even before it is durable. + // out again while it runs: a create that the object store did not + // confirm is repeated under its sequence and any other failure + // poisons the store. Every later id of the region is greater, so the + // id needs no floor even before it is durable. if self.ack_mode == AckMode::Enqueued && entry_id <= self.issued_entry_ids.get(®ion_id).copied().unwrap_or(0) { @@ -1980,12 +2160,16 @@ impl Actor { } } - /// Begins stopping. Nothing is admitted from now on; the `durable` mode - /// drops the open batch and the batches whose create has not started, - /// the `enqueued` mode seals its backlog so it is uploaded. Stop is - /// answered by [`finish_stop`](Self::finish_stop) once nothing is in flight. + /// Begins stopping. Nothing is admitted from now on and the append held + /// back is refused; the `durable` mode drops the open batch and the + /// batches whose create has not started, the `enqueued` mode seals its + /// backlog and uploads it, a failing store included. Stop is answered by + /// [`finish_stop`](Self::finish_stop) once nothing is in flight. fn handle_stop(&mut self, response: oneshot::Sender>) { self.stop.push(response); + if let Some((_, response)) = self.take_stalled() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + } match self.ack_mode { AckMode::Durable => { self.reset_open_batch(); @@ -2005,10 +2189,8 @@ impl Actor { } } AckMode::Enqueued => { - if let Some((_, response)) = self.take_stalled() { - let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); - } self.flush_open_batch(); + self.start_creates(); } } } @@ -2079,6 +2261,7 @@ impl Drop for Actor { fn drop(&mut self) { self.deleting.abandon_all(); self.take_stalled(); + self.leave_failing("its actor exited"); #[cfg(test)] LIVE_ACTORS.fetch_sub(1, Ordering::SeqCst); } @@ -2142,11 +2325,34 @@ fn publish_indexed(catalog: &ObjectCatalog) { METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.set(bytes as i64); } -/// Returns true for a storage error that a later attempt may not meet: one -/// the object store reports as temporary, or as persistent, which is how its -/// retry layer reports a temporary error that outlasted its retries. -fn is_transient(error: &Error) -> bool { - matches!(error, Error::WalObjectStore { error, .. } if !error.is_permanent()) +/// Returns true for an object store error that repeating the request cannot +/// settle until someone changes the credentials, the configuration or the +/// bucket: denied access; an invalid configuration, which is how OpenDAL +/// reports a missing S3 bucket; and a write to a location that does not +/// exist, which is how OSS, GCS and Azure Blob Storage report a missing +/// bucket or container. A read that finds nothing is not one: it is what the +/// read that reconciles a create meets when the object is gone, and a repeat +/// creates it. Every other error, an HTTP 400 or 5xx, a timeout or a refused +/// connection, may pass by itself. A create that fails with an error +/// requiring intervention is still repeated, with the backoff [`Failing`] +/// describes. +/// +/// Only the kind of the error decides. Its status cannot: on object storage +/// the store sits behind OpenDAL's retry layer, which marks every error it +/// returns as persistent, including those it never retried because they were +/// not temporary, so a 403 arrives with the status of a 503 that outlasted +/// the retries. +fn requires_intervention(error: &Error) -> bool { + match error { + Error::WalObjectStore { + operation, error, .. + } => match error.kind() { + ErrorKind::PermissionDenied | ErrorKind::ConfigInvalid => true, + ErrorKind::NotFound => *operation == "write", + _ => false, + }, + _ => false, + } } /// Wraps an error that several callers receive. @@ -2762,13 +2968,16 @@ mod tests { use common_base::readable_size::ReadableSize; use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; use futures::FutureExt; + use object_store::layers::mock::{self, MockLayerBuilder, oio}; use object_store::services::Memory; + use object_store::util::with_retry_layers; use store_api::logstore::entry::{MultiplePartEntry, MultiplePartHeader}; use tokio::time::timeout; use super::*; - use crate::error::WalObjectStoreSnafu; + use crate::error::{WalObjectConflictSnafu, WalObjectStoreSnafu}; use crate::object_store_wal::batch::{POSITION_LIMIT, entry_id}; use crate::object_store_wal::format::{FOOTER_ENTRY_LEN, decode_object}; @@ -2973,6 +3182,110 @@ mod tests { store.seal_open_batch().await.unwrap(); } + /// An in-memory object store behind the retry layer object storage has + /// in production, whose writes and reads fail with the error kind their + /// switch holds, recording when each refused write started. The retry + /// layer reports such an error as persistent, as it reports a 403 of S3. + struct RefusingStore { + store: ObjectStore, + writes: Arc>>, + reads: Arc>>, + refused_writes: Arc>>, + } + + impl RefusingStore { + fn new() -> Self { + let writes = Arc::new(Mutex::new(None::)); + let reads = Arc::new(Mutex::new(None::)); + let refused_writes = Arc::new(Mutex::new(Vec::new())); + let (refusing_writes, refused_reads) = (writes.clone(), reads.clone()); + let refused_at = refused_writes.clone(); + let mock = MockLayerBuilder::default() + .writer_factory(Arc::new(move |_, _, inner| { + match *refusing_writes.lock().unwrap() { + Some(kind) => { + refused_at.lock().unwrap().push(Instant::now()); + Box::new(RefusedWriter(kind)) + } + None => inner, + } + })) + .reader_factory(Arc::new(move |_, _, inner| { + match *refused_reads.lock().unwrap() { + Some(kind) => Box::new(RefusedReader(kind)), + None => inner, + } + })) + .build() + .unwrap(); + Self { + store: with_retry_layers(memory_store().layer(mock)), + writes, + reads, + refused_writes, + } + } + + fn refuse_writes(&self, kind: Option) { + *self.writes.lock().unwrap() = kind; + } + + fn refuse_reads(&self, kind: Option) { + *self.reads.lock().unwrap() = kind; + } + } + + struct RefusedWriter(ErrorKind); + + impl oio::Write for RefusedWriter { + async fn write(&mut self, _buffer: mock::Buffer) -> mock::Result<()> { + Ok(()) + } + + async fn close(&mut self) -> mock::Result { + Err(mock::Error::new(self.0, "refused by the object store")) + } + + async fn abort(&mut self) -> mock::Result<()> { + Ok(()) + } + } + + struct RefusedReader(ErrorKind); + + impl oio::Read for RefusedReader { + async fn open( + &self, + _range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, Box)> { + Err(mock::Error::new(self.0, "refused by the object store")) + } + + async fn read( + &self, + _range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, mock::Buffer)> { + Err(mock::Error::new(self.0, "refused by the object store")) + } + } + + /// Asserts that an append was rejected because the object store refuses + /// the backlog with `kind`: the error is retryable and carries the + /// refusal. + fn assert_rejected(error: &Error, kind: ErrorKind) { + assert!( + matches!(error, Error::ObjectStoreWalFailing { source, .. } + if matches!(&**source, Error::WalObjectStore { error, .. } if error.kind() == kind)), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + } + + fn failing_stores() -> i64 { + METRIC_OBJECT_STORE_WAL_FAILING.get() + } + #[tokio::test] async fn test_store_rejects_invalid_config() { let _serialized = SERIALIZED.lock().await; @@ -5624,26 +5937,36 @@ mod tests { let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_id = region(1); let admitted = || *store.admitted_appends.borrow(); + let stalled = METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get(); + let stalled_waits = METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.get_sample_count(); - // Every create is parked: once no two more sealed batches fit, - // further appends wait in the channel and are not admitted. + // Every create is parked: once no two more sealed batches fit, the + // next append is held back and counted as stalled, and further + // appends wait in the channel; none is admitted. let mut appends = spawn_appends(&store, region_id, MAX_SEALED_BATCHES - 1).await; for index in 0..3 { let entries = vec![entry(&store, region_id, &format!("w{index}"))]; appends.push(spawn_append_batch(&store, entries)); } let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + wait_until(|| METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get() == stalled + 1).await; round_trip_actor(&store).await; assert_eq!(MAX_SEALED_BATCHES - 1, admitted()); - // A durable object frees a place, and one more append is admitted. + // A durable object frees a place, and the held-back append is + // admitted; the next one is held back in its turn. releases.remove(&1).unwrap().send(true).unwrap(); timeout(WAIT, store.wait_for_admitted_appends(MAX_SEALED_BATCHES)) .await .unwrap() .unwrap(); + wait_until(|| METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get() == stalled + 2).await; round_trip_actor(&store).await; assert_eq!(MAX_SEALED_BATCHES, admitted()); + assert_eq!( + stalled_waits + 1, + METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.get_sample_count() + ); // Stop is handled while appends are held back: the batches whose create // has not started and the appends still waiting learn of the stop. @@ -6533,37 +6856,388 @@ mod tests { ); // An object store with a retry layer reports a temporary error that - // outlasted its retries as persistent; the create is repeated too. + // outlasted its retries as persistent; the create is repeated too, + // and appends are still admitted meanwhile. + let rejected = METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.get(); io.fail_next_put_persistently(); append(&store, region_id, "a2").await.unwrap(); - timeout(WAIT, store.wait_durable(&provider(region_id), id(2, 1))) + append(&store, region_id, "a3").await.unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(3, 1))) .await .unwrap() .unwrap(); - assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + assert_eq!( + rejected, + METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.get() + ); } #[tokio::test] - async fn test_store_enqueued_permanent_failure_poisons() { + async fn test_store_classifies_object_store_errors_by_kind() { + let refusing = RefusingStore::new(); + let io = ObjectStoreIo::new(refusing.store.clone(), PREFIX).unwrap(); + let content = Bytes::from_static(b"wal"); + // Behind the retry layer every refusal is persistent, whatever its + // kind, so only the kind tells what a repeat can settle: denied + // access, a missing S3 bucket, or a write to a missing bucket of + // another service cannot; an HTTP 400 or a 5xx that outlasted the + // retries may pass. + for (kind, requires) in [ + (ErrorKind::PermissionDenied, true), + (ErrorKind::ConfigInvalid, true), + (ErrorKind::NotFound, true), + (ErrorKind::Unexpected, false), + ] { + refusing.refuse_writes(Some(kind)); + let error = io.put_if_absent(1, content.clone()).await.unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { operation: "write", error, .. } + if error.kind() == kind && error.is_persistent()), + "unexpected error for {kind:?}: {error:?}" + ); + assert_eq!(requires, requires_intervention(&error), "{kind:?}"); + } + + // The object exists, so the create is refused by its condition and + // the read that reconciles it decides: denied access cannot pass, an + // object that is gone is created by the repeat. + refusing.refuse_writes(None); + io.put_if_absent(1, content.clone()).await.unwrap(); + for (kind, requires) in [ + (ErrorKind::PermissionDenied, true), + (ErrorKind::NotFound, false), + ] { + refusing.refuse_reads(Some(kind)); + let error = io + .put_if_absent(1, Bytes::from_static(b"other")) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { operation: "read", error, .. } + if error.kind() == kind && error.is_persistent()), + "unexpected error for {kind:?}: {error:?}" + ); + assert_eq!(requires, requires_intervention(&error), "{kind:?}"); + } + assert!(!requires_intervention( + &WalObjectConflictSnafu { + path: io.object_path(1) + } + .build() + )); + } + + #[test] + fn test_store_failing_delay_doubles_up_to_the_cap() { + let now = Instant::now(); + let mut failing = Failing { + error: Arc::new(ObjectStoreWalStoppedSnafu.build()), + since: now, + delay: CREATE_RETRY_DELAY, + probe_at: None, + failed_creates: 1, + rejected_appends: 0, + logged_at: now, + }; + let mut delays = Vec::new(); + for _ in 0..12 { + let delay = failing.delay; + let scheduled_at = Instant::now(); + failing.schedule(); + assert!(failing.probe_at.unwrap() >= scheduled_at + delay); + delays.push(delay.as_millis()); + } + assert_eq!( + vec![ + 100, 200, 400, 800, 1600, 3200, 6400, 12800, 25600, 30000, 30000, 30000 + ], + delays + ); + } + + /// Opens an `enqueued` store over a refusing object store, makes it + /// refuse writes with `kind` and appends one entry, whose create fails, + /// and waits until the store is failing. + async fn open_failing(kind: ErrorKind) -> (Arc, RefusingStore) { + let refusing = RefusingStore::new(); + let store = open(refusing.store.clone(), &enqueued(eager())).await; + let region_id = region(1); + refusing.refuse_writes(Some(kind)); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 1))]), + response.last_entry_ids + ); + wait_until(|| failing_stores() == 1).await; + (store, refusing) + } + + #[tokio::test] + async fn test_store_enqueued_failure_requiring_intervention_keeps_the_backlog() { + let _serialized = SERIALIZED.lock().await; + let region_id = region(1); + // Denied access and a missing bucket, as S3 reports them. + for kind in [ErrorKind::PermissionDenied, ErrorKind::ConfigInvalid] { + let rejected = METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.get(); + let (store, refusing) = open_failing(kind).await; + let admitted = *store.admitted_appends.borrow(); + + // A later append is rejected at once and takes no entry id; the + // acknowledged entry waits in the backlog and the store is not + // poisoned. + let error = append(&store, region_id, "a2").await.unwrap_err(); + assert_rejected(&error, kind); + assert_eq!(admitted, *store.admitted_appends.borrow()); + assert_eq!( + rejected + 1, + METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.get() + ); + let wait = { + let store = store.clone(); + tokio::spawn( + async move { store.wait_durable(&provider(region_id), id(1, 1)).await }, + ) + }; + timeout(WAIT, store.wait_for_durability_waits(1)) + .await + .unwrap() + .unwrap(); + assert!(!wait.is_finished()); + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!(0, latest(&store, region_id)); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + + // Once the object store accepts writes again, the backlog is + // uploaded under its sequence and admission resumes. + refusing.refuse_writes(None); + timeout(WAIT, wait).await.unwrap().unwrap().unwrap(); + assert_eq!(0, failing_stores()); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + let response = append(&store, region_id, "a3").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + timeout(WAIT, store.wait_durable(&provider(region_id), id(2, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_enqueued_create_retries_back_off() { + let _serialized = SERIALIZED.lock().await; + let (store, refusing) = open_failing(ErrorKind::PermissionDenied).await; + let refused = || refusing.refused_writes.lock().unwrap().clone(); + wait_until(|| refused().len() >= 5).await; + + // Each attempt starts at least the doubling delay after the previous + // one started, so the 1.6 s from the first refusal hold at most five + // attempts, where a repeat every 100 ms would make sixteen. + let refused = refused(); + for (index, delay) in [100, 200, 400, 800].into_iter().enumerate() { + let gap = refused[index + 1] - refused[index]; + assert!( + gap >= Duration::from_millis(delay), + "attempt {} started {gap:?} after the one before", + index + 1 + ); + } + let window = refused[0] + Duration::from_millis(1600); + assert!(refused.iter().filter(|at| **at < window).count() <= 5); + assert_eq!(1, failing_stores()); + + refusing.refuse_writes(None); + timeout(WAIT, store.wait_durable(&provider(region(1)), id(1, 1))) + .await + .unwrap() + .unwrap(); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_enqueued_failing_store_attempts_the_lowest_batch_alone() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let region_id = region(1); + for append in spawn_appends(&store, region_id, 3).await { + timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + } + let mut releases = parked_operations(&mut parked, 3).await; + io.creates_refused.store(true, Ordering::SeqCst); + for object_seq in [2, 1, 3] { + releases.remove(&object_seq).unwrap().send(false).unwrap(); + } + + for _ in 0..2 { + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert_eq!(1, io.in_flight.load(Ordering::SeqCst)); + release.send(false).unwrap(); + assert!(parked.try_recv().is_err()); + } + + // The attempt that succeeds ends the failing state: the other + // batches are created together and everything is acknowledged in + // order. + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + release.send(true).unwrap(); + for release in parked_operations(&mut parked, 2).await.into_values() { + release.send(true).unwrap(); + } + timeout(WAIT, store.wait_durable(&provider(region_id), id(3, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(0, failing_stores()); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + assert_eq!( + expected_entries( + region_id, + &[(id(1, 1), "a1"), (id(2, 1), "a2"), (id(3, 1), "a3")] + ), + read_entries(&store, region_id, 1).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_enqueued_rejects_the_appends_held_back_by_a_refused_backlog() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let region_id = region(1); + let stalled = METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get(); + let stalled_waits = METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.get_sample_count(); + let rejected = METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.get(); + + // The sealed batches are full: the next append is held back and the + // ones after it wait in their channel. + for append in spawn_appends(&store, region_id, MAX_SEALED_BATCHES - 1).await { + timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + } + let held = (0..3) + .map(|index| { + spawn_append_batch(&store, vec![entry(&store, region_id, &format!("h{index}"))]) + }) + .collect::>(); + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + wait_until(|| METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get() == stalled + 1).await; + assert!(held.iter().all(|append| !append.is_finished())); + + // The object store denies access: every append waiting is rejected + // instead of hanging, and so is a later one. + io.creates_refused.store(true, Ordering::SeqCst); + for release in releases.drain().map(|(_, release)| release) { + release.send(false).unwrap(); + } + for append in held { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert_rejected(&error, ErrorKind::PermissionDenied); + } + assert_rejected( + &append(&store, region_id, "a8").await.unwrap_err(), + ErrorKind::PermissionDenied, + ); + assert_eq!( + rejected + 4, + METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.get() + ); + assert_eq!( + stalled_waits + 1, + METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.get_sample_count() + ); + + // Once a create succeeds the backlog drains and the next append + // takes the sequence after it. + io.creates_refused.store(false, Ordering::SeqCst); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + release.send(true).unwrap(); + for _ in 2..MAX_SEALED_BATCHES { + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + } + let last = id(MAX_SEALED_BATCHES as u64 - 1, 1); + timeout(WAIT, store.wait_durable(&provider(region_id), last)) + .await + .unwrap() + .unwrap(); + let appended = spawn_append_batch(&store, vec![entry(&store, region_id, "a9")]); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + let response = timeout(WAIT, appended).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(MAX_SEALED_BATCHES as u64, 1))]), + response.last_entry_ids + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_enqueued_stop_uploads_a_refused_backlog_at_once() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let region_id = region(1); + for append in spawn_appends(&store, region_id, 3).await { + timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + } + io.creates_refused.store(true, Ordering::SeqCst); + for release in parked_operations(&mut parked, 3).await.into_values() { + release.send(false).unwrap(); + } + wait_until(|| failing_stores() == 1).await; + + // Stop does not wait for the next attempt of the failing store: it + // starts the creates of the whole backlog together, which succeed + // now. + let stop = begin_spawned_stop(&store).await; + let releases = parked_operations(&mut parked, 3).await; + for object_seq in 1..=3 { + assert!(releases.contains_key(&object_seq)); + } + for release in releases.into_values() { + release.send(true).unwrap(); + } + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_eq!(0, failing_stores()); + assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + assert_eq!(id(3, 1), latest(&store, region_id)); + } + + #[tokio::test] + async fn test_store_enqueued_conflict_poisons() { let _serialized = SERIALIZED.lock().await; // A conflict of any epoch poisons the store without reading the // existing object, since the acknowledged entries cannot move to - // another sequence. So does a permanent storage error. - for foreign_epoch in [Some(0), Some(2), None] { + // another sequence. + for foreign_epoch in [0, 2] { let object_store = memory_store(); let (io, reads) = RecordingIo::over(object_store.clone()); let store = open_over(io.clone(), &enqueued(eager())).await; let region_id = region(1); - match foreign_epoch { - Some(epoch) => put_foreign(&object_store, 1, epoch).await, - None => io.fail_next_put_permanently(), - } + put_foreign(&object_store, 1, foreign_epoch).await; let assert_poisoned = |error: Error| { - let poisoned = match foreign_epoch { - Some(_) => matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), - None => matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), - }; - assert!(poisoned, "unexpected error: {error:?}"); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); }; // The append was acknowledged; the failure surfaces afterwards. @@ -6642,8 +7316,8 @@ mod tests { read_entries(&store, region_id, 1).await ); - // A backlog that cannot be uploaded is reported by stop. A transient - // failure, temporary or persistent, drops it; a permanent one also + // A backlog that cannot be uploaded is dropped and reported by stop, + // whether the failure is temporary, persistent or permanent; none // poisons the store. for failure in ["temporary", "persistent", "permanent"] { let (io, _) = RecordingIo::over(memory_store()); @@ -6660,10 +7334,7 @@ mod tests { "unexpected error: {error:?}" ); assert_eq!(vec![0], object_seqs(io.as_ref()).await); - assert_eq!( - failure == "permanent", - store.latest_entry_id(&provider(region_id)).is_err() - ); + assert!(terminal(&store.terminal_error).is_none()); store.stop().await.unwrap(); } } @@ -7382,11 +8053,13 @@ mod tests { /// Object access whose whole-object reads, or conditional creates, and /// on request deletes and range reads, park until the test releases /// them, counting how many are in flight. A release of false fails the - /// operation with a transient error before it reaches the object store. + /// operation with a transient error before it reaches the object store, + /// or a create with denied access once `creates_refused` is set. struct ParkedIo { inner: ObjectStoreIo, parked: mpsc::UnboundedSender<(u64, oneshot::Sender)>, park_creates: bool, + creates_refused: AtomicBool, /// Whether creates park yet; set once the store opened. creates_parked: AtomicBool, /// Whether deletes park. @@ -7428,6 +8101,7 @@ mod tests { inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), parked, park_creates, + creates_refused: AtomicBool::new(false), creates_parked: AtomicBool::new(false), deletes_parked: AtomicBool::new(false), range_reads_parked: AtomicBool::new(false), @@ -7458,6 +8132,17 @@ mod tests { && self.creates_parked.load(Ordering::SeqCst) && !self.park(object_seq).await { + if self.creates_refused.load(Ordering::SeqCst) { + return Err(object_store::Error::new( + ErrorKind::PermissionDenied, + "access denied", + ) + .set_persistent()) + .context(WalObjectStoreSnafu { + operation: "write", + path: self.object_path(object_seq), + }); + } return injected_failure("write", self.object_path(object_seq)); } self.inner.put_if_absent(object_seq, content).await @@ -7660,8 +8345,10 @@ mod tests { EarlierEpochDurable, /// transient I/O error, `enqueued` mode TransientEnqueued, - /// transient I/O error, `enqueued` mode after `stop` began - TransientEnqueuedAfterStop, + /// I/O error that requires intervention, `enqueued` mode + InterventionEnqueued, + /// any I/O error, `enqueued` mode after `stop` began + ErrorEnqueuedAfterStop, /// an object of the same or a later epoch holds the sequence, or /// any conflicting object in the `enqueued` mode Conflict, @@ -7677,12 +8364,13 @@ mod tests { } /// Every row of the matrix, in the order the RFC lists them. - const ROWS: [Row; 10] = [ + const ROWS: [Row; 11] = [ Row::Created, Row::TransientDurable, Row::EarlierEpochDurable, Row::TransientEnqueued, - Row::TransientEnqueuedAfterStop, + Row::InterventionEnqueued, + Row::ErrorEnqueuedAfterStop, Row::Conflict, Row::CatalogError, Row::LastSequence, @@ -7746,6 +8434,7 @@ mod tests { skipped_segments: METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(), stalled_appends: METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get(), stalled_waits: METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.get_sample_count(), + rejected_appends: METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL.get(), } /// Samples of the two object shape histograms, which a sealed batch @@ -7941,11 +8630,21 @@ mod tests { row: Some(Row::TransientEnqueued), run: || Box::pin(create_fails_transiently_enqueued()), }, + Case { + fault: "creates the object store refuses with an error that requires intervention, `enqueued` mode", + row: Some(Row::InterventionEnqueued), + run: || Box::pin(creates_refused_enqueued()), + }, Case { fault: "a create that fails transiently after `stop` began, `enqueued` mode", - row: Some(Row::TransientEnqueuedAfterStop), + row: Some(Row::ErrorEnqueuedAfterStop), run: || Box::pin(create_fails_transiently_after_stop_began()), }, + Case { + fault: "a stop while the object store refuses the backlog, `enqueued` mode", + row: Some(Row::ErrorEnqueuedAfterStop), + run: || Box::pin(stop_while_the_backlog_is_refused()), + }, Case { fault: "a create that finds an object of an earlier epoch, `durable` mode", row: Some(Row::EarlierEpochDurable), @@ -8061,6 +8760,11 @@ mod tests { row: None, run: || Box::pin(create_conflicts_after_stop_began()), }, + Case { + fault: "an append held back by full sealed batches, `durable` mode", + row: None, + run: || Box::pin(append_held_back_by_full_sealed_batches()), + }, Case { fault: "a stalled append released by admission, then one released by a stop", row: None, @@ -8414,6 +9118,157 @@ mod tests { ); } + /// Objects 1 to 3 are acknowledged when the object store starts + /// denying access. The store keeps them under their sequences, + /// rejects the next append and attempts object 1 alone; once it is + /// created the others follow and appends are admitted again. + async fn creates_refused_enqueued() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let before = Counters::sample(); + let region_id = region(1); + for append in spawn_appends(&store, region_id, 3).await { + timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + } + let mut releases = parked_operations(&mut parked, 3).await; + io.creates_refused.store(true, Ordering::SeqCst); + for object_seq in 1..=3 { + releases.remove(&object_seq).unwrap().send(false).unwrap(); + } + wait_until(|| Counters::since(before).create_failures == 3).await; + assert_eq!(1, failing_stores()); + let error = append(&store, region_id, "a4").await.unwrap_err(); + assert_rejected(&error, ErrorKind::PermissionDenied); + assert_eq!(0, store.durable_entry_id(&provider(region_id)).unwrap()); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert!(terminal(&store.terminal_error).is_none()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 3, + sealed_objects: 3, + object_bytes: 3 * one_entry_object(), + object_entries: 3, + reads: 1, + rejected_appends: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert!(parked.try_recv().is_err()); + release.send(true).unwrap(); + for release in parked_operations(&mut parked, 2).await.into_values() { + release.send(true).unwrap(); + } + timeout(WAIT, store.wait_durable(&provider(region_id), id(3, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(0, failing_stores()); + let response = append(&store, region_id, "a5").await.unwrap(); + assert_eq!(Some(&id(4, 1)), response.last_entry_ids.get(®ion_id)); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(4, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![1, 2, 3, 4], collected_object_seqs(&store).await); + assert_eq!( + expected_entries( + region_id, + &[ + (id(1, 1), "a1"), + (id(2, 1), "a2"), + (id(3, 1), "a3"), + (id(4, 1), "a5") + ] + ), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[1, 2, 3, 4]).await; + assert_eq!( + Counters { + created_objects: 4, + seal_to_durable: 4, + create_failures: 3, + sealed_objects: 4, + object_bytes: 4 * one_entry_object(), + object_entries: 4, + reads: 2, + deleted_objects: 1, + rejected_appends: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// The object store denies access to objects 1 and 2 and still does + /// when `stop` begins: stop attempts both together, drops the backlog + /// when they fail and reports the refusal; the store is not poisoned. + async fn stop_while_the_backlog_is_refused() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let before = Counters::sample(); + let region_id = region(1); + for append in spawn_appends(&store, region_id, 2).await { + timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + } + io.creates_refused.store(true, Ordering::SeqCst); + for release in parked_operations(&mut parked, 2).await.into_values() { + release.send(false).unwrap(); + } + wait_until(|| Counters::since(before).create_failures == 2).await; + assert_eq!(1, failing_stores()); + let wait = { + let store = store.clone(); + tokio::spawn( + async move { store.wait_durable(&provider(region_id), id(1, 1)).await }, + ) + }; + timeout(WAIT, store.wait_for_durability_waits(1)) + .await + .unwrap() + .unwrap(); + + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + for release in parked_operations(&mut parked, 2).await.into_values() { + release.send(false).unwrap(); + } + let error = timeout(WAIT, stop).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { error, .. } + if error.kind() == ErrorKind::PermissionDenied), + "unexpected error: {error:?}" + ); + assert_stopped(&timeout(WAIT, wait).await.unwrap().unwrap().unwrap_err()); + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!(0, failing_stores()); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 4, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + } + async fn create_finds_an_earlier_epoch() { let object_store = memory_store(); let store = open(object_store.clone(), &eager()).await; @@ -9581,6 +10436,63 @@ mod tests { ); } + /// Seven sealed batches leave no room for the two an admission can + /// seal: the next append is held back and counted as stalled, and it + /// is admitted once object 1 is durable, which records its wait. + async fn append_held_back_by_full_sealed_batches() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let mut appends = spawn_appends(&store, region_id, MAX_SEALED_BATCHES - 1).await; + appends.push(spawn_append_batch( + &store, + vec![entry(&store, region_id, "h1")], + )); + let releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + wait_until(|| Counters::since(before).stalled_appends == 1).await; + assert_eq!(0, Counters::since(before).stalled_waits); + + for release in releases.into_values() { + release.send(true).unwrap(); + } + for _ in MAX_IN_FLIGHT_CREATES..MAX_SEALED_BATCHES { + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + } + for (index, append) in appends.into_iter().enumerate() { + let response = timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + assert_eq!( + Some(&id(index as u64 + 1, 1)), + response.last_entry_ids.get(®ion_id) + ); + } + let sealed = MAX_SEALED_BATCHES as u64; + assert_eq!( + (1..=sealed).collect::>(), + collected_object_seqs(&store).await + ); + assert_indexed(io.as_ref(), &(1..=sealed).collect::>()).await; + assert_eq!( + Counters { + created_objects: sealed, + seal_to_durable: sealed, + sealed_objects: sealed, + object_bytes: sealed * one_entry_object(), + object_entries: sealed, + acknowledged_appends: sealed, + deleted_objects: 1, + stalled_appends: 1, + stalled_waits: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + /// The `enqueued` configuration under which the second append stalls /// behind the create of the first. fn stalling() -> ObjectStoreWalConfig { diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs index 126e38e4b26..54626b24432 100644 --- a/src/mito2/src/engine/object_store_wal_recovery_test.rs +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -22,12 +22,12 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::Ordering; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use api::v1::Rows; use common_base::readable_size::ReadableSize; -use common_error::ext::{BoxedError, ErrorExt}; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; use common_wal::config::object_store::{ @@ -35,9 +35,12 @@ use common_wal::config::object_store::{ }; use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; use log_store::ObjectStoreLogStore; +use log_store::metrics::METRIC_OBJECT_STORE_WAL_FAILING; use log_store::object_store_wal::{WalHole, entry_id}; use object_store::ObjectStore; +use object_store::layers::mock::{self, MockLayerBuilder, oio}; use object_store::services::Memory; +use object_store::util::with_retry_layers; use rstest::rstest; use store_api::logstore::LogStore; use store_api::logstore::provider::Provider; @@ -1355,3 +1358,106 @@ async fn test_flush_does_not_publish_a_frontier_for_a_lost_enqueued_backlog() { ); assert!(store.stop().await.is_err()); } + +/// An in-memory object store behind the retry layer object storage has in +/// production, whose writes fail with denied access while `denied` is set. +fn denying_store(denied: Arc) -> ObjectStore { + let mock = MockLayerBuilder::default() + .writer_factory(Arc::new(move |_, _, inner| { + if denied.load(Ordering::SeqCst) { + Box::new(DeniedWriter) + } else { + inner + } + })) + .build() + .unwrap(); + with_retry_layers(memory_store().layer(mock)) +} + +struct DeniedWriter; + +impl oio::Write for DeniedWriter { + async fn write(&mut self, _buffer: mock::Buffer) -> mock::Result<()> { + Ok(()) + } + + async fn close(&mut self) -> mock::Result { + Err(mock::Error::new( + mock::ErrorKind::PermissionDenied, + "access denied", + )) + } + + async fn abort(&mut self) -> mock::Result<()> { + Ok(()) + } +} + +#[tokio::test] +async fn test_enqueued_writes_fail_fast_while_the_object_store_refuses_the_backlog() { + let mut env = TestEnv::with_prefix("object-store-wal-refused-backlog").await; + let denied = Arc::new(AtomicBool::new(false)); + let store = open_store(&denying_store(denied.clone()), AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema) = create_region(&engine, REGION_A, &[]).await; + + // The write is acknowledged, then the object store denies the create of + // its object, which the store keeps. + denied.store(true, Ordering::SeqCst); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + let seal = spawn_seal(&store); + tokio::time::timeout(WAIT, async { + while METRIC_OBJECT_STORE_WAL_FAILING.get() == 0 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .expect("the store must fail"); + + // A later write fails at once with a retryable error that names the + // refusal, and takes no entry id. + let error = try_put(&engine, REGION_A, rows(&schema, 3, 5)) + .await + .unwrap_err(); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + assert!( + format!("{error:?}").contains("PermissionDenied"), + "{error:?}" + ); + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + + // A flush waits for the acknowledged entry instead of failing. + let flush = spawn_flush(&engine, REGION_A); + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!flush.is_finished()); + + // Once the object store accepts writes again, the backlog is uploaded, + // the flush records the entry and writes are accepted again. + denied.store(false, Ordering::SeqCst); + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .unwrap(); + tokio::time::timeout(WAIT, flush) + .await + .unwrap() + .unwrap() + .unwrap(); + put_rows(&engine, REGION_A, rows(&schema, 3, 5)).await; + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(1, 1), + last_entry_id: entry_id(2, 1), + topic_latest_entry_id: entry_id(1, 1), + manifest_flushed_entry_id: entry_id(1, 1), + memtable_rows: 2, + }, + region_recovery_state(&engine, REGION_A).await + ); +} From 4ff3296f975c67214e0dd337540170de525f459d Mon Sep 17 00:00:00 2001 From: sun Date: Wed, 14 Oct 2026 06:14:58 +0000 Subject: [PATCH 39/46] feat: update dashboard to v0.14.2 (#9570) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index 4a29f93bbe5..b0d2e474f72 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.14.0 +v0.14.2 From 00aa65de7a9a3e72bc8674cf3a54b7f668809edf Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 14 Oct 2026 06:49:35 +0000 Subject: [PATCH 40/46] fix(query): guard deep expressions during query planning (#9564) * fix(query): avoid speculative clones in constant normalization Signed-off-by: evenyag * fix(query): reject deeply nested expressions before optimization Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/promql/src/extension_plan/empty_metric.rs | 5 + src/query/src/datafusion.rs | 5 + src/query/src/dist_plan.rs | 2 +- src/query/src/dist_plan/analyzer.rs | 3 + src/query/src/dist_plan/planner.rs | 2 + src/query/src/optimizer.rs | 1 + .../src/optimizer/const_normalization.rs | 189 +++++++-------- src/query/src/optimizer/expr_depth.rs | 223 ++++++++++++++++++ src/query/src/planner.rs | 2 + src/query/src/query_engine/state.rs | 9 +- tests-integration/tests/sql.rs | 51 ++++ .../common/filter/expression_depth.result | 24 ++ .../common/filter/expression_depth.sql | 9 + .../common/tql-explain-analyze/explain.result | 3 + 14 files changed, 433 insertions(+), 95 deletions(-) create mode 100644 src/query/src/optimizer/expr_depth.rs create mode 100644 tests/cases/standalone/common/filter/expression_depth.result create mode 100644 tests/cases/standalone/common/filter/expression_depth.sql diff --git a/src/promql/src/extension_plan/empty_metric.rs b/src/promql/src/extension_plan/empty_metric.rs index 4d2e029417a..a534bf20493 100644 --- a/src/promql/src/extension_plan/empty_metric.rs +++ b/src/promql/src/extension_plan/empty_metric.rs @@ -109,6 +109,11 @@ impl EmptyMetric { }) } + /// Returns the expression used to generate metric values. + pub fn expr(&self) -> Option<&Expr> { + self.expr.as_ref() + } + pub const fn name() -> &'static str { "EmptyMetric" } diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 970f3109804..72788461836 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -68,6 +68,7 @@ use crate::metrics::{ OnDone, QUERY_STAGE_ELAPSED, maybe_attach_region_watermark_metrics, should_collect_region_watermark_from_query_ctx, }; +use crate::optimizer::expr_depth; use crate::options::ScheduledTimeExtension; use crate::physical_wrapper::PhysicalPlanWrapperRef; use crate::planner::{DfLogicalPlanner, LogicalPlanner}; @@ -381,6 +382,8 @@ impl DatafusionQueryEngine { } } + expr_depth::validate_plan(logical_plan)?; + let mut logger = PanicLogger { input_logical_plan: logical_plan, after_analyze: None, @@ -408,6 +411,7 @@ impl DatafusionQueryEngine { |_, _| {}, )?; + expr_depth::validate_plan(&analyzed_plan)?; logger.after_analyze = Some(analyzed_plan.clone()); common_telemetry::debug!("Create physical plan, analyzed plan: {analyzed_plan}"); @@ -423,6 +427,7 @@ impl DatafusionQueryEngine { .optimize(analyzed_plan, state, |_, _| {})? }; + expr_depth::validate_plan(&optimized_plan)?; common_telemetry::debug!("Create physical plan, optimized plan: {optimized_plan}"); logger.after_optimize = Some(optimized_plan.clone()); diff --git a/src/query/src/dist_plan.rs b/src/query/src/dist_plan.rs index 969f00e7c9c..a982eb1d8bf 100644 --- a/src/query/src/dist_plan.rs +++ b/src/query/src/dist_plan.rs @@ -27,7 +27,7 @@ mod remote_dyn_filter_registry; pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions}; pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId}; pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan}; -pub(crate) use merge_sort::MergeSortExec; +pub(crate) use merge_sort::{MergeSortExec, MergeSortLogicalPlan}; pub use planner::{DistExtensionPlanner, MergeSortExtensionPlanner}; pub use predicate_extractor::PredicateExtractor; pub use region_pruner::ConstraintPruner; diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 38d386cc818..a39514a0e43 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -53,6 +53,7 @@ use crate::dist_plan::commutativity::{ use crate::dist_plan::merge_scan::MergeScanLogicalPlan; use crate::dist_plan::merge_sort::MergeSortLogicalPlan; use crate::metrics::PUSH_DOWN_FALLBACK_ERRORS_TOTAL; +use crate::optimizer::expr_depth; use crate::options::ScheduledTimeExtension; use crate::plan::ExtractExpr; use crate::query_engine::DefaultSerializer; @@ -118,6 +119,7 @@ impl AnalyzerRule for DistPlannerAnalyzer { plan: LogicalPlan, config: &ConfigOptions, ) -> datafusion_common::Result { + expr_depth::validate_plan(&plan)?; let mut config = config.clone(); // Aligned with the behavior in `datafusion_optimizer::OptimizerContext::new()`. config.optimizer.filter_null_join_keys = true; @@ -142,6 +144,7 @@ impl AnalyzerRule for DistPlannerAnalyzer { let plan = plan .rewrite_with_subqueries(&mut PlanTreeExpressionSimplifier::new(optimizer_context))? .data; + expr_depth::validate_plan(&plan)?; let fallback_plan = plan.clone(); // Run a filter-focused optimizer subset before MergeScan wraps remote diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 88b13caf388..34c0994f9e0 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -50,6 +50,7 @@ use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan}; use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan}; use crate::dist_plan::region_pruner::ConstraintPruner; use crate::error::{CatalogSnafu, PartitionRuleManagerSnafu, TableNotFoundSnafu}; +use crate::optimizer::expr_depth; use crate::region_query::RegionQueryHandlerRef; /// Planner for converting merge sort logical plan to physical plan. @@ -447,6 +448,7 @@ impl DistExtensionPlanner { session: &dyn Session, plan: &LogicalPlan, ) -> Result { + expr_depth::validate_plan(plan)?; let session_state = session .as_any() .downcast_ref::() diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 03d6ce3aa46..35a3cac2639 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -17,6 +17,7 @@ pub mod constant_term; pub mod count_nest_aggr; pub mod count_wildcard; pub mod enforce_sorting; +pub(crate) mod expr_depth; pub mod global_limit; pub(crate) mod insert_assignment; pub(crate) mod json_schema_concretize; diff --git a/src/query/src/optimizer/const_normalization.rs b/src/query/src/optimizer/const_normalization.rs index d29f0aa4f39..64d39a5f668 100644 --- a/src/query/src/optimizer/const_normalization.rs +++ b/src/query/src/optimizer/const_normalization.rs @@ -23,8 +23,6 @@ use datafusion_expr::{Between, BinaryExpr, Expr, ExprSchemable, LogicalPlan, Ope use datafusion_expr_common::casts::try_cast_literal_to_type; use datafusion_optimizer::analyzer::AnalyzerRule; -use crate::plan::ExtractExpr; - /// ConstNormalizationRule rewrites castable constants against their /// non-constant comparison operand ahead of filter pushdown. #[derive(Debug)] @@ -52,26 +50,12 @@ impl AnalyzerRule for ConstNormalizationRule { } fn rewrite_plan_exprs(plan: LogicalPlan, schema: DFSchemaRef) -> Result> { - let mut rewriter = ConstNormalizationRewriter { - schema, - transformed: false, - }; - let exprs = plan - .expressions_consider_join() - .into_iter() - .map(|expr| expr.rewrite(&mut rewriter).map(|rewritten| rewritten.data)) - .collect::>>()?; - if !rewriter.transformed { - return Ok(Transformed::no(plan)); - } - - let inputs = plan.inputs().into_iter().cloned().collect::>(); - plan.with_new_exprs(exprs, inputs).map(Transformed::yes) + let mut rewriter = ConstNormalizationRewriter { schema }; + plan.map_expressions(|expr| expr.rewrite(&mut rewriter)) } struct ConstNormalizationRewriter { schema: DFSchemaRef, - transformed: bool, } impl TreeNodeRewriter for ConstNormalizationRewriter { @@ -91,44 +75,36 @@ impl TreeNodeRewriter for ConstNormalizationRewriter { } fn f_up(&mut self, expr: Expr) -> Result> { - let rewritten = rewrite_expr_node(expr, &self.schema)?; - self.transformed |= rewritten.transformed; - Ok(rewritten) + rewrite_expr_node(expr, &self.schema) } } fn rewrite_expr_node(expr: Expr, schema: &DFSchemaRef) -> Result> { - match expr { - Expr::BinaryExpr(binary) => match rewrite_binary_expr(binary.clone(), schema)? { - Some(expr) => Ok(Transformed::yes(expr)), - None => Ok(Transformed::no(Expr::BinaryExpr(binary))), - }, - Expr::Between(between) => match rewrite_between_expr(between.clone(), schema)? { - Some(expr) => Ok(Transformed::yes(expr)), - None => Ok(Transformed::no(Expr::Between(between))), - }, - Expr::InList(in_list) => match rewrite_in_list_expr(in_list.clone(), schema)? { - Some(expr) => Ok(Transformed::yes(expr)), - None => Ok(Transformed::no(Expr::InList(in_list))), - }, - Expr::Like(like) => rewrite_like_expr(like, PatternMatchKind::Like, schema), - Expr::SimilarTo(like) => rewrite_like_expr(like, PatternMatchKind::SimilarTo, schema), - expr => Ok(Transformed::no(expr)), - } + // Borrow while checking eligibility: cloning unchanged subtrees here can overflow + // the stack even though DataFusion protects the surrounding rewrite traversal. + let rewritten = match &expr { + Expr::BinaryExpr(binary) => rewrite_binary_expr(binary, schema)?, + Expr::Between(between) => rewrite_between_expr(between, schema)?, + Expr::InList(in_list) => rewrite_in_list_expr(in_list, schema)?, + Expr::Like(like) => rewrite_like_expr(like, PatternMatchKind::Like, schema)?, + Expr::SimilarTo(like) => rewrite_like_expr(like, PatternMatchKind::SimilarTo, schema)?, + _ => None, + }; + Ok(match rewritten { + Some(rewritten) => Transformed::yes(rewritten), + None => Transformed::no(expr), + }) } -fn rewrite_between_expr(between: Between, schema: &DFSchemaRef) -> Result> { +fn rewrite_between_expr(between: &Between, schema: &DFSchemaRef) -> Result> { let Between { expr, negated, low, high, } = between; - let expr = *expr; - let low_expr = *low; - let high_expr = *high; let Some((target, constants)) = - extract_rewrite_operands(&expr, &[low_expr.clone(), high_expr.clone()], schema)? + extract_rewrite_operands(expr, [low.as_ref(), high.as_ref()], schema)? else { return Ok(None); }; @@ -142,7 +118,7 @@ fn rewrite_between_expr(between: Between, schema: &DFSchemaRef) -> Result