mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-07 22:18:57 +00:00
fix: backport v1.2.0 correctness and liveness fixes (#9023)
* fix(mito2): fence checkpoints during region transitions (#8847) * fix: fence checkpoints during region transitions Signed-off-by: WenyXu <wenymedia@gmail.com> * test(datanode): fix transient downgrade setup Signed-off-by: WenyXu <wenymedia@gmail.com> * test(mito2): fix checkpoint lifecycle test setup Signed-off-by: WenyXu <wenymedia@gmail.com> * test(mito2): cover cancelled downgrade waiter retry Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(mito2): fence direct follower transitions Signed-off-by: WenyXu <wenymedia@gmail.com> * test: trim checkpoint transition coverage Signed-off-by: WenyXu <wenymedia@gmail.com> * refactor(mito2): clarify checkpoint task lifecycle Signed-off-by: WenyXu <wenymedia@gmail.com> --------- Signed-off-by: WenyXu <wenymedia@gmail.com> (cherry picked from commita502dfdefd) * fix(mito2): split SSTs at primary key series boundaries (#8888) * fix(mito2): split SSTs at series boundaries Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * test(mito2): cover SST splitting without primary key Also document the sortedness precondition and the series boundary split semantics on write_all_flat/write_all_flat_as_primary_key and the new split helpers. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * test(mito2): avoid per-row Vec allocation for empty primary key Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * Update src/mito2/src/sst/parquet/writer.rs Co-authored-by: Yingwen <realevenyag@gmail.com> --------- Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> Co-authored-by: Yingwen <realevenyag@gmail.com> (cherry picked from commit4ac3423261) * fix(meta): allow manual migration from offline datanodes (#8934) * fix(meta): allow migration from offline datanodes Signed-off-by: WenyXu <wenymedia@gmail.com> * test: fix offline migration event actor Signed-off-by: WenyXu <wenymedia@gmail.com> * test: read migration routes from metadata Signed-off-by: WenyXu <wenymedia@gmail.com> --------- Signed-off-by: WenyXu <wenymedia@gmail.com> (cherry picked from commit8a473c5bf0) * fix(flight): bound DoGet response wait (#8943) * fix(flight): defer datanode query initialization Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(client): retain Flight stream peer context Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(client): improve Flight stream diagnostics Signed-off-by: WenyXu <wenymedia@gmail.com> --------- Signed-off-by: WenyXu <wenymedia@gmail.com> (cherry picked from commit28398138ec) * fix(query): keep INSERT timestamp conversion out of the source query (#8911) * fix(query): keep INSERT timestamp conversion out of the source query Interpreting an INSERT's string timestamps used to work by pushing the conversion down into the source query, which changed what that query means. Two consequences: - Pushing through a UNION's DISTINCT moved the dedup key from the raw strings to parsed instants, so rows spelling the same instant differently collapsed into one. On an append-only table that is a silently dropped row. - A UNION branch that needed no conversion (a NULL, or an explicit cast) made the whole column give up, leaving sibling branches on UTC while the rest of the row used the session timezone. Convert at the assignment instead, by routing its cast through a timezone-carrying timestamp type and back. Arrow applies the timezone when a cast target carries one, and stripping it afterwards preserves the value. The source query is no longer touched, so both cases go away and the tree-walking rewrite (roughly 160 lines) is deleted. The rewrite reads source types, so it now runs TypeCoercion first: a UNION still carries its loose per-branch schema before coercion, and retargeting a cast whose input later becomes a timestamp would shift the value rather than reinterpret it. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix(query): address review on INSERT assignment rewrite - Clone the input `Arc` instead of the whole subtree, and only rebuild it when a `Values` row actually changes. - Defer cloning the cast source until the literal-folding path has been ruled out. - Move the UTC check onto `Timezone::is_utc`, replacing a bare string compare. - Cover a prepared `INSERT ... VALUES (?)`: an untyped placeholder types as `Null`, so the assignment cast is left for parameter substitution. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commit1851f6bf4d) * test: renew etcd TLS certificates (#8956) Signed-off-by: WenyXu <wenymedia@gmail.com> (cherry picked from commit144f83528d) * fix: update tokio-postgres and correct explain/fetch cursor output schema (#8955) * chore(deps): update tokio-postgres * fix: describing fetch cursor and analyze (cherry picked from commitb31f05eb59) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): avoid insert select HTTP/2 stalls Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commitd4a3d88cd7) * fix(query): share record batch forwarding for DML Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commitc34bec197c) * fix(mysql): interpret prepared statement datetime params in session timezone (#8923) * fix(mysql): interpret prepared statement datetime params in session timezone Binary DATETIME parameters of server-side prepared statements were converted as if UTC, ignoring the session timezone set via SET time_zone. Convert them with the session timezone and add an integration test covering prepared inserts and predicates under Asia/Shanghai. Signed-off-by: wy471x <wy471x@gmail.com> * refactor: share naive datetime timezone policy via common-time Address review feedback on the prepared-statement timezone fix: - Expose Timestamp::from_naive_datetime in common-time so the DST policy (gap -> error, ambiguous -> earlier instant) lives in one place, shared by the text protocol (Timestamp::from_str) and the MySQL binary protocol. - Route the MySQL prepared-statement datetime conversion through it. - Match the target type before converting datetime params so PreparedStmtTypeMismatch fails fast without wasted conversion. - Use the short Timezone import form for consistency with the rest of servers. Signed-off-by: wy471x <wy471x@gmail.com> --------- Signed-off-by: wy471x <wy471x@gmail.com> Co-authored-by: Ning Sun <sunng@protonmail.com> (cherry picked from commitaaa843104b) * fix(mito2): use target sequence for foreign SSTs (#8946) * fix(mito2): use target sequence for foreign SSTs Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore(mito2): address foreign SST review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commitbd7d2c1dfa) * fix: postgres describe for more statements (#8974) * fix: postgres describe for more statements Signed-off-by: Ning Sun <sunning@greptime.com> * fix: cover more show statements Signed-off-by: Ning Sun <sunning@greptime.com> * fix: address review comments - add missing `clippy::too_many_arguments` allow on `query_from_information_schema_dataframe` (CI clippy failure) - take `&ShowKind` in the information-schema dataframe helper so `kind` is no longer cloned at every call site; only the WHERE arm (which needs an owned expression for `sql_to_expr`) clones internally - document why re-applying TQL explain formats never overwrites an existing value (per-query context state) Signed-off-by: Ning Sun <sunning@greptime.com> * chore: trim comments to essentials Signed-off-by: Ning Sun <sunning@greptime.com> --------- Signed-off-by: Ning Sun <sunning@greptime.com> (cherry picked from commitd32cd77505) * fix(promql): resolve derived labels in aggregation arithmetic (#8994) Signed-off-by: shuiyisong <xixing.sys@gmail.com> (cherry picked from commitc4dafb5826) * fix(deps): narrow postgres lockfile update Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(deps): normalize remaining x509 lock reference Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(client): isolate query and control transports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(frontend): clarify adaptive window purpose Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(grpc): satisfy clippy in transport isolation coverage Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(pipeline): coalesce concurrent pipeline cache misses (#9022) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * Revert "test(grpc): satisfy clippy in transport isolation coverage" This reverts commit888bcedc96. (cherry picked from commit e0216b4074485bcffaae53348a17091532cb3600) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * Revert "docs(frontend): clarify adaptive window purpose" This reverts commit84cf69d099. (cherry picked from commit da30e605ad7ebfb2a3fa86ea44b67405e5afde98) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * Revert "fix(client): isolate query and control transports" This reverts commita1d2dbb050. (cherry picked from commit 03f5c31f544d3d3852beddfa87cff34d04a646ae) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(json2): add bounded auto-expansion to the JSON2 vector builder (#8909) * refactor(json2): add bounded auto-expansion to the JSON2 vector builder Signed-off-by: luofucong <luofc@foxmail.com> * resolve PR comments Signed-off-by: luofucong <luofc@foxmail.com> * fix ci Signed-off-by: luofucong <luofc@foxmail.com> --------- Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commit76f08d2b3f) (cherry picked from commit 640fa63bd8ee098fac8e8d62d6fb7ca5af4ddcd0) * refactor(json2): optimize JSON2 building without auto-expanded paths (#8928) * refactor(json2): optimize JSON2 building without auto-expanded paths Signed-off-by: luofucong <luofc@foxmail.com> * resolve PR comments Signed-off-by: luofucong <luofc@foxmail.com> * avoid panicking memtable write Signed-off-by: luofucong <luofc@foxmail.com> --------- Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commit7fd0a7bb98) (cherry picked from commit 63c68ee60101e2a09e14420051f8f889b41a2a3b) * refactor(json2): support querying v2 storage layout (#8940) * feat(json2): support querying v2 storage layout - route missing JSON2 paths to the v2 remainder field - reconstruct complete values from explicit fields and remainder data - preserve root JSON2 columns across projections and filters - support nested JSON values in json_get string results - add and reorganize JSON2 sqlness coverage Signed-off-by: luofucong <luofc@foxmail.com> * resolve PR comments Signed-off-by: luofucong <luofc@foxmail.com> --------- Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commit932f87f7a8) (cherry picked from commit 5af1394a1f319bb99cafa9474b8a5a075be1e805) * feat(pipeline): support table-aware JSON2 transforms (#8964) * feat(pipeline): support table-aware JSON2 transforms Signed-off-by: shuiyisong <xixing.sys@gmail.com> * feat(pipeline): support JSON2 type hints in transforms Signed-off-by: shuiyisong <xixing.sys@gmail.com> * fix(pipeline): default failed JSON2 transforms to null Signed-off-by: shuiyisong <xixing.sys@gmail.com> * refactor(json2): distinguish invalid settings from layout errors Signed-off-by: shuiyisong <xixing.sys@gmail.com> --------- Signed-off-by: shuiyisong <xixing.sys@gmail.com> (cherry picked from commitfb86f6573e) (cherry picked from commit 18a21a9ce017b6887c857a3f3bbb9a83aa7f968b) * refactor: json2 v2 storage layout (#8979) * refactor: json2 v2 storage layout Signed-off-by: luofucong <luofc@foxmail.com> * resolve PR comments Signed-off-by: luofucong <luofc@foxmail.com> * fix ci Signed-off-by: luofucong <luofc@foxmail.com> * rethinking when "needs_remainder" Signed-off-by: luofucong <luofc@foxmail.com> * restore "ReadColumns" Signed-off-by: luofucong <luofc@foxmail.com> * resolve PR comments Signed-off-by: luofucong <luofc@foxmail.com> * fix ci Signed-off-by: luofucong <luofc@foxmail.com> --------- Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commit529f046110) (cherry picked from commit 7628d795197c786872730c3a5ef36c5c662d4322) * feat(json2): support JSON2 paths in SQL functions (#9007) feat(query): support JSON2 paths in SQL functions Update the DataFusion fork to expose scalar function planning hooks. Infer JSON2 path output types from scalar, aggregate, and window function signatures, while preserving the default Utf8View behavior for functions that accept arbitrary inputs. Add unit and sqlness coverage for type conflicts, mixed typed and untyped JSON paths, filters, aggregates, and window functions. Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commit15317a131b) (cherry picked from commit a945cbc7a4398e6bee74d4d2c67528772faf398d) * feat(json2): support empty and null JSON2 value (#9010) * feat(json2): support empty and null JSON2 value Signed-off-by: fys <fengys1996@gmail.com> * test(json2): cover explicit NULL and omitted-column inserts Signed-off-by: fys <fengys1996@gmail.com> * fix: cargo fmt Signed-off-by: fys <fengys1996@gmail.com> * fix: infer empty JSON object as object type Signed-off-by: fys <fengys1996@gmail.com> --------- Signed-off-by: fys <fengys1996@gmail.com> (cherry picked from commit05c65f54a8) (cherry picked from commit 05c27452ea5e958fab85e0ba4b82f5df205a5d08) * feat(json2): support list indexing for JSON2 columns (#9013) feat(query): support list indexing for JSON2 columns Extend JSON2 paths through DataFusion field-access planning, including nested list indexes and object fields following an index. Preserve Variant reads for bracket JSONPath expressions and normalize dot accesses after subscripts to work around the current DataFusion planner limitation. Add unit and sqlness coverage for nested indexes, type conflicts, missing paths, flushes, and compacted SSTs. Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commitd62a5a990a) (cherry picked from commit c11d8647feb30ecdd5ba1bc8ed4df78999801672) * refactor(json2): concretize JSON2 schemas at merge scan boundaries (#9016) * refactor(json2): concretize JSON2 schemas at merge scan boundaries Infer concrete JSON2 output types from remote plans and expose them on MergeScanLogicalPlan before physical planning. Recompute affected local schemas and remove the JSON2-specific rewrite from MergeScanExec. Add SQLness coverage for whole JSON2 columns in windows and joins. Signed-off-by: luofucong <luofc@foxmail.com> * fix ci Signed-off-by: luofucong <luofc@foxmail.com> --------- Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commit84bd993131) (cherry picked from commit 51c98a7e31a929fc33ff5133714bb47108f7539b) * fix(json2): keep empty structs in remainder (#9027) Signed-off-by: luofucong <luofc@foxmail.com> (cherry picked from commitd67d3501a9) (cherry picked from commit b3f2c1cc5bf0cfddb7c0670ac8f3522eb67c35a0) * fix(mito2): remove stale scan test imports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit d63bbe6aef298f5b8e960bb02eed08743a6684ee) * test(query): expect bounded merge scan partitions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit 7bcef46b1c010d2e73ec49be7581c216e6b5b1b6) * style(mito2): remove trailing whitespace Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(sqlness): refresh bounded partition plans Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(sqlness): normalize CPU-dependent partition plans Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: WenyXu <wenymedia@gmail.com> Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> Signed-off-by: Dennis Zhuang <killme2008@gmail.com> Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: wy471x <wy471x@gmail.com> Signed-off-by: Ning Sun <sunning@greptime.com> Signed-off-by: shuiyisong <xixing.sys@gmail.com> Signed-off-by: luofucong <luofc@foxmail.com> Signed-off-by: fys <fengys1996@gmail.com> Co-authored-by: WenyXu <wenymedia@gmail.com> Co-authored-by: Lei, HUANG <ratuthomm@gmail.com> Co-authored-by: Yingwen <realevenyag@gmail.com> Co-authored-by: dennis zhuang <killme2008@gmail.com> Co-authored-by: Ning Sun <sunng@protonmail.com> Co-authored-by: wy471x <wy471x@gmail.com> Co-authored-by: Ning Sun <sunning@greptime.com> Co-authored-by: shuiyisong <xixing.sys@gmail.com> Co-authored-by: luofucong <luofc@foxmail.com> Co-authored-by: fys <fengys1996@gmail.com>
This commit is contained in:
discord9
co-authored by
WenyXu
Lei, HUANG
Yingwen
dennis zhuang
Ning Sun
wy471x
Ning Sun
shuiyisong
luofucong
fys
WenyXu
Lei, HUANG
Yingwen
dennis zhuang
Ning Sun
wy471x
Ning Sun
shuiyisong
luofucong
fys
parent
460e0ec77e
commit
c60311973d
Generated
+174
-83
@@ -3740,7 +3740,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8"
|
||||
[[package]]
|
||||
name = "datafusion"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"arrow-schema 58.3.0",
|
||||
@@ -3794,7 +3794,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-catalog"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-trait",
|
||||
@@ -3818,7 +3818,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-catalog-listing"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-trait",
|
||||
@@ -3840,7 +3840,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-common"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"ahash 0.8.12",
|
||||
"arrow 58.3.0",
|
||||
@@ -3864,7 +3864,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-common-runtime"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"futures",
|
||||
"log",
|
||||
@@ -3874,7 +3874,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-datasource"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-compression",
|
||||
@@ -3908,7 +3908,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-datasource-arrow"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"arrow-ipc 58.3.0",
|
||||
@@ -3931,7 +3931,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-datasource-csv"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-trait",
|
||||
@@ -3953,7 +3953,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-datasource-json"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-trait",
|
||||
@@ -3976,7 +3976,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-datasource-parquet"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-trait",
|
||||
@@ -4005,12 +4005,12 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-doc"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-execution"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"arrow-buffer 58.3.0",
|
||||
@@ -4032,7 +4032,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-expr"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-trait",
|
||||
@@ -4054,7 +4054,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-expr-common"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"datafusion-common",
|
||||
@@ -4066,7 +4066,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-functions"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"arrow-buffer 58.3.0",
|
||||
@@ -4097,7 +4097,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-functions-aggregate"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"ahash 0.8.12",
|
||||
"arrow 58.3.0",
|
||||
@@ -4118,7 +4118,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-functions-aggregate-common"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"ahash 0.8.12",
|
||||
"arrow 58.3.0",
|
||||
@@ -4130,7 +4130,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-functions-nested"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"arrow-ord 58.3.0",
|
||||
@@ -4154,7 +4154,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-functions-table"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"async-trait",
|
||||
@@ -4169,7 +4169,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-functions-window"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"datafusion-common",
|
||||
@@ -4186,7 +4186,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-functions-window-common"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"datafusion-common",
|
||||
"datafusion-physical-expr-common",
|
||||
@@ -4195,7 +4195,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-macros"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"datafusion-doc",
|
||||
"quote",
|
||||
@@ -4205,7 +4205,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-optimizer"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"chrono",
|
||||
@@ -4254,7 +4254,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-physical-expr"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"ahash 0.8.12",
|
||||
"arrow 58.3.0",
|
||||
@@ -4277,7 +4277,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-physical-expr-adapter"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"datafusion-common",
|
||||
@@ -4291,7 +4291,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-physical-expr-common"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"ahash 0.8.12",
|
||||
"arrow 58.3.0",
|
||||
@@ -4307,7 +4307,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-physical-optimizer"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"datafusion-common",
|
||||
@@ -4325,7 +4325,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-physical-plan"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"ahash 0.8.12",
|
||||
"arrow 58.3.0",
|
||||
@@ -4356,7 +4356,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-proto"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"chrono",
|
||||
@@ -4383,7 +4383,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-proto-common"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"datafusion-common",
|
||||
@@ -4393,7 +4393,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-pruning"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"datafusion-common",
|
||||
@@ -4409,7 +4409,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-session"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"datafusion-common",
|
||||
@@ -4422,7 +4422,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-sql"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"arrow 58.3.0",
|
||||
"bigdecimal 0.4.8",
|
||||
@@ -4440,7 +4440,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "datafusion-substrait"
|
||||
version = "53.1.0"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147"
|
||||
source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a"
|
||||
dependencies = [
|
||||
"async-recursion",
|
||||
"async-trait",
|
||||
@@ -4630,10 +4630,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
|
||||
dependencies = [
|
||||
"const-oid 0.9.6",
|
||||
"der_derive",
|
||||
"flagset",
|
||||
"pem-rfc7468",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "der_derive"
|
||||
version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8034092389675178f570469e6c3b0465d3d30b4505c294a6550db47f3c17ad18"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "deranged"
|
||||
version = "0.5.8"
|
||||
@@ -5337,6 +5350,12 @@ version = "0.5.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
|
||||
|
||||
[[package]]
|
||||
name = "flagset"
|
||||
version = "0.4.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b7ac824320a75a52197e8f2d787f6a38b6718bb6897a35142d749af3c0e8f4fe"
|
||||
|
||||
[[package]]
|
||||
name = "flatbuffers"
|
||||
version = "25.2.10"
|
||||
@@ -5979,7 +5998,7 @@ dependencies = [
|
||||
"cfg-if",
|
||||
"js-sys",
|
||||
"libc",
|
||||
"wasi",
|
||||
"wasi 0.11.1+wasi-snapshot-preview1",
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
@@ -7870,13 +7889,14 @@ checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de"
|
||||
|
||||
[[package]]
|
||||
name = "libredox"
|
||||
version = "0.1.4"
|
||||
version = "0.1.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1580801010e535496706ba011c15f8532df6b42297d2e471fec38ceadd8c0638"
|
||||
checksum = "8d8f1ea3f21fd3405dcaf6c9b5c1630af9afc422d9073ea39c5f6d6c772e08ed"
|
||||
dependencies = [
|
||||
"bitflags 2.12.1",
|
||||
"libc",
|
||||
"redox_syscall 0.5.13",
|
||||
"plain",
|
||||
"redox_syscall 0.9.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8546,7 +8566,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"log",
|
||||
"wasi",
|
||||
"wasi 0.11.1+wasi-snapshot-preview1",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
@@ -9331,6 +9351,15 @@ dependencies = [
|
||||
"objc2-encode",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "objc2-core-foundation"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536"
|
||||
dependencies = [
|
||||
"bitflags 2.12.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "objc2-encode"
|
||||
version = "4.1.0"
|
||||
@@ -9347,6 +9376,15 @@ dependencies = [
|
||||
"objc2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "objc2-system-configuration"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396"
|
||||
dependencies = [
|
||||
"objc2-core-foundation",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "object"
|
||||
version = "0.36.7"
|
||||
@@ -10259,6 +10297,9 @@ dependencies = [
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
"object_store",
|
||||
"parquet-variant",
|
||||
"parquet-variant-compute",
|
||||
"parquet-variant-json",
|
||||
"paste",
|
||||
"seq-macro",
|
||||
"simdutf8",
|
||||
@@ -10576,16 +10617,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"x509-certificate 0.25.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "phf"
|
||||
version = "0.11.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078"
|
||||
dependencies = [
|
||||
"phf_shared 0.11.3",
|
||||
"x509-certificate",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -10838,6 +10870,12 @@ version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
|
||||
|
||||
[[package]]
|
||||
name = "plain"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6"
|
||||
|
||||
[[package]]
|
||||
name = "plist"
|
||||
version = "1.7.2"
|
||||
@@ -10930,34 +10968,34 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "postgres-protocol"
|
||||
version = "0.6.8"
|
||||
version = "0.6.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "76ff0abab4a9b844b93ef7b81f1efc0a366062aaef2cd702c76256b5dc075c54"
|
||||
checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"byteorder",
|
||||
"bytes",
|
||||
"fallible-iterator",
|
||||
"hmac 0.12.1",
|
||||
"md-5 0.10.6",
|
||||
"hmac 0.13.0",
|
||||
"md-5 0.11.0",
|
||||
"memchr",
|
||||
"rand 0.9.4",
|
||||
"sha2 0.10.9",
|
||||
"rand 0.10.1",
|
||||
"sha2 0.11.0",
|
||||
"stringprep",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "postgres-types"
|
||||
version = "0.2.9"
|
||||
version = "0.2.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "613283563cd90e1dfc3518d548caee47e0e725455ed619881f5cf21f36de4b48"
|
||||
checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be"
|
||||
dependencies = [
|
||||
"array-init",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"fallible-iterator",
|
||||
"postgres-protocol",
|
||||
"serde",
|
||||
"serde_core",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
@@ -12047,6 +12085,15 @@ dependencies = [
|
||||
"bitflags 2.12.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.9.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d678d17679829e73d371e96880897e98fee2ded7acc0a50bdf8af2affa4b2fe5"
|
||||
dependencies = [
|
||||
"bitflags 2.12.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_users"
|
||||
version = "0.4.6"
|
||||
@@ -14076,7 +14123,7 @@ dependencies = [
|
||||
"stringprep",
|
||||
"thiserror 2.0.17",
|
||||
"tracing",
|
||||
"whoami",
|
||||
"whoami 1.6.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -14115,7 +14162,7 @@ dependencies = [
|
||||
"stringprep",
|
||||
"thiserror 2.0.17",
|
||||
"tracing",
|
||||
"whoami",
|
||||
"whoami 1.6.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -15132,6 +15179,27 @@ version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
|
||||
|
||||
[[package]]
|
||||
name = "tls_codec"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0de2e01245e2bb89d6f05801c564fa27624dbd7b1846859876c7dad82e90bf6b"
|
||||
dependencies = [
|
||||
"tls_codec_derive",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tls_codec_derive"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2d2e76690929402faae40aebdda620a2c0e25dd6d3b9afe48867dfd95991f4bd"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.52.3"
|
||||
@@ -15198,9 +15266,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tokio-postgres"
|
||||
version = "0.7.13"
|
||||
version = "0.7.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c95d533c83082bb6490e0189acaa0bbeef9084e60471b696ca6988cd0541fb0"
|
||||
checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"byteorder",
|
||||
@@ -15211,29 +15279,29 @@ dependencies = [
|
||||
"log",
|
||||
"parking_lot 0.12.4",
|
||||
"percent-encoding",
|
||||
"phf 0.11.3",
|
||||
"phf 0.13.1",
|
||||
"pin-project-lite",
|
||||
"postgres-protocol",
|
||||
"postgres-types",
|
||||
"rand 0.9.4",
|
||||
"socket2 0.5.10",
|
||||
"rand 0.10.1",
|
||||
"socket2 0.6.4",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"whoami",
|
||||
"whoami 2.1.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-postgres-rustls"
|
||||
version = "0.12.0"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04fb792ccd6bbcd4bba408eb8a292f70fc4a3589e5d793626f45190e6454b6ab"
|
||||
checksum = "4c2ad44aa0ae96db89c4742212ed41645b2f597311ff6e1945542a4d9fadc2fb"
|
||||
dependencies = [
|
||||
"ring",
|
||||
"rustls",
|
||||
"sha2 0.11.0",
|
||||
"tokio",
|
||||
"tokio-postgres",
|
||||
"tokio-rustls",
|
||||
"x509-certificate 0.23.1",
|
||||
"x509-cert",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -16250,6 +16318,15 @@ version = "0.11.1+wasi-snapshot-preview1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
|
||||
|
||||
[[package]]
|
||||
name = "wasi"
|
||||
version = "0.14.7+wasi-0.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c"
|
||||
dependencies = [
|
||||
"wasip2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasip2"
|
||||
version = "1.0.2+wasi-0.2.9"
|
||||
@@ -16274,6 +16351,15 @@ version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b"
|
||||
|
||||
[[package]]
|
||||
name = "wasite"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42"
|
||||
dependencies = [
|
||||
"wasi 0.14.7+wasi-0.2.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen"
|
||||
version = "0.2.122"
|
||||
@@ -16484,7 +16570,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6994d13118ab492c3c80c1f81928718159254c53c472bf9ce36f8dae4add02a7"
|
||||
dependencies = [
|
||||
"redox_syscall 0.5.13",
|
||||
"wasite",
|
||||
"wasite 0.1.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "whoami"
|
||||
version = "2.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"libredox",
|
||||
"objc2-system-configuration",
|
||||
"wasite 1.0.2",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
@@ -17093,22 +17191,15 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "x509-certificate"
|
||||
version = "0.23.1"
|
||||
name = "x509-cert"
|
||||
version = "0.2.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "66534846dec7a11d7c50a74b7cdb208b9a581cad890b7866430d438455847c85"
|
||||
checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94"
|
||||
dependencies = [
|
||||
"bcder",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"const-oid 0.9.6",
|
||||
"der",
|
||||
"hex",
|
||||
"pem",
|
||||
"ring",
|
||||
"signature",
|
||||
"spki",
|
||||
"thiserror 1.0.69",
|
||||
"zeroize",
|
||||
"tls_codec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
+16
-16
@@ -251,7 +251,7 @@ strum = { version = "0.27", features = ["derive"] }
|
||||
sysinfo = "0.33"
|
||||
tempfile = "3"
|
||||
tokio = { version = "1.47", features = ["full"] }
|
||||
tokio-postgres = "0.7"
|
||||
tokio-postgres = "0.7.18"
|
||||
tokio-rustls = { version = "0.26.2", default-features = false }
|
||||
tokio-stream = "0.1"
|
||||
tokio-util = { version = "0.7", features = ["io-util", "compat"] }
|
||||
@@ -346,21 +346,21 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git"
|
||||
rev = "5618e779cf2bb4755b499c630fba4c35e91898cb"
|
||||
|
||||
[patch.crates-io]
|
||||
datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" }
|
||||
datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" }
|
||||
sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x
|
||||
|
||||
[profile.release]
|
||||
|
||||
+96
-20
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use api::region::RegionResponse;
|
||||
use api::v1::ResponseHeader;
|
||||
@@ -24,7 +25,7 @@ use arc_swap::ArcSwapOption;
|
||||
use arrow_flight::Ticket;
|
||||
use async_stream::stream;
|
||||
use async_trait::async_trait;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_error::ext::{BoxedError, ErrorExt};
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_grpc::flight::{FlightDecoder, FlightMessage};
|
||||
use common_meta::error::{self as meta_error, Result as MetaResult};
|
||||
@@ -49,6 +50,8 @@ use crate::error::{
|
||||
use crate::flight::decode_flight_data;
|
||||
use crate::{Client, metrics};
|
||||
|
||||
const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct RegionRequester {
|
||||
client: Client,
|
||||
@@ -109,20 +112,24 @@ impl RegionRequester {
|
||||
let mut flight_client = self
|
||||
.client
|
||||
.make_flight_client(self.send_compression, self.accept_compression)?;
|
||||
// Limit Flight DoGet response time without limiting query stream execution.
|
||||
let addr = flight_client.addr().to_string();
|
||||
let mut request = tonic::Request::new(ticket);
|
||||
request.set_timeout(FLIGHT_DO_GET_TIMEOUT);
|
||||
let response = flight_client
|
||||
.mut_inner()
|
||||
.do_get(ticket)
|
||||
.do_get(request)
|
||||
.await
|
||||
.or_else(|e| {
|
||||
let tonic_code = e.code();
|
||||
let e: error::Error = e.into();
|
||||
error!(
|
||||
e; "Failed to do Flight get, addr: {}, code: {}",
|
||||
flight_client.addr(),
|
||||
addr,
|
||||
tonic_code
|
||||
);
|
||||
Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu {
|
||||
addr: flight_client.addr().to_string(),
|
||||
addr: addr.clone(),
|
||||
tonic_code,
|
||||
})
|
||||
})?;
|
||||
@@ -133,7 +140,7 @@ impl RegionRequester {
|
||||
let flight_message_stream = flight_data_stream
|
||||
.filter_map(move |flight_data| decode_flight_data(&mut decoder, flight_data));
|
||||
|
||||
recordbatches_from_flight_message_stream(flight_message_stream).await
|
||||
recordbatches_from_flight_message_stream(addr, flight_message_stream).await
|
||||
}
|
||||
|
||||
async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponse> {
|
||||
@@ -195,6 +202,7 @@ impl RegionRequester {
|
||||
}
|
||||
|
||||
async fn recordbatches_from_flight_message_stream<S>(
|
||||
addr: String,
|
||||
mut flight_message_stream: S,
|
||||
) -> Result<SendableRecordBatchStream>
|
||||
where
|
||||
@@ -204,13 +212,17 @@ where
|
||||
return IllegalFlightMessagesSnafu {
|
||||
reason: "Expect the response not to be empty",
|
||||
}
|
||||
.fail();
|
||||
.fail()
|
||||
.map_err(|error| flight_stream_error(&addr, error));
|
||||
};
|
||||
let FlightMessage::Schema(schema) = first_flight_message? else {
|
||||
let FlightMessage::Schema(schema) =
|
||||
first_flight_message.map_err(|e| flight_stream_error(&addr, e))?
|
||||
else {
|
||||
return IllegalFlightMessagesSnafu {
|
||||
reason: "Expect schema to be the first flight message",
|
||||
}
|
||||
.fail();
|
||||
.fail()
|
||||
.map_err(|error| flight_stream_error(&addr, error));
|
||||
};
|
||||
|
||||
let metrics = Arc::new(ArcSwapOption::from(None));
|
||||
@@ -221,6 +233,7 @@ where
|
||||
let schema =
|
||||
Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?);
|
||||
let schema_cloned = schema.clone();
|
||||
let stream_addr = addr.clone();
|
||||
let stream = Box::pin(stream!({
|
||||
let _span = tracing_context.attach(common_telemetry::tracing::info_span!(
|
||||
"poll_flight_data_stream"
|
||||
@@ -240,7 +253,8 @@ where
|
||||
let flight_message = match flight_message_item {
|
||||
Some(Ok(message)) => message,
|
||||
Some(Err(e)) => {
|
||||
yield Err(BoxedError::new(e)).context(ExternalSnafu);
|
||||
yield Err(BoxedError::new(flight_stream_error(&stream_addr, e)))
|
||||
.context(ExternalSnafu);
|
||||
break;
|
||||
}
|
||||
None => break,
|
||||
@@ -273,7 +287,8 @@ where
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
yield Err(BoxedError::new(e)).context(ExternalSnafu);
|
||||
yield Err(BoxedError::new(flight_stream_error(&stream_addr, e)))
|
||||
.context(ExternalSnafu);
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -312,6 +327,23 @@ where
|
||||
Ok(Box::pin(record_batch_stream))
|
||||
}
|
||||
|
||||
fn flight_stream_error(addr: &str, error: error::Error) -> error::Error {
|
||||
let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown);
|
||||
if error.status_code().should_log_error() {
|
||||
error!(
|
||||
error; "Failed to receive Flight data, addr: {}, code: {}",
|
||||
addr,
|
||||
tonic_code
|
||||
);
|
||||
}
|
||||
|
||||
error::Error::FlightGet {
|
||||
addr: addr.to_string(),
|
||||
tonic_code,
|
||||
source: BoxedError::new(error),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_remote_dyn_filter_update_request(
|
||||
query_id: impl Into<String>,
|
||||
update: RemoteDynFilterUpdate,
|
||||
@@ -389,6 +421,44 @@ mod test {
|
||||
use super::*;
|
||||
use crate::Error::{self, IllegalDatabaseResponse, Server};
|
||||
|
||||
#[test]
|
||||
fn test_flight_stream_error_preserves_peer_address() {
|
||||
let error = flight_stream_error(
|
||||
"127.0.0.1:4001",
|
||||
tonic::Status::unavailable("datanode unavailable").into(),
|
||||
);
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
error::Error::FlightGet {
|
||||
addr,
|
||||
tonic_code: tonic::Code::Unavailable,
|
||||
..
|
||||
} if addr == "127.0.0.1:4001"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_flight_stream_preserves_peer_address() {
|
||||
let Err(error) = recordbatches_from_flight_message_stream(
|
||||
"127.0.0.1:4001".to_string(),
|
||||
stream::empty::<Result<FlightMessage>>(),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
panic!("expected empty Flight stream to fail");
|
||||
};
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
error::Error::FlightGet {
|
||||
addr,
|
||||
tonic_code: tonic::Code::Unknown,
|
||||
..
|
||||
} if addr == "127.0.0.1:4001"
|
||||
));
|
||||
}
|
||||
|
||||
fn test_schema() -> Arc<Schema> {
|
||||
Arc::new(Schema::new(vec![ColumnSchema::new(
|
||||
"v",
|
||||
@@ -509,11 +579,14 @@ mod test {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![
|
||||
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
|
||||
Ok(FlightMessage::Metrics(test_metrics_json())),
|
||||
Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
|
||||
]))
|
||||
let mut recordbatches = recordbatches_from_flight_message_stream(
|
||||
"test-peer".to_string(),
|
||||
stream::iter(vec![
|
||||
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
|
||||
Ok(FlightMessage::Metrics(test_metrics_json())),
|
||||
Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
|
||||
]),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -528,11 +601,14 @@ mod test {
|
||||
#[tokio::test]
|
||||
async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() {
|
||||
let schema = test_schema();
|
||||
let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![
|
||||
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
|
||||
Ok(FlightMessage::Metrics(test_metrics_json())),
|
||||
Err(Error::from(Status::internal("boom after metrics"))),
|
||||
]))
|
||||
let mut recordbatches = recordbatches_from_flight_message_stream(
|
||||
"test-peer".to_string(),
|
||||
stream::iter(vec![
|
||||
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
|
||||
Ok(FlightMessage::Metrics(test_metrics_json())),
|
||||
Err(Error::from(Status::internal("boom after metrics"))),
|
||||
]),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
||||
@@ -434,6 +434,8 @@ impl StartCommand {
|
||||
// Some queries are expected to take long time.
|
||||
let mut channel_config = opts.datanode.client.channel_config();
|
||||
channel_config.timeout = None;
|
||||
// Source Flight streams and sink unary responses share pooled connections.
|
||||
channel_config.http2_adaptive_window = Some(true);
|
||||
if opts.grpc.flight_compression.transport_compression() {
|
||||
channel_config.accept_compression = true;
|
||||
channel_config.send_compression = true;
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::borrow::Cow;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow::array::{ArrayRef, BinaryViewArray, new_null_array};
|
||||
@@ -54,12 +55,17 @@ trait JsonGetResultBuilder {
|
||||
fn build(&mut self) -> ArrayRef;
|
||||
}
|
||||
|
||||
fn result_builder(len: usize, with_type: &DataType) -> Result<Box<dyn JsonGetResultBuilder>> {
|
||||
fn result_builder(
|
||||
len: usize,
|
||||
with_type: &DataType,
|
||||
is_json2: bool,
|
||||
) -> Result<Box<dyn JsonGetResultBuilder>> {
|
||||
let builder = match with_type {
|
||||
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
|
||||
Box::new(StringResultBuilder(StringViewBuilder::with_capacity(len)))
|
||||
as Box<dyn JsonGetResultBuilder>
|
||||
}
|
||||
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Box::new(StringResultBuilder {
|
||||
inner: StringViewBuilder::with_capacity(len),
|
||||
is_json2,
|
||||
})
|
||||
as Box<dyn JsonGetResultBuilder>,
|
||||
DataType::Int64 => Box::new(IntResultBuilder(Int64Builder::with_capacity(len))),
|
||||
DataType::Float64 => Box::new(FloatResultBuilder(Float64Builder::with_capacity(len))),
|
||||
DataType::Boolean => Box::new(BoolResultBuilder(BooleanBuilder::with_capacity(len))),
|
||||
@@ -71,20 +77,30 @@ fn result_builder(len: usize, with_type: &DataType) -> Result<Box<dyn JsonGetRes
|
||||
}
|
||||
|
||||
// TODO: refactor this to StringLikeArrayBuilder from Arrow 57
|
||||
struct StringResultBuilder(StringViewBuilder);
|
||||
struct StringResultBuilder {
|
||||
inner: StringViewBuilder,
|
||||
is_json2: bool,
|
||||
}
|
||||
|
||||
impl JsonGetResultBuilder for StringResultBuilder {
|
||||
fn append_value(&mut self, value: &[u8]) -> Result<()> {
|
||||
self.0.append_option(jsonb::to_str(value).ok());
|
||||
// Scalar casts stay unquoted and map JSON null to SQL NULL; only containers
|
||||
// use `to_string` to preserve their JSON representation.
|
||||
let value = if self.is_json2 && (jsonb::is_array(value) || jsonb::is_object(value)) {
|
||||
Some(jsonb::to_string(value))
|
||||
} else {
|
||||
jsonb::to_str(value).ok()
|
||||
};
|
||||
self.inner.append_option(value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn append_null(&mut self) {
|
||||
self.0.append_null();
|
||||
self.inner.append_null();
|
||||
}
|
||||
|
||||
fn build(&mut self) -> ArrayRef {
|
||||
Arc::new(self.0.finish())
|
||||
Arc::new(self.inner.finish())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -404,18 +420,21 @@ impl Function for JsonGetWithType {
|
||||
let result = match arg0.data_type() {
|
||||
DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
|
||||
let arg0 = compute::cast(&arg0, &DataType::BinaryView)?;
|
||||
let is_json2 = args.arg_fields.first().is_some_and(is_json2_extension_type);
|
||||
|
||||
if args.arg_fields.first().is_some_and(is_json2_extension_type) {
|
||||
// Query concretization projects nested JSON2 paths as Struct arrays. A binary
|
||||
// JSON2 argument is therefore an already-selected scalar or root value that
|
||||
// only needs conversion from its JSONB representation to the requested type.
|
||||
if is_json2 && path.trim_start_matches('$').split('.').all(str::is_empty) {
|
||||
JsonArray::from(&arg0)
|
||||
.project_to(&with_type)
|
||||
.map_err(|e| exec_datafusion_err!("{e:?}"))?
|
||||
} else {
|
||||
let jsons = arg0.as_binary_view();
|
||||
let mut builder = result_builder(len, &with_type)?;
|
||||
jsonb_get(jsons, path, builder.as_mut())?;
|
||||
let path = if is_json2 && !path.starts_with('$') {
|
||||
Cow::Owned(format!("$.{path}"))
|
||||
} else {
|
||||
Cow::Borrowed(path)
|
||||
};
|
||||
let mut builder = result_builder(len, &with_type, is_json2)?;
|
||||
jsonb_get(jsons, &path, builder.as_mut())?;
|
||||
builder.build()
|
||||
}
|
||||
}
|
||||
@@ -511,6 +530,7 @@ mod tests {
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_common::arrow::array::{BinaryArray, BinaryViewArray, StringArray};
|
||||
use datafusion_common::arrow::datatypes::{Float64Type, Int64Type};
|
||||
use datatypes::extension::json::Json2ExtensionType;
|
||||
use datatypes::types::parse_string_to_jsonb;
|
||||
use serde_json::json;
|
||||
|
||||
@@ -566,6 +586,15 @@ mod tests {
|
||||
))
|
||||
}
|
||||
|
||||
fn test_json_field(json: &ArrayRef, is_json2: bool) -> Arc<Field> {
|
||||
let field = Field::new("json", json.data_type().clone(), true);
|
||||
Arc::new(if is_json2 {
|
||||
field.with_extension_type(Json2ExtensionType::default())
|
||||
} else {
|
||||
field
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_get_int() {
|
||||
let json_get_int = JsonGetInt::default();
|
||||
@@ -850,7 +879,10 @@ mod tests {
|
||||
ColumnarValue::Array(json.clone()),
|
||||
ColumnarValue::Scalar(path.into()),
|
||||
],
|
||||
arg_fields: vec![],
|
||||
arg_fields: vec![
|
||||
test_json_field(json, i >= json_strings.len()),
|
||||
Arc::new(Field::new("path", DataType::Utf8, false)),
|
||||
],
|
||||
number_rows: 1,
|
||||
return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
|
||||
config_options: Arc::new(Default::default()),
|
||||
@@ -984,7 +1016,11 @@ mod tests {
|
||||
ColumnarValue::Scalar(path.into()),
|
||||
ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
|
||||
],
|
||||
arg_fields: vec![],
|
||||
arg_fields: vec![
|
||||
test_json_field(json, i >= json_strings.len()),
|
||||
Arc::new(Field::new("path", DataType::Utf8, false)),
|
||||
Arc::new(Field::new("with_type", DataType::Utf8View, true)),
|
||||
],
|
||||
number_rows: 1,
|
||||
return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
|
||||
config_options: Arc::new(Default::default()),
|
||||
|
||||
@@ -89,7 +89,7 @@ strum.workspace = true
|
||||
table = { workspace = true, features = ["testing"] }
|
||||
tokio.workspace = true
|
||||
tokio-postgres = { workspace = true, optional = true }
|
||||
tokio-postgres-rustls = { version = "0.12", optional = true }
|
||||
tokio-postgres-rustls = { version = "0.14", optional = true }
|
||||
tonic.workspace = true
|
||||
tracing.workspace = true
|
||||
typetag.workspace = true
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use datatypes::schema::SchemaRef;
|
||||
use futures::StreamExt;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
@@ -28,12 +29,15 @@ struct Inner {
|
||||
|
||||
/// A cursor on RecordBatchStream that fetches data batch by batch
|
||||
pub struct RecordBatchStreamCursor {
|
||||
schema: SchemaRef,
|
||||
inner: Mutex<Inner>,
|
||||
}
|
||||
|
||||
impl RecordBatchStreamCursor {
|
||||
pub fn new(stream: SendableRecordBatchStream) -> RecordBatchStreamCursor {
|
||||
let schema = stream.schema();
|
||||
Self {
|
||||
schema,
|
||||
inner: Mutex::new(Inner {
|
||||
stream,
|
||||
current_row_index: 0,
|
||||
@@ -43,6 +47,10 @@ impl RecordBatchStreamCursor {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn schema(&self) -> SchemaRef {
|
||||
self.schema.clone()
|
||||
}
|
||||
|
||||
/// Take `size` of row from the `RecordBatchStream` and create a new
|
||||
/// `RecordBatch` for these rows.
|
||||
pub async fn take(&self, size: usize) -> Result<RecordBatch> {
|
||||
|
||||
@@ -483,6 +483,31 @@ impl Timestamp {
|
||||
ParseTimestampSnafu { raw: s }.fail()
|
||||
}
|
||||
|
||||
/// Interprets a timezone-less [`NaiveDateTime`] in the given timezone and
|
||||
/// returns the corresponding timestamp.
|
||||
///
|
||||
/// Datetimes that fall into a DST gap (a local time that does not exist)
|
||||
/// are rejected, while ambiguous datetimes (from a repeated local time)
|
||||
/// are resolved to the earlier instant. This policy is shared with
|
||||
/// [`Timestamp::from_str`] so that the text and binary protocols interpret
|
||||
/// datetimes consistently.
|
||||
pub fn from_naive_datetime(
|
||||
datetime: NaiveDateTime,
|
||||
timezone: &Timezone,
|
||||
) -> crate::error::Result<Timestamp> {
|
||||
match datetime_to_utc(&datetime, timezone) {
|
||||
LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => {
|
||||
Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu {
|
||||
raw: format!("{datetime} (timezone {timezone})"),
|
||||
})
|
||||
}
|
||||
LocalResult::None => ParseTimestampSnafu {
|
||||
raw: format!("{datetime} (timezone {timezone})"),
|
||||
}
|
||||
.fail(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn negative(mut self) -> Self {
|
||||
self.value = -self.value;
|
||||
self
|
||||
@@ -531,12 +556,8 @@ fn naive_datetime_to_timestamp(
|
||||
.context(ParseTimestampSnafu { raw: s });
|
||||
};
|
||||
|
||||
match datetime_to_utc(&datetime, timezone) {
|
||||
LocalResult::None => ParseTimestampSnafu { raw: s }.fail(),
|
||||
LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => {
|
||||
Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { raw: s })
|
||||
}
|
||||
}
|
||||
Timestamp::from_naive_datetime(datetime, timezone)
|
||||
.map_err(|_| ParseTimestampSnafu { raw: s }.build())
|
||||
}
|
||||
|
||||
impl From<i64> for Timestamp {
|
||||
@@ -919,6 +940,48 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_naive_datetime() {
|
||||
let datetime = NaiveDate::from_ymd_opt(2026, 8, 13)
|
||||
.unwrap()
|
||||
.and_hms_opt(8, 0, 0)
|
||||
.unwrap();
|
||||
|
||||
// A fixed-offset timezone shifts the datetime by a constant amount.
|
||||
let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap();
|
||||
assert_eq!(
|
||||
"2026-08-13 00:00:00",
|
||||
Timestamp::from_naive_datetime(datetime, &shanghai)
|
||||
.unwrap()
|
||||
.to_chrono_datetime()
|
||||
.unwrap()
|
||||
.to_string()
|
||||
);
|
||||
|
||||
// 2026-03-08 02:30 does not exist in America/New_York (DST gap).
|
||||
let new_york = Timezone::from_tz_string("America/New_York").unwrap();
|
||||
let gap = NaiveDate::from_ymd_opt(2026, 3, 8)
|
||||
.unwrap()
|
||||
.and_hms_opt(2, 30, 0)
|
||||
.unwrap();
|
||||
assert!(Timestamp::from_naive_datetime(gap, &new_york).is_err());
|
||||
|
||||
// 2026-11-01 01:30 is ambiguous in America/New_York; picks the first
|
||||
// instant (EDT, UTC-4).
|
||||
let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1)
|
||||
.unwrap()
|
||||
.and_hms_opt(1, 30, 0)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
"2026-11-01 05:30:00",
|
||||
Timestamp::from_naive_datetime(ambiguous, &new_york)
|
||||
.unwrap()
|
||||
.to_chrono_datetime()
|
||||
.unwrap()
|
||||
.to_string()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_iso8601_string() {
|
||||
set_default_timezone(Some("Asia/Shanghai")).unwrap();
|
||||
|
||||
@@ -108,6 +108,15 @@ impl Timezone {
|
||||
}
|
||||
}
|
||||
|
||||
/// A named zone merely sitting at +00:00 today is not UTC: it may have been
|
||||
/// elsewhere at the timestamp being converted.
|
||||
pub fn is_utc(&self) -> bool {
|
||||
match self {
|
||||
Self::Offset(offset) => offset.local_minus_utc() == 0,
|
||||
Self::Named(tz) => matches!(tz, Tz::UTC),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the number of seconds to add to convert from UTC to the local time.
|
||||
pub fn local_minus_utc(&self) -> i64 {
|
||||
match self {
|
||||
|
||||
@@ -98,6 +98,8 @@ mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_error::ext::{BoxedError, RetryHint};
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_meta::RegionIdent;
|
||||
use common_meta::heartbeat::handler::{HandleControl, HeartbeatResponseHandler};
|
||||
use common_meta::heartbeat::mailbox::MessageMeta;
|
||||
@@ -105,14 +107,21 @@ mod tests {
|
||||
use common_meta::kv_backend::memory::MemoryKvBackend;
|
||||
use mito2::config::MitoConfig;
|
||||
use mito2::engine::MITO_ENGINE_NAME;
|
||||
use mito2::error::ManifestDeltaNotFoundSnafu;
|
||||
use mito2::test_util::{CreateRequestBuilder, TestEnv};
|
||||
use object_store::{Error as ObjectStoreError, ErrorKind};
|
||||
use snafu::IntoError;
|
||||
use store_api::path_utils::table_dir;
|
||||
use store_api::region_request::{RegionCloseRequest, RegionRequest, RegionRequirements};
|
||||
use store_api::storage::RegionId;
|
||||
|
||||
use crate::heartbeat::handler::RegionHeartbeatResponseHandler;
|
||||
use super::OpenRegionsHandler;
|
||||
use crate::error::{self, HandleRegionRequestSnafu};
|
||||
use crate::heartbeat::handler::tests::HeartbeatResponseTestEnv;
|
||||
use crate::tests::mock_region_server;
|
||||
use crate::heartbeat::handler::{
|
||||
HandlerContext, InstructionHandler, RegionHeartbeatResponseHandler,
|
||||
};
|
||||
use crate::tests::{MockRegionEngine, mock_region_server};
|
||||
|
||||
fn open_regions_instruction(
|
||||
region_ids: impl IntoIterator<Item = RegionId>,
|
||||
@@ -198,4 +207,52 @@ mod tests {
|
||||
assert!(engine.is_region_exists(region_id));
|
||||
assert!(engine.is_region_exists(region_id1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_open_regions_preserves_manifest_delta_not_found_retry_hint() {
|
||||
let retryable_region = RegionId::new(1024, 1);
|
||||
let non_retryable_region = RegionId::new(1024, 2);
|
||||
let (engine, _) = MockRegionEngine::with_mock_fn(
|
||||
MITO_ENGINE_NAME,
|
||||
Box::new(move |region_id, _request| {
|
||||
if region_id == retryable_region {
|
||||
let manifest_error = ManifestDeltaNotFoundSnafu {
|
||||
version: 1_u64,
|
||||
path: "manifest/00000000000000000001.json",
|
||||
}
|
||||
.into_error(ObjectStoreError::new(
|
||||
ErrorKind::NotFound,
|
||||
"mock listed manifest delta not found",
|
||||
));
|
||||
return Err(HandleRegionRequestSnafu { region_id }
|
||||
.into_error(BoxedError::new(manifest_error)));
|
||||
}
|
||||
|
||||
error::RegionNotFoundSnafu { region_id }.fail()
|
||||
}),
|
||||
);
|
||||
let mut region_server = mock_region_server();
|
||||
region_server.register_engine(engine);
|
||||
let ctx = HandlerContext::new_for_test(region_server, Arc::new(MemoryKvBackend::new()));
|
||||
let Instruction::OpenRegions(open_regions) =
|
||||
open_regions_instruction([retryable_region, non_retryable_region], "test")
|
||||
else {
|
||||
unreachable!()
|
||||
};
|
||||
|
||||
// Serial execution makes the first error selection deterministic.
|
||||
let reply = OpenRegionsHandler {
|
||||
open_region_parallelism: 1,
|
||||
}
|
||||
.handle(&ctx, open_regions)
|
||||
.await
|
||||
.unwrap()
|
||||
.expect_open_regions_reply();
|
||||
|
||||
assert!(!reply.result);
|
||||
let error = reply.error.unwrap();
|
||||
assert_eq!(StatusCode::StorageUnavailable, error.code);
|
||||
assert_eq!(RetryHint::Retryable, error.retry_hint);
|
||||
assert!(error.message.contains("00000000000000000001.json"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -103,11 +103,18 @@ impl SyncRegionHandler {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_error::ext::{BoxedError, RetryHint};
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_meta::kv_backend::memory::MemoryKvBackend;
|
||||
use mito2::engine::MITO_ENGINE_NAME;
|
||||
use mito2::error::ManifestDeltaNotFoundSnafu;
|
||||
use object_store::{Error as ObjectStoreError, ErrorKind};
|
||||
use snafu::IntoError;
|
||||
use store_api::metric_engine_consts::METRIC_ENGINE_NAME;
|
||||
use store_api::region_engine::{RegionRole, SyncRegionFromRequest};
|
||||
use store_api::storage::RegionId;
|
||||
|
||||
use crate::error::HandleRegionRequestSnafu;
|
||||
use crate::heartbeat::handler::sync_region::SyncRegionHandler;
|
||||
use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
|
||||
use crate::tests::{MockRegionEngine, mock_region_server};
|
||||
@@ -201,4 +208,47 @@ mod tests {
|
||||
assert!(reply[0].ready);
|
||||
assert!(reply[0].error.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handle_sync_region_preserves_manifest_delta_not_found_retry_hint() {
|
||||
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::Leader));
|
||||
engine.handle_sync_region_mock_fn = Some(Box::new(|region_id, _request| {
|
||||
let manifest_error = ManifestDeltaNotFoundSnafu {
|
||||
version: 1_u64,
|
||||
path: "manifest/00000000000000000001.json",
|
||||
}
|
||||
.into_error(ObjectStoreError::new(
|
||||
ErrorKind::NotFound,
|
||||
"mock listed manifest delta not found",
|
||||
));
|
||||
Err(HandleRegionRequestSnafu { region_id }
|
||||
.into_error(BoxedError::new(manifest_error)))
|
||||
}));
|
||||
});
|
||||
mock_region_server.register_test_region(region_id, mock_engine);
|
||||
|
||||
let handler_context =
|
||||
HandlerContext::new_for_test(mock_region_server, Arc::new(MemoryKvBackend::new()));
|
||||
let sync_region = common_meta::instruction::SyncRegion {
|
||||
region_id,
|
||||
request: SyncRegionFromRequest::from_manifest(Default::default()),
|
||||
};
|
||||
|
||||
let reply = SyncRegionHandler
|
||||
.handle(&handler_context, vec![sync_region])
|
||||
.await
|
||||
.unwrap()
|
||||
.expect_sync_regions_reply();
|
||||
|
||||
assert_eq!(1, reply.len());
|
||||
assert!(reply[0].exists);
|
||||
assert!(!reply[0].ready);
|
||||
let error = reply[0].error.as_ref().unwrap();
|
||||
assert_eq!(StatusCode::StorageUnavailable, error.code);
|
||||
assert_eq!(RetryHint::Retryable, error.retry_hint);
|
||||
assert!(error.message.contains("00000000000000000001.json"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -66,7 +66,9 @@ use servers::error::{
|
||||
self as servers_error, ExecuteGrpcRequestSnafu, Result as ServerResult, SuspendedSnafu,
|
||||
};
|
||||
use servers::grpc::FlightCompression;
|
||||
use servers::grpc::flight::{FlightCraft, FlightRecordBatchStream, TonicStream};
|
||||
use servers::grpc::flight::{
|
||||
FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream,
|
||||
};
|
||||
use servers::grpc::region_server::RegionServerHandler;
|
||||
use session::context::{
|
||||
FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef,
|
||||
@@ -976,13 +978,19 @@ impl FlightCraft for RegionServer {
|
||||
.map(|h| Arc::new(QueryContext::from(h)))
|
||||
.unwrap_or(QueryContext::arc());
|
||||
|
||||
let result = self
|
||||
.handle_remote_read(request, query_ctx.clone())
|
||||
.trace(tracing_context.attach(info_span!("RegionServer::handle_read")))
|
||||
.await?;
|
||||
let region_server = self.clone();
|
||||
let initializer_query_ctx = query_ctx.clone();
|
||||
let initializer_tracing_context = tracing_context.clone();
|
||||
let initializer = async move {
|
||||
region_server
|
||||
.handle_remote_read(request, initializer_query_ctx)
|
||||
.trace(initializer_tracing_context.attach(info_span!("RegionServer::handle_read")))
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
};
|
||||
|
||||
let stream = Box::pin(FlightRecordBatchStream::new(
|
||||
result,
|
||||
FlightRecordBatchSource::initializer(initializer),
|
||||
tracing_context,
|
||||
self.flight_compression,
|
||||
query_ctx,
|
||||
@@ -1330,11 +1338,8 @@ impl RegionServerInner {
|
||||
}
|
||||
|
||||
if !errors.is_empty() {
|
||||
return error::UnexpectedSnafu {
|
||||
// Returns the first error.
|
||||
violated: format!("Failed to open batch regions: {:?}", errors[0]),
|
||||
}
|
||||
.fail();
|
||||
// Preserve the first region error so callers can honor its status code and retry hint.
|
||||
return Err(errors.swap_remove(0)).context(HandleBatchOpenRequestSnafu);
|
||||
}
|
||||
|
||||
Ok(open_regions)
|
||||
@@ -1979,7 +1984,7 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::{Rows, SemanticType};
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::ext::{ErrorExt, RetryHint};
|
||||
use common_recordbatch::RecordBatches;
|
||||
use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry};
|
||||
use datatypes::prelude::{ConcreteDataType, VectorRef};
|
||||
@@ -2514,7 +2519,9 @@ mod tests {
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(err.status_code(), StatusCode::Unexpected);
|
||||
assert_matches!(&err, error::Error::HandleBatchOpenRequest { .. });
|
||||
assert_eq!(err.status_code(), StatusCode::RegionNotFound);
|
||||
assert_eq!(err.retry_hint(), RetryHint::NonRetryable);
|
||||
}
|
||||
|
||||
struct CurrentEngineTest {
|
||||
|
||||
@@ -77,6 +77,13 @@ pub enum Error {
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Unimplemented: {feat}"))]
|
||||
Unimplemented {
|
||||
feat: String,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Failed to parse version in schema meta, value: {}", value))]
|
||||
ParseSchemaVersion {
|
||||
value: String,
|
||||
@@ -203,6 +210,13 @@ pub enum Error {
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Invalid JSON2 settings: {reason}"))]
|
||||
InvalidJson2Settings {
|
||||
reason: String,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Invalid Vector: {}", msg))]
|
||||
InvalidVector {
|
||||
msg: String,
|
||||
@@ -312,6 +326,7 @@ impl ErrorExt for Error {
|
||||
use Error::*;
|
||||
match self {
|
||||
UnsupportedOperation { .. }
|
||||
| Unimplemented { .. }
|
||||
| UnsupportedArrowType { .. }
|
||||
| UnsupportedJsonType { .. }
|
||||
| UnsupportedDefaultExpr { .. } => StatusCode::Unsupported,
|
||||
@@ -326,6 +341,7 @@ impl ErrorExt for Error {
|
||||
| InvalidPrecisionOrScale { .. }
|
||||
| InvalidJson { .. }
|
||||
| InvalidJson2Layout { .. }
|
||||
| InvalidJson2Settings { .. }
|
||||
| InvalidJsonb { .. }
|
||||
| InvalidVector { .. }
|
||||
| InvalidFulltextOption { .. }
|
||||
|
||||
@@ -129,27 +129,32 @@ pub struct JsonMetadata {
|
||||
}
|
||||
|
||||
impl JsonMetadata {
|
||||
/// Creates metadata for the legacy JSON2 layout.
|
||||
/// Creates metadata for the JSON2 layout version 2.
|
||||
pub fn new(json_settings: JsonSettings) -> Self {
|
||||
Self {
|
||||
json_settings,
|
||||
layout_version: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates metadata for the new JSON2 physical v2 layout.
|
||||
pub fn new_v2(json_settings: JsonSettings) -> Self {
|
||||
Self {
|
||||
json_settings,
|
||||
layout_version: Some(JSON2_LAYOUT_V2),
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates metadata for the legacy JSON2 layout.
|
||||
pub fn new_v1(json_settings: JsonSettings) -> Self {
|
||||
Self {
|
||||
json_settings,
|
||||
layout_version: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the JSON2 settings.
|
||||
pub fn json_settings(&self) -> &JsonSettings {
|
||||
&self.json_settings
|
||||
}
|
||||
|
||||
/// Consumes the metadata and returns its JSON2 settings.
|
||||
pub fn into_json_settings(self) -> JsonSettings {
|
||||
self.json_settings
|
||||
}
|
||||
|
||||
/// Returns whether this metadata describes JSON2 layout version 2.
|
||||
pub fn is_version_2(&self) -> bool {
|
||||
self.layout_version == Some(JSON2_LAYOUT_V2)
|
||||
@@ -226,7 +231,7 @@ impl ExtensionType for Json2ExtensionType {
|
||||
})?;
|
||||
Ok(Arc::new(metadata))
|
||||
} else {
|
||||
Ok(Arc::new(JsonMetadata::default()))
|
||||
Ok(Arc::new(JsonMetadata::new_v1(JsonSettings::default())))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -406,7 +411,7 @@ mod tests {
|
||||
let legacy: JsonMetadata = serde_json::from_str(r#"{"json_settings":{}}"#)?;
|
||||
assert!(!legacy.is_version_2());
|
||||
|
||||
let metadata = JsonMetadata::new_v2(JsonSettings::default());
|
||||
let metadata = JsonMetadata::new(JsonSettings::default());
|
||||
assert!(metadata.is_version_2());
|
||||
let serialized = serde_json::to_string(&metadata)?;
|
||||
let deserialized: JsonMetadata = serde_json::from_str(&serialized)?;
|
||||
@@ -418,7 +423,9 @@ mod tests {
|
||||
#[test]
|
||||
fn test_parse_json2_physical_layout() -> crate::error::Result<()> {
|
||||
let legacy = Field::new("data", DataType::Struct(Fields::empty()), true)
|
||||
.with_extension_type(Json2ExtensionType::default());
|
||||
.with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1(
|
||||
JsonSettings::default(),
|
||||
))));
|
||||
assert!(!Json2PhysicalLayout::try_from_root(&legacy)?.is_version_2());
|
||||
|
||||
let v2 = Field::new(
|
||||
@@ -432,7 +439,7 @@ mod tests {
|
||||
),
|
||||
true,
|
||||
)
|
||||
.with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(
|
||||
.with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new(
|
||||
JsonSettings::default(),
|
||||
))));
|
||||
assert!(Json2PhysicalLayout::try_from_root(&v2)?.is_version_2());
|
||||
@@ -447,7 +454,7 @@ mod tests {
|
||||
assert!(Json2PhysicalLayout::try_from_root(&field).is_err());
|
||||
|
||||
let metadata =
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default())));
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default())));
|
||||
let missing = Field::new("data", DataType::Struct(Fields::empty()), true)
|
||||
.with_extension_type(metadata.clone());
|
||||
assert!(json2_remainder_field(&missing).is_ok_and(|x| x.is_none()));
|
||||
|
||||
@@ -29,7 +29,7 @@ use serde_json::{Map, Value as Json};
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::data_type::ConcreteDataType;
|
||||
use crate::error::{self, InvalidJson2LayoutSnafu, Result, UnsupportedJsonTypeSnafu};
|
||||
use crate::error::{self, InvalidJson2SettingsSnafu, Result, UnsupportedJsonTypeSnafu};
|
||||
use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb};
|
||||
use crate::schema::ColumnDefaultConstraint;
|
||||
use crate::types::json_type::{JsonNativeType, JsonObjectType};
|
||||
@@ -39,6 +39,8 @@ use crate::value::{ListValue, StructValue, Value};
|
||||
pub const JSON2_MAX_STRUCTURED_DEPTH: usize = 50;
|
||||
/// Reserved physical field containing unexpanded JSON2 paths.
|
||||
pub const JSON2_REMAINDER_FIELD_NAME: &str = "!__remainder__!";
|
||||
/// Default maximum number of unhinted JSON leaf paths expanded into Arrow fields.
|
||||
pub const JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS: u32 = 100;
|
||||
|
||||
/// JSON2 settings stored in column schema metadata and represented through
|
||||
/// Arrow extension metadata.
|
||||
@@ -105,6 +107,14 @@ pub struct JsonContext<'a> {
|
||||
}
|
||||
|
||||
impl JsonSettings {
|
||||
/// Creates default v2 settings for newly created JSON2 columns.
|
||||
pub fn new_v2() -> Self {
|
||||
Self {
|
||||
type_hints: vec![],
|
||||
max_auto_expanded_paths: Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS),
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates and validates JSON2 settings.
|
||||
pub fn try_new(
|
||||
type_hints: Vec<JsonTypeHint>,
|
||||
@@ -140,6 +150,16 @@ impl JsonSettings {
|
||||
|
||||
/// Encode a serde_json::Value into a Value::Json using current settings.
|
||||
pub fn encode(&self, json: Json) -> Result<Value> {
|
||||
if let Json::Object(object) = &json
|
||||
&& object.contains_key(JSON2_REMAINDER_FIELD_NAME)
|
||||
{
|
||||
return error::InvalidJsonSnafu {
|
||||
value: format!(
|
||||
"root object cannot contain reserved field '{JSON2_REMAINDER_FIELD_NAME}'"
|
||||
),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
let mut context = JsonContext {
|
||||
path: Vec::new(),
|
||||
settings: self,
|
||||
@@ -149,10 +169,14 @@ impl JsonSettings {
|
||||
}
|
||||
|
||||
fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> {
|
||||
if type_hints.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut object = JsonObjectType::new();
|
||||
for hint in type_hints {
|
||||
if hint.path.len() > JSON2_MAX_STRUCTURED_DEPTH {
|
||||
return InvalidJson2LayoutSnafu {
|
||||
return InvalidJson2SettingsSnafu {
|
||||
reason: format!(
|
||||
"JSON2 type hint path cannot exceed {JSON2_MAX_STRUCTURED_DEPTH} segments"
|
||||
),
|
||||
@@ -164,7 +188,7 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> {
|
||||
.first()
|
||||
.is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME)
|
||||
{
|
||||
return InvalidJson2LayoutSnafu {
|
||||
return InvalidJson2SettingsSnafu {
|
||||
reason: format!(
|
||||
"JSON2 type hint path cannot be rooted at reserved field '{JSON2_REMAINDER_FIELD_NAME}'"
|
||||
),
|
||||
@@ -183,14 +207,30 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> {
|
||||
| ConcreteDataType::Int64(_)
|
||||
| ConcreteDataType::Float32(_)
|
||||
| ConcreteDataType::Float64(_)
|
||||
| ConcreteDataType::String(_) => (&hint.data_type).into(),
|
||||
| ConcreteDataType::String(_)
|
||||
| ConcreteDataType::List(_)
|
||||
| ConcreteDataType::Struct(_) => (&hint.data_type).into(),
|
||||
data_type => {
|
||||
return InvalidJson2LayoutSnafu {
|
||||
return InvalidJson2SettingsSnafu {
|
||||
reason: format!("unsupported JSON2 type hint data type: {data_type}"),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
};
|
||||
let non_finite_default = match &hint.default_constraint {
|
||||
Some(ColumnDefaultConstraint::Value(Value::Float32(value))) => !value.0.is_finite(),
|
||||
Some(ColumnDefaultConstraint::Value(Value::Float64(value))) => !value.0.is_finite(),
|
||||
_ => false,
|
||||
};
|
||||
if non_finite_default {
|
||||
return InvalidJson2SettingsSnafu {
|
||||
reason: format!(
|
||||
"JSON2 type hint default for '{}' must be finite",
|
||||
hint.path.join(".")
|
||||
),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
validate_type_hint(&mut object, &hint.path, data_type)?;
|
||||
}
|
||||
Ok(())
|
||||
@@ -202,14 +242,14 @@ fn validate_type_hint(
|
||||
data_type: JsonNativeType,
|
||||
) -> Result<()> {
|
||||
let Some((name, path)) = path.split_first() else {
|
||||
return InvalidJson2LayoutSnafu {
|
||||
return InvalidJson2SettingsSnafu {
|
||||
reason: "JSON2 type hint path must not be empty".to_string(),
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
if path.is_empty() {
|
||||
if object.insert(name.clone(), data_type).is_some() {
|
||||
return InvalidJson2LayoutSnafu {
|
||||
return InvalidJson2SettingsSnafu {
|
||||
reason: format!("duplicate JSON2 type hint path '{name}'"),
|
||||
}
|
||||
.fail();
|
||||
@@ -221,7 +261,7 @@ fn validate_type_hint(
|
||||
.entry(name.clone())
|
||||
.or_insert_with(|| JsonNativeType::Object(JsonObjectType::new()));
|
||||
let JsonNativeType::Object(child) = child else {
|
||||
return InvalidJson2LayoutSnafu {
|
||||
return InvalidJson2SettingsSnafu {
|
||||
reason: format!("conflicting JSON2 type hint path at '{name}'"),
|
||||
}
|
||||
.fail();
|
||||
@@ -702,7 +742,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_settings_reject_invalid_type_hint_layout() {
|
||||
fn test_json_settings_reject_invalid_type_hints() {
|
||||
for type_hints in [
|
||||
json!([{"path": [], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}]),
|
||||
json!([
|
||||
@@ -777,6 +817,21 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encode_rejects_reserved_remainder_field() -> Result<()> {
|
||||
let settings = JsonSettings::default();
|
||||
let err = settings
|
||||
.encode(json!({"!__remainder__!": "user-value"}))
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("root object cannot contain reserved field")
|
||||
);
|
||||
|
||||
settings.encode(json!({"nested": {"!__remainder__!": "user-value"}}))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encode_json_object() {
|
||||
let json = json!({
|
||||
|
||||
@@ -30,6 +30,8 @@ use crate::types::json_type::{JsonNativeType, JsonNumberType, is_include};
|
||||
use crate::types::{StructField, StructType};
|
||||
use crate::value::{ListValue, StructValue, Value};
|
||||
|
||||
pub type JsonObjectVariant = BTreeMap<String, JsonVariant>;
|
||||
|
||||
/// Number in json, can be a positive integer, a negative integer, or a floating number.
|
||||
/// Each of which is represented as `u64`, `i64` and `f64`.
|
||||
///
|
||||
@@ -124,7 +126,7 @@ pub enum JsonVariant {
|
||||
Number(JsonNumber),
|
||||
String(String),
|
||||
Array(Vec<JsonVariant>),
|
||||
Object(BTreeMap<String, JsonVariant>),
|
||||
Object(JsonObjectVariant),
|
||||
/// A special "variant" value of JSON, to represent a union result of conflict JSON type values.
|
||||
Variant(Vec<u8>),
|
||||
}
|
||||
@@ -160,7 +162,8 @@ impl JsonVariant {
|
||||
}
|
||||
}
|
||||
|
||||
fn contains_empty_object(&self) -> bool {
|
||||
/// Returns whether this value recursively contains an empty object.
|
||||
pub(crate) fn contains_empty_object(&self) -> bool {
|
||||
match self {
|
||||
JsonVariant::Array(array) => array.iter().any(JsonVariant::contains_empty_object),
|
||||
JsonVariant::Object(object) => {
|
||||
@@ -721,12 +724,7 @@ where
|
||||
I: IntoIterator<Item = (K, JsonNativeType)>,
|
||||
K: Into<String>,
|
||||
{
|
||||
let mut fields = fields.into_iter().peekable();
|
||||
if fields.peek().is_none() {
|
||||
JsonNativeType::Null
|
||||
} else {
|
||||
JsonNativeType::Object(fields.map(|(k, v)| (k.into(), v)).collect())
|
||||
}
|
||||
JsonNativeType::Object(fields.into_iter().map(|(k, v)| (k.into(), v)).collect())
|
||||
}
|
||||
|
||||
impl From<()> for JsonVariantRef<'_> {
|
||||
@@ -972,11 +970,10 @@ mod tests {
|
||||
])))
|
||||
);
|
||||
|
||||
// Empty objects have native type Null, but the value still needs alignment
|
||||
// before converting into a typed struct value.
|
||||
// Empty objects have an empty Object type and remain distinct from Null.
|
||||
let expected = JsonNativeType::Object(JsonObjectType::from([(
|
||||
"empty".to_string(),
|
||||
JsonNativeType::Null,
|
||||
JsonNativeType::Object(JsonObjectType::default()),
|
||||
)]));
|
||||
let mut value = parse_json_value(r#"{"empty":{}}"#);
|
||||
assert_eq!(value.json_type(), &expected);
|
||||
@@ -985,7 +982,7 @@ mod tests {
|
||||
value,
|
||||
JsonValue::from(JsonVariant::Object(BTreeMap::from([(
|
||||
"empty".to_string(),
|
||||
JsonVariant::Null,
|
||||
JsonVariant::Object(BTreeMap::default()),
|
||||
)])))
|
||||
);
|
||||
|
||||
|
||||
@@ -28,10 +28,12 @@ use crate::data_type::{ConcreteDataType, DataType};
|
||||
use crate::error::{
|
||||
self, ArrowMetadataSnafu, Error, InvalidFulltextOptionSnafu, ParseExtendedTypeSnafu, Result,
|
||||
};
|
||||
use crate::extension::json::Json2ExtensionType;
|
||||
use crate::schema::TYPE_KEY;
|
||||
use crate::schema::constraint::ColumnDefaultConstraint;
|
||||
use crate::value::Value;
|
||||
use crate::vectors::VectorRef;
|
||||
use crate::vectors::json::builder::JsonVectorBuilder;
|
||||
use crate::vectors::{MutableVector, VectorRef};
|
||||
|
||||
pub type Metadata = HashMap<String, String>;
|
||||
|
||||
@@ -128,6 +130,21 @@ impl ColumnSchema {
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a mutable vector using this column's extension metadata.
|
||||
pub fn create_mutable_vector(&self, capacity: usize) -> Box<dyn MutableVector> {
|
||||
if self.data_type.is_json2()
|
||||
&& let Some(extension) = self.extension_type::<Json2ExtensionType>().ok().flatten()
|
||||
&& extension.metadata().is_version_2()
|
||||
{
|
||||
Box::new(JsonVectorBuilder::with_settings(
|
||||
extension.metadata().json_settings(),
|
||||
capacity,
|
||||
))
|
||||
} else {
|
||||
self.data_type.create_mutable_vector(capacity)
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn is_time_index(&self) -> bool {
|
||||
self.is_time_index
|
||||
|
||||
@@ -50,6 +50,11 @@ pub enum JsonNumberType {
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, Default)]
|
||||
pub enum JsonNativeType {
|
||||
/// JSON null value type.
|
||||
///
|
||||
/// This variant may also appear as the initial state while merging inferred
|
||||
/// types, but it does not represent an empty object. Empty objects are
|
||||
/// represented as `Object({})`.
|
||||
#[default]
|
||||
Null,
|
||||
Bool,
|
||||
@@ -80,7 +85,7 @@ impl JsonNativeType {
|
||||
Self::Number(JsonNumberType::F64)
|
||||
}
|
||||
|
||||
fn object() -> Self {
|
||||
pub fn object() -> Self {
|
||||
Self::Object(JsonObjectType::new())
|
||||
}
|
||||
|
||||
@@ -143,6 +148,14 @@ impl JsonNativeType {
|
||||
JsonNativeType::Variant => ArrowDataType::Binary,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns whether this type is a boolean, number, or string scalar.
|
||||
pub fn is_primitive(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
JsonNativeType::Bool | JsonNativeType::Number(_) | JsonNativeType::String
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&ConcreteDataType> for JsonNativeType {
|
||||
@@ -360,6 +373,7 @@ impl DataType for JsonType {
|
||||
fn create_mutable_vector(&self, capacity: usize) -> Box<dyn MutableVector> {
|
||||
match &self.format {
|
||||
JsonFormat::Jsonb => Box::new(BinaryVectorBuilder::with_capacity(capacity)),
|
||||
// TODO(LFC): Carry JsonSettings in JsonFormat::Json2 and use with_settings here.
|
||||
JsonFormat::Json2(x) => Box::new(JsonVectorBuilder::new(x.as_ref().clone(), capacity)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -87,22 +87,6 @@ impl Helper {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn check_get_mutable_vector<T: 'static + MutableVector>(
|
||||
vector: &mut dyn MutableVector,
|
||||
) -> Result<&mut T> {
|
||||
let ty = vector.data_type();
|
||||
vector
|
||||
.as_mut_any()
|
||||
.downcast_mut()
|
||||
.with_context(|| error::UnknownVectorSnafu {
|
||||
msg: format!(
|
||||
"downcast vector error, vector type: {:?}, expected vector: {:?}",
|
||||
ty,
|
||||
std::any::type_name::<T>(),
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn check_get_scalar_vector<T: Scalar>(
|
||||
vector: &VectorRef,
|
||||
) -> Result<&<T as Scalar>::VectorType> {
|
||||
|
||||
@@ -15,3 +15,5 @@
|
||||
pub mod array;
|
||||
pub(crate) mod builder;
|
||||
pub mod variant;
|
||||
|
||||
pub use builder::json2_physical_data_type;
|
||||
|
||||
@@ -33,9 +33,12 @@ use crate::error::{
|
||||
AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result,
|
||||
};
|
||||
use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field};
|
||||
use crate::json::JsonSettings;
|
||||
use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb};
|
||||
use crate::prelude::{DataType as _, Value as GreptimeValue};
|
||||
use crate::value::{ListValue, StructValue};
|
||||
use crate::vectors::MutableVector;
|
||||
use crate::vectors::json::builder::{JsonVectorBuilder, json2_physical_data_type};
|
||||
use crate::vectors::json::variant::variant_to_json_values;
|
||||
|
||||
pub struct JsonArray<'a> {
|
||||
@@ -115,6 +118,38 @@ impl JsonArray<'_> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Rewrites a JSON2 array from the current physical layout into the specified
|
||||
/// v2 physical layout.
|
||||
pub fn rewrite_to_v2(
|
||||
&self,
|
||||
field: &Field,
|
||||
logical_settings: &JsonSettings,
|
||||
target_layout: &JsonSettings,
|
||||
) -> Result<ArrayRef> {
|
||||
let is_v2 = json2_remainder_field(field)?.is_some();
|
||||
if is_v2 && self.inner.data_type() == &json2_physical_data_type(target_layout) {
|
||||
return Ok(self.inner.clone());
|
||||
}
|
||||
|
||||
let values = if is_v2 {
|
||||
self.json2_values()?
|
||||
} else {
|
||||
(0..self.inner.len())
|
||||
.map(|i| self.try_get_value(i))
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
};
|
||||
let mut builder = JsonVectorBuilder::with_settings(target_layout, values.len());
|
||||
for value in values {
|
||||
if value.is_null() {
|
||||
builder.push_null();
|
||||
} else {
|
||||
let value = logical_settings.encode(value)?;
|
||||
builder.try_push_value_ref(&value.as_value_ref())?;
|
||||
}
|
||||
}
|
||||
Ok(builder.to_vector().to_arrow_array())
|
||||
}
|
||||
|
||||
fn json2_values(&self) -> Result<Vec<Value>> {
|
||||
let structs = self.inner.as_struct_opt().context(AlignJsonArraySnafu {
|
||||
reason: "JSON2 layout v2 root array must be a struct",
|
||||
@@ -149,7 +184,14 @@ impl JsonArray<'_> {
|
||||
if child.name() == JSON2_REMAINDER_FIELD_NAME {
|
||||
continue;
|
||||
}
|
||||
let value = JsonArray::from(column).try_get_value(i)?;
|
||||
let mut value = JsonArray::from(column).try_get_value(i)?;
|
||||
// Arrow child nulls cannot distinguish a missing path from an explicit JSON
|
||||
// null. Builders preserve explicit null presence in the remainder, so nulls
|
||||
// from the explicit branch must be discarded before merging both branches.
|
||||
remove_null_object_fields(&mut value);
|
||||
if value.is_null() {
|
||||
continue;
|
||||
}
|
||||
merge_explicit_value(&mut object, child.name().clone(), value, &mut path)?;
|
||||
}
|
||||
values.push(Value::Object(object));
|
||||
@@ -420,6 +462,16 @@ fn merge_explicit_value(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn remove_null_object_fields(value: &mut Value) {
|
||||
let Value::Object(object) = value else {
|
||||
return;
|
||||
};
|
||||
object.retain(|_, value| {
|
||||
remove_null_object_fields(value);
|
||||
!value.is_null()
|
||||
});
|
||||
}
|
||||
|
||||
/// Returns whether Arrow can cast between the types without JSON-aware projection.
|
||||
/// Binary and nested types require JSONB decoding or recursive projection.
|
||||
fn can_fast_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
|
||||
@@ -545,6 +597,8 @@ impl<'a> From<&'a ArrayRef> for JsonArray<'a> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::types::Int64Type;
|
||||
use arrow_array::{
|
||||
BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
|
||||
@@ -555,7 +609,7 @@ mod test {
|
||||
|
||||
use super::*;
|
||||
use crate::extension::json::{Json2ExtensionType, JsonMetadata};
|
||||
use crate::json::JsonSettings;
|
||||
use crate::json::{JsonSettings, JsonTypeHint};
|
||||
use crate::vectors::json::variant::{json_values_to_variant, variant_field};
|
||||
|
||||
#[test]
|
||||
@@ -1049,7 +1103,7 @@ mod test {
|
||||
None,
|
||||
));
|
||||
let field = Field::new("data", DataType::Struct(fields), true).with_extension_type(
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))),
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default()))),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
@@ -1063,12 +1117,10 @@ mod test {
|
||||
assert_eq!(
|
||||
json!({
|
||||
"!__remainder__!": "user value",
|
||||
"count": null,
|
||||
"nested": {"left": null}
|
||||
"nested": {}
|
||||
}),
|
||||
JsonArray::from(&array).json2_values()?[1]
|
||||
);
|
||||
|
||||
let target = DataType::Struct(
|
||||
vec![
|
||||
Arc::new(Field::new("cold", DataType::UInt64, true)),
|
||||
@@ -1088,6 +1140,38 @@ mod test {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rewrite_to_v2_reuses_matching_layout() -> Result<()> {
|
||||
let settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["kind".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
Some(0),
|
||||
)?;
|
||||
let value = settings.encode(json!({"kind": "access", "cold": 1}))?;
|
||||
let mut builder = JsonVectorBuilder::with_settings(&settings, 1);
|
||||
builder.try_push_value_ref(&value.as_value_ref())?;
|
||||
let array = builder.to_vector().to_arrow_array();
|
||||
let structs = array.as_struct();
|
||||
assert!(structs.column_by_name("kind").is_some());
|
||||
assert_eq!(
|
||||
vec![Some(json!({"cold": 1}))],
|
||||
variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
|
||||
);
|
||||
let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings.clone()))),
|
||||
);
|
||||
|
||||
let rewritten = JsonArray::from(&array).rewrite_to_v2(&field, &settings, &settings)?;
|
||||
|
||||
assert!(Arc::ptr_eq(&array, &rewritten));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_project_partial_json2_v2_without_remainder() -> Result<()> {
|
||||
let fields = Fields::from(vec![Arc::new(Field::new("hot", DataType::Int64, true))]);
|
||||
@@ -1097,7 +1181,7 @@ mod test {
|
||||
None,
|
||||
));
|
||||
let field = Field::new("data", DataType::Struct(fields), true).with_extension_type(
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))),
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default()))),
|
||||
);
|
||||
|
||||
let projected = JsonArray::from(&array).project_to_v2(&field, field.data_type())?;
|
||||
@@ -1124,6 +1208,18 @@ mod test {
|
||||
)
|
||||
);
|
||||
|
||||
let Value::Object(mut remainder) = json!({"count": 1}) else {
|
||||
unreachable!();
|
||||
};
|
||||
let error = merge_explicit_value(
|
||||
&mut remainder,
|
||||
"count".to_string(),
|
||||
json!(1),
|
||||
&mut Vec::new(),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(error.to_string().contains("cannot merge 'count'"));
|
||||
|
||||
let Value::Object(mut remainder) = json!({"nested": {"count": 1}}) else {
|
||||
unreachable!();
|
||||
};
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -15,10 +15,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::ArrayRef;
|
||||
#[cfg(test)]
|
||||
use arrow_schema::ArrowError;
|
||||
use arrow_schema::{DataType, Field};
|
||||
#[cfg(test)]
|
||||
use arrow_schema::{ArrowError, DataType, Field};
|
||||
use parquet_variant::{ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal16};
|
||||
#[cfg(test)]
|
||||
use parquet_variant_compute::VariantArrayBuilder;
|
||||
@@ -27,8 +24,7 @@ use parquet_variant_json::VariantToJson;
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::error::{ArrowComputeSnafu, Result};
|
||||
#[cfg(test)]
|
||||
use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant};
|
||||
use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, decode_json_variant};
|
||||
|
||||
/// Returns the canonical Arrow field for an unshredded Parquet Variant array.
|
||||
pub fn variant_field(name: impl Into<String>, nullable: bool) -> Field {
|
||||
@@ -43,6 +39,10 @@ pub fn variant_field(name: impl Into<String>, nullable: bool) -> Field {
|
||||
.with_extension_type(VariantType)
|
||||
}
|
||||
|
||||
/// Encodes JSON values as an unshredded Parquet Variant array.
|
||||
///
|
||||
/// `None` represents an Arrow null while `Some(Value::Null)` represents a JSON
|
||||
/// null, preserving the distinction required by JSON2.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn json_values_to_variant(values: &[Option<serde_json::Value>]) -> Result<ArrayRef> {
|
||||
let mut builder = VariantArrayBuilder::new(values.len());
|
||||
@@ -57,7 +57,7 @@ pub(crate) fn json_values_to_variant(values: &[Option<serde_json::Value>]) -> Re
|
||||
|
||||
/// Encodes JSON variants as an unshredded Parquet Variant array.
|
||||
#[cfg(test)]
|
||||
fn json_variants_to_variant(values: &[Option<JsonVariant>]) -> Result<ArrayRef> {
|
||||
pub(crate) fn json_variants_to_variant(values: &[Option<JsonVariant>]) -> Result<ArrayRef> {
|
||||
let mut builder = VariantArrayBuilder::new(values.len());
|
||||
for value in values {
|
||||
match value {
|
||||
@@ -68,8 +68,7 @@ fn json_variants_to_variant(values: &[Option<JsonVariant>]) -> Result<ArrayRef>
|
||||
Ok(ArrayRef::from(builder.build()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn append_json_variant(
|
||||
pub(super) fn append_json_variant(
|
||||
builder: &mut impl VariantBuilderExt,
|
||||
value: &JsonVariant,
|
||||
) -> std::result::Result<(), ArrowError> {
|
||||
@@ -115,7 +114,52 @@ fn append_json_variant(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(super) fn append_json_variant_ref(
|
||||
builder: &mut impl VariantBuilderExt,
|
||||
value: &JsonVariantRef<'_>,
|
||||
) -> std::result::Result<(), ArrowError> {
|
||||
match value {
|
||||
JsonVariantRef::Null => builder.append_value(Variant::Null),
|
||||
JsonVariantRef::Bool(value) => builder.append_value(*value),
|
||||
JsonVariantRef::Number(JsonNumber::PosInt(value)) => {
|
||||
if let Ok(value) = i64::try_from(*value) {
|
||||
builder.append_value(value);
|
||||
} else {
|
||||
append_large_u64(builder, *value)?;
|
||||
}
|
||||
}
|
||||
JsonVariantRef::Number(JsonNumber::NegInt(value)) => builder.append_value(*value),
|
||||
JsonVariantRef::Number(JsonNumber::Float(value)) => {
|
||||
if value.0.is_finite() {
|
||||
builder.append_value(value.0)
|
||||
} else {
|
||||
builder.append_value("NaN")
|
||||
}
|
||||
}
|
||||
JsonVariantRef::String(value) => builder.append_value(*value),
|
||||
JsonVariantRef::Array(values) => {
|
||||
let mut list = builder.try_new_list()?;
|
||||
for value in values {
|
||||
append_json_variant_ref(&mut list, value)?;
|
||||
}
|
||||
list.finish();
|
||||
}
|
||||
JsonVariantRef::Object(values) => {
|
||||
let mut object = builder.try_new_object()?;
|
||||
for (name, value) in values {
|
||||
append_json_variant_ref(&mut ObjectFieldBuilder::new(name, &mut object), value)?;
|
||||
}
|
||||
object.finish();
|
||||
}
|
||||
JsonVariantRef::Variant(value) => {
|
||||
let value = decode_json_variant(value)
|
||||
.map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?;
|
||||
append_json_value(builder, &value)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn append_json_value(
|
||||
builder: &mut impl VariantBuilderExt,
|
||||
value: &serde_json::Value,
|
||||
@@ -157,11 +201,11 @@ fn append_json_value(
|
||||
|
||||
/// Parquet Variant has no unsigned integer primitive. Treat u64 as i64 first, then use Decimal16
|
||||
/// to represent large (larger than i64::MAX) u64.
|
||||
#[cfg(test)]
|
||||
fn append_large_u64(
|
||||
builder: &mut impl VariantBuilderExt,
|
||||
value: u64,
|
||||
) -> std::result::Result<(), ArrowError> {
|
||||
// Parquet Variant has no unsigned integer primitive. Decimal16 preserves the full u64 range.
|
||||
let value = VariantDecimal16::try_new(value as i128, 0).map_err(|e| {
|
||||
ArrowError::InvalidArgumentError(format!(
|
||||
"Failed to encode JSON large integer as Variant Decimal16: {e}"
|
||||
|
||||
@@ -59,6 +59,7 @@ use common_recordbatch::error::StreamTimeoutSnafu;
|
||||
use common_telemetry::logging::SlowQueryOptions;
|
||||
use common_telemetry::{debug, error, tracing};
|
||||
use dashmap::DashMap;
|
||||
use datafusion::dataframe::DataFrame;
|
||||
use datafusion::physical_plan::ExecutionPlan;
|
||||
use datafusion_expr::LogicalPlan;
|
||||
use futures::{Stream, StreamExt, future};
|
||||
@@ -847,6 +848,22 @@ impl Instance {
|
||||
|
||||
query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?;
|
||||
|
||||
// TQL EXPLAIN/ANALYZE formats are consumed from the query context at
|
||||
// execution time (see `optimize_physical_plan`); re-apply the side
|
||||
// effect of `plan_tql` that was lost when the plan was built during
|
||||
// Describe. `explain_format` is per-query state, so this never
|
||||
// overwrites anything.
|
||||
if let Some(Statement::Tql(tql)) = &stmt {
|
||||
let format = match tql {
|
||||
Tql::Explain(explain) => explain.format.as_ref(),
|
||||
Tql::Analyze(analyze) => analyze.format.as_ref(),
|
||||
Tql::Eval(_) => None,
|
||||
};
|
||||
if let Some(format) = format {
|
||||
query_ctx.set_explain_format(format.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let query = stmt
|
||||
.as_ref()
|
||||
.map(|s| s.to_string())
|
||||
@@ -928,6 +945,59 @@ impl Instance {
|
||||
vec![result]
|
||||
}
|
||||
|
||||
/// Builds the [`DataFrame`] for an information-schema-backed `SHOW`
|
||||
/// statement; `None` for other statements. The future is boxed to keep
|
||||
/// `do_describe_inner`'s state machine small.
|
||||
fn show_statement_dataframe<'a>(
|
||||
&'a self,
|
||||
stmt: &'a Statement,
|
||||
query_ctx: &'a QueryContextRef,
|
||||
) -> Pin<Box<dyn Future<Output = Option<query::error::Result<DataFrame>>> + Send + 'a>> {
|
||||
Box::pin(async move {
|
||||
let engine = &self.query_engine;
|
||||
let catalog_manager = self.catalog_manager();
|
||||
let ctx = query_ctx.clone();
|
||||
let dataframe = match stmt {
|
||||
Statement::ShowDatabases(show) => {
|
||||
query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowTables(show) => {
|
||||
query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowViews(show) => {
|
||||
query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowFlows(show) => {
|
||||
query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowColumns(show) => {
|
||||
query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowTableStatus(show) => {
|
||||
query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx)
|
||||
.await
|
||||
}
|
||||
Statement::ShowCharset(kind) => {
|
||||
query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowCollation(kind) => {
|
||||
query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowIndex(show) => {
|
||||
query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowRegion(show) => {
|
||||
query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
Statement::ShowProcesslist(show) => {
|
||||
query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await
|
||||
}
|
||||
_ => return None,
|
||||
};
|
||||
Some(dataframe)
|
||||
})
|
||||
}
|
||||
|
||||
async fn do_describe_inner(
|
||||
&self,
|
||||
stmt: Statement,
|
||||
@@ -947,6 +1017,37 @@ impl Instance {
|
||||
let plannable = is_inner_plannable(&stmt)
|
||||
|| matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));
|
||||
|
||||
if let Statement::Tql(tql) = stmt {
|
||||
// TQL produces a logical plan; describe it from the plan so the
|
||||
// extended-protocol RowDescription matches the executed DataRows.
|
||||
self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx)
|
||||
.await?;
|
||||
let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?;
|
||||
return self
|
||||
.query_engine
|
||||
.describe(plan, query_ctx)
|
||||
.await
|
||||
.map(Some)
|
||||
.context(error::DescribeStatementSnafu);
|
||||
}
|
||||
|
||||
// Describe SHOW statements from the same projection the executor builds.
|
||||
if let Some(dataframe) = self
|
||||
.show_statement_dataframe(&stmt, &query_ctx)
|
||||
.await
|
||||
.transpose()
|
||||
.context(PlanStatementSnafu)?
|
||||
{
|
||||
self.check_sql_permission(&stmt, &query_ctx).await?;
|
||||
let plan = dataframe.into_unoptimized_plan();
|
||||
return self
|
||||
.query_engine
|
||||
.describe(plan, query_ctx)
|
||||
.await
|
||||
.map(Some)
|
||||
.context(error::DescribeStatementSnafu);
|
||||
}
|
||||
|
||||
if plannable {
|
||||
self.check_sql_permission(&stmt, &query_ctx).await?;
|
||||
|
||||
|
||||
@@ -12,6 +12,8 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#![recursion_limit = "256"]
|
||||
|
||||
pub mod error;
|
||||
pub mod events;
|
||||
pub mod frontend;
|
||||
|
||||
@@ -487,10 +487,14 @@ mod tests {
|
||||
.await;
|
||||
|
||||
send_mock_reply(mailbox, rx, |id| {
|
||||
Ok(new_open_region_reply(
|
||||
Ok(new_open_region_reply_with_error(
|
||||
id,
|
||||
false,
|
||||
Some("test mocked".to_string()),
|
||||
Some(InstructionError {
|
||||
code: StatusCode::StorageUnavailable,
|
||||
message: "test mocked".to_string(),
|
||||
retry_hint: RetryHint::Retryable,
|
||||
}),
|
||||
))
|
||||
});
|
||||
|
||||
|
||||
@@ -350,6 +350,9 @@ impl SyncRegion {
|
||||
mod tests {
|
||||
use std::assert_matches;
|
||||
|
||||
use common_error::ext::RetryHint;
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_meta::instruction::InstructionError;
|
||||
use common_meta::peer::Peer;
|
||||
use common_meta::rpc::router::{Region, RegionRoute};
|
||||
use store_api::region_engine::SyncRegionFromRequest;
|
||||
@@ -359,7 +362,9 @@ mod tests {
|
||||
use crate::procedure::repartition::group::GroupPrepareResult;
|
||||
use crate::procedure::repartition::group::sync_region::SyncRegion;
|
||||
use crate::procedure::repartition::test_util::{TestingEnv, new_persistent_context};
|
||||
use crate::procedure::test_util::{new_sync_region_reply, send_mock_reply};
|
||||
use crate::procedure::test_util::{
|
||||
new_sync_region_reply, new_sync_region_reply_with_error, send_mock_reply,
|
||||
};
|
||||
use crate::service::mailbox::Channel;
|
||||
|
||||
#[test]
|
||||
@@ -455,4 +460,46 @@ mod tests {
|
||||
let err = sync_region.sync_regions(&mut ctx).await.unwrap_err();
|
||||
assert_matches!(err, Error::RetryLater { .. });
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sync_regions_retryable_instruction_error() {
|
||||
let mut env = TestingEnv::new();
|
||||
let table_id = 1024;
|
||||
let region_id = RegionId::new(table_id, 3);
|
||||
let mut persistent_context = new_persistent_context(table_id, vec![], vec![]);
|
||||
persistent_context.group_prepare_result = Some(test_prepare_result(table_id));
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(1);
|
||||
env.mailbox_ctx
|
||||
.insert_heartbeat_response_receiver(Channel::Datanode(1), tx)
|
||||
.await;
|
||||
send_mock_reply(env.mailbox_ctx.mailbox().clone(), rx, move |id| {
|
||||
Ok(new_sync_region_reply_with_error(
|
||||
id,
|
||||
region_id,
|
||||
false,
|
||||
true,
|
||||
Some(InstructionError {
|
||||
code: StatusCode::StorageUnavailable,
|
||||
message: "manifest delta disappeared".to_string(),
|
||||
retry_hint: RetryHint::Retryable,
|
||||
}),
|
||||
))
|
||||
});
|
||||
|
||||
let mut ctx = env.create_context(persistent_context);
|
||||
let sync_region = SyncRegion {
|
||||
region_routes: vec![RegionRoute {
|
||||
region: Region {
|
||||
id: region_id,
|
||||
..Default::default()
|
||||
},
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
..Default::default()
|
||||
}],
|
||||
};
|
||||
|
||||
let err = sync_region.sync_regions(&mut ctx).await.unwrap_err();
|
||||
assert_matches!(err, Error::RetryLater { .. });
|
||||
}
|
||||
}
|
||||
|
||||
@@ -288,6 +288,23 @@ pub fn new_sync_region_reply(
|
||||
ready: bool,
|
||||
exists: bool,
|
||||
error: Option<String>,
|
||||
) -> MailboxMessage {
|
||||
new_sync_region_reply_with_error(
|
||||
id,
|
||||
region_id,
|
||||
ready,
|
||||
exists,
|
||||
legacy_instruction_error(error),
|
||||
)
|
||||
}
|
||||
|
||||
/// Generates a [InstructionReply::SyncRegions] reply with a structured error.
|
||||
pub fn new_sync_region_reply_with_error(
|
||||
id: u64,
|
||||
region_id: RegionId,
|
||||
ready: bool,
|
||||
exists: bool,
|
||||
error: Option<InstructionError>,
|
||||
) -> MailboxMessage {
|
||||
MailboxMessage {
|
||||
id,
|
||||
@@ -301,7 +318,7 @@ pub fn new_sync_region_reply(
|
||||
region_id,
|
||||
ready,
|
||||
exists,
|
||||
error: legacy_instruction_error(error),
|
||||
error,
|
||||
},
|
||||
])))
|
||||
.unwrap(),
|
||||
|
||||
@@ -26,6 +26,7 @@ use api::v1::meta::{
|
||||
use common_event_recorder::{PersistentEventContext, ProcedureEventInput};
|
||||
use common_meta::key::TableMetadataManagerRef;
|
||||
use common_meta::key::table_name::TableNameKey;
|
||||
use common_meta::peer::Peer;
|
||||
use common_meta::procedure_executor::ExecutorContext;
|
||||
use common_meta::rpc::ddl::{
|
||||
CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY,
|
||||
@@ -195,7 +196,7 @@ impl procedure_service_server::ProcedureService for Metasrv {
|
||||
let from_peer = self
|
||||
.lookup_datanode_peer(from_peer)
|
||||
.await?
|
||||
.context(error::PeerUnavailableSnafu { peer_id: from_peer })?;
|
||||
.unwrap_or_else(|| Peer::empty(from_peer));
|
||||
let to_peer = self
|
||||
.lookup_datanode_peer(to_peer)
|
||||
.await?
|
||||
|
||||
@@ -65,7 +65,7 @@ log-store = { workspace = true }
|
||||
mito-codec.workspace = true
|
||||
moka = { workspace = true, features = ["sync", "future"] }
|
||||
object-store = { workspace = true, features = ["testing"] }
|
||||
parquet = { workspace = true, features = ["async"] }
|
||||
parquet = { workspace = true, features = ["async", "variant_experimental"] }
|
||||
paste.workspace = true
|
||||
pin-project.workspace = true
|
||||
prometheus.workspace = true
|
||||
|
||||
@@ -34,6 +34,7 @@ use common_datasource::compression::CompressionType;
|
||||
use common_telemetry::warn;
|
||||
use datatypes::arrow::buffer::BooleanBuffer;
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::types::json_type::JsonNativeType;
|
||||
use datatypes::value::Value;
|
||||
use datatypes::vectors::VectorRef;
|
||||
use index::bloom_filter_index::{BloomFilterIndexCache, BloomFilterIndexCacheRef};
|
||||
@@ -2065,8 +2066,7 @@ impl SelectorResultValue {
|
||||
SelectorResult::Flat(batches) => batches.iter().map(record_batch_estimated_size).sum(),
|
||||
};
|
||||
result_size
|
||||
+ self.json_target_types.len()
|
||||
* (mem::size_of::<ColumnId>() + mem::size_of::<ConcreteDataType>())
|
||||
+ self.json_target_types.len() * (size_of::<ColumnId>() + size_of::<JsonNativeType>())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
|
||||
mod buckets;
|
||||
pub mod compactor;
|
||||
mod json2;
|
||||
pub mod memory_manager;
|
||||
pub mod picker;
|
||||
mod reader;
|
||||
@@ -31,6 +32,9 @@ use common_meta::key::SchemaMetadataManagerRef;
|
||||
use common_telemetry::{debug, error};
|
||||
use common_time::TimeToLive;
|
||||
use common_time::range::TimestampRange;
|
||||
pub(crate) use json2::{
|
||||
Json2RewritePlans, collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema,
|
||||
};
|
||||
pub use scheduler::CompactionRequest;
|
||||
pub(crate) use scheduler::{CompactionExecution, CompactionPickFinished, CompactionScheduler};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
@@ -0,0 +1,661 @@
|
||||
// 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 std::sync::Arc;
|
||||
|
||||
use arrow_schema::extension::ExtensionType;
|
||||
use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, SchemaRef};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::extension::json::{JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, JsonMetadata};
|
||||
use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint};
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::types::json_type::JsonNativeType;
|
||||
use datatypes::vectors::json::array::JsonArray;
|
||||
use datatypes::vectors::json::json2_physical_data_type;
|
||||
use parquet::arrow::parquet_to_arrow_schema;
|
||||
use parquet::file::metadata::ParquetMetaData;
|
||||
use snafu::{OptionExt, ResultExt, ensure};
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
|
||||
use crate::error::{
|
||||
ConvertValueSnafu, DataTypeMismatchSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result,
|
||||
};
|
||||
|
||||
/// Plan for rewriting one JSON2 column to a fixed compaction layout.
|
||||
///
|
||||
/// A compaction input may contain v1 and v2 SSTs with different physical schemas. This plan is
|
||||
/// derived only from the current region metadata and remains fixed while all input batches are
|
||||
/// decoded and rewritten. It therefore prevents source-only paths from expanding the output
|
||||
/// schema without a bound.
|
||||
pub(crate) struct Json2RewritePlan {
|
||||
/// User-defined settings used to encode logical values.
|
||||
logical_settings: JsonSettings,
|
||||
/// Fixed settings used to build the target physical layout.
|
||||
pub(super) target_layout: JsonSettings,
|
||||
}
|
||||
|
||||
/// JSON2 rewrite plans keyed by logical column name.
|
||||
pub(crate) type Json2RewritePlans = HashMap<String, Json2RewritePlan>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct Json2LeafPathStats {
|
||||
rows: u64,
|
||||
data_type: JsonNativeType,
|
||||
is_type_conflicted: bool,
|
||||
}
|
||||
|
||||
/// Builds the JSON2 rewrite plans for a compaction.
|
||||
///
|
||||
/// Type hints from current region metadata are always retained. Existing explicit dynamic paths
|
||||
/// from all input SST schemas are ranked once to produce a fixed layout; paths found only in a v2
|
||||
/// remainder are deliberately not promoted. [`rewrite_json2_batch`] decodes inputs and rewrites
|
||||
/// them according to these plans. Non-JSON2 columns are omitted from the returned map.
|
||||
///
|
||||
/// Returns an error when a JSON2 column has invalid or missing extension metadata, or when its
|
||||
/// output layout is not v2. Legacy region metadata is upgraded in memory by the region opener
|
||||
/// before compaction reaches this function.
|
||||
pub(super) fn collect_json2_rewrite_plans_from_parquet(
|
||||
metadata: &RegionMetadataRef,
|
||||
parquet_metadata: &[Arc<ParquetMetaData>],
|
||||
) -> Result<Json2RewritePlans> {
|
||||
let schemas = parquet_metadata
|
||||
.iter()
|
||||
.map(|metadata| {
|
||||
let file = metadata.file_metadata();
|
||||
let schema = parquet_to_arrow_schema(file.schema_descr(), file.key_value_metadata())
|
||||
.map_err(|error| {
|
||||
InvalidRecordBatchSnafu {
|
||||
reason: format!("Failed to read compaction input Arrow schema: {error}"),
|
||||
}
|
||||
.build()
|
||||
})?;
|
||||
let rows = metadata
|
||||
.row_groups()
|
||||
.iter()
|
||||
.map(|x| x.num_rows())
|
||||
.sum::<i64>() as u64;
|
||||
Ok((Arc::new(schema), rows))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
collect_json2_rewrite_plans(metadata, &schemas)
|
||||
}
|
||||
|
||||
/// Builds JSON2 rewrite plans from existing physical schemas.
|
||||
///
|
||||
/// Each source row count weights all of its existing explicit leaves. Ordinary paths are never
|
||||
/// discovered from the v2 remainder, so this operation cannot unexpectedly promote opaque data.
|
||||
pub(crate) fn collect_json2_rewrite_plans(
|
||||
metadata: &RegionMetadataRef,
|
||||
schemas: &[(SchemaRef, u64)],
|
||||
) -> Result<Json2RewritePlans> {
|
||||
let json2_columns = metadata
|
||||
.column_metadatas
|
||||
.iter()
|
||||
.filter_map(|x| {
|
||||
x.column_schema
|
||||
.data_type
|
||||
.is_json2()
|
||||
.then_some(&x.column_schema)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let mut plans = HashMap::with_capacity(json2_columns.len());
|
||||
for column in json2_columns {
|
||||
let extension = column
|
||||
.extension_type::<Json2ExtensionType>()
|
||||
.context(DataTypeMismatchSnafu)?
|
||||
.with_context(|| InvalidRecordBatchSnafu {
|
||||
reason: format!("JSON2 column '{}' has no extension metadata", column.name),
|
||||
})?;
|
||||
// Source SSTs may use v1, but current region metadata is copied to the rewritten output.
|
||||
// Since the target physical layout is always v2, v1 metadata would produce an
|
||||
// inconsistent persisted field. The region opener normally upgraded this metadata already.
|
||||
ensure!(
|
||||
extension.metadata().is_version_2(),
|
||||
InvalidRecordBatchSnafu {
|
||||
reason: format!("JSON2 column '{}' is not layout v2", column.name),
|
||||
}
|
||||
);
|
||||
|
||||
let settings = extension.metadata().json_settings();
|
||||
let hint_paths = settings
|
||||
.type_hints()
|
||||
.iter()
|
||||
.map(|hint| hint.path.iter().map(String::as_str).collect::<Vec<_>>())
|
||||
.collect::<HashSet<_>>();
|
||||
let mut stats = HashMap::new();
|
||||
for (schema, rows) in schemas {
|
||||
let Some((_, field)) = schema.fields().find(&column.name) else {
|
||||
continue;
|
||||
};
|
||||
collect_json2_path_stats(field, *rows, &hint_paths, &mut stats)?;
|
||||
}
|
||||
|
||||
let mut hints = settings.type_hints().to_vec();
|
||||
hints.extend(select_dynamic_hints(settings, &hint_paths, &stats));
|
||||
let target_layout = JsonSettings::try_new(hints, Some(0)).context(DataTypeMismatchSnafu)?;
|
||||
plans.insert(
|
||||
column.name.clone(),
|
||||
Json2RewritePlan {
|
||||
logical_settings: settings.clone(),
|
||||
target_layout,
|
||||
},
|
||||
);
|
||||
}
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
fn collect_json2_path_stats<'a>(
|
||||
field: &'a Field,
|
||||
rows: u64,
|
||||
hint_paths: &HashSet<Vec<&str>>,
|
||||
stats: &mut HashMap<Vec<&'a str>, Json2LeafPathStats>,
|
||||
) -> Result<()> {
|
||||
let ArrowDataType::Struct(fields) = field.data_type() else {
|
||||
return InvalidRecordBatchSnafu {
|
||||
reason: format!("JSON2 column '{}' is not a struct", field.name()),
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
let mut paths = Vec::new();
|
||||
for field in fields {
|
||||
if field.name() == JSON2_REMAINDER_FIELD_NAME {
|
||||
continue;
|
||||
}
|
||||
collect_leaf_path_types(field, &mut Vec::new(), &mut paths)?;
|
||||
}
|
||||
|
||||
for (path, data_type) in paths {
|
||||
if hint_paths.contains(path.as_slice()) {
|
||||
continue;
|
||||
}
|
||||
let Some(stat) = stats.get_mut(&path) else {
|
||||
stats.insert(
|
||||
path,
|
||||
Json2LeafPathStats {
|
||||
rows,
|
||||
data_type,
|
||||
is_type_conflicted: false,
|
||||
},
|
||||
);
|
||||
continue;
|
||||
};
|
||||
if stat.data_type != data_type {
|
||||
stat.is_type_conflicted = true;
|
||||
} else {
|
||||
stat.rows += rows;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn collect_leaf_path_types<'a>(
|
||||
field: &'a Field,
|
||||
path: &mut Vec<&'a str>,
|
||||
paths: &mut Vec<(Vec<&'a str>, JsonNativeType)>,
|
||||
) -> Result<()> {
|
||||
path.push(field.name());
|
||||
if let ArrowDataType::Struct(fields) = field.data_type()
|
||||
&& !fields.is_empty()
|
||||
{
|
||||
for field in fields {
|
||||
collect_leaf_path_types(field, path, paths)?;
|
||||
}
|
||||
} else {
|
||||
let json_type =
|
||||
JsonNativeType::try_from(field.data_type()).context(DataTypeMismatchSnafu)?;
|
||||
paths.push((path.clone(), json_type));
|
||||
}
|
||||
path.pop();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn select_dynamic_hints(
|
||||
settings: &JsonSettings,
|
||||
hint_paths: &HashSet<Vec<&str>>,
|
||||
stats: &HashMap<Vec<&str>, Json2LeafPathStats>,
|
||||
) -> Vec<JsonTypeHint> {
|
||||
let all_paths = stats
|
||||
.keys()
|
||||
.map(Vec::as_slice)
|
||||
.chain(hint_paths.iter().map(Vec::as_slice))
|
||||
.collect::<HashSet<_>>();
|
||||
let has_ancestor_path =
|
||||
|path: &[&str]| (1..path.len()).any(|len| all_paths.contains(&path[..len]));
|
||||
|
||||
let prefixes = all_paths
|
||||
.iter()
|
||||
.copied()
|
||||
.flat_map(|path| (1..path.len()).map(|len| &path[..len]))
|
||||
.collect::<HashSet<_>>();
|
||||
let has_descendant_path = |path: &[&str]| prefixes.contains(path);
|
||||
|
||||
let mut candidates = stats
|
||||
.iter()
|
||||
.filter(|(path, stat)| {
|
||||
!stat.is_type_conflicted
|
||||
// TODO(LFC): Instead of "primitive only", consider retaining stable compound types
|
||||
// that are safe to write to Parquet, as flush does. Or better, unite the two
|
||||
// selection process.
|
||||
&& stat.data_type.is_primitive()
|
||||
&& !has_ancestor_path(path)
|
||||
&& !has_descendant_path(path)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
candidates.sort_unstable_by(|(x_path, x), (y_path, y)| {
|
||||
y.rows.cmp(&x.rows).then_with(|| x_path.cmp(y_path))
|
||||
});
|
||||
candidates
|
||||
.into_iter()
|
||||
.take(
|
||||
settings
|
||||
.max_auto_expanded_paths()
|
||||
.unwrap_or(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS) as usize,
|
||||
)
|
||||
.map(|(path, stat)| JsonTypeHint {
|
||||
path: path.iter().map(|x| (*x).to_owned()).collect(),
|
||||
data_type: ConcreteDataType::from_arrow_type(&stat.data_type.as_arrow_type()),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Replaces JSON2 physical field types according to the computed rewrite plans.
|
||||
pub(crate) fn rewrite_json2_schema(schema: &SchemaRef, plans: &Json2RewritePlans) -> SchemaRef {
|
||||
if plans.is_empty() {
|
||||
return schema.clone();
|
||||
}
|
||||
let fields = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|field| {
|
||||
let Some(plan) = plans.get(field.name()) else {
|
||||
return field.clone();
|
||||
};
|
||||
let mut field = Field::clone(field);
|
||||
field.set_data_type(json2_physical_data_type(&plan.target_layout));
|
||||
field = field.with_extension_type(Json2ExtensionType::new(Arc::new(
|
||||
JsonMetadata::new(plan.logical_settings.clone()),
|
||||
)));
|
||||
Arc::new(field)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
|
||||
}
|
||||
|
||||
/// Rewrites JSON2 columns in `batch` according to the computed plans.
|
||||
pub(crate) fn rewrite_json2_batch(
|
||||
batch: RecordBatch,
|
||||
plans: &Json2RewritePlans,
|
||||
) -> Result<RecordBatch> {
|
||||
if plans.is_empty() {
|
||||
return Ok(batch);
|
||||
}
|
||||
let mut fields = Vec::with_capacity(batch.num_columns());
|
||||
let mut columns = Vec::with_capacity(batch.num_columns());
|
||||
|
||||
for (field, array) in batch.schema_ref().fields().iter().zip(batch.columns()) {
|
||||
let Some(plan) = plans.get(field.name()) else {
|
||||
fields.push(field.clone());
|
||||
columns.push(array.clone());
|
||||
continue;
|
||||
};
|
||||
|
||||
let array = JsonArray::from(array)
|
||||
.rewrite_to_v2(field, &plan.logical_settings, &plan.target_layout)
|
||||
.context(ConvertValueSnafu)?;
|
||||
debug_assert_eq!(
|
||||
&json2_physical_data_type(&plan.target_layout),
|
||||
array.data_type()
|
||||
);
|
||||
|
||||
let mut field = Field::clone(field);
|
||||
field.set_data_type(array.data_type().clone());
|
||||
field = field.with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new(
|
||||
plan.logical_settings.clone(),
|
||||
))));
|
||||
fields.push(Arc::new(field));
|
||||
columns.push(array);
|
||||
}
|
||||
|
||||
let schema = Arc::new(Schema::new_with_metadata(
|
||||
fields,
|
||||
batch.schema_ref().metadata().clone(),
|
||||
));
|
||||
RecordBatch::try_new(schema, columns).context(NewRecordBatchSnafu)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use datatypes::extension::json::{
|
||||
JSON2_REMAINDER_FIELD_NAME, Json2PhysicalLayout, JsonMetadata,
|
||||
};
|
||||
use datatypes::json::JsonTypeHint;
|
||||
use datatypes::prelude::{ConcreteDataType, DataType};
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_select_dynamic_hints_rejects_type_and_prefix_conflicts()
|
||||
-> Result<(), Box<dyn std::error::Error>> {
|
||||
let settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["hint".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
Some(2),
|
||||
)?;
|
||||
let stat = |rows, data_type, is_type_conflicted| Json2LeafPathStats {
|
||||
rows,
|
||||
data_type,
|
||||
is_type_conflicted,
|
||||
};
|
||||
let stats = HashMap::from([
|
||||
(
|
||||
vec!["hint", "nested"],
|
||||
stat(10, JsonNativeType::String, false),
|
||||
),
|
||||
(vec!["popular"], stat(9, JsonNativeType::String, false)),
|
||||
(
|
||||
vec!["popular", "nested"],
|
||||
stat(8, JsonNativeType::u64(), false),
|
||||
),
|
||||
(
|
||||
vec!["type_conflicted"],
|
||||
stat(7, JsonNativeType::String, true),
|
||||
),
|
||||
(
|
||||
vec!["array"],
|
||||
stat(
|
||||
6,
|
||||
JsonNativeType::Array(Box::new(JsonNativeType::String)),
|
||||
false,
|
||||
),
|
||||
),
|
||||
(vec!["variant"], stat(5, JsonNativeType::Variant, false)),
|
||||
(vec!["tie_a"], stat(2, JsonNativeType::String, false)),
|
||||
(vec!["tie_b"], stat(2, JsonNativeType::Bool, false)),
|
||||
]);
|
||||
|
||||
let hint_paths = settings
|
||||
.type_hints()
|
||||
.iter()
|
||||
.map(|hint| hint.path.iter().map(String::as_str).collect::<Vec<_>>())
|
||||
.collect::<HashSet<_>>();
|
||||
let hints = select_dynamic_hints(&settings, &hint_paths, &stats);
|
||||
assert_eq!(
|
||||
vec![vec!["tie_a".to_string()], vec!["tie_b".to_string()]],
|
||||
hints.into_iter().map(|x| x.path).collect::<Vec<_>>()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rewrite_json2_v1_batch_to_target_layout() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["kind".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
Some(0),
|
||||
)?;
|
||||
let target = json2_physical_data_type(&settings);
|
||||
let plans = HashMap::from([(
|
||||
"j".to_string(),
|
||||
Json2RewritePlan {
|
||||
logical_settings: settings.clone(),
|
||||
target_layout: settings,
|
||||
},
|
||||
)]);
|
||||
let values = [
|
||||
json!({"kind": "a", "extra": {"x": 1}}),
|
||||
json!({"kind": "b", "extra": {"x": 2}}),
|
||||
];
|
||||
let source_settings = JsonSettings::default();
|
||||
let source_extension =
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1(source_settings.clone())));
|
||||
let mut source_column = ColumnSchema::new(
|
||||
"j",
|
||||
ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
|
||||
true,
|
||||
);
|
||||
source_column.with_extension_type(&source_extension);
|
||||
let mut source_builder = source_column.data_type.create_mutable_vector(values.len());
|
||||
for value in &values {
|
||||
let value = source_settings.encode(value.clone())?;
|
||||
source_builder.try_push_value_ref(&value.as_value_ref())?;
|
||||
}
|
||||
let source = source_builder.to_vector().to_arrow_array();
|
||||
let field =
|
||||
Field::new("j", source.data_type().clone(), true).with_extension_type(source_extension);
|
||||
let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![source])?;
|
||||
|
||||
let batch = rewrite_json2_batch(batch, &plans)?;
|
||||
let field = batch.schema_ref().field(0);
|
||||
assert!(Json2PhysicalLayout::try_from_root(field)?.is_version_2());
|
||||
assert_eq!(&target, field.data_type());
|
||||
let projected =
|
||||
JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?;
|
||||
let projected = JsonArray::from(&projected);
|
||||
for (i, expected) in values.into_iter().enumerate() {
|
||||
assert_eq!(expected, projected.try_get_value(i)?);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rewrite_json2_v2_source_to_narrower_target() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let target_settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["kind".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
Some(0),
|
||||
)?;
|
||||
let target_extension =
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new(target_settings.clone())));
|
||||
let target_type = json2_physical_data_type(&target_settings);
|
||||
let plans = HashMap::from([(
|
||||
"j".to_string(),
|
||||
Json2RewritePlan {
|
||||
logical_settings: target_settings.clone(),
|
||||
target_layout: target_settings,
|
||||
},
|
||||
)]);
|
||||
|
||||
let source_settings = JsonSettings::try_new(
|
||||
vec![
|
||||
JsonTypeHint {
|
||||
path: vec!["kind".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
},
|
||||
JsonTypeHint {
|
||||
path: vec!["source_only".to_string()],
|
||||
data_type: ConcreteDataType::int64_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
},
|
||||
],
|
||||
None,
|
||||
)?;
|
||||
let source_extension =
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new(source_settings.clone())));
|
||||
let mut source_column = ColumnSchema::new(
|
||||
"j",
|
||||
ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
|
||||
true,
|
||||
);
|
||||
source_column.with_extension_type(&source_extension);
|
||||
let expected = json!({
|
||||
"kind": "a",
|
||||
"source_only": 7,
|
||||
"dynamic": {"nested": true}
|
||||
});
|
||||
let mut builder = source_column.create_mutable_vector(1);
|
||||
let value = source_settings.encode(expected.clone())?;
|
||||
builder.try_push_value_ref(&value.as_value_ref())?;
|
||||
let source = builder.to_vector().to_arrow_array();
|
||||
let source_field =
|
||||
Field::new("j", source.data_type().clone(), true).with_extension_type(source_extension);
|
||||
let source =
|
||||
JsonArray::from(&source).project_to_v2(&source_field, &ArrowDataType::Binary)?;
|
||||
let field = Field::new("j", ArrowDataType::Binary, true)
|
||||
.with_extension_type(target_extension.clone());
|
||||
let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![source])?;
|
||||
|
||||
let batch = rewrite_json2_batch(batch, &plans)?;
|
||||
let field = batch.schema_ref().field(0);
|
||||
assert!(Json2PhysicalLayout::try_from_root(field)?.is_version_2());
|
||||
assert_eq!(&target_type, field.data_type());
|
||||
let ArrowDataType::Struct(fields) = field.data_type() else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
vec![JSON2_REMAINDER_FIELD_NAME, "kind"],
|
||||
fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
|
||||
);
|
||||
|
||||
let projected =
|
||||
JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?;
|
||||
assert_eq!(expected, JsonArray::from(&projected).try_get_value(0)?);
|
||||
|
||||
let first_schema = batch.schema();
|
||||
let field =
|
||||
Field::new("j", ArrowDataType::Binary, true).with_extension_type(target_extension);
|
||||
let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![projected])?;
|
||||
let batch = rewrite_json2_batch(batch, &plans)?;
|
||||
assert_eq!(first_schema, batch.schema());
|
||||
|
||||
let field = batch.schema_ref().field(0);
|
||||
let projected =
|
||||
JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?;
|
||||
assert_eq!(expected, JsonArray::from(&projected).try_get_value(0)?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rewrite_json2_v2_source_to_wider_target() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let logical_settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["kind".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
Some(0),
|
||||
)?;
|
||||
let target_layout = JsonSettings::try_new(
|
||||
vec![
|
||||
JsonTypeHint {
|
||||
path: vec!["kind".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
},
|
||||
JsonTypeHint {
|
||||
path: vec!["promoted".to_string()],
|
||||
data_type: ConcreteDataType::int64_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
},
|
||||
],
|
||||
Some(0),
|
||||
)?;
|
||||
let target_type = json2_physical_data_type(&target_layout);
|
||||
let plans = HashMap::from([(
|
||||
"j".to_string(),
|
||||
Json2RewritePlan {
|
||||
logical_settings: logical_settings.clone(),
|
||||
target_layout,
|
||||
},
|
||||
)]);
|
||||
|
||||
let extension =
|
||||
Json2ExtensionType::new(Arc::new(JsonMetadata::new(logical_settings.clone())));
|
||||
let mut column = ColumnSchema::new(
|
||||
"j",
|
||||
ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
|
||||
true,
|
||||
);
|
||||
column.with_extension_type(&extension);
|
||||
let expected = json!({
|
||||
"kind": "a",
|
||||
"promoted": 7,
|
||||
"dynamic": {"nested": true}
|
||||
});
|
||||
let mut builder = column.create_mutable_vector(1);
|
||||
let value = logical_settings.encode(expected.clone())?;
|
||||
builder.try_push_value_ref(&value.as_value_ref())?;
|
||||
let source = builder.to_vector().to_arrow_array();
|
||||
assert_eq!(
|
||||
&json2_physical_data_type(&logical_settings),
|
||||
source.data_type()
|
||||
);
|
||||
let field =
|
||||
Field::new("j", source.data_type().clone(), true).with_extension_type(extension);
|
||||
let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![source])?;
|
||||
|
||||
let batch = rewrite_json2_batch(batch, &plans)?;
|
||||
let field = batch.schema_ref().field(0);
|
||||
assert_eq!(&target_type, field.data_type());
|
||||
let ArrowDataType::Struct(fields) = field.data_type() else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
vec![JSON2_REMAINDER_FIELD_NAME, "kind", "promoted"],
|
||||
fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
|
||||
);
|
||||
|
||||
let array = batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<datatypes::arrow::array::StructArray>()
|
||||
.unwrap();
|
||||
let promoted = array
|
||||
.column_by_name("promoted")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<datatypes::arrow::array::Int64Array>()
|
||||
.unwrap();
|
||||
assert_eq!(7, promoted.value(0));
|
||||
|
||||
let projected =
|
||||
JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?;
|
||||
assert_eq!(expected, JsonArray::from(&projected).try_get_value(0)?);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -12,26 +12,25 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY;
|
||||
use common_time::Timestamp;
|
||||
use common_time::range::TimestampRange;
|
||||
use common_time::timestamp::TimeUnit;
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_expr::Expr;
|
||||
use datatypes::extension::json::is_json2_extension_type;
|
||||
use datatypes::types::json_type::JsonNativeType;
|
||||
use parquet::arrow::parquet_to_arrow_schema;
|
||||
use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
use snafu::OptionExt;
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
|
||||
use crate::access_layer::AccessLayerRef;
|
||||
use crate::cache::{CacheManagerRef, CacheStrategy};
|
||||
use crate::error::{
|
||||
DataTypeMismatchSnafu, ParquetToArrowSchemaSnafu, Result, TimeRangePredicateOverflowSnafu,
|
||||
use crate::compaction::json2::{
|
||||
Json2RewritePlans, collect_json2_rewrite_plans_from_parquet, rewrite_json2_schema,
|
||||
};
|
||||
use crate::error::{InvalidRecordBatchSnafu, Result, TimeRangePredicateOverflowSnafu};
|
||||
use crate::read::FlatSource;
|
||||
use crate::read::flat_projection::FlatProjectionMapper;
|
||||
use crate::read::read_columns::ReadColumns;
|
||||
@@ -40,6 +39,7 @@ use crate::read::seq_scan::SeqScan;
|
||||
use crate::region::options::MergeMode;
|
||||
use crate::sst::file::FileHandle;
|
||||
use crate::sst::parquet::reader::MetadataCacheMetrics;
|
||||
use crate::sst::parquet::{Json2RewriteTargets, Json2TargetLayout};
|
||||
|
||||
/// Builders to create [BoxedRecordBatchStream] for compaction.
|
||||
pub(crate) struct CompactionSstReaderBuilder<'a> {
|
||||
@@ -57,20 +57,24 @@ impl CompactionSstReaderBuilder<'_> {
|
||||
/// Build a [FlatSource] that yields Arrow `RecordBatch`s from reading all the input SST files,
|
||||
/// for compaction. The schema of the [FlatSource] is unified.
|
||||
pub(crate) async fn build_flat_sst_reader(self) -> Result<FlatSource> {
|
||||
let scan_input = self.build_scan_input().await?;
|
||||
let parquet_metadata = self.collect_parquet_metadata().await?;
|
||||
let plans = collect_json2_rewrite_plans_from_parquet(&self.metadata, &parquet_metadata)?;
|
||||
let scan_input = self.build_scan_input(&parquet_metadata, &plans)?;
|
||||
|
||||
let schema = scan_input.mapper.output_schema();
|
||||
let schema = schema.arrow_schema();
|
||||
let schema = rewrite_json2_schema(schema.arrow_schema(), &plans);
|
||||
|
||||
let stream = SeqScan::new(scan_input)
|
||||
.build_flat_reader_for_compaction()
|
||||
.await?;
|
||||
Ok(FlatSource::new_stream(schema.clone(), stream))
|
||||
Ok(FlatSource::new_stream(schema, stream))
|
||||
}
|
||||
|
||||
async fn build_scan_input(self) -> Result<ScanInput> {
|
||||
let schema = self.metadata.schema.arrow_schema();
|
||||
let parquet_metadata = self.collect_parquet_metadata().await?;
|
||||
fn build_scan_input(
|
||||
self,
|
||||
parquet_metadata: &[Arc<ParquetMetaData>],
|
||||
plans: &Json2RewritePlans,
|
||||
) -> Result<ScanInput> {
|
||||
let batch_size = crate::batch_size::estimate_batch_size(
|
||||
parquet_metadata
|
||||
.iter()
|
||||
@@ -84,38 +88,6 @@ impl CompactionSstReaderBuilder<'_> {
|
||||
(row_group.num_rows() as u64, uncompressed_bytes)
|
||||
}),
|
||||
);
|
||||
let json_type_hint = if schema.fields().iter().any(is_json2_extension_type) {
|
||||
let mut json_type_hint = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.filter(|&field| is_json2_extension_type(field))
|
||||
.map(|field| (field.name().clone(), JsonNativeType::Null))
|
||||
.collect::<HashMap<_, _>>();
|
||||
|
||||
for metadata in &parquet_metadata {
|
||||
let file_metadata = metadata.file_metadata();
|
||||
let schema = parquet_to_arrow_schema(
|
||||
file_metadata.schema_descr(),
|
||||
file_metadata.key_value_metadata(),
|
||||
)
|
||||
.context(ParquetToArrowSchemaSnafu {
|
||||
file: "compaction input",
|
||||
})?;
|
||||
for field in schema.fields() {
|
||||
let Some(merged) = json_type_hint.get_mut(field.name()) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let json_type = JsonNativeType::try_from(field.data_type())
|
||||
.context(DataTypeMismatchSnafu)?;
|
||||
merged.merge(&json_type);
|
||||
}
|
||||
}
|
||||
|
||||
Some(json_type_hint)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let projection = (0..self.metadata.column_metadatas.len()).collect();
|
||||
let read_column_ids = self
|
||||
@@ -124,24 +96,39 @@ impl CompactionSstReaderBuilder<'_> {
|
||||
.iter()
|
||||
.map(|x| x.column_id)
|
||||
.collect::<Vec<_>>();
|
||||
let json_target_types = json_type_hint
|
||||
.as_ref()
|
||||
.map(|hint| {
|
||||
hint.iter()
|
||||
.filter_map(|(col_name, json_type)| {
|
||||
self.metadata
|
||||
.column_by_name(col_name)
|
||||
.map(|col| (col.column_id, json_type.clone()))
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let read_columns =
|
||||
ReadColumns::new(read_column_ids).with_json_target_types(json_target_types);
|
||||
let mapper =
|
||||
FlatProjectionMapper::new_with_read_columns(&self.metadata, projection, read_columns)?;
|
||||
|
||||
let mut json2_target_layouts = BTreeMap::new();
|
||||
for (name, plan) in plans {
|
||||
let Some(column) = self.metadata.column_by_name(name) else {
|
||||
continue;
|
||||
};
|
||||
let extension_metadata = column
|
||||
.column_schema
|
||||
.metadata()
|
||||
.get(EXTENSION_TYPE_METADATA_KEY)
|
||||
.cloned()
|
||||
.with_context(|| InvalidRecordBatchSnafu {
|
||||
reason: format!("JSON2 target column '{name}' has no extension metadata"),
|
||||
})?;
|
||||
json2_target_layouts.insert(
|
||||
column.column_id,
|
||||
Json2TargetLayout {
|
||||
extension_metadata,
|
||||
target_layout: plan.target_layout.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
let read_columns = ReadColumns::new(read_column_ids);
|
||||
let targets: Json2RewriteTargets = Arc::new(json2_target_layouts);
|
||||
let mapper = FlatProjectionMapper::new_with_json2_rewrite_targets(
|
||||
&self.metadata,
|
||||
projection,
|
||||
read_columns,
|
||||
&targets,
|
||||
)?;
|
||||
|
||||
let mut scan_input = ScanInput::new(self.sst_layer, mapper)
|
||||
.with_json2_rewrite_targets(targets)
|
||||
.with_files(self.inputs.to_vec())
|
||||
.with_compaction(true)
|
||||
.with_batch_size(batch_size)
|
||||
|
||||
@@ -12,20 +12,24 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::{assert_matches, fs};
|
||||
|
||||
use api::v1::Rows;
|
||||
use api::v1::region::{StrictWindow, compact_request};
|
||||
use common_function::utils::partition_expr_version;
|
||||
use common_recordbatch::RecordBatches;
|
||||
use datatypes::arrow::array::AsArray;
|
||||
use datatypes::arrow::datatypes::Float64Type;
|
||||
use datatypes::value::Value;
|
||||
use partition::expr::{PartitionExpr, col};
|
||||
use store_api::region_engine::{
|
||||
RegionEngine, RegionRole, RemapManifestsRequest, SettableRegionRoleState,
|
||||
};
|
||||
use store_api::region_request::{
|
||||
ApplyStagingManifestRequest, EnterStagingRequest, RegionFlushRequest, RegionPutRequest,
|
||||
RegionRequest, StagingPartitionDirective,
|
||||
ApplyStagingManifestRequest, EnterStagingRequest, RegionCompactRequest, RegionFlushRequest,
|
||||
RegionPutRequest, RegionRequest, StagingPartitionDirective,
|
||||
};
|
||||
use store_api::storage::{FileId, RegionId};
|
||||
|
||||
@@ -37,7 +41,9 @@ use crate::manifest::action::{
|
||||
};
|
||||
use crate::sst::FormatType;
|
||||
use crate::sst::file::FileMeta;
|
||||
use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema};
|
||||
use crate::test_util::{
|
||||
CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, put_rows, rows_schema,
|
||||
};
|
||||
|
||||
fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr {
|
||||
col(col_name)
|
||||
@@ -45,6 +51,209 @@ fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr {
|
||||
.and(col(col_name).lt(Value::Int64(end)))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_apply_staging_manifest_sequence_domain() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
test_apply_staging_manifest_sequence_domain_with_format(false).await;
|
||||
test_apply_staging_manifest_sequence_domain_with_format(true).await;
|
||||
}
|
||||
|
||||
async fn test_apply_staging_manifest_sequence_domain_with_format(flat_format: bool) {
|
||||
let mut env = TestEnv::with_prefix("apply-staging-sequence-domain").await;
|
||||
let engine = env
|
||||
.create_engine(MitoConfig {
|
||||
default_flat_format: flat_format,
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
let source = RegionId::new(1, 1);
|
||||
let target = RegionId::new(1, 2);
|
||||
let request = CreateRequestBuilder::new().build();
|
||||
let schema = rows_schema(&request);
|
||||
|
||||
engine
|
||||
.handle_request(source, RegionRequest::Create(request.clone()))
|
||||
.await
|
||||
.unwrap();
|
||||
for value in 0..3 {
|
||||
put_rows(
|
||||
&engine,
|
||||
source,
|
||||
Rows {
|
||||
schema: schema.clone(),
|
||||
rows: build_rows_for_key("0", 0, 1, value),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
engine
|
||||
.handle_request(source, RegionRequest::Flush(RegionFlushRequest::default()))
|
||||
.await
|
||||
.unwrap();
|
||||
let source_manifest = engine
|
||||
.get_region(source)
|
||||
.unwrap()
|
||||
.manifest_ctx
|
||||
.manifest()
|
||||
.await;
|
||||
assert_eq!(source_manifest.files.len(), 1);
|
||||
assert_eq!(
|
||||
source_manifest.files.values().next().unwrap().sequence,
|
||||
Some(std::num::NonZeroU64::new(3).unwrap())
|
||||
);
|
||||
|
||||
engine
|
||||
.set_region_role_state_gracefully(source, SettableRegionRoleState::StagingLeader)
|
||||
.await
|
||||
.unwrap();
|
||||
let partition_expr = float_range_expr("field_0", 0.1, 100.1)
|
||||
.as_json_str()
|
||||
.unwrap();
|
||||
let result = engine
|
||||
.remap_manifests(RemapManifestsRequest {
|
||||
region_id: source,
|
||||
input_regions: vec![source],
|
||||
region_mapping: [(source, vec![target])].into_iter().collect(),
|
||||
new_partition_exprs: [(target, partition_expr.clone())].into_iter().collect(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
engine
|
||||
.handle_request(target, RegionRequest::Create(request.clone()))
|
||||
.await
|
||||
.unwrap();
|
||||
engine
|
||||
.handle_request(
|
||||
target,
|
||||
RegionRequest::EnterStaging(EnterStagingRequest {
|
||||
partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
|
||||
partition_expr.clone(),
|
||||
),
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
engine
|
||||
.handle_request(
|
||||
target,
|
||||
RegionRequest::ApplyStagingManifest(ApplyStagingManifestRequest {
|
||||
partition_expr,
|
||||
central_region_id: source,
|
||||
manifest_path: result.manifest_paths[&target].clone(),
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let manifest = engine
|
||||
.get_region(target)
|
||||
.unwrap()
|
||||
.manifest_ctx
|
||||
.manifest()
|
||||
.await;
|
||||
assert_eq!(manifest.files.len(), 1);
|
||||
assert_eq!(manifest.committed_sequence, Some(1));
|
||||
let imported_file = manifest.files.values().next().unwrap();
|
||||
assert_eq!(imported_file.region_id, source);
|
||||
assert_eq!(
|
||||
imported_file.sequence,
|
||||
Some(std::num::NonZeroU64::new(1).unwrap())
|
||||
);
|
||||
|
||||
put_rows(
|
||||
&engine,
|
||||
target,
|
||||
Rows {
|
||||
schema: schema.clone(),
|
||||
rows: build_rows_for_key("0", 0, 1, 99),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
assert_target_value(&engine, target, 99.0).await;
|
||||
|
||||
engine
|
||||
.handle_request(target, RegionRequest::Flush(RegionFlushRequest::default()))
|
||||
.await
|
||||
.unwrap();
|
||||
let input_ids = current_file_ids(&engine, target);
|
||||
assert_eq!(
|
||||
input_ids.len(),
|
||||
2,
|
||||
"imported and target-write SSTs must both exist"
|
||||
);
|
||||
|
||||
engine
|
||||
.handle_request(
|
||||
target,
|
||||
RegionRequest::Compact(RegionCompactRequest {
|
||||
options: compact_request::Options::StrictWindow(StrictWindow {
|
||||
window_seconds: 60,
|
||||
}),
|
||||
parallelism: None,
|
||||
time_range: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let output_version = engine.get_region(target).unwrap().version();
|
||||
let output_files = output_version
|
||||
.ssts
|
||||
.levels()
|
||||
.iter()
|
||||
.flat_map(|level| level.files.values())
|
||||
.collect::<Vec<_>>();
|
||||
let output_ids = output_files
|
||||
.iter()
|
||||
.map(|file| file.meta_ref().file_id)
|
||||
.collect::<HashSet<_>>();
|
||||
assert_eq!(output_ids.len(), 1);
|
||||
assert!(
|
||||
output_files
|
||||
.iter()
|
||||
.all(|file| file.meta_ref().region_id == target)
|
||||
);
|
||||
assert!(
|
||||
input_ids.iter().all(|id| !output_ids.contains(id)),
|
||||
"real compaction must replace both input SSTs"
|
||||
);
|
||||
assert_target_value(&engine, target, 99.0).await;
|
||||
}
|
||||
|
||||
fn current_file_ids(engine: &crate::engine::MitoEngine, region_id: RegionId) -> HashSet<FileId> {
|
||||
engine
|
||||
.get_region(region_id)
|
||||
.unwrap()
|
||||
.version()
|
||||
.ssts
|
||||
.levels()
|
||||
.iter()
|
||||
.flat_map(|level| level.files.values())
|
||||
.map(|file| file.meta_ref().file_id)
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn assert_target_value(
|
||||
engine: &crate::engine::MitoEngine,
|
||||
region_id: RegionId,
|
||||
expected: f64,
|
||||
) {
|
||||
let scan = engine
|
||||
.scan_to_stream(region_id, ScanRequest::default())
|
||||
.await
|
||||
.unwrap();
|
||||
let batches = RecordBatches::try_collect(scan).await.unwrap();
|
||||
assert_eq!(
|
||||
batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
|
||||
1
|
||||
);
|
||||
let batch = batches.iter().next().unwrap();
|
||||
let values = batch
|
||||
.column_by_name("field_0")
|
||||
.unwrap()
|
||||
.as_primitive::<Float64Type>();
|
||||
assert_eq!(values.value(0), expected);
|
||||
}
|
||||
|
||||
fn float_range_expr(col_name: &str, start: f64, end: f64) -> PartitionExpr {
|
||||
col(col_name)
|
||||
.gt_eq(Value::Float64(start.into()))
|
||||
|
||||
@@ -15,9 +15,13 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use api::v1::Rows;
|
||||
use common_error::ext::{ErrorExt, WhateverResult};
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_recordbatch::RecordBatches;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
|
||||
use store_api::region_engine::RegionEngine;
|
||||
use store_api::region_request::{RegionCloseRequest, RegionRequest};
|
||||
use store_api::region_request::{PathType, RegionCloseRequest, RegionOpenRequest, RegionRequest};
|
||||
use store_api::storage::{RegionId, ScanRequest};
|
||||
|
||||
use crate::config::MitoConfig;
|
||||
@@ -32,6 +36,81 @@ async fn test_engine_create_new_region() {
|
||||
test_engine_create_new_region_with_format(true).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_engine_rejects_json2_with_time_series_memtable_on_create_and_open()
|
||||
-> WhateverResult<()> {
|
||||
let mut env = TestEnv::with_prefix("json2-rejects-time-series-memtable").await;
|
||||
let engine = env
|
||||
.create_engine(MitoConfig {
|
||||
default_flat_format: false,
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
let region_id = RegionId::new(1, 1);
|
||||
let request = CreateRequestBuilder::new()
|
||||
.field_datatype(ConcreteDataType::json2(JsonNativeType::Object(
|
||||
JsonObjectType::new(),
|
||||
)))
|
||||
.insert_option("append_mode", "true")
|
||||
.insert_option("memtable.type", "time_series")
|
||||
.build();
|
||||
|
||||
let err = engine
|
||||
.handle_request(region_id, RegionRequest::Create(request))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(StatusCode::InvalidArguments, err.status_code());
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("JSON2 columns only support BulkMemtable")
|
||||
);
|
||||
|
||||
let request = CreateRequestBuilder::new()
|
||||
.field_datatype(ConcreteDataType::json2(JsonNativeType::Object(
|
||||
JsonObjectType::new(),
|
||||
)))
|
||||
.insert_option("append_mode", "true")
|
||||
.insert_option("memtable.type", "bulk")
|
||||
.build();
|
||||
engine
|
||||
.handle_request(region_id, RegionRequest::Create(request))
|
||||
.await?;
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Close(RegionCloseRequest::default()),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let options = [
|
||||
("append_mode".to_string(), "true".to_string()),
|
||||
("memtable.type".to_string(), "time_series".to_string()),
|
||||
]
|
||||
.into_iter()
|
||||
.collect();
|
||||
let err = engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Open(RegionOpenRequest {
|
||||
engine: String::new(),
|
||||
table_dir: "test".to_string(),
|
||||
path_type: PathType::Bare,
|
||||
options,
|
||||
skip_wal_replay: false,
|
||||
checkpoint: None,
|
||||
requirements: Default::default(),
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(StatusCode::InvalidArguments, err.status_code());
|
||||
assert!(
|
||||
err.output_msg()
|
||||
.contains("JSON2 columns only support BulkMemtable")
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_engine_create_new_region_with_format(flat_format: bool) {
|
||||
let mut env = TestEnv::with_prefix("new-region").await;
|
||||
let engine = env
|
||||
|
||||
@@ -13,11 +13,13 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::helper::encode_json_value;
|
||||
use api::v1::Rows;
|
||||
use api::v1::helper::row;
|
||||
use api::v1::value::ValueData;
|
||||
use arrow_schema::extension::ExtensionType;
|
||||
use common_base::readable_size::ReadableSize;
|
||||
use common_error::ext::{ErrorExt, WhateverResult};
|
||||
use common_error::status_code::StatusCode;
|
||||
@@ -27,13 +29,16 @@ use datafusion_common::ScalarValue;
|
||||
use datafusion_expr::{col, lit};
|
||||
use datatypes::arrow::array::AsArray;
|
||||
use datatypes::arrow::datatypes::{Float64Type, TimestampMillisecondType};
|
||||
use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
|
||||
use datatypes::json::JsonSettings;
|
||||
use datatypes::json::value::JsonValue;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
|
||||
use datatypes::vectors::json::array::JsonArray;
|
||||
use futures::TryStreamExt;
|
||||
use serde_json::json;
|
||||
use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner};
|
||||
use store_api::region_request::RegionRequest;
|
||||
use store_api::region_request::{RegionCompactRequest, RegionRequest};
|
||||
use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution};
|
||||
|
||||
use crate::config::MitoConfig;
|
||||
@@ -41,7 +46,7 @@ use crate::error::Error;
|
||||
use crate::read::read_columns::ReadColumns;
|
||||
use crate::read::scan_region::Scanner;
|
||||
use crate::test_util;
|
||||
use crate::test_util::{CreateRequestBuilder, TestEnv};
|
||||
use crate::test_util::{CreateRequestBuilder, TestEnv, reopen_region};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResult<()> {
|
||||
@@ -73,28 +78,34 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul
|
||||
|
||||
// Write full JSON objects, then flush them so the scanner has an Parquet file where nested
|
||||
// projection can be pushed down.
|
||||
let rows = Rows {
|
||||
schema,
|
||||
rows: vec![
|
||||
row(vec![
|
||||
ValueData::StringValue("tag-1".to_string()),
|
||||
ValueData::JsonValue(encode_json_value(JsonValue::from(json!({
|
||||
"a": { "x": 10, "y": "ignored-a" },
|
||||
"b": "ignored-b"
|
||||
})))),
|
||||
ValueData::TimestampMillisecondValue(1000),
|
||||
]),
|
||||
row(vec![
|
||||
ValueData::StringValue("tag-2".to_string()),
|
||||
ValueData::JsonValue(encode_json_value(JsonValue::from(json!({
|
||||
"a": { "x": 20, "y": "ignored-c" },
|
||||
"b": "ignored-d"
|
||||
})))),
|
||||
ValueData::TimestampMillisecondValue(2000),
|
||||
]),
|
||||
for values in [
|
||||
vec![
|
||||
ValueData::StringValue("tag-1".to_string()),
|
||||
ValueData::JsonValue(encode_json_value(JsonValue::from(json!({
|
||||
"a": { "x": 10, "y": "ignored-a" },
|
||||
"b": "ignored-b"
|
||||
})))),
|
||||
ValueData::TimestampMillisecondValue(1000),
|
||||
],
|
||||
};
|
||||
test_util::put_rows(&engine, region_id, rows).await;
|
||||
vec![
|
||||
ValueData::StringValue("tag-2".to_string()),
|
||||
ValueData::JsonValue(encode_json_value(JsonValue::from(json!({
|
||||
"a": { "x": 20, "y": "ignored-c" },
|
||||
"b": "ignored-d"
|
||||
})))),
|
||||
ValueData::TimestampMillisecondValue(2000),
|
||||
],
|
||||
] {
|
||||
test_util::put_rows(
|
||||
&engine,
|
||||
region_id,
|
||||
Rows {
|
||||
schema: schema.clone(),
|
||||
rows: vec![row(values)],
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
test_util::flush_region(&engine, region_id, None).await;
|
||||
|
||||
// Without a type hint, the scanner reads the whole JSON2 root column.
|
||||
@@ -208,6 +219,231 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json2_v1_region_reopen_and_compaction() -> WhateverResult<()> {
|
||||
let mut request = CreateRequestBuilder::new()
|
||||
.field_datatype(ConcreteDataType::json2(JsonNativeType::Object(
|
||||
JsonObjectType::new(),
|
||||
)))
|
||||
.insert_option("memtable.type", "bulk")
|
||||
.build();
|
||||
let settings = JsonSettings::default();
|
||||
request.column_metadatas[1]
|
||||
.column_schema
|
||||
.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1(
|
||||
settings,
|
||||
))));
|
||||
let table_dir = request.table_dir.clone();
|
||||
let schema = test_util::rows_schema(&request);
|
||||
let mut env = TestEnv::new().await;
|
||||
let engine = env.create_engine(MitoConfig::default()).await;
|
||||
let region_id = RegionId::new(1024, 0);
|
||||
engine
|
||||
.handle_request(region_id, RegionRequest::Create(request))
|
||||
.await?;
|
||||
|
||||
let values = [
|
||||
json!({"route": 1}),
|
||||
json!({"b": {"c": "x"}}),
|
||||
json!({"route": 3, "written_after_reopen": true}),
|
||||
];
|
||||
for (i, value) in values[..2].iter().enumerate() {
|
||||
test_util::put_rows(
|
||||
&engine,
|
||||
region_id,
|
||||
Rows {
|
||||
schema: schema.clone(),
|
||||
rows: vec![row(vec![
|
||||
ValueData::StringValue("tag".to_string()),
|
||||
ValueData::JsonValue(encode_json_value(JsonValue::from(value.clone()))),
|
||||
ValueData::TimestampMillisecondValue((i as i64 + 1) * 1000),
|
||||
])],
|
||||
},
|
||||
)
|
||||
.await;
|
||||
test_util::flush_region(&engine, region_id, None).await;
|
||||
}
|
||||
|
||||
reopen_region(
|
||||
&engine,
|
||||
region_id,
|
||||
table_dir,
|
||||
true,
|
||||
HashMap::from([("memtable.type".to_string(), "bulk".to_string())]),
|
||||
)
|
||||
.await;
|
||||
let region = engine.get_region(region_id).unwrap();
|
||||
let version = region.version();
|
||||
let column = &version
|
||||
.metadata
|
||||
.column_by_name("field_0")
|
||||
.unwrap()
|
||||
.column_schema;
|
||||
let extension = column.extension_type::<Json2ExtensionType>()?.unwrap();
|
||||
assert!(extension.metadata().is_version_2());
|
||||
|
||||
test_util::put_rows(
|
||||
&engine,
|
||||
region_id,
|
||||
Rows {
|
||||
schema: schema.clone(),
|
||||
rows: vec![row(vec![
|
||||
ValueData::StringValue("tag".to_string()),
|
||||
ValueData::JsonValue(encode_json_value(JsonValue::from(values[2].clone()))),
|
||||
ValueData::TimestampMillisecondValue(3000),
|
||||
])],
|
||||
},
|
||||
)
|
||||
.await;
|
||||
test_util::flush_region(&engine, region_id, None).await;
|
||||
|
||||
let region = engine.get_region(region_id).unwrap();
|
||||
let input_files = region
|
||||
.version()
|
||||
.ssts
|
||||
.levels()
|
||||
.iter()
|
||||
.map(|level| level.files.len())
|
||||
.sum::<usize>();
|
||||
assert_eq!(3, input_files);
|
||||
|
||||
// The same requested path is stored as a v1 explicit leaf in the first SST, is absent from
|
||||
// the second SST, and lives in the v2 remainder in the third SST. A per-file route preserves
|
||||
// those differences while exposing one logical query type to the merge reader.
|
||||
let scanner = engine
|
||||
.scanner(
|
||||
region_id,
|
||||
ScanRequest {
|
||||
projection: Some(vec![1]),
|
||||
json_type_hint: HashMap::from([(
|
||||
"field_0".to_string(),
|
||||
JsonNativeType::Object(JsonObjectType::from([(
|
||||
"route".to_string(),
|
||||
JsonNativeType::i64(),
|
||||
)])),
|
||||
)]),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let batches = RecordBatches::try_collect(scanner.scan().await?).await?;
|
||||
let mut routed = Vec::new();
|
||||
for batch in batches.iter() {
|
||||
let array = batch.column_by_name("field_0").unwrap().clone();
|
||||
let json = JsonArray::from(&array);
|
||||
for i in 0..array.len() {
|
||||
routed.push(json.try_get_value(i)?);
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
[json!({"route": 1}), json!(null), json!({"route": 3})],
|
||||
routed.as_slice()
|
||||
);
|
||||
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Compact(RegionCompactRequest::default()),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let scanner = engine
|
||||
.scanner(
|
||||
region_id,
|
||||
ScanRequest {
|
||||
projection: Some(vec![1]),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let batches = RecordBatches::try_collect(scanner.scan().await?).await?;
|
||||
let mut actual = Vec::new();
|
||||
for batch in batches.iter() {
|
||||
let array = batch.column_by_name("field_0").unwrap().clone();
|
||||
let json = JsonArray::from(&array);
|
||||
for i in 0..array.len() {
|
||||
actual.push(json.try_get_value(i)?);
|
||||
}
|
||||
}
|
||||
assert_eq!(values, actual.as_slice());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_flush_aligns_different_json2_layouts() -> WhateverResult<()> {
|
||||
let mut request = CreateRequestBuilder::new()
|
||||
.field_datatype(ConcreteDataType::json2(JsonNativeType::Object(
|
||||
JsonObjectType::new(),
|
||||
)))
|
||||
.insert_option("append_mode", "true")
|
||||
.insert_option("memtable.type", "bulk")
|
||||
.build();
|
||||
let settings = JsonSettings::try_new(vec![], Some(1))?;
|
||||
request.column_metadatas[1]
|
||||
.column_schema
|
||||
.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new(
|
||||
settings,
|
||||
))));
|
||||
let schema = test_util::rows_schema(&request);
|
||||
let mut env = TestEnv::new().await;
|
||||
let engine = env.create_engine(MitoConfig::default()).await;
|
||||
let region_id = RegionId::new(1025, 0);
|
||||
engine
|
||||
.handle_request(region_id, RegionRequest::Create(request))
|
||||
.await?;
|
||||
|
||||
for (offset, name) in [(0, "a"), (1024, "b")] {
|
||||
let rows = (0..1024)
|
||||
.map(|i| {
|
||||
row(vec![
|
||||
ValueData::StringValue("tag".to_string()),
|
||||
ValueData::JsonValue(encode_json_value(JsonValue::from(json!({(name): i})))),
|
||||
ValueData::TimestampMillisecondValue(offset + i),
|
||||
])
|
||||
})
|
||||
.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 = engine
|
||||
.scanner(
|
||||
region_id,
|
||||
ScanRequest {
|
||||
projection: Some(vec![1]),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let batches = RecordBatches::try_collect(scanner.scan().await?).await?;
|
||||
let mut counts = HashMap::new();
|
||||
for batch in batches.iter() {
|
||||
let array = batch.column_by_name("field_0").unwrap().clone();
|
||||
let json = JsonArray::from(&array);
|
||||
for i in 0..array.len() {
|
||||
let value = json.try_get_value(i)?;
|
||||
let name = match (value.get("a"), value.get("b")) {
|
||||
(Some(_), None) => "a",
|
||||
(None, Some(_)) => "b",
|
||||
_ => panic!("expected exactly one dynamic JSON2 field, got {value}"),
|
||||
};
|
||||
*counts.entry(name).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
assert_eq!(Some(&1024), counts.get("a"));
|
||||
assert_eq!(Some(&1024), counts.get("b"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_incremental_query_stale_error() {
|
||||
let mut env = TestEnv::with_prefix("test_incremental_query_stale_error").await;
|
||||
|
||||
@@ -12,6 +12,8 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use api::v1::Rows;
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
@@ -20,12 +22,16 @@ use store_api::region_engine::{
|
||||
SettableRegionRoleState,
|
||||
};
|
||||
use store_api::region_request::{
|
||||
EnterStagingRequest, RegionPutRequest, RegionRequest, StagingPartitionDirective,
|
||||
EnterStagingRequest, RegionFlushRequest, RegionPutRequest, RegionRequest,
|
||||
StagingPartitionDirective,
|
||||
};
|
||||
use store_api::storage::RegionId;
|
||||
|
||||
use crate::config::MitoConfig;
|
||||
use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema};
|
||||
use crate::region::{RegionLeaderState, RegionRoleState};
|
||||
use crate::test_util::{
|
||||
CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema,
|
||||
};
|
||||
|
||||
/// Helper function to assert a successful response with expected entry id
|
||||
fn assert_success_response(response: &SetRegionRoleStateResponse, expected_entry_id: u64) {
|
||||
@@ -212,6 +218,324 @@ async fn test_write_downgrading_region_with_format(flat_format: bool) {
|
||||
assert_eq!(err.status_code(), StatusCode::RegionNotReady)
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_downgrading_waits_for_checkpoint_and_stops_new_checkpoints() {
|
||||
let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup();
|
||||
let mut env = TestEnv::new().await.with_mock_layer(mock_layer);
|
||||
let engine = env
|
||||
.create_engine(MitoConfig {
|
||||
manifest_checkpoint_distance: 1,
|
||||
..Default::default()
|
||||
})
|
||||
.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.clone(),
|
||||
rows: build_rows(0, 1),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest::default()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("checkpoint cleanup did not start");
|
||||
|
||||
// Leave one memtable for the final flush after entering Downgrading.
|
||||
put_rows(
|
||||
&engine,
|
||||
region_id,
|
||||
Rows {
|
||||
schema: column_schemas.clone(),
|
||||
rows: build_rows(1, 2),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let region = engine.get_region(region_id).unwrap();
|
||||
let wait_started = region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.pending_checkpoint_wait_started();
|
||||
// Register before starting the transition so the notification cannot be lost.
|
||||
let wait_started = wait_started.notified();
|
||||
let cloned_engine = engine.clone();
|
||||
let downgrade = tokio::spawn(async move {
|
||||
cloned_engine
|
||||
.set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader)
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(5), wait_started)
|
||||
.await
|
||||
.expect("downgrade did not start waiting for the checkpoint");
|
||||
assert_eq!(
|
||||
RegionRoleState::Leader(RegionLeaderState::Downgrading),
|
||||
region.state()
|
||||
);
|
||||
assert!(
|
||||
!downgrade.is_finished(),
|
||||
"downgrade returned before checkpoint cleanup finished"
|
||||
);
|
||||
|
||||
blocker.release();
|
||||
assert_success_response(&downgrade.await.unwrap().unwrap(), 2);
|
||||
|
||||
// Final flush publishes a normal delta, but Downgrading must not start a
|
||||
// checkpoint after the barrier.
|
||||
blocker.arm_next_close();
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest::default()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
!region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.is_doing_checkpoint(),
|
||||
"final flush started a checkpoint while Downgrading"
|
||||
);
|
||||
assert_eq!(2, region.manifest_ctx.manifest().await.manifest_version);
|
||||
assert_eq!(
|
||||
1,
|
||||
region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.last_checkpoint_version()
|
||||
);
|
||||
|
||||
// Leaving Downgrading removes the scheduling restriction. The next normal
|
||||
// manifest update can checkpoint the accumulated deltas.
|
||||
engine
|
||||
.set_region_role(region_id, RegionRole::Leader)
|
||||
.unwrap();
|
||||
put_rows(
|
||||
&engine,
|
||||
region_id,
|
||||
Rows {
|
||||
schema: column_schemas,
|
||||
rows: build_rows(2, 3),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest::default()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("checkpoint scheduling did not resume after leaving Downgrading");
|
||||
blocker.release();
|
||||
region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.write()
|
||||
.await
|
||||
.wait_for_pending_checkpoint()
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_direct_follower_waits_for_pending_checkpoint() {
|
||||
let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup();
|
||||
let mut env = TestEnv::new().await.with_mock_layer(mock_layer);
|
||||
let engine = env
|
||||
.create_engine(MitoConfig {
|
||||
manifest_checkpoint_distance: 1,
|
||||
..Default::default()
|
||||
})
|
||||
.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, 1),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest::default()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("checkpoint cleanup did not start");
|
||||
|
||||
let region = engine.get_region(region_id).unwrap();
|
||||
let wait_started = region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.pending_checkpoint_wait_started();
|
||||
// The no-flush downgrade path requests Follower directly. It must still
|
||||
// wait for checkpoint cleanup before replying to the caller.
|
||||
let wait_started = wait_started.notified();
|
||||
let cloned_engine = engine.clone();
|
||||
let set_follower = tokio::spawn(async move {
|
||||
cloned_engine
|
||||
.set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower)
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(5), wait_started)
|
||||
.await
|
||||
.expect("follower transition did not start waiting for the checkpoint");
|
||||
assert_eq!(RegionRoleState::Follower, region.state());
|
||||
assert!(
|
||||
!set_follower.is_finished(),
|
||||
"follower transition returned before checkpoint cleanup finished"
|
||||
);
|
||||
|
||||
blocker.release();
|
||||
assert_success_response(&set_follower.await.unwrap().unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_retried_downgrade_waits_after_first_request_is_cancelled() {
|
||||
let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup();
|
||||
let mut env = TestEnv::new().await.with_mock_layer(mock_layer);
|
||||
let engine = env
|
||||
.create_engine(MitoConfig {
|
||||
manifest_checkpoint_distance: 1,
|
||||
..Default::default()
|
||||
})
|
||||
.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, 1),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest::default()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("checkpoint cleanup did not start");
|
||||
|
||||
let region = engine.get_region(region_id).unwrap();
|
||||
let first_wait_started = region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.pending_checkpoint_wait_started();
|
||||
// Register before starting the transition so the notification cannot be lost.
|
||||
let first_wait_started = first_wait_started.notified();
|
||||
// Call the region transition directly so aborting this task cancels the
|
||||
// actual checkpoint waiter. The engine API submits the same transition to
|
||||
// a detached worker task, so aborting its caller only drops the reply receiver.
|
||||
let first_region = region.clone();
|
||||
let first_downgrade = tokio::spawn(async move {
|
||||
first_region
|
||||
.set_role_state_gracefully(SettableRegionRoleState::DowngradingLeader)
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(5), first_wait_started)
|
||||
.await
|
||||
.expect("first downgrade did not start waiting for the checkpoint");
|
||||
assert_eq!(
|
||||
RegionRoleState::Leader(RegionLeaderState::Downgrading),
|
||||
region.state()
|
||||
);
|
||||
assert!(
|
||||
!first_downgrade.is_finished(),
|
||||
"first downgrade did not wait for checkpoint cleanup"
|
||||
);
|
||||
|
||||
// Cancelling the first caller must not remove the checkpoint handle from
|
||||
// the manifest manager. A retried migration request observes Downgrading
|
||||
// and must wait for the same checkpoint task before it can proceed.
|
||||
first_downgrade.abort();
|
||||
assert!(first_downgrade.await.unwrap_err().is_cancelled());
|
||||
assert_eq!(
|
||||
RegionRoleState::Leader(RegionLeaderState::Downgrading),
|
||||
region.state()
|
||||
);
|
||||
|
||||
let second_wait_started = region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.pending_checkpoint_wait_started();
|
||||
// Register before spawning the retry so the notification cannot be lost.
|
||||
let second_wait_started = second_wait_started.notified();
|
||||
let retry_region = region.clone();
|
||||
let retry_downgrade = tokio::spawn(async move {
|
||||
retry_region
|
||||
.set_role_state_gracefully(SettableRegionRoleState::DowngradingLeader)
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(5), second_wait_started)
|
||||
.await
|
||||
.expect("retried downgrade did not start waiting for the checkpoint");
|
||||
assert!(
|
||||
!retry_downgrade.is_finished(),
|
||||
"retried downgrade did not wait for checkpoint cleanup"
|
||||
);
|
||||
|
||||
blocker.release();
|
||||
retry_downgrade.await.unwrap().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unified_state_transitions() {
|
||||
test_unified_state_transitions_with_format(false).await;
|
||||
|
||||
@@ -48,7 +48,9 @@ use crate::manifest::action::{
|
||||
use crate::region::{RegionLeaderState, RegionRoleState, parse_partition_expr};
|
||||
use crate::request::WorkerRequest;
|
||||
use crate::sst::FormatType;
|
||||
use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema};
|
||||
use crate::test_util::{
|
||||
CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema,
|
||||
};
|
||||
|
||||
fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr {
|
||||
col(col_name)
|
||||
@@ -1101,6 +1103,101 @@ async fn test_write_stall_on_enter_staging_with_format(flat_format: bool) {
|
||||
assert_eq!(expected, batches.pretty_print().unwrap());
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_enter_staging_waits_for_pending_checkpoint() {
|
||||
for block_last_checkpoint_write in [false, true] {
|
||||
let partition_directive =
|
||||
StagingPartitionDirective::UpdatePartitionExpr(default_partition_expr());
|
||||
let (blocker, mock_layer) = if block_last_checkpoint_write {
|
||||
CheckpointTaskBlocker::block_last_checkpoint_write()
|
||||
} else {
|
||||
CheckpointTaskBlocker::block_cleanup()
|
||||
};
|
||||
let mut env = TestEnv::new().await.with_mock_layer(mock_layer);
|
||||
let engine = env
|
||||
.create_engine(MitoConfig {
|
||||
manifest_checkpoint_distance: 1,
|
||||
..Default::default()
|
||||
})
|
||||
.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, 1),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest::default()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("checkpoint task did not reach the configured block point");
|
||||
|
||||
let region = engine.get_region(region_id).unwrap();
|
||||
let wait_started = region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.pending_checkpoint_wait_started();
|
||||
// Register before starting the transition so the notification cannot be lost.
|
||||
let wait_started = wait_started.notified();
|
||||
let cloned_engine = engine.clone();
|
||||
let enter_staging = tokio::spawn(async move {
|
||||
cloned_engine
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::EnterStaging(EnterStagingRequest {
|
||||
partition_directive,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(5), wait_started)
|
||||
.await
|
||||
.expect("EnterStaging did not start waiting for the checkpoint");
|
||||
assert_eq!(
|
||||
RegionRoleState::Leader(RegionLeaderState::EnteringStaging),
|
||||
region.state()
|
||||
);
|
||||
assert!(
|
||||
!enter_staging.is_finished(),
|
||||
"EnterStaging returned before the checkpoint task finished"
|
||||
);
|
||||
|
||||
blocker.release();
|
||||
enter_staging.await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
RegionRoleState::Leader(RegionLeaderState::Staging),
|
||||
region.state()
|
||||
);
|
||||
assert!(
|
||||
!region
|
||||
.manifest_ctx
|
||||
.manifest_manager
|
||||
.read()
|
||||
.await
|
||||
.checkpointer()
|
||||
.is_doing_checkpoint()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_enter_staging_clean_staging_manifest_error() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
+20
-2
@@ -67,6 +67,20 @@ pub enum Error {
|
||||
error: object_store::Error,
|
||||
},
|
||||
|
||||
#[snafu(display(
|
||||
"Manifest delta {} disappeared after it was listed, path: {}",
|
||||
version,
|
||||
path
|
||||
))]
|
||||
ManifestDeltaNotFound {
|
||||
version: ManifestVersion,
|
||||
path: String,
|
||||
#[snafu(source)]
|
||||
error: object_store::Error,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Fail to compress object by {}, path: {}", compress_type, path))]
|
||||
CompressObject {
|
||||
compress_type: CompressionType,
|
||||
@@ -1356,6 +1370,7 @@ impl Error {
|
||||
pub(crate) fn is_object_not_found(&self) -> bool {
|
||||
match self {
|
||||
Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound,
|
||||
Error::ManifestDeltaNotFound { .. } => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
@@ -1397,7 +1412,9 @@ impl ErrorExt for Error {
|
||||
|
||||
match self {
|
||||
DataTypeMismatch { source, .. } => source.status_code(),
|
||||
OpenDal { .. } | ReadParquet { .. } => StatusCode::StorageUnavailable,
|
||||
OpenDal { .. } | ManifestDeltaNotFound { .. } | ReadParquet { .. } => {
|
||||
StatusCode::StorageUnavailable
|
||||
}
|
||||
WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } => {
|
||||
source.status_code()
|
||||
}
|
||||
@@ -1592,7 +1609,8 @@ impl ErrorExt for Error {
|
||||
| UpdateManifest { .. }
|
||||
| RegionStopped { .. }
|
||||
| RegionBusy { .. }
|
||||
| FlushableRegionState { .. } => RetryHint::Retryable,
|
||||
| FlushableRegionState { .. }
|
||||
| ManifestDeltaNotFound { .. } => RetryHint::Retryable,
|
||||
|
||||
OpenDal { error, .. }
|
||||
| DeleteSsts { error, .. }
|
||||
|
||||
+171
-55
@@ -24,10 +24,10 @@ use bytes::Bytes;
|
||||
use common_base::cancellation::CancellableFuture;
|
||||
use common_telemetry::{debug, error, info};
|
||||
use datatypes::arrow::datatypes::SchemaRef;
|
||||
use datatypes::extension::json::is_json2_extension_type;
|
||||
use partition::expr::PartitionExpr;
|
||||
use smallvec::{SmallVec, smallvec};
|
||||
use snafu::ResultExt;
|
||||
use snafu::{ResultExt, ensure};
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
use store_api::region_request::RegionFlushReason;
|
||||
use store_api::storage::{RegionId, SequenceNumber};
|
||||
use strum::IntoStaticStr;
|
||||
@@ -37,15 +37,15 @@ use crate::access_layer::{
|
||||
AccessLayerRef, Metrics, OperationType, SstInfoArray, SstWriteRequest, WriteType,
|
||||
};
|
||||
use crate::cache::CacheManagerRef;
|
||||
use crate::compaction::{collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema};
|
||||
use crate::config::MitoConfig;
|
||||
use crate::engine::region_hook::SstFileInfo;
|
||||
use crate::error::{
|
||||
Error, FlushCancelledSnafu, FlushRegionSnafu, JoinSnafu, RegionBusySnafu, RegionClosedSnafu,
|
||||
RegionDroppedSnafu, RegionTruncatedSnafu, Result,
|
||||
RegionDroppedSnafu, RegionTruncatedSnafu, Result, UnexpectedSnafu,
|
||||
};
|
||||
use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList};
|
||||
use crate::memtable::bulk::ENCODE_ROW_THRESHOLD;
|
||||
use crate::memtable::bulk::json_align::Json2Aligner;
|
||||
use crate::memtable::{BoxedRecordBatchIterator, EncodedRange, MemtableRanges, RangesOptions};
|
||||
use crate::metrics::{
|
||||
FLUSH_BYTES_TOTAL, FLUSH_ELAPSED, FLUSH_FAILURE_TOTAL, FLUSH_FILE_TOTAL, FLUSH_REQUESTS_TOTAL,
|
||||
@@ -708,6 +708,7 @@ impl RegionFlushTask {
|
||||
let flat_sources = memtable_flat_sources(
|
||||
batch_schema,
|
||||
mem_ranges,
|
||||
&version.metadata,
|
||||
&version.options,
|
||||
field_column_start,
|
||||
)?;
|
||||
@@ -901,6 +902,7 @@ struct FlatSources {
|
||||
fn memtable_flat_sources(
|
||||
schema: SchemaRef,
|
||||
mem_ranges: MemtableRanges,
|
||||
metadata: &RegionMetadataRef,
|
||||
options: &RegionOptions,
|
||||
field_column_start: usize,
|
||||
) -> Result<FlatSources> {
|
||||
@@ -918,6 +920,7 @@ fn memtable_flat_sources(
|
||||
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.
|
||||
@@ -951,12 +954,23 @@ fn memtable_flat_sources(
|
||||
let mut input_iters = Vec::with_capacity(num_ranges);
|
||||
let mut current_ranges = Vec::new();
|
||||
|
||||
let has_json2 = schema.fields().iter().any(is_json2_extension_type);
|
||||
let mut json_align_schemas = if has_json2 {
|
||||
Some(Vec::with_capacity(num_ranges))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
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::<Vec<_>>();
|
||||
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() {
|
||||
@@ -965,15 +979,26 @@ fn memtable_flat_sources(
|
||||
continue;
|
||||
}
|
||||
|
||||
// Collect schemas if has json2 field.
|
||||
if let Some(schemas) = json_align_schemas.as_mut() {
|
||||
let schema = range
|
||||
.record_batch_schema_hint()
|
||||
.unwrap_or_else(|| schema.clone());
|
||||
schemas.push(schema);
|
||||
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;
|
||||
@@ -1007,11 +1032,6 @@ fn memtable_flat_sources(
|
||||
|
||||
let input_iters =
|
||||
std::mem::replace(&mut input_iters, Vec::with_capacity(num_ranges));
|
||||
let (schema, input_iters) = maybe_align_json2_iters(
|
||||
schema.clone(),
|
||||
json_align_schemas.take(),
|
||||
input_iters,
|
||||
)?;
|
||||
|
||||
let maybe_dedup = merge_and_dedup_with_batch_size(
|
||||
&schema,
|
||||
@@ -1022,17 +1042,12 @@ fn memtable_flat_sources(
|
||||
batch_size,
|
||||
)?;
|
||||
|
||||
flat_sources
|
||||
.sources
|
||||
.push((FlatSource::new_iter(schema, maybe_dedup), max_sequence));
|
||||
flat_sources.sources.push((
|
||||
FlatSource::new_iter(schema.clone(), maybe_dedup),
|
||||
max_sequence,
|
||||
));
|
||||
last_iter_rows = 0;
|
||||
current_ranges.clear();
|
||||
|
||||
json_align_schemas = if has_json2 {
|
||||
Some(Vec::with_capacity(num_ranges))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1046,9 +1061,6 @@ fn memtable_flat_sources(
|
||||
rows_remaining
|
||||
);
|
||||
|
||||
let (schema, input_iters) =
|
||||
maybe_align_json2_iters(schema, json_align_schemas, input_iters)?;
|
||||
|
||||
let max_sequence = current_ranges
|
||||
.iter()
|
||||
.map(|r| r.stats().max_sequence())
|
||||
@@ -1078,24 +1090,6 @@ fn memtable_flat_sources(
|
||||
Ok(flat_sources)
|
||||
}
|
||||
|
||||
fn maybe_align_json2_iters(
|
||||
schema: SchemaRef,
|
||||
schemas: Option<Vec<SchemaRef>>,
|
||||
input_iters: Vec<BoxedRecordBatchIterator>,
|
||||
) -> Result<(SchemaRef, Vec<BoxedRecordBatchIterator>)> {
|
||||
let Some(schemas) = schemas else {
|
||||
return Ok((schema, input_iters));
|
||||
};
|
||||
|
||||
let aligner = Json2Aligner::try_new(schemas)?;
|
||||
let input_iters = input_iters
|
||||
.into_iter()
|
||||
.map(|input_iter| aligner.wrap_iter(input_iter))
|
||||
.collect();
|
||||
|
||||
Ok((aligner.schema().clone(), input_iters))
|
||||
}
|
||||
|
||||
/// 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
|
||||
@@ -1646,6 +1640,8 @@ mod tests {
|
||||
use api::v1::{OpType, Rows};
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
use datatypes::arrow::datatypes::Schema;
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use mito_codec::row_converter::build_primary_key_codec;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
@@ -1654,7 +1650,9 @@ mod tests {
|
||||
use crate::error::InvalidSchedulerStateSnafu;
|
||||
use crate::memtable::bulk::part::BulkPartConverter;
|
||||
use crate::memtable::time_series::TimeSeriesMemtableBuilder;
|
||||
use crate::memtable::{Memtable, RangesOptions};
|
||||
use crate::memtable::{
|
||||
IterBuilder, Memtable, MemtableRange, MemtableRangeContext, MemtableStats, RangesOptions,
|
||||
};
|
||||
use crate::request::WriteRequest;
|
||||
use crate::schedule::scheduler::Scheduler;
|
||||
use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema};
|
||||
@@ -2243,6 +2241,7 @@ mod tests {
|
||||
let flat_sources = memtable_flat_sources(
|
||||
schema.clone(),
|
||||
mem_ranges,
|
||||
&metadata,
|
||||
&options,
|
||||
metadata.primary_key.len(),
|
||||
)
|
||||
@@ -2271,9 +2270,14 @@ mod tests {
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let flat_sources =
|
||||
memtable_flat_sources(schema, mem_ranges, &options, metadata.primary_key.len())
|
||||
.unwrap();
|
||||
let flat_sources = memtable_flat_sources(
|
||||
schema,
|
||||
mem_ranges,
|
||||
&metadata,
|
||||
&options,
|
||||
metadata.primary_key.len(),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(flat_sources.encoded.is_empty());
|
||||
assert_eq!(1, flat_sources.sources.len());
|
||||
|
||||
@@ -2288,6 +2292,118 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memtable_flat_sources_uses_non_encoded_schema() -> Result<()> {
|
||||
struct TestIterBuilder {
|
||||
schema: SchemaRef,
|
||||
batch: Option<RecordBatch>,
|
||||
}
|
||||
|
||||
impl IterBuilder for TestIterBuilder {
|
||||
fn build(
|
||||
&self,
|
||||
_metrics: Option<crate::memtable::MemScanMetrics>,
|
||||
) -> Result<crate::memtable::BoxedBatchIterator> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn is_record_batch(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn build_record_batch(
|
||||
&self,
|
||||
_time_range: Option<(common_time::Timestamp, common_time::Timestamp)>,
|
||||
_metrics: Option<crate::memtable::MemScanMetrics>,
|
||||
) -> Result<BoxedRecordBatchIterator> {
|
||||
let Some(batch) = self.batch.clone() else {
|
||||
unimplemented!()
|
||||
};
|
||||
Ok(Box::new(std::iter::once(Ok(batch))))
|
||||
}
|
||||
|
||||
fn record_batch_schema_hint(&self) -> Option<SchemaRef> {
|
||||
Some(self.schema.clone())
|
||||
}
|
||||
|
||||
fn encoded_range(&self) -> Option<EncodedRange> {
|
||||
self.batch.is_none().then(|| EncodedRange {
|
||||
data: Bytes::new(),
|
||||
sst_info: SstInfo::default(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
let metadata = metadata_for_test();
|
||||
let schema = to_flat_sst_arrow_schema(
|
||||
&metadata,
|
||||
&FlatSchemaOptions::from_encoding(metadata.primary_key_encoding),
|
||||
);
|
||||
let pk_codec = build_primary_key_codec(&metadata);
|
||||
let mut converter = BulkPartConverter::new(&metadata, schema.clone(), 1, pk_codec, true);
|
||||
let kvs = build_key_values_with_ts_seq_values(
|
||||
&metadata,
|
||||
"key".to_string(),
|
||||
1,
|
||||
std::iter::once(1000),
|
||||
std::iter::once(Some(1.0)),
|
||||
1,
|
||||
);
|
||||
converter.append_key_values(&kvs)?;
|
||||
let batch = converter.convert()?.batch;
|
||||
let encoded_schema = Arc::new(Schema::empty());
|
||||
|
||||
let new_range = |id, builder| {
|
||||
MemtableRange::new(
|
||||
Arc::new(MemtableRangeContext::new(
|
||||
id,
|
||||
Box::new(builder),
|
||||
Default::default(),
|
||||
)),
|
||||
MemtableStats {
|
||||
num_rows: 1,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
};
|
||||
let mut ranges = std::collections::BTreeMap::new();
|
||||
ranges.insert(
|
||||
0,
|
||||
new_range(
|
||||
0,
|
||||
TestIterBuilder {
|
||||
schema: encoded_schema.clone(),
|
||||
batch: None,
|
||||
},
|
||||
),
|
||||
);
|
||||
ranges.insert(
|
||||
1,
|
||||
new_range(
|
||||
0,
|
||||
TestIterBuilder {
|
||||
schema: schema.clone(),
|
||||
batch: Some(batch),
|
||||
},
|
||||
),
|
||||
);
|
||||
|
||||
let sources = memtable_flat_sources(
|
||||
encoded_schema,
|
||||
MemtableRanges { ranges },
|
||||
&metadata,
|
||||
&RegionOptions {
|
||||
append_mode: true,
|
||||
..Default::default()
|
||||
},
|
||||
metadata.primary_key.len(),
|
||||
)?;
|
||||
assert_eq!(1, sources.encoded.len());
|
||||
assert_eq!(1, sources.sources.len());
|
||||
assert_eq!(&schema, sources.sources[0].0.schema());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_schedule_pending_request_on_flush_success() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
@@ -14,11 +14,14 @@
|
||||
|
||||
use std::fmt::Debug;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use common_runtime::JoinHandle;
|
||||
use common_telemetry::{error, info, warn};
|
||||
use store_api::storage::RegionId;
|
||||
use store_api::{MIN_VERSION, ManifestVersion};
|
||||
#[cfg(test)]
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::manifest::action::{RegionCheckpoint, RegionManifest};
|
||||
@@ -31,6 +34,9 @@ use crate::metrics::MANIFEST_OP_ELAPSED;
|
||||
pub(crate) struct Checkpointer {
|
||||
manifest_options: RegionManifestOptions,
|
||||
inner: Arc<Inner>,
|
||||
checkpoint_task: Option<JoinHandle<()>>,
|
||||
#[cfg(test)]
|
||||
pending_checkpoint_wait_started: Arc<Notify>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
@@ -38,15 +44,10 @@ struct Inner {
|
||||
region_id: RegionId,
|
||||
manifest_store: ManifestObjectStore,
|
||||
last_checkpoint_version: AtomicU64,
|
||||
is_doing_checkpoint: AtomicBool,
|
||||
}
|
||||
|
||||
impl Inner {
|
||||
async fn do_checkpoint(&self, checkpoint: RegionCheckpoint) {
|
||||
let _guard = scopeguard::guard(&self.is_doing_checkpoint, |x| {
|
||||
x.store(false, Ordering::Relaxed);
|
||||
});
|
||||
|
||||
let _t = MANIFEST_OP_ELAPSED
|
||||
.with_label_values(&["checkpoint"])
|
||||
.start_timer();
|
||||
@@ -90,14 +91,6 @@ impl Inner {
|
||||
fn region_id(&self) -> RegionId {
|
||||
self.region_id
|
||||
}
|
||||
|
||||
fn is_doing_checkpoint(&self) -> bool {
|
||||
self.is_doing_checkpoint.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
fn set_doing_checkpoint(&self) {
|
||||
self.is_doing_checkpoint.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
impl Checkpointer {
|
||||
@@ -113,8 +106,10 @@ impl Checkpointer {
|
||||
region_id,
|
||||
manifest_store,
|
||||
last_checkpoint_version: AtomicU64::new(last_checkpoint_version),
|
||||
is_doing_checkpoint: AtomicBool::new(false),
|
||||
}),
|
||||
checkpoint_task: None,
|
||||
#[cfg(test)]
|
||||
pending_checkpoint_wait_started: Arc::new(Notify::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -140,7 +135,19 @@ impl Checkpointer {
|
||||
/// Check if it's needed to do checkpoint for the region by the checkpoint distance.
|
||||
/// If needed, and there's no currently running checkpoint task, it will start a new checkpoint
|
||||
/// task running in the background.
|
||||
pub(crate) fn maybe_do_checkpoint(&self, manifest: &RegionManifest) {
|
||||
pub(crate) async fn maybe_do_checkpoint(&mut self, manifest: &RegionManifest) {
|
||||
if self
|
||||
.checkpoint_task
|
||||
.as_ref()
|
||||
.is_some_and(|handle| !handle.is_finished())
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
// Reap a completed task before checking whether to start the next one.
|
||||
// This keeps the handle as the single source of truth for task state.
|
||||
self.wait_for_pending_checkpoint().await;
|
||||
|
||||
if self.manifest_options.checkpoint_distance == 0 {
|
||||
return;
|
||||
}
|
||||
@@ -152,12 +159,6 @@ impl Checkpointer {
|
||||
return;
|
||||
}
|
||||
|
||||
// We can simply check whether there's a running checkpoint task like this, all because of
|
||||
// the caller of this function is ran single threaded, inside the lock of RegionManifestManager.
|
||||
if self.inner.is_doing_checkpoint() {
|
||||
return;
|
||||
}
|
||||
|
||||
let start_version = if last_checkpoint_version == 0 {
|
||||
// Checkpoint version can't be zero by implementation.
|
||||
// So last checkpoint version is zero means no last checkpoint.
|
||||
@@ -181,17 +182,44 @@ impl Checkpointer {
|
||||
self.do_checkpoint(checkpoint);
|
||||
}
|
||||
|
||||
fn do_checkpoint(&self, checkpoint: RegionCheckpoint) {
|
||||
self.inner.set_doing_checkpoint();
|
||||
|
||||
fn do_checkpoint(&mut self, checkpoint: RegionCheckpoint) {
|
||||
let inner = self.inner.clone();
|
||||
common_runtime::spawn_global(async move {
|
||||
self.checkpoint_task = Some(common_runtime::spawn_global(async move {
|
||||
inner.do_checkpoint(checkpoint).await;
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
/// Waits for the current checkpoint task without removing its handle first.
|
||||
///
|
||||
/// Keeping the handle in `self` while awaiting is important. If the caller is
|
||||
/// cancelled, another lifecycle transition can still wait for the same task.
|
||||
pub(crate) async fn wait_for_pending_checkpoint(&mut self) {
|
||||
let Some(handle) = self.checkpoint_task.as_mut() else {
|
||||
return;
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
self.pending_checkpoint_wait_started.notify_one();
|
||||
|
||||
let result = (&mut *handle).await;
|
||||
// There is no cancellation point between observing completion and
|
||||
// clearing the handle.
|
||||
self.checkpoint_task = None;
|
||||
|
||||
if let Err(e) = result {
|
||||
warn!(e; "Failed to join checkpoint task for region {}", self.inner.region_id());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn is_doing_checkpoint(&self) -> bool {
|
||||
self.inner.is_doing_checkpoint()
|
||||
self.checkpoint_task
|
||||
.as_ref()
|
||||
.is_some_and(|handle| !handle.is_finished())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn pending_checkpoint_wait_started(&self) -> Arc<Notify> {
|
||||
self.pending_checkpoint_wait_started.clone()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -550,6 +550,28 @@ impl RegionManifestManager {
|
||||
&mut self,
|
||||
action_list: RegionMetaActionList,
|
||||
is_staging: bool,
|
||||
) -> Result<ManifestVersion> {
|
||||
self.update_inner(action_list, is_staging, !is_staging)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Updates the normal manifest without starting a checkpoint.
|
||||
///
|
||||
/// This is only used while a leader is downgrading. The final flush still
|
||||
/// publishes its manifest edit, but must not start cleanup after the
|
||||
/// downgrade checkpoint barrier.
|
||||
pub(crate) async fn update_normal_without_checkpoint(
|
||||
&mut self,
|
||||
action_list: RegionMetaActionList,
|
||||
) -> Result<ManifestVersion> {
|
||||
self.update_inner(action_list, false, false).await
|
||||
}
|
||||
|
||||
async fn update_inner(
|
||||
&mut self,
|
||||
action_list: RegionMetaActionList,
|
||||
is_staging: bool,
|
||||
allow_checkpoint: bool,
|
||||
) -> Result<ManifestVersion> {
|
||||
let _t = MANIFEST_OP_ELAPSED
|
||||
.with_label_values(&["update"])
|
||||
@@ -616,13 +638,21 @@ impl RegionManifestManager {
|
||||
.checkpointer
|
||||
.update_manifest_removed_files(new_manifest)?;
|
||||
self.manifest = Arc::new(updated_manifest);
|
||||
self.checkpointer
|
||||
.maybe_do_checkpoint(self.manifest.as_ref());
|
||||
if allow_checkpoint {
|
||||
self.checkpointer
|
||||
.maybe_do_checkpoint(self.manifest.as_ref())
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(version)
|
||||
}
|
||||
|
||||
/// Waits for an in-flight checkpoint to finish, including its cleanup.
|
||||
pub(crate) async fn wait_for_pending_checkpoint(&mut self) {
|
||||
self.checkpointer.wait_for_pending_checkpoint().await;
|
||||
}
|
||||
|
||||
/// Clear deleted files from manifest's `removed_files` field without update version. Notice if datanode exit before checkpoint then new manifest by open region may still contain these deleted files, which is acceptable for gc process.
|
||||
pub fn clear_deleted_files(&mut self, deleted_files: Vec<RemovedFile>) {
|
||||
let mut manifest = (*self.manifest()).clone();
|
||||
|
||||
@@ -26,7 +26,8 @@ use tokio::sync::Semaphore;
|
||||
|
||||
use crate::cache::manifest_cache::ManifestCache;
|
||||
use crate::error::{
|
||||
CompressObjectSnafu, DecompressObjectSnafu, InvalidScanIndexSnafu, OpenDalSnafu, Result,
|
||||
CompressObjectSnafu, DecompressObjectSnafu, InvalidScanIndexSnafu, ManifestDeltaNotFoundSnafu,
|
||||
OpenDalSnafu, Result,
|
||||
};
|
||||
use crate::manifest::storage::size_tracker::Tracker;
|
||||
use crate::manifest::storage::utils::{
|
||||
@@ -212,11 +213,16 @@ impl<T: Tracker> DeltaStorage<T> {
|
||||
|
||||
// Fetch from remote object store
|
||||
let compress_type = file_compress_type(entry.name());
|
||||
let bytes = self
|
||||
.object_store
|
||||
.read(entry.path())
|
||||
.await
|
||||
.context(OpenDalSnafu)?;
|
||||
let bytes = match self.object_store.read(entry.path()).await {
|
||||
Ok(bytes) => bytes,
|
||||
Err(error) if error.kind() == ErrorKind::NotFound => {
|
||||
return Err(error).context(ManifestDeltaNotFoundSnafu {
|
||||
version: *v,
|
||||
path: entry.path(),
|
||||
});
|
||||
}
|
||||
Err(error) => return Err(error).context(OpenDalSnafu),
|
||||
};
|
||||
let data = compress_type
|
||||
.decode(bytes)
|
||||
.await
|
||||
|
||||
@@ -18,13 +18,16 @@ use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use common_datasource::compression::CompressionType;
|
||||
use common_error::ext::{ErrorExt, RetryHint};
|
||||
use common_error::status_code::StatusCode;
|
||||
use object_store::layers::mock::{
|
||||
Error as MockError, ErrorKind, MockLayerBuilder, OpDelete, Result as MockResult, oio,
|
||||
Buffer, Error as MockError, ErrorKind, MockLayer, MockLayerBuilder, OpDelete,
|
||||
Result as MockResult, oio,
|
||||
};
|
||||
use store_api::storage::{FileId, RegionId};
|
||||
use strum::IntoEnumIterator;
|
||||
|
||||
use crate::error::Error::ChecksumMismatch;
|
||||
use crate::error::Error::{ChecksumMismatch, ManifestDeltaNotFound};
|
||||
use crate::manifest::action::{
|
||||
RegionCheckpoint, RegionEdit, RegionMetaAction, RegionMetaActionList,
|
||||
};
|
||||
@@ -33,7 +36,7 @@ use crate::manifest::storage::checkpoint::CheckpointMetadata;
|
||||
use crate::manifest::storage::is_delta_file;
|
||||
use crate::manifest::tests::utils::basic_region_metadata;
|
||||
use crate::sst::file::FileMeta;
|
||||
use crate::test_util::TestEnv;
|
||||
use crate::test_util::{CheckpointTaskBlocker, TestEnv};
|
||||
|
||||
async fn build_manager(
|
||||
checkpoint_distance: u64,
|
||||
@@ -85,6 +88,31 @@ fn nop_action() -> RegionMetaActionList {
|
||||
})])
|
||||
}
|
||||
|
||||
struct NotFoundReader;
|
||||
|
||||
impl oio::Read for NotFoundReader {
|
||||
async fn read(&mut self) -> MockResult<Buffer> {
|
||||
Err(MockError::new(
|
||||
ErrorKind::NotFound,
|
||||
"mock listed manifest delta not found",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn fail_manifest_delta_reads_layer() -> MockLayer {
|
||||
MockLayerBuilder::default()
|
||||
.reader_factory(Arc::new(|path, _args, inner| {
|
||||
let file_name = path.rsplit('/').next().unwrap_or(path);
|
||||
if is_delta_file(file_name) {
|
||||
Box::new(NotFoundReader)
|
||||
} else {
|
||||
inner
|
||||
}
|
||||
}))
|
||||
.build()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manager_without_checkpoint() {
|
||||
let (_env, mut manager) = build_manager(0, CompressionType::Uncompressed).await;
|
||||
@@ -589,3 +617,167 @@ async fn checkpoint_advances_and_recovery_works_when_delete_fails() {
|
||||
.expect("manifest should be recoverable");
|
||||
assert_eq!(reopened.manifest().manifest_version, 10);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn open_preserves_listed_delta_not_found_retry_hint() {
|
||||
let env = TestEnv::new()
|
||||
.await
|
||||
.with_mock_layer(fail_manifest_delta_reads_layer());
|
||||
let metadata = Arc::new(basic_region_metadata());
|
||||
let mut manager = env
|
||||
.create_manifest_manager(CompressionType::Uncompressed, 0, Some(metadata.clone()))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
manager.stop().await;
|
||||
|
||||
let error = env
|
||||
.create_manifest_manager(CompressionType::Uncompressed, 0, None)
|
||||
.await
|
||||
.expect_err("reopen must fail on the mocked delta read");
|
||||
assert_matches!(
|
||||
&error,
|
||||
ManifestDeltaNotFound {
|
||||
version: 0,
|
||||
path,
|
||||
error,
|
||||
..
|
||||
} if path.ends_with("00000000000000000000.json")
|
||||
&& error.kind() == object_store::ErrorKind::NotFound
|
||||
);
|
||||
assert_eq!(StatusCode::StorageUnavailable, error.status_code());
|
||||
assert_eq!(RetryHint::Retryable, error.retry_hint());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn install_preserves_listed_delta_not_found_retry_hint() {
|
||||
let env = TestEnv::new()
|
||||
.await
|
||||
.with_mock_layer(fail_manifest_delta_reads_layer());
|
||||
let metadata = Arc::new(basic_region_metadata());
|
||||
let mut manager = env
|
||||
.create_manifest_manager(CompressionType::Uncompressed, 0, Some(metadata))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut store = manager.store();
|
||||
store
|
||||
.save(1, &nop_action().encode().unwrap(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = manager.install_manifest_to(1).await.unwrap_err();
|
||||
assert_matches!(
|
||||
&error,
|
||||
ManifestDeltaNotFound {
|
||||
version: 1,
|
||||
path,
|
||||
..
|
||||
} if path.ends_with("00000000000000000001.json")
|
||||
);
|
||||
assert_eq!(RetryHint::Retryable, error.retry_hint());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_waiter_keeps_pending_checkpoint_handle() {
|
||||
let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup();
|
||||
let env = TestEnv::new().await.with_mock_layer(mock_layer);
|
||||
let metadata = Arc::new(basic_region_metadata());
|
||||
let mut manager = env
|
||||
.create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
manager.update(nop_action(), false).await.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("checkpoint cleanup did not start");
|
||||
|
||||
let wait_started = manager.checkpointer().pending_checkpoint_wait_started();
|
||||
let wait_started = wait_started.notified();
|
||||
let mut first_wait = Box::pin(manager.wait_for_pending_checkpoint());
|
||||
tokio::select! {
|
||||
_ = wait_started => {}
|
||||
_ = &mut first_wait => panic!("checkpoint waiter returned before cleanup finished"),
|
||||
}
|
||||
drop(first_wait);
|
||||
assert!(manager.checkpointer().is_doing_checkpoint());
|
||||
|
||||
blocker.release();
|
||||
manager.wait_for_pending_checkpoint().await;
|
||||
assert!(!manager.checkpointer().is_doing_checkpoint());
|
||||
|
||||
let (version, _) = manager
|
||||
.store()
|
||||
.load_last_checkpoint()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("checkpoint must be published");
|
||||
assert_eq!(1, version);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn running_checkpoint_prevents_scheduling_another_checkpoint() {
|
||||
let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup();
|
||||
let env = TestEnv::new().await.with_mock_layer(mock_layer);
|
||||
let metadata = Arc::new(basic_region_metadata());
|
||||
let mut manager = env
|
||||
.create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
manager.update(nop_action(), false).await.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("checkpoint cleanup did not start");
|
||||
|
||||
// The second update is eligible for checkpointing, but the first task still
|
||||
// owns the pending handle and must prevent another task from being scheduled.
|
||||
manager.update(nop_action(), false).await.unwrap();
|
||||
blocker.release();
|
||||
manager.wait_for_pending_checkpoint().await;
|
||||
|
||||
assert_eq!(2, manager.manifest().manifest_version);
|
||||
assert_eq!(1, manager.checkpointer().last_checkpoint_version());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn completed_checkpoint_is_reaped_before_scheduling_the_next_one() {
|
||||
let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup();
|
||||
let env = TestEnv::new().await.with_mock_layer(mock_layer);
|
||||
let metadata = Arc::new(basic_region_metadata());
|
||||
let mut manager = env
|
||||
.create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
manager.update(nop_action(), false).await.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("first checkpoint cleanup did not start");
|
||||
blocker.release();
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
while manager.checkpointer().is_doing_checkpoint() {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("first checkpoint did not finish");
|
||||
assert_eq!(1, manager.checkpointer().last_checkpoint_version());
|
||||
|
||||
// Do not explicitly wait/reap the completed handle. The next eligible
|
||||
// update must reap it before scheduling another checkpoint.
|
||||
blocker.arm_next_close();
|
||||
manager.update(nop_action(), false).await.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked())
|
||||
.await
|
||||
.expect("second checkpoint cleanup did not start");
|
||||
blocker.release();
|
||||
manager.wait_for_pending_checkpoint().await;
|
||||
|
||||
assert_eq!(2, manager.checkpointer().last_checkpoint_version());
|
||||
}
|
||||
|
||||
@@ -30,11 +30,11 @@ pub use mito_codec::key_values::KeyValues;
|
||||
use mito_codec::row_converter::{PrimaryKeyCodec, build_primary_key_codec};
|
||||
use snafu::ensure;
|
||||
use store_api::codec::PrimaryKeyEncoding;
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
use store_api::metadata::{RegionMetadata, RegionMetadataRef};
|
||||
use store_api::storage::{ColumnId, SequenceNumber, SequenceRange};
|
||||
|
||||
use crate::config::MitoConfig;
|
||||
use crate::error::{Result, UnsupportedOperationSnafu};
|
||||
use crate::error::{InvalidRegionOptionsSnafu, Result, UnsupportedOperationSnafu};
|
||||
use crate::flush::WriteBufferManagerRef;
|
||||
use crate::memtable::bulk::{BulkMemtableBuilder, CompactDispatcher};
|
||||
use crate::memtable::time_series::TimeSeriesMemtableBuilder;
|
||||
@@ -401,6 +401,26 @@ pub(crate) struct MemtableBuilderProvider {
|
||||
compact_dispatcher: Arc<CompactDispatcher>,
|
||||
}
|
||||
|
||||
/// Ensures JSON2 columns are not used with [`TimeSeriesMemtable`].
|
||||
pub(crate) fn ensure_json2_not_use_time_series_memtable(
|
||||
metadata: &RegionMetadata,
|
||||
options: &RegionOptions,
|
||||
) -> Result<()> {
|
||||
if metadata
|
||||
.column_metadatas
|
||||
.iter()
|
||||
.any(|x| x.column_schema.data_type.is_json2())
|
||||
{
|
||||
ensure!(
|
||||
!matches!(&options.memtable, Some(MemtableOptions::TimeSeries)),
|
||||
InvalidRegionOptionsSnafu {
|
||||
reason: "JSON2 columns only support BulkMemtable",
|
||||
}
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl MemtableBuilderProvider {
|
||||
pub(crate) fn new(
|
||||
write_buffer_manager: Option<WriteBufferManagerRef>,
|
||||
@@ -772,9 +792,15 @@ impl MemtableRange {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_error::ext::WhateverResult;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
|
||||
use store_api::metadata::RegionMetadataBuilder;
|
||||
|
||||
use super::*;
|
||||
use crate::flush::{WriteBufferManager, WriteBufferManagerImpl};
|
||||
use crate::memtable::bulk::BulkMemtableConfig;
|
||||
use crate::test_util::sst_util::sst_region_metadata;
|
||||
|
||||
#[test]
|
||||
fn test_alloc_tracker_without_manager() {
|
||||
@@ -848,4 +874,27 @@ mod tests {
|
||||
|
||||
assert_eq!(&config, builder.config());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json2_requires_bulk_memtable() -> WhateverResult<()> {
|
||||
let mut metadata = sst_region_metadata();
|
||||
metadata.column_metadatas[2].column_schema.data_type =
|
||||
ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new()));
|
||||
let metadata = RegionMetadataBuilder::from_existing(metadata).build()?;
|
||||
let mut options = RegionOptions {
|
||||
sst_format: Some(FormatType::PrimaryKey),
|
||||
memtable: Some(MemtableOptions::TimeSeries),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let err = ensure_json2_not_use_time_series_memtable(&metadata, &options).unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("JSON2 columns only support BulkMemtable")
|
||||
);
|
||||
|
||||
options.memtable = Some(MemtableOptions::Bulk(BulkMemtableConfig::default()));
|
||||
ensure_json2_not_use_time_series_memtable(&metadata, &options)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,8 +22,8 @@ use datatypes::arrow::array::{
|
||||
};
|
||||
use datatypes::arrow::buffer::Buffer;
|
||||
use datatypes::arrow_array::StringArray;
|
||||
use datatypes::data_type::DataType;
|
||||
use datatypes::prelude::{ConcreteDataType, MutableVector, VectorRef};
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use datatypes::value::ValueRef;
|
||||
use datatypes::vectors::StringVector;
|
||||
|
||||
@@ -35,11 +35,11 @@ pub(crate) enum FieldBuilder {
|
||||
|
||||
impl FieldBuilder {
|
||||
/// Creates a [FieldBuilder] instance with given type and capacity.
|
||||
pub fn create(data_type: &ConcreteDataType, init_cap: usize) -> Self {
|
||||
if let ConcreteDataType::String(_) = data_type {
|
||||
pub(crate) fn create(column_schema: &ColumnSchema, init_cap: usize) -> Self {
|
||||
if let ConcreteDataType::String(_) = &column_schema.data_type {
|
||||
Self::String(StringBuilder::with_capacity(init_cap / 16, init_cap))
|
||||
} else {
|
||||
Self::Other(data_type.create_mutable_vector(init_cap))
|
||||
Self::Other(column_schema.create_mutable_vector(init_cap))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+198
-26
@@ -16,7 +16,6 @@
|
||||
|
||||
pub(crate) mod chunk_reader;
|
||||
pub mod context;
|
||||
pub(crate) mod json_align;
|
||||
pub mod part;
|
||||
pub mod part_reader;
|
||||
mod row_group_reader;
|
||||
@@ -44,10 +43,12 @@ use store_api::metadata::RegionMetadataRef;
|
||||
use store_api::storage::{ColumnId, FileId, RegionId, SequenceRange};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::compaction::{
|
||||
Json2RewritePlans, collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema,
|
||||
};
|
||||
use crate::error::{Result, UnsupportedOperationSnafu};
|
||||
use crate::flush::WriteBufferManagerRef;
|
||||
use crate::memtable::bulk::context::BulkIterContext;
|
||||
use crate::memtable::bulk::json_align::Json2Aligner;
|
||||
use crate::memtable::bulk::part::{
|
||||
BulkPart, BulkPartEncodeMetrics, BulkPartEncoder, MultiBulkPart, UnorderedPart,
|
||||
should_prune_bulk_part,
|
||||
@@ -464,7 +465,8 @@ impl Memtable for BulkMemtable {
|
||||
|
||||
// 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()?
|
||||
&& let Some(bulk_part) =
|
||||
bulk_parts.unordered_part.to_bulk_part(&self.metadata)?
|
||||
{
|
||||
bulk_parts.parts.push(BulkPartWrapper {
|
||||
part: PartToMerge::Bulk {
|
||||
@@ -525,7 +527,8 @@ impl Memtable for BulkMemtable {
|
||||
|
||||
// Adds range for unordered part if not empty
|
||||
if !bulk_parts.unordered_part.is_empty()
|
||||
&& let Some(unordered_bulk_part) = bulk_parts.unordered_part.to_bulk_part()?
|
||||
&& let Some(unordered_bulk_part) =
|
||||
bulk_parts.unordered_part.to_bulk_part(&self.metadata)?
|
||||
{
|
||||
let part_stats = unordered_bulk_part.to_memtable_stats(&self.metadata);
|
||||
let range = MemtableRange::new(
|
||||
@@ -1116,8 +1119,9 @@ impl PartToMerge {
|
||||
fn create_iterator(
|
||||
self,
|
||||
context: Arc<BulkIterContext>,
|
||||
plans: Arc<Json2RewritePlans>,
|
||||
) -> Result<Option<BoxedRecordBatchIterator>> {
|
||||
match self {
|
||||
let iter = match self {
|
||||
PartToMerge::Bulk { part, .. } => {
|
||||
let series_count = part.estimated_series_count();
|
||||
let iter = BulkPartBatchIter::from_single(
|
||||
@@ -1127,10 +1131,18 @@ impl PartToMerge {
|
||||
series_count,
|
||||
None, // No metrics for merging
|
||||
);
|
||||
Ok(Some(Box::new(iter) as BoxedRecordBatchIterator))
|
||||
Some(Box::new(iter) as BoxedRecordBatchIterator)
|
||||
}
|
||||
PartToMerge::Multi { part, .. } => part.read(context, None, None),
|
||||
PartToMerge::Encoded { part, .. } => part.read(context, None, None),
|
||||
PartToMerge::Multi { part, .. } => part.read(context, None, None)?,
|
||||
PartToMerge::Encoded { part, .. } => part.read(context, None, None)?,
|
||||
};
|
||||
if plans.is_empty() {
|
||||
Ok(iter)
|
||||
} else {
|
||||
Ok(iter.map(|x| {
|
||||
Box::new(x.map(move |batch| rewrite_json2_batch(batch?, &plans)))
|
||||
as BoxedRecordBatchIterator
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1289,19 +1301,37 @@ impl MemtableCompactor {
|
||||
batch_size,
|
||||
)?);
|
||||
|
||||
let aligner = Json2Aligner::try_new(parts_to_merge.iter().map(PartToMerge::arrow_schema))?;
|
||||
let schemas = parts_to_merge
|
||||
.iter()
|
||||
.map(|part| (part.arrow_schema(), part.num_rows() as u64))
|
||||
.collect::<Vec<_>>();
|
||||
let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?);
|
||||
|
||||
debug_assert!(parts_to_merge.windows(2).all(|w| rewrite_json2_schema(
|
||||
&w[0].arrow_schema(),
|
||||
&plans
|
||||
) == rewrite_json2_schema(
|
||||
&w[1].arrow_schema(),
|
||||
&plans
|
||||
)));
|
||||
// Parts in one merge group may differ only in their JSON2 physical layouts. So every source
|
||||
// schema is therefore a valid template for producing the final target schema that has
|
||||
// the union JSON2 types (rewritten).
|
||||
let schema = rewrite_json2_schema(&parts_to_merge[0].arrow_schema(), &plans);
|
||||
|
||||
let iterators: Vec<BoxedRecordBatchIterator> = parts_to_merge
|
||||
.into_iter()
|
||||
.filter_map(|part| part.create_iterator(context.clone()).ok().flatten())
|
||||
.map(|iter| aligner.wrap_iter(iter))
|
||||
.map(|part| part.create_iterator(context.clone(), plans.clone()))
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect();
|
||||
|
||||
if iterators.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let merged_iter = FlatMergeIterator::new(aligner.schema().clone(), iterators, batch_size)?;
|
||||
let merged_iter = FlatMergeIterator::new(schema.clone(), iterators, batch_size)?;
|
||||
|
||||
let boxed_iter: BoxedRecordBatchIterator = if dedup {
|
||||
match merge_mode {
|
||||
@@ -1310,8 +1340,7 @@ impl MemtableCompactor {
|
||||
Box::new(dedup_iter)
|
||||
}
|
||||
MergeMode::LastNonNull => {
|
||||
let field_column_start =
|
||||
field_column_start(metadata, aligner.schema().fields().len());
|
||||
let field_column_start = field_column_start(metadata, schema.fields().len());
|
||||
|
||||
let dedup_iter = FlatDedupIterator::new(
|
||||
merged_iter,
|
||||
@@ -1332,7 +1361,7 @@ impl MemtableCompactor {
|
||||
let mut metrics = BulkPartEncodeMetrics::default();
|
||||
let encoded_part = encoder.encode_record_batch_iter(
|
||||
boxed_iter,
|
||||
aligner.schema().clone(),
|
||||
schema,
|
||||
min_timestamp,
|
||||
max_timestamp,
|
||||
max_sequence,
|
||||
@@ -1532,9 +1561,14 @@ mod tests {
|
||||
use api::helper::encode_json_value;
|
||||
use api::v1::value::ValueData;
|
||||
use api::v1::{Mutation, Row, Rows, SemanticType};
|
||||
use common_error::ext::WhateverResult;
|
||||
use datatypes::arrow::datatypes::DataType as ArrowDataType;
|
||||
use datatypes::data_type::ConcreteDataType;
|
||||
use datatypes::extension::json::Json2ExtensionType;
|
||||
use datatypes::extension::json::{
|
||||
JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, Json2PhysicalLayout, JsonMetadata,
|
||||
};
|
||||
use datatypes::json::value::JsonValue;
|
||||
use datatypes::json::{JsonSettings, JsonTypeHint};
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
|
||||
use mito_codec::row_converter::build_primary_key_codec;
|
||||
@@ -1678,7 +1712,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_bulk_memtable_compact_parts_with_json2() {
|
||||
let metadata = mock_metadata_with_json2();
|
||||
let metadata = mock_metadata_with_json2(JsonSettings::default());
|
||||
|
||||
let config = BulkMemtableConfig {
|
||||
merge_threshold: 2,
|
||||
@@ -1717,7 +1751,126 @@ mod tests {
|
||||
assert_eq!(4, total_rows);
|
||||
}
|
||||
|
||||
fn mock_metadata_with_json2() -> RegionMetadataRef {
|
||||
#[test]
|
||||
fn test_bulk_memtable_merge_bounds_json2_paths() -> WhateverResult<()> {
|
||||
let metadata = mock_metadata_with_json2(JsonSettings::try_new(vec![], Some(1))?);
|
||||
let first = mock_bulk_part_with_json2_values(
|
||||
&metadata,
|
||||
vec![1000, 2000],
|
||||
vec![json!({"a": 1}), json!({"a": 2})],
|
||||
100,
|
||||
)?;
|
||||
let second = mock_bulk_part_with_json2_values(
|
||||
&metadata,
|
||||
vec![3000, 4000],
|
||||
vec![json!({"b": 3}), json!({"b": 4})],
|
||||
200,
|
||||
)?;
|
||||
let parts = vec![
|
||||
PartToMerge::Bulk {
|
||||
part: first,
|
||||
file_id: FileId::random(),
|
||||
},
|
||||
PartToMerge::Bulk {
|
||||
part: second,
|
||||
file_id: FileId::random(),
|
||||
},
|
||||
];
|
||||
|
||||
let merged = MemtableCompactor::merge_parts_group(
|
||||
parts,
|
||||
&metadata,
|
||||
false,
|
||||
MergeMode::LastRow,
|
||||
usize::MAX,
|
||||
usize::MAX,
|
||||
DEFAULT_ROW_GROUP_SIZE,
|
||||
)?
|
||||
.unwrap();
|
||||
let MergedPart::Multi(part) = merged else {
|
||||
unreachable!()
|
||||
};
|
||||
let schema = part.schemas().next().unwrap();
|
||||
let ArrowDataType::Struct(fields) = schema.field(0).data_type() else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
vec![JSON2_REMAINDER_FIELD_NAME, "a"],
|
||||
fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unordered_parts_align_json2_layouts() -> WhateverResult<()> {
|
||||
let metadata = mock_metadata_with_json2(JsonSettings::try_new(vec![], Some(2))?);
|
||||
let memtable = BulkMemtable::new(
|
||||
42,
|
||||
BulkMemtableConfig::default(),
|
||||
metadata.clone(),
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
MergeMode::LastRow,
|
||||
);
|
||||
memtable.write_bulk(mock_bulk_part_with_json2_values(
|
||||
&metadata,
|
||||
vec![1000, 2000],
|
||||
vec![json!({"a": 1}), json!({"a": 2})],
|
||||
100,
|
||||
)?)?;
|
||||
memtable.write_bulk(mock_bulk_part_with_json2_values(
|
||||
&metadata,
|
||||
vec![3000, 4000],
|
||||
vec![json!({"b": 3}), json!({"b": 4})],
|
||||
200,
|
||||
)?)?;
|
||||
|
||||
let predicate = PredicateGroup::new(&metadata, &[])?;
|
||||
let ranges = memtable.ranges(None, RangesOptions::default().with_predicate(predicate))?;
|
||||
let range = ranges.ranges.values().next().unwrap();
|
||||
let batch = range.build_record_batch_iter(None, None)?.next().unwrap()?;
|
||||
let schema = batch.schema();
|
||||
let ArrowDataType::Struct(fields) = schema.field(0).data_type() else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
vec![JSON2_REMAINDER_FIELD_NAME, "a", "b"],
|
||||
fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bulk_part_converter_uses_json2_v2_layout() -> WhateverResult<()> {
|
||||
let settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["id".to_string()],
|
||||
data_type: ConcreteDataType::int64_datatype(),
|
||||
nullable: true,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
Some(0),
|
||||
)?;
|
||||
let metadata = mock_metadata_with_json2(settings);
|
||||
let part = mock_bulk_part_with_json2(&metadata, vec![1000, 2000], 100)?;
|
||||
let schema = part.batch.schema();
|
||||
let field = schema.field(0);
|
||||
let layout = Json2PhysicalLayout::try_from_root(field)?;
|
||||
|
||||
assert!(layout.is_version_2());
|
||||
let ArrowDataType::Struct(fields) = field.data_type() else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
vec![JSON2_REMAINDER_FIELD_NAME, "id"],
|
||||
fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn mock_metadata_with_json2(settings: JsonSettings) -> RegionMetadataRef {
|
||||
let col_meta_1 = ColumnMetadata {
|
||||
column_schema: ColumnSchema::new(
|
||||
"ts",
|
||||
@@ -1730,7 +1883,9 @@ mod tests {
|
||||
|
||||
let data_type = ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new()));
|
||||
let mut col_schema = ColumnSchema::new("data", data_type, true);
|
||||
col_schema.with_extension_type(&Json2ExtensionType::default());
|
||||
col_schema.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new(
|
||||
settings,
|
||||
))));
|
||||
|
||||
let col_meta_2 = ColumnMetadata {
|
||||
column_schema: col_schema,
|
||||
@@ -1749,8 +1904,29 @@ mod tests {
|
||||
metadata: &RegionMetadataRef,
|
||||
timestamps: Vec<i64>,
|
||||
sequence: u64,
|
||||
) -> Result<BulkPart> {
|
||||
let values = timestamps
|
||||
.iter()
|
||||
.map(|ts| {
|
||||
json!({
|
||||
"id": ts,
|
||||
"payload": {
|
||||
"message": format!("row-{ts}"),
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
mock_bulk_part_with_json2_values(metadata, timestamps, values, sequence)
|
||||
}
|
||||
|
||||
fn mock_bulk_part_with_json2_values(
|
||||
metadata: &RegionMetadataRef,
|
||||
timestamps: Vec<i64>,
|
||||
values: Vec<serde_json::Value>,
|
||||
sequence: u64,
|
||||
) -> Result<BulkPart> {
|
||||
let capacity = timestamps.len();
|
||||
debug_assert_eq!(capacity, values.len());
|
||||
let primary_key_codec = build_primary_key_codec(metadata);
|
||||
let json_type = JsonNativeType::Object(JsonObjectType::from([
|
||||
("id".to_string(), JsonNativeType::i64()),
|
||||
@@ -1773,16 +1949,12 @@ mod tests {
|
||||
|
||||
let rows = timestamps
|
||||
.into_iter()
|
||||
.map(|ts| {
|
||||
.zip(values)
|
||||
.map(|(ts, value)| {
|
||||
let val1 = api::v1::Value {
|
||||
value_data: Some(ValueData::TimestampMillisecondValue(ts)),
|
||||
};
|
||||
let value_data = ValueData::JsonValue(encode_json_value(JsonValue::from(json!({
|
||||
"id": ts,
|
||||
"payload": {
|
||||
"message": format!("row-{ts}"),
|
||||
},
|
||||
}))));
|
||||
let value_data = ValueData::JsonValue(encode_json_value(JsonValue::from(value)));
|
||||
let val2 = api::v1::Value {
|
||||
value_data: Some(value_data),
|
||||
};
|
||||
|
||||
@@ -1,451 +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::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use datatypes::arrow::datatypes::{DataType as ArrowDataType, Schema, SchemaRef};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::extension::json::is_json2_extension_type;
|
||||
use datatypes::types::json_type::JsonNativeType;
|
||||
use datatypes::vectors::json::array::JsonArray;
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
|
||||
use crate::error::{
|
||||
ConvertValueSnafu, DataTypeMismatchSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu,
|
||||
};
|
||||
use crate::memtable::BoxedRecordBatchIterator;
|
||||
|
||||
/// Aligns concrete JSON2 Arrow types across record batches.
|
||||
///
|
||||
/// JSON2 column concrete Arrow types are derived from data. Different memtable
|
||||
/// parts may therefore have different concrete types for the same JSON2 column.
|
||||
/// This helper merges those concrete types and aligns batches to the merged schema.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Json2Aligner {
|
||||
/// Schema after merging all JSON2 column concrete types.
|
||||
schema: SchemaRef,
|
||||
/// JSON2 columns that may need per-batch alignment.
|
||||
json_columns: Vec<(usize, ArrowDataType)>,
|
||||
}
|
||||
|
||||
impl Json2Aligner {
|
||||
/// Builds an aligner from input schemas.
|
||||
///
|
||||
/// Note: except for JSON2 columns, all input schemas must be identical.
|
||||
pub(crate) fn try_new<I>(input_schemas: I) -> Result<Self>
|
||||
where
|
||||
I: IntoIterator<Item = SchemaRef>,
|
||||
{
|
||||
let mut input_schemas = input_schemas.into_iter();
|
||||
|
||||
// Use first schema as base: it defines column order and non-JSON types.
|
||||
let base_schema = input_schemas.next().context(UnexpectedSnafu {
|
||||
reason: "Json2Aligner requires at least one input schema",
|
||||
})?;
|
||||
|
||||
// Init merged types from base schema.
|
||||
let mut merged_types = base_schema
|
||||
.fields()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|&(_idx, field)| is_json2_extension_type(field))
|
||||
.map(|(idx, field)| {
|
||||
let json_type =
|
||||
JsonNativeType::try_from(field.data_type()).context(DataTypeMismatchSnafu)?;
|
||||
Ok((idx, json_type))
|
||||
})
|
||||
.collect::<Result<HashMap<usize, JsonNativeType>>>()?;
|
||||
|
||||
// No JSON2 columns, no alignment needed.
|
||||
if merged_types.is_empty() {
|
||||
return Ok(Self {
|
||||
schema: base_schema,
|
||||
json_columns: Vec::new(),
|
||||
});
|
||||
}
|
||||
|
||||
// Merge JSON2 types from remaining schemas.
|
||||
for schema in input_schemas {
|
||||
// Input schemas should only differ in JSON2 concrete types.
|
||||
#[cfg(debug_assertions)]
|
||||
assert_columns_match_except_json2(&base_schema, &schema);
|
||||
|
||||
for (idx, merged) in &mut merged_types {
|
||||
if *idx >= schema.fields().len() {
|
||||
continue;
|
||||
}
|
||||
let json_type = JsonNativeType::try_from(schema.field(*idx).data_type())
|
||||
.context(DataTypeMismatchSnafu)?;
|
||||
merged.merge(&json_type);
|
||||
}
|
||||
}
|
||||
|
||||
// Build output schema with merged JSON2 types.
|
||||
let mut json_columns = Vec::with_capacity(merged_types.len());
|
||||
let fields: Vec<_> = base_schema
|
||||
.fields()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, field)| {
|
||||
if let Some(merged) = merged_types.get(&idx) {
|
||||
let data_type = merged.as_arrow_type();
|
||||
json_columns.push((idx, data_type.clone()));
|
||||
let mut field = (**field).clone();
|
||||
field.set_data_type(data_type);
|
||||
Arc::new(field)
|
||||
} else {
|
||||
field.clone()
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let schema = Arc::new(Schema::new_with_metadata(
|
||||
fields,
|
||||
base_schema.metadata().clone(),
|
||||
));
|
||||
|
||||
Ok(Self {
|
||||
schema,
|
||||
json_columns,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the aligned output schema.
|
||||
pub(crate) fn schema(&self) -> &SchemaRef {
|
||||
&self.schema
|
||||
}
|
||||
|
||||
/// Aligns a [`RecordBatch`] to [`Self::schema`].
|
||||
pub(crate) fn align_batch(&self, batch: RecordBatch) -> Result<RecordBatch> {
|
||||
if self.json_columns.is_empty() {
|
||||
return Ok(batch);
|
||||
}
|
||||
let mut cols = batch.columns().to_vec();
|
||||
for (idx, expected_type) in &self.json_columns {
|
||||
if batch.schema_ref().field(*idx).data_type() != expected_type {
|
||||
cols[*idx] = JsonArray::from(batch.column(*idx))
|
||||
.widen_to(expected_type)
|
||||
.context(ConvertValueSnafu)?;
|
||||
}
|
||||
}
|
||||
RecordBatch::try_new(self.schema.clone(), cols).context(NewRecordBatchSnafu)
|
||||
}
|
||||
|
||||
/// Aligns [`RecordBatch`]s to [`Self::schema`].
|
||||
pub(crate) fn align_batches<I>(&self, batches: I) -> Result<Vec<RecordBatch>>
|
||||
where
|
||||
I: IntoIterator<Item = RecordBatch>,
|
||||
{
|
||||
batches
|
||||
.into_iter()
|
||||
.map(|batch| self.align_batch(batch))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Wraps an iterator so each yielded [`RecordBatch`] is lazily aligned.
|
||||
pub(crate) fn wrap_iter(&self, iter: BoxedRecordBatchIterator) -> BoxedRecordBatchIterator {
|
||||
let aligner = self.clone();
|
||||
Box::new(iter.map(move |batch| aligner.align_batch(batch?)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(debug_assertions)]
|
||||
fn assert_columns_match_except_json2(base_schema: &Schema, schema: &Schema) {
|
||||
debug_assert_eq!(
|
||||
base_schema.fields().len(),
|
||||
schema.fields().len(),
|
||||
"input schemas for Json2Aligner must have the same column count"
|
||||
);
|
||||
for (idx, (base_field, field)) in base_schema.fields().iter().zip(schema.fields()).enumerate() {
|
||||
let base_is_json2 = is_json2_extension_type(base_field);
|
||||
let is_json2 = is_json2_extension_type(field);
|
||||
debug_assert_eq!(
|
||||
base_is_json2, is_json2,
|
||||
"column {idx} must be JSON2 in all input schemas or none"
|
||||
);
|
||||
if !base_is_json2 && !is_json2 {
|
||||
debug_assert_eq!(
|
||||
base_field, field,
|
||||
"non-JSON2 column {idx} must be identical across input schemas"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use datatypes::arrow::array::{
|
||||
Array, ArrayRef, AsArray, Int64Array, StringViewArray, StructArray, UInt64Array,
|
||||
};
|
||||
use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema};
|
||||
use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType};
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_try_new_rejects_empty_input() {
|
||||
let err = match Json2Aligner::try_new([]) {
|
||||
Ok(_) => panic!("expected empty input to fail"),
|
||||
Err(err) => err,
|
||||
};
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("Json2Aligner requires at least one input schema")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_try_new_keeps_non_json_schema_unchanged() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Arc::new(Field::new("ts", DataType::Int64, false)),
|
||||
Arc::new(Field::new("value", DataType::UInt64, true)),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values([1, 2])) as ArrayRef,
|
||||
Arc::new(UInt64Array::from(vec![Some(10), None])) as ArrayRef,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let aligner = Json2Aligner::try_new([schema.clone()]).unwrap();
|
||||
assert!(Arc::ptr_eq(aligner.schema(), &schema));
|
||||
|
||||
let aligned = aligner.align_batch(batch).unwrap();
|
||||
assert!(Arc::ptr_eq(aligned.schema_ref(), &schema));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_try_new_ignores_legacy_jsonb_extension_field() {
|
||||
let legacy_jsonb_field = Arc::new(
|
||||
Field::new("data", DataType::Binary, true).with_extension_type(JsonExtensionType),
|
||||
);
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Arc::new(Field::new("ts", DataType::Int64, false)),
|
||||
legacy_jsonb_field,
|
||||
]));
|
||||
|
||||
let aligner = Json2Aligner::try_new([schema.clone()]).unwrap();
|
||||
|
||||
assert!(Arc::ptr_eq(aligner.schema(), &schema));
|
||||
assert!(aligner.json_columns.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_try_new_merges_json2_object_fields() {
|
||||
let id_fields = Fields::from(vec![id_field()]);
|
||||
let name_fields = Fields::from(vec![name_field()]);
|
||||
let schema_with_id = schema_with_json_field(json_field("data", id_fields));
|
||||
let schema_with_name = schema_with_json_field(json_field("data", name_fields));
|
||||
|
||||
let aligner = Json2Aligner::try_new([schema_with_id, schema_with_name]).unwrap();
|
||||
let data_field = aligner.schema().field(1);
|
||||
let DataType::Struct(fields) = data_field.data_type() else {
|
||||
panic!("expected JSON2 field to be a struct");
|
||||
};
|
||||
|
||||
assert_eq!(2, fields.len());
|
||||
assert_eq!("id", fields[0].name());
|
||||
assert_eq!(&DataType::Int64, fields[0].data_type());
|
||||
assert_eq!("name", fields[1].name());
|
||||
assert_eq!(&DataType::Utf8View, fields[1].data_type());
|
||||
assert!(is_json2_extension_type(&aligner.schema().fields()[1]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_align_batch_fills_missing_json2_fields() {
|
||||
let id_fields = Fields::from(vec![id_field()]);
|
||||
let name_fields = Fields::from(vec![name_field()]);
|
||||
let schema_with_id = schema_with_json_field(json_field("data", id_fields.clone()));
|
||||
let schema_with_name = schema_with_json_field(json_field("data", name_fields.clone()));
|
||||
|
||||
let batch_with_id = RecordBatch::try_new(
|
||||
schema_with_id.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values([1, 2])) as ArrayRef,
|
||||
struct_array(
|
||||
id_fields,
|
||||
vec![Arc::new(Int64Array::from_iter_values([10, 20])) as ArrayRef],
|
||||
),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let batch_with_name = RecordBatch::try_new(
|
||||
schema_with_name.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values([3, 4])) as ArrayRef,
|
||||
struct_array(
|
||||
name_fields,
|
||||
vec![
|
||||
Arc::new(StringViewArray::from(vec![Some("alice"), Some("bob")]))
|
||||
as ArrayRef,
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let aligner = Json2Aligner::try_new([schema_with_id, schema_with_name]).unwrap();
|
||||
let aligned_with_id = aligner.align_batch(batch_with_id).unwrap();
|
||||
let aligned_with_name = aligner.align_batch(batch_with_name).unwrap();
|
||||
|
||||
let data_with_id = aligned_with_id
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<StructArray>()
|
||||
.unwrap();
|
||||
let id_values = data_with_id
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.unwrap();
|
||||
let missing_names = data_with_id.column(1);
|
||||
assert_eq!(10, id_values.value(0));
|
||||
assert_eq!(20, id_values.value(1));
|
||||
assert!(missing_names.is_null(0));
|
||||
assert!(missing_names.is_null(1));
|
||||
|
||||
let data_with_name = aligned_with_name
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<StructArray>()
|
||||
.unwrap();
|
||||
let missing_ids = data_with_name.column(0);
|
||||
let name_values = data_with_name
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<StringViewArray>()
|
||||
.unwrap();
|
||||
assert!(missing_ids.is_null(0));
|
||||
assert!(missing_ids.is_null(1));
|
||||
assert_eq!("alice", name_values.value(0));
|
||||
assert_eq!("bob", name_values.value(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_align_conflicting_number_types_as_variant() {
|
||||
let u64_fields = Fields::from(vec![Arc::new(Field::new("value", DataType::UInt64, true))]);
|
||||
let i64_fields = Fields::from(vec![Arc::new(Field::new("value", DataType::Int64, true))]);
|
||||
let u64_schema = schema_with_json_field(json_field("data", u64_fields.clone()));
|
||||
let i64_schema = schema_with_json_field(json_field("data", i64_fields.clone()));
|
||||
let u64_batch = RecordBatch::try_new(
|
||||
u64_schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values([1])) as ArrayRef,
|
||||
struct_array(
|
||||
u64_fields,
|
||||
vec![Arc::new(UInt64Array::from_iter_values([u64::MAX])) as ArrayRef],
|
||||
),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let i64_batch = RecordBatch::try_new(
|
||||
i64_schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values([2])) as ArrayRef,
|
||||
struct_array(
|
||||
i64_fields,
|
||||
vec![Arc::new(Int64Array::from_iter_values([i64::MIN])) as ArrayRef],
|
||||
),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let aligner = Json2Aligner::try_new([u64_schema, i64_schema]).unwrap();
|
||||
let DataType::Struct(fields) = aligner.schema().field(1).data_type() else {
|
||||
panic!("expected JSON2 field to be a struct");
|
||||
};
|
||||
assert_eq!(&DataType::Binary, fields[0].data_type());
|
||||
|
||||
for (batch, expected) in [(u64_batch, json!(u64::MAX)), (i64_batch, json!(i64::MIN))] {
|
||||
let aligned = aligner.align_batch(batch).unwrap();
|
||||
let data = aligned.column(1).as_struct();
|
||||
assert_eq!(
|
||||
expected,
|
||||
JsonArray::from(data.column(0)).try_get_value(0).unwrap()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrap_iter_aligns_each_batch() {
|
||||
let id_fields = Fields::from(vec![id_field()]);
|
||||
let name_fields = Fields::from(vec![name_field()]);
|
||||
let schema_with_id = schema_with_json_field(json_field("data", id_fields.clone()));
|
||||
let schema_with_name = schema_with_json_field(json_field("data", name_fields.clone()));
|
||||
|
||||
let batch_with_id = RecordBatch::try_new(
|
||||
schema_with_id.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values([1])) as ArrayRef,
|
||||
struct_array(
|
||||
id_fields,
|
||||
vec![Arc::new(Int64Array::from_iter_values([10])) as ArrayRef],
|
||||
),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let batch_with_name = RecordBatch::try_new(
|
||||
schema_with_name.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values([2])) as ArrayRef,
|
||||
struct_array(
|
||||
name_fields,
|
||||
vec![Arc::new(StringViewArray::from(vec![Some("alice")])) as ArrayRef],
|
||||
),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let aligner = Json2Aligner::try_new([schema_with_id, schema_with_name]).unwrap();
|
||||
let iter: BoxedRecordBatchIterator =
|
||||
Box::new(vec![Ok(batch_with_id), Ok(batch_with_name)].into_iter());
|
||||
let aligned = aligner.wrap_iter(iter).collect::<Result<Vec<_>>>().unwrap();
|
||||
|
||||
assert_eq!(2, aligned.len());
|
||||
assert!(Arc::ptr_eq(aligned[0].schema_ref(), aligner.schema()));
|
||||
assert!(Arc::ptr_eq(aligned[1].schema_ref(), aligner.schema()));
|
||||
}
|
||||
|
||||
fn json_field(name: &str, fields: Fields) -> Arc<Field> {
|
||||
Arc::new(
|
||||
Field::new(name, DataType::Struct(fields), true)
|
||||
.with_extension_type(Json2ExtensionType::default()),
|
||||
)
|
||||
}
|
||||
|
||||
fn schema_with_json_field(json_field: Arc<Field>) -> SchemaRef {
|
||||
Arc::new(Schema::new(vec![
|
||||
Arc::new(Field::new("ts", DataType::Int64, false)),
|
||||
json_field,
|
||||
]))
|
||||
}
|
||||
|
||||
fn id_field() -> Arc<Field> {
|
||||
Arc::new(Field::new("id", DataType::Int64, true))
|
||||
}
|
||||
|
||||
fn name_field() -> Arc<Field> {
|
||||
Arc::new(Field::new("name", DataType::Utf8View, true))
|
||||
}
|
||||
|
||||
fn struct_array(fields: Fields, columns: Vec<ArrayRef>) -> ArrayRef {
|
||||
Arc::new(StructArray::new(fields, columns, None))
|
||||
}
|
||||
}
|
||||
@@ -54,13 +54,13 @@ use store_api::metadata::{RegionMetadata, RegionMetadataRef};
|
||||
use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
|
||||
use store_api::storage::{ColumnId, FileId, SequenceNumber, SequenceRange};
|
||||
|
||||
use crate::compaction::{collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema};
|
||||
use crate::error::{
|
||||
self, ColumnNotFoundSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DataTypeMismatchSnafu,
|
||||
EncodeMemtableSnafu, EncodeSnafu, InvalidMetadataSnafu, InvalidRequestSnafu,
|
||||
NewRecordBatchSnafu, Result,
|
||||
};
|
||||
use crate::memtable::bulk::context::{BulkIterContext, BulkIterContextRef};
|
||||
use crate::memtable::bulk::json_align::Json2Aligner;
|
||||
use crate::memtable::bulk::part_reader::EncodedBulkPartIter;
|
||||
use crate::memtable::time_series::{ValueBuilder, Values};
|
||||
use crate::memtable::{BoxedRecordBatchIterator, MemScanMetrics, MemtableStats};
|
||||
@@ -440,7 +440,7 @@ impl UnorderedPart {
|
||||
|
||||
/// Concatenates and sorts all parts into a single RecordBatch.
|
||||
/// Returns None if the collection is empty.
|
||||
pub fn concat_and_sort(&self) -> Result<Option<RecordBatch>> {
|
||||
pub fn concat_and_sort(&self, metadata: &RegionMetadataRef) -> Result<Option<RecordBatch>> {
|
||||
if self.parts.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
@@ -450,17 +450,28 @@ impl UnorderedPart {
|
||||
return Ok(Some(self.parts[0].batch.clone()));
|
||||
}
|
||||
|
||||
// Get the schema from the first part
|
||||
let schema = self.parts[0].batch.schema();
|
||||
let concatenated = if schema.fields().iter().any(is_json2_extension_type) {
|
||||
let aligner = Json2Aligner::try_new(self.parts.iter().map(|part| part.batch.schema()))?;
|
||||
let aligned_batches =
|
||||
aligner.align_batches(self.parts.iter().map(|part| part.batch.clone()))?;
|
||||
concat_batches(aligner.schema(), &aligned_batches).context(ComputeArrowSnafu)?
|
||||
} else {
|
||||
concat_batches(&schema, self.parts.iter().map(|x| &x.batch))
|
||||
.context(ComputeArrowSnafu)?
|
||||
};
|
||||
let schemas = self
|
||||
.parts
|
||||
.iter()
|
||||
.map(|x| (x.batch.schema(), x.num_rows() as u64))
|
||||
.collect::<Vec<_>>();
|
||||
let plans = collect_json2_rewrite_plans(metadata, &schemas)?;
|
||||
|
||||
debug_assert!(self.parts.windows(2).all(|w| rewrite_json2_schema(
|
||||
&w[0].batch.schema(),
|
||||
&plans
|
||||
) == rewrite_json2_schema(
|
||||
&w[1].batch.schema(),
|
||||
&plans
|
||||
)));
|
||||
let schema = rewrite_json2_schema(&self.parts[0].batch.schema(), &plans);
|
||||
|
||||
let batches = self
|
||||
.parts
|
||||
.iter()
|
||||
.map(|x| rewrite_json2_batch(x.batch.clone(), &plans))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let concatenated = concat_batches(&schema, &batches).context(ComputeArrowSnafu)?;
|
||||
|
||||
// Sort the concatenated batch
|
||||
let sorted_batch = sort_primary_key_record_batch(&concatenated)?;
|
||||
@@ -470,8 +481,8 @@ impl UnorderedPart {
|
||||
|
||||
/// Converts all parts into a single sorted BulkPart.
|
||||
/// Returns None if the collection is empty.
|
||||
pub fn to_bulk_part(&self) -> Result<Option<BulkPart>> {
|
||||
let Some(sorted_batch) = self.concat_and_sort()? else {
|
||||
pub fn to_bulk_part(&self, metadata: &RegionMetadataRef) -> Result<Option<BulkPart>> {
|
||||
let Some(sorted_batch) = self.concat_and_sort(metadata)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
|
||||
@@ -29,6 +29,7 @@ use datatypes::arrow::array::ArrayRef;
|
||||
use datatypes::arrow_array::StringArray;
|
||||
use datatypes::data_type::{ConcreteDataType, DataType};
|
||||
use datatypes::prelude::{ScalarVector, Vector, VectorRef};
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use datatypes::types::TimestampType;
|
||||
use datatypes::value::{Value, ValueRef};
|
||||
use datatypes::vectors::{
|
||||
@@ -46,7 +47,7 @@ use crate::error::{
|
||||
self, ComputeArrowSnafu, ConvertVectorSnafu, EncodeSnafu, PrimaryKeyLengthMismatchSnafu, Result,
|
||||
};
|
||||
use crate::flush::WriteBufferManagerRef;
|
||||
use crate::memtable::builder::{FieldBuilder, StringBuilder};
|
||||
use crate::memtable::builder::FieldBuilder;
|
||||
use crate::memtable::bulk::part::BulkPart;
|
||||
use crate::memtable::simple_bulk_memtable::SimpleBulkMemtable;
|
||||
use crate::memtable::stats::WriteMetrics;
|
||||
@@ -918,7 +919,7 @@ pub(crate) struct ValueBuilder {
|
||||
sequence: Vec<u64>,
|
||||
op_type: Vec<u8>,
|
||||
fields: Vec<Option<FieldBuilder>>,
|
||||
field_types: Vec<ConcreteDataType>,
|
||||
field_schemas: Vec<ColumnSchema>,
|
||||
}
|
||||
|
||||
impl ValueBuilder {
|
||||
@@ -931,18 +932,18 @@ impl ValueBuilder {
|
||||
let sequence = Vec::with_capacity(capacity);
|
||||
let op_type = Vec::with_capacity(capacity);
|
||||
|
||||
let field_types = region_metadata
|
||||
let field_schemas = region_metadata
|
||||
.field_columns()
|
||||
.map(|c| c.column_schema.data_type.clone())
|
||||
.map(|c| c.column_schema.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let fields = (0..field_types.len()).map(|_| None).collect();
|
||||
let fields = (0..field_schemas.len()).map(|_| None).collect();
|
||||
Self {
|
||||
timestamp: Vec::with_capacity(capacity),
|
||||
timestamp_type,
|
||||
sequence,
|
||||
op_type,
|
||||
fields,
|
||||
field_types,
|
||||
field_schemas,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -984,15 +985,10 @@ impl ValueBuilder {
|
||||
.push(field_value)
|
||||
.unwrap_or_else(|e| panic!("Failed to push field value: {e:?}"));
|
||||
} else {
|
||||
let mut mutable_vector =
|
||||
if let ConcreteDataType::String(_) = &self.field_types[idx] {
|
||||
FieldBuilder::String(StringBuilder::with_capacity(4, 8))
|
||||
} else {
|
||||
FieldBuilder::Other(
|
||||
self.field_types[idx]
|
||||
.create_mutable_vector(num_rows.max(INITIAL_BUILDER_CAPACITY)),
|
||||
)
|
||||
};
|
||||
let mut mutable_vector = FieldBuilder::create(
|
||||
&self.field_schemas[idx],
|
||||
num_rows.max(INITIAL_BUILDER_CAPACITY),
|
||||
);
|
||||
mutable_vector.push_nulls(num_rows - 1);
|
||||
mutable_vector
|
||||
.push(field_value)
|
||||
@@ -1017,7 +1013,12 @@ impl ValueBuilder {
|
||||
/// the Arrow string array offset limit and thus can never be accommodated, not even by an
|
||||
/// empty builder.
|
||||
pub(crate) fn can_accommodate(&self, fields: &[VectorRef]) -> Result<bool> {
|
||||
scan_string_capacity(fields, &self.fields, &self.field_types, i32::MAX)
|
||||
let data_types = self
|
||||
.field_schemas
|
||||
.iter()
|
||||
.map(|x| x.data_type.clone())
|
||||
.collect::<Vec<_>>();
|
||||
scan_string_capacity(fields, &self.fields, &data_types, i32::MAX)
|
||||
}
|
||||
|
||||
pub(crate) fn extend(
|
||||
@@ -1082,7 +1083,7 @@ impl ValueBuilder {
|
||||
{
|
||||
let builder = field_dest.get_or_insert_with(|| {
|
||||
let mut field_builder =
|
||||
FieldBuilder::create(&self.field_types[field_idx], INITIAL_BUILDER_CAPACITY);
|
||||
FieldBuilder::create(&self.field_schemas[field_idx], INITIAL_BUILDER_CAPACITY);
|
||||
field_builder.push_nulls(num_rows_before);
|
||||
field_builder
|
||||
});
|
||||
@@ -1140,9 +1141,9 @@ impl ValueBuilder {
|
||||
MEMTABLE_ACTIVE_FIELD_BUILDER_COUNT.dec();
|
||||
v.finish_cloned()
|
||||
} else {
|
||||
let mut single_null = self.field_types[i].create_mutable_vector(num_rows);
|
||||
single_null.push_nulls(num_rows);
|
||||
single_null.to_vector()
|
||||
let mut builder = FieldBuilder::create(&self.field_schemas[i], num_rows);
|
||||
builder.push_nulls(num_rows);
|
||||
builder.finish()
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
@@ -1282,9 +1283,9 @@ impl From<ValueBuilder> for Values {
|
||||
MEMTABLE_ACTIVE_FIELD_BUILDER_COUNT.dec();
|
||||
v.finish()
|
||||
} else {
|
||||
let mut single_null = value.field_types[i].create_mutable_vector(num_rows);
|
||||
single_null.push_nulls(num_rows);
|
||||
single_null.to_vector()
|
||||
let mut builder = FieldBuilder::create(&value.field_schemas[i], num_rows);
|
||||
builder.push_nulls(num_rows);
|
||||
builder.finish()
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
@@ -1390,6 +1391,7 @@ mod tests {
|
||||
use store_api::storage::RegionId;
|
||||
|
||||
use super::*;
|
||||
use crate::memtable::builder::StringBuilder;
|
||||
use crate::test_util::column_metadata_to_column_schema;
|
||||
|
||||
fn schema_for_test() -> RegionMetadataRef {
|
||||
|
||||
@@ -95,16 +95,11 @@ impl FlatCompatBatch {
|
||||
compaction: bool,
|
||||
) -> Result<Option<Self>> {
|
||||
let actual = read_format.metadata();
|
||||
let format_projection = read_format.format_projection();
|
||||
let mut actual_schema = flat_projected_columns(actual, format_projection);
|
||||
for (column_id, target_type) in read_format.json_target_types().iter() {
|
||||
if let Some(i) = actual_schema
|
||||
.iter()
|
||||
.position(|(actual_column_id, _)| actual_column_id == column_id)
|
||||
{
|
||||
actual_schema[i].1 = ConcreteDataType::json2(target_type.clone());
|
||||
}
|
||||
}
|
||||
let actual_schema = flat_projected_columns(
|
||||
actual,
|
||||
read_format.format_projection(),
|
||||
read_format.json_target_types(),
|
||||
);
|
||||
|
||||
let expect_schema = mapper.batch_schema();
|
||||
if expect_schema == actual_schema
|
||||
@@ -176,7 +171,19 @@ impl FlatCompatBatch {
|
||||
|
||||
// Same column different type.
|
||||
if expect_data_type != *actual_data_type {
|
||||
cast_type = Some(expect_data_type.clone())
|
||||
ensure!(
|
||||
!expect_data_type.is_json2() && !actual_data_type.is_json2(),
|
||||
CompatReaderSnafu {
|
||||
region_id: expect_metadata.region_id,
|
||||
reason: format!(
|
||||
"JSON2 column '{}' must be aligned before FlatCompatBatch, actual: {}, expected: {}",
|
||||
expect_column.column_schema.name,
|
||||
actual_data_type,
|
||||
expect_data_type,
|
||||
),
|
||||
}
|
||||
);
|
||||
cast_type = Some(expect_data_type.clone());
|
||||
}
|
||||
// Source has this column.
|
||||
index_or_defaults.push(IndexOrDefault::Index {
|
||||
@@ -634,6 +641,7 @@ impl FlatCompatPrimaryKey {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::{OpType, SemanticType};
|
||||
@@ -645,6 +653,7 @@ mod tests {
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use datatypes::types::json_type::JsonNativeType;
|
||||
use datatypes::value::ValueRef;
|
||||
use mito_codec::row_converter::{
|
||||
DensePrimaryKeyCodec, PrimaryKeyCodecExt, SparsePrimaryKeyCodec,
|
||||
@@ -814,6 +823,57 @@ mod tests {
|
||||
assert_eq!(expected_batch, result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flat_compat_batch_uses_projected_json2_type() -> Result<()> {
|
||||
let json2 = ConcreteDataType::json2(JsonNativeType::object());
|
||||
let actual_metadata = Arc::new(new_metadata(
|
||||
&[
|
||||
(
|
||||
0,
|
||||
SemanticType::Timestamp,
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
),
|
||||
(1, SemanticType::Field, json2.clone()),
|
||||
],
|
||||
&[],
|
||||
));
|
||||
let expected_metadata = Arc::new(new_metadata(
|
||||
&[
|
||||
(
|
||||
0,
|
||||
SemanticType::Timestamp,
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
),
|
||||
(1, SemanticType::Field, json2),
|
||||
(2, SemanticType::Field, ConcreteDataType::int64_datatype()),
|
||||
],
|
||||
&[],
|
||||
));
|
||||
let read_columns = ReadColumns::new([0, 1, 2])
|
||||
.with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)]));
|
||||
let mapper = FlatProjectionMapper::new_with_read_columns(
|
||||
&expected_metadata,
|
||||
vec![0, 1, 2],
|
||||
read_columns.clone(),
|
||||
)?;
|
||||
let read_format = FlatReadFormat::new(actual_metadata, read_columns, None, "test", false)?;
|
||||
|
||||
let compat = FlatCompatBatch::try_new(&mapper, &read_format, false)?.unwrap();
|
||||
let json_index = mapper
|
||||
.batch_schema()
|
||||
.iter()
|
||||
.position(|(id, _)| *id == 1)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
&compat.index_or_defaults[json_index],
|
||||
IndexOrDefault::Index {
|
||||
cast_type: None,
|
||||
..
|
||||
}
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flat_compat_batch_with_read_projection_superset() {
|
||||
let actual_metadata = Arc::new(new_metadata(
|
||||
|
||||
@@ -27,11 +27,10 @@ use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field};
|
||||
use datatypes::extension::json::is_json2_extension_type;
|
||||
use datatypes::prelude::{ConcreteDataType, DataType};
|
||||
use datatypes::schema::{Schema, SchemaRef};
|
||||
use datatypes::types::JsonType;
|
||||
use datatypes::types::json_type::JsonNativeType;
|
||||
use datatypes::value::Value;
|
||||
use datatypes::vectors::Helper;
|
||||
use datatypes::vectors::json::array::JsonArray;
|
||||
use datatypes::vectors::json::json2_physical_data_type;
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
use store_api::metadata::{RegionMetadata, RegionMetadataRef};
|
||||
use store_api::storage::ColumnId;
|
||||
@@ -39,7 +38,8 @@ use store_api::storage::ColumnId;
|
||||
use crate::cache::CacheStrategy;
|
||||
use crate::error::{InvalidRequestSnafu, RecordBatchSnafu, Result};
|
||||
use crate::read::projection::{read_column_ids_from_projection, repeated_vector_with_cache};
|
||||
use crate::read::read_columns::ReadColumns;
|
||||
use crate::read::read_columns::{JsonTargetTypes, ReadColumns};
|
||||
use crate::sst::parquet::Json2RewriteTargets;
|
||||
use crate::sst::parquet::flat_format::sst_column_id_indices;
|
||||
use crate::sst::parquet::format::FormatProjection;
|
||||
use crate::sst::{
|
||||
@@ -94,6 +94,21 @@ impl FlatProjectionMapper {
|
||||
metadata: &RegionMetadataRef,
|
||||
projection: Vec<usize>,
|
||||
read_cols: ReadColumns,
|
||||
) -> Result<Self> {
|
||||
Self::new_with_json2_rewrite_targets(
|
||||
metadata,
|
||||
projection,
|
||||
read_cols,
|
||||
&Json2RewriteTargets::default(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Returns a mapper for a compaction read with fixed JSON2 output layouts.
|
||||
pub(crate) fn new_with_json2_rewrite_targets(
|
||||
metadata: &RegionMetadataRef,
|
||||
projection: Vec<usize>,
|
||||
read_cols: ReadColumns,
|
||||
json2_rewrite_targets: &Json2RewriteTargets,
|
||||
) -> Result<Self> {
|
||||
// If the original projection is empty.
|
||||
let is_empty_projection = projection.is_empty();
|
||||
@@ -113,10 +128,8 @@ impl FlatProjectionMapper {
|
||||
output_col_ids.push(col.column_id);
|
||||
|
||||
let mut schema = col.column_schema.clone();
|
||||
if let Some(data_type) =
|
||||
json2_read_datatype(col.column_id, &schema.data_type, &read_cols)
|
||||
{
|
||||
schema.data_type = data_type;
|
||||
if let Some(data_type) = read_cols.json_target_type(col.column_id) {
|
||||
schema.data_type = ConcreteDataType::json2(data_type.clone());
|
||||
}
|
||||
col_schemas.push(schema);
|
||||
}
|
||||
@@ -133,16 +146,11 @@ impl FlatProjectionMapper {
|
||||
read_cols.clone(),
|
||||
);
|
||||
|
||||
let mut batch_schema = flat_projected_columns(metadata, &format_projection);
|
||||
let batch_schema =
|
||||
flat_projected_columns(metadata, &format_projection, read_cols.json_target_types());
|
||||
|
||||
for (column_id, data_type) in batch_schema.iter_mut() {
|
||||
if let Some(updated) = json2_read_datatype(*column_id, data_type, &read_cols) {
|
||||
*data_type = updated;
|
||||
}
|
||||
}
|
||||
|
||||
// Safety: We get the column id from the metadata.
|
||||
let input_arrow_schema = compute_input_arrow_schema(metadata, &batch_schema);
|
||||
let input_arrow_schema =
|
||||
compute_input_arrow_schema(metadata, &batch_schema, &read_cols, json2_rewrite_targets);
|
||||
|
||||
// If projection is empty, we don't output any column.
|
||||
let output_schema = if is_empty_projection {
|
||||
@@ -348,9 +356,9 @@ impl FlatProjectionMapper {
|
||||
}
|
||||
|
||||
let field = &self.output_schema.arrow_schema().fields()[output_idx];
|
||||
if is_json2_extension_type(field) {
|
||||
if is_json2_extension_type(field) && array.data_type() != field.data_type() {
|
||||
array = JsonArray::from(&array)
|
||||
.project_to(field.data_type())
|
||||
.project_to_v2(batch.schema_ref().field(*index), field.data_type())
|
||||
.context(DataTypesSnafu)?;
|
||||
}
|
||||
|
||||
@@ -391,35 +399,6 @@ impl FlatProjectionMapper {
|
||||
}
|
||||
}
|
||||
|
||||
fn json2_read_datatype(
|
||||
column_id: ColumnId,
|
||||
data_type: &ConcreteDataType,
|
||||
read_cols: &ReadColumns,
|
||||
) -> Option<ConcreteDataType> {
|
||||
let json_type = data_type.as_json()?;
|
||||
if !json_type.is_json2() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(concretized) = read_cols.json_target_type(column_id).cloned() {
|
||||
return Some(ConcreteDataType::json2(concretized));
|
||||
}
|
||||
|
||||
if is_empty_json2_type(json_type) {
|
||||
return Some(ConcreteDataType::json2(JsonNativeType::Variant));
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn is_empty_json2_type(json_type: &JsonType) -> bool {
|
||||
match json_type.native_type() {
|
||||
JsonNativeType::Null => true,
|
||||
JsonNativeType::Object(fields) if fields.is_empty() => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn single_value_string_dictionary<'a>(
|
||||
array: &'a Arc<dyn Array>,
|
||||
output_type: &ConcreteDataType,
|
||||
@@ -442,12 +421,13 @@ fn single_value_string_dictionary<'a>(
|
||||
(dict_array.values().len() == 1 && dict_array.null_count() == 0).then_some(dict_array)
|
||||
}
|
||||
|
||||
/// Returns ids and datatypes of columns of the output batch after applying the `projection`.
|
||||
/// Returns ids and datatypes of columns after applying the projection and JSON2 target types.
|
||||
///
|
||||
/// It adds the time index column if it doesn't present in the projection.
|
||||
pub(crate) fn flat_projected_columns(
|
||||
metadata: &RegionMetadata,
|
||||
format_projection: &FormatProjection,
|
||||
json_target_types: &JsonTargetTypes,
|
||||
) -> Vec<(ColumnId, ConcreteDataType)> {
|
||||
let time_index = metadata.time_index_column();
|
||||
let num_columns = if format_projection
|
||||
@@ -460,16 +440,18 @@ pub(crate) fn flat_projected_columns(
|
||||
};
|
||||
let mut schema = vec![None; num_columns];
|
||||
for (column_id, index) in &format_projection.column_id_to_projected_index {
|
||||
// Safety: FormatProjection ensures the id is valid.
|
||||
schema[*index] = Some((
|
||||
*column_id,
|
||||
let data_type = if let Some(json_type) = json_target_types.get(column_id) {
|
||||
ConcreteDataType::json2(json_type.clone())
|
||||
} else {
|
||||
// Safety: FormatProjection ensures the id is valid.
|
||||
metadata
|
||||
.column_by_id(*column_id)
|
||||
.unwrap()
|
||||
.column_schema
|
||||
.data_type
|
||||
.clone(),
|
||||
));
|
||||
.clone()
|
||||
};
|
||||
schema[*index] = Some((*column_id, data_type));
|
||||
}
|
||||
if num_columns != format_projection.column_id_to_projected_index.len() {
|
||||
schema[num_columns - 1] = Some((
|
||||
@@ -489,13 +471,25 @@ pub(crate) fn flat_projected_columns(
|
||||
pub(crate) fn compute_input_arrow_schema(
|
||||
metadata: &RegionMetadata,
|
||||
batch_schema: &[(ColumnId, ConcreteDataType)],
|
||||
read_cols: &ReadColumns,
|
||||
json2_rewrite_targets: &Json2RewriteTargets,
|
||||
) -> datatypes::arrow::datatypes::SchemaRef {
|
||||
let mut new_fields = Vec::with_capacity(batch_schema.len() + 3);
|
||||
for (column_id, data_type) in batch_schema {
|
||||
let data_type = json2_rewrite_targets
|
||||
.get(column_id)
|
||||
.map(|x| json2_physical_data_type(&x.target_layout))
|
||||
.or_else(|| {
|
||||
read_cols
|
||||
.json_target_type(*column_id)
|
||||
.map(|x| x.as_arrow_type())
|
||||
})
|
||||
.unwrap_or_else(|| data_type.as_arrow_type());
|
||||
|
||||
let column_metadata = metadata.column_by_id(*column_id).unwrap();
|
||||
let field = Field::new(
|
||||
&column_metadata.column_schema.name,
|
||||
data_type.as_arrow_type(),
|
||||
data_type,
|
||||
column_metadata.column_schema.is_nullable(),
|
||||
)
|
||||
.with_metadata(column_metadata.column_schema.metadata().clone());
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::hash::Hash;
|
||||
use std::mem;
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -46,6 +47,7 @@ impl ReadColumns {
|
||||
}
|
||||
}
|
||||
|
||||
/// Attaches query-time JSON2 projection types.
|
||||
pub fn with_json_target_types(
|
||||
mut self,
|
||||
json_target_types: BTreeMap<ColumnId, JsonNativeType>,
|
||||
@@ -66,10 +68,11 @@ impl ReadColumns {
|
||||
self.column_ids_iter().collect()
|
||||
}
|
||||
|
||||
pub fn json_target_types(&self) -> &JsonTargetTypes {
|
||||
pub(crate) fn json_target_types(&self) -> &JsonTargetTypes {
|
||||
&self.json_target_types
|
||||
}
|
||||
|
||||
/// Returns the query-time JSON2 projection type for a column.
|
||||
pub fn json_target_type(&self, column_id: ColumnId) -> Option<&JsonNativeType> {
|
||||
self.json_target_types.get(&column_id)
|
||||
}
|
||||
@@ -77,7 +80,6 @@ impl ReadColumns {
|
||||
pub fn estimated_size(&self) -> usize {
|
||||
self.col_ids.capacity() * mem::size_of::<ColumnId>()
|
||||
+ self.col_ids.len() * mem::size_of::<ColumnId>()
|
||||
+ self.json_target_types.len()
|
||||
* (mem::size_of::<ColumnId>() + mem::size_of::<JsonNativeType>())
|
||||
+ self.json_target_types.len() * (size_of::<ColumnId>() + size_of::<JsonNativeType>())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -82,6 +82,7 @@ use crate::sst::index::inverted_index::applier::InvertedIndexApplierRef;
|
||||
use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
|
||||
#[cfg(feature = "vector_index")]
|
||||
use crate::sst::index::vector_index::applier::{VectorIndexApplier, VectorIndexApplierRef};
|
||||
use crate::sst::parquet::Json2RewriteTargets;
|
||||
use crate::sst::parquet::file_range::PreFilterMode;
|
||||
use crate::sst::parquet::reader::ReaderMetrics;
|
||||
|
||||
@@ -927,6 +928,8 @@ pub struct ScanInput {
|
||||
pub(crate) snapshot_sequence: Option<SequenceNumber>,
|
||||
/// Whether this scan is for compaction.
|
||||
pub(crate) compaction: bool,
|
||||
/// Compaction-only JSON2 physical rewrite targets.
|
||||
json2_rewrite_targets: Json2RewriteTargets,
|
||||
/// Counters that should receive query-load metrics.
|
||||
pub(crate) query_stat_counters: Option<RegionQueryStatCounters>,
|
||||
#[cfg(feature = "enterprise")]
|
||||
@@ -966,6 +969,7 @@ impl ScanInput {
|
||||
explain_flat_format: false,
|
||||
snapshot_sequence: None,
|
||||
compaction: false,
|
||||
json2_rewrite_targets: Arc::default(),
|
||||
query_stat_counters: None,
|
||||
#[cfg(feature = "enterprise")]
|
||||
extension_ranges: Vec::new(),
|
||||
@@ -1164,6 +1168,13 @@ impl ScanInput {
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets compaction-only JSON2 physical rewrite targets.
|
||||
#[must_use]
|
||||
pub(crate) fn with_json2_rewrite_targets(mut self, targets: Json2RewriteTargets) -> Self {
|
||||
self.json2_rewrite_targets = targets;
|
||||
self
|
||||
}
|
||||
|
||||
/// Builds memtable ranges to scan by `index`.
|
||||
pub(crate) fn build_mem_ranges(&self, index: RowGroupIndex) -> SmallVec<[MemtableRange; 2]> {
|
||||
let memtable = &self.memtables[index.index];
|
||||
@@ -1291,6 +1302,7 @@ impl ScanInput {
|
||||
.read_sst(file.clone())
|
||||
.predicate(predicate)
|
||||
.projection(Some(self.read_cols.clone()))
|
||||
.json2_rewrite_targets(self.json2_rewrite_targets.clone())
|
||||
.cache(self.cache_strategy.clone())
|
||||
.inverted_index_appliers(self.inverted_index_appliers.clone())
|
||||
.bloom_filter_index_appliers(self.bloom_filter_index_appliers.clone())
|
||||
|
||||
+21
-4
@@ -483,6 +483,7 @@ impl MitoRegion {
|
||||
let mut manager: RwLockWriteGuard<'_, RegionManifestManager> =
|
||||
self.manifest_ctx.manifest_manager.write().await;
|
||||
let current_state = self.state();
|
||||
let mut wait_for_checkpoint = false;
|
||||
|
||||
let hook_payload: Option<PendingManifestHook> = match state {
|
||||
SettableRegionRoleState::Leader => {
|
||||
@@ -545,10 +546,12 @@ impl MitoRegion {
|
||||
);
|
||||
self.exit_staging()?;
|
||||
self.set_role(RegionRole::Follower);
|
||||
wait_for_checkpoint = true;
|
||||
}
|
||||
RegionRoleState::Leader(_) => {
|
||||
info!("Demoting region {} from leader to follower", self.region_id);
|
||||
self.set_role(RegionRole::Follower);
|
||||
wait_for_checkpoint = true;
|
||||
}
|
||||
RegionRoleState::Follower => {
|
||||
// Already in desired state - no-op
|
||||
@@ -568,14 +571,17 @@ impl MitoRegion {
|
||||
);
|
||||
self.exit_staging()?;
|
||||
self.set_role(RegionRole::DowngradingLeader);
|
||||
wait_for_checkpoint = true;
|
||||
}
|
||||
RegionRoleState::Leader(RegionLeaderState::Writable) => {
|
||||
info!("Starting downgrade for region {}", self.region_id);
|
||||
self.set_role(RegionRole::DowngradingLeader);
|
||||
wait_for_checkpoint = true;
|
||||
}
|
||||
RegionRoleState::Leader(RegionLeaderState::Downgrading) => {
|
||||
// Already in desired state - no-op
|
||||
info!("Region {} already in downgrading mode", self.region_id);
|
||||
wait_for_checkpoint = true;
|
||||
}
|
||||
_ => {
|
||||
warn!(
|
||||
@@ -588,6 +594,13 @@ impl MitoRegion {
|
||||
}
|
||||
};
|
||||
|
||||
// The state is changed before waiting, so no new writable-leader work
|
||||
// can race with the barrier. Keep the manager lock while joining to
|
||||
// serialize the barrier with checkpoint scheduling.
|
||||
if wait_for_checkpoint {
|
||||
manager.wait_for_pending_checkpoint().await;
|
||||
}
|
||||
|
||||
// Hack(zhongzc): If we have just become leader (writable), persist any backfilled metadata.
|
||||
let mut backfill_hook_payload: Option<PendingManifestHook> = None;
|
||||
if self.state() == RegionRoleState::Leader(RegionLeaderState::Writable) {
|
||||
@@ -1371,10 +1384,14 @@ impl ManifestContext {
|
||||
// Clone before `action_list` is moved into `update` so the hook still
|
||||
// sees what was written.
|
||||
let action_list_for_hook = self.hook.as_ref().map(|_| action_list.clone());
|
||||
let version = manager
|
||||
.update(action_list, is_staging)
|
||||
.await
|
||||
.inspect_err(|e| error!(e; "Failed to update manifest, region_id: {}", region_id))?;
|
||||
let version = if !is_staging
|
||||
&& self.state.load() == RegionRoleState::Leader(RegionLeaderState::Downgrading)
|
||||
{
|
||||
manager.update_normal_without_checkpoint(action_list).await
|
||||
} else {
|
||||
manager.update(action_list, is_staging).await
|
||||
}
|
||||
.inspect_err(|e| error!(e; "Failed to update manifest, region_id: {}", region_id))?;
|
||||
|
||||
Ok(PendingManifestHook::new(
|
||||
region_id,
|
||||
|
||||
@@ -20,8 +20,10 @@ use std::sync::atomic::{AtomicI64, AtomicU64};
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use std::time::Instant;
|
||||
|
||||
use arrow_schema::extension::ExtensionType;
|
||||
use common_telemetry::{debug, error, info, warn};
|
||||
use common_wal::options::WalOptions;
|
||||
use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
|
||||
use futures::StreamExt;
|
||||
use futures::future::BoxFuture;
|
||||
use log_store::kafka::log_store::KafkaLogStore;
|
||||
@@ -50,14 +52,14 @@ use crate::config::MitoConfig;
|
||||
use crate::engine::region_hook::RegionHookRef;
|
||||
use crate::error;
|
||||
use crate::error::{
|
||||
EmptyRegionDirSnafu, InvalidMetadataSnafu, InvalidRegionOptionsSnafu, ObjectStoreNotFoundSnafu,
|
||||
RegionCorruptedSnafu, Result, StaleLogEntrySnafu,
|
||||
DataTypeMismatchSnafu, EmptyRegionDirSnafu, InvalidMetadataSnafu, InvalidRegionOptionsSnafu,
|
||||
ObjectStoreNotFoundSnafu, RegionCorruptedSnafu, Result, StaleLogEntrySnafu,
|
||||
};
|
||||
use crate::manifest::action::RegionManifest;
|
||||
use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
|
||||
use crate::memtable::MemtableBuilderProvider;
|
||||
use crate::memtable::bulk::part::BulkPart;
|
||||
use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef};
|
||||
use crate::memtable::{MemtableBuilderProvider, ensure_json2_not_use_time_series_memtable};
|
||||
use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES};
|
||||
use crate::region::options::RegionOptions;
|
||||
use crate::region::version::{VersionBuilder, VersionControl, VersionControlRef};
|
||||
@@ -93,6 +95,41 @@ fn initial_pruned_entry_id(wal_options: &WalOptions) -> EntryId {
|
||||
}
|
||||
}
|
||||
|
||||
fn maybe_upgrade_json2_layout(metadata: RegionMetadataRef) -> Result<RegionMetadataRef> {
|
||||
let mut upgrades = Vec::new();
|
||||
for (index, column) in metadata.column_metadatas.iter().enumerate() {
|
||||
if !column.column_schema.data_type.is_json2() {
|
||||
continue;
|
||||
}
|
||||
let Some(extension) = column
|
||||
.column_schema
|
||||
.extension_type::<Json2ExtensionType>()
|
||||
.context(DataTypeMismatchSnafu)?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if extension.metadata().is_version_2() {
|
||||
continue;
|
||||
}
|
||||
upgrades.push((index, extension.metadata().json_settings().clone()));
|
||||
}
|
||||
|
||||
if upgrades.is_empty() {
|
||||
return Ok(metadata);
|
||||
}
|
||||
|
||||
let mut upgraded = metadata.as_ref().clone();
|
||||
for (index, settings) in upgrades {
|
||||
let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings)));
|
||||
upgraded.column_metadatas[index]
|
||||
.column_schema
|
||||
.with_extension_type(&extension);
|
||||
}
|
||||
|
||||
let builder = RegionMetadataBuilder::from_existing(upgraded);
|
||||
Ok(Arc::new(builder.build().context(InvalidMetadataSnafu)?))
|
||||
}
|
||||
|
||||
/// A fetcher to retrieve partition expr for a region.
|
||||
///
|
||||
/// Compatibility: older regions didn't persist `partition_expr` in engine metadata,
|
||||
@@ -324,6 +361,7 @@ impl RegionOpener {
|
||||
options.sst_format = Some(FormatType::PrimaryKey);
|
||||
FormatType::PrimaryKey
|
||||
};
|
||||
ensure_json2_not_use_time_series_memtable(&metadata, &options)?;
|
||||
// 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);
|
||||
@@ -475,8 +513,10 @@ impl RegionOpener {
|
||||
} else {
|
||||
manifest.metadata.clone()
|
||||
};
|
||||
let metadata = maybe_upgrade_json2_layout(metadata)?;
|
||||
// Updates the region options with the manifest.
|
||||
sanitize_region_options(&manifest, &mut region_options);
|
||||
ensure_json2_not_use_time_series_memtable(&metadata, ®ion_options)?;
|
||||
|
||||
let region_id = self.region_id;
|
||||
let provider = self.provider::<S>(®ion_options.wal_options)?;
|
||||
@@ -1316,23 +1356,31 @@ mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::extension::ExtensionType;
|
||||
use common_base::readable_size::ReadableSize;
|
||||
use common_error::ext::WhateverResult;
|
||||
use common_test_util::temp_dir::create_temp_dir;
|
||||
use common_time::Timestamp;
|
||||
use common_wal::options::{KafkaWalOptions, WalOptions};
|
||||
use datatypes::arrow::array::{ArrayRef, BinaryArray, Int64Array};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
|
||||
use datatypes::json::JsonSettings;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
|
||||
use object_store::ObjectStore;
|
||||
use object_store::services::{Fs, Memory, S3};
|
||||
use parquet::arrow::ArrowWriter;
|
||||
use parquet::file::metadata::{KeyValue, PageIndexPolicy};
|
||||
use parquet::file::properties::WriterProperties;
|
||||
use store_api::metadata::RegionMetadataBuilder;
|
||||
use store_api::region_request::PathType;
|
||||
use store_api::storage::{FileId, RegionId};
|
||||
|
||||
use super::{
|
||||
initial_pruned_entry_id, preload_parquet_meta_cache_for_files, sanitize_region_options,
|
||||
supports_open_region_object_storage_requirement,
|
||||
initial_pruned_entry_id, maybe_upgrade_json2_layout, preload_parquet_meta_cache_for_files,
|
||||
sanitize_region_options, supports_open_region_object_storage_requirement,
|
||||
};
|
||||
use crate::cache::CacheManager;
|
||||
use crate::cache::file_cache::{FileType, IndexKey};
|
||||
@@ -1386,6 +1434,38 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_upgrade_json2_layout() -> WhateverResult<()> {
|
||||
let settings = JsonSettings::try_new(vec![], Some(3))?;
|
||||
let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1(settings.clone())));
|
||||
let mut column = ColumnSchema::new(
|
||||
"field_0",
|
||||
ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
|
||||
true,
|
||||
);
|
||||
column.with_extension_type(&extension);
|
||||
|
||||
let mut metadata = sst_region_metadata();
|
||||
metadata.column_metadatas[2].column_schema = column;
|
||||
let builder = RegionMetadataBuilder::from_existing(metadata);
|
||||
let metadata = Arc::new(builder.build()?);
|
||||
|
||||
let upgraded = maybe_upgrade_json2_layout(metadata)?;
|
||||
let column = &upgraded.column_metadatas[2].column_schema;
|
||||
let extension = column.extension_type::<Json2ExtensionType>()?.unwrap();
|
||||
assert!(extension.metadata().is_version_2());
|
||||
assert_eq!(&settings, extension.metadata().json_settings());
|
||||
|
||||
let arrow_schema = upgraded.schema.arrow_schema();
|
||||
let field = arrow_schema.field_with_name("field_0").unwrap();
|
||||
let extension = field.try_extension_type::<Json2ExtensionType>().unwrap();
|
||||
assert!(extension.metadata().is_version_2());
|
||||
|
||||
let unchanged = maybe_upgrade_json2_layout(upgraded.clone())?;
|
||||
assert!(Arc::ptr_eq(&upgraded, &unchanged));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(not(feature = "test-shared-fs-region-migration"))]
|
||||
fn test_open_requirement_rejects_fs_object_store() {
|
||||
|
||||
+130
-2
@@ -505,13 +505,20 @@ impl SeriesEstimator {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use ::parquet::arrow::AsyncArrowWriter;
|
||||
use ::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||
use ::parquet::basic::LogicalType;
|
||||
use ::parquet::variant::{VariantArray, VariantType, json_to_variant};
|
||||
use common_query::prelude::greptime_native_histogram;
|
||||
use datatypes::arrow::array::{
|
||||
BinaryArray, DictionaryArray, TimestampMillisecondArray, UInt8Array, UInt32Array,
|
||||
UInt64Array,
|
||||
ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, StructArray,
|
||||
TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array,
|
||||
};
|
||||
use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, TimeUnit};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::extension::json::{Json2ExtensionType, Json2PhysicalLayout};
|
||||
use datatypes::vectors::json::array::JsonArray;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -1004,4 +1011,125 @@ mod tests {
|
||||
err
|
||||
);
|
||||
}
|
||||
|
||||
fn json2_v2_test_type() -> ArrowDataType {
|
||||
ArrowDataType::Struct(
|
||||
vec![
|
||||
Arc::new(Field::new("active", ArrowDataType::Boolean, true)),
|
||||
Arc::new(Field::new("hot", ArrowDataType::Int64, true)),
|
||||
Arc::new(Field::new("name", ArrowDataType::Utf8, true)),
|
||||
]
|
||||
.into(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Validates the persisted-format foundation for the JSON2 v2 remainder.
|
||||
///
|
||||
/// JSON2 will store `!__remainder__!` as a nested Variant child of its root
|
||||
/// Struct. Before enabling that layout in production, this test ensures the
|
||||
/// SST schema wrapper and Arrow writer preserve the Variant extension,
|
||||
/// encode the Parquet Variant logical type, and round-trip the values
|
||||
/// without changing the surrounding Struct.
|
||||
#[tokio::test]
|
||||
async fn test_nested_variant_survives_sst_writer_schema_roundtrip()
|
||||
-> Result<(), Box<dyn std::error::Error>> {
|
||||
let json: ArrayRef = Arc::new(StringArray::from(vec![
|
||||
Some(r#"{}"#),
|
||||
Some(r#"{"name":"Alice","active":true}"#),
|
||||
Some(r#"{"nested":{"count":42},"items":[1,"two",null]}"#),
|
||||
Some(r#"{"\u5b57\u6bb5":"\u503c"}"#),
|
||||
None,
|
||||
]));
|
||||
let remainder = json_to_variant(&json)?;
|
||||
let remainder_field = remainder.field("!__remainder__!");
|
||||
let remainder_array = ArrayRef::from(remainder);
|
||||
let hot_field = Field::new("hot", ArrowDataType::Int64, true);
|
||||
let data_array = Arc::new(StructArray::new(
|
||||
vec![remainder_field.clone(), hot_field.clone()].into(),
|
||||
vec![
|
||||
remainder_array,
|
||||
Arc::new(Int64Array::from(vec![
|
||||
Some(1),
|
||||
Some(2),
|
||||
Some(3),
|
||||
Some(4),
|
||||
None,
|
||||
])),
|
||||
],
|
||||
None,
|
||||
));
|
||||
let data_field = Field::new(
|
||||
"data",
|
||||
ArrowDataType::Struct(vec![remainder_field, hot_field].into()),
|
||||
true,
|
||||
)
|
||||
.with_extension_type(Json2ExtensionType::default());
|
||||
let schema = Arc::new(Schema::new(vec![data_field]));
|
||||
let source = RecordBatch::try_new(schema.clone(), vec![data_array])?;
|
||||
|
||||
let wrapped = maybe_wrap_schema(&schema)?;
|
||||
let mut buffer = Vec::new();
|
||||
let mut writer = AsyncArrowWriter::try_new(&mut buffer, wrapped, None)?;
|
||||
writer.write(&source).await?;
|
||||
writer.close().await?;
|
||||
|
||||
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(buffer))?;
|
||||
let parquet_remainder =
|
||||
&builder.parquet_schema().root_schema().get_fields()[0].get_fields()[0];
|
||||
assert_eq!(
|
||||
parquet_remainder.get_basic_info().logical_type_ref(),
|
||||
Some(&LogicalType::Variant {
|
||||
specification_version: None,
|
||||
})
|
||||
);
|
||||
|
||||
let ArrowDataType::Struct(children) = builder.schema().field_with_name("data")?.data_type()
|
||||
else {
|
||||
unreachable!();
|
||||
};
|
||||
assert!(children[0].has_valid_extension_type::<VariantType>());
|
||||
|
||||
let mut reader = builder.build()?;
|
||||
let result = reader.next().unwrap()?;
|
||||
assert_eq!(source, result);
|
||||
let result_field = result.schema().field(0).clone();
|
||||
let result = result
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<StructArray>()
|
||||
.unwrap();
|
||||
VariantArray::try_new(result.column(0))?;
|
||||
let result: ArrayRef = Arc::new(result.clone());
|
||||
let result =
|
||||
JsonArray::from(&result).project_to_v2(&result_field, &json2_v2_test_type())?;
|
||||
assert_eq!(
|
||||
json!({"active": true, "hot": 2, "name": "Alice"}),
|
||||
JsonArray::from(&result).try_get_value(1)?
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Ensures future readers retain compatibility with the first JSON2 v2 layout.
|
||||
#[test]
|
||||
fn test_read_json2_v2_fixture() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let bytes = bytes::Bytes::from_static(include_bytes!("../test-data/json2-v2.parquet"));
|
||||
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes)?;
|
||||
let field = builder.schema().field(0).clone();
|
||||
assert!(Json2PhysicalLayout::try_from_root(&field)?.is_version_2());
|
||||
|
||||
let batch = builder.build()?.next().unwrap()?;
|
||||
let data = batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<StructArray>()
|
||||
.unwrap();
|
||||
VariantArray::try_new(data.column(0))?;
|
||||
let data: ArrayRef = Arc::new(data.clone());
|
||||
let data = JsonArray::from(&data).project_to_v2(&field, &json2_v2_test_type())?;
|
||||
assert_eq!(
|
||||
json!({"active": true, "hot": 2, "name": "Alice"}),
|
||||
JsonArray::from(&data).try_get_value(1)?
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -706,6 +706,14 @@ pub struct IndexBuildTask {
|
||||
pub region_id: RegionId,
|
||||
/// The SST file handle to build index for.
|
||||
pub file: FileHandle,
|
||||
/// The target region metadata used to decode rows from the SST.
|
||||
///
|
||||
/// An SST may originate in another region while being visible in the target
|
||||
/// manifest. This metadata defines the target schema and sequence domain;
|
||||
/// applying the staging manifest only makes imported files visible. Index
|
||||
/// rebuild happens later when a flush, compaction, schema change, or manual
|
||||
/// index build request schedules it.
|
||||
pub(crate) target_region_metadata: RegionMetadataRef,
|
||||
/// The manifest state this build is based on.
|
||||
pub(crate) source: IndexBuildSource,
|
||||
pub reason: IndexBuildType,
|
||||
@@ -846,6 +854,7 @@ impl IndexBuildTask {
|
||||
let mut parquet_reader = self
|
||||
.access_layer
|
||||
.read_sst(self.file.clone()) // use the latest file handle instead of creating a new one
|
||||
.expected_metadata(Some(self.target_region_metadata.clone()))
|
||||
.build()
|
||||
.await?;
|
||||
|
||||
@@ -1494,8 +1503,12 @@ mod tests {
|
||||
use datatypes::schema::{
|
||||
ColumnSchema, FulltextOptions, SkippingIndexOptions, SkippingIndexType,
|
||||
};
|
||||
use datatypes::value::Value;
|
||||
use index::inverted_index::format::reader::InvertedIndexReader;
|
||||
use object_store::ObjectStore;
|
||||
use object_store::services::Memory;
|
||||
use partition::expr::col;
|
||||
use puffin::puffin_manager::{PuffinManager, PuffinReader};
|
||||
use puffin_manager::PuffinManagerFactory;
|
||||
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
|
||||
use tokio::sync::mpsc;
|
||||
@@ -2088,6 +2101,7 @@ mod tests {
|
||||
let task = IndexBuildTask {
|
||||
region_id,
|
||||
file,
|
||||
target_region_metadata: version_control.current().version.metadata.clone(),
|
||||
source: IndexBuildSource::new(
|
||||
file_meta,
|
||||
version_control.current().version.metadata.schema_version,
|
||||
@@ -2120,16 +2134,34 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_index_build_task_increments_legacy_index_version() {
|
||||
async fn test_index_build_task_foreign_file_uses_target_metadata() {
|
||||
let env = SchedulerEnv::new().await;
|
||||
let mut scheduler = env.mock_index_build_scheduler(4);
|
||||
let metadata = Arc::new(sst_region_metadata());
|
||||
let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
|
||||
let region_id = metadata.region_id;
|
||||
let source_metadata = Arc::new(sst_region_metadata());
|
||||
let mut target_metadata = (*source_metadata).clone();
|
||||
target_metadata.region_id = RegionId::new(1, 3);
|
||||
let mut target_builder = RegionMetadataBuilder::new(target_metadata.region_id);
|
||||
for mut column_metadata in target_metadata.column_metadatas.clone() {
|
||||
if column_metadata.column_id == 2 {
|
||||
column_metadata.column_schema =
|
||||
column_metadata.column_schema.with_inverted_index(true);
|
||||
}
|
||||
target_builder.push_column_metadata(column_metadata);
|
||||
}
|
||||
let partition_expr = col("field_0")
|
||||
.gt_eq(Value::UInt64(100))
|
||||
.and(col("field_0").lt(Value::UInt64(200)));
|
||||
target_builder
|
||||
.primary_key(target_metadata.primary_key.clone())
|
||||
.partition_expr_json(Some(partition_expr.as_json_str().unwrap()))
|
||||
.bump_version();
|
||||
let target_metadata = Arc::new(target_builder.build().unwrap());
|
||||
let manifest_ctx = env.mock_manifest_context(target_metadata.clone()).await;
|
||||
let region_id = target_metadata.region_id;
|
||||
let file_purger = Arc::new(NoopFilePurger {});
|
||||
let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await;
|
||||
let sst_info = mock_sst_file(source_metadata.clone(), &env, IndexBuildMode::Async).await;
|
||||
let file_meta = FileMeta {
|
||||
region_id,
|
||||
region_id: source_metadata.region_id,
|
||||
file_id: sst_info.file_id,
|
||||
file_size: sst_info.file_size,
|
||||
max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size,
|
||||
@@ -2144,8 +2176,8 @@ mod tests {
|
||||
seed_manifest_file(&manifest_ctx, &file_meta).await;
|
||||
let files = HashMap::from([(file_meta.file_id, file_meta.clone())]);
|
||||
let version_control =
|
||||
mock_version_control(metadata.clone(), file_purger.clone(), files).await;
|
||||
let indexer_builder = mock_indexer_builder(metadata.clone(), &env).await;
|
||||
mock_version_control(target_metadata.clone(), file_purger.clone(), files).await;
|
||||
let indexer_builder = mock_indexer_builder(target_metadata.clone(), &env).await;
|
||||
|
||||
let file = FileHandle::new(file_meta.clone(), file_purger.clone());
|
||||
|
||||
@@ -2155,6 +2187,7 @@ mod tests {
|
||||
let task = IndexBuildTask {
|
||||
region_id,
|
||||
file,
|
||||
target_region_metadata: version_control.current().version.metadata.clone(),
|
||||
source: IndexBuildSource::new(
|
||||
file_meta.clone(),
|
||||
version_control.current().version.metadata.schema_version,
|
||||
@@ -2199,9 +2232,40 @@ mod tests {
|
||||
assert!(updated_meta.index_file_size > 0);
|
||||
assert_eq!(updated_meta.file_id, file_meta.file_id);
|
||||
assert_eq!(updated_meta.index_version, 1);
|
||||
let field_0_index = updated_meta
|
||||
.indexes
|
||||
.iter()
|
||||
.find(|index| index.column_id == 2)
|
||||
.expect("field_0 should have an inverted index");
|
||||
assert_eq!(
|
||||
field_0_index.created_indexes.as_slice(),
|
||||
[IndexType::InvertedIndex]
|
||||
);
|
||||
}
|
||||
_ => panic!("Unexpected worker request: {:?}", worker_req),
|
||||
}
|
||||
|
||||
let puffin_reader = env
|
||||
.access_layer
|
||||
.build_puffin_manager()
|
||||
.reader(&RegionIndexId::new(
|
||||
RegionFileId::new(source_metadata.region_id, file_meta.file_id),
|
||||
1,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let blob = puffin_reader
|
||||
.blob(inverted_index::INDEX_BLOB_TYPE)
|
||||
.await
|
||||
.unwrap();
|
||||
let blob_reader = blob.reader().await.unwrap();
|
||||
let index_metadata =
|
||||
index::inverted_index::format::reader::InvertedIndexBlobReader::new(blob_reader)
|
||||
.metadata(None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(index_metadata.metas.contains_key("2"));
|
||||
assert_eq!(index_metadata.total_row_count, 100);
|
||||
}
|
||||
|
||||
async fn schedule_index_build_task_with_mode(build_mode: IndexBuildMode) {
|
||||
@@ -2236,6 +2300,7 @@ mod tests {
|
||||
let task = IndexBuildTask {
|
||||
region_id,
|
||||
file,
|
||||
target_region_metadata: version_control.current().version.metadata.clone(),
|
||||
source: IndexBuildSource::new(
|
||||
file_meta.clone(),
|
||||
version_control.current().version.metadata.schema_version,
|
||||
@@ -2346,6 +2411,7 @@ mod tests {
|
||||
let task = IndexBuildTask {
|
||||
region_id,
|
||||
file,
|
||||
target_region_metadata: version_control.current().version.metadata.clone(),
|
||||
source: IndexBuildSource::new(
|
||||
file_meta.clone(),
|
||||
version_control.current().version.metadata.schema_version,
|
||||
@@ -2444,6 +2510,7 @@ mod tests {
|
||||
let task = IndexBuildTask {
|
||||
region_id,
|
||||
file,
|
||||
target_region_metadata: version_control.current().version.metadata.clone(),
|
||||
source: IndexBuildSource::new(
|
||||
file_meta.clone(),
|
||||
version_control.current().version.metadata.schema_version,
|
||||
@@ -2505,7 +2572,7 @@ mod tests {
|
||||
let schema_version = metadata.schema_version;
|
||||
let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
|
||||
let file_purger = Arc::new(NoopFilePurger {});
|
||||
let indexer_builder = mock_indexer_builder(metadata, env).await;
|
||||
let indexer_builder = mock_indexer_builder(metadata.clone(), env).await;
|
||||
let (tx, _rx) = mpsc::channel(4);
|
||||
let (result_tx, result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);
|
||||
|
||||
@@ -2521,6 +2588,7 @@ mod tests {
|
||||
let task = IndexBuildTask {
|
||||
region_id,
|
||||
file,
|
||||
target_region_metadata: metadata,
|
||||
source: IndexBuildSource::new(file_meta, schema_version),
|
||||
reason,
|
||||
access_layer: env.access_layer.clone(),
|
||||
|
||||
+390
-77
@@ -14,11 +14,13 @@
|
||||
|
||||
//! SST in parquet format.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_base::readable_size::ReadableSize;
|
||||
use datatypes::json::JsonSettings;
|
||||
use parquet::file::metadata::ParquetMetaData;
|
||||
use store_api::storage::FileId;
|
||||
use store_api::storage::{ColumnId, FileId};
|
||||
|
||||
use crate::sst::DEFAULT_WRITE_BUFFER_SIZE;
|
||||
use crate::sst::file::FileTimeRange;
|
||||
@@ -48,6 +50,19 @@ pub const PARQUET_METADATA_KEY: &str = "greptime:metadata";
|
||||
/// default execution batch size to reduce rebatching and concatenation in the
|
||||
/// query pipeline.
|
||||
pub(crate) const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024;
|
||||
|
||||
/// JSON2 physical layouts requested by a compaction read.
|
||||
pub(crate) type Json2RewriteTargets = Arc<BTreeMap<ColumnId, Json2TargetLayout>>;
|
||||
|
||||
/// Fixed JSON2 physical layout used while rewriting compaction input.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct Json2TargetLayout {
|
||||
/// Logical JSON2 extension metadata attached to the rewritten field.
|
||||
pub(crate) extension_metadata: String,
|
||||
/// Settings used to build the fixed physical layout.
|
||||
pub(crate) target_layout: JsonSettings,
|
||||
}
|
||||
|
||||
/// Default row group size for parquet files.
|
||||
///
|
||||
/// Keep the existing persisted/on-disk default stable. It intentionally stays
|
||||
@@ -118,8 +133,8 @@ mod tests {
|
||||
use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit};
|
||||
use datatypes::arrow;
|
||||
use datatypes::arrow::array::{
|
||||
ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringArray, StringDictionaryBuilder,
|
||||
TimestampMillisecondArray, UInt8Array, UInt64Array,
|
||||
ArrayRef, AsArray, BinaryDictionaryBuilder, RecordBatch, StringArray,
|
||||
StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array,
|
||||
};
|
||||
use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type};
|
||||
use datatypes::arrow::util::pretty::pretty_format_batches;
|
||||
@@ -152,6 +167,7 @@ mod tests {
|
||||
use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
|
||||
use crate::sst::index::{IndexBuildType, Indexer, IndexerBuilder, IndexerBuilderImpl};
|
||||
use crate::sst::parquet::flat_format::FlatWriteFormat;
|
||||
use crate::sst::parquet::metadata::extract_primary_key_range;
|
||||
use crate::sst::parquet::reader::{ParquetReader, ParquetReaderBuilder, ReaderMetrics};
|
||||
use crate::sst::parquet::row_selection::RowGroupSelection;
|
||||
use crate::sst::parquet::writer::ParquetWriter;
|
||||
@@ -687,12 +703,13 @@ mod tests {
|
||||
let object_store = env.init_object_store_manager();
|
||||
let metadata = Arc::new(sst_region_metadata());
|
||||
let batches = vec![
|
||||
new_record_batch_by_range(&["a", "d"], 0, 1000),
|
||||
new_record_batch_by_range(&["b", "f"], 0, 1000),
|
||||
new_record_batch_by_range(&["c", "g"], 0, 1000),
|
||||
new_record_batch_by_range(&["b", "h"], 100, 200),
|
||||
new_record_batch_by_range(&["b", "h"], 200, 300),
|
||||
new_record_batch_by_range(&["b", "h"], 300, 1000),
|
||||
new_record_batch_by_range(&["a", "a"], 0, 1000),
|
||||
new_record_batch_by_range(&["b", "b"], 0, 1000),
|
||||
new_record_batch_by_range(&["c", "c"], 0, 1000),
|
||||
new_record_batch_by_range(&["d", "d"], 100, 200),
|
||||
new_record_batch_by_range(&["d", "d"], 200, 300),
|
||||
new_record_batch_by_range(&["d", "d"], 300, 1000),
|
||||
new_record_batch_by_range(&["e", "e"], 0, 100),
|
||||
];
|
||||
let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
|
||||
|
||||
@@ -745,6 +762,159 @@ mod tests {
|
||||
assert_eq!(total_rows, rows_read);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_split_file_at_series_boundary_inside_batch() {
|
||||
let mut env = TestEnv::new().await;
|
||||
let object_store = env.init_object_store_manager();
|
||||
let metadata = Arc::new(sst_region_metadata());
|
||||
let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
|
||||
let second_batch_rows = (0..1000).map(|ts| ("b", "b", ts)).collect::<Vec<_>>();
|
||||
let mut third_batch_rows = (1000..2000).map(|ts| ("b", "b", ts)).collect::<Vec<_>>();
|
||||
third_batch_rows.extend((0..1000).map(|ts| ("c", "c", ts)));
|
||||
let batches = vec![
|
||||
new_record_batch_from_rows(&first_batch_rows),
|
||||
new_record_batch_from_rows(&second_batch_rows),
|
||||
new_record_batch_from_rows(&third_batch_rows),
|
||||
];
|
||||
let total_rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
|
||||
let source = new_flat_source_from_record_batches(batches);
|
||||
let write_opts = WriteOptions {
|
||||
row_group_size: 50,
|
||||
max_file_size: Some(1),
|
||||
..Default::default()
|
||||
};
|
||||
let path_provider = RegionFilePathFactory {
|
||||
table_dir: "test_series_boundary".to_string(),
|
||||
path_type: PathType::Bare,
|
||||
};
|
||||
let mut metrics = Metrics::new(WriteType::Compaction);
|
||||
let mut writer = ParquetWriter::new_with_object_store(
|
||||
object_store,
|
||||
metadata.clone(),
|
||||
IndexConfig::default(),
|
||||
NoopIndexBuilder,
|
||||
path_provider,
|
||||
&mut metrics,
|
||||
)
|
||||
.await;
|
||||
|
||||
let files = writer
|
||||
.write_all_flat(source, None, &write_opts)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(files.len() > 1);
|
||||
assert_eq!(
|
||||
total_rows,
|
||||
files.iter().map(|file| file.num_rows).sum::<usize>()
|
||||
);
|
||||
let primary_key_ranges = files
|
||||
.iter()
|
||||
.map(|file| {
|
||||
extract_primary_key_range(file.file_metadata.as_ref().unwrap(), metadata.as_ref())
|
||||
.unwrap()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert!(
|
||||
primary_key_ranges
|
||||
.windows(2)
|
||||
.all(|ranges| ranges[0].1 < ranges[1].0)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oversized_single_series_stays_in_one_file() {
|
||||
let mut env = TestEnv::new().await;
|
||||
let object_store = env.init_object_store_manager();
|
||||
let metadata = Arc::new(sst_region_metadata());
|
||||
let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
|
||||
let second_batch_rows = (1000..2000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
|
||||
let third_batch_rows = (2000..3000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
|
||||
let batches = vec![
|
||||
new_record_batch_from_rows(&first_batch_rows),
|
||||
new_record_batch_from_rows(&second_batch_rows),
|
||||
new_record_batch_from_rows(&third_batch_rows),
|
||||
];
|
||||
let total_rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
|
||||
let source = new_flat_source_from_record_batches(batches);
|
||||
let write_opts = WriteOptions {
|
||||
row_group_size: 50,
|
||||
max_file_size: Some(1),
|
||||
..Default::default()
|
||||
};
|
||||
let path_provider = RegionFilePathFactory {
|
||||
table_dir: "test_oversized_series".to_string(),
|
||||
path_type: PathType::Bare,
|
||||
};
|
||||
let mut metrics = Metrics::new(WriteType::Compaction);
|
||||
let mut writer = ParquetWriter::new_with_object_store(
|
||||
object_store,
|
||||
metadata,
|
||||
IndexConfig::default(),
|
||||
NoopIndexBuilder,
|
||||
path_provider,
|
||||
&mut metrics,
|
||||
)
|
||||
.await;
|
||||
|
||||
let files = writer
|
||||
.write_all_flat_as_primary_key(source, None, &write_opts)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(1, files.len());
|
||||
assert_eq!(total_rows, files[0].num_rows);
|
||||
assert!(files[0].file_size > write_opts.max_file_size.unwrap() as u64);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_write_multiple_files_without_primary_key() {
|
||||
let mut env = TestEnv::new().await;
|
||||
let object_store = env.init_object_store_manager();
|
||||
let metadata = Arc::new(sst_region_metadata_without_primary_key());
|
||||
let batch_rows = 1000;
|
||||
let batches = vec![
|
||||
new_record_batch_without_primary_key(0, batch_rows),
|
||||
new_record_batch_without_primary_key(batch_rows, 2 * batch_rows),
|
||||
new_record_batch_without_primary_key(2 * batch_rows, 3 * batch_rows),
|
||||
];
|
||||
let total_rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
|
||||
let source = new_flat_source_from_record_batches(batches);
|
||||
let write_opts = WriteOptions {
|
||||
row_group_size: 50,
|
||||
max_file_size: Some(1),
|
||||
..Default::default()
|
||||
};
|
||||
let path_provider = RegionFilePathFactory {
|
||||
table_dir: "test_no_primary_key".to_string(),
|
||||
path_type: PathType::Bare,
|
||||
};
|
||||
let mut metrics = Metrics::new(WriteType::Compaction);
|
||||
let mut writer = ParquetWriter::new_with_object_store(
|
||||
object_store,
|
||||
metadata,
|
||||
IndexConfig::default(),
|
||||
NoopIndexBuilder,
|
||||
path_provider,
|
||||
&mut metrics,
|
||||
)
|
||||
.await;
|
||||
|
||||
let files = writer
|
||||
.write_all_flat(source, None, &write_opts)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Regions without a primary key keep splitting at batch boundaries:
|
||||
// the limit produces multiple files and no batch is sliced.
|
||||
assert!(files.len() > 1);
|
||||
assert!(files.iter().all(|file| file.num_rows % batch_rows == 0));
|
||||
assert_eq!(
|
||||
total_rows,
|
||||
files.iter().map(|file| file.num_rows).sum::<usize>()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_write_read_with_index() {
|
||||
let mut env = TestEnv::new().await;
|
||||
@@ -1145,6 +1315,62 @@ mod tests {
|
||||
assert!(reader.next_record_batch().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
/// Creates a new region metadata without primary key for testing SSTs.
|
||||
///
|
||||
/// Schema: field_0, ts
|
||||
fn sst_region_metadata_without_primary_key() -> RegionMetadata {
|
||||
let mut builder = RegionMetadataBuilder::new(REGION_ID);
|
||||
builder
|
||||
.push_column_metadata(ColumnMetadata {
|
||||
column_schema: ColumnSchema::new(
|
||||
"field_0".to_string(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
true,
|
||||
),
|
||||
semantic_type: SemanticType::Field,
|
||||
column_id: 0,
|
||||
})
|
||||
.push_column_metadata(ColumnMetadata {
|
||||
column_schema: ColumnSchema::new(
|
||||
"ts".to_string(),
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
semantic_type: SemanticType::Timestamp,
|
||||
column_id: 1,
|
||||
})
|
||||
.primary_key(vec![]);
|
||||
builder.build().unwrap()
|
||||
}
|
||||
|
||||
/// Creates a flat format RecordBatch for regions without a primary key.
|
||||
fn new_record_batch_without_primary_key(start: usize, end: usize) -> RecordBatch {
|
||||
assert!(end >= start);
|
||||
let metadata = Arc::new(sst_region_metadata_without_primary_key());
|
||||
let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
|
||||
|
||||
let num_rows = end - start;
|
||||
let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::new();
|
||||
// Regions without a primary key encode it as empty bytes.
|
||||
for _ in 0..num_rows {
|
||||
pk_builder.append([]).unwrap();
|
||||
}
|
||||
|
||||
RecordBatch::try_new(
|
||||
flat_schema,
|
||||
vec![
|
||||
Arc::new(UInt64Array::from_iter_values(start as u64..end as u64)) as ArrayRef,
|
||||
Arc::new(TimestampMillisecondArray::from_iter_values(
|
||||
start as i64..end as i64,
|
||||
)) as ArrayRef,
|
||||
Arc::new(pk_builder.finish()) as ArrayRef,
|
||||
Arc::new(UInt64Array::from_value(1000, num_rows)) as ArrayRef,
|
||||
Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)) as ArrayRef,
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn new_record_batch_from_rows(rows: &[(&str, &str, i64)]) -> RecordBatch {
|
||||
let metadata = Arc::new(sst_region_metadata());
|
||||
let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
|
||||
@@ -1414,91 +1640,178 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_with_override_sequence() {
|
||||
test_read_with_override_sequence_with_format(false).await;
|
||||
test_read_with_override_sequence_with_format(true).await;
|
||||
}
|
||||
|
||||
async fn test_read_with_override_sequence_with_format(flat_format: bool) {
|
||||
let mut env = TestEnv::new().await;
|
||||
let object_store = env.init_object_store_manager();
|
||||
let handle = sst_file_handle(0, 1000);
|
||||
let file_path = FixedPathProvider {
|
||||
region_file_id: handle.file_id(),
|
||||
};
|
||||
let metadata = Arc::new(sst_region_metadata());
|
||||
|
||||
// Create batches with sequence 0 to trigger override functionality.
|
||||
let source = new_flat_source_from_record_batches(vec![
|
||||
new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, 0),
|
||||
new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, 0),
|
||||
]);
|
||||
async fn read_sequences(builder: ParquetReaderBuilder) -> Vec<u64> {
|
||||
let mut reader = builder.build().await.unwrap().unwrap();
|
||||
let mut sequences = Vec::new();
|
||||
while let Some(batch) = reader.next_record_batch().await.unwrap() {
|
||||
let sequence = batch
|
||||
.column(batch.num_columns() - 2)
|
||||
.as_primitive::<datatypes::arrow::datatypes::UInt64Type>();
|
||||
sequences.extend((0..sequence.len()).map(|idx| sequence.value(idx)));
|
||||
}
|
||||
sequences
|
||||
}
|
||||
|
||||
let write_opts = WriteOptions {
|
||||
row_group_size: 50,
|
||||
..Default::default()
|
||||
};
|
||||
async fn write_sst(
|
||||
object_store: ObjectStore,
|
||||
metadata: Arc<RegionMetadata>,
|
||||
handle: FileHandle,
|
||||
flat_format: bool,
|
||||
sequence: u64,
|
||||
) {
|
||||
let file_path = FixedPathProvider {
|
||||
region_file_id: handle.file_id(),
|
||||
};
|
||||
let source = new_flat_source_from_record_batches(vec![
|
||||
new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, sequence),
|
||||
new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, sequence),
|
||||
]);
|
||||
let write_opts = WriteOptions {
|
||||
row_group_size: 50,
|
||||
..Default::default()
|
||||
};
|
||||
let mut metrics = Metrics::new(WriteType::Flush);
|
||||
let mut writer = ParquetWriter::new_with_object_store(
|
||||
object_store,
|
||||
metadata,
|
||||
IndexConfig::default(),
|
||||
NoopIndexBuilder,
|
||||
file_path,
|
||||
&mut metrics,
|
||||
)
|
||||
.await;
|
||||
if flat_format {
|
||||
writer
|
||||
.write_all_flat(source, None, &write_opts)
|
||||
.await
|
||||
.unwrap();
|
||||
} else {
|
||||
writer
|
||||
.write_all_flat_as_primary_key(source, None, &write_opts)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
let mut metrics = Metrics::new(WriteType::Flush);
|
||||
let mut writer = ParquetWriter::new_with_object_store(
|
||||
let custom_sequence = 12345u64;
|
||||
let local_zero_handle = sst_file_handle(0, 1000);
|
||||
write_sst(
|
||||
object_store.clone(),
|
||||
metadata.clone(),
|
||||
IndexConfig::default(),
|
||||
NoopIndexBuilder,
|
||||
file_path,
|
||||
&mut metrics,
|
||||
local_zero_handle.clone(),
|
||||
flat_format,
|
||||
0,
|
||||
)
|
||||
.await;
|
||||
|
||||
writer
|
||||
.write_all_flat_as_primary_key(source, None, &write_opts)
|
||||
.await
|
||||
.unwrap()
|
||||
.remove(0);
|
||||
// Local all-zero SSTs retain the compatibility override.
|
||||
let local_zero_none = read_sequences(
|
||||
ParquetReaderBuilder::new(
|
||||
FILE_DIR.to_string(),
|
||||
PathType::Bare,
|
||||
local_zero_handle.clone(),
|
||||
object_store.clone(),
|
||||
)
|
||||
.expected_metadata(Some(metadata.clone())),
|
||||
)
|
||||
.await;
|
||||
assert!(local_zero_none.iter().all(|sequence| *sequence == 0));
|
||||
|
||||
// Read without override sequence (should read sequence 0)
|
||||
let builder = ParquetReaderBuilder::new(
|
||||
FILE_DIR.to_string(),
|
||||
PathType::Bare,
|
||||
handle.clone(),
|
||||
object_store.clone(),
|
||||
);
|
||||
let mut reader = builder.build().await.unwrap().unwrap();
|
||||
let mut normal_batches = Vec::new();
|
||||
while let Some(batch) = reader.next_record_batch().await.unwrap() {
|
||||
normal_batches.push(batch);
|
||||
}
|
||||
|
||||
// Read with override sequence using FileMeta.sequence
|
||||
let custom_sequence = 12345u64;
|
||||
let file_meta = handle.meta_ref();
|
||||
let mut override_file_meta = file_meta.clone();
|
||||
override_file_meta.sequence = Some(std::num::NonZero::new(custom_sequence).unwrap());
|
||||
let override_handle = FileHandle::new(
|
||||
override_file_meta,
|
||||
let mut local_zero_meta = local_zero_handle.meta_ref().clone();
|
||||
local_zero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap());
|
||||
let local_zero_override_handle = FileHandle::new(
|
||||
local_zero_meta,
|
||||
Arc::new(crate::sst::file_purger::NoopFilePurger),
|
||||
);
|
||||
|
||||
let builder = ParquetReaderBuilder::new(
|
||||
FILE_DIR.to_string(),
|
||||
PathType::Bare,
|
||||
override_handle,
|
||||
object_store.clone(),
|
||||
let local_zero_override = read_sequences(
|
||||
ParquetReaderBuilder::new(
|
||||
FILE_DIR.to_string(),
|
||||
PathType::Bare,
|
||||
local_zero_override_handle,
|
||||
object_store.clone(),
|
||||
)
|
||||
.expected_metadata(Some(metadata.clone())),
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
local_zero_override
|
||||
.iter()
|
||||
.all(|sequence| *sequence == custom_sequence)
|
||||
);
|
||||
let mut reader = builder.build().await.unwrap().unwrap();
|
||||
let mut override_batches = Vec::new();
|
||||
while let Some(batch) = reader.next_record_batch().await.unwrap() {
|
||||
override_batches.push(batch);
|
||||
}
|
||||
|
||||
// Compare the results
|
||||
assert_eq!(normal_batches.len(), override_batches.len());
|
||||
for (normal, override_batch) in normal_batches.into_iter().zip(override_batches.iter()) {
|
||||
let expected_batch = {
|
||||
let mut columns = normal.columns().to_vec();
|
||||
let num_cols = columns.len();
|
||||
columns[num_cols - 2] =
|
||||
Arc::new(UInt64Array::from_value(custom_sequence, normal.num_rows()));
|
||||
RecordBatch::try_new(normal.schema(), columns).unwrap()
|
||||
};
|
||||
let local_nonzero_handle = sst_file_handle(0, 1000);
|
||||
write_sst(
|
||||
object_store.clone(),
|
||||
metadata.clone(),
|
||||
local_nonzero_handle.clone(),
|
||||
flat_format,
|
||||
7,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Override batch should match expected batch
|
||||
assert_eq!(*override_batch, expected_batch);
|
||||
}
|
||||
// Local nonzero SSTs retain physical per-row sequences, even with FileMeta.sequence.
|
||||
let mut local_nonzero_meta = local_nonzero_handle.meta_ref().clone();
|
||||
local_nonzero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap());
|
||||
let local_nonzero_override_handle = FileHandle::new(
|
||||
local_nonzero_meta,
|
||||
Arc::new(crate::sst::file_purger::NoopFilePurger),
|
||||
);
|
||||
let local_nonzero_override = read_sequences(
|
||||
ParquetReaderBuilder::new(
|
||||
FILE_DIR.to_string(),
|
||||
PathType::Bare,
|
||||
local_nonzero_override_handle,
|
||||
object_store.clone(),
|
||||
)
|
||||
.expected_metadata(Some(metadata.clone())),
|
||||
)
|
||||
.await;
|
||||
assert!(local_nonzero_override.iter().all(|sequence| *sequence == 7));
|
||||
|
||||
// None never overrides a local nonzero physical sequence.
|
||||
let local_nonzero_none = read_sequences(
|
||||
ParquetReaderBuilder::new(
|
||||
FILE_DIR.to_string(),
|
||||
PathType::Bare,
|
||||
local_nonzero_handle.clone(),
|
||||
object_store.clone(),
|
||||
)
|
||||
.expected_metadata(Some(metadata.clone())),
|
||||
)
|
||||
.await;
|
||||
assert!(local_nonzero_none.iter().all(|sequence| *sequence == 7));
|
||||
|
||||
// A source-owned handle is foreign when read against target metadata, so the
|
||||
// target-local manifest barrier is applied even for nonzero physical sequences.
|
||||
let mut target_metadata = (*metadata).clone();
|
||||
target_metadata.region_id = RegionId::new(0, 1);
|
||||
let target_metadata = Arc::new(target_metadata);
|
||||
let mut foreign_meta = local_nonzero_handle.meta_ref().clone();
|
||||
foreign_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap());
|
||||
let foreign_handle = FileHandle::new(
|
||||
foreign_meta,
|
||||
Arc::new(crate::sst::file_purger::NoopFilePurger),
|
||||
);
|
||||
let foreign = read_sequences(
|
||||
ParquetReaderBuilder::new(
|
||||
FILE_DIR.to_string(),
|
||||
PathType::Bare,
|
||||
foreign_handle,
|
||||
object_store,
|
||||
)
|
||||
.expected_metadata(Some(target_metadata)),
|
||||
)
|
||||
.await;
|
||||
assert!(foreign.iter().all(|sequence| *sequence == custom_sequence));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -241,9 +241,8 @@ impl FlatReadFormat {
|
||||
}
|
||||
|
||||
/// Enables wrapping binary `__primary_key` batches back to a dictionary in [`Self::convert_batch`].
|
||||
pub(crate) fn set_pk_as_binary(&mut self) -> Result<()> {
|
||||
self.pk_dict_wrap_schema = Some(self.output_arrow_schema()?);
|
||||
Ok(())
|
||||
pub(crate) fn set_pk_as_binary(&mut self, output_schema: SchemaRef) {
|
||||
self.pk_dict_wrap_schema = Some(output_schema);
|
||||
}
|
||||
|
||||
/// Index of a column in the projected batch by its column id.
|
||||
@@ -306,19 +305,16 @@ impl FlatReadFormat {
|
||||
.project(projection)
|
||||
.context(ComputeArrowSnafu)?;
|
||||
let mut fields = schema.fields().iter().cloned().collect::<Vec<_>>();
|
||||
for (column_id, target_type) in self.json_target_types().iter() {
|
||||
for (column_id, target) in self.json_target_types().iter() {
|
||||
let Some(index) = self.parquet_projected_index_by_id(*column_id) else {
|
||||
continue;
|
||||
};
|
||||
let Some(field) = schema.fields().get(index) else {
|
||||
continue;
|
||||
};
|
||||
fields[index] = Arc::new(
|
||||
field
|
||||
.as_ref()
|
||||
.clone()
|
||||
.with_data_type(ConcreteDataType::json2(target_type.clone()).as_arrow_type()),
|
||||
);
|
||||
let mut field = field.as_ref().clone();
|
||||
field.set_data_type(ConcreteDataType::json2(target.clone()).as_arrow_type());
|
||||
fields[index] = Arc::new(field);
|
||||
}
|
||||
schema.fields = fields.into();
|
||||
Ok(Arc::new(schema))
|
||||
@@ -326,7 +322,7 @@ impl FlatReadFormat {
|
||||
|
||||
/// Index of a column in the projected schema produced directly by parquet
|
||||
/// reading, before any primary-key-to-flat conversion.
|
||||
fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option<usize> {
|
||||
pub(crate) fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option<usize> {
|
||||
match &self.parquet_adapter {
|
||||
ParquetAdapter::Flat(p) => p
|
||||
.format_projection
|
||||
@@ -359,7 +355,7 @@ impl FlatReadFormat {
|
||||
}
|
||||
}
|
||||
|
||||
/// Gets JSON2 target types keyed by column id.
|
||||
/// Gets JSON2 read targets.
|
||||
pub(crate) fn json_target_types(&self) -> &JsonTargetTypes {
|
||||
self.read_cols.json_target_types()
|
||||
}
|
||||
@@ -1012,9 +1008,8 @@ mod tests {
|
||||
false,
|
||||
)
|
||||
.unwrap();
|
||||
read_format.set_pk_as_binary().unwrap();
|
||||
|
||||
let output_schema = read_format.output_arrow_schema().unwrap();
|
||||
read_format.set_pk_as_binary(output_schema.clone());
|
||||
let binary_schema = override_pk_field_to_binary(&output_schema);
|
||||
|
||||
// The __primary_key field must preserve its field_id metadata after
|
||||
|
||||
@@ -53,7 +53,7 @@ use store_api::storage::{ColumnId, NestedPath, SequenceNumber};
|
||||
use crate::error::{
|
||||
ConvertVectorSnafu, DecodeSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result,
|
||||
};
|
||||
use crate::read::read_columns::{JsonTargetTypes, ReadColumns};
|
||||
use crate::read::read_columns::ReadColumns;
|
||||
use crate::read::{Batch, BatchBuilder, BatchColumn};
|
||||
use crate::sst::file::{FileMeta, FileTimeRange};
|
||||
use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns};
|
||||
@@ -599,14 +599,13 @@ impl FormatProjection {
|
||||
sst_column_num: usize,
|
||||
cols: ReadColumns,
|
||||
) -> Self {
|
||||
let json_target_types = cols.json_target_types().clone();
|
||||
let mut projected_columns: Vec<_> = cols
|
||||
.col_ids
|
||||
.into_iter()
|
||||
.iter()
|
||||
.copied()
|
||||
.filter_map(|col_id| {
|
||||
id_to_index.get(&col_id).copied().map(|index_of_sst| {
|
||||
let nested_paths =
|
||||
json_target_nested_paths(metadata, &json_target_types, col_id);
|
||||
let nested_paths = json_target_nested_paths(metadata, &cols, col_id);
|
||||
(col_id, index_of_sst, nested_paths)
|
||||
})
|
||||
})
|
||||
@@ -688,10 +687,10 @@ impl FormatProjection {
|
||||
|
||||
fn json_target_nested_paths(
|
||||
metadata: &RegionMetadataRef,
|
||||
json_target_types: &JsonTargetTypes,
|
||||
read_columns: &ReadColumns,
|
||||
column_id: ColumnId,
|
||||
) -> Vec<NestedPath> {
|
||||
let Some(target_type) = json_target_types.get(&column_id) else {
|
||||
let Some(target_type) = read_columns.json_target_type(column_id) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let Some(column) = metadata.column_by_id(column_id) else {
|
||||
|
||||
@@ -12,15 +12,17 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use datafusion_common::cast_column;
|
||||
use datafusion_common::format::DEFAULT_CAST_OPTIONS;
|
||||
use datatypes::arrow::array::{ArrayRef, new_null_array};
|
||||
use datatypes::arrow::datatypes::{DataType, FieldRef, SchemaRef};
|
||||
use datatypes::arrow::datatypes::{DataType, Field, FieldRef, SchemaRef};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::extension::json::is_json2_extension_type;
|
||||
use datatypes::extension::json::{JsonMetadata, is_json2_extension_type};
|
||||
use datatypes::json::JsonSettings;
|
||||
use datatypes::vectors::json::array::JsonArray;
|
||||
use futures::Stream;
|
||||
use snafu::{ResultExt, ensure};
|
||||
@@ -28,6 +30,13 @@ use snafu::{ResultExt, ensure};
|
||||
use crate::error::{
|
||||
CastColumnSnafu, DataTypeMismatchSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu,
|
||||
};
|
||||
use crate::sst::parquet::Json2TargetLayout;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Json2RewriteSettings {
|
||||
logical_settings: JsonSettings,
|
||||
target_layout: JsonSettings,
|
||||
}
|
||||
|
||||
/// Aligns projected batches to the expected output schema for nested projections.
|
||||
///
|
||||
@@ -61,6 +70,8 @@ pub struct NestedSchemaAligner<S> {
|
||||
/// Whether all projected roots are present and the stream can pass batches
|
||||
/// through.
|
||||
all_roots_present: bool,
|
||||
/// JSON2 columns that require semantic source-to-target layout rewriting.
|
||||
json2_rewrite_targets: HashMap<String, Json2RewriteSettings>,
|
||||
/// The cache for whether incoming batches already match output schema.
|
||||
is_schema_matched: Option<bool>,
|
||||
}
|
||||
@@ -96,9 +107,39 @@ where
|
||||
projected_root_presence,
|
||||
expected_input_col_num,
|
||||
all_roots_present,
|
||||
json2_rewrite_targets: HashMap::new(),
|
||||
is_schema_matched: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Sets JSON2 columns that must be rewritten into the output field layout.
|
||||
pub(crate) fn with_json2_rewrite_targets(
|
||||
mut self,
|
||||
targets: &HashMap<String, Json2TargetLayout>,
|
||||
) -> Result<Self> {
|
||||
self.json2_rewrite_targets = targets
|
||||
.iter()
|
||||
.map(|(name, layout)| {
|
||||
let metadata = serde_json::from_str::<JsonMetadata>(&layout.extension_metadata)
|
||||
.map_err(|e| {
|
||||
UnexpectedSnafu {
|
||||
reason: format!(
|
||||
"invalid JSON2 extension metadata for column '{name}': {e}"
|
||||
),
|
||||
}
|
||||
.build()
|
||||
})?;
|
||||
Ok((
|
||||
name.clone(),
|
||||
Json2RewriteSettings {
|
||||
logical_settings: metadata.into_json_settings(),
|
||||
target_layout: layout.target_layout.clone(),
|
||||
},
|
||||
))
|
||||
})
|
||||
.collect::<Result<_>>()?;
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Stream for NestedSchemaAligner<S>
|
||||
@@ -125,6 +166,7 @@ where
|
||||
&this.output_schema,
|
||||
&this.projected_root_presence,
|
||||
this.expected_input_col_num,
|
||||
&this.json2_rewrite_targets,
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -140,6 +182,7 @@ fn align_projected_batch(
|
||||
output_schema: &SchemaRef,
|
||||
projected_root_presence: &[bool],
|
||||
expected_input_col_num: usize,
|
||||
json2_rewrite_targets: &HashMap<String, Json2RewriteSettings>,
|
||||
) -> Result<RecordBatch> {
|
||||
ensure!(
|
||||
rb.columns().len() == expected_input_col_num,
|
||||
@@ -154,6 +197,7 @@ fn align_projected_batch(
|
||||
|
||||
let mut cols = Vec::with_capacity(projected_root_presence.len());
|
||||
let mut idx = 0;
|
||||
let input_schema = rb.schema_ref();
|
||||
|
||||
for (field, present) in output_schema.fields().iter().zip(projected_root_presence) {
|
||||
if !present {
|
||||
@@ -161,19 +205,39 @@ fn align_projected_batch(
|
||||
continue;
|
||||
}
|
||||
|
||||
cols.push(align_array(rb.column(idx), field)?);
|
||||
cols.push(align_array(
|
||||
rb.column(idx),
|
||||
input_schema.field(idx),
|
||||
field,
|
||||
json2_rewrite_targets.get(field.name()),
|
||||
)?);
|
||||
idx += 1;
|
||||
}
|
||||
|
||||
RecordBatch::try_new(output_schema.clone(), cols).context(NewRecordBatchSnafu)
|
||||
}
|
||||
|
||||
fn align_array(array: &ArrayRef, field: &FieldRef) -> Result<ArrayRef> {
|
||||
fn align_array(
|
||||
array: &ArrayRef,
|
||||
source: &Field,
|
||||
field: &FieldRef,
|
||||
rewrite_settings: Option<&Json2RewriteSettings>,
|
||||
) -> Result<ArrayRef> {
|
||||
if let Some(settings) = rewrite_settings {
|
||||
return JsonArray::from(array)
|
||||
.rewrite_to_v2(source, &settings.logical_settings, &settings.target_layout)
|
||||
.context(DataTypeMismatchSnafu);
|
||||
}
|
||||
if array.data_type() == field.data_type() {
|
||||
return Ok(array.clone());
|
||||
}
|
||||
|
||||
if is_json2_extension_type(field) {
|
||||
if is_json2_extension_type(source) {
|
||||
return JsonArray::from(array)
|
||||
.project_to_v2(source, field.data_type())
|
||||
.context(DataTypeMismatchSnafu);
|
||||
}
|
||||
return JsonArray::from(array)
|
||||
.project_to(field.data_type())
|
||||
.context(DataTypeMismatchSnafu);
|
||||
@@ -188,6 +252,7 @@ fn align_array(array: &ArrayRef, field: &FieldRef) -> Result<ArrayRef> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use datatypes::arrow::array::{
|
||||
@@ -200,6 +265,33 @@ mod tests {
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_aligner_resolves_json2_rewrite_settings()
|
||||
-> std::result::Result<(), Box<dyn std::error::Error>> {
|
||||
let logical_settings = JsonSettings::default();
|
||||
let target_layout = JsonSettings::try_new(vec![], Some(0))?;
|
||||
let rewrite_targets = HashMap::from([(
|
||||
"j".to_string(),
|
||||
Json2TargetLayout {
|
||||
extension_metadata: serde_json::to_string(&JsonMetadata::new(
|
||||
logical_settings.clone(),
|
||||
))?,
|
||||
target_layout: target_layout.clone(),
|
||||
},
|
||||
)]);
|
||||
let aligner = NestedSchemaAligner::new(
|
||||
stream::empty::<Result<RecordBatch>>(),
|
||||
vec![],
|
||||
schema(Vec::<Field>::new()),
|
||||
)?
|
||||
.with_json2_rewrite_targets(&rewrite_targets)?;
|
||||
|
||||
let settings = &aligner.json2_rewrite_targets["j"];
|
||||
assert_eq!(logical_settings, settings.logical_settings);
|
||||
assert_eq!(target_layout, settings.target_layout);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_aligner_with_all_projected_roots_match() {
|
||||
let output_schema = schema([
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME;
|
||||
use parquet::arrow::ProjectionMask;
|
||||
use parquet::basic::{ConvertedType, Type as PhysicalType};
|
||||
use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor};
|
||||
@@ -179,7 +180,7 @@ pub struct ProjectionMaskPlan {
|
||||
/// returned plan keeps `k` in the projection mask and marks `j` as
|
||||
/// not present in the output, so it can be synthesized during
|
||||
/// post-processing.
|
||||
pub fn build_projection_plan(
|
||||
pub(crate) fn build_projection_plan(
|
||||
parquet_read_cols: &ParquetReadColumns,
|
||||
parquet_schema_desc: &SchemaDescriptor,
|
||||
) -> ProjectionMaskPlan {
|
||||
@@ -261,24 +262,38 @@ fn build_parquet_leaves_indices(
|
||||
}
|
||||
}
|
||||
|
||||
// Then fallback prefix misses to their nearest variant parent.
|
||||
// Then include v2 remainder leaves or fallback prefix misses to their nearest variant parent.
|
||||
// TODO(fys): Gate fallback planning on the root being JSON2. A raw Binary
|
||||
// leaf is a JSONB variant only under a JSON2 root; plain struct Binary
|
||||
// children should not enter this fallback path.
|
||||
for col in &projection.cols {
|
||||
for (path_idx, nested_path) in col.nested_paths.iter().enumerate() {
|
||||
if prefix_matched[&col.root_index][path_idx] {
|
||||
let path_matches = &prefix_matched[&col.root_index];
|
||||
let mut needs_remainder = false;
|
||||
for (matched, nested_path) in path_matches.iter().zip(&col.nested_paths) {
|
||||
if *matched {
|
||||
if !needs_remainder {
|
||||
needs_remainder =
|
||||
path_points_to_struct(parquet_schema_desc, col.root_index, nested_path);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(leaf_idx) =
|
||||
if let Some(leaf_idx) =
|
||||
find_nearest_variant_parent(parquet_schema_desc, col.root_index, nested_path)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
{
|
||||
matched_leaves.insert(leaf_idx);
|
||||
matched_roots.insert(col.root_index);
|
||||
} else {
|
||||
needs_remainder = true;
|
||||
}
|
||||
}
|
||||
|
||||
matched_leaves.insert(leaf_idx);
|
||||
matched_roots.insert(col.root_index);
|
||||
if needs_remainder {
|
||||
let remainder_leaves = find_remainder_leaves(parquet_schema_desc, col.root_index);
|
||||
if !remainder_leaves.is_empty() {
|
||||
matched_leaves.extend(remainder_leaves);
|
||||
matched_roots.insert(col.root_index);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -287,6 +302,49 @@ fn build_parquet_leaves_indices(
|
||||
(matched_leaves, matched_roots)
|
||||
}
|
||||
|
||||
/// Returns whether a nested path points to an explicitly materialized object.
|
||||
///
|
||||
/// JSON2 v2 can split an object's children between its Struct field and the remainder,
|
||||
/// so reading the Struct leaves alone may produce an incomplete object.
|
||||
fn path_points_to_struct(
|
||||
parquet_schema_desc: &SchemaDescriptor,
|
||||
root_idx: usize,
|
||||
path: &[String],
|
||||
) -> bool {
|
||||
let Some(mut field) = parquet_schema_desc.root_schema().get_fields().get(root_idx) else {
|
||||
return false;
|
||||
};
|
||||
for name in path.iter().skip(1) {
|
||||
if !field.is_group() {
|
||||
return false;
|
||||
}
|
||||
let Some(child) = field.get_fields().iter().find(|field| field.name() == name) else {
|
||||
return false;
|
||||
};
|
||||
field = child;
|
||||
}
|
||||
field.is_group()
|
||||
}
|
||||
|
||||
/// Finds the Parquet leaves backing a JSON2 v2 remainder field.
|
||||
///
|
||||
/// The remainder is a sibling of explicitly materialized fields, so prefix matching a
|
||||
/// requested path cannot find it. These leaves are needed when an explicit path is absent
|
||||
/// or an explicitly materialized object may have additional children in the remainder.
|
||||
fn find_remainder_leaves(parquet_schema_desc: &SchemaDescriptor, root_idx: usize) -> Vec<usize> {
|
||||
parquet_schema_desc
|
||||
.columns()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(i, column)| {
|
||||
let path = column.path().parts();
|
||||
(parquet_schema_desc.get_column_root_idx(i) == root_idx
|
||||
&& path.get(1).is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME))
|
||||
.then_some(i)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
}
|
||||
|
||||
fn find_nearest_variant_parent(
|
||||
parquet_schema_desc: &SchemaDescriptor,
|
||||
root_idx: usize,
|
||||
@@ -329,6 +387,7 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use parquet::basic::{ConvertedType, LogicalType, Repetition};
|
||||
use parquet::errors::ParquetError;
|
||||
use parquet::schema::types::Type;
|
||||
|
||||
use super::*;
|
||||
@@ -441,6 +500,76 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_v2_routes_missing_path_to_remainder() -> Result<(), ParquetError> {
|
||||
let parquet = build_test_v2_schema()?;
|
||||
let projection =
|
||||
ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
|
||||
vec![
|
||||
vec!["j".to_string(), "cold".to_string()],
|
||||
vec!["j".to_string(), "another".to_string()],
|
||||
],
|
||||
)]);
|
||||
|
||||
let plan = build_projection_plan(&projection, &parquet);
|
||||
|
||||
assert_eq!(vec![true], plan.projected_root_presence);
|
||||
assert_eq!(ProjectionMask::leaves(&parquet, [0, 1]), plan.mask);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_v2_explicit_path_does_not_read_remainder() -> Result<(), ParquetError> {
|
||||
let parquet = build_test_v2_schema()?;
|
||||
let projection = ParquetReadColumns::from_deduped(vec![
|
||||
ParquetReadColumn::new(0)
|
||||
.with_nested_paths(vec![vec!["j".to_string(), "hot".to_string()]]),
|
||||
]);
|
||||
|
||||
let plan = build_projection_plan(&projection, &parquet);
|
||||
|
||||
assert_eq!(vec![true], plan.projected_root_presence);
|
||||
assert_eq!(ProjectionMask::leaves(&parquet, [3]), plan.mask);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_v2_container_path_reads_remainder() -> Result<(), ParquetError> {
|
||||
let parquet = build_test_v2_schema()?;
|
||||
let projection = ParquetReadColumns::from_deduped(vec![
|
||||
ParquetReadColumn::new(0)
|
||||
.with_nested_paths(vec![vec!["j".to_string(), "commit".to_string()]]),
|
||||
]);
|
||||
|
||||
let plan = build_projection_plan(&projection, &parquet);
|
||||
|
||||
assert_eq!(vec![true], plan.projected_root_presence);
|
||||
assert_eq!(ProjectionMask::leaves(&parquet, [0, 1, 2]), plan.mask);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// A nested path under an explicit Variant is stored entirely in that Variant parent. The
|
||||
// remainder may preserve `opaque: null`, but it cannot contain `opaque.leaf`, so reading the
|
||||
// nearest Variant parent is sufficient.
|
||||
#[test]
|
||||
fn test_v2_variant_parent_path_reads_parent() -> Result<(), ParquetError> {
|
||||
let parquet = build_test_v2_schema()?;
|
||||
let projection =
|
||||
ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
|
||||
vec![vec![
|
||||
"j".to_string(),
|
||||
"opaque".to_string(),
|
||||
"leaf".to_string(),
|
||||
]],
|
||||
)]);
|
||||
|
||||
let plan = build_projection_plan(&projection, &parquet);
|
||||
|
||||
assert_eq!(vec![true], plan.projected_root_presence);
|
||||
assert_eq!(ProjectionMask::leaves(&parquet, [4]), plan.mask);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merges_mixed_paths() {
|
||||
let parquet_schema_desc = build_test_nested_parquet_schema();
|
||||
@@ -673,6 +802,63 @@ mod tests {
|
||||
SchemaDescriptor::new(schema)
|
||||
}
|
||||
|
||||
fn build_test_v2_schema() -> Result<SchemaDescriptor, ParquetError> {
|
||||
let metadata = Arc::new(
|
||||
Type::primitive_type_builder("metadata", parquet::basic::Type::BYTE_ARRAY)
|
||||
.with_repetition(Repetition::REQUIRED)
|
||||
.build()?,
|
||||
);
|
||||
let value = Arc::new(
|
||||
Type::primitive_type_builder("value", parquet::basic::Type::BYTE_ARRAY)
|
||||
.with_repetition(Repetition::REQUIRED)
|
||||
.build()?,
|
||||
);
|
||||
let remainder = Arc::new(
|
||||
Type::group_type_builder(JSON2_REMAINDER_FIELD_NAME)
|
||||
.with_repetition(Repetition::OPTIONAL)
|
||||
.with_logical_type(Some(LogicalType::Variant {
|
||||
specification_version: None,
|
||||
}))
|
||||
.with_fields(vec![metadata, value])
|
||||
.build()?,
|
||||
);
|
||||
let operation = Arc::new(
|
||||
Type::primitive_type_builder("operation", parquet::basic::Type::INT64)
|
||||
.with_repetition(Repetition::OPTIONAL)
|
||||
.build()?,
|
||||
);
|
||||
let commit = Arc::new(
|
||||
Type::group_type_builder("commit")
|
||||
.with_repetition(Repetition::OPTIONAL)
|
||||
.with_fields(vec![operation])
|
||||
.build()?,
|
||||
);
|
||||
let hot = Arc::new(
|
||||
Type::primitive_type_builder("hot", parquet::basic::Type::INT64)
|
||||
.with_repetition(Repetition::OPTIONAL)
|
||||
.build()?,
|
||||
);
|
||||
// Normally there are no other explicit Variant fields exist if a remainder field is present.
|
||||
// However, when structured values reach JSON2_MAX_STRUCTURED_DEPTH, there are. `opaque`
|
||||
// models such a deep leaf without building a deeply nested test schema.
|
||||
let opaque = Arc::new(
|
||||
Type::primitive_type_builder("opaque", parquet::basic::Type::BYTE_ARRAY)
|
||||
.with_repetition(Repetition::OPTIONAL)
|
||||
.build()?,
|
||||
);
|
||||
let root = Arc::new(
|
||||
Type::group_type_builder("j")
|
||||
.with_repetition(Repetition::OPTIONAL)
|
||||
.with_fields(vec![remainder, commit, hot, opaque])
|
||||
.build()?,
|
||||
);
|
||||
Ok(SchemaDescriptor::new(Arc::new(
|
||||
Type::group_type_builder("schema")
|
||||
.with_fields(vec![root])
|
||||
.build()?,
|
||||
)))
|
||||
}
|
||||
|
||||
// Test schema:
|
||||
// schema
|
||||
// `- j
|
||||
|
||||
@@ -16,13 +16,16 @@
|
||||
|
||||
#[cfg(feature = "vector_index")]
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::HashSet;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use api::v1::SemanticType;
|
||||
use arrow_schema::extension::{
|
||||
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
|
||||
};
|
||||
use common_recordbatch::filter::SimpleFilterEvaluator;
|
||||
use common_telemetry::{error, tracing, warn};
|
||||
use common_telemetry::{debug, error, tracing, warn};
|
||||
use datafusion::physical_plan::PhysicalExpr;
|
||||
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
|
||||
use datafusion_expr::utils::expr_to_columns;
|
||||
@@ -31,8 +34,9 @@ use datatypes::arrow::array::ArrayRef;
|
||||
use datatypes::arrow::datatypes::{Field, Schema as ArrowSchema, SchemaRef};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::data_type::ConcreteDataType;
|
||||
use datatypes::extension::json::is_json2_extension_type;
|
||||
use datatypes::extension::json::{Json2ExtensionType, is_json2_extension_type};
|
||||
use datatypes::prelude::DataType;
|
||||
use datatypes::vectors::json::json2_physical_data_type;
|
||||
use futures::StreamExt;
|
||||
use mito_codec::row_converter::build_primary_key_codec;
|
||||
use object_store::ObjectStore;
|
||||
@@ -76,7 +80,6 @@ use crate::sst::index::inverted_index::applier::{
|
||||
};
|
||||
#[cfg(feature = "vector_index")]
|
||||
use crate::sst::index::vector_index::applier::VectorIndexApplierRef;
|
||||
use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
|
||||
use crate::sst::parquet::file_range::{
|
||||
FileRangeContext, FileRangeContextRef, PartitionFilterContext, PreFilterMode, RangeBase,
|
||||
};
|
||||
@@ -94,6 +97,7 @@ use crate::sst::parquet::read_columns::{ProjectionMaskPlan, build_projection_pla
|
||||
use crate::sst::parquet::row_group::ParquetFetchMetrics;
|
||||
use crate::sst::parquet::row_selection::RowGroupSelection;
|
||||
use crate::sst::parquet::stats::RowGroupPruningStats;
|
||||
use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, Json2RewriteTargets, Json2TargetLayout};
|
||||
use crate::sst::{override_pk_field_to_binary, tag_maybe_to_dictionary_field};
|
||||
|
||||
const INDEX_TYPE_FULLTEXT: &str = "fulltext";
|
||||
@@ -108,6 +112,43 @@ fn should_read_pk_as_binary(parquet_meta: &ParquetMetaData) -> bool {
|
||||
should_read_pk_as_binary_with_limit(parquet_meta, DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT)
|
||||
}
|
||||
|
||||
fn apply_json2_rewrite_targets(
|
||||
read_format: &FlatReadFormat,
|
||||
targets: &Json2RewriteTargets,
|
||||
) -> Result<SchemaRef> {
|
||||
let schema = read_format.output_arrow_schema()?;
|
||||
if targets.is_empty() {
|
||||
return Ok(schema);
|
||||
}
|
||||
|
||||
let mut schema = schema.as_ref().clone();
|
||||
let mut fields = schema.fields().iter().cloned().collect::<Vec<_>>();
|
||||
for (column_id, layout) in targets.iter() {
|
||||
let Some(index) = read_format.parquet_projected_index_by_id(*column_id) else {
|
||||
continue;
|
||||
};
|
||||
let Some(field) = fields.get(index) else {
|
||||
continue;
|
||||
};
|
||||
let mut field = field.as_ref().clone();
|
||||
field.set_data_type(json2_physical_data_type(&layout.target_layout));
|
||||
|
||||
let mut metadata = field.metadata().clone();
|
||||
metadata.insert(
|
||||
EXTENSION_TYPE_NAME_KEY.to_string(),
|
||||
Json2ExtensionType::NAME.to_string(),
|
||||
);
|
||||
metadata.insert(
|
||||
EXTENSION_TYPE_METADATA_KEY.to_string(),
|
||||
layout.extension_metadata.clone(),
|
||||
);
|
||||
field.set_metadata(metadata);
|
||||
fields[index] = Arc::new(field);
|
||||
}
|
||||
schema.fields = fields.into();
|
||||
Ok(Arc::new(schema))
|
||||
}
|
||||
|
||||
fn should_read_pk_as_binary_with_limit(
|
||||
parquet_meta: &ParquetMetaData,
|
||||
dict_page_size_limit: usize,
|
||||
@@ -163,6 +204,8 @@ pub struct ParquetReaderBuilder {
|
||||
/// `None` reads all columns. Due to schema change, the projection
|
||||
/// can contain columns not in the parquet file.
|
||||
read_cols: Option<ReadColumns>,
|
||||
/// Compaction-only JSON2 physical rewrite targets.
|
||||
json2_rewrite_targets: Json2RewriteTargets,
|
||||
/// Strategy to cache SST data.
|
||||
cache_strategy: CacheStrategy,
|
||||
/// Index appliers.
|
||||
@@ -206,6 +249,7 @@ impl ParquetReaderBuilder {
|
||||
object_store,
|
||||
predicate: None,
|
||||
read_cols: None,
|
||||
json2_rewrite_targets: Arc::default(),
|
||||
cache_strategy: CacheStrategy::Disabled,
|
||||
inverted_index_appliers: [None, None],
|
||||
bloom_filter_index_appliers: [None, None],
|
||||
@@ -247,6 +291,13 @@ impl ParquetReaderBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Attaches fixed JSON2 physical layouts used by compaction readers.
|
||||
#[must_use]
|
||||
pub(crate) fn json2_rewrite_targets(mut self, targets: Json2RewriteTargets) -> Self {
|
||||
self.json2_rewrite_targets = targets;
|
||||
self
|
||||
}
|
||||
|
||||
/// Attaches the cache to the builder.
|
||||
#[must_use]
|
||||
pub fn cache(mut self, cache: CacheStrategy) -> ParquetReaderBuilder {
|
||||
@@ -461,7 +512,24 @@ impl ParquetReaderBuilder {
|
||||
&file_path,
|
||||
skip_auto_convert,
|
||||
)?;
|
||||
if need_override_sequence(&parquet_meta) {
|
||||
// `region_meta` comes from the Parquet/source file and must not be used as the
|
||||
// target identity. When the caller has no current metadata, the handle is the
|
||||
// only local identity available and therefore denotes a local read.
|
||||
let expected_region_id = self
|
||||
.expected_metadata
|
||||
.as_ref()
|
||||
.map(|metadata| metadata.region_id)
|
||||
.unwrap_or(self.file_handle.region_id());
|
||||
let is_foreign = self.file_handle.region_id() != expected_region_id;
|
||||
if is_foreign {
|
||||
debug!(
|
||||
"Reading foreign SST, file_id: {}, source_region_id: {}, expected_region_id: {}",
|
||||
self.file_handle.file_id().file_id(),
|
||||
self.file_handle.region_id(),
|
||||
expected_region_id,
|
||||
);
|
||||
}
|
||||
if is_foreign || need_override_sequence(&parquet_meta) {
|
||||
read_format
|
||||
.set_override_sequence(self.file_handle.meta_ref().sequence.map(|x| x.get()));
|
||||
}
|
||||
@@ -535,6 +603,8 @@ impl ParquetReaderBuilder {
|
||||
);
|
||||
}
|
||||
|
||||
let output_schema = apply_json2_rewrite_targets(&read_format, &self.json2_rewrite_targets)?;
|
||||
|
||||
// Create ArrowReaderMetadata for async stream building.
|
||||
let mut arrow_reader_options = ArrowReaderOptions::new();
|
||||
if !read_format
|
||||
@@ -546,7 +616,7 @@ impl ParquetReaderBuilder {
|
||||
// Read `__primary_key` as Binary when it's too large for dictionary
|
||||
// encoding; convert_batch wraps it back to a DictionaryArray.
|
||||
let schema_for_reader = if should_read_pk_as_binary(&parquet_meta) {
|
||||
read_format.set_pk_as_binary()?;
|
||||
read_format.set_pk_as_binary(output_schema.clone());
|
||||
override_pk_field_to_binary(read_format.arrow_schema())
|
||||
} else {
|
||||
read_format.arrow_schema().clone()
|
||||
@@ -557,7 +627,18 @@ impl ParquetReaderBuilder {
|
||||
ArrowReaderMetadata::try_new(parquet_meta.clone(), arrow_reader_options)
|
||||
.context(ReadDataPartSnafu)?;
|
||||
|
||||
let output_schema = read_format.output_arrow_schema()?;
|
||||
let json2_rewrite_targets = self
|
||||
.json2_rewrite_targets
|
||||
.iter()
|
||||
.map(|(column_id, layout)| {
|
||||
let column = region_meta
|
||||
.column_by_id(*column_id)
|
||||
.context(UnexpectedSnafu {
|
||||
reason: format!("JSON2 target column by id {column_id} does not exist"),
|
||||
})?;
|
||||
Ok((column.column_schema.name.clone(), layout.clone()))
|
||||
})
|
||||
.collect::<Result<HashMap<_, _>>>()?;
|
||||
|
||||
let reader_builder = RowGroupReaderBuilder {
|
||||
file_handle: self.file_handle.clone(),
|
||||
@@ -566,6 +647,7 @@ impl ParquetReaderBuilder {
|
||||
parquet_metadata_size,
|
||||
arrow_metadata,
|
||||
output_schema,
|
||||
json2_rewrite_targets,
|
||||
object_store: self.object_store.clone(),
|
||||
projection: projection_plan,
|
||||
has_nested_projection,
|
||||
@@ -1786,6 +1868,8 @@ pub(crate) struct RowGroupReaderBuilder {
|
||||
arrow_metadata: ArrowReaderMetadata,
|
||||
/// Projected output schema aligned with `projection.projected_root_presence`.
|
||||
output_schema: SchemaRef,
|
||||
/// JSON2 columns that must be semantically rewritten into the projected target layout.
|
||||
json2_rewrite_targets: HashMap<String, Json2TargetLayout>,
|
||||
/// Object store as an Operator.
|
||||
object_store: ObjectStore,
|
||||
/// Projection mask.
|
||||
@@ -1955,7 +2039,7 @@ impl RowGroupReaderBuilder {
|
||||
&self,
|
||||
stream: ProjectedRecordBatchStream,
|
||||
) -> Result<ProjectedRecordBatchStream> {
|
||||
if !self.has_nested_projection {
|
||||
if !self.has_nested_projection && self.json2_rewrite_targets.is_empty() {
|
||||
return Ok(stream);
|
||||
}
|
||||
|
||||
@@ -1964,6 +2048,7 @@ impl RowGroupReaderBuilder {
|
||||
self.projection.projected_root_presence.clone(),
|
||||
self.output_schema.clone(),
|
||||
)?
|
||||
.with_json2_rewrite_targets(&self.json2_rewrite_targets)?
|
||||
.boxed())
|
||||
}
|
||||
|
||||
|
||||
@@ -23,11 +23,12 @@ use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Instant;
|
||||
|
||||
use bytes::Bytes;
|
||||
use common_telemetry::debug;
|
||||
use common_time::Timestamp;
|
||||
use datatypes::arrow::array::{
|
||||
ArrayRef, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
|
||||
TimestampSecondArray,
|
||||
ArrayRef, BinaryArray, TimestampMicrosecondArray, TimestampMillisecondArray,
|
||||
TimestampNanosecondArray, TimestampSecondArray, UInt32Array,
|
||||
};
|
||||
use datatypes::arrow::compute::{max, min};
|
||||
use datatypes::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
|
||||
@@ -40,7 +41,7 @@ use parquet::file::metadata::KeyValue;
|
||||
use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder};
|
||||
use parquet::schema::types::ColumnPath;
|
||||
use smallvec::smallvec;
|
||||
use snafu::ResultExt;
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME};
|
||||
use store_api::storage::{FileId, SequenceNumber};
|
||||
@@ -50,13 +51,16 @@ use tokio_util::compat::{Compat, FuturesAsyncWriteCompatExt};
|
||||
use crate::access_layer::{FilePathProvider, Metrics, SstInfoArray, TempFileCleaner};
|
||||
use crate::config::{IndexBuildMode, IndexConfig};
|
||||
use crate::error::{
|
||||
InvalidMetadataSnafu, OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu,
|
||||
InvalidMetadataSnafu, InvalidRecordBatchSnafu, OpenDalSnafu, Result, UnexpectedSnafu,
|
||||
WriteParquetSnafu,
|
||||
};
|
||||
use crate::read::FlatSource;
|
||||
use crate::sst::file::RegionFileId;
|
||||
use crate::sst::index::{IndexOutput, Indexer, IndexerBuilder};
|
||||
use crate::sst::parquet::flat_format::{FlatWriteFormat, time_index_column_index};
|
||||
use crate::sst::parquet::format::PrimaryKeyWriteFormat;
|
||||
use crate::sst::parquet::flat_format::{
|
||||
FlatWriteFormat, primary_key_column_index, time_index_column_index,
|
||||
};
|
||||
use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat};
|
||||
use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions};
|
||||
use crate::sst::{
|
||||
DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator,
|
||||
@@ -85,6 +89,71 @@ impl FlatBatchConverter {
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of splitting a batch at the next series boundary.
|
||||
enum SeriesBoundarySplit {
|
||||
/// The whole batch belongs to the current series and stays in the current file.
|
||||
Continue(RecordBatch),
|
||||
/// The batch crosses a series boundary.
|
||||
Split {
|
||||
/// Remaining rows of the current series, appended to the current file.
|
||||
current_file_tail: Option<RecordBatch>,
|
||||
/// Rows of the following series that start the next file.
|
||||
next_file_head: RecordBatch,
|
||||
},
|
||||
}
|
||||
|
||||
/// Returns the dictionary keys and values of the encoded `__primary_key` column.
|
||||
fn encoded_primary_keys(batch: &RecordBatch) -> Result<(&UInt32Array, &BinaryArray)> {
|
||||
let column = batch.column(primary_key_column_index(batch.num_columns()));
|
||||
let primary_keys = column
|
||||
.as_any()
|
||||
.downcast_ref::<PrimaryKeyArray>()
|
||||
.with_context(|| InvalidRecordBatchSnafu {
|
||||
reason: format!(
|
||||
"expected dictionary primary key column, got {:?}",
|
||||
column.data_type()
|
||||
),
|
||||
})?;
|
||||
let values = primary_keys
|
||||
.values()
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.with_context(|| InvalidRecordBatchSnafu {
|
||||
reason: format!(
|
||||
"expected binary primary key values, got {:?}",
|
||||
primary_keys.values().data_type()
|
||||
),
|
||||
})?;
|
||||
Ok((primary_keys.keys(), values))
|
||||
}
|
||||
|
||||
/// Splits `batch` at the first row whose encoded primary key differs from
|
||||
/// `current_primary_key`, the last primary key written to the current file.
|
||||
///
|
||||
/// The writer finishes an oversized file at such a series boundary so that the
|
||||
/// primary key ranges of output files never overlap, which allows pickers like
|
||||
/// TWCS to detect overlapping files by their time and primary key ranges. A
|
||||
/// series is never split in the middle: if no row starts a new series, the whole
|
||||
/// batch stays in the current file even if it already exceeds the size limit.
|
||||
fn split_at_next_series(
|
||||
batch: RecordBatch,
|
||||
current_primary_key: &[u8],
|
||||
) -> Result<SeriesBoundarySplit> {
|
||||
let (keys, values) = encoded_primary_keys(&batch)?;
|
||||
let Some(offset) = (0..batch.num_rows())
|
||||
.find(|&row| values.value(keys.value(row) as usize) != current_primary_key)
|
||||
else {
|
||||
return Ok(SeriesBoundarySplit::Continue(batch));
|
||||
};
|
||||
|
||||
let next_file_head = batch.slice(offset, batch.num_rows() - offset);
|
||||
let current_file_tail = (offset > 0).then(|| batch.slice(0, offset));
|
||||
Ok(SeriesBoundarySplit::Split {
|
||||
current_file_tail,
|
||||
next_file_head,
|
||||
})
|
||||
}
|
||||
|
||||
/// Parquet SST writer.
|
||||
pub struct ParquetWriter<'a, F: WriterFactory, I: IndexerBuilder, P: FilePathProvider> {
|
||||
/// Path provider that creates SST and index file paths according to file id.
|
||||
@@ -269,6 +338,13 @@ where
|
||||
|
||||
/// Iterates FlatSource and writes all RecordBatch in flat format to Parquet file.
|
||||
///
|
||||
/// The source must yield batches globally sorted by the encoded primary key.
|
||||
/// `opts.max_file_size` is a soft limit for regions with a primary key: the
|
||||
/// current file is finished at the next series boundary after exceeding the
|
||||
/// limit (see [split_at_next_series]), so a series larger than the limit
|
||||
/// stays in one file. Regions without a primary key are split at batch
|
||||
/// boundaries.
|
||||
///
|
||||
/// Returns the [SstInfo] if the SST is written.
|
||||
pub async fn write_all_flat(
|
||||
&mut self,
|
||||
@@ -304,6 +380,9 @@ where
|
||||
|
||||
/// Iterates FlatSource and writes all RecordBatch in primary-key format to Parquet file.
|
||||
///
|
||||
/// The source must yield batches globally sorted by the encoded primary key.
|
||||
/// See [Self::write_all_flat] for the `opts.max_file_size` semantics.
|
||||
///
|
||||
/// Returns the [SstInfo] if the SST is written.
|
||||
pub async fn write_all_flat_as_primary_key(
|
||||
&mut self,
|
||||
@@ -336,36 +415,55 @@ where
|
||||
let mut results = smallvec![];
|
||||
let mut stats = SourceStats::default();
|
||||
|
||||
while let Some(record_batch) = self
|
||||
.write_next_flat_batch(&mut source, converter, opts)
|
||||
.await
|
||||
.transpose()
|
||||
{
|
||||
match record_batch {
|
||||
Ok(batch) => {
|
||||
stats.update_flat(&batch)?;
|
||||
if matches!(self.index_config.build_mode, IndexBuildMode::Sync) {
|
||||
let start = Instant::now();
|
||||
// safety: self.current_indexer must be set when first batch has been written.
|
||||
self.current_indexer
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.update_flat(&batch)
|
||||
.await;
|
||||
self.metrics.update_index += start.elapsed();
|
||||
}
|
||||
if let Some(max_file_size) = opts.max_file_size
|
||||
&& self.bytes_written.load(Ordering::Relaxed) > max_file_size
|
||||
{
|
||||
self.finish_current_file(&mut results, &mut stats).await?;
|
||||
}
|
||||
}
|
||||
loop {
|
||||
let start = Instant::now();
|
||||
let batch = match source.next_batch().await {
|
||||
Ok(Some(batch)) => batch,
|
||||
Ok(None) => break,
|
||||
Err(e) => {
|
||||
if let Some(indexer) = &mut self.current_indexer {
|
||||
indexer.abort().await;
|
||||
}
|
||||
self.abort_current_indexer().await;
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
self.metrics.iter_source += start.elapsed();
|
||||
|
||||
if self.metadata.primary_key.is_empty() {
|
||||
self.append_flat_batch(&batch, converter, opts, &mut stats)
|
||||
.await?;
|
||||
if self.exceeds_max_file_size(opts) {
|
||||
self.finish_current_file(&mut results, &mut stats).await?;
|
||||
}
|
||||
} else if self.exceeds_max_file_size(opts)
|
||||
&& let Some(current_primary_key) = stats.last_primary_key.as_deref()
|
||||
{
|
||||
let series_split = match split_at_next_series(batch, current_primary_key) {
|
||||
Ok(series_split) => series_split,
|
||||
Err(e) => {
|
||||
self.abort_current_indexer().await;
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
match series_split {
|
||||
SeriesBoundarySplit::Continue(batch) => {
|
||||
self.append_flat_batch(&batch, converter, opts, &mut stats)
|
||||
.await?;
|
||||
}
|
||||
SeriesBoundarySplit::Split {
|
||||
current_file_tail,
|
||||
next_file_head,
|
||||
} => {
|
||||
if let Some(tail) = current_file_tail {
|
||||
self.append_flat_batch(&tail, converter, opts, &mut stats)
|
||||
.await?;
|
||||
}
|
||||
self.finish_current_file(&mut results, &mut stats).await?;
|
||||
self.append_flat_batch(&next_file_head, converter, opts, &mut stats)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.append_flat_batch(&batch, converter, opts, &mut stats)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -398,29 +496,53 @@ where
|
||||
.set_column_compression(op_type_col, Compression::UNCOMPRESSED)
|
||||
}
|
||||
|
||||
async fn write_next_flat_batch(
|
||||
async fn append_flat_batch(
|
||||
&mut self,
|
||||
source: &mut FlatSource,
|
||||
batch: &RecordBatch,
|
||||
converter: &FlatBatchConverter,
|
||||
opts: &WriteOptions,
|
||||
) -> Result<Option<RecordBatch>> {
|
||||
let start = Instant::now();
|
||||
let Some(record_batch) = source.next_batch().await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
self.metrics.iter_source += start.elapsed();
|
||||
stats: &mut SourceStats,
|
||||
) -> Result<()> {
|
||||
let result = async {
|
||||
let arrow_batch = converter.convert_batch(batch)?;
|
||||
let start = Instant::now();
|
||||
self.maybe_init_writer(arrow_batch.schema_ref(), opts)
|
||||
.await?
|
||||
.write(&arrow_batch)
|
||||
.await
|
||||
.context(WriteParquetSnafu)?;
|
||||
self.metrics.write_batch += start.elapsed();
|
||||
|
||||
let arrow_batch = converter.convert_batch(&record_batch)?;
|
||||
stats.update_flat(batch)?;
|
||||
if matches!(self.index_config.build_mode, IndexBuildMode::Sync) {
|
||||
let start = Instant::now();
|
||||
// safety: self.current_indexer must be set when first batch has been written.
|
||||
self.current_indexer
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.update_flat(batch)
|
||||
.await;
|
||||
self.metrics.update_index += start.elapsed();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
|
||||
let start = Instant::now();
|
||||
self.maybe_init_writer(arrow_batch.schema_ref(), opts)
|
||||
.await?
|
||||
.write(&arrow_batch)
|
||||
.await
|
||||
.context(WriteParquetSnafu)?;
|
||||
self.metrics.write_batch += start.elapsed();
|
||||
// Return original flat batch for stats/indexer which use flat layout.
|
||||
Ok(Some(record_batch))
|
||||
if result.is_err() {
|
||||
self.abort_current_indexer().await;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn exceeds_max_file_size(&self, opts: &WriteOptions) -> bool {
|
||||
opts.max_file_size
|
||||
.is_some_and(|max_size| self.bytes_written.load(Ordering::Relaxed) >= max_size)
|
||||
}
|
||||
|
||||
async fn abort_current_indexer(&mut self) {
|
||||
if let Some(indexer) = &mut self.current_indexer {
|
||||
indexer.abort().await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_init_writer(
|
||||
@@ -481,6 +603,8 @@ struct SourceStats {
|
||||
num_rows: usize,
|
||||
/// Time range of fetched batches.
|
||||
time_range: Option<(Timestamp, Timestamp)>,
|
||||
/// Last primary key written to the current file.
|
||||
last_primary_key: Option<Bytes>,
|
||||
/// Series estimator for computing num_series.
|
||||
series_estimator: SeriesEstimator,
|
||||
}
|
||||
@@ -493,6 +617,9 @@ impl SourceStats {
|
||||
|
||||
self.num_rows += record_batch.num_rows();
|
||||
self.series_estimator.update_flat(record_batch);
|
||||
let (keys, values) = encoded_primary_keys(record_batch)?;
|
||||
let key = keys.value(record_batch.num_rows() - 1);
|
||||
self.last_primary_key = Some(Bytes::copy_from_slice(values.value(key as usize)));
|
||||
|
||||
// Get the timestamp column by index
|
||||
let time_index_col_idx = time_index_column_index(record_batch.num_columns());
|
||||
|
||||
+132
-1
@@ -52,7 +52,9 @@ use log_store::raft_engine::log_store::RaftEngineLogStore;
|
||||
use log_store::test_util::log_store_util;
|
||||
use moka::future::CacheBuilder;
|
||||
use object_store::ObjectStore;
|
||||
use object_store::layers::mock::MockLayer;
|
||||
use object_store::layers::mock::{
|
||||
Buffer, Deleter, Metadata, MockLayer, MockLayerBuilder, OpDelete, Result as MockResult, Writer,
|
||||
};
|
||||
use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef};
|
||||
use object_store::services::Fs;
|
||||
use rskafka::client::partition::{Compression, UnknownTopicHandling};
|
||||
@@ -67,6 +69,7 @@ use store_api::region_request::{
|
||||
RegionOpenRequest, RegionPutRequest, RegionRequest,
|
||||
};
|
||||
use store_api::storage::{ColumnId, RegionId};
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use crate::cache::write_cache::{WriteCache, WriteCacheRef};
|
||||
use crate::config::MitoConfig;
|
||||
@@ -75,6 +78,7 @@ use crate::engine::{MITO_ENGINE_NAME, MitoEngine};
|
||||
use crate::error::Result;
|
||||
use crate::flush::{WriteBufferManager, WriteBufferManagerRef};
|
||||
use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
|
||||
use crate::manifest::storage::{is_checkpoint_file, is_delta_file};
|
||||
use crate::read::{Batch, BatchBuilder, BatchReader};
|
||||
use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef};
|
||||
use crate::sst::FormatType;
|
||||
@@ -85,6 +89,133 @@ use crate::sst::index::puffin_manager::PuffinManagerFactory;
|
||||
use crate::time_provider::{StdTimeProvider, TimeProviderRef};
|
||||
use crate::worker::WorkerGroup;
|
||||
|
||||
/// Controls a mock object-store layer that blocks a checkpoint task once.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CheckpointTaskBlocker {
|
||||
entered: Arc<Notify>,
|
||||
release: Arc<Notify>,
|
||||
armed: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl CheckpointTaskBlocker {
|
||||
/// Blocks normal manifest checkpoint cleanup at batch-delete close.
|
||||
pub(crate) fn block_cleanup() -> (Self, MockLayer) {
|
||||
let blocker = Self {
|
||||
entered: Arc::new(Notify::new()),
|
||||
release: Arc::new(Notify::new()),
|
||||
armed: Arc::new(AtomicBool::new(true)),
|
||||
};
|
||||
let factory_blocker = blocker.clone();
|
||||
let layer = MockLayerBuilder::default()
|
||||
.deleter_factory(Arc::new(move |inner| {
|
||||
Box::new(BlockingCheckpointDeleter {
|
||||
inner,
|
||||
blocker: factory_blocker.clone(),
|
||||
has_manifest_cleanup_target: false,
|
||||
})
|
||||
}))
|
||||
.build()
|
||||
.unwrap();
|
||||
(blocker, layer)
|
||||
}
|
||||
|
||||
/// Blocks publication of `_last_checkpoint` at writer close.
|
||||
pub(crate) fn block_last_checkpoint_write() -> (Self, MockLayer) {
|
||||
let blocker = Self {
|
||||
entered: Arc::new(Notify::new()),
|
||||
release: Arc::new(Notify::new()),
|
||||
armed: Arc::new(AtomicBool::new(true)),
|
||||
};
|
||||
let factory_blocker = blocker.clone();
|
||||
let layer = MockLayerBuilder::default()
|
||||
.writer_factory(Arc::new(move |path, _args, inner| {
|
||||
Box::new(BlockingCheckpointWriter {
|
||||
path: path.to_string(),
|
||||
inner,
|
||||
blocker: factory_blocker.clone(),
|
||||
})
|
||||
}))
|
||||
.build()
|
||||
.unwrap();
|
||||
(blocker, layer)
|
||||
}
|
||||
|
||||
pub(crate) async fn wait_until_blocked(&self) {
|
||||
self.entered.notified().await;
|
||||
}
|
||||
|
||||
pub(crate) fn release(&self) {
|
||||
self.release.notify_one();
|
||||
}
|
||||
|
||||
pub(crate) fn arm_next_close(&self) {
|
||||
self.armed.store(true, Ordering::Release);
|
||||
}
|
||||
|
||||
async fn block_once(&self) {
|
||||
if self.armed.swap(false, Ordering::AcqRel) {
|
||||
self.entered.notify_one();
|
||||
self.release.notified().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct BlockingCheckpointDeleter {
|
||||
inner: Deleter,
|
||||
blocker: CheckpointTaskBlocker,
|
||||
has_manifest_cleanup_target: bool,
|
||||
}
|
||||
|
||||
impl object_store::layers::mock::Delete for BlockingCheckpointDeleter {
|
||||
async fn delete(&mut self, path: &str, args: OpDelete) -> MockResult<()> {
|
||||
self.inner.delete(path, args).await?;
|
||||
if is_manifest_checkpoint_file(path) {
|
||||
self.has_manifest_cleanup_target = true;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close(&mut self) -> MockResult<()> {
|
||||
if self.has_manifest_cleanup_target {
|
||||
self.blocker.block_once().await;
|
||||
}
|
||||
self.inner.close().await
|
||||
}
|
||||
}
|
||||
|
||||
fn is_manifest_checkpoint_file(path: &str) -> bool {
|
||||
// The mock deleter receives paths relative to the listed manifest
|
||||
// directory, so the normal/staging directory segments are unavailable.
|
||||
let path = Path::new(path);
|
||||
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
|
||||
return false;
|
||||
};
|
||||
is_delta_file(file_name) || is_checkpoint_file(file_name)
|
||||
}
|
||||
|
||||
struct BlockingCheckpointWriter {
|
||||
path: String,
|
||||
inner: Writer,
|
||||
blocker: CheckpointTaskBlocker,
|
||||
}
|
||||
|
||||
impl object_store::layers::mock::Write for BlockingCheckpointWriter {
|
||||
async fn write(&mut self, bs: Buffer) -> MockResult<()> {
|
||||
self.inner.write(bs).await
|
||||
}
|
||||
|
||||
async fn close(&mut self) -> MockResult<Metadata> {
|
||||
if self.path.ends_with("_last_checkpoint") {
|
||||
self.blocker.block_once().await;
|
||||
}
|
||||
self.inner.close().await
|
||||
}
|
||||
|
||||
async fn abort(&mut self) -> MockResult<()> {
|
||||
self.inner.abort().await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn new_noop_file_purger() -> FilePurgerRef {
|
||||
Arc::new(NoopFilePurger)
|
||||
}
|
||||
|
||||
@@ -126,6 +126,10 @@ impl<S: LogStore> RegionWorkerLoop<S> {
|
||||
// First step: clear all staging manifest files.
|
||||
{
|
||||
let mut manager = region.manifest_ctx.manifest_manager.write().await;
|
||||
// `set_entering_staging` has already stopped new normal manifest
|
||||
// publications. Wait for an older checkpoint to finish cleanup
|
||||
// before exposing staging to repartition remap readers.
|
||||
manager.wait_for_pending_checkpoint().await;
|
||||
manager
|
||||
.clear_staging_manifest_and_dir()
|
||||
.await
|
||||
|
||||
@@ -79,6 +79,7 @@ impl<S> RegionWorkerLoop<S> {
|
||||
IndexBuildTask {
|
||||
region_id: region.region_id,
|
||||
file: file.clone(),
|
||||
target_region_metadata: version.metadata.clone(),
|
||||
source: IndexBuildSource::new(file_meta, version.metadata.schema_version),
|
||||
reason: build_type,
|
||||
access_layer: access_layer.clone(),
|
||||
|
||||
Binary file not shown.
@@ -31,8 +31,8 @@ use table::metadata::TableInfoRef;
|
||||
|
||||
use crate::error::{
|
||||
CatalogSnafu, ColumnDataTypeSnafu, ColumnDefaultValueSnafu, ColumnNoneDefaultValueSnafu,
|
||||
ColumnNotFoundSnafu, InvalidInsertRequestSnafu, InvalidSqlSnafu, MissingInsertBodySnafu,
|
||||
ParseSqlSnafu, Result, SchemaReadOnlySnafu, TableNotFoundSnafu,
|
||||
ColumnNotFoundSnafu, InvalidSqlSnafu, MissingInsertBodySnafu, ParseSqlSnafu, Result,
|
||||
SchemaReadOnlySnafu, TableNotFoundSnafu,
|
||||
};
|
||||
use crate::insert::InstantAndNormalInsertRequests;
|
||||
use crate::req_convert::common::partitioner::Partitioner;
|
||||
@@ -73,7 +73,6 @@ impl<'a> StatementToRegion<'a> {
|
||||
!common_catalog::consts::is_readonly_schema(&schema),
|
||||
SchemaReadOnlySnafu { name: schema }
|
||||
);
|
||||
|
||||
let column_names = column_names(stmt, &table_schema);
|
||||
let column_count = column_names.len();
|
||||
|
||||
@@ -280,27 +279,9 @@ fn sql_value_to_value(
|
||||
)
|
||||
.context(crate::error::SqlCommonSnafu)?
|
||||
};
|
||||
validate(&value)?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn validate(value: &Value) -> Result<()> {
|
||||
match value {
|
||||
Value::Json(value) => {
|
||||
// Json object will be stored as Arrow struct in parquet, and it has the restriction:
|
||||
// "Parquet does not support writing empty structs".
|
||||
ensure!(
|
||||
!value.is_empty_object(),
|
||||
InvalidInsertRequestSnafu {
|
||||
reason: "empty json object is not supported, consider adding a dummy field"
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
fn replace_default(sql_val: &SqlValue) -> bool {
|
||||
matches!(sql_val, SqlValue::Placeholder(s) if s.to_lowercase() == DEFAULT_PLACEHOLDER_VALUE)
|
||||
}
|
||||
|
||||
@@ -83,7 +83,7 @@ use table::table_reference::TableReference;
|
||||
pub use self::admin::{
|
||||
AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef,
|
||||
AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService,
|
||||
AdminFunctionServiceRef,
|
||||
AdminFunctionServiceRef, admin_output_schema,
|
||||
};
|
||||
use self::set::{
|
||||
set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding,
|
||||
|
||||
@@ -19,6 +19,7 @@ use std::sync::Arc;
|
||||
|
||||
use common_function::function::FunctionContext;
|
||||
use common_function::function_registry::{FUNCTION_REGISTRY, get_admin_function};
|
||||
use common_function::state::FunctionState;
|
||||
use common_query::Output;
|
||||
use common_recordbatch::{RecordBatch, RecordBatches};
|
||||
use common_sql::convert::sql_value_to_value;
|
||||
@@ -77,6 +78,103 @@ struct CoreAdminFunctionService {
|
||||
query_engine: query::QueryEngineRef,
|
||||
}
|
||||
|
||||
/// Parts of an `ADMIN` call needed both for execution and schema derivation.
|
||||
struct ResolvedAdminFunction {
|
||||
admin_udf: datafusion_expr::ScalarUDF,
|
||||
fn_name: String,
|
||||
args: Vec<VectorRef>,
|
||||
arg_types: Vec<ArrowDataType>,
|
||||
ret_type: ArrowDataType,
|
||||
}
|
||||
|
||||
/// Resolves the function, parses its literal arguments and derives its
|
||||
/// return type, without executing it.
|
||||
fn resolve_admin_function(
|
||||
stmt: &Admin,
|
||||
query_ctx: &QueryContextRef,
|
||||
state: Arc<FunctionState>,
|
||||
) -> Result<ResolvedAdminFunction> {
|
||||
let Admin::Func(func) = stmt;
|
||||
// the function name should be in lower case.
|
||||
let func_name = func.name.to_string().to_lowercase();
|
||||
let factory = get_admin_function(&func_name)
|
||||
.or_else(|| FUNCTION_REGISTRY.get_function(&func_name))
|
||||
.context(error::AdminFunctionNotFoundSnafu {
|
||||
name: func_name.clone(),
|
||||
})?;
|
||||
|
||||
let func_ctx = FunctionContext {
|
||||
query_ctx: query_ctx.clone(),
|
||||
state,
|
||||
};
|
||||
|
||||
let admin_udf = factory.provide(func_ctx);
|
||||
admin_udf
|
||||
.as_async()
|
||||
.context(error::AdminFunctionNotFoundSnafu { name: func_name })?;
|
||||
|
||||
let fn_name = admin_udf.name().to_string();
|
||||
let signature = admin_udf.signature();
|
||||
|
||||
// Parse function arguments
|
||||
let FunctionArguments::List(args) = &func.args else {
|
||||
return error::BuildAdminFunctionArgsSnafu {
|
||||
msg: format!("unsupported function args {} for {}", func.args, fn_name),
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
let arg_values = args
|
||||
.args
|
||||
.iter()
|
||||
.map(|arg| {
|
||||
let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else {
|
||||
return error::BuildAdminFunctionArgsSnafu {
|
||||
msg: format!("unsupported function arg {arg} for {}", fn_name),
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
Ok(&value.value)
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let args = args_to_vector(&signature.type_signature, &arg_values, query_ctx)?;
|
||||
let arg_types = args
|
||||
.iter()
|
||||
.map(|arg| arg.data_type().as_arrow_type())
|
||||
.collect::<Vec<_>>();
|
||||
let ret_type =
|
||||
admin_udf
|
||||
.return_type(&arg_types)
|
||||
.map_err(|e| error::Error::BuildAdminFunctionArgs {
|
||||
msg: format!(
|
||||
"Failed to get return type of admin function {}: {}",
|
||||
fn_name, e
|
||||
),
|
||||
})?;
|
||||
|
||||
Ok(ResolvedAdminFunction {
|
||||
admin_udf,
|
||||
fn_name,
|
||||
args,
|
||||
arg_types,
|
||||
ret_type,
|
||||
})
|
||||
}
|
||||
|
||||
/// Output schema of an `ADMIN` statement, mirroring what
|
||||
/// [`CoreAdminFunctionService::execute`] produces. `None` if the function
|
||||
/// or arguments are unresolvable; execution will then surface the error.
|
||||
pub fn admin_output_schema(stmt: &Admin, query_ctx: &QueryContextRef) -> Option<Schema> {
|
||||
let resolved =
|
||||
resolve_admin_function(stmt, query_ctx, Arc::new(FunctionState::default())).ok()?;
|
||||
Some(Schema::new(vec![ColumnSchema::new(
|
||||
// Use statement as the result column name
|
||||
stmt.to_string(),
|
||||
ConcreteDataType::from_arrow_type(&resolved.ret_type),
|
||||
false,
|
||||
)]))
|
||||
}
|
||||
|
||||
impl CoreAdminFunctionService {
|
||||
fn new(query_engine: query::QueryEngineRef) -> Self {
|
||||
Self { query_engine }
|
||||
@@ -88,62 +186,23 @@ impl CoreAdminFunctionService {
|
||||
query_ctx,
|
||||
} = request;
|
||||
|
||||
let Admin::Func(func) = &stmt;
|
||||
// the function name should be in lower case.
|
||||
let func_name = func.name.to_string().to_lowercase();
|
||||
let factory = get_admin_function(&func_name)
|
||||
.or_else(|| FUNCTION_REGISTRY.get_function(&func_name))
|
||||
.context(error::AdminFunctionNotFoundSnafu {
|
||||
name: func_name.clone(),
|
||||
})?;
|
||||
|
||||
let func_ctx = FunctionContext {
|
||||
query_ctx: query_ctx.clone(),
|
||||
state: self.query_engine.engine_state().function_state(),
|
||||
};
|
||||
|
||||
let admin_udf = factory.provide(func_ctx);
|
||||
let resolved = resolve_admin_function(
|
||||
&stmt,
|
||||
&query_ctx,
|
||||
self.query_engine.engine_state().function_state(),
|
||||
)?;
|
||||
let ResolvedAdminFunction {
|
||||
admin_udf,
|
||||
fn_name,
|
||||
args,
|
||||
arg_types,
|
||||
ret_type,
|
||||
} = resolved;
|
||||
let admin_async_fn = admin_udf
|
||||
.as_async()
|
||||
.context(error::AdminFunctionNotFoundSnafu { name: func_name })?;
|
||||
|
||||
let fn_name = admin_udf.name();
|
||||
let signature = admin_udf.signature();
|
||||
|
||||
// Parse function arguments
|
||||
let FunctionArguments::List(args) = &func.args else {
|
||||
return error::BuildAdminFunctionArgsSnafu {
|
||||
msg: format!("unsupported function args {} for {}", func.args, fn_name),
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
let arg_values = args
|
||||
.args
|
||||
.iter()
|
||||
.map(|arg| {
|
||||
let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else {
|
||||
return error::BuildAdminFunctionArgsSnafu {
|
||||
msg: format!("unsupported function arg {arg} for {}", fn_name),
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
Ok(&value.value)
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let args = args_to_vector(&signature.type_signature, &arg_values, &query_ctx)?;
|
||||
let arg_types = args
|
||||
.iter()
|
||||
.map(|arg| arg.data_type().as_arrow_type())
|
||||
.collect::<Vec<_>>();
|
||||
let ret_type = admin_udf.return_type(&arg_types).map_err(|e| {
|
||||
error::Error::BuildAdminFunctionArgs {
|
||||
msg: format!(
|
||||
"Failed to get return type of admin function {}: {}",
|
||||
fn_name, e
|
||||
),
|
||||
}
|
||||
})?;
|
||||
.context(error::AdminFunctionNotFoundSnafu {
|
||||
name: fn_name.clone(),
|
||||
})?;
|
||||
|
||||
// Convert arguments to DataFusion ColumnarValue format
|
||||
let columnar_args: Vec<datafusion_expr::ColumnarValue> = args
|
||||
@@ -174,9 +233,7 @@ impl CoreAdminFunctionService {
|
||||
let result_columnar = admin_async_fn
|
||||
.invoke_async_with_args(func_args)
|
||||
.await
|
||||
.with_context(|_| ExecuteAdminFunctionSnafu {
|
||||
msg: fn_name.to_string(),
|
||||
})?;
|
||||
.with_context(|_| ExecuteAdminFunctionSnafu { msg: fn_name })?;
|
||||
|
||||
// Convert result back to VectorRef
|
||||
let result_columnar: common_query::prelude::ColumnarValue =
|
||||
|
||||
@@ -45,7 +45,7 @@ itertools.workspace = true
|
||||
jsonb.workspace = true
|
||||
jsonpath-rust = "0.7.5"
|
||||
lazy_static.workspace = true
|
||||
moka = { workspace = true, features = ["sync"] }
|
||||
moka = { workspace = true, features = ["future"] }
|
||||
once_cell.workspace = true
|
||||
operator.workspace = true
|
||||
ordered-float.workspace = true
|
||||
|
||||
@@ -389,6 +389,20 @@ pub enum Error {
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
#[snafu(display("Invalid JSON2 type hint: {reason}"))]
|
||||
InvalidJson2TypeHint {
|
||||
reason: String,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
#[snafu(display("Invalid JSON2 type hint path '{path}'"))]
|
||||
ParseJson2TypeHintPath {
|
||||
path: String,
|
||||
#[snafu(source)]
|
||||
source: sql::error::Error,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
#[snafu(display("Transform index `type` must be set."))]
|
||||
TransformIndexTypeMustBeSet {
|
||||
#[snafu(implicit)]
|
||||
@@ -693,6 +707,15 @@ pub enum Error {
|
||||
source: common_recordbatch::error::Error,
|
||||
},
|
||||
|
||||
/// `try_get_with` shares one loader across concurrent misses, so its error
|
||||
/// arrives behind an `Arc`.
|
||||
#[snafu(display("Failed to load pipeline into cache: {}", error))]
|
||||
CacheLoad {
|
||||
error: std::sync::Arc<Error>,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("A valid table suffix template is required for tablesuffix section"))]
|
||||
RequiredTableSuffixTemplate,
|
||||
|
||||
@@ -881,6 +904,7 @@ impl ErrorExt for Error {
|
||||
fn status_code(&self) -> StatusCode {
|
||||
use Error::*;
|
||||
match self {
|
||||
CacheLoad { error, .. } => error.status_code(),
|
||||
CastType { .. } => StatusCode::Unexpected,
|
||||
PipelineTableNotFound { .. } => StatusCode::TableNotFound,
|
||||
InsertPipeline { source, .. } => source.status_code(),
|
||||
@@ -952,6 +976,8 @@ impl ErrorExt for Error {
|
||||
| TransformElementMustBeMap { .. }
|
||||
| TransformFieldMustBeSet { .. }
|
||||
| TransformTypeMustBeSet { .. }
|
||||
| InvalidJson2TypeHint { .. }
|
||||
| ParseJson2TypeHintPath { .. }
|
||||
| TransformIndexTypeMustBeSet { .. }
|
||||
| TransformIndexUnsupportedField { .. }
|
||||
| TransformIndexOptionMustBeScalar { .. }
|
||||
|
||||
+66
-27
@@ -37,7 +37,9 @@ use crate::error::{
|
||||
YamlLoadSnafu, YamlParseSnafu,
|
||||
};
|
||||
use crate::etl::processor::ProcessorKind;
|
||||
use crate::etl::transform::transformer::greptime::{RowWithTableSuffix, values_to_rows};
|
||||
use crate::etl::transform::transformer::greptime::{
|
||||
RowWithTableSuffix, values_to_row, values_to_rows,
|
||||
};
|
||||
use crate::tablesuffix::TableSuffixTemplate;
|
||||
use crate::{
|
||||
ContextOpt, GreptimeTransformer, IdentityTimeIndex, PipelineContext, SchemaInfo,
|
||||
@@ -236,6 +238,14 @@ pub enum PipelineExecOutput {
|
||||
Filtered,
|
||||
}
|
||||
|
||||
/// The result after processors and dispatcher rules have run.
|
||||
#[derive(Debug)]
|
||||
pub enum PipelineProcessOutput {
|
||||
Processed(VrlValue),
|
||||
DispatchedTo(DispatchedTo, VrlValue),
|
||||
Filtered,
|
||||
}
|
||||
|
||||
/// Output from a successful pipeline transformation.
|
||||
///
|
||||
/// Rows are grouped by their ContextOpt, with each row having its own optional
|
||||
@@ -303,24 +313,42 @@ impl Pipeline {
|
||||
|
||||
pub fn exec_mut(
|
||||
&self,
|
||||
mut val: VrlValue,
|
||||
val: VrlValue,
|
||||
pipeline_ctx: &PipelineContext<'_>,
|
||||
schema_info: &mut SchemaInfo,
|
||||
) -> Result<PipelineExecOutput> {
|
||||
// process
|
||||
match self.process_mut(val)? {
|
||||
PipelineProcessOutput::Processed(val) => self
|
||||
.transform_mut(val, pipeline_ctx, schema_info)
|
||||
.map(PipelineExecOutput::Transformed),
|
||||
PipelineProcessOutput::DispatchedTo(dispatched_to, val) => {
|
||||
Ok(PipelineExecOutput::DispatchedTo(dispatched_to, val))
|
||||
}
|
||||
PipelineProcessOutput::Filtered => Ok(PipelineExecOutput::Filtered),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn process_mut(&self, mut val: VrlValue) -> Result<PipelineProcessOutput> {
|
||||
for processor in self.processors.iter() {
|
||||
val = processor.exec_mut(val)?;
|
||||
if val.is_null() {
|
||||
// line is filtered
|
||||
return Ok(PipelineExecOutput::Filtered);
|
||||
return Ok(PipelineProcessOutput::Filtered);
|
||||
}
|
||||
}
|
||||
|
||||
// dispatch, fast return if matched
|
||||
if let Some(rule) = self.dispatcher.as_ref().and_then(|d| d.exec(&val)) {
|
||||
return Ok(PipelineExecOutput::DispatchedTo(rule.into(), val));
|
||||
return Ok(PipelineProcessOutput::DispatchedTo(rule.into(), val));
|
||||
}
|
||||
|
||||
Ok(PipelineProcessOutput::Processed(val))
|
||||
}
|
||||
|
||||
pub fn transform_mut(
|
||||
&self,
|
||||
val: VrlValue,
|
||||
pipeline_ctx: &PipelineContext<'_>,
|
||||
schema_info: &mut SchemaInfo,
|
||||
) -> Result<TransformedOutput> {
|
||||
let mut val = if val.is_array() {
|
||||
val
|
||||
} else {
|
||||
@@ -356,9 +384,7 @@ impl Pipeline {
|
||||
}
|
||||
};
|
||||
|
||||
Ok(PipelineExecOutput::Transformed(TransformedOutput {
|
||||
rows_by_context,
|
||||
}))
|
||||
Ok(TransformedOutput { rows_by_context })
|
||||
}
|
||||
|
||||
pub fn processors(&self) -> &processor::Processors {
|
||||
@@ -369,6 +395,10 @@ impl Pipeline {
|
||||
&self.transformer
|
||||
}
|
||||
|
||||
pub fn resolve_table_suffix(&self, value: &VrlValue) -> Option<String> {
|
||||
ContextOpt::resolve_table_suffix(self.tablesuffix.as_ref(), value)
|
||||
}
|
||||
|
||||
// the method is for test purpose
|
||||
pub fn schemas(&self) -> Option<&Vec<greptime_proto::v1::ColumnSchema>> {
|
||||
match &self.transformer {
|
||||
@@ -409,34 +439,43 @@ fn transform_array_elements_by_ctx(
|
||||
);
|
||||
}
|
||||
|
||||
let values =
|
||||
unwrap_or_continue_if_err!(transformer.transform_mut(element, is_v1), skip_error);
|
||||
let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, element);
|
||||
let values = unwrap_or_continue_if_err!(
|
||||
transformer.transform_mut_with_schema(
|
||||
element,
|
||||
is_v1,
|
||||
schema_info,
|
||||
table_suffix.as_deref(),
|
||||
),
|
||||
skip_error
|
||||
);
|
||||
if is_v1 {
|
||||
// v1 mode: just use transformer output directly
|
||||
let mut opt = unwrap_or_continue_if_err!(
|
||||
let opt = unwrap_or_continue_if_err!(
|
||||
ContextOpt::from_pipeline_map_to_opt(element),
|
||||
skip_error
|
||||
);
|
||||
let table_suffix = opt.resolve_table_suffix(tablesuffix_template, element);
|
||||
rows_by_context
|
||||
.entry(opt)
|
||||
.or_insert_with(Vec::new)
|
||||
.push((Row { values }, table_suffix));
|
||||
} else {
|
||||
// v2 mode: combine with auto-transform for remaining fields
|
||||
let element_rows_map = values_to_rows(
|
||||
schema_info,
|
||||
element.clone(),
|
||||
pipeline_ctx,
|
||||
Some(values),
|
||||
false,
|
||||
tablesuffix_template,
|
||||
)
|
||||
.map_err(Box::new)
|
||||
.context(TransformArrayElementSnafu { index })?;
|
||||
for (k, v) in element_rows_map {
|
||||
rows_by_context.entry(k).or_default().extend(v);
|
||||
}
|
||||
let mut value = element.clone();
|
||||
let opt = unwrap_or_continue_if_err!(
|
||||
ContextOpt::from_pipeline_map_to_opt(&mut value),
|
||||
skip_error
|
||||
);
|
||||
let row = unwrap_or_continue_if_err!(
|
||||
values_to_row(schema_info, value, pipeline_ctx, Some(values), false,)
|
||||
.map_err(Box::new)
|
||||
.context(TransformArrayElementSnafu { index }),
|
||||
skip_error
|
||||
);
|
||||
rows_by_context
|
||||
.entry(opt)
|
||||
.or_default()
|
||||
.push((row, table_suffix));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -69,10 +69,6 @@ pub struct ContextOpt {
|
||||
|
||||
// reset the schema in query context
|
||||
schema: Option<String>,
|
||||
|
||||
// pipeline options, not set in query context
|
||||
// can be removed before the end of the pipeline execution
|
||||
table_suffix: Option<String>,
|
||||
}
|
||||
|
||||
impl ContextOpt {
|
||||
@@ -112,9 +108,7 @@ impl ContextOpt {
|
||||
GREPTIME_SKIP_WAL => {
|
||||
opt.skip_wal = Some(v);
|
||||
}
|
||||
GREPTIME_TABLE_SUFFIX => {
|
||||
opt.table_suffix = Some(v);
|
||||
}
|
||||
GREPTIME_TABLE_SUFFIX => {}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
@@ -123,12 +117,13 @@ impl ContextOpt {
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_table_suffix(
|
||||
&mut self,
|
||||
table_suffix: Option<&TableSuffixTemplate>,
|
||||
pipeline_map: &VrlValue,
|
||||
) -> Option<String> {
|
||||
self.table_suffix
|
||||
.take()
|
||||
pipeline_map
|
||||
.as_object()
|
||||
.and_then(|map| map.get(GREPTIME_TABLE_SUFFIX))
|
||||
.map(|value| value.to_string_lossy().to_string())
|
||||
.or_else(|| table_suffix.and_then(|s| s.apply(pipeline_map)))
|
||||
}
|
||||
|
||||
|
||||
@@ -17,17 +17,21 @@ pub mod transformer;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use api::helper::ColumnDataTypeWrapper;
|
||||
use api::v1::ColumnDataType;
|
||||
use api::v1::value::ValueData;
|
||||
use chrono::Utc;
|
||||
use datatypes::schema::{FulltextOptions, SkippingIndexOptions};
|
||||
use datatypes::json::{JsonSettings, JsonTypeHint};
|
||||
use datatypes::schema::{ColumnDefaultConstraint, FulltextOptions, SkippingIndexOptions};
|
||||
use datatypes::value::Value;
|
||||
use snafu::{OptionExt, ResultExt, ensure};
|
||||
use sql::parsers::utils::{
|
||||
validate_column_fulltext_create_option, validate_column_skipping_index_create_option,
|
||||
};
|
||||
|
||||
use crate::error::{
|
||||
Error, FieldMustBeTypeSnafu, KeyMustBeStringSnafu, Result, TransformElementMustBeMapSnafu,
|
||||
Error, FieldMustBeTypeSnafu, InvalidJson2TypeHintSnafu, KeyMustBeStringSnafu,
|
||||
ParseJson2TypeHintPathSnafu, Result, TransformElementMustBeMapSnafu,
|
||||
TransformFieldMustBeSetSnafu, TransformIndexOptionMustBeScalarSnafu, TransformIndexOptionSnafu,
|
||||
TransformIndexOptionUnsupportedSnafu, TransformIndexOptionsUnsupportedSnafu,
|
||||
TransformIndexTypeMismatchSnafu, TransformIndexTypeMustBeSetSnafu,
|
||||
@@ -49,6 +53,10 @@ const TRANSFORM_INDEX_OPTIONS_FIELD: &str = "index.options";
|
||||
const TRANSFORM_TAG: &str = "tag";
|
||||
const TRANSFORM_DEFAULT: &str = "default";
|
||||
const TRANSFORM_ON_FAILURE: &str = "on_failure";
|
||||
const JSON2_TYPE: &str = "json2";
|
||||
const JSON2_TYPE_HINT: &str = "type.json2[]";
|
||||
const JSON2_TYPE_HINT_PATH: &str = "path";
|
||||
const JSON2_TYPE_HINT_NULLABLE: &str = "nullable";
|
||||
|
||||
pub use transformer::greptime::GreptimeTransformer;
|
||||
|
||||
@@ -141,6 +149,7 @@ impl TryFrom<&Vec<yaml_rust::Yaml>> for Transforms {
|
||||
pub struct Transform {
|
||||
pub fields: Fields,
|
||||
pub type_: ColumnDataType,
|
||||
pub(crate) json_settings: Option<JsonSettings>,
|
||||
pub default: Option<ValueData>,
|
||||
pub index: Option<Index>,
|
||||
pub index_options: Option<TransformIndexOptions>,
|
||||
@@ -196,7 +205,8 @@ impl TransformIndexOptions {
|
||||
// ColumnDataType::TimestampMicrosecond
|
||||
// ColumnDataType::TimestampMillisecond
|
||||
// ColumnDataType::TimestampSecond
|
||||
// ColumnDataType::Binary
|
||||
// ColumnDataType::Binary (JSONB)
|
||||
// ColumnDataType::Json (JSON2)
|
||||
|
||||
impl Transform {
|
||||
pub(crate) fn get_default(&self) -> Option<&ValueData> {
|
||||
@@ -260,6 +270,7 @@ fn get_default_for_type(ty: &ColumnDataType) -> Result<ValueData> {
|
||||
ColumnDataType::Float32 => ValueData::F32Value(0.0),
|
||||
ColumnDataType::Float64 => ValueData::F64Value(0.0),
|
||||
ColumnDataType::Binary => ValueData::BinaryValue(jsonb::Value::Null.to_vec()),
|
||||
ColumnDataType::Json => ValueData::JsonValue(Default::default()),
|
||||
ColumnDataType::String => ValueData::StringValue(String::new()),
|
||||
|
||||
ColumnDataType::TimestampSecond => ValueData::TimestampSecondValue(0),
|
||||
@@ -419,6 +430,162 @@ fn lower_transform_index_options(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_transform_type(value: &yaml_rust::Yaml) -> Result<(ColumnDataType, Option<JsonSettings>)> {
|
||||
if let Some(type_name) = value.as_str() {
|
||||
return Ok((parse_str_type(type_name)?, None));
|
||||
}
|
||||
|
||||
let config = value.as_hash().context(FieldMustBeTypeSnafu {
|
||||
field: TRANSFORM_TYPE,
|
||||
ty: "string or map",
|
||||
})?;
|
||||
ensure!(
|
||||
config.len() == 1,
|
||||
InvalidJson2TypeHintSnafu {
|
||||
reason: "transform type map must contain exactly one `json2` field".to_string()
|
||||
}
|
||||
);
|
||||
let (type_name, hints) = config.iter().next().context(InvalidJson2TypeHintSnafu {
|
||||
reason: "transform type map must contain a `json2` field".to_string(),
|
||||
})?;
|
||||
let type_name = type_name.as_str().with_context(|| KeyMustBeStringSnafu {
|
||||
k: type_name.clone(),
|
||||
})?;
|
||||
ensure!(
|
||||
type_name.eq_ignore_ascii_case(JSON2_TYPE),
|
||||
InvalidJson2TypeHintSnafu {
|
||||
reason: format!("unsupported transform type map `{type_name}`")
|
||||
}
|
||||
);
|
||||
|
||||
let hints = hints.as_vec().context(FieldMustBeTypeSnafu {
|
||||
field: JSON2_TYPE,
|
||||
ty: "list",
|
||||
})?;
|
||||
let hints = hints
|
||||
.iter()
|
||||
.map(parse_json2_type_hint)
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
Ok((
|
||||
ColumnDataType::Json,
|
||||
Some(JsonSettings::try_new(hints, None)?),
|
||||
))
|
||||
}
|
||||
|
||||
fn parse_json2_type_hint(value: &yaml_rust::Yaml) -> Result<JsonTypeHint> {
|
||||
let config = value.as_hash().context(FieldMustBeTypeSnafu {
|
||||
field: JSON2_TYPE_HINT,
|
||||
ty: "map",
|
||||
})?;
|
||||
let mut path = None;
|
||||
let mut type_name = None;
|
||||
let mut nullable = true;
|
||||
let mut default = None;
|
||||
let mut index = None;
|
||||
|
||||
for (key, value) in config {
|
||||
let key = key
|
||||
.as_str()
|
||||
.with_context(|| KeyMustBeStringSnafu { k: key.clone() })?;
|
||||
match key {
|
||||
JSON2_TYPE_HINT_PATH => path = Some(yaml_string(value, JSON2_TYPE_HINT_PATH)?),
|
||||
TRANSFORM_TYPE => type_name = Some(yaml_string(value, TRANSFORM_TYPE)?),
|
||||
JSON2_TYPE_HINT_NULLABLE => {
|
||||
nullable = yaml_bool(value, JSON2_TYPE_HINT_NULLABLE)?;
|
||||
}
|
||||
TRANSFORM_DEFAULT => default = Some(value),
|
||||
TRANSFORM_INDEX => index = Some(value),
|
||||
_ => {
|
||||
return InvalidJson2TypeHintSnafu {
|
||||
reason: format!("unsupported field `{key}`"),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let path = path.context(InvalidJson2TypeHintSnafu {
|
||||
reason: "`path` must be set".to_string(),
|
||||
})?;
|
||||
let path = sql::parse_json2_type_hint_path(&path)
|
||||
.with_context(|_| ParseJson2TypeHintPathSnafu { path: path.clone() })?;
|
||||
let type_name = type_name.context(InvalidJson2TypeHintSnafu {
|
||||
reason: "`type` must be set".to_string(),
|
||||
})?;
|
||||
let type_ = parse_str_type(&type_name)?;
|
||||
ensure!(
|
||||
matches!(
|
||||
type_,
|
||||
ColumnDataType::String
|
||||
| ColumnDataType::Int64
|
||||
| ColumnDataType::Uint64
|
||||
| ColumnDataType::Float64
|
||||
| ColumnDataType::Boolean
|
||||
),
|
||||
InvalidJson2TypeHintSnafu {
|
||||
reason: format!("unsupported type `{type_name}`")
|
||||
}
|
||||
);
|
||||
let data_type = ColumnDataTypeWrapper::new(type_, None).into();
|
||||
let default_constraint = default
|
||||
.map(|value| parse_json2_type_hint_default(value, &type_))
|
||||
.transpose()?;
|
||||
if let Some(default_constraint) = &default_constraint {
|
||||
default_constraint.validate(&data_type, nullable)?;
|
||||
}
|
||||
|
||||
let inverted_index = if let Some(value) = index {
|
||||
let (index, options) = parse_transform_index(value)?;
|
||||
ensure!(
|
||||
index == Index::Inverted,
|
||||
InvalidJson2TypeHintSnafu {
|
||||
reason: format!("unsupported index `{index}`")
|
||||
}
|
||||
);
|
||||
lower_transform_index_options(index, &ColumnDataType::Json, options)?;
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
Ok(JsonTypeHint {
|
||||
path,
|
||||
data_type,
|
||||
nullable,
|
||||
default_constraint,
|
||||
inverted_index,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_json2_type_hint_default(
|
||||
value: &yaml_rust::Yaml,
|
||||
type_: &ColumnDataType,
|
||||
) -> Result<ColumnDefaultConstraint> {
|
||||
if value.is_null() {
|
||||
return Ok(ColumnDefaultConstraint::Value(Value::Null));
|
||||
}
|
||||
|
||||
let value = match value {
|
||||
yaml_rust::Yaml::Real(value) | yaml_rust::Yaml::String(value) => value.clone(),
|
||||
yaml_rust::Yaml::Integer(value) => value.to_string(),
|
||||
yaml_rust::Yaml::Boolean(value) => value.to_string(),
|
||||
_ => {
|
||||
return FieldMustBeTypeSnafu {
|
||||
field: TRANSFORM_DEFAULT,
|
||||
ty: "scalar",
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
};
|
||||
let value = api::v1::Value {
|
||||
value_data: Some(parse_str_value(type_, &value)?),
|
||||
};
|
||||
Ok(ColumnDefaultConstraint::Value(
|
||||
api::helper::pb_value_to_value_ref(&value, None).into(),
|
||||
))
|
||||
}
|
||||
|
||||
impl TryFrom<&yaml_rust::yaml::Hash> for Transform {
|
||||
type Error = Error;
|
||||
|
||||
@@ -430,6 +597,7 @@ impl TryFrom<&yaml_rust::yaml::Hash> for Transform {
|
||||
let mut on_failure = None;
|
||||
|
||||
let mut type_ = None;
|
||||
let mut json_settings = None;
|
||||
|
||||
for (k, v) in hash {
|
||||
let key = k
|
||||
@@ -445,8 +613,9 @@ impl TryFrom<&yaml_rust::yaml::Hash> for Transform {
|
||||
}
|
||||
|
||||
TRANSFORM_TYPE => {
|
||||
let t = yaml_string(v, TRANSFORM_TYPE)?;
|
||||
type_ = Some(parse_str_type(&t)?);
|
||||
let (parsed_type, parsed_json_settings) = parse_transform_type(v)?;
|
||||
type_ = Some(parsed_type);
|
||||
json_settings = parsed_json_settings;
|
||||
}
|
||||
|
||||
TRANSFORM_INDEX => {
|
||||
@@ -507,6 +676,7 @@ impl TryFrom<&yaml_rust::yaml::Hash> for Transform {
|
||||
let builder = Transform {
|
||||
fields,
|
||||
type_,
|
||||
json_settings,
|
||||
default: final_default,
|
||||
index,
|
||||
index_options,
|
||||
@@ -529,6 +699,72 @@ mod tests {
|
||||
docs[0].as_hash().unwrap().try_into()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_parses_json2_type_hints() {
|
||||
let transform = parse_transform(
|
||||
r#"
|
||||
field: payload
|
||||
type:
|
||||
json2:
|
||||
- path: "user.id"
|
||||
type: int64
|
||||
nullable: false
|
||||
default: 7
|
||||
index:
|
||||
type: inverted
|
||||
- path: 'attrs."http.status_code"'
|
||||
type: string
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(transform.type_, ColumnDataType::Json);
|
||||
let hints = transform.json_settings.as_ref().unwrap().type_hints();
|
||||
assert_eq!(hints.len(), 2);
|
||||
assert_eq!(hints[0].path, ["user", "id"]);
|
||||
assert_eq!(
|
||||
hints[0].data_type,
|
||||
datatypes::prelude::ConcreteDataType::int64_datatype()
|
||||
);
|
||||
assert!(!hints[0].nullable);
|
||||
assert_eq!(
|
||||
hints[0].default_constraint,
|
||||
Some(ColumnDefaultConstraint::Value(Value::Int64(7)))
|
||||
);
|
||||
assert!(hints[0].inverted_index);
|
||||
assert_eq!(hints[1].path, ["attrs", "http.status_code"]);
|
||||
assert!(hints[1].nullable);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_rejects_non_finite_json2_default() {
|
||||
for default in ["NaN", "1e9999"] {
|
||||
let err = parse_transform(&format!(
|
||||
r#"
|
||||
field: payload
|
||||
type:
|
||||
json2:
|
||||
- path: score
|
||||
type: float64
|
||||
default: {default}
|
||||
"#,
|
||||
))
|
||||
.unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(
|
||||
&err,
|
||||
Error::Datatypes {
|
||||
source: datatypes::error::Error::InvalidJson2Settings { .. },
|
||||
..
|
||||
}
|
||||
),
|
||||
"{err:?}"
|
||||
);
|
||||
assert!(err.to_string().contains("must be finite"), "{err}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_parses_legacy_string_index() {
|
||||
let transform = parse_transform(
|
||||
|
||||
@@ -150,6 +150,7 @@ impl GreptimeTransformer {
|
||||
let transform = Transform {
|
||||
fields: Fields::one(Field::new(greptime_timestamp().to_string(), None)),
|
||||
type_,
|
||||
json_settings: None,
|
||||
default,
|
||||
index: Some(Index::Time),
|
||||
index_options: None,
|
||||
@@ -225,6 +226,16 @@ impl GreptimeTransformer {
|
||||
&self,
|
||||
pipeline_map: &mut VrlValue,
|
||||
is_v1: bool,
|
||||
) -> Result<Vec<GreptimeValue>> {
|
||||
self.transform_mut_with_schema(pipeline_map, is_v1, &SchemaInfo::default(), None)
|
||||
}
|
||||
|
||||
pub(crate) fn transform_mut_with_schema(
|
||||
&self,
|
||||
pipeline_map: &mut VrlValue,
|
||||
is_v1: bool,
|
||||
schema_info: &SchemaInfo,
|
||||
table_suffix: Option<&str>,
|
||||
) -> Result<Vec<GreptimeValue>> {
|
||||
let mut values = vec![GreptimeValue { value_data: None }; self.schema.len()];
|
||||
let mut output_index = 0;
|
||||
@@ -236,7 +247,17 @@ impl GreptimeTransformer {
|
||||
// let keep us `get` here to be compatible with v1
|
||||
match pipeline_map.get(column_name) {
|
||||
Some(v) => {
|
||||
let value_data = coerce_value(v, transform)?;
|
||||
let json_settings = if transform.type_ == ColumnDataType::Json
|
||||
&& matches!(v, VrlValue::Array(_) | VrlValue::Object(_))
|
||||
{
|
||||
schema_info.json_settings_for_column(
|
||||
field.target_or_input_field(),
|
||||
table_suffix,
|
||||
)?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let value_data = coerce_value(v, transform, json_settings.as_ref())?;
|
||||
// every transform fields has only one output field
|
||||
values[output_index] = GreptimeValue { value_data };
|
||||
}
|
||||
@@ -276,6 +297,12 @@ impl GreptimeTransformer {
|
||||
&self.schema
|
||||
}
|
||||
|
||||
pub fn has_json_transform(&self) -> bool {
|
||||
self.transforms
|
||||
.iter()
|
||||
.any(|transform| transform.type_ == ColumnDataType::Json)
|
||||
}
|
||||
|
||||
pub fn transforms_mut(&mut self) -> &mut Transforms {
|
||||
&mut self.transforms
|
||||
}
|
||||
@@ -285,6 +312,7 @@ impl GreptimeTransformer {
|
||||
pub struct ColumnMetadata {
|
||||
column_schema: datatypes::schema::ColumnSchema,
|
||||
semantic_type: SemanticType,
|
||||
json_settings: OnceCell<JsonSettings>,
|
||||
}
|
||||
|
||||
impl From<ColumnSchema> for ColumnMetadata {
|
||||
@@ -311,10 +339,23 @@ impl From<ColumnSchema> for ColumnMetadata {
|
||||
Self {
|
||||
column_schema,
|
||||
semantic_type,
|
||||
json_settings: OnceCell::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ColumnMetadata {
|
||||
fn json_settings(&self) -> Result<&JsonSettings> {
|
||||
self.json_settings.get_or_try_init(|| {
|
||||
if let Some(extension) = self.column_schema.extension_type::<Json2ExtensionType>()? {
|
||||
Ok(extension.metadata().json_settings().clone())
|
||||
} else {
|
||||
Ok(parse_legacy_json2_settings(self.column_schema.metadata())?.unwrap_or_default())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<ColumnMetadata> for ColumnSchema {
|
||||
type Error = api::error::Error;
|
||||
|
||||
@@ -322,6 +363,7 @@ impl TryFrom<ColumnMetadata> for ColumnSchema {
|
||||
let ColumnMetadata {
|
||||
column_schema,
|
||||
semantic_type,
|
||||
..
|
||||
} = value;
|
||||
|
||||
let options = options_from_column_schema(&column_schema);
|
||||
@@ -348,8 +390,8 @@ pub struct SchemaInfo {
|
||||
pub schema: Vec<ColumnMetadata>,
|
||||
/// index of the column name
|
||||
pub index: HashMap<String, usize>,
|
||||
/// The pipeline's corresponding table (if already created). Useful to retrieve column schemas.
|
||||
table: Option<Arc<Table>>,
|
||||
/// Tables already looked up, keyed by their resolved suffix. Missing tables are cached as None.
|
||||
tables: HashMap<String, Option<Arc<Table>>>,
|
||||
}
|
||||
|
||||
impl SchemaInfo {
|
||||
@@ -357,7 +399,7 @@ impl SchemaInfo {
|
||||
Self {
|
||||
schema: Vec::with_capacity(capacity),
|
||||
index: HashMap::with_capacity(capacity),
|
||||
table: None,
|
||||
tables: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -369,16 +411,33 @@ impl SchemaInfo {
|
||||
Self {
|
||||
schema: schema_list.into_iter().map(Into::into).collect(),
|
||||
index,
|
||||
table: None,
|
||||
tables: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_table(&mut self, table: Option<Arc<Table>>) {
|
||||
self.table = table;
|
||||
self.set_table_for_suffix(String::new(), table);
|
||||
}
|
||||
|
||||
pub fn has_table_for_suffix(&self, table_suffix: &str) -> bool {
|
||||
self.tables.contains_key(table_suffix)
|
||||
}
|
||||
|
||||
pub fn set_table_for_suffix(&mut self, table_suffix: String, table: Option<Arc<Table>>) {
|
||||
self.tables.insert(table_suffix, table);
|
||||
}
|
||||
|
||||
fn table_for_suffix(&self, table_suffix: Option<&str>) -> Option<&Arc<Table>> {
|
||||
let table_suffix = table_suffix.unwrap_or_default();
|
||||
match self.tables.get(table_suffix) {
|
||||
Some(table) => table.as_ref(),
|
||||
None if !table_suffix.is_empty() => self.tables.get("").and_then(Option::as_ref),
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn find_column_schema_in_table(&self, column_name: &str) -> Option<ColumnMetadata> {
|
||||
if let Some(table) = &self.table
|
||||
if let Some(table) = self.table_for_suffix(None)
|
||||
&& let Some(i) = table.schema_ref().column_index_by_name(column_name)
|
||||
{
|
||||
let column_schema = table.schema_ref().column_schemas()[i].clone();
|
||||
@@ -394,12 +453,37 @@ impl SchemaInfo {
|
||||
Some(ColumnMetadata {
|
||||
column_schema,
|
||||
semantic_type,
|
||||
json_settings: OnceCell::new(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn json_settings_for_column(
|
||||
&self,
|
||||
column_name: &str,
|
||||
table_suffix: Option<&str>,
|
||||
) -> Result<Option<JsonSettings>> {
|
||||
let Some(column_schema) = self
|
||||
.table_for_suffix(table_suffix)
|
||||
.and_then(|table| table.schema_ref().column_schema_by_name(column_name))
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !column_schema.data_type.is_json2() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Some(extension) = column_schema.extension_type::<Json2ExtensionType>()? {
|
||||
Ok(Some(extension.metadata().json_settings().clone()))
|
||||
} else {
|
||||
Ok(Some(
|
||||
parse_legacy_json2_settings(column_schema.metadata())?.unwrap_or_default(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn column_schemas(&self) -> api::error::Result<Vec<ColumnSchema>> {
|
||||
self.schema
|
||||
.iter()
|
||||
@@ -444,6 +528,7 @@ fn resolve_schema(
|
||||
ColumnMetadata {
|
||||
column_schema,
|
||||
semantic_type,
|
||||
json_settings: OnceCell::new(),
|
||||
}
|
||||
});
|
||||
let key = column.to_string();
|
||||
@@ -499,12 +584,12 @@ pub(crate) fn values_to_rows(
|
||||
// Single object: extract ContextOpt and table_suffix
|
||||
let mut result = std::collections::HashMap::new();
|
||||
|
||||
let mut opt = match ContextOpt::from_pipeline_map_to_opt(&mut values) {
|
||||
let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, &values);
|
||||
let opt = match ContextOpt::from_pipeline_map_to_opt(&mut values) {
|
||||
Ok(r) => r,
|
||||
Err(e) => return if skip_error { Ok(result) } else { Err(e) },
|
||||
};
|
||||
|
||||
let table_suffix = opt.resolve_table_suffix(tablesuffix_template, &values);
|
||||
let row = match values_to_row(schema_info, values, pipeline_ctx, row, need_calc_ts) {
|
||||
Ok(r) => r,
|
||||
Err(e) => return if skip_error { Ok(result) } else { Err(e) },
|
||||
@@ -528,11 +613,11 @@ pub(crate) fn values_to_rows(
|
||||
}
|
||||
|
||||
// Extract ContextOpt and table_suffix for this element
|
||||
let mut opt = unwrap_or_continue_if_err!(
|
||||
let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, &value);
|
||||
let opt = unwrap_or_continue_if_err!(
|
||||
ContextOpt::from_pipeline_map_to_opt(&mut value),
|
||||
skip_error
|
||||
);
|
||||
let table_suffix = opt.resolve_table_suffix(tablesuffix_template, &value);
|
||||
let transformed_row = unwrap_or_continue_if_err!(
|
||||
values_to_row(schema_info, value, pipeline_ctx, row.clone(), need_calc_ts),
|
||||
skip_error
|
||||
@@ -691,38 +776,29 @@ fn resolve_value(
|
||||
}
|
||||
|
||||
VrlValue::Array(_) | VrlValue::Object(_) => {
|
||||
let is_json2 = schema_info
|
||||
.find_column_schema_in_table(&column_name)
|
||||
// TODO(LFC): Default to JSON2 for auto-created tables.
|
||||
.is_some_and(|x| {
|
||||
matches!(
|
||||
&x.column_schema.data_type,
|
||||
ConcreteDataType::Json(column_type) if column_type.is_json2()
|
||||
)
|
||||
});
|
||||
let index = index.or_else(|| {
|
||||
let column = schema_info.find_column_schema_in_table(&column_name)?;
|
||||
let index = schema_info.schema.len();
|
||||
schema_info.schema.push(column);
|
||||
schema_info.index.insert(column_name.clone(), index);
|
||||
Some(index)
|
||||
});
|
||||
// TODO(LFC): Default to JSON2 for auto-created tables.
|
||||
let json2_index = index.filter(|&index| {
|
||||
matches!(
|
||||
&schema_info.schema[index].column_schema.data_type,
|
||||
ConcreteDataType::Json(column_type) if column_type.is_json2()
|
||||
)
|
||||
});
|
||||
|
||||
let value = if is_json2 {
|
||||
let value = if let Some(index) = json2_index {
|
||||
let value: serde_json::Value = value.try_into().map_err(|e: StdError| {
|
||||
CoerceIncompatibleTypesSnafu { msg: e.to_string() }.build()
|
||||
})?;
|
||||
let value =
|
||||
if let Some(column) = schema_info.find_column_schema_in_table(&column_name) {
|
||||
if let Some(extension) = column
|
||||
.column_schema
|
||||
.extension_type::<Json2ExtensionType>()?
|
||||
{
|
||||
extension.metadata().json_settings().encode(value)?
|
||||
} else {
|
||||
parse_legacy_json2_settings(column.column_schema.metadata())?
|
||||
.unwrap_or_default()
|
||||
.encode(value)?
|
||||
}
|
||||
} else {
|
||||
JsonSettings::default().encode(value)?
|
||||
};
|
||||
let value = schema_info.schema[index].json_settings()?.encode(value)?;
|
||||
|
||||
resolve_schema(
|
||||
index,
|
||||
Some(index),
|
||||
p_ctx,
|
||||
&column_name,
|
||||
&ConcreteDataType::json2(Default::default()),
|
||||
@@ -810,6 +886,7 @@ fn identity_pipeline_inner(
|
||||
schema_info.schema.push(ColumnMetadata {
|
||||
column_schema,
|
||||
semantic_type: SemanticType::Timestamp,
|
||||
json_settings: OnceCell::new(),
|
||||
});
|
||||
|
||||
let mut opt_map = HashMap::new();
|
||||
@@ -907,7 +984,7 @@ pub fn flatten_object(object: VrlValue, max_nested_levels: usize) -> Result<VrlV
|
||||
Ok(VrlValue::Object(flattened))
|
||||
}
|
||||
|
||||
fn vrl_value_to_serde_json(value: &VrlValue) -> serde_json_crate::Value {
|
||||
pub(crate) fn vrl_value_to_serde_json(value: &VrlValue) -> serde_json_crate::Value {
|
||||
match value {
|
||||
VrlValue::Null => serde_json_crate::Value::Null,
|
||||
VrlValue::Boolean(b) => serde_json_crate::Value::Bool(*b),
|
||||
@@ -980,10 +1057,125 @@ fn do_flatten_object(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use api::v1::SemanticType;
|
||||
use common_recordbatch::RecordBatch;
|
||||
use datatypes::extension::json::JsonMetadata;
|
||||
use datatypes::json::JsonTypeHint;
|
||||
use datatypes::schema::{ColumnSchema as DatatypeColumnSchema, Schema};
|
||||
use table::test_util::MemTable;
|
||||
|
||||
use super::*;
|
||||
use crate::{PipelineDefinition, identity_pipeline};
|
||||
|
||||
#[test]
|
||||
fn test_column_metadata_caches_json_settings() -> Result<()> {
|
||||
let column = ColumnMetadata {
|
||||
column_schema: datatypes::schema::ColumnSchema::new(
|
||||
"data",
|
||||
ConcreteDataType::json2(Default::default()),
|
||||
true,
|
||||
),
|
||||
semantic_type: SemanticType::Field,
|
||||
json_settings: OnceCell::new(),
|
||||
};
|
||||
|
||||
let first = column.json_settings()?;
|
||||
let second = column.json_settings()?;
|
||||
assert!(std::ptr::eq(first, second));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_json2_uses_destination_table_settings() {
|
||||
let table = |name: &str, settings: JsonSettings, sample: serde_json::Value| {
|
||||
let data_type = settings.encode(sample).unwrap().data_type();
|
||||
let mut column_schema = DatatypeColumnSchema::new("payload", data_type, true);
|
||||
column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new(
|
||||
JsonMetadata::new(settings),
|
||||
)));
|
||||
MemTable::table(
|
||||
name,
|
||||
RecordBatch::new_empty(Arc::new(Schema::new(vec![column_schema]))),
|
||||
)
|
||||
};
|
||||
let int_settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["age".to_string()],
|
||||
data_type: ConcreteDataType::int64_datatype(),
|
||||
nullable: false,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let mut schema_info = SchemaInfo::default();
|
||||
schema_info.set_table(Some(table(
|
||||
"events",
|
||||
int_settings,
|
||||
serde_json::json!({"age": 42}),
|
||||
)));
|
||||
let string_settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["age".to_string()],
|
||||
data_type: ConcreteDataType::string_datatype(),
|
||||
nullable: false,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
schema_info.set_table_for_suffix(
|
||||
"_mobile".to_string(),
|
||||
Some(table(
|
||||
"events_mobile",
|
||||
string_settings,
|
||||
serde_json::json!({"age": "42"}),
|
||||
)),
|
||||
);
|
||||
|
||||
let pipeline = crate::parse(&crate::Content::Yaml(
|
||||
r#"
|
||||
transform:
|
||||
- field: source, payload
|
||||
type: json2
|
||||
table_suffix: _${device}
|
||||
"#,
|
||||
))
|
||||
.unwrap();
|
||||
let (pipeline, _, pipeline_definition, pipeline_params) = crate::setup_pipeline!(pipeline);
|
||||
let pipeline_context =
|
||||
PipelineContext::new(&pipeline_definition, &pipeline_params, Channel::Unknown);
|
||||
|
||||
let error = pipeline
|
||||
.exec_mut(
|
||||
serde_json::json!({"source": {"age": "42"}}).into(),
|
||||
&pipeline_context,
|
||||
&mut schema_info,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
error.to_string().contains("does not match JSON2 type hint"),
|
||||
"{error:?}"
|
||||
);
|
||||
|
||||
let mut rows = pipeline
|
||||
.exec_mut(
|
||||
serde_json::json!({"source": {"age": "42"}, "device": "mobile"}).into(),
|
||||
&pipeline_context,
|
||||
&mut schema_info,
|
||||
)
|
||||
.unwrap()
|
||||
.into_transformed()
|
||||
.unwrap();
|
||||
let (row, table_suffix) = rows.swap_remove(0);
|
||||
assert_eq!(table_suffix.as_deref(), Some("_mobile"));
|
||||
assert!(matches!(
|
||||
&row.values[0].value_data,
|
||||
Some(ValueData::JsonValue(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identify_pipeline() {
|
||||
let params = GreptimePipelineParams::default();
|
||||
|
||||
@@ -12,10 +12,18 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::column_data_type_extension::TypeExt;
|
||||
use api::v1::column_def::{options_from_fulltext, options_from_inverted, options_from_skipping};
|
||||
use api::v1::{ColumnDataTypeExtension, ColumnOptions, JsonTypeExtension};
|
||||
use arrow_schema::extension::{
|
||||
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
|
||||
};
|
||||
use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
|
||||
use datatypes::json::JsonSettings;
|
||||
use datatypes::schema::{FulltextOptions, SkippingIndexOptions};
|
||||
use datatypes::value::Value;
|
||||
use greptime_proto::v1::value::ValueData;
|
||||
use greptime_proto::v1::{ColumnDataType, ColumnSchema, SemanticType};
|
||||
use snafu::{OptionExt, ResultExt, ensure};
|
||||
@@ -28,7 +36,9 @@ use crate::error::{
|
||||
UnsupportedTypeInPipelineSnafu, VrlRegexValueSnafu,
|
||||
};
|
||||
use crate::etl::transform::index::Index;
|
||||
use crate::etl::transform::transformer::greptime::vrl_value_to_jsonb_value;
|
||||
use crate::etl::transform::transformer::greptime::{
|
||||
vrl_value_to_jsonb_value, vrl_value_to_serde_json,
|
||||
};
|
||||
use crate::etl::transform::{OnFailure, Transform, TransformIndexOptions};
|
||||
|
||||
pub(crate) fn coerce_columns(transform: &Transform) -> Result<Vec<ColumnSchema>> {
|
||||
@@ -128,7 +138,7 @@ fn build_skipping_index_options(transform: &Transform) -> Result<SkippingIndexOp
|
||||
fn coerce_options(transform: &Transform) -> Result<Option<ColumnOptions>> {
|
||||
validate_transform_index_state(transform)?;
|
||||
|
||||
match transform.index {
|
||||
let mut options = match transform.index {
|
||||
Some(Index::Fulltext) => {
|
||||
let options = build_fulltext_index_options(transform)?;
|
||||
options_from_fulltext(&options).context(ColumnOptionsSnafu)
|
||||
@@ -139,10 +149,32 @@ fn coerce_options(transform: &Transform) -> Result<Option<ColumnOptions>> {
|
||||
}
|
||||
Some(Index::Inverted) => Ok(Some(options_from_inverted())),
|
||||
_ => Ok(None),
|
||||
}?;
|
||||
|
||||
if transform.type_ == ColumnDataType::Json {
|
||||
let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(
|
||||
transform.json_settings.clone().unwrap_or_default(),
|
||||
)));
|
||||
let options = options.get_or_insert_default();
|
||||
options.options.insert(
|
||||
EXTENSION_TYPE_NAME_KEY.to_string(),
|
||||
Json2ExtensionType::NAME.to_string(),
|
||||
);
|
||||
if let Some(metadata) = extension.serialize_metadata() {
|
||||
options
|
||||
.options
|
||||
.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), metadata);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(options)
|
||||
}
|
||||
|
||||
pub(crate) fn coerce_value(val: &VrlValue, transform: &Transform) -> Result<Option<ValueData>> {
|
||||
pub(crate) fn coerce_value(
|
||||
val: &VrlValue,
|
||||
transform: &Transform,
|
||||
json_settings: Option<&JsonSettings>,
|
||||
) -> Result<Option<ValueData>> {
|
||||
match val {
|
||||
VrlValue::Null => Ok(None),
|
||||
VrlValue::Integer(n) => coerce_i64_value(*n, transform),
|
||||
@@ -169,7 +201,9 @@ pub(crate) fn coerce_value(val: &VrlValue, transform: &Transform) -> Result<Opti
|
||||
}
|
||||
.fail(),
|
||||
},
|
||||
VrlValue::Array(_) | VrlValue::Object(_) => coerce_json_value(val, transform),
|
||||
VrlValue::Array(_) | VrlValue::Object(_) => {
|
||||
coerce_json_value(val, transform, json_settings)
|
||||
}
|
||||
VrlValue::Regex(_) => VrlRegexValueSnafu.fail(),
|
||||
}
|
||||
}
|
||||
@@ -205,7 +239,7 @@ fn coerce_bool_value(b: bool, transform: &Transform) -> Result<Option<ValueData>
|
||||
}
|
||||
},
|
||||
|
||||
ColumnDataType::Binary => {
|
||||
ColumnDataType::Binary | ColumnDataType::Json => {
|
||||
return CoerceJsonTypeToSnafu {
|
||||
ty: transform.type_.as_str_name(),
|
||||
}
|
||||
@@ -294,7 +328,7 @@ fn coerce_i64_value(n: i64, transform: &Transform) -> Result<Option<ValueData>>
|
||||
ColumnDataType::TimestampMillisecond => ValueData::TimestampMillisecondValue(n),
|
||||
ColumnDataType::TimestampSecond => ValueData::TimestampSecondValue(n),
|
||||
|
||||
ColumnDataType::Binary => {
|
||||
ColumnDataType::Binary | ColumnDataType::Json => {
|
||||
return CoerceJsonTypeToSnafu {
|
||||
ty: transform.type_.as_str_name(),
|
||||
}
|
||||
@@ -363,7 +397,7 @@ fn coerce_u64_value(n: u64, transform: &Transform) -> Result<Option<ValueData>>
|
||||
Err(_) => return integer_out_of_range(n, transform),
|
||||
},
|
||||
|
||||
ColumnDataType::Binary => {
|
||||
ColumnDataType::Binary | ColumnDataType::Json => {
|
||||
return CoerceJsonTypeToSnafu {
|
||||
ty: transform.type_.as_str_name(),
|
||||
}
|
||||
@@ -407,7 +441,7 @@ fn coerce_f64_value(n: f64, transform: &Transform) -> Result<Option<ValueData>>
|
||||
}
|
||||
},
|
||||
|
||||
ColumnDataType::Binary => {
|
||||
ColumnDataType::Binary | ColumnDataType::Json => {
|
||||
return CoerceJsonTypeToSnafu {
|
||||
ty: transform.type_.as_str_name(),
|
||||
}
|
||||
@@ -486,7 +520,7 @@ fn coerce_string_value(s: &str, transform: &Transform) -> Result<Option<ValueDat
|
||||
None => CoerceUnsupportedEpochTypeSnafu { ty: "String" }.fail(),
|
||||
},
|
||||
|
||||
ColumnDataType::Binary => CoerceStringToTypeSnafu {
|
||||
ColumnDataType::Binary | ColumnDataType::Json => CoerceStringToTypeSnafu {
|
||||
s,
|
||||
ty: transform.type_.as_str_name(),
|
||||
}
|
||||
@@ -496,23 +530,48 @@ fn coerce_string_value(s: &str, transform: &Transform) -> Result<Option<ValueDat
|
||||
}
|
||||
}
|
||||
|
||||
fn coerce_json_value(v: &VrlValue, transform: &Transform) -> Result<Option<ValueData>> {
|
||||
match &transform.type_ {
|
||||
ColumnDataType::Binary => (),
|
||||
fn coerce_json_value(
|
||||
v: &VrlValue,
|
||||
transform: &Transform,
|
||||
json_settings: Option<&JsonSettings>,
|
||||
) -> Result<Option<ValueData>> {
|
||||
let value = match transform.type_ {
|
||||
ColumnDataType::Binary => {
|
||||
let data: jsonb::Value = vrl_value_to_jsonb_value(v);
|
||||
ValueData::BinaryValue(data.to_vec())
|
||||
}
|
||||
ColumnDataType::Json => {
|
||||
let json = vrl_value_to_serde_json(v);
|
||||
let encoded = if let Some(settings) = json_settings.or(transform.json_settings.as_ref())
|
||||
{
|
||||
settings.encode(json)
|
||||
} else {
|
||||
JsonSettings::default().encode(json)
|
||||
};
|
||||
let value = match encoded {
|
||||
Ok(value) => value,
|
||||
Err(error) => return handle_coercion_failure(transform, error.into()),
|
||||
};
|
||||
let Value::Json(value) = value else {
|
||||
unreachable!()
|
||||
};
|
||||
ValueData::JsonValue(api::helper::encode_json_value(*value))
|
||||
}
|
||||
t => {
|
||||
return CoerceTypeToJsonSnafu {
|
||||
ty: t.as_str_name(),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
}
|
||||
let data: jsonb::Value = vrl_value_to_jsonb_value(v);
|
||||
Ok(Some(ValueData::BinaryValue(data.to_vec())))
|
||||
};
|
||||
Ok(Some(value))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use datatypes::data_type::ConcreteDataType;
|
||||
use datatypes::json::JsonTypeHint;
|
||||
use datatypes::schema::{FulltextAnalyzer, FulltextBackend, SkippingIndexType};
|
||||
use vrl::prelude::Bytes;
|
||||
|
||||
@@ -523,6 +582,7 @@ mod tests {
|
||||
Transform {
|
||||
fields: Fields::default(),
|
||||
type_,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: None,
|
||||
index_options: None,
|
||||
@@ -666,6 +726,7 @@ mod tests {
|
||||
let transform = Transform {
|
||||
fields: Fields::default(),
|
||||
type_: ColumnDataType::Int32,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: None,
|
||||
index_options: None,
|
||||
@@ -676,14 +737,14 @@ mod tests {
|
||||
// valid string
|
||||
{
|
||||
let val = VrlValue::Integer(123);
|
||||
let result = coerce_value(&val, &transform).unwrap();
|
||||
let result = coerce_value(&val, &transform, None).unwrap();
|
||||
assert_eq!(result, Some(ValueData::I32Value(123)));
|
||||
}
|
||||
|
||||
// invalid string
|
||||
{
|
||||
let val = VrlValue::Bytes(Bytes::from("hello"));
|
||||
let result = coerce_value(&val, &transform);
|
||||
let result = coerce_value(&val, &transform, None);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
@@ -693,6 +754,7 @@ mod tests {
|
||||
let transform = Transform {
|
||||
fields: Fields::default(),
|
||||
type_: ColumnDataType::Int32,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: None,
|
||||
index_options: None,
|
||||
@@ -701,15 +763,43 @@ mod tests {
|
||||
};
|
||||
|
||||
let val = VrlValue::Bytes(Bytes::from("hello"));
|
||||
let result = coerce_value(&val, &transform).unwrap();
|
||||
let result = coerce_value(&val, &transform, None).unwrap();
|
||||
assert_eq!(result, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_json2_with_on_failure() {
|
||||
let settings = JsonSettings::try_new(
|
||||
vec![JsonTypeHint {
|
||||
path: vec!["age".to_string()],
|
||||
data_type: ConcreteDataType::int64_datatype(),
|
||||
nullable: false,
|
||||
default_constraint: None,
|
||||
inverted_index: false,
|
||||
}],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let mut transform = transform(ColumnDataType::Json);
|
||||
transform.json_settings = Some(settings);
|
||||
transform.on_failure = Some(OnFailure::Ignore);
|
||||
let value: VrlValue = serde_json::json!({"age": "42"}).into();
|
||||
|
||||
assert_eq!(coerce_value(&value, &transform, None).unwrap(), None);
|
||||
|
||||
transform.on_failure = Some(OnFailure::Default);
|
||||
assert_eq!(
|
||||
coerce_value(&value, &transform, None).unwrap(),
|
||||
Some(ValueData::JsonValue(Default::default()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_string_with_on_failure_default() {
|
||||
let mut transform = Transform {
|
||||
fields: Fields::default(),
|
||||
type_: ColumnDataType::Int32,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: None,
|
||||
index_options: None,
|
||||
@@ -720,7 +810,7 @@ mod tests {
|
||||
// with no explicit default value
|
||||
{
|
||||
let val = VrlValue::Bytes(Bytes::from("hello"));
|
||||
let result = coerce_value(&val, &transform).unwrap();
|
||||
let result = coerce_value(&val, &transform, None).unwrap();
|
||||
assert_eq!(result, Some(ValueData::I32Value(0)));
|
||||
}
|
||||
|
||||
@@ -728,7 +818,7 @@ mod tests {
|
||||
{
|
||||
transform.default = Some(ValueData::I32Value(42));
|
||||
let val = VrlValue::Bytes(Bytes::from("hello"));
|
||||
let result = coerce_value(&val, &transform).unwrap();
|
||||
let result = coerce_value(&val, &transform, None).unwrap();
|
||||
assert_eq!(result, Some(ValueData::I32Value(42)));
|
||||
}
|
||||
}
|
||||
@@ -738,6 +828,7 @@ mod tests {
|
||||
let transform = Transform {
|
||||
fields: Fields::default(),
|
||||
type_: ColumnDataType::String,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: Some(Index::Fulltext),
|
||||
index_options: Some(TransformIndexOptions::Fulltext(
|
||||
@@ -769,6 +860,7 @@ mod tests {
|
||||
let transform = Transform {
|
||||
fields: Fields::default(),
|
||||
type_: ColumnDataType::Int64,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: Some(Index::Skipping),
|
||||
index_options: Some(TransformIndexOptions::Skipping(
|
||||
@@ -792,6 +884,7 @@ mod tests {
|
||||
let transform = Transform {
|
||||
fields: Fields::default(),
|
||||
type_: ColumnDataType::String,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: Some(Index::Fulltext),
|
||||
index_options: Some(TransformIndexOptions::Skipping(
|
||||
@@ -809,6 +902,7 @@ mod tests {
|
||||
let transform = Transform {
|
||||
fields: Fields::default(),
|
||||
type_: ColumnDataType::String,
|
||||
json_settings: None,
|
||||
default: None,
|
||||
index: None,
|
||||
index_options: Some(TransformIndexOptions::Fulltext(
|
||||
|
||||
@@ -102,6 +102,7 @@ pub fn parse_str_type(t: &str) -> Result<ColumnDataType> {
|
||||
// We only consider object and array to be json types. and use Map to represent json
|
||||
// TODO(qtang): Needs to be defined with better semantics
|
||||
"json" => Ok(ColumnDataType::Binary),
|
||||
"json2" => Ok(ColumnDataType::Json),
|
||||
|
||||
_ => ValueParseTypeSnafu { t }.fail(),
|
||||
}
|
||||
|
||||
@@ -27,7 +27,8 @@ pub use etl::transform::GreptimeTransformer;
|
||||
pub use etl::transform::transformer::greptime::{GreptimePipelineParams, SchemaInfo};
|
||||
pub use etl::transform::transformer::identity_pipeline;
|
||||
pub use etl::{
|
||||
Content, DispatchedTo, Pipeline, PipelineExecOutput, TransformedOutput, TransformerMode, parse,
|
||||
Content, DispatchedTo, Pipeline, PipelineExecOutput, PipelineProcessOutput, TransformedOutput,
|
||||
TransformerMode, parse,
|
||||
};
|
||||
pub use manager::{
|
||||
GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME, GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME,
|
||||
|
||||
@@ -12,14 +12,14 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use common_telemetry::debug;
|
||||
use datatypes::timestamp::TimestampNanosecond;
|
||||
use moka::sync::Cache;
|
||||
use moka::future::Cache;
|
||||
|
||||
use crate::error::{MultiPipelineWithDiffSchemaSnafu, Result};
|
||||
use crate::error::{CacheLoadSnafu, MultiPipelineWithDiffSchemaSnafu, Result};
|
||||
use crate::etl::Pipeline;
|
||||
use crate::manager::PipelineVersion;
|
||||
use crate::table::EMPTY_SCHEMA_NAME;
|
||||
@@ -32,6 +32,11 @@ const PIPELINES_CACHE_TTL: Duration = Duration::from_secs(10);
|
||||
|
||||
/// Pipeline cache is located on a separate file on purpose,
|
||||
/// to encapsulate inner cache. Only public methods are exposed.
|
||||
///
|
||||
/// `pipelines` and `original_pipelines` are keyed by the *requested* schema so
|
||||
/// a lookup is a single key probe, as [`Cache::try_get_with`] requires;
|
||||
/// resolving it to a stored schema is the loader's job. `failover_cache` has no
|
||||
/// loader and keeps the stored-schema key.
|
||||
pub(crate) struct PipelineCache {
|
||||
pipelines: Cache<String, Arc<Pipeline>>,
|
||||
original_pipelines: Cache<String, PipelineContent>,
|
||||
@@ -68,163 +73,242 @@ impl PipelineCache {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn insert_pipeline_cache(
|
||||
/// Concurrent misses on the same key share one `init` call.
|
||||
pub(crate) async fn get_pipeline_with(
|
||||
&self,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
version: PipelineVersion,
|
||||
pipeline: Arc<Pipeline>,
|
||||
with_latest: bool,
|
||||
) {
|
||||
insert_cache_generic(
|
||||
&self.pipelines,
|
||||
schema,
|
||||
name,
|
||||
version,
|
||||
pipeline.clone(),
|
||||
with_latest,
|
||||
);
|
||||
init: impl Future<Output = Result<Arc<Pipeline>>>,
|
||||
) -> Result<Arc<Pipeline>> {
|
||||
let key = generate_pipeline_cache_key(schema, name, version);
|
||||
self.pipelines
|
||||
.try_get_with(key, init)
|
||||
.await
|
||||
.map_err(|error| CacheLoadSnafu { error }.build())
|
||||
}
|
||||
|
||||
pub(crate) fn insert_pipeline_str_cache(&self, pipeline: &PipelineContent, with_latest: bool) {
|
||||
let schema = pipeline.schema.as_str();
|
||||
let name = pipeline.name.as_str();
|
||||
let version = pipeline.version;
|
||||
insert_cache_generic(
|
||||
&self.original_pipelines,
|
||||
schema,
|
||||
name,
|
||||
Some(version),
|
||||
pipeline.clone(),
|
||||
with_latest,
|
||||
);
|
||||
insert_cache_generic(
|
||||
&self.failover_cache,
|
||||
schema,
|
||||
name,
|
||||
Some(version),
|
||||
pipeline.clone(),
|
||||
with_latest,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn get_pipeline_cache(
|
||||
/// Concurrent misses on the same key share one `init` call.
|
||||
pub(crate) async fn get_pipeline_str_with(
|
||||
&self,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
version: PipelineVersion,
|
||||
) -> Result<Option<Arc<Pipeline>>> {
|
||||
get_cache_generic(&self.pipelines, schema, name, version)
|
||||
init: impl Future<Output = Result<PipelineContent>>,
|
||||
) -> Result<PipelineContent> {
|
||||
let key = generate_pipeline_cache_key(schema, name, version);
|
||||
self.original_pipelines
|
||||
.try_get_with(key, init)
|
||||
.await
|
||||
.map_err(|error| CacheLoadSnafu { error }.build())
|
||||
}
|
||||
|
||||
pub(crate) fn get_failover_cache(
|
||||
/// Resolves across schemas, unlike the loaded caches: a pipeline stored
|
||||
/// under the empty schema is reachable from any schema.
|
||||
pub(crate) async fn get_failover_cache(
|
||||
&self,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
version: PipelineVersion,
|
||||
) -> Result<Option<PipelineContent>> {
|
||||
get_cache_generic(&self.failover_cache, schema, name, version)
|
||||
}
|
||||
for key in [
|
||||
generate_pipeline_cache_key(EMPTY_SCHEMA_NAME, name, version),
|
||||
generate_pipeline_cache_key(schema, name, version),
|
||||
] {
|
||||
if let Some(content) = self.failover_cache.get(&key).await {
|
||||
return Ok(Some(content));
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn get_pipeline_str_cache(
|
||||
&self,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
version: PipelineVersion,
|
||||
) -> Result<Option<PipelineContent>> {
|
||||
get_cache_generic(&self.original_pipelines, schema, name, version)
|
||||
}
|
||||
|
||||
// remove cache with version and latest in all schemas
|
||||
pub(crate) fn remove_cache(&self, name: &str, version: PipelineVersion) {
|
||||
let version_suffix = generate_pipeline_cache_key_suffix(name, version);
|
||||
let latest_suffix = generate_pipeline_cache_key_suffix(name, None);
|
||||
|
||||
let ks = self
|
||||
.pipelines
|
||||
// Stored under some other schema; unambiguous only if exactly one has it.
|
||||
let suffix = generate_pipeline_cache_key_suffix(name, version);
|
||||
let mut found = self
|
||||
.failover_cache
|
||||
.iter()
|
||||
.filter_map(|(k, _)| {
|
||||
if k.ends_with(&version_suffix) || k.ends_with(&latest_suffix) {
|
||||
Some(k.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.filter(|(k, _)| k.ends_with(&suffix))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
for k in ks {
|
||||
let k = k.as_str();
|
||||
self.pipelines.remove(k);
|
||||
self.original_pipelines.remove(k);
|
||||
self.failover_cache.remove(k);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn insert_cache_generic<T: Clone + Send + Sync + 'static>(
|
||||
cache: &Cache<String, T>,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
version: PipelineVersion,
|
||||
value: T,
|
||||
with_latest: bool,
|
||||
) {
|
||||
let k = generate_pipeline_cache_key(schema, name, version);
|
||||
cache.insert(k, value.clone());
|
||||
if with_latest {
|
||||
let k = generate_pipeline_cache_key(schema, name, None);
|
||||
cache.insert(k, value);
|
||||
}
|
||||
}
|
||||
|
||||
fn get_cache_generic<T: Clone + Send + Sync + 'static>(
|
||||
cache: &Cache<String, T>,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
version: PipelineVersion,
|
||||
) -> Result<Option<T>> {
|
||||
// lets try empty schema first
|
||||
let emp_key = generate_pipeline_cache_key(EMPTY_SCHEMA_NAME, name, version);
|
||||
if let Some(value) = cache.get(&emp_key) {
|
||||
return Ok(Some(value));
|
||||
}
|
||||
// use input schema
|
||||
let schema_k = generate_pipeline_cache_key(schema, name, version);
|
||||
if let Some(value) = cache.get(&schema_k) {
|
||||
return Ok(Some(value));
|
||||
}
|
||||
|
||||
// try all schemas
|
||||
let suffix_key = generate_pipeline_cache_key_suffix(name, version);
|
||||
let mut ks = cache
|
||||
.iter()
|
||||
.filter(|e| e.0.ends_with(&suffix_key))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
match ks.len() {
|
||||
0 => Ok(None),
|
||||
1 => {
|
||||
let (_, value) = ks.remove(0);
|
||||
Ok(Some(value))
|
||||
}
|
||||
_ => {
|
||||
debug!(
|
||||
"caches keys: {:?}, emp key: {:?}, schema key: {:?}, suffix key: {:?}",
|
||||
cache.iter().map(|e| e.0).collect::<Vec<_>>(),
|
||||
emp_key,
|
||||
schema_k,
|
||||
suffix_key
|
||||
);
|
||||
MultiPipelineWithDiffSchemaSnafu {
|
||||
match found.len() {
|
||||
0 => Ok(None),
|
||||
1 => Ok(Some(found.remove(0).1)),
|
||||
_ => MultiPipelineWithDiffSchemaSnafu {
|
||||
name: name.to_string(),
|
||||
current_schema: schema.to_string(),
|
||||
schemas: ks
|
||||
schemas: found
|
||||
.iter()
|
||||
.filter_map(|(k, _)| k.split_once('/').map(|k| k.0))
|
||||
.collect::<Vec<_>>()
|
||||
.join(","),
|
||||
}
|
||||
.fail()?
|
||||
.fail(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn insert_failover_cache(&self, content: PipelineContent, with_latest: bool) {
|
||||
let versioned =
|
||||
generate_pipeline_cache_key(&content.schema, &content.name, Some(content.version));
|
||||
let latest = generate_pipeline_cache_key(&content.schema, &content.name, None);
|
||||
|
||||
self.failover_cache.insert(versioned, content.clone()).await;
|
||||
if with_latest {
|
||||
self.failover_cache.insert(latest, content).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Dropping the stale `latest` aliases also clears the failover entries, so
|
||||
/// the new version is written back: an outage before the first read-back
|
||||
/// would otherwise have nothing to fall back on.
|
||||
pub(crate) async fn on_pipeline_created(&self, content: PipelineContent) {
|
||||
self.invalidate(&content.name, None).await;
|
||||
self.insert_failover_cache(content, true).await;
|
||||
}
|
||||
|
||||
/// Sweeps every schema and all three caches: the `latest` alias always,
|
||||
/// plus `version` when given.
|
||||
pub(crate) async fn invalidate(&self, name: &str, version: PipelineVersion) {
|
||||
let mut suffixes = vec![generate_pipeline_cache_key_suffix(name, None)];
|
||||
if version.is_some() {
|
||||
suffixes.push(generate_pipeline_cache_key_suffix(name, version));
|
||||
}
|
||||
|
||||
let ks = self
|
||||
.pipelines
|
||||
.iter()
|
||||
.map(|(k, _)| k)
|
||||
.chain(self.original_pipelines.iter().map(|(k, _)| k))
|
||||
.chain(self.failover_cache.iter().map(|(k, _)| k))
|
||||
.filter(|k| suffixes.iter().any(|suffix| k.ends_with(suffix)))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
for k in ks {
|
||||
let k = k.as_str();
|
||||
self.pipelines.invalidate(k).await;
|
||||
self.original_pipelines.invalidate(k).await;
|
||||
self.failover_cache.invalidate(k).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use tokio::sync::Barrier;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Stored under the empty schema, i.e. visible from every schema.
|
||||
fn content_at(version: i64) -> PipelineContent {
|
||||
PipelineContent {
|
||||
name: "p".to_string(),
|
||||
content: "transform:".to_string(),
|
||||
version: TimestampNanosecond::new(version),
|
||||
schema: EMPTY_SCHEMA_NAME.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn test_concurrent_misses_run_one_loader() {
|
||||
const CONCURRENCY: usize = 8;
|
||||
|
||||
let cache = Arc::new(PipelineCache::new());
|
||||
let loads = Arc::new(AtomicUsize::new(0));
|
||||
let barrier = Arc::new(Barrier::new(CONCURRENCY));
|
||||
|
||||
let handles = (0..CONCURRENCY)
|
||||
.map(|_| {
|
||||
let (cache, loads, barrier) = (cache.clone(), loads.clone(), barrier.clone());
|
||||
tokio::spawn(async move {
|
||||
barrier.wait().await;
|
||||
cache
|
||||
.get_pipeline_str_with("db", "p", None, async {
|
||||
loads.fetch_add(1, Ordering::SeqCst);
|
||||
// Hold the loader open so every caller is waiting on it.
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
Ok(content_at(1))
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
for handle in handles {
|
||||
assert_eq!(handle.await.unwrap(), content_at(1));
|
||||
}
|
||||
assert_eq!(loads.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_delete_drops_version_pinned_entry() {
|
||||
let cache = PipelineCache::new();
|
||||
let content = content_at(1);
|
||||
let version = Some(content.version);
|
||||
|
||||
cache
|
||||
.get_pipeline_str_with("db", "p", version, async { Ok(content.clone()) })
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
cache.invalidate("p", version).await;
|
||||
|
||||
let loads = AtomicUsize::new(0);
|
||||
cache
|
||||
.get_pipeline_str_with("db", "p", version, async {
|
||||
loads.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(content.clone())
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(loads.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_drops_stale_latest_and_primes_failover() {
|
||||
let cache = PipelineCache::new();
|
||||
let v2 = content_at(2);
|
||||
|
||||
cache
|
||||
.get_pipeline_str_with("a", "p", None, async { Ok(content_at(1)) })
|
||||
.await
|
||||
.unwrap();
|
||||
cache.insert_failover_cache(content_at(1), true).await;
|
||||
|
||||
cache.on_pipeline_created(v2.clone()).await;
|
||||
|
||||
let loads = AtomicUsize::new(0);
|
||||
let cached = cache
|
||||
.get_pipeline_str_with("a", "p", None, async {
|
||||
loads.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(v2.clone())
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(loads.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(cached.version, v2.version);
|
||||
|
||||
let failover = cache.get_failover_cache("b", "p", None).await.unwrap();
|
||||
assert_eq!(failover.map(|c| c.version), Some(v2.version));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_failover_serves_global_pipeline_to_unwarmed_schema() {
|
||||
let cache = PipelineCache::new();
|
||||
let content = content_at(1);
|
||||
|
||||
cache.insert_failover_cache(content.clone(), true).await;
|
||||
|
||||
let found = cache.get_failover_cache("b", "p", None).await.unwrap();
|
||||
assert_eq!(found, Some(content.clone()));
|
||||
|
||||
// A same-named pipeline under another schema must not shadow the global one.
|
||||
let schema_local = PipelineContent {
|
||||
schema: "x".to_string(),
|
||||
..content_at(2)
|
||||
};
|
||||
cache.insert_failover_cache(schema_local, true).await;
|
||||
|
||||
let found = cache.get_failover_cache("b", "p", None).await.unwrap();
|
||||
assert_eq!(found, Some(content));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -262,21 +262,12 @@ impl PipelineTable {
|
||||
name: &str,
|
||||
input_version: PipelineVersion,
|
||||
) -> Result<Arc<Pipeline>> {
|
||||
if let Some(pipeline) = self.cache.get_pipeline_cache(schema, name, input_version)? {
|
||||
return Ok(pipeline);
|
||||
}
|
||||
|
||||
let pipeline_content = self.get_pipeline_str(schema, name, input_version).await?;
|
||||
let compiled_pipeline = Arc::new(Self::compile_pipeline(&pipeline_content.content)?);
|
||||
|
||||
self.cache.insert_pipeline_cache(
|
||||
&pipeline_content.schema,
|
||||
name,
|
||||
Some(pipeline_content.version),
|
||||
compiled_pipeline.clone(),
|
||||
input_version.is_none(),
|
||||
);
|
||||
Ok(compiled_pipeline)
|
||||
self.cache
|
||||
.get_pipeline_with(schema, name, input_version, async {
|
||||
let pipeline_content = self.get_pipeline_str(schema, name, input_version).await?;
|
||||
Ok(Arc::new(Self::compile_pipeline(&pipeline_content.content)?))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Get a original pipeline by name.
|
||||
@@ -287,13 +278,19 @@ impl PipelineTable {
|
||||
name: &str,
|
||||
input_version: PipelineVersion,
|
||||
) -> Result<PipelineContent> {
|
||||
if let Some(pipeline) = self
|
||||
.cache
|
||||
.get_pipeline_str_cache(schema, name, input_version)?
|
||||
{
|
||||
return Ok(pipeline);
|
||||
}
|
||||
self.cache
|
||||
.get_pipeline_str_with(schema, name, input_version, async {
|
||||
self.load_pipeline_str(schema, name, input_version).await
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn load_pipeline_str(
|
||||
&self,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
input_version: PipelineVersion,
|
||||
) -> Result<PipelineContent> {
|
||||
let mut pipeline_vec;
|
||||
match self.find_pipeline(name, input_version).await {
|
||||
Ok(p) => {
|
||||
@@ -312,7 +309,8 @@ impl PipelineTable {
|
||||
.inc();
|
||||
return self
|
||||
.cache
|
||||
.get_failover_cache(schema, name, input_version)?
|
||||
.get_failover_cache(schema, name, input_version)
|
||||
.await?
|
||||
.context(PipelineNotFoundSnafu {
|
||||
name,
|
||||
version: input_version,
|
||||
@@ -338,7 +336,8 @@ impl PipelineTable {
|
||||
let pipeline_content = pipeline_vec.remove(0);
|
||||
|
||||
self.cache
|
||||
.insert_pipeline_str_cache(&pipeline_content, input_version.is_none());
|
||||
.insert_failover_cache(pipeline_content.clone(), input_version.is_none())
|
||||
.await;
|
||||
return Ok(pipeline_content);
|
||||
}
|
||||
|
||||
@@ -359,12 +358,12 @@ impl PipelineTable {
|
||||
})?;
|
||||
|
||||
self.cache
|
||||
.insert_pipeline_str_cache(&pipeline_content, input_version.is_none());
|
||||
.insert_failover_cache(pipeline_content.clone(), input_version.is_none())
|
||||
.await;
|
||||
Ok(pipeline_content)
|
||||
}
|
||||
|
||||
/// Insert a pipeline into the pipeline table and compile it.
|
||||
/// The compiled pipeline will be inserted into the cache.
|
||||
/// Newly created pipelines will be saved under empty schema.
|
||||
pub async fn insert_and_compile(
|
||||
&self,
|
||||
@@ -378,25 +377,14 @@ impl PipelineTable {
|
||||
.insert_pipeline_to_pipeline_table(name, content_type, pipeline)
|
||||
.await?;
|
||||
|
||||
{
|
||||
self.cache.insert_pipeline_cache(
|
||||
EMPTY_SCHEMA_NAME,
|
||||
name,
|
||||
Some(TimestampNanosecond(version)),
|
||||
compiled_pipeline.clone(),
|
||||
true,
|
||||
);
|
||||
|
||||
let pipeline_content = PipelineContent {
|
||||
self.cache
|
||||
.on_pipeline_created(PipelineContent {
|
||||
name: name.to_string(),
|
||||
content: pipeline.to_string(),
|
||||
version: TimestampNanosecond(version),
|
||||
schema: EMPTY_SCHEMA_NAME.to_string(),
|
||||
};
|
||||
|
||||
self.cache
|
||||
.insert_pipeline_str_cache(&pipeline_content, true);
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
Ok((version, compiled_pipeline))
|
||||
}
|
||||
@@ -464,8 +452,7 @@ impl PipelineTable {
|
||||
output
|
||||
);
|
||||
|
||||
// remove cache with version and latest
|
||||
self.cache.remove_cache(name, version);
|
||||
self.cache.invalidate(name, version).await;
|
||||
|
||||
Ok(Some(()))
|
||||
}
|
||||
|
||||
@@ -17,7 +17,12 @@ mod common;
|
||||
use std::borrow::Cow;
|
||||
|
||||
use api::v1::ColumnDataType;
|
||||
use api::v1::json_value::Value as JsonValue;
|
||||
use api::v1::value::ValueData;
|
||||
use arrow_schema::extension::{
|
||||
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
|
||||
};
|
||||
use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
|
||||
|
||||
const INPUT_VALUE_OBJ: &str = r#"
|
||||
[
|
||||
@@ -113,6 +118,52 @@ transform:
|
||||
assert_eq!(v, jsonb::Value::Object(expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json2_parse() {
|
||||
let pipeline_yaml = r#"
|
||||
---
|
||||
processors:
|
||||
- json_parse:
|
||||
field: commit
|
||||
|
||||
transform:
|
||||
- field: commit
|
||||
type:
|
||||
json2:
|
||||
- path: "commitAuthor"
|
||||
type: string
|
||||
nullable: false
|
||||
"#;
|
||||
|
||||
let output = common::parse_and_exec(INPUT_VALUE_OBJ, pipeline_yaml);
|
||||
|
||||
assert_eq!(output.schema[0].datatype, ColumnDataType::Json as i32);
|
||||
assert!(output.schema[0].datatype_extension.is_none());
|
||||
assert_eq!(
|
||||
output.schema[0]
|
||||
.options
|
||||
.as_ref()
|
||||
.and_then(|options| options.options.get(EXTENSION_TYPE_NAME_KEY))
|
||||
.map(String::as_str),
|
||||
Some(Json2ExtensionType::NAME)
|
||||
);
|
||||
let metadata = output.schema[0]
|
||||
.options
|
||||
.as_ref()
|
||||
.and_then(|options| options.options.get(EXTENSION_TYPE_METADATA_KEY))
|
||||
.unwrap();
|
||||
let metadata: JsonMetadata = serde_json::from_str(metadata).unwrap();
|
||||
assert_eq!(
|
||||
metadata.json_settings().type_hints()[0].path,
|
||||
["commitAuthor"]
|
||||
);
|
||||
|
||||
let ValueData::JsonValue(value) = output.rows[0].values[0].value_data.as_ref().unwrap() else {
|
||||
panic!("expect JSON2 value");
|
||||
};
|
||||
assert!(matches!(value.value, Some(JsonValue::Object(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_parse_with_simple_extractor() {
|
||||
let pipeline_yaml = r#"
|
||||
|
||||
@@ -46,6 +46,17 @@ const STAGE: &str = "stage";
|
||||
const NODE: &str = "node";
|
||||
const PLAN: &str = "plan";
|
||||
|
||||
/// Fixed output schema of [`DistAnalyzeExec`], for `Describe` handlers:
|
||||
/// execution rewrites the plan in `optimize_physical_plan`, so this schema
|
||||
/// differs from the logical `Analyze` plan's.
|
||||
pub fn dist_analyze_output_schema() -> SchemaRef {
|
||||
SchemaRef::new(Schema::new(vec![
|
||||
Field::new(STAGE, DataType::UInt32, true),
|
||||
Field::new(NODE, DataType::UInt32, true),
|
||||
Field::new(PLAN, DataType::Utf8, true),
|
||||
]))
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct DistAnalyzeExec {
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
@@ -58,11 +69,7 @@ pub struct DistAnalyzeExec {
|
||||
impl DistAnalyzeExec {
|
||||
/// Create a new DistAnalyzeExec
|
||||
pub fn new(input: Arc<dyn ExecutionPlan>, verbose: bool, format: AnalyzeFormat) -> Self {
|
||||
let schema = SchemaRef::new(Schema::new(vec![
|
||||
Field::new(STAGE, DataType::UInt32, true),
|
||||
Field::new(NODE, DataType::UInt32, true),
|
||||
Field::new(PLAN, DataType::Utf8, true),
|
||||
]));
|
||||
let schema = dist_analyze_output_schema();
|
||||
let properties = Arc::new(Self::compute_properties(&input, schema.clone()));
|
||||
Self {
|
||||
input,
|
||||
|
||||
+101
-25
@@ -30,7 +30,7 @@ use common_function::function::FunctionContext;
|
||||
use common_function::function_factory::ScalarFunctionFactory;
|
||||
use common_query::{Output, OutputData, OutputMeta};
|
||||
use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters};
|
||||
use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream};
|
||||
use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream};
|
||||
use common_telemetry::tracing;
|
||||
use datafusion::catalog::TableFunction;
|
||||
use datafusion::dataframe::DataFrame;
|
||||
@@ -44,6 +44,7 @@ use datafusion_expr::{
|
||||
use datatypes::prelude::VectorRef;
|
||||
use datatypes::schema::Schema;
|
||||
use futures_util::StreamExt;
|
||||
use futures_util::future::try_join;
|
||||
use session::context::QueryContextRef;
|
||||
use snafu::{OptionExt, ResultExt, ensure};
|
||||
use sqlparser::ast::AnalyzeFormat;
|
||||
@@ -80,6 +81,29 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism";
|
||||
/// Whether to fallback to the original plan when failed to push down.
|
||||
pub const QUERY_FALLBACK_HINT: &str = "query_fallback";
|
||||
|
||||
// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared
|
||||
// HTTP/2 connection are pending, trading bounded memory for request liveness.
|
||||
async fn forward_record_batches(
|
||||
mut stream: SendableRecordBatchStream,
|
||||
batch_tx: tokio::sync::mpsc::UnboundedSender<Result<RecordBatch>>,
|
||||
) -> Result<()> {
|
||||
while let Some(batch) = stream.next().await {
|
||||
match batch.context(CreateRecordBatchSnafu) {
|
||||
Ok(batch) => {
|
||||
if batch_tx.send(Ok(batch)).is_err() {
|
||||
break;
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = batch_tx.send(Err(error));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn query_load_region_id(plan: &Arc<dyn ExecutionPlan>) -> Option<u64> {
|
||||
let mut region_id = None;
|
||||
let mut stack = vec![plan.clone()];
|
||||
@@ -205,7 +229,7 @@ impl DatafusionQueryEngine {
|
||||
let Output { data, meta } = self
|
||||
.exec_query_plan((*dml.input).clone(), query_ctx.clone())
|
||||
.await?;
|
||||
let mut stream = match data {
|
||||
let stream = match data {
|
||||
OutputData::RecordBatches(batches) => batches.as_stream(),
|
||||
OutputData::Stream(stream) => stream,
|
||||
_ => unreachable!(),
|
||||
@@ -214,30 +238,35 @@ impl DatafusionQueryEngine {
|
||||
let mut affected_rows = 0;
|
||||
let mut insert_cost = 0;
|
||||
|
||||
while let Some(batch) = stream.next().await {
|
||||
let batch = batch.context(CreateRecordBatchSnafu)?;
|
||||
let column_vectors = batch
|
||||
.column_vectors(&table_name.to_string(), table.schema())
|
||||
.map_err(BoxedError::new)
|
||||
.context(QueryExecutionSnafu)?;
|
||||
|
||||
match dml.op {
|
||||
WriteOp::Insert(_) => {
|
||||
// We ignore the insert op.
|
||||
let output = self
|
||||
.insert(&table_name, column_vectors, query_ctx.clone())
|
||||
.await?;
|
||||
let (rows, cost) = output.extract_rows_and_cost();
|
||||
affected_rows += rows;
|
||||
insert_cost += cost;
|
||||
}
|
||||
WriteOp::Delete => {
|
||||
affected_rows += self
|
||||
.delete(&table_name, &table, column_vectors, query_ctx.clone())
|
||||
.await?;
|
||||
}
|
||||
_ => unreachable!("guarded by the 'ensure!' at the beginning"),
|
||||
match dml.op {
|
||||
WriteOp::Insert(_) => {
|
||||
let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let producer = forward_record_batches(stream, batch_tx);
|
||||
let consumer = self.consume_insert_record_batches(
|
||||
batch_rx,
|
||||
&table_name,
|
||||
table.schema(),
|
||||
query_ctx.clone(),
|
||||
);
|
||||
let ((), (rows, cost)) = try_join(producer, consumer).await?;
|
||||
affected_rows += rows;
|
||||
insert_cost += cost;
|
||||
}
|
||||
WriteOp::Delete => {
|
||||
// Keep DELETE on the same producer/consumer schedule as INSERT so the source
|
||||
// stream can continue draining while mutation RPCs are pending.
|
||||
let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let producer = forward_record_batches(stream, batch_tx);
|
||||
let consumer = self.consume_delete_record_batches(
|
||||
batch_rx,
|
||||
&table_name,
|
||||
&table,
|
||||
query_ctx.clone(),
|
||||
);
|
||||
let (rows, ()) = try_join(consumer, producer).await?;
|
||||
affected_rows += rows;
|
||||
}
|
||||
_ => unreachable!("guarded by the 'ensure!' at the beginning"),
|
||||
}
|
||||
Ok(Output::new(
|
||||
OutputData::AffectedRows(affected_rows),
|
||||
@@ -245,6 +274,53 @@ impl DatafusionQueryEngine {
|
||||
))
|
||||
}
|
||||
|
||||
async fn consume_insert_record_batches(
|
||||
&self,
|
||||
mut batch_rx: tokio::sync::mpsc::UnboundedReceiver<Result<RecordBatch>>,
|
||||
table_name: &ResolvedTableReference,
|
||||
table_schema: Arc<Schema>,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<(usize, usize)> {
|
||||
let mut affected_rows = 0;
|
||||
let mut insert_cost = 0;
|
||||
while let Some(batch) = batch_rx.recv().await {
|
||||
let batch = batch?;
|
||||
let column_vectors = batch
|
||||
.column_vectors(&table_name.to_string(), table_schema.clone())
|
||||
.map_err(BoxedError::new)
|
||||
.context(QueryExecutionSnafu)?;
|
||||
// We ignore the insert op.
|
||||
let output = self
|
||||
.insert(table_name, column_vectors, query_ctx.clone())
|
||||
.await?;
|
||||
let (rows, cost) = output.extract_rows_and_cost();
|
||||
affected_rows += rows;
|
||||
insert_cost += cost;
|
||||
}
|
||||
Ok((affected_rows, insert_cost))
|
||||
}
|
||||
|
||||
async fn consume_delete_record_batches(
|
||||
&self,
|
||||
mut batch_rx: tokio::sync::mpsc::UnboundedReceiver<Result<RecordBatch>>,
|
||||
table_name: &ResolvedTableReference,
|
||||
table: &TableRef,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<usize> {
|
||||
let mut affected_rows = 0;
|
||||
while let Some(batch) = batch_rx.recv().await {
|
||||
let batch = batch?;
|
||||
let column_vectors = batch
|
||||
.column_vectors(&table_name.to_string(), table.schema())
|
||||
.map_err(BoxedError::new)
|
||||
.context(QueryExecutionSnafu)?;
|
||||
affected_rows += self
|
||||
.delete(table_name, table, column_vectors, query_ctx.clone())
|
||||
.await?;
|
||||
}
|
||||
Ok(affected_rows)
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
async fn delete(
|
||||
&self,
|
||||
|
||||
@@ -18,23 +18,32 @@ use arrow_schema::Field;
|
||||
use common_function::scalars::json::json_get::JsonGetWithType;
|
||||
use common_function::scalars::udf::create_udf;
|
||||
use datafusion_common::arrow::datatypes::DataType;
|
||||
use datafusion_common::{Column, DFSchema, Result, ScalarValue, TableReference};
|
||||
use datafusion_common::{Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference};
|
||||
use datafusion_expr::expr::{BinaryExpr, ScalarFunction};
|
||||
use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
|
||||
use datafusion_expr::{Expr, ExprSchemable, Operator, ScalarUDF};
|
||||
use datafusion_expr::planner::{
|
||||
ExprPlanner, PlannerResult, RawAggregateExpr, RawBinaryExpr, RawFieldAccessExpr, RawScalarExpr,
|
||||
RawWindowExpr,
|
||||
};
|
||||
use datafusion_expr::type_coercion::functions::{UDFCoercionExt, fields_with_udf};
|
||||
use datafusion_expr::{
|
||||
Expr, ExprSchemable, GetFieldAccess, Operator, ScalarUDF, WindowFunctionDefinition,
|
||||
};
|
||||
use datatypes::extension::json::is_json2_extension_type;
|
||||
use either::Either;
|
||||
use sqlparser::ast::BinaryOperator;
|
||||
|
||||
/// Rewrites JSON-aware SQL expressions into DataFusion expressions.
|
||||
///
|
||||
/// This planner handles two cases:
|
||||
/// This planner handles three cases:
|
||||
/// - Rewrites compound identifiers on JSON extension columns into `json_get` function.
|
||||
/// 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\"]")`.
|
||||
/// - 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.
|
||||
/// For example, `select json_get(a, "b.c") + 1` => `select json_get(a, "b.c", NULL::Int64) + 1`.
|
||||
/// - Infers the expected type from scalar, aggregate, and window function signatures.
|
||||
/// For example, `select abs(a.b.c)` => `select abs(json_get(a, "b.c", NULL::Float64))`.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct JsonExprPlanner;
|
||||
|
||||
@@ -50,9 +59,7 @@ impl ExprPlanner for JsonExprPlanner {
|
||||
mut right,
|
||||
} = expr;
|
||||
|
||||
if extract_untyped_json_get(&mut left).is_none()
|
||||
&& extract_untyped_json_get(&mut right).is_none()
|
||||
{
|
||||
if !is_untyped_json_get(&left) && !is_untyped_json_get(&right) {
|
||||
return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
|
||||
}
|
||||
|
||||
@@ -62,20 +69,64 @@ impl ExprPlanner for JsonExprPlanner {
|
||||
|
||||
let left_type = left.get_type(schema)?;
|
||||
let right_type = right.get_type(schema)?;
|
||||
let left = push_json_get_type_arg(left, right_type)?;
|
||||
let right = push_json_get_type_arg(right, left_type)?;
|
||||
match (left, right) {
|
||||
(Either::Left(left), Either::Left(right)) => {
|
||||
Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }))
|
||||
}
|
||||
(left, right) => Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr::new(
|
||||
Box::new(left.into_inner()),
|
||||
let left_changed = push_json_get_type_arg(&mut left, &right_type)?;
|
||||
let right_changed = push_json_get_type_arg(&mut right, &left_type)?;
|
||||
if left_changed || right_changed {
|
||||
Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr::new(
|
||||
Box::new(left),
|
||||
expr_op,
|
||||
Box::new(right.into_inner()),
|
||||
)))),
|
||||
Box::new(right),
|
||||
))))
|
||||
} else {
|
||||
Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }))
|
||||
}
|
||||
}
|
||||
|
||||
/// Extends the path of an untyped `json_get` with one field access.
|
||||
///
|
||||
/// For `j.o.l[1].inner.l[2]`, `plan_compound_identifier` first produces
|
||||
/// `json_get(j, "o.l")`. DataFusion then calls this method successively
|
||||
/// with a list index, two named fields, and another list index, producing
|
||||
/// the final path `o.l[1]["inner"]["l"][2]`.
|
||||
fn plan_field_access(
|
||||
&self,
|
||||
mut expr: RawFieldAccessExpr,
|
||||
_schema: &DFSchema,
|
||||
) -> Result<PlannerResult<RawFieldAccessExpr>> {
|
||||
// See `normalize_field_access_after_subscript` for the reason why we construct the
|
||||
// "suffix" like this.
|
||||
let suffix = match &expr.field_access {
|
||||
GetFieldAccess::ListIndex { key } => {
|
||||
// DataFusion parses ordinary integer literals within the i64 range as Int64.
|
||||
let Expr::Literal(ScalarValue::Int64(Some(index)), _) = key.as_ref() else {
|
||||
return Ok(PlannerResult::Original(expr));
|
||||
};
|
||||
format!("[{index}]")
|
||||
}
|
||||
GetFieldAccess::NamedStructField { name } => {
|
||||
let Some(name) = name.try_as_str().flatten() else {
|
||||
return Ok(PlannerResult::Original(expr));
|
||||
};
|
||||
// Encode the field name as a JSON string before embedding it in the
|
||||
// bracket accessor. This preserves dots as literal field-name characters
|
||||
// and escapes quotes, backslashes, and control characters correctly.
|
||||
let name = serde_json::to_string(name)
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
format!("[{name}]")
|
||||
}
|
||||
GetFieldAccess::ListRange { .. } => return Ok(PlannerResult::Original(expr)),
|
||||
};
|
||||
let Some(json_get) = extract_untyped_json_get(&mut expr.expr) else {
|
||||
return Ok(PlannerResult::Original(expr));
|
||||
};
|
||||
let Some(Expr::Literal(ScalarValue::Utf8(Some(path)), _)) = json_get.args.get_mut(1) else {
|
||||
return Ok(PlannerResult::Original(expr));
|
||||
};
|
||||
|
||||
path.push_str(&suffix);
|
||||
Ok(PlannerResult::Planned(expr.expr))
|
||||
}
|
||||
|
||||
fn plan_compound_identifier(
|
||||
&self,
|
||||
field: &Field,
|
||||
@@ -101,31 +152,235 @@ impl ExprPlanner for JsonExprPlanner {
|
||||
),
|
||||
)))
|
||||
}
|
||||
|
||||
/// Rewrites JSON2 arguments without taking over the final function planning.
|
||||
///
|
||||
/// `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<PlannerResult<RawScalarExpr>> {
|
||||
push_function_arg_types(expr.func.as_ref(), &mut expr.args)?;
|
||||
Ok(PlannerResult::Original(expr))
|
||||
}
|
||||
|
||||
/// Rewrites JSON2 arguments while preserving subsequent aggregate planning.
|
||||
fn plan_aggregate(
|
||||
&self,
|
||||
mut expr: RawAggregateExpr,
|
||||
) -> Result<PlannerResult<RawAggregateExpr>> {
|
||||
push_function_arg_types(expr.func.as_ref(), &mut expr.args)?;
|
||||
Ok(PlannerResult::Original(expr))
|
||||
}
|
||||
|
||||
/// Rewrites JSON2 arguments while preserving subsequent window planning.
|
||||
fn plan_window(&self, mut expr: RawWindowExpr) -> Result<PlannerResult<RawWindowExpr>> {
|
||||
match &expr.func_def {
|
||||
WindowFunctionDefinition::AggregateUDF(func) => {
|
||||
push_function_arg_types(func.as_ref(), &mut expr.args)?;
|
||||
}
|
||||
WindowFunctionDefinition::WindowUDF(func) => {
|
||||
push_function_arg_types(func.as_ref(), &mut expr.args)?;
|
||||
}
|
||||
}
|
||||
Ok(PlannerResult::Original(expr))
|
||||
}
|
||||
}
|
||||
|
||||
enum JsonGetTypeResolution {
|
||||
Fallback,
|
||||
Typed(Vec<(usize, DataType)>),
|
||||
}
|
||||
|
||||
/// Infers static output types for untyped `json_get` arguments from a function signature.
|
||||
///
|
||||
/// DataFusion requires every expression to have one Arrow data type during planning. A JSON path
|
||||
/// may contain heterogeneous values across rows, but it cannot expose those values as different
|
||||
/// Arrow types in one result column. Preserving their runtime types would require a single
|
||||
/// Variant-like data type and Variant-aware functions instead. Maybe we can wait for
|
||||
/// https://github.com/apache/datafusion/issues/16116
|
||||
///
|
||||
/// This helper uses the function's coercion rules to select a supported output type, then appends
|
||||
/// a typed NULL argument to each relevant `json_get`. The typed argument makes `json_get` project
|
||||
/// compatible JSON values to that type and return NULL for incompatible values. Functions that
|
||||
/// accept json_get's default `Utf8View` output keep the two-argument form so later rewrites can
|
||||
/// still push down an outer cast.
|
||||
fn push_function_arg_types<F>(func: &F, args: &mut [Expr]) -> Result<()>
|
||||
where
|
||||
F: UDFCoercionExt,
|
||||
{
|
||||
if !args.iter().any(is_untyped_json_get) {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let fields = args.iter().map(function_arg_field).collect::<Vec<_>>();
|
||||
match infer_json_get_types(func, args, &fields) {
|
||||
JsonGetTypeResolution::Fallback => {
|
||||
let Some(data_type) = fallback_json_get_type(func, args, &fields) else {
|
||||
return Ok(());
|
||||
};
|
||||
for arg in args.iter_mut() {
|
||||
if is_untyped_json_get(arg) {
|
||||
let _ = push_json_get_type_arg(arg, &data_type)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
JsonGetTypeResolution::Typed(types) => {
|
||||
for (index, data_type) in types {
|
||||
let _ = push_json_get_type_arg(&mut args[index], &data_type)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn infer_json_get_types<F>(func: &F, args: &[Expr], fields: &[Arc<Field>]) -> JsonGetTypeResolution
|
||||
where
|
||||
F: UDFCoercionExt,
|
||||
{
|
||||
// Only untyped json_get arguments use Null placeholders; preserve every other known argument
|
||||
// type. fields_with_udf performs contextual coercion rather than reverse inference from a
|
||||
// signature alone. Numeric signatures may preserve all-Null inputs, while Comparable
|
||||
// signatures may default them to Utf8. For example, retaining the Float64 peer in
|
||||
// coalesce(json_get(...), 1.0) lets DataFusion resolve json_get to Float64 instead of Utf8.
|
||||
//
|
||||
// This is a best-effort probe: a failure does not mean the actual function call is invalid, so
|
||||
// try concrete JSON types before leaving final validation to DataFusion's default planner.
|
||||
let Ok(coerced) = fields_with_udf(fields, func) else {
|
||||
return JsonGetTypeResolution::Fallback;
|
||||
};
|
||||
|
||||
let mut inferred_types = Vec::with_capacity(coerced.len());
|
||||
for (index, (arg, field)) in args.iter().zip(coerced).enumerate() {
|
||||
if !is_untyped_json_get(arg) || field.data_type().is_null() {
|
||||
continue;
|
||||
}
|
||||
let Some(data_type) = json_get_output_type(field.data_type()) else {
|
||||
return JsonGetTypeResolution::Fallback;
|
||||
};
|
||||
inferred_types.push((index, data_type));
|
||||
}
|
||||
if inferred_types.is_empty() {
|
||||
JsonGetTypeResolution::Fallback
|
||||
} else {
|
||||
JsonGetTypeResolution::Typed(inferred_types)
|
||||
}
|
||||
}
|
||||
|
||||
fn fallback_json_get_type<F>(func: &F, args: &[Expr], fields: &[Arc<Field>]) -> Option<DataType>
|
||||
where
|
||||
F: UDFCoercionExt,
|
||||
{
|
||||
// Prefer json_get's default Utf8View type. If the function rejects strings but accepts numeric
|
||||
// values, prefer Float64 so both integers and fractions remain usable.
|
||||
let mut candidate_fields = fields.to_vec();
|
||||
for data_type in [
|
||||
DataType::Utf8View,
|
||||
DataType::Float64,
|
||||
DataType::Int64,
|
||||
DataType::Boolean,
|
||||
] {
|
||||
for (index, arg) in args.iter().enumerate() {
|
||||
if is_untyped_json_get(arg) {
|
||||
candidate_fields[index] = Arc::new(
|
||||
fields[index]
|
||||
.as_ref()
|
||||
.clone()
|
||||
.with_data_type(data_type.clone()),
|
||||
);
|
||||
}
|
||||
}
|
||||
if fields_with_udf(&candidate_fields, func).is_ok() {
|
||||
return Some(data_type);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn function_arg_field(expr: &Expr) -> Arc<Field> {
|
||||
let data_type = if is_untyped_json_get(expr) {
|
||||
DataType::Null
|
||||
} else if let Some(data_type) = extract_json_get_type(expr) {
|
||||
data_type
|
||||
} else {
|
||||
// Treat unresolved expressions as untyped NULL. This lets signatures such as `power`
|
||||
// infer a JSON type, while functions such as `coalesce` can leave it untyped for default
|
||||
// planning. This is only best-effort: overloaded or user-defined functions may select a
|
||||
// different signature for NULL than for the expression's actual type.
|
||||
// TODO(LFC): Use the input schema once DataFusion passes it to ExprPlanner::plan_*().
|
||||
expr.get_type(&DFSchema::empty()).unwrap_or(DataType::Null)
|
||||
};
|
||||
Arc::new(Field::new("", data_type, true))
|
||||
}
|
||||
|
||||
fn json_get_output_type(data_type: &DataType) -> Option<DataType> {
|
||||
let output_type = match data_type {
|
||||
DataType::Boolean => DataType::Boolean,
|
||||
data_type if data_type.is_integer() => DataType::Int64,
|
||||
data_type if data_type.is_floating() => DataType::Float64,
|
||||
DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => DataType::Float64,
|
||||
data_type if data_type.is_string() => DataType::Utf8View,
|
||||
_ => return None,
|
||||
};
|
||||
Some(output_type)
|
||||
}
|
||||
|
||||
macro_rules! is_untyped_json_get_func {
|
||||
($func:expr) => {
|
||||
$func
|
||||
.func
|
||||
.name()
|
||||
.eq_ignore_ascii_case(JsonGetWithType::NAME)
|
||||
&& $func.args.len() == 2
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! is_typed_json_get_func {
|
||||
($func:expr) => {
|
||||
$func
|
||||
.func
|
||||
.name()
|
||||
.eq_ignore_ascii_case(JsonGetWithType::NAME)
|
||||
&& $func.args.len() == 3
|
||||
};
|
||||
}
|
||||
|
||||
fn extract_untyped_json_get(expr: &mut Expr) -> Option<&mut ScalarFunction> {
|
||||
match expr {
|
||||
Expr::ScalarFunction(f)
|
||||
if f.func.name().eq_ignore_ascii_case(JsonGetWithType::NAME) && f.args.len() == 2 =>
|
||||
{
|
||||
Some(f)
|
||||
}
|
||||
Expr::ScalarFunction(f) if is_untyped_json_get_func!(f) => Some(f),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn push_json_get_type_arg(mut expr: Expr, mut data_type: DataType) -> Result<Either<Expr, Expr>> {
|
||||
let Some(json_get) = extract_untyped_json_get(&mut expr) else {
|
||||
return Ok(Either::Left(expr));
|
||||
fn extract_json_get_type(expr: &Expr) -> Option<DataType> {
|
||||
match expr {
|
||||
Expr::ScalarFunction(f) if is_typed_json_get_func!(f) => f
|
||||
.args
|
||||
.get(2)
|
||||
.and_then(|x| x.as_literal())
|
||||
.map(|x| x.data_type()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_untyped_json_get(expr: &Expr) -> bool {
|
||||
matches!(
|
||||
expr,
|
||||
Expr::ScalarFunction(f) if is_untyped_json_get_func!(f)
|
||||
)
|
||||
}
|
||||
|
||||
fn push_json_get_type_arg(expr: &mut Expr, data_type: &DataType) -> Result<bool> {
|
||||
let Some(json_get) = extract_untyped_json_get(expr) else {
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
// The two-argument form already returns Utf8View. Keep it so JsonGetRewriter can still absorb
|
||||
// a cast added by subsequent function coercion.
|
||||
if data_type.is_string() {
|
||||
data_type = DataType::Utf8View;
|
||||
return Ok(false);
|
||||
}
|
||||
let with_type = ScalarValue::try_new_null(&data_type).map(|x| Expr::Literal(x, None))?;
|
||||
let with_type = ScalarValue::try_new_null(data_type).map(|x| Expr::Literal(x, None))?;
|
||||
json_get.args.push(with_type);
|
||||
|
||||
Ok(Either::Right(expr))
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn parse_sql_op(op: &BinaryOperator) -> Option<Operator> {
|
||||
@@ -153,6 +408,11 @@ fn parse_sql_op(op: &BinaryOperator) -> Option<Operator> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use arrow_schema::Fields;
|
||||
use datafusion::functions_aggregate::count::count_udaf;
|
||||
use datafusion::functions_aggregate::sum::sum_udaf;
|
||||
use datafusion_expr::WindowFrame;
|
||||
use datafusion_functions::core::coalesce;
|
||||
use datafusion_functions::math::{abs, power};
|
||||
use datatypes::extension::json::Json2ExtensionType;
|
||||
|
||||
use super::*;
|
||||
@@ -236,6 +496,55 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_plan_list_index() -> Result<()> {
|
||||
let planner = JsonExprPlanner;
|
||||
let planned = planner.plan_field_access(
|
||||
RawFieldAccessExpr {
|
||||
field_access: GetFieldAccess::ListIndex {
|
||||
key: Box::new(Expr::Literal(ScalarValue::Int64(Some(0)), None)),
|
||||
},
|
||||
expr: json_get_expr(Expr::Column(Column::new_unqualified("j")), "list"),
|
||||
},
|
||||
&DFSchema::empty(),
|
||||
)?;
|
||||
let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(func.func.name(), JsonGetWithType::NAME);
|
||||
assert_eq!(func.args.len(), 2);
|
||||
assert_eq!(
|
||||
func.args[1],
|
||||
Expr::Literal(ScalarValue::Utf8(Some("list[0]".to_string())), None)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_plan_field_after_list_index() -> Result<()> {
|
||||
let planner = JsonExprPlanner;
|
||||
let planned = planner.plan_field_access(
|
||||
RawFieldAccessExpr {
|
||||
field_access: GetFieldAccess::NamedStructField {
|
||||
name: ScalarValue::Utf8(Some("a.b".to_string())),
|
||||
},
|
||||
expr: json_get_expr(Expr::Column(Column::new_unqualified("j")), "list[0]"),
|
||||
},
|
||||
&DFSchema::empty(),
|
||||
)?;
|
||||
let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
func.args[1],
|
||||
Expr::Literal(
|
||||
ScalarValue::Utf8(Some("list[0][\"a.b\"]".to_string())),
|
||||
None
|
||||
)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_plan_compound_identifier() -> Result<()> {
|
||||
let planner = JsonExprPlanner;
|
||||
@@ -281,4 +590,110 @@ mod tests {
|
||||
|
||||
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()],
|
||||
})?
|
||||
else {
|
||||
unreachable!();
|
||||
};
|
||||
assert_eq!(
|
||||
Some(DataType::Float64),
|
||||
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")),
|
||||
],
|
||||
})?
|
||||
else {
|
||||
unreachable!();
|
||||
};
|
||||
assert_eq!(
|
||||
Some(DataType::Float64),
|
||||
extract_json_get_type(&scalar.args[0])
|
||||
);
|
||||
|
||||
let PlannerResult::Original(aggregate) = planner.plan_aggregate(RawAggregateExpr {
|
||||
func: sum_udaf(),
|
||||
args: vec![json_get()],
|
||||
distinct: false,
|
||||
filter: None,
|
||||
order_by: vec![],
|
||||
null_treatment: None,
|
||||
})?
|
||||
else {
|
||||
unreachable!();
|
||||
};
|
||||
assert_eq!(
|
||||
Some(DataType::Float64),
|
||||
extract_json_get_type(&aggregate.args[0])
|
||||
);
|
||||
|
||||
let PlannerResult::Original(count) = planner.plan_aggregate(RawAggregateExpr {
|
||||
func: count_udaf(),
|
||||
args: vec![json_get()],
|
||||
distinct: false,
|
||||
filter: None,
|
||||
order_by: vec![],
|
||||
null_treatment: None,
|
||||
})?
|
||||
else {
|
||||
unreachable!();
|
||||
};
|
||||
assert_eq!(None, extract_json_get_type(&count.args[0]));
|
||||
|
||||
let PlannerResult::Original(window) = planner.plan_window(RawWindowExpr {
|
||||
func_def: WindowFunctionDefinition::AggregateUDF(sum_udaf()),
|
||||
args: vec![json_get()],
|
||||
partition_by: vec![],
|
||||
order_by: vec![],
|
||||
window_frame: WindowFrame::new(None),
|
||||
filter: None,
|
||||
null_treatment: None,
|
||||
distinct: false,
|
||||
})?
|
||||
else {
|
||||
unreachable!();
|
||||
};
|
||||
assert_eq!(
|
||||
Some(DataType::Float64),
|
||||
extract_json_get_type(&window.args[0])
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_plan_function_with_mixed_json_get_types() -> Result<()> {
|
||||
let planner = JsonExprPlanner;
|
||||
let json_get = || json_get_expr(Expr::Column(Column::new_unqualified("j")), "a.b");
|
||||
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],
|
||||
})?
|
||||
else {
|
||||
unreachable!();
|
||||
};
|
||||
assert_eq!(
|
||||
Some(DataType::Float64),
|
||||
extract_json_get_type(&scalar.args[0])
|
||||
);
|
||||
assert_eq!(
|
||||
Some(DataType::Float64),
|
||||
extract_json_get_type(&scalar.args[1])
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,7 +20,7 @@ use std::time::Duration;
|
||||
|
||||
use ahash::{HashMap, HashSet};
|
||||
use arrow_schema::{
|
||||
ArrowError, DataType, DataType as ArrowDataType, Field, Schema as ArrowSchema,
|
||||
ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema,
|
||||
SchemaRef as ArrowSchemaRef, SortOptions,
|
||||
};
|
||||
use async_stream::stream;
|
||||
@@ -42,14 +42,12 @@ use datafusion::physical_plan::{
|
||||
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
|
||||
SendableRecordBatchStream,
|
||||
};
|
||||
use datafusion_common::{Column as ColumnExpr, DataFusionError, Result};
|
||||
use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore};
|
||||
use datafusion_common::stats::Precision;
|
||||
use datafusion_common::{Column as ColumnExpr, 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 datatypes::extension::json::{
|
||||
Json2ExtensionType, is_any_json_extension_type, is_json2_extension_type,
|
||||
is_legacy_json2_extension_type,
|
||||
};
|
||||
use datatypes::extension::json::is_any_json_extension_type;
|
||||
use futures_util::StreamExt;
|
||||
use greptime_proto::v1::region::RegionRequestHeader;
|
||||
use meter_core::data::ReadItem;
|
||||
@@ -81,6 +79,50 @@ fn query_engine_state_from_task_context(context: &TaskContext) -> Option<Arc<Que
|
||||
context.session_config().get_extension()
|
||||
}
|
||||
|
||||
/// Returns a deterministic upper bound on rows emitted by the remote plan for one region.
|
||||
///
|
||||
/// Only explicit caps and row-non-increasing wrappers are followed. Nodes that may multiply
|
||||
/// rows fail open, so the result remains a sound upper bound.
|
||||
fn remote_plan_row_bound(plan: &LogicalPlan) -> Option<usize> {
|
||||
match plan {
|
||||
LogicalPlan::Limit(limit) => {
|
||||
let input_bound = remote_plan_row_bound(&limit.input);
|
||||
match limit.get_fetch_type() {
|
||||
Ok(FetchType::Literal(Some(fetch))) => {
|
||||
Some(input_bound.map_or(fetch, |bound| bound.min(fetch)))
|
||||
}
|
||||
_ => input_bound,
|
||||
}
|
||||
}
|
||||
LogicalPlan::Sort(sort) => {
|
||||
let input_bound = remote_plan_row_bound(&sort.input);
|
||||
sort.fetch
|
||||
.map(|fetch| input_bound.map_or(fetch, |bound| bound.min(fetch)))
|
||||
.or(input_bound)
|
||||
}
|
||||
LogicalPlan::Projection(projection) => remote_plan_row_bound(&projection.input),
|
||||
LogicalPlan::Filter(filter) => remote_plan_row_bound(&filter.input),
|
||||
LogicalPlan::SubqueryAlias(alias) => remote_plan_row_bound(&alias.input),
|
||||
LogicalPlan::Window(window) => remote_plan_row_bound(&window.input),
|
||||
LogicalPlan::Repartition(repartition) => remote_plan_row_bound(&repartition.input),
|
||||
LogicalPlan::Distinct(distinct) => remote_plan_row_bound(distinct.input()),
|
||||
LogicalPlan::Aggregate(aggregate) => {
|
||||
if aggregate
|
||||
.group_expr
|
||||
.iter()
|
||||
.any(|expr| matches!(expr, Expr::GroupingSet(_)))
|
||||
{
|
||||
None
|
||||
} else if aggregate.group_expr.is_empty() {
|
||||
Some(1)
|
||||
} else {
|
||||
remote_plan_row_bound(&aggregate.input)
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn remote_dyn_filter_enabled(query_ctx: &QueryContextRef) -> Result<bool> {
|
||||
remote_dyn_filter_pushdown_enabled_from_extensions(&query_ctx.extensions())
|
||||
.map_err(|err| DataFusionError::External(Box::new(err)))
|
||||
@@ -232,10 +274,12 @@ fn query_context_for_remote_dyn_filter_region(
|
||||
query_context_with_initial_dyn_filter_regs(query_ctx, region_id, captured_dyn_filters)
|
||||
}
|
||||
|
||||
#[derive(Debug, Hash, PartialOrd, PartialEq, Eq, Clone)]
|
||||
#[derive(Debug, Hash, PartialEq, Eq, Clone)]
|
||||
pub struct MergeScanLogicalPlan {
|
||||
/// In logical plan phase it only contains one input
|
||||
input: LogicalPlan,
|
||||
/// Schema exposed to the local stage.
|
||||
output_schema: DFSchemaRef,
|
||||
/// If this plan is a placeholder
|
||||
is_placeholder: bool,
|
||||
partition_cols: AliasMapping,
|
||||
@@ -243,6 +287,42 @@ pub struct MergeScanLogicalPlan {
|
||||
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
|
||||
}
|
||||
|
||||
impl PartialOrd for MergeScanLogicalPlan {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
||||
let Self {
|
||||
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,
|
||||
));
|
||||
match ordering {
|
||||
Some(std::cmp::Ordering::Equal) if output_schema != other_output_schema => None,
|
||||
ordering => ordering,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UserDefinedLogicalNodeCore for MergeScanLogicalPlan {
|
||||
fn name(&self) -> &str {
|
||||
Self::name()
|
||||
@@ -255,7 +335,7 @@ impl UserDefinedLogicalNodeCore for MergeScanLogicalPlan {
|
||||
}
|
||||
|
||||
fn schema(&self) -> &datafusion_common::DFSchemaRef {
|
||||
self.input.schema()
|
||||
&self.output_schema
|
||||
}
|
||||
|
||||
// Prevent further optimization
|
||||
@@ -281,8 +361,10 @@ 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 {
|
||||
Self {
|
||||
output_schema: input.schema().clone(),
|
||||
input,
|
||||
is_placeholder,
|
||||
partition_cols,
|
||||
@@ -290,6 +372,12 @@ impl MergeScanLogicalPlan {
|
||||
}
|
||||
}
|
||||
|
||||
/// Replaces the schema exposed to the local stage.
|
||||
pub(crate) fn with_output_schema(mut self, output_schema: DFSchemaRef) -> Self {
|
||||
self.output_schema = output_schema;
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn with_remote_dyn_filter_producer_id(
|
||||
mut self,
|
||||
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
|
||||
@@ -373,7 +461,11 @@ impl MergeScanExec {
|
||||
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
|
||||
enable_per_region_metrics: bool,
|
||||
) -> Result<Self> {
|
||||
let arrow_schema = maybe_amend_json2_field(arrow_schema);
|
||||
// JSON2 schemas are concretized by the analyzer before physical planning.
|
||||
// Keep the selected boundary schema unchanged here so the physical plan,
|
||||
// remote batches, and local consumers share the same contract.
|
||||
let arrow_schema = Arc::new(arrow_schema.clone());
|
||||
let output_partition_count = Self::output_partition_count(regions.len(), target_partition);
|
||||
|
||||
// States the output ordering of the plan.
|
||||
//
|
||||
@@ -383,7 +475,7 @@ impl MergeScanExec {
|
||||
//
|
||||
// Otherwise, we need to use the default ordering.
|
||||
let eq_properties = if let LogicalPlan::Sort(sort) = &plan
|
||||
&& target_partition >= regions.len()
|
||||
&& output_partition_count >= regions.len()
|
||||
{
|
||||
let lex_ordering = sort
|
||||
.expr
|
||||
@@ -422,7 +514,7 @@ impl MergeScanExec {
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let partitioning = Partitioning::Hash(partition_exprs, target_partition);
|
||||
let partitioning = Partitioning::Hash(partition_exprs, output_partition_count);
|
||||
|
||||
let properties = Arc::new(PlanProperties::new(
|
||||
eq_properties,
|
||||
@@ -449,6 +541,25 @@ impl MergeScanExec {
|
||||
})
|
||||
}
|
||||
|
||||
/// Conservative row-count upper bound for all selected regions.
|
||||
fn estimated_num_rows(&self) -> Precision<usize> {
|
||||
if self.regions.is_empty() {
|
||||
return Precision::Inexact(0);
|
||||
}
|
||||
|
||||
let Some(rows_per_region) = remote_plan_row_bound(&self.plan) else {
|
||||
return Precision::Absent;
|
||||
};
|
||||
rows_per_region
|
||||
.checked_mul(self.regions.len())
|
||||
.map_or(Precision::Absent, Precision::Inexact)
|
||||
}
|
||||
|
||||
/// Number of partitions populated by the region striping in [`Self::to_stream`].
|
||||
fn output_partition_count(num_regions: usize, target_partition: usize) -> usize {
|
||||
num_regions.max(1).min(target_partition.max(1))
|
||||
}
|
||||
|
||||
pub fn to_stream(
|
||||
&self,
|
||||
context: Arc<TaskContext>,
|
||||
@@ -463,7 +574,8 @@ impl MergeScanExec {
|
||||
let sub_stage_metrics_moved = self.sub_stage_metrics.clone();
|
||||
let partition_metrics_moved = self.partition_metrics.clone();
|
||||
let plan = self.plan.clone();
|
||||
let target_partition = self.target_partition;
|
||||
let target_partition =
|
||||
Self::output_partition_count(self.regions.len(), self.target_partition);
|
||||
let remote_dyn_filter_enabled = remote_dyn_filter_enabled(&self.query_ctx)?;
|
||||
let captured_remote_dyn_filters = if remote_dyn_filter_enabled {
|
||||
self.captured_remote_dyn_filters()
|
||||
@@ -769,7 +881,7 @@ impl MergeScanExec {
|
||||
metric: self.metric.clone(),
|
||||
properties: Arc::new(PlanProperties::new(
|
||||
self.properties.eq_properties.clone(),
|
||||
Partitioning::Hash(overlaps, self.target_partition),
|
||||
Partitioning::Hash(overlaps, self.partition_count()),
|
||||
self.properties.emission_type,
|
||||
self.properties.boundedness,
|
||||
)),
|
||||
@@ -820,7 +932,7 @@ impl MergeScanExec {
|
||||
}
|
||||
|
||||
pub fn partition_count(&self) -> usize {
|
||||
self.target_partition
|
||||
Self::output_partition_count(self.regions.len(), self.target_partition)
|
||||
}
|
||||
|
||||
pub fn region_count(&self) -> usize {
|
||||
@@ -837,41 +949,6 @@ impl MergeScanExec {
|
||||
}
|
||||
}
|
||||
|
||||
// If the schema has JSON2 field, AND the field is of empty Struct datatype, amend it with Binary
|
||||
// datatype.
|
||||
// This is a very hacky way to make it possible to query the whole JSON2 column. Because when
|
||||
// querying a whole JSON2 column, like in the SQL `select * from ...`, we can't concretize the JSON2
|
||||
// datatype from the query. Hence, the JSON2 datatype remains what in the column schema, i.e., empty
|
||||
// Struct. An empty Struct is not alignable like any other concretized JSON2 datatypes, so to make
|
||||
// the query work, we amend(rewrite) it to Binary datatype.
|
||||
// Why the Binary datatype? Because underlying the scan and projection stage, the JSON2 data are
|
||||
// variant shape, will be all converted to bytes.
|
||||
// Anyway, this is not clean nor elegant. TODO(LFC) Maybe make it into some plan analyzer rule?
|
||||
fn maybe_amend_json2_field(schema: &ArrowSchema) -> ArrowSchemaRef {
|
||||
let schema = schema.clone();
|
||||
let mut new_fields = Vec::with_capacity(schema.fields().len());
|
||||
for field in schema.fields().iter() {
|
||||
let new_field = if is_json2_extension_type(field)
|
||||
&& matches!(field.data_type(), DataType::Struct(fields) if fields.is_empty())
|
||||
{
|
||||
let is_legacy_json2 = is_legacy_json2_extension_type(field);
|
||||
let mut new_field = field.as_ref().clone();
|
||||
new_field.set_data_type(DataType::Binary);
|
||||
if is_legacy_json2 {
|
||||
new_field = new_field.with_extension_type(Json2ExtensionType::default());
|
||||
}
|
||||
Arc::new(new_field)
|
||||
} else {
|
||||
field.clone()
|
||||
};
|
||||
new_fields.push(new_field);
|
||||
}
|
||||
Arc::new(ArrowSchema::new_with_metadata(
|
||||
new_fields,
|
||||
schema.metadata().clone(),
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl MergeScanExec {
|
||||
fn remote_dyn_filter_producer_id(&self) -> Option<RemoteDynFilterProducerId> {
|
||||
@@ -1085,6 +1162,16 @@ impl ExecutionPlan for MergeScanExec {
|
||||
Some(self.metric.clone_inner())
|
||||
}
|
||||
|
||||
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
|
||||
if partition.is_some() {
|
||||
return Ok(Statistics::new_unknown(&self.arrow_schema));
|
||||
}
|
||||
|
||||
let mut statistics = Statistics::new_unknown(&self.arrow_schema);
|
||||
statistics.num_rows = self.estimated_num_rows();
|
||||
Ok(statistics)
|
||||
}
|
||||
|
||||
fn name(&self) -> &str {
|
||||
"MergeScanExec"
|
||||
}
|
||||
@@ -1231,10 +1318,7 @@ mod tests {
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use arrow_schema::extension::{
|
||||
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
|
||||
};
|
||||
use arrow_schema::{DataType as TestArrowDataType, Field, Fields};
|
||||
use arrow_schema::{DataType as TestArrowDataType, Field};
|
||||
use async_trait::async_trait;
|
||||
use common_query::request::INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY;
|
||||
use common_recordbatch::adapter::{PlanMetrics, RecordBatchMetrics};
|
||||
@@ -1303,6 +1387,30 @@ mod tests {
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn merge_scan_exec_with_plan(
|
||||
regions: Vec<RegionId>,
|
||||
plan: LogicalPlan,
|
||||
target_partition: usize,
|
||||
) -> MergeScanExec {
|
||||
let session_state = SessionStateBuilder::new().build();
|
||||
let schema = plan.schema().as_arrow().clone();
|
||||
|
||||
MergeScanExec::new(
|
||||
&session_state,
|
||||
TableName::new("catalog", "schema", "table"),
|
||||
regions,
|
||||
plan,
|
||||
&schema,
|
||||
Arc::new(TestRegionQueryHandler::default()),
|
||||
QueryContext::arc(),
|
||||
target_partition,
|
||||
AliasMapping::new(),
|
||||
None,
|
||||
false,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn collect_merge_scan(
|
||||
exec: MergeScanExec,
|
||||
) -> datafusion_common::Result<Vec<DfRecordBatch>> {
|
||||
@@ -1328,43 +1436,6 @@ mod tests {
|
||||
metadata
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_amend_legacy_json2_field_preserves_json2_identity() {
|
||||
let field = Field::new("j", TestArrowDataType::Struct(Fields::empty()), true)
|
||||
.with_metadata(StdHashMap::from([
|
||||
(
|
||||
EXTENSION_TYPE_NAME_KEY.to_string(),
|
||||
"greptime.json".to_string(),
|
||||
),
|
||||
(
|
||||
EXTENSION_TYPE_METADATA_KEY.to_string(),
|
||||
serde_json::json!({
|
||||
"json_structure_settings": { "Structured": null }
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
]));
|
||||
|
||||
let legacy_schema = ArrowSchema::new(vec![field]);
|
||||
let amended = maybe_amend_json2_field(&legacy_schema);
|
||||
let amended_field = amended.field(0);
|
||||
assert_eq!(&TestArrowDataType::Binary, amended_field.data_type());
|
||||
assert_eq!(
|
||||
Some(Json2ExtensionType::NAME),
|
||||
amended_field.extension_type_name()
|
||||
);
|
||||
assert!(is_json2_extension_type(amended_field));
|
||||
|
||||
// The remote wire schema is Binary, while the legacy advertised schema
|
||||
// still carries the greptime.json identity. MergeScan must accept the
|
||||
// amended JSON2 field as the corresponding remote column.
|
||||
let wire_schema = ArrowSchema::new(vec![
|
||||
Field::new("j", TestArrowDataType::Binary, true)
|
||||
.with_metadata(legacy_schema.field(0).metadata().clone()),
|
||||
]);
|
||||
assert!(validate_remote_schema(&wire_schema, amended.as_ref(), "legacy json2").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_scan_validates_remote_schema_semantics() {
|
||||
let expected = expected_int64_schema();
|
||||
@@ -1882,6 +1953,118 @@ mod tests {
|
||||
assert!(exec.properties().output_ordering().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_scan_reports_populated_partition_count() {
|
||||
let cases = [(0, 10, 1), (1, 10, 1), (3, 2, 2), (5, 10, 5), (3, 0, 1)];
|
||||
|
||||
for (region_count, target, expected) in cases {
|
||||
let exec = merge_scan_exec_with_sorted_input(region_count, target);
|
||||
assert_eq!(exec.partition_count(), expected);
|
||||
assert_eq!(
|
||||
exec.properties().output_partitioning().partition_count(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_scan_reports_only_deterministic_plan_bounds() {
|
||||
use datafusion::functions_aggregate::expr_fn::count;
|
||||
use datafusion_expr::GroupingSet;
|
||||
|
||||
let regions = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)];
|
||||
let limited = LogicalPlanBuilder::empty(true)
|
||||
.project(vec![lit(1i64).alias("col")])
|
||||
.unwrap()
|
||||
.limit(0, Some(50))
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
merge_scan_exec_with_plan(regions.clone(), limited, 10)
|
||||
.partition_statistics(None)
|
||||
.unwrap()
|
||||
.num_rows,
|
||||
Precision::Inexact(100)
|
||||
);
|
||||
|
||||
let large_bound = i32::MAX as usize + 1;
|
||||
let large_limit = LogicalPlanBuilder::empty(true)
|
||||
.project(vec![lit(1i64).alias("col")])
|
||||
.unwrap()
|
||||
.limit(0, Some(large_bound))
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
merge_scan_exec_with_plan(vec![RegionId::new(1024, 1)], large_limit, 10)
|
||||
.partition_statistics(None)
|
||||
.unwrap()
|
||||
.num_rows,
|
||||
Precision::Inexact(large_bound)
|
||||
);
|
||||
|
||||
let uncapped = LogicalPlanBuilder::empty(true)
|
||||
.project(vec![lit(1i64).alias("col")])
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
merge_scan_exec_with_plan(regions.clone(), uncapped.clone(), 10)
|
||||
.partition_statistics(None)
|
||||
.unwrap()
|
||||
.num_rows,
|
||||
Precision::Absent
|
||||
);
|
||||
assert_eq!(
|
||||
merge_scan_exec_with_plan(Vec::new(), uncapped, 10)
|
||||
.partition_statistics(None)
|
||||
.unwrap()
|
||||
.num_rows,
|
||||
Precision::Inexact(0)
|
||||
);
|
||||
|
||||
let global_aggregate = LogicalPlanBuilder::empty(true)
|
||||
.project(vec![lit(1i64).alias("col")])
|
||||
.unwrap()
|
||||
.limit(0, Some(0))
|
||||
.unwrap()
|
||||
.aggregate(Vec::<Expr>::new(), vec![count(lit(1))])
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
merge_scan_exec_with_plan(regions.clone(), global_aggregate, 10)
|
||||
.partition_statistics(None)
|
||||
.unwrap()
|
||||
.num_rows,
|
||||
Precision::Inexact(2)
|
||||
);
|
||||
|
||||
let grouping_sets = LogicalPlanBuilder::empty(true)
|
||||
.project(vec![lit(1i64).alias("col")])
|
||||
.unwrap()
|
||||
.limit(0, Some(50))
|
||||
.unwrap()
|
||||
.aggregate(
|
||||
vec![Expr::GroupingSet(GroupingSet::GroupingSets(vec![
|
||||
vec![],
|
||||
vec![col("col")],
|
||||
]))],
|
||||
Vec::<Expr>::new(),
|
||||
)
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
merge_scan_exec_with_plan(regions, grouping_sets, 10)
|
||||
.partition_statistics(None)
|
||||
.unwrap()
|
||||
.num_rows,
|
||||
Precision::Absent
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sub_stage_metrics_are_sorted_by_region_id() {
|
||||
let exec = merge_scan_exec_with_sorted_input(0, 1);
|
||||
|
||||
@@ -195,7 +195,7 @@ impl ExtensionPlanner for DistExtensionPlanner {
|
||||
};
|
||||
|
||||
// TODO(ruihang): generate different execution plans for different variant merge operation
|
||||
let schema = optimized_plan.schema().as_arrow();
|
||||
let schema = merge_scan.schema().as_arrow();
|
||||
let query_ctx = session_state
|
||||
.config()
|
||||
.get_extension()
|
||||
|
||||
@@ -45,7 +45,7 @@ pub(crate) mod test_util;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
pub use crate::analyze::analyze_plan_metrics_to_json_value;
|
||||
pub use crate::analyze::{analyze_plan_metrics_to_json_value, dist_analyze_output_schema};
|
||||
pub use crate::datafusion::DfContextProviderAdapter;
|
||||
pub use crate::query_engine::{
|
||||
QueryEngine, QueryEngineContext, QueryEngineFactory, QueryEngineRef,
|
||||
|
||||
@@ -17,6 +17,8 @@ pub mod constant_term;
|
||||
pub mod count_nest_aggr;
|
||||
pub mod count_wildcard;
|
||||
pub mod global_limit;
|
||||
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;
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
// 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 common_time::Timezone;
|
||||
use datafusion::config::ConfigOptions;
|
||||
use datafusion_common::{DFSchemaRef, Result, ScalarValue};
|
||||
use datafusion_expr::expr::{Alias, Cast};
|
||||
use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values};
|
||||
use datafusion_optimizer::analyzer::AnalyzerRule;
|
||||
use datafusion_optimizer::analyzer::type_coercion::TypeCoercion;
|
||||
use datatypes::arrow::datatypes::{DataType, TimeUnit};
|
||||
use session::context::QueryContextRef;
|
||||
|
||||
use crate::optimizer::type_conversion::cast_string_to_timestamp;
|
||||
|
||||
/// Interprets strings assigned to timestamp columns at an `INSERT` boundary
|
||||
/// using the session timezone. `plan` is the assignment projection under a
|
||||
/// `WriteOp::Insert`.
|
||||
///
|
||||
/// DataFusion plans `INSERT` as a projection casting each source column to its
|
||||
/// target column type, and that cast reads a naive string as UTC. Arrow does
|
||||
/// apply a timezone when the cast target carries one, so the assignment is
|
||||
/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone
|
||||
/// afterwards is value-preserving — arrow only shifts values in the opposite
|
||||
/// direction.
|
||||
///
|
||||
/// The source query is left untouched. Reinterpreting a value where it is
|
||||
/// *produced* would change what the source query means: pushing the conversion
|
||||
/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw
|
||||
/// strings to parsed instants and silently drops rows.
|
||||
///
|
||||
/// # Why `TypeCoercion` runs here
|
||||
///
|
||||
/// The rewrite reads source types, and those are only settled once a `UNION`'s
|
||||
/// branch types have been reconciled: before coercion a union carries its loose
|
||||
/// schema (the first branch's types), so a mixed
|
||||
/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a
|
||||
/// string. Retargeting that cast would leave `Timestamp(None) ->
|
||||
/// Timestamp(Some(tz))` behind once coercion retypes the union — the one
|
||||
/// direction in which arrow shifts the value instead of relabelling it.
|
||||
///
|
||||
/// Coercing here rather than deferring to the analyzer is forced by where an
|
||||
/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and
|
||||
/// executes its input, so by the time the analyzer runs, an assignment
|
||||
/// projection is indistinguishable from any other projection.
|
||||
///
|
||||
/// Explicit casts stay out of this: the SQL layer turns a user's
|
||||
/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an
|
||||
/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are
|
||||
/// therefore the only `Expr::Cast` reaching a timestamp column here.
|
||||
///
|
||||
/// # Reach
|
||||
///
|
||||
/// The emitted cast only carries its timezone where the expression is evaluated
|
||||
/// on this node. Substrait drops the timezone name when a plan is pushed down —
|
||||
/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back
|
||||
/// as UTC — so a source reading from a table falls back to UTC, the behaviour it
|
||||
/// had before this rule existed. Sources that never leave this node (literals,
|
||||
/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what
|
||||
/// an INSERT's timestamp assignment is in practice.
|
||||
pub(crate) fn rewrite_insert_assignments(
|
||||
plan: LogicalPlan,
|
||||
query_ctx: &QueryContextRef,
|
||||
config: &ConfigOptions,
|
||||
) -> Result<LogicalPlan> {
|
||||
let Some(timezone) = session_timezone(query_ctx) else {
|
||||
return Ok(plan);
|
||||
};
|
||||
|
||||
let plan = TypeCoercion::new().analyze(plan, config)?;
|
||||
rewrite_assignment(plan, &timezone)
|
||||
}
|
||||
|
||||
/// Session timezone, in both forms the rewrite needs.
|
||||
struct SessionTimezone {
|
||||
/// Parses literals, matching the plain `INSERT ... VALUES` path.
|
||||
parsed: Timezone,
|
||||
/// Names the intermediate arrow cast target.
|
||||
name: Arc<str>,
|
||||
}
|
||||
|
||||
fn session_timezone(query_ctx: &QueryContextRef) -> Option<SessionTimezone> {
|
||||
let parsed = query_ctx.timezone();
|
||||
|
||||
// A UTC session already gets UTC semantics from the plain assignment cast.
|
||||
if parsed.is_utc() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(SessionTimezone {
|
||||
name: Arc::from(parsed.to_string()),
|
||||
parsed,
|
||||
})
|
||||
}
|
||||
|
||||
fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result<LogicalPlan> {
|
||||
let LogicalPlan::Projection(assignment) = plan else {
|
||||
return Ok(plan);
|
||||
};
|
||||
|
||||
let mut exprs = assignment.expr.clone();
|
||||
let mut changed = false;
|
||||
for expr in &mut exprs {
|
||||
changed |= retarget_assignment_cast(
|
||||
expr,
|
||||
assignment.input.schema(),
|
||||
Some(assignment.input.as_ref()),
|
||||
timezone,
|
||||
)?;
|
||||
}
|
||||
|
||||
// The planner types `VALUES` against the target table, so the assignment
|
||||
// cast lands inside the `Values` rows instead of on the projection above.
|
||||
let mut input = assignment.input.clone();
|
||||
if let LogicalPlan::Values(values) = assignment.input.as_ref()
|
||||
&& let Some(rewritten) = rewrite_values(values, timezone)?
|
||||
{
|
||||
input = Arc::new(LogicalPlan::Values(rewritten));
|
||||
changed = true;
|
||||
}
|
||||
|
||||
if !changed {
|
||||
return Ok(LogicalPlan::Projection(assignment));
|
||||
}
|
||||
Projection::try_new(exprs, input).map(LogicalPlan::Projection)
|
||||
}
|
||||
|
||||
fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result<Option<Values>> {
|
||||
let mut rewritten = values.clone();
|
||||
let mut changed = false;
|
||||
for row in &mut rewritten.values {
|
||||
for expr in row.iter_mut() {
|
||||
changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(changed.then_some(rewritten))
|
||||
}
|
||||
|
||||
/// Reinterprets one assignment cast, returning whether it was rewritten.
|
||||
///
|
||||
/// `source_plan` is the projection's input, used to resolve a literal behind a
|
||||
/// column reference; `Values` rows carry their expression inline and pass `None`.
|
||||
fn retarget_assignment_cast(
|
||||
expr: &mut Expr,
|
||||
schema: &DFSchemaRef,
|
||||
source_plan: Option<&LogicalPlan>,
|
||||
timezone: &SessionTimezone,
|
||||
) -> Result<bool> {
|
||||
let expr = unalias_mut(expr);
|
||||
let Expr::Cast(Cast {
|
||||
expr: source,
|
||||
data_type: DataType::Timestamp(unit, None),
|
||||
}) = expr
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
let unit = *unit;
|
||||
|
||||
if !matches!(
|
||||
source.get_type(schema)?,
|
||||
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
|
||||
) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// Fold literals with the same parser the plain `INSERT ... VALUES` path
|
||||
// uses, so a given string means the same thing however it reaches a column.
|
||||
// The parsers disagree on ambiguous local times: this one resolves them,
|
||||
// arrow rejects them.
|
||||
let folded = source_literal(source.as_ref(), source_plan)
|
||||
.and_then(|literal| convert_literal(&literal, unit, &timezone.parsed));
|
||||
if let Some(folded) = folded {
|
||||
*expr = folded;
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let source = source.as_ref().clone();
|
||||
*expr = Expr::Cast(Cast::new(
|
||||
Box::new(Expr::Cast(Cast::new(
|
||||
Box::new(source),
|
||||
DataType::Timestamp(unit, Some(timezone.name.clone())),
|
||||
))),
|
||||
DataType::Timestamp(unit, None),
|
||||
));
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option<ScalarValue> {
|
||||
match source {
|
||||
Expr::Literal(value, _) => Some(value.clone()),
|
||||
Expr::Column(column) => {
|
||||
let plan = source_plan?;
|
||||
let index = plan.schema().maybe_index_of_column(column)?;
|
||||
lineage_literal(plan, index).cloned()
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolves a literal when every row carries the same value at `output_idx`.
|
||||
///
|
||||
/// Read-only: the literal is folded into the assignment above, so nodes that
|
||||
/// drop, reorder or deduplicate rows can be traversed — none of them changes
|
||||
/// the value a surviving row carries, and folding above them leaves their keys
|
||||
/// on the original strings.
|
||||
fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> {
|
||||
if output_idx >= plan.schema().fields().len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
match plan {
|
||||
LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) {
|
||||
Expr::Literal(value, _) => Some(value),
|
||||
Expr::Column(column) => {
|
||||
let input_idx = projection.input.schema().maybe_index_of_column(column)?;
|
||||
lineage_literal(projection.input.as_ref(), input_idx)
|
||||
}
|
||||
_ => None,
|
||||
},
|
||||
LogicalPlan::Filter(_)
|
||||
| LogicalPlan::Sort(_)
|
||||
| LogicalPlan::Limit(_)
|
||||
| LogicalPlan::SubqueryAlias(_)
|
||||
| LogicalPlan::Distinct(Distinct::All(_)) => {
|
||||
let inputs = plan.inputs();
|
||||
let [input] = inputs.as_slice() else {
|
||||
return None;
|
||||
};
|
||||
lineage_literal(input, output_idx)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option<Expr> {
|
||||
let ScalarValue::Utf8(Some(value)) = value else {
|
||||
return None;
|
||||
};
|
||||
cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone))
|
||||
.ok()
|
||||
.filter(|value| !value.is_null())
|
||||
.map(|value| Expr::Literal(value, None))
|
||||
}
|
||||
|
||||
fn unalias(expr: &Expr) -> &Expr {
|
||||
match expr {
|
||||
Expr::Alias(Alias { expr, .. }) => unalias(expr),
|
||||
expr => expr,
|
||||
}
|
||||
}
|
||||
|
||||
fn unalias_mut(expr: &mut Expr) -> &mut Expr {
|
||||
match expr {
|
||||
Expr::Alias(Alias { expr, .. }) => unalias_mut(expr),
|
||||
expr => expr,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use datafusion_common::DFSchema;
|
||||
use datafusion_expr::expr::Placeholder;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn shanghai() -> SessionTimezone {
|
||||
let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap();
|
||||
SessionTimezone {
|
||||
name: Arc::from(parsed.to_string()),
|
||||
parsed,
|
||||
}
|
||||
}
|
||||
|
||||
/// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped
|
||||
/// placeholder, which must survive for parameter substitution.
|
||||
#[test]
|
||||
fn test_untyped_placeholder_assignment_is_left_alone() {
|
||||
let schema = Arc::new(DFSchema::empty());
|
||||
let mut expr = Expr::Cast(Cast::new(
|
||||
Box::new(Expr::Placeholder(Placeholder::new_with_field(
|
||||
"$1".to_string(),
|
||||
None,
|
||||
))),
|
||||
DataType::Timestamp(TimeUnit::Millisecond, None),
|
||||
));
|
||||
let original = expr.clone();
|
||||
|
||||
let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap();
|
||||
|
||||
assert!(!changed);
|
||||
assert_eq!(expr, original);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
// 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::sync::Arc;
|
||||
|
||||
use datafusion::config::ConfigOptions;
|
||||
use datafusion_common::tree_node::Transformed;
|
||||
use datafusion_common::{DFSchema, DFSchemaRef, Result};
|
||||
use datafusion_expr::{LogicalPlan, UserDefinedLogicalNodeCore};
|
||||
use datafusion_optimizer::analyzer::AnalyzerRule;
|
||||
use datatypes::extension::json::{
|
||||
Json2ExtensionType, is_json2_extension_type, is_legacy_json2_extension_type,
|
||||
};
|
||||
use datatypes::types::json_type::JsonNativeType;
|
||||
|
||||
use crate::dist_plan::MergeScanLogicalPlan;
|
||||
use crate::optimizer::json_type_concretize::deduce_json_types;
|
||||
|
||||
/// Keeps JSON2 schemas consistent across distributed query boundaries.
|
||||
///
|
||||
/// An unresolved JSON2 column is represented as an empty `Struct`, while a remote stage always
|
||||
/// emits a concrete Arrow array. DataFusion expects the schema declared by a logical node, the
|
||||
/// schema used to build its physical plan, and the schema of its record batches to agree. This rule
|
||||
/// gives each [`MergeScanLogicalPlan`] the concrete schema emitted by its remote stage and propagates
|
||||
/// that schema through the local plan.
|
||||
///
|
||||
/// For example:
|
||||
///
|
||||
/// - `SELECT j FROM t` transfers the complete JSON2 value as `Binary` (`Variant`). A projection or
|
||||
/// window above the MergeScan must therefore also describe `j` as `Binary`, not an empty `Struct`.
|
||||
/// - `SELECT j.a FROM t` may transfer the complete `j` as `Binary` and extract `a` locally, or
|
||||
/// transfer only the extracted scalar when the expression runs remotely. The boundary schema
|
||||
/// must describe the remote output rather than the local expression that consumes it.
|
||||
/// - `SELECT l.j, r.j FROM l JOIN r ON l.k = r.k` has two independent boundaries. Each input and the
|
||||
/// join schema must agree on the concrete type of its JSON2 column.
|
||||
///
|
||||
/// Correcting only the physical MergeScan schema can make simple queries work because many
|
||||
/// operators access columns by position, but it leaves the logical plan describing a different
|
||||
/// type. Keeping the schemas consistent lets optimizers, physical planners, validators, and future
|
||||
/// type-aware operators rely on the normal DataFusion contract. It also keeps the generic physical
|
||||
/// MergeScan implementation independent of JSON2.
|
||||
///
|
||||
/// The boundary schema and the storage read layout answer different questions. For
|
||||
/// `SELECT j.a::BIGINT FROM t`, the storage scan may use a structured `{a: Int64}` layout, while the
|
||||
/// distributed boundary can still emit either the complete `j` as `Binary` or only `a` as `Int64`.
|
||||
/// Therefore this rule cannot replace the separate JSON2 scan-type inference.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct JsonSchemaConcretizeRule;
|
||||
|
||||
impl AnalyzerRule for JsonSchemaConcretizeRule {
|
||||
fn analyze(&self, plan: LogicalPlan, _config: &ConfigOptions) -> Result<LogicalPlan> {
|
||||
let plan = plan.transform_up_with_subqueries(|plan| {
|
||||
let LogicalPlan::Extension(mut extension) = plan else {
|
||||
return Ok(Transformed::no(plan));
|
||||
};
|
||||
let Some(merge_scan) = extension
|
||||
.node
|
||||
.as_any()
|
||||
.downcast_ref::<MergeScanLogicalPlan>()
|
||||
else {
|
||||
return Ok(Transformed::no(LogicalPlan::Extension(extension)));
|
||||
};
|
||||
|
||||
// Infer the boundary schema from the hidden remote plan, not its local consumer.
|
||||
let json_types = deduce_json_types(merge_scan.input())?;
|
||||
if json_types.is_empty() {
|
||||
return Ok(Transformed::no(LogicalPlan::Extension(extension)));
|
||||
}
|
||||
let schema = concretize_json2_schema(merge_scan.schema(), &json_types)?;
|
||||
if schema.as_ref() == merge_scan.schema().as_ref() {
|
||||
return Ok(Transformed::no(LogicalPlan::Extension(extension)));
|
||||
}
|
||||
|
||||
extension.node = Arc::new(merge_scan.clone().with_output_schema(schema));
|
||||
Ok(Transformed::yes(LogicalPlan::Extension(extension)))
|
||||
})?;
|
||||
|
||||
if plan.transformed {
|
||||
plan.data
|
||||
.transform_up_with_subqueries(|plan| {
|
||||
if matches!(plan, LogicalPlan::Extension(_)) || plan.inputs().is_empty() {
|
||||
Ok(Transformed::no(plan))
|
||||
} else {
|
||||
plan.recompute_schema().map(Transformed::yes)
|
||||
}
|
||||
})
|
||||
.map(|x| x.data)
|
||||
} else {
|
||||
Ok(plan.data)
|
||||
}
|
||||
}
|
||||
|
||||
fn name(&self) -> &str {
|
||||
"JsonSchemaConcretizeRule"
|
||||
}
|
||||
}
|
||||
|
||||
fn concretize_json2_schema(
|
||||
schema: &DFSchemaRef,
|
||||
json_types: &HashMap<String, JsonNativeType>,
|
||||
) -> Result<DFSchemaRef> {
|
||||
if !schema
|
||||
.iter()
|
||||
.any(|(_, field)| json_types.contains_key(field.name()) && is_json2_extension_type(field))
|
||||
{
|
||||
return Ok(schema.clone());
|
||||
}
|
||||
|
||||
let mut changed = false;
|
||||
let fields = schema
|
||||
.iter()
|
||||
.map(|(qualifier, field)| {
|
||||
let Some(json_type) = json_types
|
||||
.get(field.name())
|
||||
.filter(|_| is_json2_extension_type(field))
|
||||
else {
|
||||
return (qualifier.cloned(), field.clone());
|
||||
};
|
||||
let data_type = json_type.as_arrow_type();
|
||||
if field.data_type() == &data_type {
|
||||
return (qualifier.cloned(), field.clone());
|
||||
}
|
||||
|
||||
changed = true;
|
||||
|
||||
// Before type hints, JSON2 used the `greptime.json` marker together with
|
||||
// `json_structure_settings`. Once concretized to `Binary`, that field no longer
|
||||
// matches the legacy JSON2 shape and could be mistaken for JSONB, so upgrade its
|
||||
// marker. Do not replace modern markers because that would discard their JSON
|
||||
// settings and layout version.
|
||||
let legacy = is_legacy_json2_extension_type(field);
|
||||
let mut field = field.as_ref().clone().with_data_type(data_type);
|
||||
if legacy {
|
||||
field = field.with_extension_type(Json2ExtensionType::default());
|
||||
}
|
||||
(qualifier.cloned(), Arc::new(field))
|
||||
})
|
||||
.collect();
|
||||
|
||||
if changed {
|
||||
let schema = DFSchema::new_with_metadata(fields, schema.metadata().clone())?
|
||||
.with_functional_dependencies(schema.functional_dependencies().clone())?;
|
||||
Ok(Arc::new(schema))
|
||||
} else {
|
||||
Ok(schema.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use arrow_schema::extension::{
|
||||
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
use datafusion_common::DFSchema;
|
||||
use datafusion_expr::{LogicalPlanBuilder, col};
|
||||
use datatypes::extension::json::{JsonExtensionType, is_json2_extension_type};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_json_schema_concretize_rule_updates_merge_scan() -> Result<()> {
|
||||
let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata(
|
||||
HashMap::from([
|
||||
(
|
||||
EXTENSION_TYPE_NAME_KEY.to_string(),
|
||||
JsonExtensionType::NAME.to_string(),
|
||||
),
|
||||
(
|
||||
EXTENSION_TYPE_METADATA_KEY.to_string(),
|
||||
serde_json::json!({
|
||||
"json_structure_settings": { "Structured": null }
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
]),
|
||||
);
|
||||
let schema = Arc::new(DFSchema::try_from(Schema::new(vec![field]))?);
|
||||
let input = LogicalPlan::EmptyRelation(datafusion_expr::logical_plan::EmptyRelation {
|
||||
produce_one_row: false,
|
||||
schema,
|
||||
});
|
||||
let merge_scan =
|
||||
MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan();
|
||||
let plan = LogicalPlanBuilder::from(merge_scan)
|
||||
.project(vec![col("j")])?
|
||||
.build()?;
|
||||
|
||||
let plan = JsonSchemaConcretizeRule.analyze(plan, &ConfigOptions::default())?;
|
||||
let field = plan.schema().field(0);
|
||||
assert_eq!(&DataType::Binary, field.data_type());
|
||||
assert_eq!(Some(Json2ExtensionType::NAME), field.extension_type_name());
|
||||
assert!(is_json2_extension_type(field));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -108,16 +108,36 @@ fn apply_json_type_hint(
|
||||
false
|
||||
}
|
||||
|
||||
fn deduce_json_types(plan: &LogicalPlan) -> Result<HashMap<String, JsonNativeType>> {
|
||||
pub(crate) fn deduce_json_types(plan: &LogicalPlan) -> Result<HashMap<String, JsonNativeType>> {
|
||||
let mut json_types = HashMap::<String, JsonNativeType>::new();
|
||||
|
||||
// JSON2 columns in the final output must retain their complete values even when
|
||||
// predicates or other expressions access only specific paths.
|
||||
// For example, `SELECT j FROM t WHERE json_get(j, 'a') = 1`.
|
||||
plan.schema()
|
||||
.fields()
|
||||
.iter()
|
||||
.filter(|field| is_json2_extension_type(field))
|
||||
.for_each(|field| {
|
||||
json_types.insert(field.name().clone(), JsonNativeType::Variant);
|
||||
});
|
||||
|
||||
plan.apply(|plan| {
|
||||
for expr in plan.expressions() {
|
||||
// Optimizer-generated projections may keep the JSON root only so later json_get
|
||||
// expressions can access another path. A same-name pass-through does not require the
|
||||
// complete root by itself; any real whole-column consumer above it is visited
|
||||
// separately, and a whole root in the final output is captured from the plan schema.
|
||||
if matches!(plan, LogicalPlan::Projection(_)) && is_same_name_column_projection(&expr) {
|
||||
continue;
|
||||
}
|
||||
expr.apply(|expr| {
|
||||
if let Some((column, json_type)) = deduce_json_type(expr)? {
|
||||
json_types.entry(column).or_default().merge(&json_type);
|
||||
Ok(TreeNodeRecursion::Jump)
|
||||
} else {
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
}
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
})?;
|
||||
}
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
@@ -125,9 +145,20 @@ fn deduce_json_types(plan: &LogicalPlan) -> Result<HashMap<String, JsonNativeTyp
|
||||
Ok(json_types)
|
||||
}
|
||||
|
||||
fn is_same_name_column_projection(expr: &Expr) -> bool {
|
||||
match expr {
|
||||
Expr::Column(_) => true,
|
||||
Expr::Alias(alias) => {
|
||||
matches!(alias.expr.as_ref(), Expr::Column(column) if column.name == alias.name)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn deduce_json_type(expr: &Expr) -> Result<Option<(String, JsonNativeType)>> {
|
||||
let f = match expr {
|
||||
Expr::ScalarFunction(f) if f.name().eq_ignore_ascii_case(JsonGetWithType::NAME) => f,
|
||||
Expr::Column(c) => return Ok(Some((c.name.clone(), JsonNativeType::Variant))),
|
||||
_ => return Ok(None),
|
||||
};
|
||||
|
||||
@@ -153,6 +184,12 @@ fn deduce_json_type(expr: &Expr) -> Result<Option<(String, JsonNativeType)>> {
|
||||
);
|
||||
};
|
||||
|
||||
// Object-only type deduction cannot represent bracket JSONPath access, so preserve the
|
||||
// full Variant and let json_get apply the expression.
|
||||
if path.contains('[') {
|
||||
return Ok(Some((column.name.clone(), JsonNativeType::Variant)));
|
||||
}
|
||||
|
||||
let with_type = f
|
||||
.args
|
||||
.get(2)
|
||||
@@ -293,6 +330,17 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deduce_json_type_with_list_index() -> Result<()> {
|
||||
let expr = json_get_expr(col("j"), path_expr("l[0]"), Some(DataType::Int64))?;
|
||||
|
||||
assert_eq!(
|
||||
Some(("j".to_string(), JsonNativeType::Variant)),
|
||||
deduce_json_type(&expr)?
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_type_concretize_rule_conflict_to_variant() -> Result<()> {
|
||||
let exprs = vec![
|
||||
@@ -360,25 +408,13 @@ mod tests {
|
||||
.rewrite(plan, &OptimizerContext::default())?
|
||||
.transformed
|
||||
);
|
||||
assert!(provider.scan_request().json_type_hint.contains_key("j"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allow_json2_passthrough_for_later_projection() -> Result<()> {
|
||||
let json_get = json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?;
|
||||
let (provider, plan) = build_json2_scan()?;
|
||||
let plan = plan
|
||||
.project(vec![json_get.alias("__common_expr"), col("j")])?
|
||||
.aggregate(Vec::<Expr>::new(), vec![count(lit(1))])?
|
||||
.build()?;
|
||||
|
||||
assert!(
|
||||
JsonTypeConcretizeRule
|
||||
.rewrite(plan, &OptimizerContext::default())?
|
||||
.transformed
|
||||
assert_eq!(
|
||||
Some(&JsonNativeType::Object(JsonObjectType::from([(
|
||||
"a".to_string(),
|
||||
JsonNativeType::i64(),
|
||||
)]))),
|
||||
provider.scan_request().json_type_hint.get("j")
|
||||
);
|
||||
assert!(provider.scan_request().json_type_hint.contains_key("j"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -402,6 +438,25 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allow_json2_filter_with_root_projection() -> Result<()> {
|
||||
let predicate =
|
||||
json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.eq(lit(1_i64));
|
||||
let (provider, plan) = build_json2_scan()?;
|
||||
let plan = plan.filter(predicate)?.build()?;
|
||||
|
||||
assert!(
|
||||
JsonTypeConcretizeRule
|
||||
.rewrite(plan, &OptimizerContext::default())?
|
||||
.transformed
|
||||
);
|
||||
assert_eq!(
|
||||
Some(&JsonNativeType::Variant),
|
||||
provider.scan_request().json_type_hint.get("j")
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deduce_json_type_with_non_column_base() -> Result<()> {
|
||||
let expr = json_get_expr(
|
||||
|
||||
@@ -270,8 +270,8 @@ mod tests {
|
||||
panic!("expected right merge scan hash partitioning");
|
||||
};
|
||||
|
||||
assert_eq!(*left_count, 32);
|
||||
assert_eq!(*right_count, 32);
|
||||
assert_eq!(*left_count, 2);
|
||||
assert_eq!(*right_count, 2);
|
||||
assert_eq!(
|
||||
column_names(left_exprs),
|
||||
vec![DATA_SCHEMA_TSID_COLUMN_NAME, "greptime_timestamp"]
|
||||
|
||||
@@ -12,8 +12,6 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
mod insert_assignment;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_time::Timezone;
|
||||
@@ -31,7 +29,7 @@ use session::context::QueryContextRef;
|
||||
|
||||
use crate::QueryEngineContext;
|
||||
use crate::optimizer::ExtensionAnalyzerRule;
|
||||
use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments;
|
||||
use crate::optimizer::insert_assignment::rewrite_insert_assignments;
|
||||
use crate::plan::ExtractExpr;
|
||||
|
||||
/// TypeConversionRule converts some literal values in logical plan to other types according
|
||||
@@ -46,7 +44,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
|
||||
&self,
|
||||
plan: LogicalPlan,
|
||||
ctx: &QueryEngineContext,
|
||||
_config: &ConfigOptions,
|
||||
config: &ConfigOptions,
|
||||
) -> Result<LogicalPlan> {
|
||||
plan.transform_up_with_subqueries(|plan| match plan {
|
||||
LogicalPlan::Filter(filter) => {
|
||||
@@ -129,7 +127,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
|
||||
LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => {
|
||||
dml.input = Arc::new(rewrite_insert_assignments(
|
||||
dml.input.as_ref().clone(),
|
||||
ctx.query_ctx(),
|
||||
&ctx.query_ctx(),
|
||||
config,
|
||||
)?);
|
||||
Ok(Transformed::yes(LogicalPlan::Dml(dml)))
|
||||
}
|
||||
@@ -324,7 +323,7 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr {
|
||||
)
|
||||
}
|
||||
|
||||
fn cast_string_to_timestamp(
|
||||
pub(crate) fn cast_string_to_timestamp(
|
||||
string: &str,
|
||||
target_type: &DataType,
|
||||
timezone: Option<&Timezone>,
|
||||
|
||||
@@ -1,336 +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_common::{Column, Result, ScalarValue};
|
||||
use datafusion_expr::expr::{Alias, Cast};
|
||||
use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values};
|
||||
use datatypes::arrow::datatypes::DataType;
|
||||
use session::context::QueryContextRef;
|
||||
|
||||
use crate::optimizer::type_conversion::cast_string_to_timestamp;
|
||||
use crate::plan::ExtractExpr;
|
||||
|
||||
/// Rewrites string literals that feed timestamp columns at an INSERT boundary.
|
||||
/// Constants are folded at the assignment to avoid changing source types;
|
||||
/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts
|
||||
/// stay on DataFusion's existing path.
|
||||
pub(super) fn rewrite_insert_assignments(
|
||||
plan: LogicalPlan,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<LogicalPlan> {
|
||||
let LogicalPlan::Projection(assignment) = plan else {
|
||||
return Ok(plan);
|
||||
};
|
||||
|
||||
let converter = InsertAssignmentConverter { query_ctx };
|
||||
let mut exprs = assignment.expr.clone();
|
||||
let mut input = assignment.input.as_ref().clone();
|
||||
let mut changed = false;
|
||||
for (output_idx, expr) in assignment.expr.iter().enumerate() {
|
||||
let target_type = assignment.schema.field(output_idx).data_type();
|
||||
if !matches!(target_type, DataType::Timestamp(_, _)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(column) = assignment_input_column(expr) else {
|
||||
continue;
|
||||
};
|
||||
let Some(input_idx) = input.schema().maybe_index_of_column(column) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let literal = lineage_literal(&input, input_idx).cloned();
|
||||
if let Some(literal) = literal
|
||||
&& let Some(folded) = converter.convert_literal(&literal, target_type)
|
||||
{
|
||||
let (qualifier, field) = assignment.schema.qualified_field(output_idx);
|
||||
exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name());
|
||||
changed = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
// A hand-built DML plan (e.g. from flow) may share one source column
|
||||
// between targets; rewriting it in place would retype every reader.
|
||||
if assignment
|
||||
.expr
|
||||
.iter()
|
||||
.enumerate()
|
||||
.any(|(idx, other)| idx != output_idx && other.column_refs().contains(column))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? {
|
||||
input = rewritten;
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
if !changed {
|
||||
return Ok(LogicalPlan::Projection(assignment));
|
||||
}
|
||||
Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection)
|
||||
}
|
||||
|
||||
/// Resolves a literal when every row carries the same value at `output_idx`.
|
||||
fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> {
|
||||
if output_idx >= plan.schema().fields().len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
match plan {
|
||||
LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) {
|
||||
Expr::Literal(value, _) => Some(value),
|
||||
Expr::Column(column) => {
|
||||
let input_idx = projection.input.schema().maybe_index_of_column(column)?;
|
||||
lineage_literal(projection.input.as_ref(), input_idx)
|
||||
}
|
||||
_ => None,
|
||||
},
|
||||
// These nodes drop, reorder or deduplicate rows without touching the
|
||||
// value a surviving row carries, and their schema stays positional.
|
||||
LogicalPlan::Filter(_)
|
||||
| LogicalPlan::Sort(_)
|
||||
| LogicalPlan::Limit(_)
|
||||
| LogicalPlan::SubqueryAlias(_)
|
||||
| LogicalPlan::Distinct(Distinct::All(_)) => {
|
||||
let inputs = plan.inputs();
|
||||
let [input] = inputs.as_slice() else {
|
||||
return None;
|
||||
};
|
||||
lineage_literal(input, output_idx)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
struct InsertAssignmentConverter {
|
||||
query_ctx: QueryContextRef,
|
||||
}
|
||||
|
||||
impl InsertAssignmentConverter {
|
||||
fn rewrite_output_column(
|
||||
&self,
|
||||
plan: &LogicalPlan,
|
||||
output_idx: usize,
|
||||
target_type: &DataType,
|
||||
) -> Result<Option<LogicalPlan>> {
|
||||
if output_idx >= plan.schema().fields().len() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
match plan {
|
||||
// DataFusion pushes INSERT assignment casts into Values, so inspect
|
||||
// them even though the Values schema already has the target type.
|
||||
LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type),
|
||||
// Once the source query has produced the target type, its casts and
|
||||
// coercions belong to the source query rather than INSERT assignment.
|
||||
_ if plan.schema().field(output_idx).data_type() == target_type => Ok(None),
|
||||
LogicalPlan::Projection(projection) => {
|
||||
self.rewrite_projection(projection, output_idx, target_type)
|
||||
}
|
||||
LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type),
|
||||
// Filter and Sort are excluded here: their predicates and keys read
|
||||
// the retyped column, which would change the source query. Constants
|
||||
// still reach the assignment through `lineage_literal`. Distinct::All
|
||||
// holds no expressions; allowing it shifts dedup keys from raw
|
||||
// strings to parsed instants.
|
||||
LogicalPlan::Limit(_)
|
||||
| LogicalPlan::SubqueryAlias(_)
|
||||
| LogicalPlan::Distinct(Distinct::All(_)) => {
|
||||
self.rewrite_passthrough(plan, output_idx, target_type)
|
||||
}
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn rewrite_projection(
|
||||
&self,
|
||||
projection: &Projection,
|
||||
output_idx: usize,
|
||||
target_type: &DataType,
|
||||
) -> Result<Option<LogicalPlan>> {
|
||||
match unalias(&projection.expr[output_idx]) {
|
||||
Expr::Literal(value, _) => {
|
||||
let Some(converted) = self.convert_literal(value, target_type) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let (qualifier, field) = projection.schema.qualified_field(output_idx);
|
||||
let mut exprs = projection.expr.clone();
|
||||
exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name());
|
||||
|
||||
Projection::try_new(exprs, projection.input.clone())
|
||||
.map(LogicalPlan::Projection)
|
||||
.map(Some)
|
||||
}
|
||||
Expr::Column(column) => {
|
||||
// Retyping the input column affects every output column reading
|
||||
// it, so only follow lineage with a single consumer.
|
||||
if projection
|
||||
.expr
|
||||
.iter()
|
||||
.enumerate()
|
||||
.any(|(idx, other)| idx != output_idx && other.column_refs().contains(column))
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let input = projection.input.as_ref();
|
||||
let Some(input_idx) = input.schema().maybe_index_of_column(column) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
Projection::try_new(projection.expr.clone(), Arc::new(rewritten))
|
||||
.map(LogicalPlan::Projection)
|
||||
.map(Some)
|
||||
}
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn rewrite_values(
|
||||
&self,
|
||||
values: &Values,
|
||||
output_idx: usize,
|
||||
target_type: &DataType,
|
||||
) -> Result<Option<LogicalPlan>> {
|
||||
let mut rewritten = values.clone();
|
||||
let mut changed = false;
|
||||
for row in &mut rewritten.values {
|
||||
let Some(expr) = row.get_mut(output_idx) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if let Expr::Cast(Cast {
|
||||
expr: inner,
|
||||
data_type,
|
||||
}) = expr
|
||||
&& data_type == target_type
|
||||
&& let Expr::Literal(value, _) = inner.as_ref()
|
||||
&& let Some(value) = self.convert_literal(value, target_type)
|
||||
{
|
||||
*expr = value;
|
||||
changed = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(changed.then_some(LogicalPlan::Values(rewritten)))
|
||||
}
|
||||
|
||||
fn rewrite_union(
|
||||
&self,
|
||||
union: &Union,
|
||||
output_idx: usize,
|
||||
target_type: &DataType,
|
||||
) -> Result<Option<LogicalPlan>> {
|
||||
let mut inputs = Vec::with_capacity(union.inputs.len());
|
||||
// A union has one shared schema, so rewrite it only when the same output
|
||||
// column is an INSERT-assigned literal in every branch.
|
||||
for input in &union.inputs {
|
||||
let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
inputs.push(Arc::new(rewritten));
|
||||
}
|
||||
|
||||
// Untouched columns may still have mismatched branch types before
|
||||
// TypeCoercion runs, so rebuild loosely like the SQL planner does.
|
||||
Union::try_new_with_loose_types(inputs)
|
||||
.map(LogicalPlan::Union)
|
||||
.map(Some)
|
||||
}
|
||||
|
||||
fn rewrite_passthrough(
|
||||
&self,
|
||||
plan: &LogicalPlan,
|
||||
output_idx: usize,
|
||||
target_type: &DataType,
|
||||
) -> Result<Option<LogicalPlan>> {
|
||||
let inputs = plan.inputs();
|
||||
let [input] = inputs.as_slice() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input])
|
||||
.map(Some)
|
||||
}
|
||||
|
||||
fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option<Expr> {
|
||||
let ScalarValue::Utf8(Some(value)) = value else {
|
||||
return None;
|
||||
};
|
||||
cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone()))
|
||||
.ok()
|
||||
.filter(|value| !value.is_null())
|
||||
.map(|value| Expr::Literal(value, None))
|
||||
}
|
||||
}
|
||||
|
||||
fn assignment_input_column(expr: &Expr) -> Option<&Column> {
|
||||
let expr = match unalias(expr) {
|
||||
Expr::Cast(Cast { expr, .. }) => expr.as_ref(),
|
||||
expr => expr,
|
||||
};
|
||||
let Expr::Column(column) = expr else {
|
||||
return None;
|
||||
};
|
||||
Some(column)
|
||||
}
|
||||
|
||||
fn unalias(expr: &Expr) -> &Expr {
|
||||
match expr {
|
||||
Expr::Alias(Alias { expr, .. }) => unalias(expr),
|
||||
expr => expr,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_common::arrow::datatypes::TimeUnit;
|
||||
use session::context::QueryContext;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_convert_literal_falls_back_for_unsupported_literal() {
|
||||
let converter = InsertAssignmentConverter {
|
||||
query_ctx: QueryContext::arc(),
|
||||
};
|
||||
let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None);
|
||||
|
||||
for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] {
|
||||
assert_eq!(
|
||||
converter
|
||||
.convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
+141
-63
@@ -15,6 +15,7 @@
|
||||
use std::any::Any;
|
||||
use std::borrow::Cow;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::ops::ControlFlow;
|
||||
use std::str::FromStr;
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -34,7 +35,8 @@ use datafusion_expr::{
|
||||
Analyze, Explain, ExplainFormat, Expr as DfExpr, LogicalPlan, LogicalPlanBuilder, PlanType,
|
||||
ToStringifiedPlan, col,
|
||||
};
|
||||
use datafusion_sql::planner::{ParserOptions, SqlToRel};
|
||||
use datafusion_sql::parser::Statement as DfStatement;
|
||||
use datafusion_sql::planner::{IdentNormalizer, ParserOptions, SqlToRel};
|
||||
use log_query::LogQuery;
|
||||
use promql_parser::parser::EvalStmt;
|
||||
use session::context::QueryContextRef;
|
||||
@@ -45,6 +47,7 @@ use sql::statements::explain::ExplainStatement;
|
||||
use sql::statements::query::Query;
|
||||
use sql::statements::statement::Statement;
|
||||
use sql::statements::tql::Tql;
|
||||
use sqlparser::ast::{AccessExpr, Value, visit_expressions_mut};
|
||||
|
||||
use crate::error::{
|
||||
CteColumnSchemaMismatchSnafu, PlanSqlSnafu, QueryPlanSnafu, Result, SqlSnafu,
|
||||
@@ -192,6 +195,13 @@ impl DfLogicalPlanner {
|
||||
}
|
||||
|
||||
let mut df_stmt = stmt.as_ref().try_into().context(SqlSnafu)?;
|
||||
normalize_field_access_after_subscript(
|
||||
&mut df_stmt,
|
||||
self.session_state
|
||||
.config_options()
|
||||
.sql_parser
|
||||
.enable_ident_normalization,
|
||||
);
|
||||
|
||||
// TODO(LFC): Remove this when Datafusion supports **both** the syntax and implementation of "explain with format".
|
||||
if let datafusion::sql::parser::Statement::Statement(
|
||||
@@ -595,6 +605,72 @@ impl DfLogicalPlanner {
|
||||
}
|
||||
}
|
||||
|
||||
/// Normalizes dot field accesses that follow a subscript for DataFusion.
|
||||
///
|
||||
/// sqlparser represents `j.o.l[1].inner.l[2]` as a compound field access with
|
||||
/// the following access chain:
|
||||
///
|
||||
/// ```text
|
||||
/// Dot(Identifier("o")),
|
||||
/// Dot(Identifier("l")),
|
||||
/// Subscript(1),
|
||||
/// Dot(Identifier("inner")),
|
||||
/// Dot(Identifier("l")),
|
||||
/// Subscript(2)
|
||||
/// ```
|
||||
///
|
||||
/// DataFusion first resolves the leading `j.o.l` through
|
||||
/// `JsonExprPlanner::plan_compound_identifier`, which produces an untyped
|
||||
/// `json_get` with path `o.l`. Before invoking `JsonExprPlanner::plan_field_access`,
|
||||
/// however, DataFusion eagerly converts every remaining access into a
|
||||
/// `GetFieldAccess`. It accepts string values but not [`SqlExpr::Identifier`]s
|
||||
/// in [`AccessExpr::Dot`] after a subscript. Without this normalization, that
|
||||
/// conversion fails at `.inner`, and `plan_field_access` is never called, even
|
||||
/// for the preceding `[1]`.
|
||||
///
|
||||
/// This function converts dot identifiers after the first subscript into
|
||||
/// `Dot(Value(SingleQuotedString(...)))`, applying DataFusion's identifier
|
||||
/// normalization before discarding whether each identifier was quoted. It
|
||||
/// changes neither the SQL text nor the dot accesses into subscript nodes: the
|
||||
/// resulting AST is conceptually `j.o.l[1].'inner'.'l'[2]`. DataFusion converts
|
||||
/// the string-valued dot accesses into named field accesses, which
|
||||
/// `plan_field_access` safely encodes as bracket members. It can then extend the
|
||||
/// JSON path to `o.l[1]["inner"]["l"][2]`.
|
||||
///
|
||||
/// This behavior is unchanged in the latest upstream releases checked here:
|
||||
/// DataFusion 55.0.0 and sqlparser 0.62.0.
|
||||
///
|
||||
/// TODO(LFC): Remove this workaround after upstream supports dot identifiers after subscripts.
|
||||
fn normalize_field_access_after_subscript(stmt: &mut DfStatement, normalize_ident: bool) {
|
||||
let DfStatement::Statement(stmt) = stmt else {
|
||||
return;
|
||||
};
|
||||
let normalizer = IdentNormalizer::new(normalize_ident);
|
||||
|
||||
let _ = visit_expressions_mut(stmt.as_mut(), |expr| {
|
||||
let SqlExpr::CompoundFieldAccess { access_chain, .. } = expr else {
|
||||
return ControlFlow::<()>::Continue(());
|
||||
};
|
||||
let Some(index) = access_chain
|
||||
.iter()
|
||||
.position(|x| matches!(x, AccessExpr::Subscript(_)))
|
||||
else {
|
||||
return ControlFlow::Continue(());
|
||||
};
|
||||
|
||||
for access in &mut access_chain[index + 1..] {
|
||||
let AccessExpr::Dot(SqlExpr::Identifier(ident)) = access else {
|
||||
continue;
|
||||
};
|
||||
let value = normalizer.normalize(ident.clone());
|
||||
*access = AccessExpr::Dot(SqlExpr::Value(
|
||||
Value::SingleQuotedString(value).with_span(ident.span),
|
||||
));
|
||||
}
|
||||
ControlFlow::Continue(())
|
||||
});
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LogicalPlanner for DfLogicalPlanner {
|
||||
#[tracing::instrument(skip_all)]
|
||||
@@ -852,6 +928,30 @@ mod tests {
|
||||
engine.planner().plan(&stmt, query_ctx).await.unwrap()
|
||||
}
|
||||
|
||||
/// Plans `sql` and runs the DataFusion analyzer, which is where
|
||||
/// `InsertAssignmentRule` sits. Planning alone stops short of it, so these
|
||||
/// assertions would not see the assignment rewrite at all.
|
||||
async fn analyze_insert(
|
||||
engine: &QueryEngineRef,
|
||||
sql: &str,
|
||||
query_ctx: &QueryContextRef,
|
||||
) -> String {
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap();
|
||||
let plan = engine
|
||||
.planner()
|
||||
.plan(&stmt, query_ctx.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
let context = engine.engine_context(query_ctx.clone());
|
||||
let state = context.state();
|
||||
state
|
||||
.analyzer()
|
||||
.execute_and_check(plan, state.config_options(), |_, _| {})
|
||||
.unwrap()
|
||||
.display_indent_schema()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_insert_timestamp_literals_use_query_timezone() {
|
||||
let query_ctx = Arc::new(
|
||||
@@ -884,20 +984,6 @@ mod tests {
|
||||
) AS source",
|
||||
&[1_785_902_400_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT '2026-08-06 12:00:00.001', now() \
|
||||
UNION ALL \
|
||||
SELECT '2026-08-07 12:00:00.001', now()",
|
||||
&[1_785_988_800_001_i64, 1_786_075_200_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT '2026-08-16 12:00:00.001', now() \
|
||||
UNION \
|
||||
SELECT '2026-08-17 12:00:00.001', now()",
|
||||
&[1_786_852_800_001_i64, 1_786_939_200_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT '2026-08-18 12:00:00.001', max(st) \
|
||||
@@ -926,14 +1012,7 @@ mod tests {
|
||||
&[1_786_766_400_001_i64][..],
|
||||
),
|
||||
] {
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
|
||||
let plan = engine
|
||||
.planner()
|
||||
.plan(&stmt, query_ctx.clone())
|
||||
.await
|
||||
.unwrap()
|
||||
.display_indent()
|
||||
.to_string();
|
||||
let plan = analyze_insert(&engine, sql, &query_ctx).await;
|
||||
|
||||
for expected_timestamp in expected_timestamps {
|
||||
assert!(
|
||||
@@ -954,19 +1033,19 @@ mod tests {
|
||||
let engine = create_timestamp_test_engine().await;
|
||||
let sql = "INSERT INTO timestamps (ts, st) \
|
||||
VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())";
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
|
||||
let plan = engine
|
||||
.planner()
|
||||
.plan(&stmt, query_ctx)
|
||||
.await
|
||||
.unwrap()
|
||||
.display_indent()
|
||||
.to_string();
|
||||
let plan = analyze_insert(&engine, sql, &query_ctx).await;
|
||||
|
||||
// An explicit cast reaches the analyzer as an `arrow_cast` call rather
|
||||
// than an `Expr::Cast`, which is how it stays out of the rewrite.
|
||||
assert!(
|
||||
plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"),
|
||||
"{plan}"
|
||||
);
|
||||
// 12:00:00.001 read as Shanghai local time; the source query keeps UTC.
|
||||
assert!(
|
||||
!plan.contains("TimestampMillisecond(1786104000001, None)"),
|
||||
"{plan}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1003,14 +1082,7 @@ mod tests {
|
||||
][..],
|
||||
),
|
||||
] {
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
|
||||
let plan = engine
|
||||
.planner()
|
||||
.plan(&stmt, query_ctx.clone())
|
||||
.await
|
||||
.unwrap()
|
||||
.display_indent_schema()
|
||||
.to_string();
|
||||
let plan = analyze_insert(&engine, sql, &query_ctx).await;
|
||||
|
||||
for expected in expected {
|
||||
assert!(plan.contains(expected), "{plan}");
|
||||
@@ -1019,30 +1091,37 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_insert_union_tolerates_uncoerced_untouched_column() {
|
||||
async fn test_insert_union_converts_via_assignment_cast() {
|
||||
let query_ctx = Arc::new(
|
||||
QueryContextBuilder::default()
|
||||
.timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
|
||||
.build(),
|
||||
);
|
||||
let engine = create_timestamp_test_engine().await;
|
||||
// `st` stays Timestamp vs Null across branches until TypeCoercion runs.
|
||||
let sql = "INSERT INTO timestamps (ts, st) \
|
||||
SELECT '2026-08-06 12:00:00.001', now() \
|
||||
UNION ALL \
|
||||
SELECT '2026-08-07 12:00:00.001', NULL";
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
|
||||
let plan = engine
|
||||
.planner()
|
||||
.plan(&stmt, query_ctx)
|
||||
.await
|
||||
.unwrap()
|
||||
.display_indent()
|
||||
.to_string();
|
||||
// Branches disagree, so the conversion stays as a cast on the
|
||||
// assignment instead of folding. One cast covers every branch, which is
|
||||
// why a NULL branch no longer cancels the conversion for the column and
|
||||
// why UNION's dedup keys stay on the original strings.
|
||||
for sql in [
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT '2026-08-06 12:00:00.001', now() \
|
||||
UNION ALL \
|
||||
SELECT '2026-08-07 12:00:00.001', NULL",
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT '2026-08-16 12:00:00.001', now() \
|
||||
UNION \
|
||||
SELECT '2026-08-17 12:00:00.001', now()",
|
||||
] {
|
||||
let plan = analyze_insert(&engine, sql, &query_ctx).await;
|
||||
|
||||
for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] {
|
||||
assert!(
|
||||
plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")),
|
||||
plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"),
|
||||
"{plan}"
|
||||
);
|
||||
// The branches themselves are untouched.
|
||||
assert!(
|
||||
plan.contains("Utf8(\"2026-08-07 12:00:00.001\")")
|
||||
|| plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"),
|
||||
"{plan}"
|
||||
);
|
||||
}
|
||||
@@ -1060,14 +1139,8 @@ mod tests {
|
||||
SELECT '2026-08-10 12:00:00.001', now() \
|
||||
UNION ALL \
|
||||
SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()";
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
|
||||
let plan = engine
|
||||
.planner()
|
||||
.plan(&stmt, query_ctx)
|
||||
.await
|
||||
.unwrap()
|
||||
.display_indent_schema()
|
||||
.to_string();
|
||||
let plan = analyze_insert(&engine, sql, &query_ctx).await;
|
||||
|
||||
assert!(
|
||||
!plan.contains("TimestampMillisecond(1786334400001, None)"),
|
||||
"{plan}"
|
||||
@@ -1076,6 +1149,11 @@ mod tests {
|
||||
plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"),
|
||||
"{plan}"
|
||||
);
|
||||
// TypeCoercion has already settled this union to timestamp, so the
|
||||
// assignment has nothing left to reinterpret. Retargeting the cast here
|
||||
// would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind,
|
||||
// which shifts the value instead of relabelling it.
|
||||
assert!(!plan.contains("Asia/Shanghai"), "{plan}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -4925,12 +4925,20 @@ impl PromPlanner {
|
||||
F: FnMut(&String) -> Result<DfExpr>,
|
||||
{
|
||||
let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
|
||||
// Derived labels can be unqualified even when the context still names the source table.
|
||||
let input_schema = input.schema().clone();
|
||||
let non_field_columns_iter = self
|
||||
.ctx
|
||||
.tag_columns
|
||||
.iter()
|
||||
.chain(self.ctx.time_index_column.iter())
|
||||
.map(|col| Ok(DfExpr::Column(Column::new(table_ref.clone(), col))));
|
||||
.map(|col| {
|
||||
input_schema
|
||||
.qualified_field_with_name(table_ref.as_ref(), col)
|
||||
.or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
|
||||
.map(|field| DfExpr::Column(field.into()))
|
||||
.context(DataFusionPlanningSnafu)
|
||||
});
|
||||
let tsid_iter =
|
||||
Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
|
||||
.into_iter()
|
||||
@@ -8743,6 +8751,37 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:
|
||||
assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn label_replace_aggregation_queries_plan_successfully() {
|
||||
let aggregate =
|
||||
r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
|
||||
let queries = [
|
||||
aggregate.to_string(),
|
||||
format!("{aggregate} <= 10"),
|
||||
format!("{aggregate} * 0.8"),
|
||||
format!("0.8 * {aggregate}"),
|
||||
format!("{aggregate} <= {aggregate} * 0.8"),
|
||||
];
|
||||
let state = build_query_engine_state();
|
||||
let mut failures = Vec::new();
|
||||
|
||||
for query in queries {
|
||||
let table_provider = build_test_table_provider(
|
||||
&[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
|
||||
1,
|
||||
1,
|
||||
)
|
||||
.await;
|
||||
if let Err(error) =
|
||||
PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
|
||||
{
|
||||
failures.push(format!("{query}: {error:?}"));
|
||||
}
|
||||
}
|
||||
|
||||
assert!(failures.is_empty(), "{}", failures.join("\n"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_matchers_to_expr() {
|
||||
let mut eval_stmt = EvalStmt {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user