starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-05 00:09:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: embed promql (#700)

Browse Source
This commit is contained in:
Hengfei Yang
2023-05-10 23:18:23 +08:00
committed by GitHub
parent 77cce569a5
commit c713aa9c2c
58 changed files with 4115 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+43 -31
View File
@@ -199,7 +199,7 @@ dependencies = [
"serde_urlencoded",
"smallvec",
"socket2",
"time 0.3.20",
"time 0.3.21",
"url",
]
@@ -706,6 +706,17 @@ dependencies = [
"zstd-safe 5.0.2+zstd.1.5.2",
]
[[package]]
name = "async-recursion"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0e97ce7de6cf12de5d7226c73f5ba9811622f4db3a5b91b55c53e987e5f91cba"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.15",
]
[[package]]
name = "async-stream"
version = "0.3.5"
@@ -807,7 +818,7 @@ dependencies = [
"http",
"hyper",
"ring",
"time 0.3.20",
"time 0.3.21",
"tokio",
"tower",
"tracing",
@@ -939,7 +950,7 @@ dependencies = [
"percent-encoding",
"regex",
"sha2",
"time 0.3.20",
"time 0.3.21",
"tracing",
]
@@ -1043,7 +1054,7 @@ dependencies = [
"itoa",
"num-integer",
"ryu",
"time 0.3.20",
"time 0.3.21",
]
[[package]]
@@ -1645,7 +1656,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e859cd57d0710d9e06c381b550c06e76992472a8c6d527aecd2fc673dcc231fb"
dependencies = [
"percent-encoding",
"time 0.3.20",
"time 0.3.21",
"version_check",
]
@@ -2164,9 +2175,9 @@ checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10"
[[package]]
name = "dotenv_config"
version = "0.1.5"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc8db7cd099cfe5d7b33f170b3d19054045425fe6d0d876defb7912c31d2fb67"
checksum = "72ef519d2aabc15ca3dba4729976066fa23f80187bf2b19d623d24fe1a0ec3ea"
dependencies = [
"anyhow",
"askama",
@@ -3000,9 +3011,9 @@ dependencies = [
[[package]]
name = "keccak"
version = "0.1.3"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3afef3b6eff9ce9d8ff9b3601125eec7f0c8cbac7abd14f355d053fa56c98768"
checksum = "8f6d5ed8676d904364de097082f4e7d240b571b67989ced0240f08b7f966f940"
dependencies = [
"cpufeatures",
]
@@ -3113,9 +3124,9 @@ dependencies = [
[[package]]
name = "libc"
version = "0.2.142"
version = "0.2.144"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6a987beff54b60ffa6d51982e1aa1146bc42f19bd26be28b0586f252fccf5317"
checksum = "2b00cc1c228a6782d0f076e7b232802e0c5689d41bb5df366f2a6b6621cfdfe1"
[[package]]
name = "libm"
@@ -4395,9 +4406,9 @@ dependencies = [
[[package]]
name = "quote"
version = "1.0.26"
version = "1.0.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4424af4bf778aae2051a77b60283332f386554255d722233d09fbfc7e30da2fc"
checksum = "8f4f29d145265ec1c483c7c654450edde0bfe043d3938d6972630663356d9500"
dependencies = [
"proc-macro2",
]
@@ -4894,7 +4905,7 @@ dependencies = [
"serde",
"serde_json",
"thiserror",
"time 0.3.20",
"time 0.3.21",
]
[[package]]
@@ -5022,7 +5033,7 @@ dependencies = [
"serde",
"serde_json",
"thiserror",
"time 0.3.20",
"time 0.3.21",
"url",
"uuid",
]
@@ -5035,18 +5046,18 @@ checksum = "e6b44e8fc93a14e66336d230954dda83d18b4605ccace8fe09bc7514a71ad0bc"
[[package]]
name = "serde"
version = "1.0.160"
version = "1.0.162"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bb2f3770c8bce3bcda7e149193a069a0f4365bda1fa5cd88e03bca26afc1216c"
checksum = "71b2f6e1ab5c2b98c05f0f35b236b22e8df7ead6ffbf51d7808da7f8817e7ab6"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.160"
version = "1.0.162"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "291a097c63d8497e00160b166a967a4a79c64f3facdd01cbd7502231688d77df"
checksum = "a2a0814352fd64b58489904a44ea8d90cb1a91dcb6b4f5ebabc32c8318e93cb6"
dependencies = [
"proc-macro2",
"quote",
@@ -5123,9 +5134,9 @@ dependencies = [
[[package]]
name = "sha3"
version = "0.10.7"
version = "0.10.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "54c2bb1a323307527314a36bfb73f24febb08ce2b8a554bf4ffd6f51ad15198c"
checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60"
dependencies = [
"digest",
"keccak",
@@ -5489,9 +5500,9 @@ dependencies = [
[[package]]
name = "time"
version = "0.3.20"
version = "0.3.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cd0cbfecb4d19b5ea75bb31ad904eb5b9fa13f21079c3b92017ebdf4999a5890"
checksum = "8f3403384eaacbca9923fa06940178ac13e4edb725486d70e8e15881d0c836cc"
dependencies = [
"itoa",
"serde",
@@ -5501,15 +5512,15 @@ dependencies = [
[[package]]
name = "time-core"
version = "0.1.0"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2e153e1f1acaef8acc537e68b44906d2db6436e2b35ac2c6b42640fff91f00fd"
checksum = "7300fbefb4dadc1af235a9cef3737cea692a9d97e1b9cbcd4ebdae6f8868e6fb"
[[package]]
name = "time-macros"
version = "0.2.8"
version = "0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fd80a657e71da814b8e5d60d3374fc6d35045062245d80224748ae522dd76f36"
checksum = "372950940a5f07bf38dbe211d7283c9e6d7327df53794992d293e534c733d09b"
dependencies = [
"time-core",
]
@@ -6075,7 +6086,7 @@ dependencies = [
"getset",
"rustversion",
"thiserror",
"time 0.3.20",
"time 0.3.21",
]
[[package]]
@@ -6677,6 +6688,7 @@ dependencies = [
"argon2",
"arrow",
"arrow-schema",
"async-recursion",
"async-trait",
"async_once",
"awc",
@@ -6737,7 +6749,7 @@ dependencies = [
"strum",
"sys-info",
"thiserror",
"time 0.3.20",
"time 0.3.21",
"tokio",
"tonic",
"tonic-build",
@@ -6757,9 +6769,9 @@ dependencies = [
[[package]]
name = "zip"
version = "0.6.4"
version = "0.6.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0445d0fbc924bb93539b4316c11afb121ea39296f99a3c4c9edad09e3658cdef"
checksum = "7e92305c174683d78035cbf1b70e18db6329cc0f1b9cae0a52ca90bf5bfe7125"
dependencies = [
"byteorder",
"crc32fast",
Cargo.toml
+4 -3
View File
@@ -48,6 +48,7 @@ arrow = { version = "36", features = ["simd", "ipc_compression"] }
arrow-schema = { version = "36", features = ["serde"] }
async-trait = "0.1"
async_once = "0.2"
async-recursion = "1.0"
awc = "3.1"
base64 = "0.21"
blake3 = "1.3"
@@ -117,15 +118,17 @@ reqwest = { version = "0.11", default-features = false, features = [
"rustls-tls",
] }
rs-snowflake = "0.6"
rustc-hash = "1.1.0"
rust-embed-for-web = "11.1"
rustc-hash = "1.1.0"
segment = "0.2"
sentry = { version = "0.31", features = ["log"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
simd-json = "0.7"
sled = "0.34"
snap = "1"
sqlparser = { version = "0.32", features = ["serde"] }
strum = { version = "0.24", features = ["derive"] }
sys-info = "0.9"
thiserror = "1.0"
time = "0.3"
@@ -146,8 +149,6 @@ vrl-value = { package = "value", git = "https://github.com/zinclabs/vrl", rev =
vrl-stdlib = { package = "vrl-stdlib", git = "https://github.com/zinclabs/vrl", rev = "e723f96c9305d77c853d8f95a1955ff393c8aff1" }
walkdir = "2"
zstd = "0.12"
sentry = { version = "0.31" ,features = ["log"] }
strum = { version = "0.24", features = ["derive"] }
[build-dependencies]
chrono = { version = "0.4", default-features = false, features = ["clock"] }
build.rs
+7
View File
@@ -37,10 +37,17 @@ fn main() -> Result<()> {
.type_attribute("SearchAggRequest", "#[derive(serde::Serialize)]")
.type_attribute("SearchAggResponse", "#[derive(Eq)]")
.type_attribute("SearchAggResponse", "#[derive(serde::Serialize)]")
.type_attribute("Series", "#[derive(serde::Serialize)]")
.type_attribute("Label", "#[derive(serde::Serialize)]")
.type_attribute("Sample", "#[derive(serde::Serialize)]")
.type_attribute("MetricsQueryStmt", "#[derive(serde::Serialize)]")
.type_attribute("MetricsQueryRequest", "#[derive(serde::Serialize)]")
.type_attribute("MetricsQueryResponse", "#[derive(serde::Serialize)]")
.compile(
&[
"proto/cluster/common.proto",
"proto/cluster/event.proto",
"proto/cluster/metrics.proto",
"proto/cluster/search.proto",
],
&["proto"],
deploy/build/buildspec-test.yml
+3 -1
View File
@@ -5,7 +5,9 @@ phases:
commands:
- apt-get update
- apt-get install -y bc
- apt-get install -y protobuf-compiler
- apt-get install -y zip unzip wget
- wget https://github.com/protocolbuffers/protobuf/releases/download/v21.12/protoc-21.12-linux-x86_64.zip
- unzip protoc-21.12-linux-x86_64.zip -d protoc && cp protoc/bin/protoc /usr/local/bin/ && cp -r protoc/include/google /usr/local/include/
- curl https://sh.rustup.rs -sSf | sh -s -- -y
- . "$HOME/.cargo/env"
- rustup toolchain install nightly-2023-03-28
deploy/build/buildspec.yml
+6 -6
View File
@@ -32,12 +32,12 @@ phases:
- docker push public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64
# build zincobserve-mimalloc
- docker pull public.ecr.aws/zinclabs/rust:bullseye-sccache
- docker build -t zincobserve:latest-amd64-mimalloc -f deploy/build/Dockerfile.mimalloc.amd64 .
- docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc
- docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc
- docker push public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc
- docker push public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc
# - docker pull public.ecr.aws/zinclabs/rust:bullseye-sccache
# - docker build -t zincobserve:latest-amd64-mimalloc -f deploy/build/Dockerfile.mimalloc.amd64 .
# - docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc
# - docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc
# - docker push public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc
# - docker push public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc
cache:
paths:
proto/cluster/common.proto
+17 -2
View File
@@ -1,7 +1,7 @@
syntax = "proto3";
option java_multiple_files = true;
option java_package = "org.zinc.oxide.cluster";
option java_package = "org.zinc.observe.cluster";
option java_outer_classname = "clusterProto";
package cluster;
@@ -27,4 +27,19 @@ enum StreamType {
FILELIST = 4;
}
message EmptyResponse {}
// Search request type
enum SearchType {
USER = 0; // user input search request
CLUSTER = 1; // cluster dispatch search request
CACHE_ONLY = 2; // ingester node just search local cache
}
// Job information for a request
message Job {
string session_id = 1;
string job = 2;
int32 stage = 3;
int32 partition = 4;
}
message EmptyResponse {}
proto/cluster/event.proto
+1 -3
View File
@@ -1,7 +1,7 @@
syntax = "proto3";
option java_multiple_files = true;
option java_package = "org.zinc.oxide.cluster";
option java_package = "org.zinc.observe.cluster";
option java_outer_classname = "clusterProto";
package cluster;
@@ -20,6 +20,4 @@ message FileKey {
string key = 1;
FileMeta meta = 2;
bool deleted = 3;
}
proto/cluster/metrics.proto
+57
View File
@@ -0,0 +1,57 @@
syntax = "proto3";
option java_multiple_files = true;
option java_package = "org.zinc.observe.cluster";
option java_outer_classname = "clusterProto";
package cluster;
import "cluster/common.proto";
service Metrics {
rpc Query (MetricsQueryRequest) returns (MetricsQueryResponse) {}
}
// Search request query
message MetricsQueryStmt {
string query = 1;
bool is_range_query = 2;
int64 start = 3;
int64 end = 4;
int64 step = 5;
}
// Search request
message MetricsQueryRequest {
Job job = 1;
string org_id = 2;
SearchType stype = 3;
MetricsQueryStmt query = 4;
}
// The response message containing the greetings
message MetricsQueryResponse {
Job job = 1;
int32 took = 2;
int32 file_count = 3;
int32 scan_size = 4; // unit: mb
string result_type = 5; // vector, matrix, scalar
repeated Series result = 6;
}
message Series {
repeated Label metric = 1;
repeated Sample values = 2;
optional Sample value = 3;
optional double scalar = 4;
}
message Label {
string name = 1;
string value = 2;
}
message Sample {
int64 time = 1;
double value = 2;
}
proto/cluster/search.proto
+9 -23
View File
@@ -1,31 +1,17 @@
syntax = "proto3";
option java_multiple_files = true;
option java_package = "org.zinc.oxide.cluster";
option java_package = "org.zinc.observe.cluster";
option java_outer_classname = "clusterProto";
package cluster;
import "cluster/common.proto";
service Search {
rpc Search (SearchRequest) returns (SearchResponse) {}
}
// Job information for a request
message Job {
string session_id = 1;
string job = 2;
int32 stage = 3;
int32 partition = 4;
}
// Search request type
enum SearchType {
USER = 0; // user input search request
CLUSTER = 1; // cluster dispatch search request
CACHE_ONLY = 2; // ingester node just search local cache
}
// Search request query
message SearchQuery {
string sql = 1;
@@ -43,13 +29,13 @@ message SearchQuery {
// Search request
message SearchRequest {
Job job = 1;
string org_id = 2;
string stream_type = 3;
SearchType stype = 4;
SearchQuery query = 5;
repeated string file_list = 6;
repeated SearchAggRequest aggs = 7;
Job job = 1;
string org_id = 2;
string stream_type = 3;
SearchType stype = 4;
SearchQuery query = 5;
repeated string file_list = 6;
repeated SearchAggRequest aggs = 7;
}
// The response message containing the greetings
src/handler/grpc/mod.rs
+62
View File
@@ -15,6 +15,7 @@
use uuid::Uuid;
use crate::meta;
use crate::service::promql;
pub mod auth;
pub mod request;
@@ -97,6 +98,67 @@ impl From<&meta::common::FileKey> for cluster_rpc::FileKey {
}
}
impl From<promql::MetricsQueryRequest> for cluster_rpc::MetricsQueryRequest {
fn from(req: promql::MetricsQueryRequest) -> Self {
let req_query = cluster_rpc::MetricsQueryStmt {
query: req.query.to_owned(),
is_range_query: req.is_range_query,
start: req.start,
end: req.end,
step: req.step,
};
let job = cluster_rpc::Job {
session_id: Uuid::new_v4().to_string(),
job: "".to_string(),
stage: 0,
partition: 0,
};
cluster_rpc::MetricsQueryRequest {
job: Some(job),
org_id: "".to_string(),
stype: cluster_rpc::SearchType::User.into(),
query: Some(req_query),
}
}
}
impl From<&cluster_rpc::Label> for promql::value::Label {
fn from(req: &cluster_rpc::Label) -> Self {
promql::value::Label {
name: req.name.to_owned(),
value: req.value.to_owned(),
}
}
}
impl From<&promql::value::Label> for cluster_rpc::Label {
fn from(req: &promql::value::Label) -> Self {
cluster_rpc::Label {
name: req.name.to_owned(),
value: req.value.to_owned(),
}
}
}
impl From<&cluster_rpc::Sample> for promql::value::Sample {
fn from(req: &cluster_rpc::Sample) -> Self {
promql::value::Sample {
timestamp: req.time,
value: req.value,
}
}
}
impl From<&promql::value::Sample> for cluster_rpc::Sample {
fn from(req: &promql::value::Sample) -> Self {
cluster_rpc::Sample {
time: req.timestamp,
value: req.value,
}
}
}
#[cfg(test)]
mod test {
use std::collections::HashMap;
src/handler/grpc/request/event.rs
+1 -2
View File
@@ -13,7 +13,6 @@
// limitations under the License.
use opentelemetry::global;
use std::time::Instant;
use tonic::{Request, Response, Status};
use tracing_opentelemetry::OpenTelemetrySpanExt;
@@ -34,7 +33,7 @@ impl Event for Eventer {
&self,
req: Request<FileList>,
) -> Result<Response<EmptyResponse>, Status> {
let start = Instant::now();
let start = std::time::Instant::now();
let parent_cx = global::get_text_map_propagator(|prop| {
prop.extract(&super::MetadataMap(req.metadata()))
});
src/handler/grpc/request/metrics.rs
+75
View File
@@ -0,0 +1,75 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 opentelemetry::global;
use tonic::{Request, Response, Status};
use tracing_opentelemetry::OpenTelemetrySpanExt;
use crate::handler::grpc::cluster_rpc::metrics_server::Metrics;
use crate::handler::grpc::cluster_rpc::MetricsQueryRequest;
use crate::handler::grpc::cluster_rpc::MetricsQueryResponse;
use crate::infra::errors;
use crate::infra::metrics;
use crate::meta;
use crate::service::promql::search as SearchService;
#[derive(Default)]
pub struct Querier {}
#[tonic::async_trait]
impl Metrics for Querier {
#[tracing::instrument(name = "grpc:metrics:enter", skip_all)]
async fn query(
&self,
req: Request<MetricsQueryRequest>,
) -> Result<Response<MetricsQueryResponse>, Status> {
let start = std::time::Instant::now();
let parent_cx = global::get_text_map_propagator(|prop| {
prop.extract(&super::MetadataMap(req.metadata()))
});
tracing::Span::current().set_parent(parent_cx);
let req = req.get_ref();
let org_id = req.org_id.clone();
let stream_type = meta::StreamType::Metrics.to_string();
let result = match SearchService::grpc::search(req).await {
Ok(res) => res,
Err(err) => {
// metrics
let time = start.elapsed().as_secs_f64();
metrics::GRPC_RESPONSE_TIME
.with_label_values(&["/prom/v1/query", "500", &org_id, "", &stream_type])
.observe(time);
metrics::GRPC_INCOMING_REQUESTS
.with_label_values(&["/prom/v1/query", "500", &org_id, "", &stream_type])
.inc();
return Err(match err {
errors::Error::ErrorCode(code) => Status::internal(code.to_json()),
_ => Status::internal(err.to_string()),
});
}
};
// metrics
let time = start.elapsed().as_secs_f64();
metrics::GRPC_RESPONSE_TIME
.with_label_values(&["/prom/v1/query", "200", &org_id, "", &stream_type])
.observe(time);
metrics::GRPC_INCOMING_REQUESTS
.with_label_values(&["/prom/v1/query", "200", &org_id, "", &stream_type])
.inc();
Ok(Response::new(result))
}
}
src/handler/grpc/request/mod.rs
+1
View File
@@ -15,6 +15,7 @@
use opentelemetry::propagation::Extractor;
pub mod event;
pub mod metrics;
pub mod search;
pub mod traces;
src/handler/grpc/request/search.rs
+3 -4
View File
@@ -13,7 +13,6 @@
// limitations under the License.
use opentelemetry::global;
use std::time::Instant;
use tonic::{Request, Response, Status};
use tracing_opentelemetry::OpenTelemetrySpanExt;
@@ -29,12 +28,12 @@ pub struct Searcher {}
#[tonic::async_trait]
impl Search for Searcher {
#[tracing::instrument(name = "grpc:search:enter", skip(self, req))]
#[tracing::instrument(name = "grpc:search:enter", skip_all)]
async fn search(
&self,
req: Request<SearchRequest>,
) -> Result<Response<SearchResponse>, Status> {
let start = Instant::now();
let start = std::time::Instant::now();
let parent_cx = global::get_text_map_propagator(|prop| {
prop.extract(&super::MetadataMap(req.metadata()))
});
@@ -43,7 +42,7 @@ impl Search for Searcher {
let req = req.get_ref();
let org_id = req.org_id.clone();
let stream_type = req.stream_type.as_str();
let result = match SearchService::exec::search(req).await {
let result = match SearchService::grpc::search(req).await {
Ok(res) => res,
Err(err) => {
// metrics
src/handler/http/request/prom/mod.rs
+124 -8
View File
@@ -15,10 +15,11 @@
use actix_web::{get, http, post, web, HttpRequest, HttpResponse};
use std::io::Error;
use crate::{
meta::{self, http::HttpResponse as MetaHttpResponse},
service::metrics,
};
use crate::common::time::{parse_milliseconds, parse_str_to_timestamp_micros};
use crate::infra::errors;
use crate::meta::{self, http::HttpResponse as MetaHttpResponse};
use crate::service::metrics;
use crate::service::promql;
/** prometheus remote-write endpoint for metrics */
#[utoipa::path(
@@ -105,8 +106,54 @@ pub async fn query(
query: web::Query<meta::prom::RequestQuery>,
) -> Result<HttpResponse, Error> {
let query = query.into_inner();
tracing::trace!(%org_id, ?query);
Ok(HttpResponse::NotImplemented().body("Not Implemented"))
let start = match query.time {
None => chrono::Utc::now().timestamp_micros(),
Some(v) => match parse_str_to_timestamp_micros(&v) {
Ok(v) => v,
Err(e) => {
log::error!("parse time error: {}", e);
return Ok(HttpResponse::BadRequest().json(promql::QueryResponse {
status: promql::Status::Error,
data: None,
error_type: Some("bad_data".to_string()),
error: Some(e.to_string()),
}));
}
},
};
let end = start;
let req = promql::MetricsQueryRequest {
query: query.query,
is_range_query: false,
start,
end,
step: 300_000_000,
};
match promql::search::search(&org_id, &req).await {
Ok(data) => Ok(HttpResponse::Ok().json(promql::QueryResponse {
status: promql::Status::Success,
data: Some(promql::QueryResult {
result_type: data.get_type().to_string(),
result: data,
}),
error_type: None,
error: None,
})),
Err(err) => {
let err = match err {
errors::Error::ErrorCode(code) => code.get_error_detail(),
_ => err.to_string(),
};
Ok(HttpResponse::BadRequest().json(promql::QueryResponse {
status: promql::Status::Error,
data: None,
error_type: Some("bad_data".to_string()),
error: Some(err),
}))
}
}
}
/** prometheus range queries */
@@ -168,8 +215,77 @@ pub async fn query_range(
range_query: web::Query<meta::prom::RequestRangeQuery>,
) -> Result<HttpResponse, Error> {
let range_query = range_query.into_inner();
tracing::trace!(%org_id, ?range_query);
Ok(HttpResponse::NotImplemented().body("Not Implemented"))
let start = match parse_str_to_timestamp_micros(&range_query.start) {
Ok(v) => v,
Err(e) => {
log::error!("parse time error: {}", e);
return Ok(HttpResponse::BadRequest().json(promql::QueryResponse {
status: promql::Status::Error,
data: None,
error_type: Some("bad_data".to_string()),
error: Some(e.to_string()),
}));
}
};
let end = match parse_str_to_timestamp_micros(&range_query.end) {
Ok(v) => v,
Err(e) => {
log::error!("parse time error: {}", e);
return Ok(HttpResponse::BadRequest().json(promql::QueryResponse {
status: promql::Status::Error,
data: None,
error_type: Some("bad_data".to_string()),
error: Some(e.to_string()),
}));
}
};
let step = match range_query.step {
None => 300_000_000,
Some(v) => match parse_milliseconds(&v) {
Ok(v) => (v * 1_000) as i64,
Err(e) => {
log::error!("parse time error: {}", e);
return Ok(HttpResponse::BadRequest().json(promql::QueryResponse {
status: promql::Status::Error,
data: None,
error_type: Some("bad_data".to_string()),
error: Some(e.to_string()),
}));
}
},
};
let req = promql::MetricsQueryRequest {
query: range_query.query,
is_range_query: true,
start,
end,
step,
};
match promql::search::search(&org_id, &req).await {
Ok(data) => Ok(HttpResponse::Ok().json(promql::QueryResponse {
status: promql::Status::Success,
data: Some(promql::QueryResult {
result_type: data.get_type().to_string(),
result: data,
}),
error_type: None,
error: None,
})),
Err(err) => {
let err = match err {
errors::Error::ErrorCode(code) => code.get_error_detail(),
_ => err.to_string(),
};
Ok(HttpResponse::BadRequest().json(promql::QueryResponse {
status: promql::Status::Error,
data: None,
error_type: Some("bad_data".to_string()),
error: Some(err),
}))
}
}
}
/** prometheus query metric metadata */
src/handler/http/request/search/mod.rs
+3 -4
View File
@@ -17,7 +17,6 @@ use actix_web::{get, post, web, HttpRequest, HttpResponse};
use chrono::Duration;
use std::collections::HashMap;
use std::io::Error;
use std::time::Instant;
use crate::common::base64;
use crate::common::functions;
@@ -100,7 +99,7 @@ pub async fn search(
in_req: HttpRequest,
body: actix_web::web::Bytes,
) -> Result<HttpResponse, Error> {
let start = Instant::now();
let start = std::time::Instant::now();
let org_id = org_id.into_inner();
let query = web::Query::<HashMap<String, String>>::from_query(in_req.query_string()).unwrap();
let stream_type = match get_stream_type_from_request(&query) {
@@ -237,7 +236,7 @@ pub async fn around(
path: web::Path<(String, String)>,
in_req: HttpRequest,
) -> Result<HttpResponse, Error> {
let start = Instant::now();
let start = std::time::Instant::now();
let mut uses_fn = false;
let (org_id, stream_name) = path.into_inner();
let query = web::Query::<HashMap<String, String>>::from_query(in_req.query_string()).unwrap();
@@ -465,7 +464,7 @@ pub async fn values(
path: web::Path<(String, String)>,
in_req: HttpRequest,
) -> Result<HttpResponse, Error> {
let start = Instant::now();
let start = std::time::Instant::now();
let mut uses_fn = false;
let (org_id, stream_name) = path.into_inner();
let query = web::Query::<HashMap<String, String>>::from_query(in_req.query_string()).unwrap();
src/infra/config.rs
+1 -1
View File
@@ -37,7 +37,7 @@ pub static BUILD_DATE: &str = env!("GIT_BUILD_DATE");
pub static HAS_FUNCTIONS: bool = true;
#[cfg(not(feature = "zo_functions"))]
pub static HAS_FUNCTIONS: bool = false;
pub static SEARCHING_IN_CACHE: AtomicU8 = AtomicU8::new(0);
pub static SEARCHING_IN_WAL: AtomicU8 = AtomicU8::new(0);
pub static CONFIG: Lazy<Config> = Lazy::new(init);
pub static INSTANCE_ID: Lazy<DashMap<String, String>> = Lazy::new(DashMap::new);
src/job/files/disk.rs
+1 -1
View File
@@ -151,7 +151,7 @@ async fn move_files_to_storage() -> Result<(), anyhow::Error> {
match db::file_list::local::set(&key, meta, false).await {
Ok(_) => {
loop {
let searching = config::SEARCHING_IN_CACHE.load(Ordering::Relaxed);
let searching = config::SEARCHING_IN_WAL.load(Ordering::Relaxed);
if searching == 0 {
break;
} else {
src/main.rs
+12 -6
View File
@@ -30,8 +30,11 @@ use tracing_subscriber::prelude::*;
use tracing_subscriber::Registry;
use zincobserve::handler::grpc::auth::check_auth;
use zincobserve::handler::grpc::cluster_rpc::event_server::EventServer;
use zincobserve::handler::grpc::cluster_rpc::metrics_server::MetricsServer;
use zincobserve::handler::grpc::cluster_rpc::search_server::SearchServer;
use zincobserve::handler::grpc::request::{event::Eventer, search::Searcher, traces::TraceServer};
use zincobserve::handler::grpc::request::{
event::Eventer, metrics::Querier, search::Searcher, traces::TraceServer,
};
use zincobserve::handler::http::router::{
get_basic_routes, get_other_service_routes, get_service_routes,
};
@@ -53,6 +56,7 @@ async fn main() -> Result<(), anyhow::Error> {
if cli().await? {
return Ok(());
}
let _guard;
if CONFIG.common.tracing_enabled {
let service_name = format!("zo-{}", CONFIG.common.instance_name);
@@ -110,12 +114,13 @@ async fn main() -> Result<(), anyhow::Error> {
rx.await?;
if !router::is_router() {
let gaddr: SocketAddr = format!("0.0.0.0:{}", CONFIG.grpc.port).parse()?;
let searcher = Searcher::default();
let search_svc = SearchServer::new(searcher)
let event_svc = EventServer::new(Eventer::default())
.send_compressed(CompressionEncoding::Gzip)
.accept_compressed(CompressionEncoding::Gzip);
let eventer = Eventer::default();
let event_svc = EventServer::new(eventer)
let search_svc = SearchServer::new(Searcher::default())
.send_compressed(CompressionEncoding::Gzip)
.accept_compressed(CompressionEncoding::Gzip);
let metrics_svc = MetricsServer::new(Querier::default())
.send_compressed(CompressionEncoding::Gzip)
.accept_compressed(CompressionEncoding::Gzip);
let tracer = TraceServer::default();
@@ -125,8 +130,9 @@ async fn main() -> Result<(), anyhow::Error> {
log::info!("starting gRPC server at {}", gaddr);
tonic::transport::Server::builder()
.layer(tonic::service::interceptor(check_auth))
.add_service(search_svc)
.add_service(event_svc)
.add_service(search_svc)
.add_service(metrics_svc)
.add_service(trace_svc)
.serve(gaddr)
.await
src/service/logs/mod.rs
-7
View File
@@ -59,13 +59,6 @@ fn parse_bulk_index(v: &Value) -> Option<(String, String, String)> {
None
}
// generate partition key for query
pub fn get_partition_key_query(s: &str) -> String {
let mut s = s.replace(['/', '.'], "_");
s.truncate(100);
s
}
pub fn cast_to_type(mut value: Value, delta: Vec<Field>) -> (Option<String>, Option<String>) {
let local_map = value.as_object_mut().unwrap();
//let mut error_msg = String::new();
src/service/mod.rs
+8
View File
@@ -24,6 +24,7 @@ pub mod kv;
pub mod logs;
pub mod metrics;
pub mod organization;
pub mod promql;
pub mod router;
pub mod schema;
pub mod search;
@@ -31,3 +32,10 @@ pub mod stream;
pub mod traces;
pub mod triggers;
pub mod users;
// generate partition key for query
pub fn get_partition_key_query(s: &str) -> String {
let mut s = s.replace(['/', '.'], "_");
s.truncate(100);
s
}
src/service/promql/aggregations/avg.rs
+37
View File
@@ -0,0 +1,37 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use promql_parser::parser::LabelModifier;
use crate::service::promql::value::{InstantValue, Sample, Value};
pub fn avg(timestamp: i64, param: &Option<LabelModifier>, data: &Value) -> Result<Value> {
let score_values = super::eval_arithmetic(param, data, "avg", |total, val| total + val)?;
if score_values.is_none() {
return Ok(Value::None);
}
let values = score_values
.unwrap()
.values()
.map(|v| InstantValue {
labels: v.labels.clone(),
value: Sample {
timestamp,
value: v.value / v.num as f64,
},
})
.collect::<Vec<_>>();
Ok(Value::Vector(values))
}
src/service/promql/aggregations/bottomk.rs
+23
View File
@@ -0,0 +1,23 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use promql_parser::parser::Expr as PromExpr;
use super::QueryEngine;
use crate::service::promql::value::Value;
pub async fn bottomk(ctx: &mut QueryEngine, param: Box<PromExpr>, data: &Value) -> Result<Value> {
super::eval_top(ctx, param, data, true).await
}
src/service/promql/aggregations/count.rs
+37
View File
@@ -0,0 +1,37 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use promql_parser::parser::LabelModifier;
use crate::service::promql::value::{InstantValue, Sample, Value};
pub fn count(timestamp: i64, param: &Option<LabelModifier>, data: &Value) -> Result<Value> {
let score_values = super::eval_arithmetic(param, data, "count", |_prev, _val| 0.0)?;
if score_values.is_none() {
return Ok(Value::None);
}
let values = score_values
.unwrap()
.values()
.map(|v| InstantValue {
labels: v.labels.clone(),
value: Sample {
timestamp,
value: v.num as f64,
},
})
.collect::<Vec<_>>();
Ok(Value::Vector(values))
}
src/service/promql/aggregations/max.rs
+49
View File
@@ -0,0 +1,49 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use promql_parser::parser::LabelModifier;
use crate::service::promql::value::{InstantValue, Sample, Value};
pub fn max(timestamp: i64, param: &Option<LabelModifier>, data: &Value) -> Result<Value> {
let score_values =
super::eval_arithmetic(
param,
data,
"max",
|prev, val| {
if prev >= val {
prev
} else {
val
}
},
)?;
if score_values.is_none() {
return Ok(Value::None);
}
let values = score_values
.unwrap()
.values()
.map(|v| InstantValue {
labels: v.labels.clone(),
value: Sample {
timestamp,
value: v.value,
},
})
.collect::<Vec<_>>();
Ok(Value::Vector(values))
}
src/service/promql/aggregations/min.rs
+43
View File
@@ -0,0 +1,43 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use promql_parser::parser::LabelModifier;
use crate::service::promql::value::{InstantValue, Sample, Value};
pub fn min(timestamp: i64, param: &Option<LabelModifier>, data: &Value) -> Result<Value> {
let score_values = super::eval_arithmetic(param, data, "min", |prev, val| {
if prev > 0.0 && prev <= val {
prev
} else {
val
}
})?;
if score_values.is_none() {
return Ok(Value::None);
}
let values = score_values
.unwrap()
.values()
.map(|v| InstantValue {
labels: v.labels.clone(),
value: Sample {
timestamp,
value: v.value,
},
})
.collect::<Vec<_>>();
Ok(Value::Vector(values))
}
src/service/promql/aggregations/mod.rs
+174
View File
@@ -0,0 +1,174 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::{DataFusionError, Result};
use promql_parser::parser::Expr as PromExpr;
use promql_parser::parser::LabelModifier;
use rustc_hash::FxHashMap;
use super::QueryEngine;
use crate::service::promql::value::{signature, Labels, Signature, Value};
mod avg;
mod bottomk;
mod count;
mod max;
mod min;
mod sum;
mod topk;
pub(crate) use avg::avg;
pub(crate) use bottomk::bottomk;
pub(crate) use count::count;
pub(crate) use max::max;
pub(crate) use min::min;
pub(crate) use sum::sum;
pub(crate) use topk::topk;
pub(crate) struct ArithmeticItem {
pub(crate) labels: Labels,
pub(crate) value: f64,
pub(crate) num: usize,
}
#[derive(Debug, Clone)]
pub(crate) struct TopItem {
pub(crate) index: usize,
pub(crate) value: f64,
}
pub(crate) fn eval_arithmetic(
param: &Option<LabelModifier>,
data: &Value,
f_name: &str,
f_handler: fn(total: f64, val: f64) -> f64,
) -> Result<Option<FxHashMap<Signature, ArithmeticItem>>> {
let data = match data {
Value::Vector(v) => v,
Value::None => return Ok(None),
_ => {
return Err(DataFusionError::Plan(format!(
"[{f_name}] function only accept vector values"
)))
}
};
let mut score_values = FxHashMap::default();
match param {
Some(v) => match v {
LabelModifier::Include(labels) => {
for item in data.iter() {
let mut sum_labels = Labels::default();
for label in item.labels.iter() {
if labels.contains(&label.name) {
sum_labels.push(label.clone());
}
}
let sum_hash = signature(&sum_labels);
let entry = score_values.entry(sum_hash).or_insert(ArithmeticItem {
labels: sum_labels,
value: 0.0,
num: 0,
});
entry.value = f_handler(entry.value, item.value.value);
entry.num += 1;
}
}
LabelModifier::Exclude(labels) => {
for item in data.iter() {
let mut sum_labels = Labels::default();
for label in item.labels.iter() {
if !labels.contains(&label.name) {
sum_labels.push(label.clone());
}
}
let sum_hash = signature(&sum_labels);
let entry = score_values.entry(sum_hash).or_insert(ArithmeticItem {
labels: sum_labels,
value: 0.0,
num: 0,
});
entry.value = f_handler(entry.value, item.value.value);
entry.num += 1;
}
}
},
None => {
for item in data.iter() {
let entry = score_values
.entry(Signature::default())
.or_insert(ArithmeticItem {
labels: Labels::default(),
value: 0.0,
num: 0,
});
entry.value = f_handler(entry.value, item.value.value);
entry.num += 1;
}
}
}
Ok(Some(score_values))
}
pub async fn eval_top(
ctx: &mut QueryEngine,
param: Box<PromExpr>,
data: &Value,
is_bottom: bool,
) -> Result<Value> {
let fn_name = if is_bottom { "bottomk" } else { "topk" };
let param = ctx.exec_expr(&param).await?;
let n = match param {
Value::Float(v) => v as usize,
_ => {
return Err(DataFusionError::Plan(format!(
"[{fn_name}] param must be NumberLiteral"
)))
}
};
let data = match data {
Value::Vector(v) => v,
Value::None => return Ok(Value::None),
_ => {
return Err(DataFusionError::Plan(format!(
"[{fn_name}] function only accept vector values"
)))
}
};
let mut score_values = Vec::new();
for (i, item) in data.iter().enumerate() {
if item.value.value.is_nan() {
continue;
}
score_values.push(TopItem {
index: i,
value: item.value.value,
});
}
if is_bottom {
score_values.sort_by(|a, b| a.value.partial_cmp(&b.value).unwrap());
} else {
score_values.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap());
}
let values = score_values
.iter()
.take(n)
.map(|v| data[v.index].clone())
.collect();
Ok(Value::Vector(values))
}
src/service/promql/aggregations/sum.rs
+37
View File
@@ -0,0 +1,37 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use promql_parser::parser::LabelModifier;
use crate::service::promql::value::{InstantValue, Sample, Value};
pub fn sum(timestamp: i64, param: &Option<LabelModifier>, data: &Value) -> Result<Value> {
let score_values = super::eval_arithmetic(param, data, "sum", |total, val| total + val)?;
if score_values.is_none() {
return Ok(Value::None);
}
let values = score_values
.unwrap()
.values()
.map(|v| InstantValue {
labels: v.labels.clone(),
value: Sample {
timestamp,
value: v.value,
},
})
.collect::<Vec<_>>();
Ok(Value::Vector(values))
}
src/service/promql/aggregations/topk.rs
+23
View File
@@ -0,0 +1,23 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use promql_parser::parser::Expr as PromExpr;
use super::QueryEngine;
use crate::service::promql::value::Value;
pub async fn topk(ctx: &mut QueryEngine, param: Box<PromExpr>, data: &Value) -> Result<Value> {
super::eval_top(ctx, param, data, false).await
}
src/service/promql/engine.rs
+805
View File
@@ -0,0 +1,805 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 async_recursion::async_recursion;
use async_trait::async_trait;
use datafusion::{
arrow::array::{Float64Array, Int64Array, StringArray},
error::{DataFusionError, Result},
prelude::{col, lit, SessionContext},
};
use promql_parser::{
label::MatchOp,
parser::{
token, AggregateExpr, Call, EvalStmt, Expr as PromExpr, Function, FunctionArgs,
LabelModifier, MatrixSelector, NumberLiteral, ParenExpr, TokenType, UnaryExpr,
VectorSelector,
},
};
use rustc_hash::FxHashMap;
use std::{
str::FromStr,
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use super::{aggregations, functions, value::*};
#[async_trait]
pub trait TableProvider: Sync + Send + 'static {
async fn create_context(
&self,
stream_name: &str,
time_range: (i64, i64),
filters: &[(String, String)],
) -> Result<SessionContext>;
}
pub struct QueryEngine {
table_provider: Box<dyn TableProvider>,
/// The time boundaries for the evaluation. If start equals end an instant
/// is evaluated.
start: i64,
end: i64,
/// Time between two evaluated instants for the range [start:end].
interval: i64,
/// Default look back from sample search.
lookback_delta: i64,
/// The index of the current time window. Used when evaluating a [range query].
///
/// [range query]: https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#range-queries
time_window_idx: i64,
/// key — metric name; value — time series data
metrics_cache: FxHashMap<String, Value>,
}
impl QueryEngine {
pub fn new<P>(provider: P) -> Self
where
P: TableProvider,
{
let now = micros_since_epoch(SystemTime::now());
let five_min = micros(Duration::from_secs(300));
Self {
table_provider: Box::new(provider),
start: now,
end: now,
interval: five_min,
lookback_delta: five_min,
time_window_idx: 0,
metrics_cache: FxHashMap::default(),
}
}
pub async fn exec(&mut self, stmt: EvalStmt) -> Result<Value> {
self.start = micros_since_epoch(stmt.start);
self.end = micros_since_epoch(stmt.end);
if stmt.interval > Duration::ZERO {
self.interval = micros(stmt.interval);
}
if stmt.lookback_delta > Duration::ZERO {
self.lookback_delta = micros(stmt.lookback_delta);
}
if self.start == self.end {
// Instant query
let mut value = self.exec_expr(&stmt.expr).await?;
if let Value::Float(v) = value {
value = Value::Sample(Sample {
timestamp: self.end,
value: v,
});
}
value.sort();
return Ok(value);
}
// Range query
// See https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#range-queries
let mut instant_vectors = Vec::new();
let nr_steps = ((self.end - self.start) / self.interval) + 1;
for i in 0..nr_steps {
self.time_window_idx = i;
match self.exec_expr(&stmt.expr).await? {
Value::Instant(v) => instant_vectors.push(RangeValue {
labels: v.labels.to_owned(),
time_range: None,
values: vec![v.value],
}),
Value::Vector(vs) => instant_vectors.extend(vs.into_iter().map(|v| RangeValue {
labels: v.labels.to_owned(),
time_range: None,
values: vec![v.value],
})),
Value::Range(v) => instant_vectors.push(v),
Value::Matrix(v) => instant_vectors.extend(v),
Value::Sample(v) => instant_vectors.push(RangeValue {
labels: Labels::default(),
time_range: None,
values: vec![v],
}),
Value::Float(v) => instant_vectors.push(RangeValue {
labels: Labels::default(),
time_range: None,
values: vec![Sample {
timestamp: self.start + (self.interval * self.time_window_idx),
value: v,
}],
}),
Value::None => continue,
};
}
// empty result quick return
if instant_vectors.is_empty() {
return Ok(Value::None);
}
// merge data
let mut merged_data = FxHashMap::default();
let mut merged_metrics = FxHashMap::default();
for value in instant_vectors {
merged_data
.entry(signature(&value.labels))
.or_insert_with(Vec::new)
.extend(value.values);
merged_metrics.insert(signature(&value.labels), value.labels);
}
let merged_data = merged_data
.into_iter()
.map(|(sig, values)| RangeValue {
labels: merged_metrics.get(&sig).unwrap().to_owned(),
time_range: None,
values,
})
.collect::<Vec<_>>();
// sort data
let mut value = Value::Matrix(merged_data);
value.sort();
Ok(value)
}
#[async_recursion]
pub async fn exec_expr(&mut self, prom_expr: &PromExpr) -> Result<Value> {
Ok(match &prom_expr {
PromExpr::Aggregate(AggregateExpr {
op,
expr,
param,
modifier,
}) => self.aggregate_exprs(op, expr, param, modifier).await?,
PromExpr::Unary(UnaryExpr { expr }) => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Unary: {:?}",
expr
)));
}
PromExpr::Binary(expr) => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Paren: {:?}",
expr
)));
}
PromExpr::Paren(ParenExpr { expr }) => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Paren: {:?}",
expr
)));
}
PromExpr::Subquery(expr) => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Subquery: {:?}",
expr
)));
}
PromExpr::NumberLiteral(NumberLiteral { val }) => Value::Float(*val),
PromExpr::StringLiteral(expr) => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported StringLiteral: {:?}",
expr
)));
}
PromExpr::VectorSelector(v) => {
let data = self.eval_vector_selector(v).await?;
if data.is_empty() {
Value::None
} else {
Value::Vector(data)
}
}
PromExpr::MatrixSelector(MatrixSelector {
vector_selector,
range,
}) => {
let data = self.eval_matrix_selector(vector_selector, *range).await?;
if data.is_empty() {
Value::None
} else {
Value::Matrix(data)
}
}
PromExpr::Call(Call { func, args }) => self.call_expr(func, args).await?,
PromExpr::Extension(expr) => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Extension: {:?}",
expr
)));
}
})
}
/// MatrixSelector is a special case of VectorSelector that returns a matrix of samples.
async fn eval_vector_selector(
&mut self,
selector: &VectorSelector,
) -> Result<Vec<InstantValue>> {
let metrics_name = selector.name.as_ref().unwrap();
if !self.metrics_cache.contains_key(metrics_name) {
self.selector_load_data(selector, None).await?;
}
let metrics_cache = match self.metrics_cache.get(metrics_name) {
Some(v) => match v.get_ref_matrix_values() {
Some(v) => v,
None => return Ok(vec![]),
},
None => return Ok(vec![]),
};
let mut values = vec![];
for metric in metrics_cache {
let mut value = match metric.values.last() {
Some(v) => *v,
None => continue, // have no sample
};
if value.timestamp != self.end && metric.values.len() > 1 {
let mut metric = metric.clone();
metric.time_range = Some((metric.values.first().unwrap().timestamp, self.end));
if let Some(extra) = metric.extrapolate() {
value = extra.1;
}
}
values.push(
// See https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#instant-queries
InstantValue {
labels: metric.labels.clone(),
value,
},
);
}
Ok(values)
}
/// MatrixSelector is a special case of VectorSelector that returns a matrix of samples.
async fn eval_matrix_selector(
&mut self,
selector: &VectorSelector,
range: Duration,
) -> Result<Vec<RangeValue>> {
let metrics_name = selector.name.as_ref().unwrap();
if !self.metrics_cache.contains_key(metrics_name) {
self.selector_load_data(selector, Some(range)).await?;
}
let metrics_cache = match self.metrics_cache.get(metrics_name) {
Some(v) => match v.get_ref_matrix_values() {
Some(v) => v,
None => return Ok(vec![]),
},
None => return Ok(vec![]),
};
let end = self.start + (self.interval * self.time_window_idx); // 15s
let start = end - micros(range); // 5m
let mut values = Vec::with_capacity(metrics_cache.len());
for metric in metrics_cache {
let metric_data = metric
.values
.iter()
.filter(|v| v.timestamp > start && v.timestamp <= end)
.cloned()
.collect::<Vec<_>>();
values.push(RangeValue {
labels: metric.labels.clone(),
time_range: Some((start, end)),
values: metric_data,
});
}
Ok(values)
}
async fn selector_load_data(
&mut self,
selector: &VectorSelector,
range: Option<Duration>,
) -> Result<()> {
let start = {
// https://promlabs.com/blog/2020/07/02/selecting-data-in-promql/#lookback-delta
let lookback_delta = range.map_or(self.lookback_delta, micros);
self.start + (self.interval * self.time_window_idx) - lookback_delta
};
let end = self.end; // 30 minutes + 5m = 35m
// 1. Group by metrics (sets of label name-value pairs)
let table_name = selector.name.as_ref().unwrap();
let filters: Vec<(String, String)> = selector
.matchers
.matchers
.iter()
.filter(|mat| mat.op == MatchOp::Equal)
.map(|mat| (mat.name.clone(), mat.value.clone()))
.collect();
let ctx = self
.table_provider
.create_context(table_name, (start, end), &filters)
.await?;
let table = match ctx.table(table_name).await {
Ok(v) => v,
Err(_) => {
self.metrics_cache
.insert(table_name.to_string(), Value::None);
return Ok(()); // table not found
}
};
let mut df_group = table.clone().filter(
col(FIELD_TIME)
.gt(lit(start))
.and(col(FIELD_TIME).lt_eq(lit(end))),
)?;
let regexp_match_udf =
crate::service::search::datafusion::regexp_udf::REGEX_MATCH_UDF.clone();
let regexp_not_match_udf =
crate::service::search::datafusion::regexp_udf::REGEX_NOT_MATCH_UDF.clone();
for mat in selector.matchers.matchers.iter() {
match &mat.op {
MatchOp::Equal => {
df_group = df_group.filter(col(mat.name.clone()).eq(lit(mat.value.clone())))?
}
MatchOp::NotEqual => {
df_group =
df_group.filter(col(mat.name.clone()).not_eq(lit(mat.value.clone())))?
}
MatchOp::Re(_re) => {
df_group = df_group.filter(
regexp_match_udf.call(vec![col(mat.name.clone()), lit(mat.value.clone())]),
)?
}
MatchOp::NotRe(_re) => {
df_group = df_group.filter(
regexp_not_match_udf
.call(vec![col(mat.name.clone()), lit(mat.value.clone())]),
)?
}
}
}
let batches = df_group
.sort(vec![col(FIELD_TIME).sort(true, true)])?
.collect()
.await?;
let mut metrics: FxHashMap<String, RangeValue> = FxHashMap::default();
for batch in &batches {
let hash_values = batch
.column_by_name(FIELD_HASH)
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
let time_values = batch
.column_by_name(FIELD_TIME)
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let value_values = batch
.column_by_name(FIELD_VALUE)
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
for i in 0..batch.num_rows() {
let hash = hash_values.value(i).to_string();
let entry = metrics.entry(hash).or_insert_with(|| {
let mut labels = Vec::with_capacity(batch.num_columns());
for (k, v) in batch.schema().fields().iter().zip(batch.columns()) {
let name = k.name();
if name == FIELD_HASH || name == FIELD_TIME || name == FIELD_VALUE {
continue;
}
let value = v.as_any().downcast_ref::<StringArray>().unwrap();
labels.push(Arc::new(Label {
name: name.to_string(),
value: value.value(i).to_string(),
}));
}
labels.sort_by(|a, b| a.name.cmp(&b.name));
RangeValue {
labels,
time_range: None,
values: Vec::with_capacity(20),
}
});
entry.values.push(Sample {
timestamp: time_values.value(i),
value: value_values.value(i),
});
}
}
// We don't need the primary key (FIELD_HASH) any more
let mut metric_values = metrics.into_values().collect::<Vec<_>>();
// TODO fix Extrapolation, it needs reduce by distance
// https://ihac.xyz/2018/12/11/Prometheus-Extrapolation%E5%8E%9F%E7%90%86%E8%A7%A3%E6%9E%90/
// Fix data about app restart
for series in metric_values.iter_mut() {
if let Some(metric_type) = labels_value(&series.labels, FIELD_TYPE) {
if metric_type != TYPE_COUNTER {
continue;
}
}
let mut delta: f64 = 0.0;
let mut last_value = 0.0;
for sample in series.values.iter_mut() {
if last_value > sample.value {
delta += last_value;
}
last_value = sample.value;
if delta > 0.0 {
sample.value += delta;
}
}
}
// cache data
let values = if metric_values.is_empty() {
Value::None
} else {
Value::Matrix(metric_values)
};
self.metrics_cache.insert(table_name.to_string(), values);
Ok(())
}
async fn aggregate_exprs(
&mut self,
op: &TokenType,
expr: &PromExpr,
param: &Option<Box<PromExpr>>,
modifier: &Option<LabelModifier>,
) -> Result<Value> {
let sample_time = self.start + (self.interval * self.time_window_idx);
let input = self.exec_expr(expr).await?;
Ok(match op.id() {
token::T_SUM => aggregations::sum(sample_time, modifier, &input)?,
token::T_AVG => aggregations::avg(sample_time, modifier, &input)?,
token::T_COUNT => aggregations::count(sample_time, modifier, &input)?,
token::T_MIN => aggregations::min(sample_time, modifier, &input)?,
token::T_MAX => aggregations::max(sample_time, modifier, &input)?,
token::T_GROUP => Value::None,
token::T_STDDEV => Value::None,
token::T_STDVAR => Value::None,
token::T_TOPK => aggregations::topk(self, param.clone().unwrap(), &input).await?,
token::T_BOTTOMK => aggregations::bottomk(self, param.clone().unwrap(), &input).await?,
token::T_COUNT_VALUES => Value::None,
token::T_QUANTILE => Value::None,
_ => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Aggregate: {:?}",
op
)));
}
})
}
async fn call_expr(&mut self, func: &Function, args: &FunctionArgs) -> Result<Value> {
use crate::service::promql::functions::Func;
let func_name = Func::from_str(func.name).map_err(|_| {
DataFusionError::NotImplemented(format!("Unsupported function: {}", func.name))
})?;
let last_arg = args
.last()
.expect("BUG: promql-parser should have validated function arguments");
let input = self.exec_expr(&last_arg).await?;
Ok(match func_name {
Func::Abs => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Absent => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::AbsentOverTime => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::AvgOverTime => functions::avg_over_time(&input)?,
Func::Ceil => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Changes => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Clamp => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::ClampMax => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::ClampMin => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::CountOverTime => functions::count_over_time(&input)?,
Func::DayOfMonth => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::DayOfWeek => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::DayOfYear => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::DaysInMonth => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Delta => functions::delta(&input)?,
Func::Deriv => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Exp => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Floor => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::HistogramCount => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::HistogramFraction => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::HistogramQuantile => {
let args = &args.args;
if args.len() != 2 {
return Err(DataFusionError::Plan(format!(
"{}: expected 2 arguments, got {}",
func.name,
args.len()
)));
}
let phi = {
match *args[0] {
PromExpr::NumberLiteral(ref num) => num.val,
_ => {
return Err(DataFusionError::Plan(format!(
"{}: the first argument must be a number",
func.name
)))
}
}
};
let sample_time = self.start + (self.interval * self.time_window_idx);
functions::histogram_quantile(sample_time, phi, input)?
}
Func::HistogramSum => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::HoltWinters => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Hour => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Idelta => functions::idelta(&input)?,
Func::Increase => functions::increase(&input)?,
Func::Irate => functions::irate(&input)?,
Func::LabelJoin => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::LabelReplace => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Ln => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Log10 => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Log2 => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::MaxOverTime => functions::max_over_time(&input)?,
Func::MinOverTime => functions::min_over_time(&input)?,
Func::Minute => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Month => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::PredictLinear => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::QuantileOverTime => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Rate => functions::rate(&input)?,
Func::Resets => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Round => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Scalar => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Sgn => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Sort => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::SortDesc => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::SumOverTime => functions::sum_over_time(&input)?,
Func::Time => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Timestamp => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Vector => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
Func::Year => {
return Err(DataFusionError::NotImplemented(format!(
"Unsupported Function: {:?}",
func_name
)));
}
})
}
}
/// Converts `t` to the number of microseconds elapsed since the beginning of the Unix epoch.
fn micros_since_epoch(t: SystemTime) -> i64 {
micros(
t.duration_since(UNIX_EPOCH)
.expect("BUG: {t} is earlier than Unix epoch"),
)
}
fn micros(t: Duration) -> i64 {
t.as_micros()
.try_into()
.expect("BUG: time value is too large to fit in i64")
}
src/service/promql/functions/avg_over_time.rs
+28
View File
@@ -0,0 +1,28 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn avg_over_time(data: &Value) -> Result<Value> {
super::eval_idelta(data, "avg_over_time", exec)
}
fn exec(data: &RangeValue) -> Option<f64> {
if data.values.is_empty() {
return None;
}
Some(data.values.iter().map(|s| s.value).sum::<f64>() / data.values.len() as f64)
}
src/service/promql/functions/count_over_time.rs
+28
View File
@@ -0,0 +1,28 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn count_over_time(data: &Value) -> Result<Value> {
super::eval_idelta(data, "count_over_time", exec)
}
fn exec(data: &RangeValue) -> Option<f64> {
if data.values.is_empty() {
return None;
}
Some(data.values.len() as f64)
}
src/service/promql/functions/delta.rs
+27
View File
@@ -0,0 +1,27 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn delta(data: &Value) -> Result<Value> {
super::eval_idelta(data, "delta", exec)
}
fn exec(range: &RangeValue) -> Option<f64> {
range
.extrapolate()
.map(|(first, last)| last.value - first.value)
}
src/service/promql/functions/histogram.rs
+302
View File
@@ -0,0 +1,302 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::{DataFusionError, Result};
use rustc_hash::FxHashMap;
use crate::service::promql::value::{self, InstantValue, Labels, Signature, Value};
// https://github.com/prometheus/prometheus/blob/cf1bea344a3c390a90c35ea8764c4a468b345d5e/promql/quantile.go#L33
#[derive(Debug, Clone, Copy, PartialEq)]
struct Bucket {
upper_bound: f64,
count: f64,
}
// https://github.com/prometheus/prometheus/blob/cf1bea344a3c390a90c35ea8764c4a468b345d5e/promql/quantile.go#L45
#[derive(Debug)]
struct MetricWithBuckets {
labels: Labels,
buckets: Vec<Bucket>,
}
/// TODO: support native histograms; see [`histogramQuantile`]
///
/// [`histogramQuantile`]: https://github.com/prometheus/prometheus/blob/f7c6130ff27a2a12412c02cce223f7a8abc59e49/promql/quantile.go#L146
pub(crate) fn histogram_quantile(sample_time: i64, phi: f64, data: Value) -> Result<Value> {
let in_vec = match data {
Value::Vector(v) => v,
Value::None => return Ok(Value::None),
_ => {
return Err(DataFusionError::Plan(
"histogram_quantile: vector argument expected".to_owned(),
))
}
};
let mut metrics_with_buckets: FxHashMap<Signature, MetricWithBuckets> = FxHashMap::default();
for InstantValue {
mut labels,
value: sample,
} in in_vec
{
// [https://prometheus.io/docs/prometheus/latest/querying/functions/#histogram_quantile]:
//
// The conventional float samples in `in_vec` are considered the counts
// of observations in each bucket of one or more conventional
// histograms. Each float sample must have a label `le` where the label
// value denotes the inclusive upper bound of the bucket. *Float samples
// without such a label are silently ignored.*
let upper_bound: f64 = match labels
.iter()
.find(|v| v.name == "le")
.map(|s| s.value.parse())
{
Some(Ok(u)) => u,
None | Some(Err(_)) => continue,
};
let sig = value::signature_without_labels(
&labels,
&[value::FIELD_HASH, value::FIELD_NAME, value::FIELD_BUCKET],
);
let entry = metrics_with_buckets.entry(sig).or_insert_with(|| {
labels.retain(|l| {
l.name != value::FIELD_HASH
&& l.name != value::FIELD_NAME
&& l.name != value::FIELD_BUCKET
});
MetricWithBuckets {
labels,
buckets: Vec::new(),
}
});
entry.buckets.push(Bucket {
upper_bound,
count: sample.value,
});
}
let values = metrics_with_buckets
.into_values()
.map(|mb| value::InstantValue {
labels: mb.labels,
value: value::Sample {
timestamp: sample_time,
value: bucket_quantile(phi, mb.buckets),
},
})
.collect();
Ok(Value::Vector(values))
}
// cf. https://github.com/prometheus/prometheus/blob/cf1bea344a3c390a90c35ea8764c4a468b345d5e/promql/quantile.go#L76
fn bucket_quantile(phi: f64, mut buckets: Vec<Bucket>) -> f64 {
if phi.is_nan() {
return f64::NAN;
}
if phi < 0.0 {
return f64::NEG_INFINITY;
}
if phi > 1.0 {
return f64::INFINITY;
}
buckets.sort_by(|a, b| a.upper_bound.partial_cmp(&b.upper_bound).unwrap());
// The caller of `bucket_quantile` guarantees that `buckets` is non-empty.
let highest_bucket = buckets[buckets.len() - 1];
if !(highest_bucket.upper_bound.is_infinite() && highest_bucket.upper_bound.is_sign_positive())
{
return f64::NAN;
}
let mut buckets = coalesce_buckets(buckets);
ensure_monotonic(&mut buckets);
let buckets = buckets;
if buckets.len() < 2 {
return f64::NAN;
}
let observations = highest_bucket.count;
if observations == 0.0 {
return f64::NAN;
}
let mut rank = phi * observations;
let b = match buckets[..buckets.len() - 1]
.binary_search_by(|b| b.count.partial_cmp(&rank).unwrap())
{
Ok(b) => b,
Err(b) => b,
};
if b == buckets.len() - 1 {
return buckets[buckets.len() - 2].upper_bound;
}
if b == 0 && buckets[0].upper_bound <= 0.0 {
return buckets[0].upper_bound;
}
let bucket_end = buckets[b].upper_bound;
let mut count = buckets[b].count;
let bucket_start = if b == 0 {
0.0
} else {
count -= buckets[b - 1].count;
rank -= buckets[b - 1].count;
buckets[b - 1].upper_bound
};
bucket_start + (bucket_end - bucket_start) * rank / count
}
/// `coalesce_buckets` merges buckets with the same upper bound.
/// The input buckets must be sorted.
fn coalesce_buckets(buckets: Vec<Bucket>) -> Vec<Bucket> {
let mut st = None;
let mut buckets = buckets
.into_iter()
.filter_map(|b| match st {
None => {
st = Some(b);
None
}
Some(last) => {
if b.upper_bound == last.upper_bound {
st = Some(Bucket {
count: last.count + b.count,
..last
});
None
} else {
st = Some(b);
Some(last)
}
}
})
.collect::<Vec<_>>();
if let Some(last) = st {
buckets.push(last);
}
buckets
}
// For the rationale behind this function, see
// https://github.com/prometheus/prometheus/blob/0bf707e288eaa8694105e53c81a102017529793d/promql/quantile.go#L314-L347
fn ensure_monotonic(buckets: &mut [Bucket]) {
let mut max = buckets[0].count;
for bucket in &mut buckets[1..] {
if bucket.count > max {
max = bucket.count;
} else if bucket.count < max {
bucket.count = max;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use expect_test::expect;
#[test]
fn test_coalesce_buckets() {
let buckets = vec![
Bucket {
upper_bound: 1.0,
count: 2.0,
},
Bucket {
upper_bound: 1.0,
count: 3.0,
},
Bucket {
upper_bound: 2.0,
count: 4.0,
},
Bucket {
upper_bound: 3.0,
count: 1.0,
},
Bucket {
upper_bound: 3.0,
count: 1.0,
},
];
expect![[r#"
[
Bucket {
upper_bound: 1.0,
count: 5.0,
},
Bucket {
upper_bound: 2.0,
count: 4.0,
},
Bucket {
upper_bound: 3.0,
count: 2.0,
},
]
"#]]
.assert_debug_eq(&coalesce_buckets(buckets));
}
#[test]
fn test_ensure_monotonic() {
let mut buckets = vec![
Bucket {
upper_bound: 1.0,
count: 2.0,
},
Bucket {
upper_bound: 2.0,
count: 1.0,
},
Bucket {
upper_bound: 3.0,
count: 4.0,
},
Bucket {
upper_bound: 4.0,
count: 3.0,
},
Bucket {
upper_bound: 5.0,
count: 5.0,
},
];
ensure_monotonic(&mut buckets);
expect![[r#"
[
Bucket {
upper_bound: 1.0,
count: 2.0,
},
Bucket {
upper_bound: 2.0,
count: 2.0,
},
Bucket {
upper_bound: 3.0,
count: 4.0,
},
Bucket {
upper_bound: 4.0,
count: 4.0,
},
Bucket {
upper_bound: 5.0,
count: 5.0,
},
]
"#]]
.assert_debug_eq(&buckets);
}
}
src/service/promql/functions/idelta.rs
+32
View File
@@ -0,0 +1,32 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn idelta(data: &Value) -> Result<Value> {
super::eval_idelta(data, "idelta", exec)
}
fn exec(data: &RangeValue) -> Option<f64> {
if data.values.is_empty() {
return None;
} else if data.values.len() == 1 {
return Some(0.0);
}
let (last, data) = data.values.split_last().unwrap();
let previous = data.last().unwrap();
Some(last.value - previous.value)
}
src/service/promql/functions/increase.rs
+27
View File
@@ -0,0 +1,27 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn increase(data: &Value) -> Result<Value> {
super::eval_idelta(data, "increase", exec)
}
fn exec(range: &RangeValue) -> Option<f64> {
range
.extrapolate()
.map(|(first, last)| last.value - first.value)
}
src/service/promql/functions/irate.rs
+36
View File
@@ -0,0 +1,36 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn irate(data: &Value) -> Result<Value> {
super::eval_idelta(data, "irate", exec)
}
fn exec(data: &RangeValue) -> Option<f64> {
if data.values.is_empty() {
return None;
} else if data.values.len() == 1 {
return Some(0.0);
}
let (last, data) = data.values.split_last().unwrap();
let previous = data.last().unwrap();
let dt_seconds = ((last.timestamp - previous.timestamp) / 1_000_000) as f64;
if dt_seconds == 0.0 {
return Some(0.0);
}
Some((last.value - previous.value) / dt_seconds)
}
src/service/promql/functions/max_over_time.rs
+34
View File
@@ -0,0 +1,34 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn max_over_time(data: &Value) -> Result<Value> {
super::eval_idelta(data, "max_over_time", exec)
}
fn exec(data: &RangeValue) -> Option<f64> {
if data.values.is_empty() {
return None;
}
Some(
data.values
.iter()
.map(|s| s.value)
.max_by(|a, b| a.partial_cmp(b).unwrap())
.unwrap(),
)
}
src/service/promql/functions/min_over_time.rs
+34
View File
@@ -0,0 +1,34 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn min_over_time(data: &Value) -> Result<Value> {
super::eval_idelta(data, "min_over_time", exec)
}
fn exec(data: &RangeValue) -> Option<f64> {
if data.values.is_empty() {
return None;
}
Some(
data.values
.iter()
.map(|s| s.value)
.min_by(|a, b| a.partial_cmp(b).unwrap())
.unwrap(),
)
}
src/service/promql/functions/mod.rs
+133
View File
@@ -0,0 +1,133 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::{DataFusionError, Result};
use crate::service::promql::value::{InstantValue, RangeValue, Sample, Value};
mod avg_over_time;
mod count_over_time;
mod delta;
mod histogram;
mod idelta;
mod increase;
mod irate;
mod max_over_time;
mod min_over_time;
mod rate;
mod sum_over_time;
pub(crate) use avg_over_time::avg_over_time;
pub(crate) use count_over_time::count_over_time;
pub(crate) use delta::delta;
pub(crate) use histogram::histogram_quantile;
pub(crate) use idelta::idelta;
pub(crate) use increase::increase;
pub(crate) use irate::irate;
pub(crate) use max_over_time::max_over_time;
pub(crate) use min_over_time::min_over_time;
pub(crate) use rate::rate;
pub(crate) use sum_over_time::sum_over_time;
use strum::EnumString;
/// Reference: https://prometheus.io/docs/prometheus/latest/querying/functions/
#[derive(Debug, Clone, Copy, PartialEq, EnumString)]
#[strum(serialize_all = "snake_case")]
pub(crate) enum Func {
Abs,
Absent,
AbsentOverTime,
AvgOverTime,
Ceil,
Changes,
Clamp,
ClampMax,
ClampMin,
CountOverTime,
DayOfMonth,
DayOfWeek,
DayOfYear,
DaysInMonth,
Delta,
Deriv,
Exp,
Floor,
HistogramCount,
HistogramFraction,
HistogramQuantile,
HistogramSum,
HoltWinters,
Hour,
Idelta,
Increase,
Irate,
LabelJoin,
LabelReplace,
Ln,
Log10,
Log2,
MaxOverTime,
MinOverTime,
Minute,
Month,
PredictLinear,
QuantileOverTime,
Rate,
Resets,
Round,
Scalar,
Sgn,
Sort,
SortDesc,
SumOverTime,
Time,
Timestamp,
Vector,
Year,
}
pub(crate) fn eval_idelta(
data: &Value,
fn_name: &str,
fn_handler: fn(&RangeValue) -> Option<f64>,
) -> Result<Value> {
let data = match data {
Value::Matrix(v) => v,
Value::None => return Ok(Value::None),
_ => {
return Err(DataFusionError::Plan(format!(
"{fn_name}: matrix argument expected"
)))
}
};
let mut rate_values = Vec::with_capacity(data.len());
for metric in data.iter() {
let mut labels = metric.labels.clone();
labels.retain(|l| l.name != super::value::FIELD_NAME);
let value = fn_handler(metric);
if value.is_none() {
continue;
}
rate_values.push(InstantValue {
labels,
value: Sample {
timestamp: metric.time_range.unwrap().1,
value: value.unwrap(),
},
});
}
Ok(Value::Vector(rate_values))
}
src/service/promql/functions/rate.rs
+32
View File
@@ -0,0 +1,32 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn rate(data: &Value) -> Result<Value> {
super::eval_idelta(data, "rate", exec)
}
fn exec(range: &RangeValue) -> Option<f64> {
range.extrapolate().map(|(first, last)| {
let dt_seconds = ((last.timestamp - first.timestamp) / 1_000_000) as f64;
if dt_seconds == 0.0 {
0.0
} else {
(last.value - first.value) / dt_seconds
}
})
}
src/service/promql/functions/sum_over_time.rs
+28
View File
@@ -0,0 +1,28 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::error::Result;
use crate::service::promql::value::{RangeValue, Value};
pub(crate) fn sum_over_time(data: &Value) -> Result<Value> {
super::eval_idelta(data, "sum_over_time", exec)
}
fn exec(data: &RangeValue) -> Option<f64> {
if data.values.is_empty() {
return None;
}
Some(data.values.iter().map(|s| s.value).sum())
}
src/service/promql/mod.rs
+59
View File
@@ -0,0 +1,59 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 serde::{Deserialize, Serialize};
use utoipa::ToSchema;
mod aggregations;
mod engine;
mod functions;
pub mod search;
pub mod value;
pub use engine::QueryEngine;
#[derive(Clone, Debug, Serialize, Deserialize, ToSchema)]
pub struct MetricsQueryRequest {
pub query: String,
pub is_range_query: bool,
pub start: i64,
pub end: i64,
pub step: i64,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Status {
Success,
Error,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct QueryResult {
pub result_type: String, // vector, matrix, scalar, string
pub result: value::Value,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct QueryResponse {
pub status: Status,
#[serde(skip_serializing_if = "Option::is_none")]
pub data: Option<QueryResult>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
src/service/promql/search/grpc/mod.rs
+432
View File
@@ -0,0 +1,432 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 datafusion::arrow::datatypes::Schema;
use datafusion::datasource::file_format::file_type::{FileType, GetExt};
use datafusion::datasource::file_format::json::JsonFormat;
use datafusion::datasource::file_format::parquet::ParquetFormat;
use datafusion::datasource::listing::{ListingOptions, ListingTable};
use datafusion::datasource::listing::{ListingTableConfig, ListingTableUrl};
use datafusion::datasource::TableProvider;
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_common::DataFusionError;
use std::sync::Arc;
use tracing::{info_span, Instrument};
use crate::common::str::find;
use crate::handler::grpc::cluster_rpc;
use crate::infra::cluster;
use crate::infra::config::CONFIG;
use crate::infra::errors::{Error, ErrorCodes, Result};
use crate::meta::{self, StreamType};
use crate::service::promql::value;
use crate::service::search;
use crate::service::{file_list, get_partition_key_query};
mod storage;
mod wal;
pub async fn search(
req: &cluster_rpc::MetricsQueryRequest,
) -> Result<cluster_rpc::MetricsQueryResponse> {
let start = std::time::Instant::now();
let session_id = req.job.as_ref().unwrap().session_id.to_string();
let session_id = Arc::new(session_id);
let is_range_query = req.query.as_ref().unwrap().is_range_query;
// result
let mut results = Vec::new();
let span1 = info_span!("service:promql:search:grpc:in_cache");
// search in cache
let session_id1 = session_id.clone();
let org_id1 = req.org_id.clone();
let req1: cluster_rpc::MetricsQueryStmt = req.query.as_ref().unwrap().clone();
let task1 = tokio::task::spawn(
async move {
if cluster::is_ingester(&cluster::LOCAL_NODE_ROLE) {
wal::search(&session_id1, &org_id1, &req1).await
} else {
Ok(value::Value::None)
}
}
.instrument(span1),
);
let span2 = info_span!("service:promql:search:grpc:in_storage");
// search in object storage
let req_stype = req.stype;
let session_id2 = session_id.clone();
let org_id2 = req.org_id.clone();
let req2 = req.query.as_ref().unwrap().clone();
let task2 = tokio::task::spawn(
async move {
if req_stype == cluster_rpc::SearchType::CacheOnly as i32 {
Ok(value::Value::None)
} else {
storage::search(&session_id2, &org_id2, &req2).await
}
}
.instrument(span2),
);
// merge local wal
let value1 = match task1.await {
Ok(result) => match result {
Ok(val) => val,
Err(err) => {
log::error!("datafusion execute error: {}", err);
return Err(handle_datafusion_error(err));
}
},
Err(err) => {
return Err(Error::ErrorCode(ErrorCodes::ServerInternalError(
err.to_string(),
)))
}
};
match value1 {
value::Value::None => {}
_ => {
results.push(value1);
}
}
// merge object storage search
let value2 = match task2.await {
Ok(result) => match result {
Ok(val) => val,
Err(err) => {
log::error!("datafusion execute error: {}", err);
return Err(handle_datafusion_error(err));
}
},
Err(err) => {
return Err(Error::ErrorCode(ErrorCodes::ServerInternalError(
err.to_string(),
)))
}
};
match value2 {
value::Value::None => {}
_ => {
results.push(value2);
}
}
let mut result_type = if is_range_query { "matrix" } else { "vector" };
let mut resp = cluster_rpc::MetricsQueryResponse {
job: req.job.clone(),
took: start.elapsed().as_millis() as i32,
result_type: result_type.to_string(),
..Default::default()
};
if results.is_empty() {
return Ok(resp);
}
for result in results.iter() {
result_type = result.get_type();
match result {
value::Value::None => {
continue;
}
value::Value::Instant(v) => {
resp.result.push(cluster_rpc::Series {
metric: v.labels.iter().map(|l| l.as_ref().into()).collect(),
values: vec![],
value: Some((&v.value).into()),
scalar: None,
});
}
value::Value::Range(v) => {
resp.result.push(cluster_rpc::Series {
metric: v.labels.iter().map(|l| l.as_ref().into()).collect(),
values: v.values.iter().map(|l| l.into()).collect(),
value: None,
scalar: None,
});
}
value::Value::Vector(v) => {
v.iter().for_each(|v| {
resp.result.push(cluster_rpc::Series {
metric: v.labels.iter().map(|l| l.as_ref().into()).collect(),
values: vec![],
value: Some((&v.value).into()),
scalar: None,
});
});
}
value::Value::Matrix(v) => {
v.iter().for_each(|v| {
resp.result.push(cluster_rpc::Series {
metric: v.labels.iter().map(|l| l.as_ref().into()).collect(),
values: v.values.iter().map(|l| l.into()).collect(),
value: None,
scalar: None,
});
});
}
value::Value::Sample(v) => {
resp.result.push(cluster_rpc::Series {
metric: vec![],
values: vec![],
value: Some(v.into()),
scalar: None,
});
}
value::Value::Float(v) => {
resp.result.push(cluster_rpc::Series {
metric: vec![],
values: vec![],
value: None,
scalar: Some(*v),
});
}
}
}
// reset result type
resp.result_type = result_type.to_string();
Ok(resp)
}
pub fn handle_datafusion_error(err: DataFusionError) -> Error {
let err = err.to_string();
if err.contains("Schema error: No field named") {
let pos = err.find("Schema error: No field named").unwrap();
return match get_key_from_error(&err, pos) {
Some(key) => Error::ErrorCode(ErrorCodes::SearchFieldNotFound(key)),
None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)),
};
}
if err.contains("parquet not found") {
return Error::ErrorCode(ErrorCodes::SearchParquetFileNotFound);
}
if err.contains("Invalid function ") {
let pos = err.find("Invalid function ").unwrap();
return match get_key_from_error(&err, pos) {
Some(key) => Error::ErrorCode(ErrorCodes::SearchFunctionNotDefined(key)),
None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)),
};
}
if err.contains("Incompatible data types") {
let pos = err.find("for field").unwrap();
let pos_start = err[pos..].find(' ').unwrap();
let pos_end = err[pos + pos_start + 1..].find('.').unwrap();
let field = err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string();
return Error::ErrorCode(ErrorCodes::SearchFieldHasNoCompatibleDataType(field));
}
Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err))
}
fn get_key_from_error(err: &str, pos: usize) -> Option<String> {
for ponct in ['\'', '"'] {
let pos_start = err[pos..].find(ponct);
if pos_start.is_none() {
continue;
}
let pos_start = pos_start.unwrap();
let pos_end = err[pos + pos_start + 1..].find(ponct);
if pos_end.is_none() {
continue;
}
let pos_end = pos_end.unwrap();
return Some(err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string());
}
None
}
/// match a source is a valid file or not
pub async fn match_source(
org_id: &str,
stream_name: &str,
time_range: Option<(i64, i64)>,
filters: &[(String, String)],
source: &str,
is_wal: bool,
match_min_ts_only: bool,
) -> bool {
let stream_type = StreamType::Metrics;
// match org_id & table
if !source.starts_with(format!("files/{}/{}/{}/", org_id, stream_type, stream_name).as_str()) {
return false;
}
// check partition key
if !filter_source_by_partition_key(source, filters) {
return false;
}
if is_wal {
return true;
}
// check time range
let file_meta = file_list::get_file_meta(source).await.unwrap_or_default();
if file_meta.min_ts == 0 || file_meta.max_ts == 0 {
return true;
}
log::trace!(
"time range: {:?}, file time: {}-{}, {}",
time_range,
file_meta.min_ts,
file_meta.max_ts,
source
);
// match partation clause
if time_range.is_some() {
let (time_min, time_max) = time_range.unwrap();
if match_min_ts_only && time_min > 0 {
return file_meta.min_ts >= time_min && file_meta.min_ts < time_max;
}
if time_min > 0 && time_min > file_meta.max_ts {
return false;
}
if time_max > 0 && time_max < file_meta.min_ts {
return false;
}
}
true
}
fn filter_source_by_partition_key(source: &str, filters: &[(String, String)]) -> bool {
!filters.iter().any(|(k, v)| {
let field = get_partition_key_query(&format!("{k}="));
let value = get_partition_key_query(&format!("{k}={v}"));
find(source, &format!("/{field}")) && !find(source, &format!("/{value}/"))
})
}
pub async fn register_table(
session: &meta::search::Session,
org_id: &str,
stream_name: &str,
stream_type: StreamType,
schema: Option<Arc<Schema>>,
files: &Vec<String>,
file_type: FileType,
) -> datafusion::error::Result<SessionContext> {
if files.is_empty() {
return Ok(SessionContext::new());
}
let start = std::time::Instant::now();
let runtime_env = search::datafusion::exec::create_runtime_env()?;
let session_config = SessionConfig::new()
.with_information_schema(schema.is_none())
.with_batch_size(8192);
let ctx = SessionContext::with_config_rt(session_config.clone(), Arc::new(runtime_env));
// Configure listing options
let listing_options = match file_type {
FileType::PARQUET => {
let file_format = ParquetFormat::default().with_enable_pruning(Some(false));
ListingOptions::new(Arc::new(file_format))
.with_file_extension(FileType::PARQUET.get_ext())
.with_target_partitions(CONFIG.limit.cpu_num)
}
FileType::JSON => {
let file_format = JsonFormat::default();
ListingOptions::new(Arc::new(file_format))
.with_file_extension(FileType::JSON.get_ext())
.with_target_partitions(CONFIG.limit.cpu_num)
}
_ => {
return Err(DataFusionError::Execution(format!(
"Unsupported file type scheme {file_type:?}",
)));
}
};
let prefix = if session
.data_type
.eq(&search::datafusion::storage::file_list::SessionType::Wal)
{
format!(
"{}files/{}/{stream_type}/{}/",
&CONFIG.common.data_wal_dir, org_id, stream_name
)
} else {
search::datafusion::storage::file_list::set(&session.id, files)
.await
.unwrap();
format!("mem:///{}/", session.id)
};
let prefix = match ListingTableUrl::parse(prefix) {
Ok(url) => url,
Err(e) => {
return Err(datafusion::error::DataFusionError::Execution(format!(
"ListingTableUrl error: {e}",
)));
}
};
let prefixes = vec![prefix];
log::info!(
"Prepare table took {:.3} seconds.",
start.elapsed().as_secs_f64()
);
let mut config =
ListingTableConfig::new_with_multi_paths(prefixes).with_listing_options(listing_options);
let schema = if schema.is_none()
|| session
.data_type
.eq(&search::datafusion::storage::file_list::SessionType::Wal)
{
config = config.infer_schema(&ctx.state()).await.unwrap();
let table = ListingTable::try_new(config.clone())?;
let infered_schema = table.schema();
if schema.is_none() {
infered_schema
} else {
match Schema::try_merge(vec![
schema.unwrap().as_ref().to_owned(),
infered_schema.as_ref().to_owned(),
]) {
Ok(schema) => Arc::new(schema),
Err(e) => {
return Err(datafusion::error::DataFusionError::Execution(format!(
"ListingTable Merge schema error: {e}"
)));
}
}
}
} else {
schema.as_ref().unwrap().clone()
};
config = config.with_schema(schema);
log::info!(
"infer schema took {:.3} seconds.",
start.elapsed().as_secs_f64()
);
let table = ListingTable::try_new(config)?;
ctx.register_table(stream_name, Arc::new(table))?;
log::info!(
"Register table took {:.3} seconds.",
start.elapsed().as_secs_f64()
);
Ok(ctx)
}
src/service/promql/search/grpc/storage.rs
+203
View File
@@ -0,0 +1,203 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 async_trait::async_trait;
use datafusion::datasource::file_format::file_type::FileType;
use datafusion::error::{DataFusionError, Result};
use datafusion::prelude::SessionContext;
use promql_parser::parser;
use std::sync::Arc;
use std::time::{Duration, UNIX_EPOCH};
use tokio::sync::Semaphore;
use crate::handler::grpc::cluster_rpc;
use crate::infra::cache::file_data;
use crate::infra::config::CONFIG;
use crate::meta;
use crate::service::promql::{engine, value};
use crate::service::search::datafusion::storage::file_list::SessionType;
use crate::service::{db, file_list, promql, search};
struct StorageProvider {
org_id: String,
session_id: String,
}
#[async_trait]
impl engine::TableProvider for StorageProvider {
async fn create_context(
&self,
stream_name: &str,
time_range: (i64, i64),
filters: &[(String, String)],
) -> Result<SessionContext> {
// get file list
let files = get_file_list(&self.org_id, stream_name, time_range, filters).await?;
let file_count = files.len();
// load files to local cache
let mut tasks = Vec::new();
let semaphore = std::sync::Arc::new(Semaphore::new(CONFIG.limit.query_thread_num));
for file in files.iter() {
let file = file.clone();
let permit = semaphore.clone().acquire_owned().await.unwrap();
let task: tokio::task::JoinHandle<Result<(), anyhow::Error>> =
tokio::task::spawn(async move {
if !file_data::exist(&file).unwrap_or_default() {
if let Err(e) = file_data::download(&file).await {
log::error!("storage->search: load file {}, err: {}", &file, e);
}
};
drop(permit);
Ok(())
});
tasks.push(task);
}
for task in tasks {
match task.await {
Ok(ret) => {
if let Err(err) = ret {
return Err(DataFusionError::Execution(err.to_string()));
}
}
Err(err) => {
return Err(DataFusionError::Execution(err.to_string()));
}
};
}
log::info!(
"[TRACE] promql->search->storage: load files {} done",
file_count
);
// fetch all schema versions, get latest schema
let stream_type = meta::StreamType::Metrics;
let schema = match db::schema::get(&self.org_id, stream_name, Some(stream_type)).await {
Ok(schema) => schema,
Err(err) => {
log::error!("get schema error: {}", err);
return Err(datafusion::error::DataFusionError::Execution(
err.to_string(),
));
}
};
let schema = Arc::new(
schema
.to_owned()
.with_metadata(std::collections::HashMap::new()),
);
let session = meta::search::Session {
id: self.session_id.clone(),
data_type: SessionType::Storage,
};
super::register_table(
&session,
&self.org_id,
stream_name,
stream_type,
Some(schema),
&files,
FileType::PARQUET,
)
.await
}
}
/// search in remote object storage
/// search in local wal, which haven't been sync to object storage
pub async fn search(
session_id: &str,
org_id: &str,
query: &cluster_rpc::MetricsQueryStmt,
) -> Result<value::Value> {
let prom_expr = match parser::parse(&query.query) {
Ok(expr) => expr,
Err(e) => {
return Err(DataFusionError::Execution(e));
}
};
let eval_stmt = parser::EvalStmt {
expr: prom_expr,
start: UNIX_EPOCH
.checked_add(Duration::from_micros(query.start as u64))
.unwrap(),
end: UNIX_EPOCH
.checked_add(Duration::from_micros(query.end as u64))
.unwrap(),
interval: Duration::from_micros(query.step as u64),
lookback_delta: Duration::from_secs(300),
};
let mut engine = promql::QueryEngine::new(StorageProvider {
org_id: org_id.to_string(),
session_id: session_id.to_string(),
});
let data = engine.exec(eval_stmt).await?;
// clear session
search::datafusion::storage::file_list::clear(session_id)
.await
.unwrap();
Ok(data)
}
#[inline]
async fn get_file_list(
org_id: &str,
stream_name: &str,
time_range: (i64, i64),
filters: &[(String, String)],
) -> Result<Vec<String>> {
let (time_min, time_max) = time_range;
let results = match file_list::get_file_list(
org_id,
stream_name,
Some(meta::StreamType::Metrics),
time_min,
time_max,
)
.await
{
Ok(results) => results,
Err(err) => {
log::error!("get file list error: {}", err);
return Err(DataFusionError::Execution(
"get file list error".to_string(),
));
}
};
let mut files = Vec::new();
for file in results {
if super::match_source(
org_id,
stream_name,
Some(time_range),
filters,
&file,
false,
false,
)
.await
{
files.push(file.clone());
}
}
Ok(files)
}
src/service/promql/search/grpc/wal.rs
+188
View File
@@ -0,0 +1,188 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 async_trait::async_trait;
use datafusion::datasource::file_format::file_type::FileType;
use datafusion::error::{DataFusionError, Result};
use datafusion::prelude::SessionContext;
use promql_parser::parser;
use std::path::Path;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, UNIX_EPOCH};
use crate::common::file::scan_files;
use crate::handler::grpc::cluster_rpc;
use crate::infra::config::{self, CONFIG};
use crate::meta;
use crate::service::db;
use crate::service::promql;
use crate::service::promql::engine;
use crate::service::promql::value;
use crate::service::search::datafusion::storage::file_list::SessionType;
struct WalProvider {
org_id: String,
session_id: String,
}
#[async_trait]
impl engine::TableProvider for WalProvider {
async fn create_context(
&self,
stream_name: &str,
time_range: (i64, i64),
filters: &[(String, String)],
) -> Result<SessionContext> {
// get file list
let files = get_file_list(&self.org_id, stream_name, time_range, filters).await?;
// fetch all schema versions, get latest schema
let stream_type = meta::StreamType::Metrics;
let schema = match db::schema::get(&self.org_id, stream_name, Some(stream_type)).await {
Ok(schema) => schema,
Err(err) => {
log::error!("get schema error: {}", err);
return Err(datafusion::error::DataFusionError::Execution(
err.to_string(),
));
}
};
let schema = Arc::new(
schema
.to_owned()
.with_metadata(std::collections::HashMap::new()),
);
let session = meta::search::Session {
id: self.session_id.clone(),
data_type: SessionType::Wal,
};
super::register_table(
&session,
&self.org_id,
stream_name,
stream_type,
Some(schema),
&files,
FileType::JSON,
)
.await
}
}
/// search in local wal, which haven't been sync to object storage
pub async fn search(
session_id: &str,
org_id: &str,
query: &cluster_rpc::MetricsQueryStmt,
) -> Result<value::Value> {
// mark searching in wal
let searching = Searching::new();
let prom_expr = match parser::parse(&query.query) {
Ok(expr) => expr,
Err(e) => {
return Err(DataFusionError::Execution(e));
}
};
let eval_stmt = parser::EvalStmt {
expr: prom_expr,
start: UNIX_EPOCH
.checked_add(Duration::from_micros(query.start as u64))
.unwrap(),
end: UNIX_EPOCH
.checked_add(Duration::from_micros(query.end as u64))
.unwrap(),
interval: Duration::from_micros(query.step as u64),
lookback_delta: Duration::from_secs(300),
};
let mut engine = promql::QueryEngine::new(WalProvider {
org_id: org_id.to_string(),
session_id: session_id.to_string(),
});
let data = engine.exec(eval_stmt).await?;
// searching done.
drop(searching);
Ok(data)
}
/// get file list from local cache, no need match_source, each file will be searched
#[inline]
async fn get_file_list(
org_id: &str,
stream_name: &str,
time_range: (i64, i64),
filters: &[(String, String)],
) -> Result<Vec<String>> {
let pattern = format!(
"{}/files/{}/{}/{}/*.json",
&CONFIG.common.data_wal_dir,
org_id,
meta::StreamType::Metrics,
stream_name
);
let files = scan_files(&pattern);
let mut result = Vec::new();
let data_dir = match Path::new(&CONFIG.common.data_wal_dir).canonicalize() {
Ok(path) => path,
Err(_) => {
return Ok(result);
}
};
for file in files {
let file = Path::new(&file).canonicalize().unwrap();
let file = file.strip_prefix(&data_dir).unwrap();
let local_file = file.to_str().unwrap();
let file_path = file.parent().unwrap().to_str().unwrap().replace('\\', "/");
let file_name = file.file_name().unwrap().to_str().unwrap();
let file_name = file_name.replace('_', "/");
let source_file = format!("{file_path}/{file_name}");
if super::match_source(
org_id,
stream_name,
Some(time_range),
filters,
&source_file,
false,
true,
)
.await
{
result.push(format!("{}{local_file}", &CONFIG.common.data_wal_dir));
}
}
Ok(result)
}
/// Searching for marking searching in wal
struct Searching;
impl Searching {
pub fn new() -> Self {
config::SEARCHING_IN_WAL.store(1, Ordering::Relaxed);
Searching
}
}
impl Drop for Searching {
fn drop(&mut self) {
config::SEARCHING_IN_WAL.store(0, Ordering::Relaxed);
}
}
src/service/promql/search/mod.rs
+350
View File
@@ -0,0 +1,350 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 rustc_hash::FxHashMap;
use std::sync::Arc;
use std::{cmp::min, time::Duration};
use tonic::{codec::CompressionEncoding, metadata::MetadataValue, transport::Channel, Request};
use tracing::{info_span, instrument, Instrument};
use tracing_opentelemetry::OpenTelemetrySpanExt;
use uuid::Uuid;
use super::value::{InstantValue, Labels, Value};
use crate::handler::grpc::cluster_rpc;
use crate::infra::cluster::{self, get_internal_grpc_token};
use crate::infra::config::CONFIG;
use crate::infra::db::etcd;
use crate::infra::errors::{Error, ErrorCodes, Result};
use crate::service::promql::value::*;
pub mod grpc;
pub async fn search(org_id: &str, req: &super::MetricsQueryRequest) -> Result<Value> {
let root_span = info_span!("service:promql:search:enter");
let mut req: cluster_rpc::MetricsQueryRequest = req.to_owned().into();
req.org_id = org_id.to_string();
req.stype = cluster_rpc::SearchType::User as i32;
search_in_cluster(req).instrument(root_span).await
}
#[inline(always)]
async fn get_queue_lock() -> Result<etcd::Locker> {
let mut lock = etcd::Locker::new("search/cluster_queue");
lock.lock(0).await.map_err(server_internal_error)?;
Ok(lock)
}
#[instrument(name = "service:promql:search:cluster", skip(req))]
async fn search_in_cluster(req: cluster_rpc::MetricsQueryRequest) -> Result<Value> {
let start = std::time::Instant::now();
// get a cluster search queue lock
let locker = if CONFIG.common.local_mode {
None
} else {
Some(get_queue_lock().await?)
};
let took_wait = start.elapsed().as_millis() as usize;
// get nodes from cluster
let mut nodes = cluster::get_cached_online_query_nodes().unwrap();
// sort nodes by node_id this will improve hit cache ratio
nodes.sort_by_key(|x| x.id);
let nodes = nodes;
let querier_num = match nodes
.iter()
.filter(|node| cluster::is_querier(&node.role))
.count()
{
0 => 1,
n => n as i64,
};
let is_range_query = req.query.as_ref().unwrap().is_range_query;
let req_start: i64 = req.query.as_ref().unwrap().start;
let req_end: i64 = req.query.as_ref().unwrap().end;
let req_step: i64 = req.query.as_ref().unwrap().step;
let point_num = if !is_range_query || req_start == req_end {
1
} else {
(req_end - req_start) / req_step
};
let offset = if querier_num >= point_num {
req_step
} else {
((point_num / querier_num) + 1) * req_step
};
// partition request, here plus 1 second, because division is integer, maybe lose some precision
let mut session_id = Uuid::new_v4().to_string();
let job = cluster_rpc::Job {
session_id: session_id.clone(),
job: session_id.split_off(30), // take the last 6 characters as job id
stage: 0,
partition: 0,
};
// make cluster request
let mut tasks = Vec::new();
let mut offset_start = req_start;
for (partition_no, node) in nodes.iter().cloned().enumerate() {
let mut req = req.clone();
let mut req_query = req.query.as_mut().unwrap();
let mut job = job.clone();
job.partition = partition_no as i32;
req.job = Some(job);
req.stype = cluster_rpc::SearchType::CacheOnly as i32;
let is_querier = cluster::is_querier(&node.role);
if is_querier {
if offset_start <= req_end {
req.stype = cluster_rpc::SearchType::Cluster as i32;
req_query.start = offset_start;
req_query.end = min(req_end, offset_start + offset);
offset_start += offset;
} else if !cluster::is_ingester(&node.role) {
continue; // no need more querier
}
}
log::info!(
"[TRACE] promql->search->partition: node: {}, is_querier: {}, start: {}, end: {}",
node.id,
is_querier,
req_query.start,
req_query.end,
);
let grpc_span = info_span!("service:promql:search:cluster:grpc_search");
let node_addr = node.grpc_addr.clone();
let task = tokio::task::spawn(
async move {
let org_id: MetadataValue<_> = match req.org_id.parse() {
Ok(org_id) => org_id,
Err(_) => return Err(Error::Message("invalid org_id".to_string())),
};
let mut request = tonic::Request::new(req);
request.set_timeout(Duration::from_secs(CONFIG.grpc.timeout));
opentelemetry::global::get_text_map_propagator(|propagator| {
propagator.inject_context(
&tracing::Span::current().context(),
&mut MetadataMap(request.metadata_mut()),
)
});
let token: MetadataValue<_> = match get_internal_grpc_token().parse() {
Ok(token) => token,
Err(_) => return Err(Error::Message("invalid token".to_string())),
};
let channel = match Channel::from_shared(node_addr).unwrap().connect().await {
Ok(channel) => channel,
Err(err) => {
log::error!(
"[TRACE] promql->search->grpc: node: {}, connect err: {:?}",
node.id,
err
);
return Err(server_internal_error("connect search node error"));
}
};
let mut client = cluster_rpc::metrics_client::MetricsClient::with_interceptor(
channel,
move |mut req: Request<()>| {
req.metadata_mut().insert("authorization", token.clone());
req.metadata_mut()
.insert(CONFIG.grpc.org_header_key.as_str(), org_id.clone());
Ok(req)
},
);
client = client
.send_compressed(CompressionEncoding::Gzip)
.accept_compressed(CompressionEncoding::Gzip);
let response: cluster_rpc::MetricsQueryResponse = match client.query(request).await {
Ok(res) => res.into_inner(),
Err(err) => {
log::error!(
"[TRACE] promql->search->grpc: node: {}, search err: {:?}",
node.id,
err
);
if err.code() == tonic::Code::Internal {
let err = ErrorCodes::from_json(err.message())?;
return Err(Error::ErrorCode(err));
}
return Err(server_internal_error("search node error"));
}
};
log::info!(
"[TRACE] promql->search->grpc: result node: {}, is_querier: {}, took: {}, files: {}",
node.id,
is_querier,
response.took,
response.file_count
);
Ok(response)
}
.instrument(grpc_span),
);
tasks.push(task);
}
let mut results = Vec::new();
for task in tasks {
let result = task
.await
.map_err(|err| Error::ErrorCode(ErrorCodes::ServerInternalError(err.to_string())))?;
match result {
Ok(res) => results.push(res),
Err(err) => {
// search done, release lock
if locker.is_some() {
if let Err(e) = locker.unwrap().unlock().await {
log::error!("search in cluster unlock error: {}", e);
}
}
return Err(err);
}
}
}
// search done, release lock
if locker.is_some() {
if let Err(e) = locker.unwrap().unlock().await {
log::error!("search in cluster unlock error: {}", e);
}
}
// merge multiple instances data
let mut file_count = 0;
let mut scan_size = 0;
let mut result_type = String::new();
let mut series_data: Vec<cluster_rpc::Series> = Vec::new();
for resp in results {
file_count += resp.file_count;
scan_size += resp.scan_size;
if result_type.is_empty() {
result_type = resp.result_type.clone();
}
resp.result.into_iter().for_each(|series| {
series_data.push(series);
});
}
// merge result
let values = if result_type == "matrix" {
merge_matrix_query(&series_data)
} else if result_type == "vector" {
merge_vector_query(&series_data)
} else if result_type == "scalar" {
merge_scalar_query(&series_data)
} else {
return Err(server_internal_error("invalid result type"));
};
log::info!(
"[TRACE] promql->search->result: took: {}, took_wait: {}, file_count: {}, scan_size: {}",
start.elapsed().as_millis(),
took_wait,
file_count,
scan_size,
);
Ok(values)
}
fn merge_matrix_query(series: &[cluster_rpc::Series]) -> Value {
let mut merged_data = FxHashMap::default();
let mut merged_metrics = FxHashMap::default();
series.iter().for_each(|v| {
let labels: Labels = v.metric.iter().map(|v| Arc::new(Label::from(v))).collect();
let values: Vec<Sample> = v.values.iter().map(Sample::from).collect();
merged_data
.entry(signature(&labels))
.or_insert_with(Vec::new)
.extend(values);
merged_metrics.insert(signature(&labels), labels);
});
let merged_data = merged_data
.into_iter()
.map(|(sig, values)| RangeValue {
labels: merged_metrics.get(&sig).unwrap().to_owned(),
values,
time_range: None,
})
.collect::<Vec<_>>();
// sort data
let mut value = Value::Matrix(merged_data);
value.sort();
value
}
fn merge_vector_query(series: &[cluster_rpc::Series]) -> Value {
let mut merged_data = FxHashMap::default();
let mut merged_metrics: std::collections::HashMap<
Signature,
Vec<Arc<Label>>,
std::hash::BuildHasherDefault<rustc_hash::FxHasher>,
> = FxHashMap::default();
series.iter().for_each(|v| {
let labels: Labels = v.metric.iter().map(|v| Arc::new(Label::from(v))).collect();
let value: Sample = v.value.as_ref().unwrap().into();
merged_data.insert(signature(&labels), value);
merged_metrics.insert(signature(&labels), labels);
});
let merged_data = merged_data
.into_iter()
.map(|(sig, value)| InstantValue {
labels: merged_metrics.get(&sig).unwrap().to_owned(),
value,
})
.collect::<Vec<_>>();
// sort data
let mut value = Value::Vector(merged_data);
value.sort();
value
}
fn merge_scalar_query(series: &[cluster_rpc::Series]) -> Value {
let mut sample: Sample = Default::default();
series.iter().for_each(|v| {
if v.scalar.is_some() {
sample.value = v.scalar.unwrap();
}
if v.value.is_some() {
sample = v.value.as_ref().unwrap().into();
}
});
Value::Sample(sample)
}
struct MetadataMap<'a>(&'a mut tonic::metadata::MetadataMap);
impl<'a> opentelemetry::propagation::Injector for MetadataMap<'a> {
/// Set a key and value in the MetadataMap. Does nothing if the key or value are not valid inputs
fn set(&mut self, key: &str, value: String) {
if let Ok(key) = tonic::metadata::MetadataKey::from_bytes(key.as_bytes()) {
if let Ok(val) = tonic::metadata::MetadataValue::try_from(&value) {
self.0.insert(key, val);
}
}
}
}
fn server_internal_error(error: impl ToString) -> Error {
Error::ErrorCode(ErrorCodes::ServerInternalError(error.to_string()))
}
src/service/promql/value.rs
+325
View File
@@ -0,0 +1,325 @@
// Copyright 2022 Zinc Labs Inc. and Contributors
//
// 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 serde::{
ser::{SerializeSeq, SerializeStruct, Serializer},
Serialize,
};
use std::{cmp::Ordering, sync::Arc};
pub const FIELD_HASH: &str = "__hash__";
pub const FIELD_NAME: &str = "__name__";
pub const FIELD_TYPE: &str = "metric_type";
pub const FIELD_TIME: &str = "_timestamp";
pub const FIELD_VALUE: &str = "value";
pub const FIELD_BUCKET: &str = "le";
pub const TYPE_COUNTER: &str = "counter";
pub const TYPE_GAUGE: &str = "gauge";
pub const TYPE_HISTOGRAM: &str = "histogram";
pub const TYPE_SUMMARY: &str = "summary";
// See https://docs.rs/indexmap/latest/indexmap/#alternate-hashers
type FxIndexMap<K, V> =
indexmap::IndexMap<K, V, std::hash::BuildHasherDefault<rustc_hash::FxHasher>>;
pub type Labels = Vec<Arc<Label>>;
#[derive(Debug, Clone)]
pub struct Label {
pub name: String,
pub value: String,
}
#[derive(Debug, Default, Clone, Copy)]
pub struct Sample {
/// Time in microseconds
pub timestamp: i64,
pub value: f64,
}
impl Serialize for Sample {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut seq = serializer.serialize_seq(Some(2))?;
seq.serialize_element(&(self.timestamp / 1_000_000))?;
seq.serialize_element(&self.value.to_string())?;
seq.end()
}
}
#[derive(Debug, Clone)]
pub struct InstantValue {
pub labels: Labels,
pub value: Sample,
}
impl Serialize for InstantValue {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut seq = serializer.serialize_struct("instant_value", 2)?;
let labels_map = self
.labels
.iter()
.map(|l| (l.name.as_str(), l.value.as_str()))
.collect::<FxIndexMap<_, _>>();
seq.serialize_field("metric", &labels_map)?;
seq.serialize_field("value", &self.value)?;
seq.end()
}
}
#[derive(Debug, Clone)]
pub struct RangeValue {
pub labels: Labels,
pub time_range: Option<(i64, i64)>, // start, end
pub values: Vec<Sample>,
}
impl Serialize for RangeValue {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut seq = serializer.serialize_struct("range_value", 2)?;
let labels_map = self
.labels
.iter()
.map(|l| (l.name.as_str(), l.value.as_str()))
.collect::<FxIndexMap<_, _>>();
seq.serialize_field("metric", &labels_map)?;
seq.serialize_field("values", &self.values)?;
seq.end()
}
}
impl RangeValue {
/// Returns first and last data points, [extrapolated] to the time window
/// boundaries.
///
/// [extrapolated]: https://promlabs.com/blog/2021/01/29/how-exactly-does-promql-calculate-rates/#extrapolation-of-data
pub(crate) fn extrapolate(&self) -> Option<(Sample, Sample)> {
let samples = &self.values;
if samples.len() < 2 {
return None;
}
let first = samples.first().unwrap();
let last = samples.last().unwrap();
let (t_start, t_end) = self.time_range.unwrap();
assert!(t_start < t_end);
assert!(t_start <= first.timestamp);
assert!(first.timestamp <= last.timestamp);
assert!(last.timestamp <= t_end);
Some((
if first.timestamp == t_start {
*first
} else {
extrapolated_sample(first, last, t_start)
},
if last.timestamp == t_end {
*last
} else {
extrapolated_sample(first, last, t_end)
},
))
}
}
// https://promlabs.com/blog/2021/01/29/how-exactly-does-promql-calculate-rates/#extrapolation-of-data
fn extrapolated_sample(p1: &Sample, p2: &Sample, t: i64) -> Sample {
let dt = p2.timestamp - p1.timestamp;
let dv = p2.value - p1.value;
let dt2 = t - p1.timestamp;
let dv2 = dv * dt2 as f64 / dt as f64;
Sample {
timestamp: t,
value: p1.value + dv2,
}
}
pub fn labels_value(labels: &Labels, name: &str) -> Option<String> {
labels
.binary_search_by_key(&name, |label| label.name.as_str())
.ok()
.map(|index| labels[index].value.clone())
}
#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
pub enum Value {
Instant(InstantValue),
Range(RangeValue),
Vector(Vec<InstantValue>),
Matrix(Vec<RangeValue>),
Sample(Sample), // only used for return literal value
Float(f64),
None,
}
impl Value {
pub(crate) fn get_ref_matrix_values(&self) -> Option<&Vec<RangeValue>> {
match self {
Value::Matrix(values) => Some(values),
_ => None,
}
}
pub fn get_type(&self) -> &str {
match self {
Value::Instant(_) => "vector",
Value::Range(_) => "matrix",
Value::Vector(_) => "vector",
Value::Matrix(_) => "matrix",
Value::Sample(_) => "scalar",
Value::Float(_) => "scalar",
Value::None => "scalar",
}
}
pub fn sort(&mut self) {
match self {
Value::Vector(v) => {
v.sort_by(|a, b| {
b.value
.value
.partial_cmp(&a.value.value)
.unwrap_or(Ordering::Equal)
});
}
Value::Matrix(v) => {
v.sort_by(|a, b| {
let a = a.values.iter().map(|x| x.value).sum::<f64>();
let b = b.values.iter().map(|x| x.value).sum::<f64>();
b.partial_cmp(&a).unwrap_or(Ordering::Equal)
});
}
_ => {}
}
}
}
#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
pub struct Signature([u8; 32]);
impl From<Signature> for String {
fn from(sig: Signature) -> Self {
hex::encode(sig.0)
}
}
// REFACTORME: make this a method of `Metric`
pub fn signature(labels: &Labels) -> Signature {
signature_without_labels(labels, &[])
}
/// `signature_without_labels` is just as [`signature`], but only for labels not matching `names`.
// REFACTORME: make this a method of `Metric`
pub fn signature_without_labels(labels: &Labels, exclude_names: &[&str]) -> Signature {
let mut hasher = blake3::Hasher::new();
labels
.iter()
.filter(|item| !exclude_names.contains(&item.name.as_str()))
.for_each(|item| {
hasher.update(item.name.as_bytes());
hasher.update(item.value.as_bytes());
});
Signature(hasher.finalize().into())
}
#[cfg(test)]
mod tests {
use super::*;
use expect_test::expect;
#[test]
fn test_signature_without_labels() {
let mut labels: Labels = Default::default();
labels.push(Arc::new(Label {
name: "a".to_owned(),
value: "1".to_owned(),
}));
labels.push(Arc::new(Label {
name: "b".to_owned(),
value: "2".to_owned(),
}));
labels.push(Arc::new(Label {
name: "c".to_owned(),
value: "3".to_owned(),
}));
labels.push(Arc::new(Label {
name: "d".to_owned(),
value: "4".to_owned(),
}));
let sig = signature(&labels);
expect![[r#"
"f287fde2994111abd7740b5c7c28b0eeabe3f813ae65397bb6acb684e2ab6b22"
"#]]
.assert_debug_eq(&String::from(sig));
let sig: String = signature_without_labels(&labels, &["a", "c"]).into();
expect![[r#"
"ec9c3a0c9c03420d330ab62021551cffe993c07b20189c5ed831dad22f54c0c7"
"#]]
.assert_debug_eq(&sig);
assert_eq!(sig.len(), 64);
assert_eq!(signature(&labels), signature_without_labels(&labels, &[]));
}
#[test]
fn test_extrapolated_sample() {
let p1 = Sample {
timestamp: 100,
value: 10.0,
};
let p2 = Sample {
timestamp: 200,
value: 20.0,
};
let p3 = extrapolated_sample(&p1, &p2, 300);
assert_eq!(p3.timestamp, 300);
assert_eq!(p3.value, 30.0);
let p1 = Sample {
timestamp: 225,
value: 1.0,
};
let p2 = Sample {
timestamp: 675,
value: 2.0,
};
let p3 = extrapolated_sample(&p1, &p2, 750);
let p4 = extrapolated_sample(&p1, &p2, 150);
assert_eq!(format!("{:.2}", p3.value - p4.value), "1.33");
let p1 = Sample {
timestamp: 375,
value: 1.0,
};
let p2 = Sample {
timestamp: 675,
value: 2.0,
};
let p3 = extrapolated_sample(&p1, &p2, 750);
let p4 = extrapolated_sample(&p1, &p2, 300);
assert_eq!(format!("{:.2}", p3.value - p4.value), "1.50");
}
}
src/service/router/mod.rs
+11 -1
View File
@@ -21,7 +21,17 @@ use rand::thread_rng;
use crate::infra::cluster;
use crate::infra::config::CONFIG;
const QUERIER_ROUTES: [&str; 3] = ["/_search", "/_around", "/_values"];
const QUERIER_ROUTES: [&str; 9] = [
"/_search",
"/_around",
"/_values",
"/prometheus/v1/series",
"/prometheus/v1/query_range",
"/prometheus/v1/query",
"/prometheus/v1/metadata",
"/prometheus/v1/labels",
"/prometheus/v1/label/",
];
#[inline]
pub fn is_router() -> bool {
src/service/search/datafusion/exec.rs
+3 -3
View File
@@ -89,7 +89,7 @@ pub async fn sql(
}
};
let prefix = if session.data_type.eq(&file_list::SessionType::Cache) {
let prefix = if session.data_type.eq(&file_list::SessionType::Wal) {
format!(
"{}files/{}/{stream_type}/{}/",
&CONFIG.common.data_wal_dir, sql.org_id, sql.stream_name
@@ -114,7 +114,7 @@ pub async fn sql(
let mut config =
ListingTableConfig::new_with_multi_paths(prefixes).with_listing_options(listing_options);
let schema = if schema.is_none() || session.data_type.eq(&file_list::SessionType::Cache) {
let schema = if schema.is_none() || session.data_type.eq(&file_list::SessionType::Wal) {
config = config.infer_schema(&ctx.state()).await.unwrap();
let table = ListingTable::try_new(config.clone())?;
let infered_schema = table.schema();
@@ -886,7 +886,7 @@ pub async fn merge_parquet_files(
Ok(file_meta)
}
fn create_runtime_env() -> Result<RuntimeEnv> {
pub fn create_runtime_env() -> Result<RuntimeEnv> {
let object_store_registry = DefaultObjectStoreRegistry::new();
let mem = super::storage::memory::InMemory::new();
src/service/search/datafusion/mod.rs
+2 -2
View File
@@ -16,8 +16,8 @@ use crate::meta::functions::ZoFunction;
mod date_format_udf;
pub mod exec;
mod match_udf;
mod regexp_udf;
pub mod match_udf;
pub mod regexp_udf;
pub mod storage;
mod time_range_udf;
#[cfg(feature = "zo_functions")]
src/service/search/datafusion/regexp_udf.rs
+4 -5
View File
@@ -81,7 +81,7 @@ pub fn regex_match_expr_impl(matches: bool) -> ScalarFunctionImplementation {
}
ColumnarValue::Scalar(ScalarValue::Utf8(pattern)) => pattern,
ColumnarValue::Scalar(arg) => {
return Err(DataFusionError::Internal(format!(
return Err(DataFusionError::Plan(format!(
"Expected string pattern to regex_match({matches}), got: {arg:?}"
)))
}
@@ -97,9 +97,8 @@ pub fn regex_match_expr_impl(matches: bool) -> ScalarFunctionImplementation {
// the golang regexp library which is different than Rust's regexp
let pattern = clean_non_meta_escapes(pattern);
let pattern = regex::Regex::new(&pattern).map_err(|e| {
DataFusionError::Internal(format!("error compiling regex pattern: {e}"))
})?;
let pattern = regex::Regex::new(&pattern)
.map_err(|e| DataFusionError::Plan(format!("error compiling regex pattern: {e}")))?;
match &args[0] {
ColumnarValue::Array(arr) => {
@@ -118,7 +117,7 @@ pub fn regex_match_expr_impl(matches: bool) -> ScalarFunctionImplementation {
let res = row.as_ref().map(|v| pattern.is_match(v) == matches);
Ok(ColumnarValue::Scalar(ScalarValue::Boolean(res)))
}
ColumnarValue::Scalar(v) => Err(DataFusionError::Internal(format!(
ColumnarValue::Scalar(v) => Err(DataFusionError::Plan(format!(
"regex_match({matches}) expected first argument to be utf8, got ('{v}')"
))),
}
src/service/search/datafusion/storage/file_list.rs
+2 -3
View File
@@ -24,9 +24,8 @@ lazy_static! {
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SessionType {
Cache,
Local,
Remote,
Wal,
Storage,
}
pub async fn get(session_id: &str) -> Result<Vec<ObjectMeta>, anyhow::Error> {
src/service/search/cache.rs → src/service/search/grpc/cache.rs
+11 -8
View File
@@ -19,14 +19,14 @@ use std::sync::atomic::Ordering;
use std::sync::Arc;
use tracing::info_span;
use super::datafusion::storage::file_list::SessionType;
use super::sql::Sql;
use crate::common::file::scan_files;
use crate::infra::config::{self, CONFIG};
use crate::infra::errors::{Error, ErrorCodes};
use crate::meta;
use crate::service::db;
use crate::service::file_list::calculate_local_files_size;
use crate::service::search::datafusion::storage::file_list::SessionType;
use crate::service::search::sql::Sql;
/// search in local cache, which haven't been sync to object storage
#[tracing::instrument(
@@ -92,7 +92,7 @@ pub async fn search(
let session = meta::search::Session {
id: session_id.to_string(),
data_type: SessionType::Cache,
data_type: SessionType::Wal,
};
let result = match super::datafusion::exec::sql(
&session,
@@ -118,7 +118,7 @@ pub async fn search(
Ok((result, file_count, scan_size as usize))
}
/// get file list from local cache, no need match_source, each file will be searched
/// get file list from local wal, no need match_source, each file will be searched
#[inline]
async fn get_file_list(sql: &Sql, stream_type: meta::StreamType) -> Result<Vec<String>, Error> {
let pattern = format!(
@@ -142,25 +142,28 @@ async fn get_file_list(sql: &Sql, stream_type: meta::StreamType) -> Result<Vec<S
let file_name = file.file_name().unwrap().to_str().unwrap();
let file_name = file_name.replace('_', "/");
let source_file = format!("{file_path}/{file_name}");
if sql.match_source(&source_file, false, stream_type).await {
if sql
.match_source(&source_file, false, true, stream_type)
.await
{
result.push(format!("{}{local_file}", &CONFIG.common.data_wal_dir));
}
}
Ok(result)
}
/// Searching for marking searching in cache
/// Searching for marking searching in wal
struct Searching;
impl Searching {
pub fn new() -> Self {
config::SEARCHING_IN_CACHE.store(1, Ordering::Relaxed);
config::SEARCHING_IN_WAL.store(1, Ordering::Relaxed);
Searching
}
}
impl Drop for Searching {
fn drop(&mut self) {
config::SEARCHING_IN_CACHE.store(0, Ordering::Relaxed);
config::SEARCHING_IN_WAL.store(0, Ordering::Relaxed);
}
}
src/service/search/exec.rs → src/service/search/grpc/mod.rs
+71 -11
View File
@@ -12,16 +12,24 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use ::datafusion::arrow::ipc;
use ::datafusion::arrow::{ipc, record_batch::RecordBatch};
use ahash::AHashMap as HashMap;
use datafusion_common::DataFusionError;
use std::sync::Arc;
use tracing::{info_span, Instrument};
use super::datafusion;
use crate::handler::grpc::cluster_rpc;
use crate::infra::cluster;
use crate::infra::errors::{Error, ErrorCodes};
use crate::meta::StreamType;
#[tracing::instrument(name = "service:search:exec", skip(req))]
mod cache;
mod storage;
pub type SearchResult = Result<(HashMap<String, Vec<RecordBatch>>, usize, usize), Error>;
#[tracing::instrument(name = "service:search:grpc:search", skip(req))]
pub async fn search(
req: &cluster_rpc::SearchRequest,
) -> Result<cluster_rpc::SearchResponse, Error> {
@@ -36,17 +44,23 @@ pub async fn search(
let mut file_count = 0;
let mut scan_size = 0;
let span1 = info_span!("service:search:exec:in_cache");
let span1 = info_span!("service:search:grpc:in_cache");
// search in cache
let session_id1 = session_id.clone();
let sql1 = sql.clone();
let task1 = tokio::task::spawn(
async move { super::cache::search(&session_id1, sql1, stream_type).await }
.instrument(span1),
async move {
if cluster::is_ingester(&cluster::LOCAL_NODE_ROLE) {
cache::search(&session_id1, sql1, stream_type).await
} else {
Ok((HashMap::new(), 0, 0))
}
}
.instrument(span1),
);
let span2 = info_span!("service:search:exec:in_storage");
let span2 = info_span!("service:search:grpc:in_storage");
// search in object storage
let req_stype = req.stype;
@@ -58,7 +72,7 @@ pub async fn search(
if req_stype == cluster_rpc::SearchType::CacheOnly as i32 {
Ok((HashMap::new(), 0, 0))
} else {
super::storage::search(&session_id2, sql2, file_list.as_slice(), stream_type).await
storage::search(&session_id2, sql2, file_list.as_slice(), stream_type).await
}
}
.instrument(span2),
@@ -112,7 +126,7 @@ pub async fn search(
file_count += file_count2;
scan_size += scan_size2;
let span3 = info_span!("service:search:exec:merge");
let span3 = info_span!("service:search:grpc:merge");
let guard3 = span3.enter();
// merge all batches
@@ -134,18 +148,18 @@ pub async fn search(
Ok(res) => res,
Err(err) => {
log::error!("datafusion merge error: {}", err);
return Err(super::handle_datafusion_error(err));
return Err(handle_datafusion_error(err));
}
};
}
// clear session data
super::datafusion::storage::file_list::clear(&session_id)
datafusion::storage::file_list::clear(&session_id)
.await
.unwrap();
drop(guard3);
let span4 = info_span!("service:search:exec:response");
let span4 = info_span!("service:search:grpc:response");
let _guard4 = span4.enter();
// final result
@@ -210,3 +224,49 @@ pub async fn search(
Ok(result)
}
pub fn handle_datafusion_error(err: DataFusionError) -> Error {
let err = err.to_string();
if err.contains("Schema error: No field named") {
let pos = err.find("Schema error: No field named").unwrap();
return match get_key_from_error(&err, pos) {
Some(key) => Error::ErrorCode(ErrorCodes::SearchFieldNotFound(key)),
None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)),
};
}
if err.contains("parquet not found") {
return Error::ErrorCode(ErrorCodes::SearchParquetFileNotFound);
}
if err.contains("Invalid function ") {
let pos = err.find("Invalid function ").unwrap();
return match get_key_from_error(&err, pos) {
Some(key) => Error::ErrorCode(ErrorCodes::SearchFunctionNotDefined(key)),
None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)),
};
}
if err.contains("Incompatible data types") {
let pos = err.find("for field").unwrap();
let pos_start = err[pos..].find(' ').unwrap();
let pos_end = err[pos + pos_start + 1..].find('.').unwrap();
let field = err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string();
return Error::ErrorCode(ErrorCodes::SearchFieldHasNoCompatibleDataType(field));
}
Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err))
}
fn get_key_from_error(err: &str, pos: usize) -> Option<String> {
for punctuation in ['\'', '"'] {
let pos_start = err[pos..].find(punctuation);
if pos_start.is_none() {
continue;
}
let pos_start = pos_start.unwrap();
let pos_end = err[pos + pos_start + 1..].find(punctuation);
if pos_end.is_none() {
continue;
}
let pos_end = pos_end.unwrap();
return Some(err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string());
}
None
}
src/service/search/storage.rs → src/service/search/grpc/storage.rs
+4 -4
View File
@@ -19,12 +19,12 @@ use std::sync::Arc;
use tokio::sync::Semaphore;
use tracing::{info_span, Instrument};
use super::datafusion::storage::file_list::SessionType;
use super::sql::Sql;
use crate::infra::cache::file_data;
use crate::infra::config::CONFIG;
use crate::infra::errors::{Error, ErrorCodes};
use crate::meta;
use crate::service::search::datafusion::storage::file_list::SessionType;
use crate::service::search::sql::Sql;
use crate::service::{db, file_list};
/// search in remote object storage
@@ -171,7 +171,7 @@ pub async fn search(
let sql = sql.clone();
let session = meta::search::Session {
id: format!("{session_id}-{ver}"),
data_type: SessionType::Remote,
data_type: SessionType::Storage,
};
// cacluate the diff between latest schema and group schema
let mut diff_fields = HashMap::new();
@@ -252,7 +252,7 @@ async fn get_file_list(sql: &Sql, stream_type: meta::StreamType) -> Result<Vec<S
let mut files = Vec::new();
for file in results {
if sql.match_source(&file, false, stream_type).await {
if sql.match_source(&file, false, false, stream_type).await {
files.push(file.clone());
}
}
src/service/search/mod.rs
+35 -88
View File
@@ -14,7 +14,6 @@
use ::datafusion::arrow::{datatypes::Schema, ipc, json as arrow_json, record_batch::RecordBatch};
use ahash::AHashMap as HashMap;
use datafusion_common::DataFusionError;
use std::io::Cursor;
use std::sync::Arc;
use std::{cmp::min, time::Duration};
@@ -26,7 +25,6 @@ use uuid::Uuid;
use crate::common::json;
use crate::handler::grpc::cluster_rpc;
use crate::infra::cluster::{self, get_internal_grpc_token};
use crate::infra::config::CONFIG;
use crate::infra::db::etcd;
use crate::infra::errors::{Error, ErrorCodes};
@@ -35,20 +33,16 @@ use crate::meta::search::Response;
use crate::meta::StreamType;
use crate::service::{db, file_list};
mod cache;
pub mod datafusion;
pub mod exec;
pub mod grpc;
pub mod sql;
mod storage;
pub type SearchResult = Result<(HashMap<String, Vec<RecordBatch>>, usize, usize), Error>;
pub async fn search(
org_id: &str,
stream_type: StreamType,
req: &meta::search::Request,
) -> Result<Response, Error> {
let root_span = info_span!("srv:search:enter");
let root_span = info_span!("service:search:enter");
let mut req: cluster_rpc::SearchRequest = req.to_owned().into();
req.org_id = org_id.to_string();
@@ -57,6 +51,7 @@ pub async fn search(
search_in_cluster(req).instrument(root_span).await
}
#[inline(always)]
async fn get_queue_lock() -> Result<etcd::Locker, Error> {
let mut lock = etcd::Locker::new("search/cluster_queue");
lock.lock(0).await.map_err(server_internal_error)?;
@@ -84,7 +79,7 @@ async fn get_times(sql: &sql::Sql, stream_type: StreamType) -> (i64, i64) {
}
#[instrument(skip_all)]
async fn partition_by_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec<String> {
async fn get_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec<String> {
let (time_min, time_max) = get_times(sql, stream_type).await;
match file_list::get_file_list(
&sql.org_id,
@@ -108,7 +103,7 @@ async fn partition_by_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec<
} else {
// less than 1 hour, use file meta reduce file list
for file in file_list {
if sql.match_source(&file, false, stream_type).await {
if sql.match_source(&file, false, false, stream_type).await {
files.push(file.clone());
}
}
@@ -119,7 +114,7 @@ async fn partition_by_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec<
}
}
#[instrument(name = "srv:search:cluster", skip(req))]
#[instrument(name = "service:search:cluster", skip(req))]
async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response, Error> {
let start = std::time::Instant::now();
@@ -132,7 +127,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
let took_wait = start.elapsed().as_millis() as usize;
//XXX span1.exit(); // drop span1
//XXX let span2 = info_span!("srv:search:cluster:prepare_base").entered();
//XXX let span2 = info_span!("service:search:cluster:prepare_base").entered();
// get nodes from cluster
let mut nodes = cluster::get_cached_online_query_nodes().unwrap();
@@ -164,7 +159,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
}
};
let file_list = partition_by_file_list(&meta, stream_type).await;
let file_list = get_file_list(&meta, stream_type).await;
let file_num = file_list.len();
let offset = if querier_num >= file_num {
1
@@ -183,7 +178,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
};
//XXX span3.exit(); // drop span3
//XXX let span4 = info_span!("srv:search:cluster:do_search").entered();
//XXX let span4 = info_span!("service:search:cluster:do_search").entered();
// make grpc auth token
/* let root_user = ROOT_USER.clone();
@@ -224,7 +219,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
req_file_list_end,
);
let grpc_span = info_span!("srv:search:cluster:grpc_search");
let grpc_span = info_span!("service:search:cluster:grpc_search");
let node_addr = node.grpc_addr.clone();
//let credentials_str = C.clone();
@@ -323,7 +318,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
}
//XXX span4.exit(); // drop span4
//XXX let span6 = info_span!("srv:search:cluster:release_queue_lock").entered();
//XXX let span6 = info_span!("service:search:cluster:release_queue_lock").entered();
// search done, release lock
if locker.is_some() {
@@ -333,44 +328,42 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
}
//XXX span6.exit(); // drop span6
//XXX let span7 = info_span!("srv:search:cluster:merge_result").entered();
//XXX let span7 = info_span!("service:search:cluster:merge_result").entered();
// merge multiple instances data
let mut file_count = 0;
let mut scan_size = 0;
let mut batches: HashMap<String, Vec<Vec<RecordBatch>>> = HashMap::new();
let sql = Arc::new(meta);
if !results.is_empty() {
for resp in results {
file_count += resp.file_count;
scan_size += resp.scan_size;
// handle hits
let value = batches.entry("query".to_string()).or_default();
if !resp.hits.is_empty() {
let buf = Cursor::new(resp.hits);
for resp in results {
file_count += resp.file_count;
scan_size += resp.scan_size;
// handle hits
let value = batches.entry("query".to_string()).or_default();
if !resp.hits.is_empty() {
let buf = Cursor::new(resp.hits);
let reader = ipc::reader::FileReader::try_new(buf, None).unwrap();
log::info!(
"search_in_cluster: query num_batches: {:?}",
reader.num_batches()
);
let batch = reader.into_iter().map(Result::unwrap).collect::<Vec<_>>();
value.push(batch);
}
// handle aggs
for agg in resp.aggs {
let value = batches.entry(format!("agg_{}", agg.name)).or_default();
if !agg.hits.is_empty() {
let buf = Cursor::new(agg.hits);
let reader = ipc::reader::FileReader::try_new(buf, None).unwrap();
log::info!(
"search_in_cluster: query num_batches: {:?}",
"search_in_cluster: agg:{} num_batches: {:?}",
agg.name,
reader.num_batches()
);
let batch = reader.into_iter().map(Result::unwrap).collect::<Vec<_>>();
value.push(batch);
}
// handle aggs
for agg in resp.aggs {
let value = batches.entry(format!("agg_{}", agg.name)).or_default();
if !agg.hits.is_empty() {
let buf = Cursor::new(agg.hits);
let reader = ipc::reader::FileReader::try_new(buf, None).unwrap();
log::info!(
"search_in_cluster: agg:{} num_batches: {:?}",
agg.name,
reader.num_batches()
);
let batch = reader.into_iter().map(Result::unwrap).collect::<Vec<_>>();
value.push(batch);
}
}
}
// merge all batches
@@ -404,7 +397,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
}
//XXX span7.exit(); // drop span7
//XXX let _span8 = info_span!("srv:search:cluster:response").entered();
//XXX let _span8 = info_span!("service:search:cluster:response").entered();
// final result
let mut result = Response::new(sql.meta.offset, sql.meta.limit);
@@ -534,52 +527,6 @@ fn handle_metrics_response(sources: Vec<json::Value>) -> Vec<json::Value> {
new_sources
}
pub fn handle_datafusion_error(err: DataFusionError) -> Error {
let err = err.to_string();
if err.contains("Schema error: No field named") {
let pos = err.find("Schema error: No field named").unwrap();
return match get_key_from_error(&err, pos) {
Some(key) => Error::ErrorCode(ErrorCodes::SearchFieldNotFound(key)),
None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)),
};
}
if err.contains("parquet not found") {
return Error::ErrorCode(ErrorCodes::SearchParquetFileNotFound);
}
if err.contains("Invalid function ") {
let pos = err.find("Invalid function ").unwrap();
return match get_key_from_error(&err, pos) {
Some(key) => Error::ErrorCode(ErrorCodes::SearchFunctionNotDefined(key)),
None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)),
};
}
if err.contains("Incompatible data types") {
let pos = err.find("for field").unwrap();
let pos_start = err[pos..].find(' ').unwrap();
let pos_end = err[pos + pos_start + 1..].find('.').unwrap();
let field = err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string();
return Error::ErrorCode(ErrorCodes::SearchFieldHasNoCompatibleDataType(field));
}
Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err))
}
fn get_key_from_error(err: &str, pos: usize) -> Option<String> {
for ponct in ['\'', '"'] {
let pos_start = err[pos..].find(ponct);
if pos_start.is_none() {
continue;
}
let pos_start = pos_start.unwrap();
let pos_end = err[pos + pos_start + 1..].find(ponct);
if pos_end.is_none() {
continue;
}
let pos_end = pos_end.unwrap();
return Some(err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string());
}
None
}
struct MetadataMap<'a>(&'a mut tonic::metadata::MetadataMap);
impl<'a> opentelemetry::propagation::Injector for MetadataMap<'a> {
src/service/search/sql.rs
+8 -3
View File
@@ -27,7 +27,7 @@ use crate::infra::errors::{Error, ErrorCodes};
use crate::meta::sql::Sql as MetaSql;
use crate::meta::StreamType;
use crate::service::stream::get_stream_setting_fts_fields;
use crate::service::{db, file_list, logs};
use crate::service::{db, file_list, get_partition_key_query};
const SQL_DELIMITERS: [u8; 12] = [
b' ', b'*', b'(', b')', b'<', b'>', b',', b';', b'=', b'!', b'\r', b'\n',
@@ -587,6 +587,7 @@ impl Sql {
&self,
source: &str,
match_min_ts_only: bool,
is_wal: bool,
stream_type: StreamType,
) -> bool {
// match org_id & table
@@ -605,6 +606,10 @@ impl Sql {
return false;
}
if is_wal {
return true;
}
// check time range
let file_meta = file_list::get_file_meta(source).await.unwrap_or_default();
if file_meta.min_ts == 0 || file_meta.max_ts == 0 {
@@ -797,8 +802,8 @@ fn path_matches_by_partition_key<'a>(
partitions: impl IntoIterator<Item = (&'a str, &'a str)>,
) -> bool {
!partitions.into_iter().any(|(k, v)| {
let field = logs::get_partition_key_query(&format!("{k}="));
let value = logs::get_partition_key_query(&format!("{k}={v}"));
let field = get_partition_key_query(&format!("{k}="));
let value = get_partition_key_query(&format!("{k}={v}"));
find(path, &format!("/{field}")) && !find(path, &format!("/{value}/"))
})
}