rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2025-12-22 22:20:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into create-view

Browse Source
This commit is contained in:
Ruihang Xia
2024-04-22 21:08:22 +08:00
parent 7503992d61 54432df92f
commit 94409967be
260 changed files with 6025 additions and 4050 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/apidoc.yml vendored
View File

@@ -13,7 +13,7 @@ on:
name: Build API docs
env:
RUST_TOOLCHAIN: nightly-2023-12-19
RUST_TOOLCHAIN: nightly-2024-04-18
jobs:
apidoc:

2
.github/workflows/develop.yml vendored
View File

@@ -30,7 +30,7 @@ concurrency:
cancel-in-progress: true
env:
RUST_TOOLCHAIN: nightly-2023-12-19
RUST_TOOLCHAIN: nightly-2024-04-18
jobs:
check-typos-and-docs:

2
.github/workflows/nightly-ci.yml vendored
View File

@@ -12,7 +12,7 @@ concurrency:
cancel-in-progress: true
env:
RUST_TOOLCHAIN: nightly-2023-12-19
RUST_TOOLCHAIN: nightly-2024-04-18
jobs:
sqlness:

2
.github/workflows/release.yml vendored
View File

@@ -82,7 +82,7 @@ on:
# Use env variables to control all the release process.
env:
# The arguments of building greptime.
RUST_TOOLCHAIN: nightly-2023-12-19
RUST_TOOLCHAIN: nightly-2024-04-18
CARGO_PROFILE: nightly
# Controls whether to run tests, include unit-test, integration-test and sqlness.

2
CONTRIBUTING.md
View File

@@ -50,7 +50,7 @@ GreptimeDB uses the [Apache 2.0 license](https://github.com/GreptimeTeam/greptim
- To ensure that community is free and confident in its ability to use your contributions, please sign the Contributor License Agreement (CLA) which will be incorporated in the pull request process.
- Make sure all files have proper license header (running `docker run --rm -v $(pwd):/github/workspace ghcr.io/korandoru/hawkeye-native:v3 format` from the project root).
- Make sure all your codes are formatted and follow the [coding style](https://pingcap.github.io/style-guide/rust/).
- Make sure all your codes are formatted and follow the [coding style](https://pingcap.github.io/style-guide/rust/) and [style guide](http://github.com/greptimeTeam/docs/style-guide.md).
- Make sure all unit tests are passed (using `cargo test --workspace` or [nextest](https://nexte.st/index.html) `cargo nextest run`).
- Make sure all clippy warnings are fixed (you can check it locally by running `cargo clippy --workspace --all-targets -- -D warnings`).

2452
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

52
Cargo.toml
View File

@@ -70,16 +70,24 @@ license = "Apache-2.0"
clippy.print_stdout = "warn"
clippy.print_stderr = "warn"
clippy.implicit_clone = "warn"
clippy.readonly_write_lock = "allow"
rust.unknown_lints = "deny"
# Remove this after https://github.com/PyO3/pyo3/issues/4094
rust.non_local_definitions = "allow"
[workspace.dependencies]
# We turn off default-features for some dependencies here so the workspaces which inherit them can
# selectively turn them on if needed, since we can override default-features = true (from false)
# for the inherited dependency but cannot do the reverse (override from true to false).
#
# See for more detaiils: https://github.com/rust-lang/cargo/issues/11329
ahash = { version = "0.8", features = ["compile-time-rng"] }
aquamarine = "0.3"
arrow = { version = "47.0" }
arrow-array = "47.0"
arrow-flight = "47.0"
arrow-ipc = { version = "47.0", features = ["lz4"] }
arrow-schema = { version = "47.0", features = ["serde"] }
arrow = { version = "51.0.0", features = ["prettyprint"] }
arrow-array = { version = "51.0.0", default-features = false, features = ["chrono-tz"] }
arrow-flight = "51.0"
arrow-ipc = { version = "51.0.0", default-features = false, features = ["lz4"] }
arrow-schema = { version = "51.0", features = ["serde"] }
async-stream = "0.3"
async-trait = "0.1"
axum = { version = "0.6", features = ["headers"] }
@@ -91,20 +99,22 @@ bytes = { version = "1.5", features = ["serde"] }
chrono = { version = "0.4", features = ["serde"] }
clap = { version = "4.4", features = ["derive"] }
dashmap = "5.4"
datafusion = { git = "https://github.com/apache/arrow-datafusion.git", rev = "26e43acac3a96cec8dd4c8365f22dfb1a84306e9" }
datafusion-common = { git = "https://github.com/apache/arrow-datafusion.git", rev = "26e43acac3a96cec8dd4c8365f22dfb1a84306e9" }
datafusion-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "26e43acac3a96cec8dd4c8365f22dfb1a84306e9" }
datafusion-optimizer = { git = "https://github.com/apache/arrow-datafusion.git", rev = "26e43acac3a96cec8dd4c8365f22dfb1a84306e9" }
datafusion-physical-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "26e43acac3a96cec8dd4c8365f22dfb1a84306e9" }
datafusion-sql = { git = "https://github.com/apache/arrow-datafusion.git", rev = "26e43acac3a96cec8dd4c8365f22dfb1a84306e9" }
datafusion-substrait = { git = "https://github.com/apache/arrow-datafusion.git", rev = "26e43acac3a96cec8dd4c8365f22dfb1a84306e9" }
datafusion = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
datafusion-common = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
datafusion-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
datafusion-functions = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
datafusion-optimizer = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
datafusion-physical-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
datafusion-sql = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
datafusion-substrait = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
derive_builder = "0.12"
dotenv = "0.15"
etcd-client = "0.12"
# TODO(LFC): Wait for https://github.com/etcdv3/etcd-client/pull/76
etcd-client = { git = "https://github.com/MichaelScofield/etcd-client.git", rev = "4c371e9b3ea8e0a8ee2f9cbd7ded26e54a45df3b" }
fst = "0.4.7"
futures = "0.3"
futures-util = "0.3"
greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "04d78b6e025ceb518040fdd10858c2a9d9345820" }
greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "73ac0207ab71dfea48f30259ffdb611501b5ecb8" }
humantime = "2.1"
humantime-serde = "1.1"
itertools = "0.10"
@@ -115,12 +125,12 @@ moka = "0.12"
notify = "6.1"
num_cpus = "1.16"
once_cell = "1.18"
opentelemetry-proto = { git = "https://github.com/waynexia/opentelemetry-rust.git", rev = "33841b38dda79b15f2024952be5f32533325ca02", features = [
opentelemetry-proto = { version = "0.5", features = [
"gen-tonic",
"metrics",
"trace",
] }
parquet = "47.0"
parquet = { version = "51.0.0", default-features = false, features = ["arrow", "async", "object_store"] }
paste = "1.0"
pin-project = "1.0"
prometheus = { version = "0.13.3", features = ["process"] }
@@ -144,18 +154,18 @@ serde_with = "3"
smallvec = { version = "1", features = ["serde"] }
snafu = "0.7"
sysinfo = "0.30"
# on branch v0.38.x
sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "6a93567ae38d42be5c8d08b13c8ff4dde26502ef", features = [
# on branch v0.44.x
sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "c919990bf62ad38d2b0c0a3bc90b26ad919d51b0", features = [
"visitor",
] }
strum = { version = "0.25", features = ["derive"] }
tempfile = "3"
tokio = { version = "1.28", features = ["full"] }
tokio = { version = "1.36", features = ["full"] }
tokio-stream = { version = "0.1" }
tokio-util = { version = "0.7", features = ["io-util", "compat"] }
toml = "0.8.8"
tonic = { version = "0.10", features = ["tls"] }
uuid = { version = "1", features = ["serde", "v4", "fast-rng"] }
tonic = { version = "0.11", features = ["tls"] }
uuid = { version = "1.7", features = ["serde", "v4", "fast-rng"] }
zstd = "0.13"
## workspaces members

34
benchmarks/src/bin/nyc-taxi.rs
View File

@@ -215,37 +215,7 @@ fn build_values(column: &ArrayRef) -> (Values, ColumnDataType) {
ColumnDataType::String,
)
}
DataType::Null
| DataType::Boolean
| DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float16
| DataType::Float32
| DataType::Date32
| DataType::Date64
| DataType::Time32(_)
| DataType::Time64(_)
| DataType::Duration(_)
| DataType::Interval(_)
| DataType::Binary
| DataType::FixedSizeBinary(_)
| DataType::LargeBinary
| DataType::LargeUtf8
| DataType::List(_)
| DataType::FixedSizeList(_, _)
| DataType::LargeList(_)
| DataType::Struct(_)
| DataType::Union(_, _)
| DataType::Dictionary(_, _)
| DataType::Decimal128(_, _)
| DataType::Decimal256(_, _)
| DataType::RunEndEncoded(_, _)
| DataType::Map(_, _) => todo!(),
_ => unimplemented!(),
}
}
@@ -444,7 +414,7 @@ fn create_table_expr(table_name: &str) -> CreateTableExpr {
fn query_set(table_name: &str) -> HashMap<String, String> {
HashMap::from([
(
"count_all".to_string(),
"count_all".to_string(),
format!("SELECT COUNT(*) FROM {table_name};"),
),
(

46
docs/style-guide.md Normal file
View File

@@ -0,0 +1,46 @@
# GreptimeDB Style Guide
This style guide is intended to help contributors to GreptimeDB write code that is consistent with the rest of the codebase. It is a living document and will be updated as the codebase evolves.
It's mainly an complement to the [Rust Style Guide](https://pingcap.github.io/style-guide/rust/).
## Table of Contents
- Formatting
- Modules
- Comments
## Formatting
- Place all `mod` declaration before any `use`.
- Use `unimplemented!()` instead of `todo!()` for things that aren't likely to be implemented.
- Add an empty line before and after declaration blocks.
- Place comment before attributes (`#[]`) and derive (`#[derive]`).
## Modules
- Use the file with same name instead of `mod.rs` to define a module. E.g.:
```
.
├── cache
│ ├── cache_size.rs
│ └── write_cache.rs
└── cache.rs
```
## Comments
- Add comments for public functions and structs.
- Prefer document comment (`///`) over normal comment (`//`) for structs, fields, functions etc.
- Add link (`[]`) to struct, method, or any other reference. And make sure that link works.
## Error handling
- Define a custom error type for the module if needed.
- Prefer `with_context()` over `context()` when allocation is needed to construct an error.
- Use `error!()` or `warn!()` macros in the `common_telemetry` crate to log errors. E.g.:
```rust
error!(e; "Failed to do something");
```

2
rust-toolchain.toml
View File

@@ -1,2 +1,2 @@
[toolchain]
channel = "nightly-2023-12-19"
channel = "nightly-2024-04-18"

1
src/api/src/lib.rs
View File

@@ -21,6 +21,7 @@ pub mod prom_store {
}
}
pub mod region;
pub mod v1;
pub use greptime_proto;

42
src/api/src/region.rs Normal file
View File

@@ -0,0 +1,42 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use common_base::AffectedRows;
use greptime_proto::v1::region::RegionResponse as RegionResponseV1;
/// This result struct is derived from [RegionResponseV1]
#[derive(Debug)]
pub struct RegionResponse {
pub affected_rows: AffectedRows,
pub extension: HashMap<String, Vec<u8>>,
}
impl RegionResponse {
pub fn from_region_response(region_response: RegionResponseV1) -> Self {
Self {
affected_rows: region_response.affected_rows as _,
extension: region_response.extension,
}
}
/// Creates one response without extension
pub fn new(affected_rows: AffectedRows) -> Self {
Self {
affected_rows,
extension: Default::default(),
}
}
}

6
src/auth/src/tests.rs
View File

@@ -45,9 +45,9 @@ impl Default for MockUserProvider {
impl MockUserProvider {
pub fn set_authorization_info(&mut self, info: DatabaseAuthInfo) {
self.catalog = info.catalog.to_owned();
self.schema = info.schema.to_owned();
self.username = info.username.to_owned();
info.catalog.clone_into(&mut self.catalog);
info.schema.clone_into(&mut self.schema);
info.username.clone_into(&mut self.username);
}
}

32
src/catalog/src/information_schema/predicate.rs
View File

@@ -109,11 +109,7 @@ impl Predicate {
};
}
Predicate::Not(p) => {
let Some(b) = p.eval(row) else {
return None;
};
return Some(!b);
return Some(!p.eval(row)?);
}
}
@@ -125,13 +121,7 @@ impl Predicate {
fn from_expr(expr: DfExpr) -> Option<Predicate> {
match expr {
// NOT expr
DfExpr::Not(expr) => {
let Some(p) = Self::from_expr(*expr) else {
return None;
};
Some(Predicate::Not(Box::new(p)))
}
DfExpr::Not(expr) => Some(Predicate::Not(Box::new(Self::from_expr(*expr)?))),
// expr LIKE pattern
DfExpr::Like(Like {
negated,
@@ -178,25 +168,15 @@ impl Predicate {
}
// left AND right
(left, Operator::And, right) => {
let Some(left) = Self::from_expr(left) else {
return None;
};
let Some(right) = Self::from_expr(right) else {
return None;
};
let left = Self::from_expr(left)?;
let right = Self::from_expr(right)?;
Some(Predicate::And(Box::new(left), Box::new(right)))
}
// left OR right
(left, Operator::Or, right) => {
let Some(left) = Self::from_expr(left) else {
return None;
};
let Some(right) = Self::from_expr(right) else {
return None;
};
let left = Self::from_expr(left)?;
let right = Self::from_expr(right)?;
Some(Predicate::Or(Box::new(left), Box::new(right)))
}

1
src/catalog/src/kvbackend/client.rs
View File

@@ -17,7 +17,6 @@ use std::fmt::Debug;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::usize;
use common_error::ext::BoxedError;
use common_meta::cache_invalidator::KvCacheInvalidator;

58
src/catalog/src/table_source.rs
View File

@@ -49,10 +49,7 @@ impl DfTableSourceProvider {
}
}
pub fn resolve_table_ref<'a>(
&'a self,
table_ref: TableReference<'a>,
) -> Result<ResolvedTableReference<'a>> {
pub fn resolve_table_ref(&self, table_ref: TableReference) -> Result<ResolvedTableReference> {
if self.disallow_cross_catalog_query {
match &table_ref {
TableReference::Bare { .. } => (),
@@ -76,7 +73,7 @@ impl DfTableSourceProvider {
pub async fn resolve_table(
&mut self,
table_ref: TableReference<'_>,
table_ref: TableReference,
) -> Result<Arc<dyn TableSource>> {
let table_ref = self.resolve_table_ref(table_ref)?;
@@ -106,8 +103,6 @@ impl DfTableSourceProvider {
#[cfg(test)]
mod tests {
use std::borrow::Cow;
use session::context::QueryContext;
use super::*;
@@ -120,68 +115,37 @@ mod tests {
let table_provider =
DfTableSourceProvider::new(MemoryCatalogManager::with_default_setup(), true, query_ctx);
let table_ref = TableReference::Bare {
table: Cow::Borrowed("table_name"),
};
let table_ref = TableReference::bare("table_name");
let result = table_provider.resolve_table_ref(table_ref);
assert!(result.is_ok());
let table_ref = TableReference::Partial {
schema: Cow::Borrowed("public"),
table: Cow::Borrowed("table_name"),
};
let table_ref = TableReference::partial("public", "table_name");
let result = table_provider.resolve_table_ref(table_ref);
assert!(result.is_ok());
let table_ref = TableReference::Partial {
schema: Cow::Borrowed("wrong_schema"),
table: Cow::Borrowed("table_name"),
};
let table_ref = TableReference::partial("wrong_schema", "table_name");
let result = table_provider.resolve_table_ref(table_ref);
assert!(result.is_ok());
let table_ref = TableReference::Full {
catalog: Cow::Borrowed("greptime"),
schema: Cow::Borrowed("public"),
table: Cow::Borrowed("table_name"),
};
let table_ref = TableReference::full("greptime", "public", "table_name");
let result = table_provider.resolve_table_ref(table_ref);
assert!(result.is_ok());
let table_ref = TableReference::Full {
catalog: Cow::Borrowed("wrong_catalog"),
schema: Cow::Borrowed("public"),
table: Cow::Borrowed("table_name"),
};
let table_ref = TableReference::full("wrong_catalog", "public", "table_name");
let result = table_provider.resolve_table_ref(table_ref);
assert!(result.is_err());
let table_ref = TableReference::Partial {
schema: Cow::Borrowed("information_schema"),
table: Cow::Borrowed("columns"),
};
let table_ref = TableReference::partial("information_schema", "columns");
let result = table_provider.resolve_table_ref(table_ref);
assert!(result.is_ok());
let table_ref = TableReference::Full {
catalog: Cow::Borrowed("greptime"),
schema: Cow::Borrowed("information_schema"),
table: Cow::Borrowed("columns"),
};
let table_ref = TableReference::full("greptime", "information_schema", "columns");
assert!(table_provider.resolve_table_ref(table_ref).is_ok());
let table_ref = TableReference::Full {
catalog: Cow::Borrowed("dummy"),
schema: Cow::Borrowed("information_schema"),
table: Cow::Borrowed("columns"),
};
let table_ref = TableReference::full("dummy", "information_schema", "columns");
assert!(table_provider.resolve_table_ref(table_ref).is_err());
let table_ref = TableReference::Full {
catalog: Cow::Borrowed("greptime"),
schema: Cow::Borrowed("greptime_private"),
table: Cow::Borrowed("columns"),
};
let table_ref = TableReference::full("greptime", "greptime_private", "columns");
assert!(table_provider.resolve_table_ref(table_ref).is_ok());
}
}

11
src/client/src/region.rs
View File

@@ -14,6 +14,7 @@
use std::sync::Arc;
use api::region::RegionResponse;
use api::v1::region::{QueryRequest, RegionRequest};
use api::v1::ResponseHeader;
use arc_swap::ArcSwapOption;
@@ -23,7 +24,7 @@ use async_trait::async_trait;
use common_error::ext::{BoxedError, ErrorExt};
use common_error::status_code::StatusCode;
use common_grpc::flight::{FlightDecoder, FlightMessage};
use common_meta::datanode_manager::{Datanode, HandleResponse};
use common_meta::datanode_manager::Datanode;
use common_meta::error::{self as meta_error, Result as MetaResult};
use common_recordbatch::error::ExternalSnafu;
use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream};
@@ -46,7 +47,7 @@ pub struct RegionRequester {
#[async_trait]
impl Datanode for RegionRequester {
async fn handle(&self, request: RegionRequest) -> MetaResult<HandleResponse> {
async fn handle(&self, request: RegionRequest) -> MetaResult<RegionResponse> {
self.handle_inner(request).await.map_err(|err| {
if err.should_retry() {
meta_error::Error::RetryLater {
@@ -165,7 +166,7 @@ impl RegionRequester {
Ok(Box::pin(record_batch_stream))
}
async fn handle_inner(&self, request: RegionRequest) -> Result<HandleResponse> {
async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponse> {
let request_type = request
.body
.as_ref()
@@ -194,10 +195,10 @@ impl RegionRequester {
check_response_header(&response.header)?;
Ok(HandleResponse::from_region_response(response))
Ok(RegionResponse::from_region_response(response))
}
pub async fn handle(&self, request: RegionRequest) -> Result<HandleResponse> {
pub async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
self.handle_inner(request).await
}
}

4
src/cmd/src/cli.rs
View File

@@ -84,10 +84,10 @@ impl Command {
let mut logging_opts = LoggingOptions::default();
if let Some(dir) = &cli_options.log_dir {
logging_opts.dir = dir.clone();
logging_opts.dir.clone_from(dir);
}
logging_opts.level = cli_options.log_level.clone();
logging_opts.level.clone_from(&cli_options.log_level);
Ok(Options::Cli(Box::new(logging_opts)))
}

15
src/cmd/src/datanode.rs
View File

@@ -139,19 +139,19 @@ impl StartCommand {
)?;
if let Some(dir) = &cli_options.log_dir {
opts.logging.dir = dir.clone();
opts.logging.dir.clone_from(dir);
}
if cli_options.log_level.is_some() {
opts.logging.level = cli_options.log_level.clone();
opts.logging.level.clone_from(&cli_options.log_level);
}
if let Some(addr) = &self.rpc_addr {
opts.rpc_addr = addr.clone();
opts.rpc_addr.clone_from(addr);
}
if self.rpc_hostname.is_some() {
opts.rpc_hostname = self.rpc_hostname.clone();
opts.rpc_hostname.clone_from(&self.rpc_hostname);
}
if let Some(node_id) = self.node_id {
@@ -161,7 +161,8 @@ impl StartCommand {
if let Some(metasrv_addrs) = &self.metasrv_addr {
opts.meta_client
.get_or_insert_with(MetaClientOptions::default)
.metasrv_addrs = metasrv_addrs.clone();
.metasrv_addrs
.clone_from(metasrv_addrs);
opts.mode = Mode::Distributed;
}
@@ -173,7 +174,7 @@ impl StartCommand {
}
if let Some(data_home) = &self.data_home {
opts.storage.data_home = data_home.clone();
opts.storage.data_home.clone_from(data_home);
}
// `wal_dir` only affects raft-engine config.
@@ -191,7 +192,7 @@ impl StartCommand {
}
if let Some(http_addr) = &self.http_addr {
opts.http.addr = http_addr.clone();
opts.http.addr.clone_from(http_addr);
}
if let Some(http_timeout) = self.http_timeout {

19
src/cmd/src/frontend.rs
View File

@@ -157,11 +157,11 @@ impl StartCommand {
)?;
if let Some(dir) = &cli_options.log_dir {
opts.logging.dir = dir.clone();
opts.logging.dir.clone_from(dir);
}
if cli_options.log_level.is_some() {
opts.logging.level = cli_options.log_level.clone();
opts.logging.level.clone_from(&cli_options.log_level);
}
let tls_opts = TlsOption::new(
@@ -171,7 +171,7 @@ impl StartCommand {
);
if let Some(addr) = &self.http_addr {
opts.http.addr = addr.clone()
opts.http.addr.clone_from(addr);
}
if let Some(http_timeout) = self.http_timeout {
@@ -183,24 +183,24 @@ impl StartCommand {
}
if let Some(addr) = &self.rpc_addr {
opts.grpc.addr = addr.clone()
opts.grpc.addr.clone_from(addr);
}
if let Some(addr) = &self.mysql_addr {
opts.mysql.enable = true;
opts.mysql.addr = addr.clone();
opts.mysql.addr.clone_from(addr);
opts.mysql.tls = tls_opts.clone();
}
if let Some(addr) = &self.postgres_addr {
opts.postgres.enable = true;
opts.postgres.addr = addr.clone();
opts.postgres.addr.clone_from(addr);
opts.postgres.tls = tls_opts;
}
if let Some(addr) = &self.opentsdb_addr {
opts.opentsdb.enable = true;
opts.opentsdb.addr = addr.clone();
opts.opentsdb.addr.clone_from(addr);
}
if let Some(enable) = self.influxdb_enable {
@@ -210,11 +210,12 @@ impl StartCommand {
if let Some(metasrv_addrs) = &self.metasrv_addr {
opts.meta_client
.get_or_insert_with(MetaClientOptions::default)
.metasrv_addrs = metasrv_addrs.clone();
.metasrv_addrs
.clone_from(metasrv_addrs);
opts.mode = Mode::Distributed;
}
opts.user_provider = self.user_provider.clone();
opts.user_provider.clone_from(&self.user_provider);
Ok(Options::Frontend(Box::new(opts)))
}

16
src/cmd/src/metasrv.rs
View File

@@ -134,23 +134,23 @@ impl StartCommand {
)?;
if let Some(dir) = &cli_options.log_dir {
opts.logging.dir = dir.clone();
opts.logging.dir.clone_from(dir);
}
if cli_options.log_level.is_some() {
opts.logging.level = cli_options.log_level.clone();
opts.logging.level.clone_from(&cli_options.log_level);
}
if let Some(addr) = &self.bind_addr {
opts.bind_addr = addr.clone();
opts.bind_addr.clone_from(addr);
}
if let Some(addr) = &self.server_addr {
opts.server_addr = addr.clone();
opts.server_addr.clone_from(addr);
}
if let Some(addr) = &self.store_addr {
opts.store_addr = addr.clone();
opts.store_addr.clone_from(addr);
}
if let Some(selector_type) = &self.selector {
@@ -168,7 +168,7 @@ impl StartCommand {
}
if let Some(http_addr) = &self.http_addr {
opts.http.addr = http_addr.clone();
opts.http.addr.clone_from(http_addr);
}
if let Some(http_timeout) = self.http_timeout {
@@ -176,11 +176,11 @@ impl StartCommand {
}
if let Some(data_home) = &self.data_home {
opts.data_home = data_home.clone();
opts.data_home.clone_from(data_home);
}
if !self.store_key_prefix.is_empty() {
opts.store_key_prefix = self.store_key_prefix.clone()
opts.store_key_prefix.clone_from(&self.store_key_prefix)
}
if let Some(max_txn_ops) = self.max_txn_ops {

18
src/cmd/src/standalone.rs
View File

@@ -293,11 +293,11 @@ impl StartCommand {
opts.mode = Mode::Standalone;
if let Some(dir) = &cli_options.log_dir {
opts.logging.dir = dir.clone();
opts.logging.dir.clone_from(dir);
}
if cli_options.log_level.is_some() {
opts.logging.level = cli_options.log_level.clone();
opts.logging.level.clone_from(&cli_options.log_level);
}
let tls_opts = TlsOption::new(
@@ -307,11 +307,11 @@ impl StartCommand {
);
if let Some(addr) = &self.http_addr {
opts.http.addr = addr.clone()
opts.http.addr.clone_from(addr);
}
if let Some(data_home) = &self.data_home {
opts.storage.data_home = data_home.clone();
opts.storage.data_home.clone_from(data_home);
}
if let Some(addr) = &self.rpc_addr {
@@ -325,31 +325,31 @@ impl StartCommand {
}
.fail();
}
opts.grpc.addr = addr.clone()
opts.grpc.addr.clone_from(addr)
}
if let Some(addr) = &self.mysql_addr {
opts.mysql.enable = true;
opts.mysql.addr = addr.clone();
opts.mysql.addr.clone_from(addr);
opts.mysql.tls = tls_opts.clone();
}
if let Some(addr) = &self.postgres_addr {
opts.postgres.enable = true;
opts.postgres.addr = addr.clone();
opts.postgres.addr.clone_from(addr);
opts.postgres.tls = tls_opts;
}
if let Some(addr) = &self.opentsdb_addr {
opts.opentsdb.enable = true;
opts.opentsdb.addr = addr.clone();
opts.opentsdb.addr.clone_from(addr);
}
if self.influxdb_enable {
opts.influxdb.enable = self.influxdb_enable;
}
opts.user_provider = self.user_provider.clone();
opts.user_provider.clone_from(&self.user_provider);
let metadata_store = opts.metadata_store.clone();
let procedure = opts.procedure.clone();

2
src/common/datasource/Cargo.toml
View File

@@ -30,7 +30,7 @@ derive_builder.workspace = true
futures.workspace = true
lazy_static.workspace = true
object-store.workspace = true
orc-rust = "0.2"
orc-rust = { git = "https://github.com/MichaelScofield/orc-rs.git", rev = "17347f5f084ac937863317df882218055c4ea8c1" }
parquet.workspace = true
paste = "1.0"
regex = "1.7"

2
src/common/datasource/src/file_format/csv.rs
View File

@@ -117,7 +117,7 @@ impl CsvConfig {
let mut builder = csv::ReaderBuilder::new(self.file_schema.clone())
.with_delimiter(self.delimiter)
.with_batch_size(self.batch_size)
.has_header(self.has_header);
.with_header(self.has_header);
if let Some(proj) = &self.file_projection {
builder = builder.with_projection(proj.clone());

3
src/common/datasource/src/file_format/tests.rs
View File

@@ -19,6 +19,7 @@ use std::vec;
use common_test_util::find_workspace_path;
use datafusion::assert_batches_eq;
use datafusion::config::TableParquetOptions;
use datafusion::datasource::physical_plan::{FileOpener, FileScanConfig, FileStream, ParquetExec};
use datafusion::execution::context::TaskContext;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
@@ -166,7 +167,7 @@ async fn test_parquet_exec() {
.to_string();
let base_config = scan_config(schema.clone(), None, path);
let exec = ParquetExec::new(base_config, None, None)
let exec = ParquetExec::new(base_config, None, None, TableParquetOptions::default())
.with_parquet_file_reader_factory(Arc::new(DefaultParquetFileReaderFactory::new(store)));
let ctx = SessionContext::new();

6
src/common/datasource/src/test_util.rs
View File

@@ -16,6 +16,7 @@ use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use common_test_util::temp_dir::{create_temp_dir, TempDir};
use datafusion::common::Statistics;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{FileScanConfig, FileStream};
@@ -72,17 +73,16 @@ pub fn test_basic_schema() -> SchemaRef {
pub fn scan_config(file_schema: SchemaRef, limit: Option<usize>, filename: &str) -> FileScanConfig {
// object_store only recognize the Unix style path, so make it happy.
let filename = &filename.replace('\\', "/");
let statistics = Statistics::new_unknown(file_schema.as_ref());
FileScanConfig {
object_store_url: ObjectStoreUrl::parse("empty://").unwrap(), // won't be used
file_schema,
file_groups: vec![vec![PartitionedFile::new(filename.to_string(), 10)]],
statistics: Default::default(),
statistics,
projection: None,
limit,
table_partition_cols: vec![],
output_ordering: vec![],
infinite_source: false,
}
}

22
src/common/function/src/scalars/aggregate/diff.rs
View File

@@ -56,7 +56,7 @@ where
.map(|&n| n.into())
.collect::<Vec<Value>>();
Ok(vec![Value::List(ListValue::new(
Some(Box::new(nums)),
nums,
I::LogicalType::build_data_type(),
))])
}
@@ -120,10 +120,7 @@ where
O::from_native(native).into()
})
.collect::<Vec<Value>>();
let diff = Value::List(ListValue::new(
Some(Box::new(diff)),
O::LogicalType::build_data_type(),
));
let diff = Value::List(ListValue::new(diff, O::LogicalType::build_data_type()));
Ok(diff)
}
}
@@ -218,10 +215,7 @@ mod test {
let values = vec![Value::from(2_i64), Value::from(1_i64)];
diff.update_batch(&v).unwrap();
assert_eq!(
Value::List(ListValue::new(
Some(Box::new(values)),
ConcreteDataType::int64_datatype()
)),
Value::List(ListValue::new(values, ConcreteDataType::int64_datatype())),
diff.evaluate().unwrap()
);
@@ -236,10 +230,7 @@ mod test {
let values = vec![Value::from(5_i64), Value::from(1_i64)];
diff.update_batch(&v).unwrap();
assert_eq!(
Value::List(ListValue::new(
Some(Box::new(values)),
ConcreteDataType::int64_datatype()
)),
Value::List(ListValue::new(values, ConcreteDataType::int64_datatype())),
diff.evaluate().unwrap()
);
@@ -252,10 +243,7 @@ mod test {
let values = vec![Value::from(0_i64), Value::from(0_i64), Value::from(0_i64)];
diff.update_batch(&v).unwrap();
assert_eq!(
Value::List(ListValue::new(
Some(Box::new(values)),
ConcreteDataType::int64_datatype()
)),
Value::List(ListValue::new(values, ConcreteDataType::int64_datatype())),
diff.evaluate().unwrap()
);
}

5
src/common/function/src/scalars/aggregate/percentile.rs
View File

@@ -104,10 +104,7 @@ where
.map(|&n| n.into())
.collect::<Vec<Value>>();
Ok(vec![
Value::List(ListValue::new(
Some(Box::new(nums)),
T::LogicalType::build_data_type(),
)),
Value::List(ListValue::new(nums, T::LogicalType::build_data_type())),
self.p.into(),
])
}

5
src/common/function/src/scalars/aggregate/polyval.rs
View File

@@ -72,10 +72,7 @@ where
.map(|&n| n.into())
.collect::<Vec<Value>>();
Ok(vec![
Value::List(ListValue::new(
Some(Box::new(nums)),
T::LogicalType::build_data_type(),
)),
Value::List(ListValue::new(nums, T::LogicalType::build_data_type())),
self.x.into(),
])
}

5
src/common/function/src/scalars/aggregate/scipy_stats_norm_cdf.rs
View File

@@ -56,10 +56,7 @@ where
.map(|&x| x.into())
.collect::<Vec<Value>>();
Ok(vec![
Value::List(ListValue::new(
Some(Box::new(nums)),
T::LogicalType::build_data_type(),
)),
Value::List(ListValue::new(nums, T::LogicalType::build_data_type())),
self.x.into(),
])
}

5
src/common/function/src/scalars/aggregate/scipy_stats_norm_pdf.rs
View File

@@ -56,10 +56,7 @@ where
.map(|&x| x.into())
.collect::<Vec<Value>>();
Ok(vec![
Value::List(ListValue::new(
Some(Box::new(nums)),
T::LogicalType::build_data_type(),
)),
Value::List(ListValue::new(nums, T::LogicalType::build_data_type())),
self.x.into(),
])
}

2
src/common/function/src/scalars/math.rs
View File

@@ -77,7 +77,7 @@ impl Function for RangeFunction {
/// `range_fn` will never been used. As long as a legal signature is returned, the specific content of the signature does not matter.
/// In fact, the arguments loaded by `range_fn` are very complicated, and it is difficult to use `Signature` to describe
fn signature(&self) -> Signature {
Signature::any(0, Volatility::Immutable)
Signature::variadic_any(Volatility::Immutable)
}
fn eval(&self, _func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {

14
src/common/macro/src/range_fn.rs
View File

@@ -119,15 +119,17 @@ fn build_struct(
}
pub fn scalar_udf() -> ScalarUDF {
ScalarUDF {
name: Self::name().to_string(),
signature: Signature::new(
// TODO(LFC): Use the new Datafusion UDF impl.
#[allow(deprecated)]
ScalarUDF::new(
Self::name(),
&Signature::new(
TypeSignature::Exact(Self::input_type()),
Volatility::Immutable,
),
return_type: Arc::new(|_| Ok(Arc::new(Self::return_type()))),
fun: Arc::new(Self::calc),
}
&(Arc::new(|_: &_| Ok(Arc::new(Self::return_type()))) as _),
&(Arc::new(Self::calc) as _),
)
}
fn input_type() -> Vec<DataType> {

5
src/common/meta/src/cache_invalidator.rs
View File

@@ -18,6 +18,7 @@ use tokio::sync::RwLock;
use crate::error::Result;
use crate::instruction::CacheIdent;
use crate::key::schema_name::SchemaNameKey;
use crate::key::table_info::TableInfoKey;
use crate::key::table_name::TableNameKey;
use crate::key::table_route::TableRouteKey;
@@ -107,6 +108,10 @@ where
let key: TableNameKey = (&table_name).into();
self.invalidate_key(&key.as_raw_key()).await
}
CacheIdent::SchemaName(schema_name) => {
let key: SchemaNameKey = (&schema_name).into();
self.invalidate_key(&key.as_raw_key()).await;
}
}
}
Ok(())

30
src/common/meta/src/datanode_manager.rs
View File

@@ -12,10 +12,10 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::sync::Arc;
use api::v1::region::{QueryRequest, RegionRequest, RegionResponse};
use api::region::RegionResponse;
use api::v1::region::{QueryRequest, RegionRequest};
pub use common_base::AffectedRows;
use common_recordbatch::SendableRecordBatchStream;
@@ -26,7 +26,7 @@ use crate::peer::Peer;
#[async_trait::async_trait]
pub trait Datanode: Send + Sync {
/// Handles DML, and DDL requests.
async fn handle(&self, request: RegionRequest) -> Result<HandleResponse>;
async fn handle(&self, request: RegionRequest) -> Result<RegionResponse>;
/// Handles query requests
async fn handle_query(&self, request: QueryRequest) -> Result<SendableRecordBatchStream>;
@@ -42,27 +42,3 @@ pub trait DatanodeManager: Send + Sync {
}
pub type DatanodeManagerRef = Arc<dyn DatanodeManager>;
/// This result struct is derived from [RegionResponse]
#[derive(Debug)]
pub struct HandleResponse {
pub affected_rows: AffectedRows,
pub extension: HashMap<String, Vec<u8>>,
}
impl HandleResponse {
pub fn from_region_response(region_response: RegionResponse) -> Self {
Self {
affected_rows: region_response.affected_rows as _,
extension: region_response.extension,
}
}
/// Creates one response without extension
pub fn new(affected_rows: AffectedRows) -> Self {
Self {
affected_rows,
extension: Default::default(),
}
}
}

8
src/common/meta/src/ddl/drop_database/executor.rs
View File

@@ -122,12 +122,12 @@ impl State for DropDatabaseExecutor {
mod tests {
use std::sync::Arc;
use api::region::RegionResponse;
use api::v1::region::{QueryRequest, RegionRequest};
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use common_error::ext::BoxedError;
use common_recordbatch::SendableRecordBatchStream;
use crate::datanode_manager::HandleResponse;
use crate::ddl::drop_database::cursor::DropDatabaseCursor;
use crate::ddl::drop_database::executor::DropDatabaseExecutor;
use crate::ddl::drop_database::{DropDatabaseContext, DropTableTarget, State};
@@ -144,8 +144,8 @@ mod tests {
#[async_trait::async_trait]
impl MockDatanodeHandler for NaiveDatanodeHandler {
async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> Result<HandleResponse> {
Ok(HandleResponse::new(0))
async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> Result<RegionResponse> {
Ok(RegionResponse::new(0))
}
async fn handle_query(
@@ -291,7 +291,7 @@ mod tests {
#[async_trait::async_trait]
impl MockDatanodeHandler for RetryErrorDatanodeHandler {
async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> Result<RegionResponse> {
Err(Error::RetryLater {
source: BoxedError::new(
error::UnexpectedSnafu {

74
src/common/meta/src/ddl/drop_database/metadata.rs
View File

@@ -18,10 +18,12 @@ use common_procedure::Status;
use serde::{Deserialize, Serialize};
use super::end::DropDatabaseEnd;
use crate::cache_invalidator::Context;
use crate::ddl::drop_database::{DropDatabaseContext, State};
use crate::ddl::DdlContext;
use crate::error::Result;
use crate::key::schema_name::SchemaNameKey;
use crate::instruction::CacheIdent;
use crate::key::schema_name::{SchemaName, SchemaNameKey};
#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct DropDatabaseRemoveMetadata;
@@ -40,7 +42,53 @@ impl State for DropDatabaseRemoveMetadata {
.delete(SchemaNameKey::new(&ctx.catalog, &ctx.schema))
.await?;
return Ok((Box::new(DropDatabaseEnd), Status::done()));
return Ok((Box::new(DropMetadataBroadcast), Status::executing(true)));
}
fn as_any(&self) -> &dyn Any {
self
}
}
#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct DropMetadataBroadcast;
impl DropMetadataBroadcast {
/// Invalidates frontend caches
async fn invalidate_schema_cache(
&self,
ddl_ctx: &DdlContext,
db_ctx: &mut DropDatabaseContext,
) -> Result<()> {
let cache_invalidator = &ddl_ctx.cache_invalidator;
let ctx = Context {
subject: Some("Invalidate schema cache by dropping database".to_string()),
};
cache_invalidator
.invalidate(
&ctx,
vec![CacheIdent::SchemaName(SchemaName {
catalog_name: db_ctx.catalog.clone(),
schema_name: db_ctx.schema.clone(),
})],
)
.await?;
Ok(())
}
}
#[async_trait::async_trait]
#[typetag::serde]
impl State for DropMetadataBroadcast {
async fn next(
&mut self,
ddl_ctx: &DdlContext,
ctx: &mut DropDatabaseContext,
) -> Result<(Box<dyn State>, Status)> {
self.invalidate_schema_cache(ddl_ctx, ctx).await?;
Ok((Box::new(DropDatabaseEnd), Status::done()))
}
fn as_any(&self) -> &dyn Any {
@@ -53,7 +101,7 @@ mod tests {
use std::sync::Arc;
use crate::ddl::drop_database::end::DropDatabaseEnd;
use crate::ddl::drop_database::metadata::DropDatabaseRemoveMetadata;
use crate::ddl::drop_database::metadata::{DropDatabaseRemoveMetadata, DropMetadataBroadcast};
use crate::ddl::drop_database::{DropDatabaseContext, State};
use crate::key::schema_name::SchemaNameKey;
use crate::test_util::{new_ddl_context, MockDatanodeManager};
@@ -76,14 +124,23 @@ mod tests {
tables: None,
};
let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap();
state.as_any().downcast_ref::<DropDatabaseEnd>().unwrap();
assert!(status.is_done());
state
.as_any()
.downcast_ref::<DropMetadataBroadcast>()
.unwrap();
assert!(!status.is_done());
assert!(!ddl_context
.table_metadata_manager
.schema_manager()
.exists(SchemaNameKey::new("foo", "bar"))
.await
.unwrap());
let mut state = DropMetadataBroadcast;
let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap();
state.as_any().downcast_ref::<DropDatabaseEnd>().unwrap();
assert!(status.is_done());
// Schema not exists
let mut state = DropDatabaseRemoveMetadata;
let mut ctx = DropDatabaseContext {
@@ -93,7 +150,10 @@ mod tests {
tables: None,
};
let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap();
state.as_any().downcast_ref::<DropDatabaseEnd>().unwrap();
assert!(status.is_done());
state
.as_any()
.downcast_ref::<DropMetadataBroadcast>()
.unwrap();
assert!(!status.is_done());
}
}

26
src/common/meta/src/ddl/drop_table/executor.rs
View File

@@ -106,19 +106,6 @@ impl DropTableExecutor {
ctx: &DdlContext,
table_route_value: &TableRouteValue,
) -> Result<()> {
let table_name_key = TableNameKey::new(
&self.table.catalog_name,
&self.table.schema_name,
&self.table.table_name,
);
if !ctx
.table_metadata_manager
.table_name_manager()
.exists(table_name_key)
.await?
{
return Ok(());
}
ctx.table_metadata_manager
.delete_table_metadata(self.table_id, &self.table, table_route_value)
.await
@@ -152,19 +139,6 @@ impl DropTableExecutor {
ctx: &DdlContext,
table_route_value: &TableRouteValue,
) -> Result<()> {
let table_name_key = TableNameKey::new(
&self.table.catalog_name,
&self.table.schema_name,
&self.table.table_name,
);
if ctx
.table_metadata_manager
.table_name_manager()
.exists(table_name_key)
.await?
{
return Ok(());
}
ctx.table_metadata_manager
.restore_table_metadata(self.table_id, &self.table, table_route_value)
.await

18
src/common/meta/src/ddl/test_util/datanode_handler.rs
View File

@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use api::region::RegionResponse;
use api::v1::region::{QueryRequest, RegionRequest};
use common_error::ext::{BoxedError, ErrorExt, StackError};
use common_error::status_code::StatusCode;
@@ -20,14 +21,13 @@ use common_telemetry::debug;
use snafu::{ResultExt, Snafu};
use tokio::sync::mpsc;
use crate::datanode_manager::HandleResponse;
use crate::error::{self, Error, Result};
use crate::peer::Peer;
use crate::test_util::MockDatanodeHandler;
#[async_trait::async_trait]
impl MockDatanodeHandler for () {
async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> Result<RegionResponse> {
unreachable!()
}
@@ -45,10 +45,10 @@ pub struct DatanodeWatcher(pub mpsc::Sender<(Peer, RegionRequest)>);
#[async_trait::async_trait]
impl MockDatanodeHandler for DatanodeWatcher {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<RegionResponse> {
debug!("Returning Ok(0) for request: {request:?}, peer: {peer:?}");
self.0.send((peer.clone(), request)).await.unwrap();
Ok(HandleResponse::new(0))
Ok(RegionResponse::new(0))
}
async fn handle_query(
@@ -65,7 +65,7 @@ pub struct RetryErrorDatanodeHandler;
#[async_trait::async_trait]
impl MockDatanodeHandler for RetryErrorDatanodeHandler {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<RegionResponse> {
debug!("Returning retry later for request: {request:?}, peer: {peer:?}");
Err(Error::RetryLater {
source: BoxedError::new(
@@ -91,7 +91,7 @@ pub struct UnexpectedErrorDatanodeHandler;
#[async_trait::async_trait]
impl MockDatanodeHandler for UnexpectedErrorDatanodeHandler {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<RegionResponse> {
debug!("Returning mock error for request: {request:?}, peer: {peer:?}");
error::UnexpectedSnafu {
err_msg: "mock error",
@@ -135,7 +135,7 @@ impl ErrorExt for MockRequestOutdatedError {
#[async_trait::async_trait]
impl MockDatanodeHandler for RequestOutdatedErrorDatanodeHandler {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<RegionResponse> {
debug!("Returning mock error for request: {request:?}, peer: {peer:?}");
Err(BoxedError::new(MockRequestOutdatedError)).context(error::ExternalSnafu)
}
@@ -154,9 +154,9 @@ pub struct NaiveDatanodeHandler;
#[async_trait::async_trait]
impl MockDatanodeHandler for NaiveDatanodeHandler {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<RegionResponse> {
debug!("Returning Ok(0) for request: {request:?}, peer: {peer:?}");
Ok(HandleResponse::new(0))
Ok(RegionResponse::new(0))
}
async fn handle_query(

6
src/common/meta/src/error.rs
View File

@@ -421,8 +421,8 @@ pub enum Error {
#[snafu(display("Invalid role: {}", role))]
InvalidRole { role: i32, location: Location },
#[snafu(display("Atomic key changed: {err_msg}"))]
CasKeyChanged { err_msg: String, location: Location },
#[snafu(display("Failed to move values: {err_msg}"))]
MoveValues { err_msg: String, location: Location },
#[snafu(display("Failed to parse {} from utf8", name))]
FromUtf8 {
@@ -444,7 +444,7 @@ impl ErrorExt for Error {
| EtcdFailed { .. }
| EtcdTxnFailed { .. }
| ConnectEtcd { .. }
| CasKeyChanged { .. } => StatusCode::Internal,
| MoveValues { .. } => StatusCode::Internal,
SerdeJson { .. }
| ParseOption { .. }

2
src/common/meta/src/instruction.rs
View File

@@ -21,6 +21,7 @@ use store_api::storage::{RegionId, RegionNumber};
use strum::Display;
use table::metadata::TableId;
use crate::key::schema_name::SchemaName;
use crate::table_name::TableName;
use crate::{ClusterId, DatanodeId};
@@ -156,6 +157,7 @@ pub struct UpgradeRegion {
pub enum CacheIdent {
TableId(TableId),
TableName(TableName),
SchemaName(SchemaName),
}
#[derive(Debug, Clone, Serialize, Deserialize, Display, PartialEq)]

109
src/common/meta/src/key.rs
View File

@@ -88,13 +88,13 @@ use self::schema_name::{SchemaManager, SchemaNameKey, SchemaNameValue};
use self::table_route::{TableRouteManager, TableRouteValue};
use self::tombstone::TombstoneManager;
use crate::ddl::utils::region_storage_path;
use crate::error::{self, Result, SerdeJsonSnafu, UnexpectedSnafu};
use crate::error::{self, Result, SerdeJsonSnafu};
use crate::key::table_route::TableRouteKey;
use crate::key::tombstone::Key;
use crate::key::txn_helper::TxnOpGetResponseSet;
use crate::kv_backend::txn::{Txn, TxnOp, TxnOpResponse};
use crate::kv_backend::txn::{Txn, TxnOp};
use crate::kv_backend::KvBackendRef;
use crate::rpc::router::{region_distribution, RegionRoute, RegionStatus};
use crate::rpc::store::BatchDeleteRequest;
use crate::table_name::TableName;
use crate::DatanodeId;
@@ -466,17 +466,18 @@ impl TableMetadataManager {
txn = txn.merge(create_datanode_table_txn);
}
let r = self.kv_backend.txn(txn).await?;
let mut r = self.kv_backend.txn(txn).await?;
// Checks whether metadata was already created.
if !r.succeeded {
let remote_table_info = on_create_table_info_failure(&r.responses)?
let mut set = TxnOpGetResponseSet::from(&mut r.responses);
let remote_table_info = on_create_table_info_failure(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table info during the create table metadata",
})?
.into_inner();
let remote_table_route = on_create_table_route_failure(&r.responses)?
let remote_table_route = on_create_table_route_failure(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table route during the create table metadata",
})?
@@ -505,8 +506,8 @@ impl TableMetadataManager {
let mut txns = Vec::with_capacity(3 * len);
struct OnFailure<F1, R1, F2, R2>
where
F1: FnOnce(&Vec<TxnOpResponse>) -> R1,
F2: FnOnce(&Vec<TxnOpResponse>) -> R2,
F1: FnOnce(&mut TxnOpGetResponseSet) -> R1,
F2: FnOnce(&mut TxnOpGetResponseSet) -> R2,
{
table_info_value: TableInfoValue,
on_create_table_info_failure: F1,
@@ -551,18 +552,19 @@ impl TableMetadataManager {
}
let txn = Txn::merge_all(txns);
let r = self.kv_backend.txn(txn).await?;
let mut r = self.kv_backend.txn(txn).await?;
// Checks whether metadata was already created.
if !r.succeeded {
let mut set = TxnOpGetResponseSet::from(&mut r.responses);
for on_failure in on_failures {
let remote_table_info = (on_failure.on_create_table_info_failure)(&r.responses)?
let remote_table_info = (on_failure.on_create_table_info_failure)(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table info during the create table metadata",
})?
.into_inner();
let remote_table_route = (on_failure.on_create_table_route_failure)(&r.responses)?
let remote_table_route = (on_failure.on_create_table_route_failure)(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table route during the create table metadata",
})?
@@ -582,7 +584,7 @@ impl TableMetadataManager {
table_id: TableId,
table_name: &TableName,
table_route_value: &TableRouteValue,
) -> Result<Vec<Key>> {
) -> Result<Vec<Vec<u8>>> {
// Builds keys
let datanode_ids = if table_route_value.is_physical() {
region_distribution(table_route_value.region_routes()?)
@@ -604,11 +606,11 @@ impl TableMetadataManager {
.map(|datanode_id| DatanodeTableKey::new(datanode_id, table_id))
.collect::<HashSet<_>>();
keys.push(Key::compare_and_swap(table_name.as_raw_key()));
keys.push(Key::new(table_info_key.as_raw_key()));
keys.push(Key::new(table_route_key.as_raw_key()));
keys.push(table_name.as_raw_key());
keys.push(table_info_key.as_raw_key());
keys.push(table_route_key.as_raw_key());
for key in &datanode_table_keys {
keys.push(Key::new(key.as_raw_key()));
keys.push(key.as_raw_key());
}
Ok(keys)
}
@@ -622,8 +624,7 @@ impl TableMetadataManager {
table_route_value: &TableRouteValue,
) -> Result<()> {
let keys = self.table_metadata_keys(table_id, table_name, table_route_value)?;
self.tombstone_manager.create(keys).await?;
Ok(())
self.tombstone_manager.create(keys).await
}
/// Deletes metadata tombstone for table **permanently**.
@@ -634,11 +635,7 @@ impl TableMetadataManager {
table_name: &TableName,
table_route_value: &TableRouteValue,
) -> Result<()> {
let keys = self
.table_metadata_keys(table_id, table_name, table_route_value)?
.into_iter()
.map(|key| key.into_bytes())
.collect::<Vec<_>>();
let keys = self.table_metadata_keys(table_id, table_name, table_route_value)?;
self.tombstone_manager.delete(keys).await
}
@@ -651,8 +648,7 @@ impl TableMetadataManager {
table_route_value: &TableRouteValue,
) -> Result<()> {
let keys = self.table_metadata_keys(table_id, table_name, table_route_value)?;
self.tombstone_manager.restore(keys).await?;
Ok(())
self.tombstone_manager.restore(keys).await
}
/// Deletes metadata for table **permanently**.
@@ -663,21 +659,11 @@ impl TableMetadataManager {
table_name: &TableName,
table_route_value: &TableRouteValue,
) -> Result<()> {
let operations = self
.table_metadata_keys(table_id, table_name, table_route_value)?
.into_iter()
.map(|key| TxnOp::Delete(key.into_bytes()))
.collect::<Vec<_>>();
let txn = Txn::new().and_then(operations);
let resp = self.kv_backend.txn(txn).await?;
ensure!(
resp.succeeded,
UnexpectedSnafu {
err_msg: format!("Failed to destroy table metadata: {table_id}")
}
);
let keys = self.table_metadata_keys(table_id, table_name, table_route_value)?;
let _ = self
.kv_backend
.batch_delete(BatchDeleteRequest::new().with_keys(keys))
.await?;
Ok(())
}
@@ -724,11 +710,12 @@ impl TableMetadataManager {
let txn = Txn::merge_all(vec![update_table_name_txn, update_table_info_txn]);
let r = self.kv_backend.txn(txn).await?;
let mut r = self.kv_backend.txn(txn).await?;
// Checks whether metadata was already updated.
if !r.succeeded {
let remote_table_info = on_update_table_info_failure(&r.responses)?
let mut set = TxnOpGetResponseSet::from(&mut r.responses);
let remote_table_info = on_update_table_info_failure(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table info during the rename table metadata",
})?
@@ -756,11 +743,12 @@ impl TableMetadataManager {
.table_info_manager()
.build_update_txn(table_id, current_table_info_value, &new_table_info_value)?;
let r = self.kv_backend.txn(update_table_info_txn).await?;
let mut r = self.kv_backend.txn(update_table_info_txn).await?;
// Checks whether metadata was already updated.
if !r.succeeded {
let remote_table_info = on_update_table_info_failure(&r.responses)?
let mut set = TxnOpGetResponseSet::from(&mut r.responses);
let remote_table_info = on_update_table_info_failure(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table info during the updating table info",
})?
@@ -784,7 +772,7 @@ impl TableMetadataManager {
let mut txns = Vec::with_capacity(len);
struct OnFailure<F, R>
where
F: FnOnce(&Vec<TxnOpResponse>) -> R,
F: FnOnce(&mut TxnOpGetResponseSet) -> R,
{
table_info_value: TableInfoValue,
on_update_table_info_failure: F,
@@ -812,11 +800,12 @@ impl TableMetadataManager {
}
let txn = Txn::merge_all(txns);
let r = self.kv_backend.txn(txn).await?;
let mut r = self.kv_backend.txn(txn).await?;
if !r.succeeded {
let mut set = TxnOpGetResponseSet::from(&mut r.responses);
for on_failure in on_failures {
let remote_table_info = (on_failure.on_update_table_info_failure)(&r.responses)?
let remote_table_info = (on_failure.on_update_table_info_failure)(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table info during the updating table info",
})?
@@ -863,11 +852,12 @@ impl TableMetadataManager {
let txn = Txn::merge_all(vec![update_datanode_table_txn, update_table_route_txn]);
let r = self.kv_backend.txn(txn).await?;
let mut r = self.kv_backend.txn(txn).await?;
// Checks whether metadata was already updated.
if !r.succeeded {
let remote_table_route = on_update_table_route_failure(&r.responses)?
let mut set = TxnOpGetResponseSet::from(&mut r.responses);
let remote_table_route = on_update_table_route_failure(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table route during the updating table route",
})?
@@ -914,11 +904,12 @@ impl TableMetadataManager {
.table_route_storage()
.build_update_txn(table_id, current_table_route_value, &new_table_route_value)?;
let r = self.kv_backend.txn(update_table_route_txn).await?;
let mut r = self.kv_backend.txn(update_table_route_txn).await?;
// Checks whether metadata was already updated.
if !r.succeeded {
let remote_table_route = on_update_table_route_failure(&r.responses)?
let mut set = TxnOpGetResponseSet::from(&mut r.responses);
let remote_table_route = on_update_table_route_failure(&mut set)?
.context(error::UnexpectedSnafu {
err_msg: "Reads the empty table route during the updating leader region status",
})?
@@ -1286,14 +1277,17 @@ mod tests {
.delete_table_metadata(table_id, &table_name, table_route_value)
.await
.unwrap();
// Should be ignored.
table_metadata_manager
.delete_table_metadata(table_id, &table_name, table_route_value)
.await
.unwrap();
assert!(table_metadata_manager
.table_info_manager()
.get(table_id)
.await
.unwrap()
.is_none());
assert!(table_metadata_manager
.table_route_manager()
.table_route_storage()
@@ -1301,7 +1295,6 @@ mod tests {
.await
.unwrap()
.is_none());
assert!(table_metadata_manager
.datanode_table_manager()
.tables(datanode_id)
@@ -1316,7 +1309,6 @@ mod tests {
.await
.unwrap();
assert!(table_info.is_none());
let table_route = table_metadata_manager
.table_route_manager()
.table_route_storage()
@@ -1792,5 +1784,12 @@ mod tests {
.unwrap();
let kvs = mem_kv.dump();
assert_eq!(kvs, expected_result);
// Should be ignored.
table_metadata_manager
.restore_table_metadata(table_id, &table_name, &table_route_value)
.await
.unwrap();
let kvs = mem_kv.dump();
assert_eq!(kvs, expected_result);
}
}

8
src/common/meta/src/key/datanode_table.rs
View File

@@ -240,10 +240,14 @@ impl DatanodeTableManager {
// FIXME(weny): add unit tests.
let mut new_region_info = region_info.clone();
if need_update_options {
new_region_info.region_options = new_region_options.clone();
new_region_info
.region_options
.clone_from(new_region_options);
}
if need_update_wal_options {
new_region_info.region_wal_options = new_region_wal_options.clone();
new_region_info
.region_wal_options
.clone_from(new_region_wal_options);
}
let val = DatanodeTableValue::new(table_id, regions, new_region_info)
.try_as_raw_value()?;

15
src/common/meta/src/key/schema_name.rs
View File

@@ -197,6 +197,21 @@ impl SchemaManager {
}
}
#[derive(Debug, Clone, Hash, Eq, PartialEq, Deserialize, Serialize)]
pub struct SchemaName {
pub catalog_name: String,
pub schema_name: String,
}
impl<'a> From<&'a SchemaName> for SchemaNameKey<'a> {
fn from(value: &'a SchemaName) -> Self {
Self {
catalog: &value.catalog_name,
schema: &value.schema_name,
}
}
}
#[cfg(test)]
mod tests {

21
src/common/meta/src/key/table_info.rs
View File

@@ -18,11 +18,12 @@ use serde::{Deserialize, Serialize};
use table::metadata::{RawTableInfo, TableId};
use table::table_reference::TableReference;
use super::txn_helper::TxnOpGetResponseSet;
use crate::error::Result;
use crate::key::{
txn_helper, DeserializedValueWithBytes, TableMetaKey, TableMetaValue, TABLE_INFO_KEY_PREFIX,
};
use crate::kv_backend::txn::{Txn, TxnOpResponse};
use crate::kv_backend::txn::Txn;
use crate::kv_backend::KvBackendRef;
use crate::rpc::store::BatchGetRequest;
use crate::table_name::TableName;
@@ -109,7 +110,9 @@ impl TableInfoManager {
table_info_value: &TableInfoValue,
) -> Result<(
Txn,
impl FnOnce(&Vec<TxnOpResponse>) -> Result<Option<DeserializedValueWithBytes<TableInfoValue>>>,
impl FnOnce(
&mut TxnOpGetResponseSet,
) -> Result<Option<DeserializedValueWithBytes<TableInfoValue>>>,
)> {
let key = TableInfoKey::new(table_id);
let raw_key = key.as_raw_key();
@@ -119,7 +122,10 @@ impl TableInfoManager {
table_info_value.try_as_raw_value()?,
);
Ok((txn, txn_helper::build_txn_response_decoder_fn(raw_key)))
Ok((
txn,
TxnOpGetResponseSet::decode_with(TxnOpGetResponseSet::filter(raw_key)),
))
}
/// Builds a update table info transaction, it expected the remote value equals the `current_current_table_info_value`.
@@ -131,7 +137,9 @@ impl TableInfoManager {
new_table_info_value: &TableInfoValue,
) -> Result<(
Txn,
impl FnOnce(&Vec<TxnOpResponse>) -> Result<Option<DeserializedValueWithBytes<TableInfoValue>>>,
impl FnOnce(
&mut TxnOpGetResponseSet,
) -> Result<Option<DeserializedValueWithBytes<TableInfoValue>>>,
)> {
let key = TableInfoKey::new(table_id);
let raw_key = key.as_raw_key();
@@ -140,7 +148,10 @@ impl TableInfoManager {
let txn = txn_helper::build_compare_and_put_txn(raw_key.clone(), raw_value, new_raw_value);
Ok((txn, txn_helper::build_txn_response_decoder_fn(raw_key)))
Ok((
txn,
TxnOpGetResponseSet::decode_with(TxnOpGetResponseSet::filter(raw_key)),
))
}
pub async fn get(

27
src/common/meta/src/key/table_route.rs
View File

@@ -20,13 +20,16 @@ use snafu::{ensure, OptionExt, ResultExt};
use store_api::storage::{RegionId, RegionNumber};
use table::metadata::TableId;
use super::{txn_helper, DeserializedValueWithBytes, TableMetaValue};
use crate::error::{
self, MetadataCorruptionSnafu, Result, SerdeJsonSnafu, TableRouteNotFoundSnafu,
UnexpectedLogicalRouteTableSnafu,
};
use crate::key::{RegionDistribution, TableMetaKey, TABLE_ROUTE_PREFIX};
use crate::kv_backend::txn::{Txn, TxnOpResponse};
use crate::key::txn_helper::TxnOpGetResponseSet;
use crate::key::{
txn_helper, DeserializedValueWithBytes, RegionDistribution, TableMetaKey, TableMetaValue,
TABLE_ROUTE_PREFIX,
};
use crate::kv_backend::txn::Txn;
use crate::kv_backend::KvBackendRef;
use crate::rpc::router::{region_distribution, RegionRoute};
use crate::rpc::store::BatchGetRequest;
@@ -454,7 +457,9 @@ impl TableRouteStorage {
table_route_value: &TableRouteValue,
) -> Result<(
Txn,
impl FnOnce(&Vec<TxnOpResponse>) -> Result<Option<DeserializedValueWithBytes<TableRouteValue>>>,
impl FnOnce(
&mut TxnOpGetResponseSet,
) -> Result<Option<DeserializedValueWithBytes<TableRouteValue>>>,
)> {
let key = TableRouteKey::new(table_id);
let raw_key = key.as_raw_key();
@@ -464,7 +469,10 @@ impl TableRouteStorage {
table_route_value.try_as_raw_value()?,
);
Ok((txn, txn_helper::build_txn_response_decoder_fn(raw_key)))
Ok((
txn,
TxnOpGetResponseSet::decode_with(TxnOpGetResponseSet::filter(raw_key)),
))
}
/// Builds a update table route transaction,
@@ -477,7 +485,9 @@ impl TableRouteStorage {
new_table_route_value: &TableRouteValue,
) -> Result<(
Txn,
impl FnOnce(&Vec<TxnOpResponse>) -> Result<Option<DeserializedValueWithBytes<TableRouteValue>>>,
impl FnOnce(
&mut TxnOpGetResponseSet,
) -> Result<Option<DeserializedValueWithBytes<TableRouteValue>>>,
)> {
let key = TableRouteKey::new(table_id);
let raw_key = key.as_raw_key();
@@ -486,7 +496,10 @@ impl TableRouteStorage {
let txn = txn_helper::build_compare_and_put_txn(raw_key.clone(), raw_value, new_raw_value);
Ok((txn, txn_helper::build_txn_response_decoder_fn(raw_key)))
Ok((
txn,
TxnOpGetResponseSet::decode_with(TxnOpGetResponseSet::filter(raw_key)),
))
}
/// Returns the [`TableRouteValue`].

765
src/common/meta/src/key/tombstone.rs
View File

@@ -12,13 +12,15 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use snafu::{ensure, OptionExt};
use std::collections::HashMap;
use snafu::ensure;
use super::TableMetaKeyGetTxnOp;
use crate::error::{self, Result};
use crate::key::txn_helper::TxnOpGetResponseSet;
use crate::kv_backend::txn::{Compare, CompareOp, Txn, TxnOp};
use crate::kv_backend::KvBackendRef;
use crate::rpc::store::BatchGetRequest;
/// [TombstoneManager] provides the ability to:
/// - logically delete values
@@ -29,307 +31,160 @@ pub(crate) struct TombstoneManager {
const TOMBSTONE_PREFIX: &str = "__tombstone/";
pub(crate) struct TombstoneKey<T>(T);
pub(crate) struct TombstoneKeyValue {
pub(crate) origin_key: Vec<u8>,
pub(crate) tombstone_key: Vec<u8>,
pub(crate) value: Vec<u8>,
}
fn to_tombstone(key: &[u8]) -> Vec<u8> {
[TOMBSTONE_PREFIX.as_bytes(), key].concat()
}
impl TombstoneKey<&Vec<u8>> {
/// Returns the origin key and tombstone key.
fn to_keys(&self) -> (Vec<u8>, Vec<u8>) {
let key = self.0;
let tombstone_key = to_tombstone(key);
(key.clone(), tombstone_key)
}
/// Returns the origin key and tombstone key.
fn into_keys(self) -> (Vec<u8>, Vec<u8>) {
self.to_keys()
}
/// Returns the tombstone key.
fn to_tombstone_key(&self) -> Vec<u8> {
let key = self.0;
to_tombstone(key)
}
}
impl TableMetaKeyGetTxnOp for TombstoneKey<&Vec<u8>> {
fn build_get_op(
&self,
) -> (
TxnOp,
impl FnMut(&'_ mut TxnOpGetResponseSet) -> Option<Vec<u8>>,
) {
TxnOpGetResponseSet::build_get_op(to_tombstone(self.0))
}
}
/// The key used in the [TombstoneManager].
pub(crate) struct Key {
bytes: Vec<u8>,
// Atomic Key:
// The value corresponding to the key remains consistent between two transactions.
atomic: bool,
}
impl Key {
/// Returns a new atomic key.
pub(crate) fn compare_and_swap<T: Into<Vec<u8>>>(key: T) -> Self {
Self {
bytes: key.into(),
atomic: true,
}
}
/// Returns a new normal key.
pub(crate) fn new<T: Into<Vec<u8>>>(key: T) -> Self {
Self {
bytes: key.into(),
atomic: false,
}
}
/// Into bytes
pub(crate) fn into_bytes(self) -> Vec<u8> {
self.bytes
}
fn get_inner(&self) -> &Vec<u8> {
&self.bytes
}
fn is_atomic(&self) -> bool {
self.atomic
}
}
impl TableMetaKeyGetTxnOp for Key {
fn build_get_op(
&self,
) -> (
TxnOp,
impl FnMut(&'_ mut TxnOpGetResponseSet) -> Option<Vec<u8>>,
) {
let key = self.get_inner().clone();
(TxnOp::Get(key.clone()), TxnOpGetResponseSet::filter(key))
}
}
fn format_on_failure_error_message<F: FnMut(&mut TxnOpGetResponseSet) -> Option<Vec<u8>>>(
mut set: TxnOpGetResponseSet,
on_failure_kv_and_filters: Vec<(Vec<u8>, Vec<u8>, F)>,
) -> String {
on_failure_kv_and_filters
.into_iter()
.flat_map(|(key, value, mut filter)| {
let got = filter(&mut set);
let Some(got) = got else {
return Some(format!(
"For key: {} was expected: {}, but value does not exists",
String::from_utf8_lossy(&key),
String::from_utf8_lossy(&value),
));
};
if got != value {
Some(format!(
"For key: {} was expected: {}, but got: {}",
String::from_utf8_lossy(&key),
String::from_utf8_lossy(&value),
String::from_utf8_lossy(&got),
))
} else {
None
}
})
.collect::<Vec<_>>()
.join("; ")
}
fn format_keys(keys: &[Key]) -> String {
keys.iter()
.map(|key| String::from_utf8_lossy(&key.bytes))
.collect::<Vec<_>>()
.join(", ")
}
impl TombstoneManager {
/// Returns [TombstoneManager].
pub fn new(kv_backend: KvBackendRef) -> Self {
Self { kv_backend }
}
/// Moves value to `dest_key`.
///
/// Puts `value` to `dest_key` if the value of `src_key` equals `value`.
///
/// Otherwise retrieves the value of `src_key`.
fn build_move_value_txn(
&self,
src_key: Vec<u8>,
value: Vec<u8>,
dest_key: Vec<u8>,
) -> (Txn, impl FnMut(&mut TxnOpGetResponseSet) -> Option<Vec<u8>>) {
let txn = Txn::new()
.when(vec![Compare::with_value(
src_key.clone(),
CompareOp::Equal,
value.clone(),
)])
.and_then(vec![
TxnOp::Put(dest_key.clone(), value.clone()),
TxnOp::Delete(src_key.clone()),
])
.or_else(vec![TxnOp::Get(src_key.clone())]);
(txn, TxnOpGetResponseSet::filter(src_key))
}
async fn move_values_inner(&self, keys: &[Vec<u8>], dest_keys: &[Vec<u8>]) -> Result<()> {
ensure!(
keys.len() == dest_keys.len(),
error::UnexpectedSnafu {
err_msg: "The length of keys does not match the length of dest_keys."
}
);
// The key -> dest key mapping.
let lookup_table = keys.iter().zip(dest_keys.iter()).collect::<HashMap<_, _>>();
let resp = self
.kv_backend
.batch_get(BatchGetRequest::new().with_keys(keys.to_vec()))
.await?;
let mut results = resp
.kvs
.into_iter()
.map(|kv| (kv.key, kv.value))
.collect::<HashMap<_, _>>();
const MAX_RETRIES: usize = 8;
for _ in 0..MAX_RETRIES {
let (txns, (keys, filters)): (Vec<_>, (Vec<_>, Vec<_>)) = results
.iter()
.map(|(key, value)| {
let (txn, filter) = self.build_move_value_txn(
key.clone(),
value.clone(),
lookup_table[&key].clone(),
);
(txn, (key.clone(), filter))
})
.unzip();
let mut resp = self.kv_backend.txn(Txn::merge_all(txns)).await?;
if resp.succeeded {
return Ok(());
}
let mut set = TxnOpGetResponseSet::from(&mut resp.responses);
// Updates results.
for (idx, mut filter) in filters.into_iter().enumerate() {
if let Some(value) = filter(&mut set) {
results.insert(keys[idx].clone(), value);
} else {
results.remove(&keys[idx]);
}
}
}
error::MoveValuesSnafu {
err_msg: format!(
"keys: {:?}",
keys.iter().map(|key| String::from_utf8_lossy(key)),
),
}
.fail()
}
/// Moves values to `dest_key`.
async fn move_values(&self, keys: Vec<Vec<u8>>, dest_keys: Vec<Vec<u8>>) -> Result<()> {
let chunk_size = self.kv_backend.max_txn_ops() / 2;
if keys.len() > chunk_size {
let keys_chunks = keys.chunks(chunk_size).collect::<Vec<_>>();
let dest_keys_chunks = keys.chunks(chunk_size).collect::<Vec<_>>();
for (keys, dest_keys) in keys_chunks.into_iter().zip(dest_keys_chunks) {
self.move_values_inner(keys, dest_keys).await?;
}
Ok(())
} else {
self.move_values_inner(&keys, &dest_keys).await
}
}
/// Creates tombstones for keys.
///
/// Preforms to:
/// - retrieve all values corresponding `keys`.
/// - deletes origin values.
/// - stores tombstone values.
pub(crate) async fn create(&self, keys: Vec<Key>) -> Result<()> {
// Builds transaction to retrieve all values
let (operations, mut filters): (Vec<_>, Vec<_>) =
keys.iter().map(|key| key.build_get_op()).unzip();
pub(crate) async fn create(&self, keys: Vec<Vec<u8>>) -> Result<()> {
let (keys, dest_keys): (Vec<_>, Vec<_>) = keys
.into_iter()
.map(|key| {
let tombstone_key = to_tombstone(&key);
(key, tombstone_key)
})
.unzip();
let txn = Txn::new().and_then(operations);
let mut resp = self.kv_backend.txn(txn).await?;
ensure!(
resp.succeeded,
error::UnexpectedSnafu {
err_msg: format!(
"Failed to retrieves the metadata, keys: {}",
format_keys(&keys)
),
}
);
let mut set = TxnOpGetResponseSet::from(&mut resp.responses);
// Builds the create tombstone transaction.
let mut tombstone_operations = Vec::with_capacity(keys.len() * 2);
let mut tombstone_comparison = vec![];
let mut on_failure_operations = vec![];
let mut on_failure_kv_and_filters = vec![];
for (idx, key) in keys.iter().enumerate() {
let filter = &mut filters[idx];
let value = filter(&mut set).with_context(|| error::UnexpectedSnafu {
err_msg: format!(
"Missing value, key: {}",
String::from_utf8_lossy(key.get_inner())
),
})?;
let (origin_key, tombstone_key) = TombstoneKey(key.get_inner()).into_keys();
// Compares the atomic key.
if key.is_atomic() {
tombstone_comparison.push(Compare::with_not_exist_value(
tombstone_key.clone(),
CompareOp::Equal,
));
tombstone_comparison.push(Compare::with_value(
origin_key.clone(),
CompareOp::Equal,
value.clone(),
));
let (op, filter) = TxnOpGetResponseSet::build_get_op(origin_key.clone());
on_failure_operations.push(op);
on_failure_kv_and_filters.push((origin_key.clone(), value.clone(), filter));
}
tombstone_operations.push(TxnOp::Delete(origin_key));
tombstone_operations.push(TxnOp::Put(tombstone_key, value));
}
let txn = if !tombstone_comparison.is_empty() {
Txn::new().when(tombstone_comparison)
} else {
Txn::new()
}
.and_then(tombstone_operations);
let txn = if !on_failure_operations.is_empty() {
txn.or_else(on_failure_operations)
} else {
txn
};
let mut resp = self.kv_backend.txn(txn).await?;
// TODO(weny): add tests for atomic key changed.
if !resp.succeeded {
let set = TxnOpGetResponseSet::from(&mut resp.responses);
let err_msg = format_on_failure_error_message(set, on_failure_kv_and_filters);
return error::CasKeyChangedSnafu { err_msg }.fail();
}
Ok(())
self.move_values(keys, dest_keys).await
}
/// Restores tombstones for keys.
///
/// Preforms to:
/// - retrieve all tombstone values corresponding `keys`.
/// - stores tombstone values.
pub(crate) async fn restore(&self, keys: Vec<Key>) -> Result<()> {
// Builds transaction to retrieve all tombstone values
let tombstone_keys = keys
.iter()
.map(|key| TombstoneKey(key.get_inner()))
.collect::<Vec<_>>();
let (operations, mut filters): (Vec<_>, Vec<_>) =
tombstone_keys.iter().map(|key| key.build_get_op()).unzip();
/// - restore origin value.
/// - deletes tombstone values.
pub(crate) async fn restore(&self, keys: Vec<Vec<u8>>) -> Result<()> {
let (keys, dest_keys): (Vec<_>, Vec<_>) = keys
.into_iter()
.map(|key| {
let tombstone_key = to_tombstone(&key);
(tombstone_key, key)
})
.unzip();
let txn = Txn::new().and_then(operations);
let mut resp = self.kv_backend.txn(txn).await?;
ensure!(
resp.succeeded,
error::UnexpectedSnafu {
err_msg: format!(
"Failed to retrieves the metadata, keys: {}",
format_keys(&keys)
),
}
);
let mut set = TxnOpGetResponseSet::from(&mut resp.responses);
// Builds the restore tombstone transaction.
let mut tombstone_operations = Vec::with_capacity(keys.len() * 2);
let mut tombstone_comparison = vec![];
let mut on_failure_operations = vec![];
let mut on_failure_kv_and_filters = vec![];
for (idx, key) in keys.iter().enumerate() {
let filter = &mut filters[idx];
let value = filter(&mut set).with_context(|| error::UnexpectedSnafu {
err_msg: format!(
"Missing value, key: {}",
String::from_utf8_lossy(key.get_inner())
),
})?;
let (origin_key, tombstone_key) = tombstone_keys[idx].to_keys();
// Compares the atomic key.
if key.is_atomic() {
tombstone_comparison.push(Compare::with_not_exist_value(
origin_key.clone(),
CompareOp::Equal,
));
tombstone_comparison.push(Compare::with_value(
tombstone_key.clone(),
CompareOp::Equal,
value.clone(),
));
let (op, filter) = tombstone_keys[idx].build_get_op();
on_failure_operations.push(op);
on_failure_kv_and_filters.push((tombstone_key.clone(), value.clone(), filter));
}
tombstone_operations.push(TxnOp::Delete(tombstone_key));
tombstone_operations.push(TxnOp::Put(origin_key, value));
}
let txn = if !tombstone_comparison.is_empty() {
Txn::new().when(tombstone_comparison)
} else {
Txn::new()
}
.and_then(tombstone_operations);
let txn = if !on_failure_operations.is_empty() {
txn.or_else(on_failure_operations)
} else {
txn
};
let mut resp = self.kv_backend.txn(txn).await?;
// TODO(weny): add tests for atomic key changed.
if !resp.succeeded {
let set = TxnOpGetResponseSet::from(&mut resp.responses);
let err_msg = format_on_failure_error_message(set, on_failure_kv_and_filters);
return error::CasKeyChangedSnafu { err_msg }.fail();
}
Ok(())
self.move_values(keys, dest_keys).await
}
/// Deletes tombstones for keys.
/// Deletes tombstones values for the specified `keys`.
pub(crate) async fn delete(&self, keys: Vec<Vec<u8>>) -> Result<()> {
let operations = keys
.iter()
.map(|key| TxnOp::Delete(TombstoneKey(key).to_tombstone_key()))
.map(|key| TxnOp::Delete(to_tombstone(key)))
.collect::<Vec<_>>();
let txn = Txn::new().and_then(operations);
@@ -342,13 +197,41 @@ impl TombstoneManager {
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use crate::key::tombstone::{Key, TombstoneKey, TombstoneManager};
use super::to_tombstone;
use crate::error::Error;
use crate::key::tombstone::TombstoneManager;
use crate::kv_backend::memory::MemoryKvBackend;
use crate::kv_backend::KvBackend;
use crate::rpc::store::PutRequest;
#[derive(Debug, Clone)]
struct MoveValue {
key: Vec<u8>,
dest_key: Vec<u8>,
value: Vec<u8>,
}
async fn check_moved_values(
kv_backend: Arc<MemoryKvBackend<Error>>,
move_values: &[MoveValue],
) {
for MoveValue {
key,
dest_key,
value,
} in move_values
{
assert!(kv_backend.get(key).await.unwrap().is_none());
assert_eq!(
&kv_backend.get(dest_key).await.unwrap().unwrap().value,
value,
);
}
}
#[tokio::test]
async fn test_create_tombstone() {
let kv_backend = Arc::new(MemoryKvBackend::default());
@@ -362,14 +245,14 @@ mod tests {
.await
.unwrap();
tombstone_manager
.create(vec![Key::compare_and_swap("bar"), Key::new("foo")])
.create(vec![b"bar".to_vec(), b"foo".to_vec()])
.await
.unwrap();
assert!(!kv_backend.exists(b"bar").await.unwrap());
assert!(!kv_backend.exists(b"foo").await.unwrap());
assert_eq!(
kv_backend
.get(&TombstoneKey(&"bar".into()).to_tombstone_key())
.get(&to_tombstone(b"bar"))
.await
.unwrap()
.unwrap()
@@ -378,7 +261,7 @@ mod tests {
);
assert_eq!(
kv_backend
.get(&TombstoneKey(&"foo".into()).to_tombstone_key())
.get(&to_tombstone(b"foo"))
.await
.unwrap()
.unwrap()
@@ -389,9 +272,10 @@ mod tests {
}
#[tokio::test]
async fn test_create_tombstone_without_atomic_key() {
async fn test_create_tombstone_with_non_exist_values() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
kv_backend
.put(PutRequest::new().with_key("bar").with_value("baz"))
.await
@@ -400,52 +284,20 @@ mod tests {
.put(PutRequest::new().with_key("foo").with_value("hi"))
.await
.unwrap();
tombstone_manager
.create(vec![Key::new("bar"), Key::new("foo")])
.create(vec![b"bar".to_vec(), b"baz".to_vec()])
.await
.unwrap();
assert!(!kv_backend.exists(b"bar").await.unwrap());
assert!(!kv_backend.exists(b"foo").await.unwrap());
assert_eq!(
kv_backend
.get(&TombstoneKey(&"bar".into()).to_tombstone_key())
.await
.unwrap()
.unwrap()
.value,
b"baz"
);
assert_eq!(
kv_backend
.get(&TombstoneKey(&"foo".into()).to_tombstone_key())
.await
.unwrap()
.unwrap()
.value,
b"hi"
);
assert_eq!(kv_backend.len(), 2);
}
#[tokio::test]
async fn test_create_tombstone_origin_value_not_found_err() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
kv_backend
.put(PutRequest::new().with_key("bar").with_value("baz"))
.await
.unwrap();
kv_backend
.put(PutRequest::new().with_key("foo").with_value("hi"))
.await
.unwrap();
let err = tombstone_manager
.create(vec![Key::compare_and_swap("bar"), Key::new("baz")])
.await
.unwrap_err();
assert!(err.to_string().contains("Missing value"));
check_moved_values(
kv_backend.clone(),
&[MoveValue {
key: b"bar".to_vec(),
dest_key: to_tombstone(b"bar"),
value: b"baz".to_vec(),
}],
)
.await;
}
#[tokio::test]
@@ -462,63 +314,16 @@ mod tests {
.unwrap();
let expected_kvs = kv_backend.dump();
tombstone_manager
.create(vec![Key::compare_and_swap("bar"), Key::new("foo")])
.create(vec![b"bar".to_vec(), b"foo".to_vec()])
.await
.unwrap();
tombstone_manager
.restore(vec![Key::compare_and_swap("bar"), Key::new("foo")])
.restore(vec![b"bar".to_vec(), b"foo".to_vec()])
.await
.unwrap();
assert_eq!(expected_kvs, kv_backend.dump());
}
#[tokio::test]
async fn test_restore_tombstone_without_atomic_key() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
kv_backend
.put(PutRequest::new().with_key("bar").with_value("baz"))
.await
.unwrap();
kv_backend
.put(PutRequest::new().with_key("foo").with_value("hi"))
.await
.unwrap();
let expected_kvs = kv_backend.dump();
tombstone_manager
.create(vec![Key::compare_and_swap("bar"), Key::new("foo")])
.await
.unwrap();
tombstone_manager
.restore(vec![Key::new("bar"), Key::new("foo")])
.await
.unwrap();
assert_eq!(expected_kvs, kv_backend.dump());
}
#[tokio::test]
async fn test_restore_tombstone_origin_value_not_found_err() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
kv_backend
.put(PutRequest::new().with_key("bar").with_value("baz"))
.await
.unwrap();
kv_backend
.put(PutRequest::new().with_key("foo").with_value("hi"))
.await
.unwrap();
tombstone_manager
.create(vec![Key::compare_and_swap("bar"), Key::new("foo")])
.await
.unwrap();
let err = tombstone_manager
.restore(vec![Key::new("bar"), Key::new("baz")])
.await
.unwrap_err();
assert!(err.to_string().contains("Missing value"));
}
#[tokio::test]
async fn test_delete_tombstone() {
let kv_backend = Arc::new(MemoryKvBackend::default());
@@ -532,7 +337,7 @@ mod tests {
.await
.unwrap();
tombstone_manager
.create(vec![Key::compare_and_swap("bar"), Key::new("foo")])
.create(vec![b"bar".to_vec(), b"foo".to_vec()])
.await
.unwrap();
tombstone_manager
@@ -541,4 +346,216 @@ mod tests {
.unwrap();
assert!(kv_backend.is_empty());
}
#[tokio::test]
async fn test_move_values() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
let kvs = HashMap::from([
(b"bar".to_vec(), b"baz".to_vec()),
(b"foo".to_vec(), b"hi".to_vec()),
(b"baz".to_vec(), b"hello".to_vec()),
]);
for (key, value) in &kvs {
kv_backend
.put(
PutRequest::new()
.with_key(key.clone())
.with_value(value.clone()),
)
.await
.unwrap();
}
let move_values = kvs
.iter()
.map(|(key, value)| MoveValue {
key: key.clone(),
dest_key: to_tombstone(key),
value: value.clone(),
})
.collect::<Vec<_>>();
let (keys, dest_keys): (Vec<_>, Vec<_>) = move_values
.clone()
.into_iter()
.map(|kv| (kv.key, kv.dest_key))
.unzip();
tombstone_manager
.move_values(keys.clone(), dest_keys.clone())
.await
.unwrap();
check_moved_values(kv_backend.clone(), &move_values).await;
// Moves again
tombstone_manager
.move_values(keys.clone(), dest_keys.clone())
.await
.unwrap();
check_moved_values(kv_backend.clone(), &move_values).await;
}
#[tokio::test]
async fn test_move_values_with_non_exists_values() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
let kvs = HashMap::from([
(b"bar".to_vec(), b"baz".to_vec()),
(b"foo".to_vec(), b"hi".to_vec()),
(b"baz".to_vec(), b"hello".to_vec()),
]);
for (key, value) in &kvs {
kv_backend
.put(
PutRequest::new()
.with_key(key.clone())
.with_value(value.clone()),
)
.await
.unwrap();
}
let move_values = kvs
.iter()
.map(|(key, value)| MoveValue {
key: key.clone(),
dest_key: to_tombstone(key),
value: value.clone(),
})
.collect::<Vec<_>>();
let (mut keys, mut dest_keys): (Vec<_>, Vec<_>) = move_values
.clone()
.into_iter()
.map(|kv| (kv.key, kv.dest_key))
.unzip();
keys.push(b"non-exists".to_vec());
dest_keys.push(b"hi/non-exists".to_vec());
tombstone_manager
.move_values(keys.clone(), dest_keys.clone())
.await
.unwrap();
check_moved_values(kv_backend.clone(), &move_values).await;
// Moves again
tombstone_manager
.move_values(keys.clone(), dest_keys.clone())
.await
.unwrap();
check_moved_values(kv_backend.clone(), &move_values).await;
}
#[tokio::test]
async fn test_move_values_changed() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
let kvs = HashMap::from([
(b"bar".to_vec(), b"baz".to_vec()),
(b"foo".to_vec(), b"hi".to_vec()),
(b"baz".to_vec(), b"hello".to_vec()),
]);
for (key, value) in &kvs {
kv_backend
.put(
PutRequest::new()
.with_key(key.clone())
.with_value(value.clone()),
)
.await
.unwrap();
}
kv_backend
.put(PutRequest::new().with_key("baz").with_value("changed"))
.await
.unwrap();
let move_values = kvs
.iter()
.map(|(key, value)| MoveValue {
key: key.clone(),
dest_key: to_tombstone(key),
value: value.clone(),
})
.collect::<Vec<_>>();
let (keys, dest_keys): (Vec<_>, Vec<_>) = move_values
.clone()
.into_iter()
.map(|kv| (kv.key, kv.dest_key))
.unzip();
tombstone_manager
.move_values(keys, dest_keys)
.await
.unwrap();
}
#[tokio::test]
async fn test_move_values_overwrite_dest_values() {
let kv_backend = Arc::new(MemoryKvBackend::default());
let tombstone_manager = TombstoneManager::new(kv_backend.clone());
let kvs = HashMap::from([
(b"bar".to_vec(), b"baz".to_vec()),
(b"foo".to_vec(), b"hi".to_vec()),
(b"baz".to_vec(), b"hello".to_vec()),
]);
for (key, value) in &kvs {
kv_backend
.put(
PutRequest::new()
.with_key(key.clone())
.with_value(value.clone()),
)
.await
.unwrap();
}
// Prepares
let move_values = kvs
.iter()
.map(|(key, value)| MoveValue {
key: key.clone(),
dest_key: to_tombstone(key),
value: value.clone(),
})
.collect::<Vec<_>>();
let (keys, dest_keys): (Vec<_>, Vec<_>) = move_values
.clone()
.into_iter()
.map(|kv| (kv.key, kv.dest_key))
.unzip();
tombstone_manager
.move_values(keys, dest_keys)
.await
.unwrap();
check_moved_values(kv_backend.clone(), &move_values).await;
// Overwrites existing dest keys.
let kvs = HashMap::from([
(b"bar".to_vec(), b"new baz".to_vec()),
(b"foo".to_vec(), b"new hi".to_vec()),
(b"baz".to_vec(), b"new baz".to_vec()),
]);
for (key, value) in &kvs {
kv_backend
.put(
PutRequest::new()
.with_key(key.clone())
.with_value(value.clone()),
)
.await
.unwrap();
}
let move_values = kvs
.iter()
.map(|(key, value)| MoveValue {
key: key.clone(),
dest_key: to_tombstone(key),
value: value.clone(),
})
.collect::<Vec<_>>();
let (keys, dest_keys): (Vec<_>, Vec<_>) = move_values
.clone()
.into_iter()
.map(|kv| (kv.key, kv.dest_key))
.unzip();
tombstone_manager
.move_values(keys, dest_keys)
.await
.unwrap();
check_moved_values(kv_backend.clone(), &move_values).await;
}
}

38
src/common/meta/src/key/txn_helper.rs
View File

@@ -21,21 +21,9 @@ use crate::kv_backend::txn::{Compare, CompareOp, Txn, TxnOp, TxnOpResponse};
use crate::rpc::KeyValue;
/// The response set of [TxnOpResponse::ResponseGet]
pub(crate) struct TxnOpGetResponseSet(Vec<KeyValue>);
pub struct TxnOpGetResponseSet(Vec<KeyValue>);
impl TxnOpGetResponseSet {
/// Returns a [TxnOp] to retrieve the value corresponding `key` and
/// a filter to consume corresponding [KeyValue] from [TxnOpGetResponseSet].
pub(crate) fn build_get_op<T: Into<Vec<u8>>>(
key: T,
) -> (
TxnOp,
impl FnMut(&'_ mut TxnOpGetResponseSet) -> Option<Vec<u8>>,
) {
let key = key.into();
(TxnOp::Get(key.clone()), TxnOpGetResponseSet::filter(key))
}
/// Returns a filter to consume a [KeyValue] where the key equals `key`.
pub(crate) fn filter(key: Vec<u8>) -> impl FnMut(&mut TxnOpGetResponseSet) -> Option<Vec<u8>> {
move |set| {
@@ -80,30 +68,6 @@ impl From<&mut Vec<TxnOpResponse>> for TxnOpGetResponseSet {
}
}
// TODO(weny): using `TxnOpGetResponseSet`.
pub(crate) fn build_txn_response_decoder_fn<T>(
raw_key: Vec<u8>,
) -> impl FnOnce(&Vec<TxnOpResponse>) -> Result<Option<DeserializedValueWithBytes<T>>>
where
T: Serialize + DeserializeOwned + TableMetaValue,
{
move |txn_res: &Vec<TxnOpResponse>| {
txn_res
.iter()
.filter_map(|resp| {
if let TxnOpResponse::ResponseGet(r) = resp {
Some(r)
} else {
None
}
})
.flat_map(|r| &r.kvs)
.find(|kv| kv.key == raw_key)
.map(|kv| DeserializedValueWithBytes::from_inner_slice(&kv.value))
.transpose()
}
}
pub(crate) fn build_put_if_absent_txn(key: Vec<u8>, value: Vec<u8>) -> Txn {
Txn::new()
.when(vec![Compare::with_not_exist_value(

40
src/common/meta/src/rpc/store.rs
View File

@@ -97,7 +97,6 @@ impl Display for RangeRequest {
}
impl RangeRequest {
#[inline]
pub fn new() -> Self {
Self {
key: vec![],
@@ -114,7 +113,6 @@ impl RangeRequest {
/// key is the first key for the range, If range_end is not given, the
/// request only looks up key.
#[inline]
pub fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self
@@ -129,7 +127,6 @@ impl RangeRequest {
/// then the range request gets all keys prefixed with key.
/// If both key and range_end are '\0', then the range request returns all
/// keys.
#[inline]
pub fn with_range(mut self, key: impl Into<Vec<u8>>, range_end: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self.range_end = range_end.into();
@@ -138,7 +135,6 @@ impl RangeRequest {
/// Gets all keys prefixed with key.
/// range_end is the key plus one (e.g., "aa"+1 == "ab", "a\xff"+1 == "b"),
#[inline]
pub fn with_prefix(mut self, key: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self.range_end = util::get_prefix_end_key(&self.key);
@@ -147,14 +143,12 @@ impl RangeRequest {
/// limit is a limit on the number of keys returned for the request. When
/// limit is set to 0, it is treated as no limit.
#[inline]
pub fn with_limit(mut self, limit: i64) -> Self {
self.limit = limit;
self
}
/// keys_only when set returns only the keys and not the values.
#[inline]
pub fn with_keys_only(mut self) -> Self {
self.keys_only = true;
self
@@ -204,7 +198,6 @@ impl RangeResponse {
}
}
#[inline]
pub fn take_kvs(&mut self) -> Vec<KeyValue> {
self.kvs.drain(..).collect()
}
@@ -244,7 +237,6 @@ impl From<PbPutRequest> for PutRequest {
}
impl PutRequest {
#[inline]
pub fn new() -> Self {
Self {
key: vec![],
@@ -254,7 +246,6 @@ impl PutRequest {
}
/// key is the key, in bytes, to put into the key-value store.
#[inline]
pub fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self
@@ -262,7 +253,6 @@ impl PutRequest {
/// value is the value, in bytes, to associate with the key in the
/// key-value store.
#[inline]
pub fn with_value(mut self, value: impl Into<Vec<u8>>) -> Self {
self.value = value.into();
self
@@ -270,7 +260,6 @@ impl PutRequest {
/// If prev_kv is set, gets the previous key-value pair before changing it.
/// The previous key-value pair will be returned in the put response.
#[inline]
pub fn with_prev_kv(mut self) -> Self {
self.prev_kv = true;
self
@@ -330,18 +319,15 @@ impl Default for BatchGetRequest {
}
impl BatchGetRequest {
#[inline]
pub fn new() -> Self {
Self { keys: vec![] }
}
#[inline]
pub fn with_keys(mut self, keys: Vec<Vec<u8>>) -> Self {
self.keys = keys;
self
}
#[inline]
pub fn add_key(mut self, key: impl Into<Vec<u8>>) -> Self {
self.keys.push(key.into());
self
@@ -416,7 +402,6 @@ impl From<PbBatchPutRequest> for BatchPutRequest {
}
impl BatchPutRequest {
#[inline]
pub fn new() -> Self {
Self {
kvs: vec![],
@@ -424,7 +409,6 @@ impl BatchPutRequest {
}
}
#[inline]
pub fn add_kv(mut self, key: impl Into<Vec<u8>>, value: impl Into<Vec<u8>>) -> Self {
self.kvs.push(KeyValue {
key: key.into(),
@@ -435,7 +419,6 @@ impl BatchPutRequest {
/// If prev_kv is set, gets the previous key-value pair before changing it.
/// The previous key-value pair will be returned in the put response.
#[inline]
pub fn with_prev_kv(mut self) -> Self {
self.prev_kv = true;
self
@@ -467,7 +450,6 @@ impl BatchPutResponse {
}
}
#[inline]
pub fn take_prev_kvs(&mut self) -> Vec<KeyValue> {
self.prev_kvs.drain(..).collect()
}
@@ -501,7 +483,6 @@ impl From<PbBatchDeleteRequest> for BatchDeleteRequest {
}
impl BatchDeleteRequest {
#[inline]
pub fn new() -> Self {
Self {
keys: vec![],
@@ -509,7 +490,12 @@ impl BatchDeleteRequest {
}
}
#[inline]
/// Sets `keys`.
pub fn with_keys(mut self, keys: Vec<Vec<u8>>) -> Self {
self.keys = keys;
self
}
pub fn add_key(mut self, key: impl Into<Vec<u8>>) -> Self {
self.keys.push(key.into());
self
@@ -517,7 +503,6 @@ impl BatchDeleteRequest {
/// If prev_kv is set, gets the previous key-value pair before deleting it.
/// The previous key-value pair will be returned in the batch delete response.
#[inline]
pub fn with_prev_kv(mut self) -> Self {
self.prev_kv = true;
self
@@ -582,7 +567,6 @@ impl From<PbCompareAndPutRequest> for CompareAndPutRequest {
}
impl CompareAndPutRequest {
#[inline]
pub fn new() -> Self {
Self {
key: vec![],
@@ -592,14 +576,12 @@ impl CompareAndPutRequest {
}
/// key is the key, in bytes, to put into the key-value store.
#[inline]
pub fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self
}
/// expect is the previous value, in bytes
#[inline]
pub fn with_expect(mut self, expect: impl Into<Vec<u8>>) -> Self {
self.expect = expect.into();
self
@@ -607,7 +589,6 @@ impl CompareAndPutRequest {
/// value is the value, in bytes, to associate with the key in the
/// key-value store.
#[inline]
pub fn with_value(mut self, value: impl Into<Vec<u8>>) -> Self {
self.value = value.into();
self
@@ -649,12 +630,10 @@ impl CompareAndPutResponse {
}
}
#[inline]
pub fn is_success(&self) -> bool {
self.success
}
#[inline]
pub fn take_prev_kv(&mut self) -> Option<KeyValue> {
self.prev_kv.take()
}
@@ -703,7 +682,6 @@ impl From<PbDeleteRangeRequest> for DeleteRangeRequest {
}
impl DeleteRangeRequest {
#[inline]
pub fn new() -> Self {
Self {
key: vec![],
@@ -719,7 +697,6 @@ impl DeleteRangeRequest {
/// key is the first key to delete in the range. If range_end is not given,
/// the range is defined to contain only the key argument.
#[inline]
pub fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self
@@ -735,7 +712,6 @@ impl DeleteRangeRequest {
/// the keys with the prefix (the given key).
/// If range_end is '\0', the range is all keys greater than or equal to the
/// key argument.
#[inline]
pub fn with_range(mut self, key: impl Into<Vec<u8>>, range_end: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self.range_end = range_end.into();
@@ -744,7 +720,6 @@ impl DeleteRangeRequest {
/// Deletes all keys prefixed with key.
/// range_end is one bit larger than the given key.
#[inline]
pub fn with_prefix(mut self, key: impl Into<Vec<u8>>) -> Self {
self.key = key.into();
self.range_end = util::get_prefix_end_key(&self.key);
@@ -753,7 +728,6 @@ impl DeleteRangeRequest {
/// If prev_kv is set, gets the previous key-value pairs before deleting it.
/// The previous key-value pairs will be returned in the delete response.
#[inline]
pub fn with_prev_kv(mut self) -> Self {
self.prev_kv = true;
self
@@ -788,12 +762,10 @@ impl DeleteRangeResponse {
}
}
#[inline]
pub fn deleted(&self) -> i64 {
self.deleted
}
#[inline]
pub fn take_prev_kvs(&mut self) -> Vec<KeyValue> {
self.prev_kvs.drain(..).collect()
}

9
src/common/meta/src/sequence.rs
View File

@@ -172,9 +172,7 @@ impl Inner {
if !res.success {
if let Some(kv) = res.prev_kv {
expect = kv.value.clone();
let v: [u8; 8] = match kv.value.try_into() {
let v: [u8; 8] = match kv.value.clone().try_into() {
Ok(a) => a,
Err(v) => {
return error::UnexpectedSequenceValueSnafu {
@@ -184,13 +182,12 @@ impl Inner {
}
};
let v = u64::from_le_bytes(v);
// If the existed value is smaller than the initial, we should start from the initial.
start = v.max(self.initial);
expect = kv.value;
} else {
expect = vec![];
start = self.initial;
expect = vec![];
}
continue;
}

19
src/common/meta/src/state_store.rs
View File

@@ -22,6 +22,7 @@ use common_procedure::store::util::multiple_value_stream;
use common_procedure::Result as ProcedureResult;
use futures::future::try_join_all;
use futures::StreamExt;
use itertools::Itertools;
use snafu::ResultExt;
use crate::error::Result;
@@ -79,17 +80,21 @@ fn decode_kv(kv: KeyValue) -> Result<(String, Vec<u8>)> {
Ok((key, value))
}
enum SplitValue<'a> {
Single(&'a [u8]),
Multiple(Vec<&'a [u8]>),
enum SplitValue {
Single(Vec<u8>),
Multiple(Vec<Vec<u8>>),
}
fn split_value(value: &[u8], max_value_size: Option<usize>) -> SplitValue<'_> {
fn split_value(value: Vec<u8>, max_value_size: Option<usize>) -> SplitValue {
if let Some(max_value_size) = max_value_size {
if value.len() <= max_value_size {
SplitValue::Single(value)
} else {
SplitValue::Multiple(value.chunks(max_value_size).collect::<Vec<_>>())
let mut values = vec![];
for chunk in value.into_iter().chunks(max_value_size).into_iter() {
values.push(chunk.collect());
}
SplitValue::Multiple(values)
}
} else {
SplitValue::Single(value)
@@ -99,10 +104,10 @@ fn split_value(value: &[u8], max_value_size: Option<usize>) -> SplitValue<'_> {
#[async_trait]
impl StateStore for KvStateStore {
async fn put(&self, key: &str, value: Vec<u8>) -> ProcedureResult<()> {
let split = split_value(&value, self.max_value_size);
let split = split_value(value, self.max_value_size);
let key = with_prefix(key);
match split {
SplitValue::Single(_) => {
SplitValue::Single(value) => {
self.kv_backend
.put(
PutRequest::new()

9
src/common/meta/src/test_util.rs
View File

@@ -14,14 +14,13 @@
use std::sync::Arc;
use api::region::RegionResponse;
use api::v1::region::{QueryRequest, RegionRequest};
pub use common_base::AffectedRows;
use common_recordbatch::SendableRecordBatchStream;
use crate::cache_invalidator::DummyCacheInvalidator;
use crate::datanode_manager::{
Datanode, DatanodeManager, DatanodeManagerRef, DatanodeRef, HandleResponse,
};
use crate::datanode_manager::{Datanode, DatanodeManager, DatanodeManagerRef, DatanodeRef};
use crate::ddl::table_meta::TableMetadataAllocator;
use crate::ddl::DdlContext;
use crate::error::Result;
@@ -35,7 +34,7 @@ use crate::wal_options_allocator::WalOptionsAllocator;
#[async_trait::async_trait]
pub trait MockDatanodeHandler: Sync + Send + Clone {
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<HandleResponse>;
async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<RegionResponse>;
async fn handle_query(
&self,
@@ -65,7 +64,7 @@ struct MockDatanode<T> {
#[async_trait::async_trait]
impl<T: MockDatanodeHandler> Datanode for MockDatanode<T> {
async fn handle(&self, request: RegionRequest) -> Result<HandleResponse> {
async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
self.handler.handle(&self.peer, request).await
}

2
src/common/meta/src/wal_options_allocator.rs
View File

@@ -163,7 +163,7 @@ mod tests {
let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
let mut topic_manager = KafkaTopicManager::new(config.clone(), kv_backend);
// Replaces the default topic pool with the constructed topics.
topic_manager.topic_pool = topics.clone();
topic_manager.topic_pool.clone_from(&topics);
// Replaces the default selector with a round-robin selector without shuffled.
topic_manager.topic_selector = Arc::new(RoundRobinTopicSelector::default());

2
src/common/meta/src/wal_options_allocator/kafka/topic_manager.rs
View File

@@ -291,7 +291,7 @@ mod tests {
let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
let mut manager = TopicManager::new(config.clone(), kv_backend);
// Replaces the default topic pool with the constructed topics.
manager.topic_pool = topics.clone();
manager.topic_pool.clone_from(&topics);
// Replaces the default selector with a round-robin selector without shuffled.
manager.topic_selector = Arc::new(RoundRobinTopicSelector::default());
manager.start().await.unwrap();

8
src/common/procedure/src/local/rwlock.rs
View File

@@ -19,19 +19,19 @@ use std::sync::{Arc, Mutex};
use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};
pub enum OwnedKeyRwLockGuard {
Read(OwnedRwLockReadGuard<()>),
Write(OwnedRwLockWriteGuard<()>),
Read { _guard: OwnedRwLockReadGuard<()> },
Write { _guard: OwnedRwLockWriteGuard<()> },
}
impl From<OwnedRwLockReadGuard<()>> for OwnedKeyRwLockGuard {
fn from(guard: OwnedRwLockReadGuard<()>) -> Self {
OwnedKeyRwLockGuard::Read(guard)
OwnedKeyRwLockGuard::Read { _guard: guard }
}
}
impl From<OwnedRwLockWriteGuard<()>> for OwnedKeyRwLockGuard {
fn from(guard: OwnedRwLockWriteGuard<()>) -> Self {
OwnedKeyRwLockGuard::Write(guard)
OwnedKeyRwLockGuard::Write { _guard: guard }
}
}

6
src/common/query/src/columnar_value.rs
View File

@@ -17,7 +17,7 @@ use datatypes::prelude::ConcreteDataType;
use datatypes::vectors::{Helper, VectorRef};
use snafu::ResultExt;
use crate::error::{self, IntoVectorSnafu, Result};
use crate::error::{self, GeneralDataFusionSnafu, IntoVectorSnafu, Result};
use crate::prelude::ScalarValue;
/// Represents the result from an expression
@@ -43,7 +43,9 @@ impl ColumnarValue {
Ok(match self {
ColumnarValue::Vector(v) => v,
ColumnarValue::Scalar(s) => {
let v = s.to_array_of_size(num_rows);
let v = s
.to_array_of_size(num_rows)
.context(GeneralDataFusionSnafu)?;
let data_type = v.data_type().clone();
Helper::try_into_vector(v).context(IntoVectorSnafu { data_type })?
}

4
src/common/query/src/lib.rs
View File

@@ -94,8 +94,8 @@ impl Debug for OutputData {
OutputData::RecordBatches(recordbatches) => {
write!(f, "OutputData::RecordBatches({recordbatches:?})")
}
OutputData::Stream(_) => {
write!(f, "OutputData::Stream(<stream>)")
OutputData::Stream(s) => {
write!(f, "OutputData::Stream(<{}>)", s.name())
}
}
}

13
src/common/query/src/logical_plan.rs
View File

@@ -72,6 +72,7 @@ pub fn create_aggregate_function(
mod tests {
use std::sync::Arc;
use datafusion_common::DFSchema;
use datafusion_expr::{
ColumnarValue as DfColumnarValue, ScalarUDF as DfScalarUDF,
TypeSignature as DfTypeSignature,
@@ -135,15 +136,17 @@ mod tests {
// test into_df_udf
let df_udf: DfScalarUDF = udf.into();
assert_eq!("and", df_udf.name);
assert_eq!("and", df_udf.name());
let types = vec![DataType::Boolean, DataType::Boolean];
assert!(
matches!(&df_udf.signature.type_signature, DfTypeSignature::Exact(ts) if ts.clone() == types)
matches!(&df_udf.signature().type_signature, DfTypeSignature::Exact(ts) if ts.clone() == types)
);
assert_eq!(
Arc::new(DataType::Boolean),
(df_udf.return_type)(&[]).unwrap()
DataType::Boolean,
df_udf
.return_type_from_exprs(&[], &DFSchema::empty(), &[])
.unwrap()
);
let args = vec![
@@ -152,7 +155,7 @@ mod tests {
];
// call the function
let result = (df_udf.fun)(&args).unwrap();
let result = (df_udf.fun())(&args).unwrap();
match result {
DfColumnarValue::Array(arr) => {

4
src/common/query/src/logical_plan/accumulator.rs
View File

@@ -126,7 +126,7 @@ impl DfAccumulatorAdaptor {
}
impl DfAccumulator for DfAccumulatorAdaptor {
fn state(&self) -> DfResult<Vec<ScalarValue>> {
fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
let state_values = self.accumulator.state()?;
let state_types = self.creator.state_types()?;
if state_values.len() != state_types.len() {
@@ -161,7 +161,7 @@ impl DfAccumulator for DfAccumulatorAdaptor {
Ok(())
}
fn evaluate(&self) -> DfResult<ScalarValue> {
fn evaluate(&mut self) -> DfResult<ScalarValue> {
let value = self.accumulator.evaluate()?;
let output_type = self.creator.output_type()?;
let scalar_value = value

8
src/common/query/src/logical_plan/expr.rs
View File

@@ -44,9 +44,7 @@ pub fn build_filter_from_timestamp(
ts_col_name: &str,
time_range: Option<&TimestampRange>,
) -> Option<Expr> {
let Some(time_range) = time_range else {
return None;
};
let time_range = time_range?;
let ts_col_expr = DfExpr::Column(Column {
relation: None,
name: ts_col_name.to_string(),
@@ -94,10 +92,10 @@ mod tests {
#[test]
fn test_from_df_expr() {
let df_expr = DfExpr::Wildcard;
let df_expr = DfExpr::Wildcard { qualifier: None };
let expr: Expr = df_expr.into();
assert_eq!(DfExpr::Wildcard, *expr.df_expr());
assert_eq!(DfExpr::Wildcard { qualifier: None }, *expr.df_expr());
}
}

98
src/common/query/src/logical_plan/udaf.rs
View File

@@ -16,15 +16,18 @@
//!
//! Modified from DataFusion.
use std::any::Any;
use std::fmt::{self, Debug, Formatter};
use std::sync::Arc;
use datafusion::arrow::datatypes::Field;
use datafusion_common::Result;
use datafusion_expr::function::AccumulatorArgs;
use datafusion_expr::{
AccumulatorFactoryFunction, AggregateUDF as DfAggregateUdf,
StateTypeFunction as DfStateTypeFunction,
Accumulator, AccumulatorFactoryFunction, AggregateUDF as DfAggregateUdf, AggregateUDFImpl,
};
use datatypes::arrow::datatypes::DataType as ArrowDataType;
use datatypes::prelude::*;
use datatypes::data_type::DataType;
use crate::function::{
to_df_return_type, AccumulatorFunctionImpl, ReturnTypeFunction, StateTypeFunction,
@@ -90,13 +93,72 @@ impl AggregateFunction {
impl From<AggregateFunction> for DfAggregateUdf {
fn from(udaf: AggregateFunction) -> Self {
DfAggregateUdf::new(
&udaf.name,
&udaf.signature.into(),
&to_df_return_type(udaf.return_type),
&to_df_accumulator_func(udaf.accumulator, udaf.creator.clone()),
&to_df_state_type(udaf.state_type),
)
struct DfUdafAdapter {
name: String,
signature: datafusion_expr::Signature,
return_type_func: datafusion_expr::ReturnTypeFunction,
accumulator: AccumulatorFactoryFunction,
creator: AggregateFunctionCreatorRef,
}
impl Debug for DfUdafAdapter {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_struct("DfUdafAdapter")
.field("name", &self.name)
.field("signature", &self.signature)
.finish()
}
}
impl AggregateUDFImpl for DfUdafAdapter {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
&self.name
}
fn signature(&self) -> &datafusion_expr::Signature {
&self.signature
}
fn return_type(&self, arg_types: &[ArrowDataType]) -> Result<ArrowDataType> {
(self.return_type_func)(arg_types).map(|x| x.as_ref().clone())
}
fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
(self.accumulator)(acc_args)
}
fn state_fields(
&self,
name: &str,
_value_type: ArrowDataType,
_ordering_fields: Vec<Field>,
) -> Result<Vec<Field>> {
self.creator
.state_types()
.map(|x| {
(0..x.len())
.zip(x)
.map(|(i, t)| {
Field::new(format!("{}_{}", name, i), t.as_arrow_type(), true)
})
.collect::<Vec<_>>()
})
.map_err(|e| e.into())
}
}
DfUdafAdapter {
name: udaf.name,
signature: udaf.signature.into(),
return_type_func: to_df_return_type(udaf.return_type),
accumulator: to_df_accumulator_func(udaf.accumulator, udaf.creator.clone()),
creator: udaf.creator,
}
.into()
}
}
@@ -110,19 +172,3 @@ fn to_df_accumulator_func(
Ok(Box::new(DfAccumulatorAdaptor::new(accumulator, creator)) as _)
})
}
fn to_df_state_type(func: StateTypeFunction) -> DfStateTypeFunction {
let df_func = move |data_type: &ArrowDataType| {
// DataFusion DataType -> ConcreteDataType
let concrete_data_type = ConcreteDataType::from_arrow_type(data_type);
// evaluate ConcreteDataType
let eval_result = (func)(&concrete_data_type);
// ConcreteDataType -> DataFusion DataType
eval_result
.map(|ts| Arc::new(ts.iter().map(|t| t.as_arrow_type()).collect()))
.map_err(|e| e.into())
};
Arc::new(df_func)
}

2
src/common/query/src/logical_plan/udf.rs
View File

@@ -70,6 +70,8 @@ impl ScalarUdf {
impl From<ScalarUdf> for DfScalarUDF {
fn from(udf: ScalarUdf) -> Self {
// TODO(LFC): remove deprecated
#[allow(deprecated)]
DfScalarUDF::new(
&udf.name,
&udf.signature.into(),

59
src/common/query/src/physical_plan.rs
View File

@@ -21,10 +21,9 @@ use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream}
use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
use datafusion::error::Result as DfResult;
pub use datafusion::execution::context::{SessionContext, TaskContext};
use datafusion::physical_plan::expressions::PhysicalSortExpr;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
pub use datafusion::physical_plan::Partitioning;
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, Statistics};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, PlanProperties};
use datatypes::schema::SchemaRef;
use snafu::ResultExt;
@@ -47,13 +46,9 @@ pub trait PhysicalPlan: Debug + Send + Sync {
/// Get the schema for this physical plan
fn schema(&self) -> SchemaRef;
/// Specifies the output partitioning scheme of this plan
fn output_partitioning(&self) -> Partitioning;
/// returns `Some(keys)` that describes how the output was sorted.
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
None
}
/// Return properties of the output of the [PhysicalPlan], such as output
/// ordering(s), partitioning information etc.
fn properties(&self) -> &PlanProperties;
/// Get a list of child physical plans that provide the input for this plan. The returned list
/// will be empty for leaf nodes, will contain a single value for unary nodes, or two
@@ -107,8 +102,8 @@ impl PhysicalPlan for PhysicalPlanAdapter {
self.schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
self.df_plan.output_partitioning()
fn properties(&self) -> &PlanProperties {
self.df_plan.properties()
}
fn children(&self) -> Vec<PhysicalPlanRef> {
@@ -170,14 +165,6 @@ impl DfPhysicalPlan for DfPhysicalPlanAdapter {
self.0.schema().arrow_schema().clone()
}
fn output_partitioning(&self) -> Partitioning {
self.0.output_partitioning()
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.0.output_ordering()
}
fn children(&self) -> Vec<Arc<dyn DfPhysicalPlan>> {
self.0
.children()
@@ -213,13 +200,13 @@ impl DfPhysicalPlan for DfPhysicalPlanAdapter {
Ok(Box::pin(DfRecordBatchStreamAdapter::new(stream)))
}
fn statistics(&self) -> Statistics {
Statistics::default()
}
fn metrics(&self) -> Option<MetricsSet> {
self.0.metrics()
}
fn properties(&self) -> &PlanProperties {
self.0.properties()
}
}
impl DisplayAs for DfPhysicalPlanAdapter {
@@ -232,10 +219,12 @@ impl DisplayAs for DfPhysicalPlanAdapter {
mod test {
use async_trait::async_trait;
use common_recordbatch::{RecordBatch, RecordBatches};
use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
use datafusion::datasource::{DefaultTableSource, TableProvider as DfTableProvider, TableType};
use datafusion::execution::context::{SessionContext, SessionState};
use datafusion::physical_plan::collect;
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::{collect, ExecutionMode};
use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
use datafusion_expr::{Expr, TableSource};
use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
@@ -272,10 +261,13 @@ mod test {
_filters: &[Expr],
_limit: Option<usize>,
) -> DfResult<Arc<dyn DfPhysicalPlan>> {
let schema = Schema::try_from(self.schema()).unwrap();
let my_plan = Arc::new(MyExecutionPlan {
schema: Arc::new(schema),
});
let schema = Arc::new(Schema::try_from(self.schema()).unwrap());
let properties = PlanProperties::new(
EquivalenceProperties::new(schema.arrow_schema().clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
);
let my_plan = Arc::new(MyExecutionPlan { schema, properties });
let df_plan = DfPhysicalPlanAdapter(my_plan);
Ok(Arc::new(df_plan))
}
@@ -289,9 +281,10 @@ mod test {
}
}
#[derive(Debug)]
#[derive(Debug, Clone)]
struct MyExecutionPlan {
schema: SchemaRef,
properties: PlanProperties,
}
impl PhysicalPlan for MyExecutionPlan {
@@ -303,8 +296,8 @@ mod test {
self.schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
Partitioning::UnknownPartitioning(1)
fn properties(&self) -> &PlanProperties {
&self.properties
}
fn children(&self) -> Vec<PhysicalPlanRef> {
@@ -312,7 +305,7 @@ mod test {
}
fn with_new_children(&self, _children: Vec<PhysicalPlanRef>) -> Result<PhysicalPlanRef> {
unimplemented!()
Ok(Arc::new(self.clone()))
}
fn execute(
@@ -381,7 +374,7 @@ mod test {
let plan = PhysicalPlanAdapter::new(
Arc::new(Schema::try_from(df_schema.clone()).unwrap()),
Arc::new(EmptyExec::new(true, df_schema.clone())),
Arc::new(EmptyExec::new(df_schema.clone())),
);
let _ = plan.df_plan.as_any().downcast_ref::<EmptyExec>().unwrap();

12
src/common/query/src/signature.rs
View File

@@ -31,6 +31,8 @@ pub enum TypeSignature {
// A function such as `array` is `VariadicEqual`
// The first argument decides the type used for coercion
VariadicEqual,
/// One or more arguments with arbitrary types
VariadicAny,
/// fixed number of arguments of an arbitrary but equal type out of a list of valid types
// A function of one argument of f64 is `Uniform(1, vec![ConcreteDataType::Float64])`
// A function of one argument of f64 or f32 is `Uniform(1, vec![ConcreteDataType::Float32, ConcreteDataType::Float64])`
@@ -79,6 +81,15 @@ impl Signature {
volatility,
}
}
/// variadic_any - Creates a variadic signature that represents an arbitrary number of arguments of any type.
pub fn variadic_any(volatility: Volatility) -> Self {
Self {
type_signature: TypeSignature::VariadicAny,
volatility,
}
}
/// uniform - Creates a function with a fixed number of arguments of the same type, which must be from valid_types.
pub fn uniform(
arg_count: usize,
@@ -131,6 +142,7 @@ impl From<TypeSignature> for DfTypeSignature {
TypeSignature::OneOf(ts) => {
DfTypeSignature::OneOf(ts.into_iter().map(Into::into).collect())
}
TypeSignature::VariadicAny => DfTypeSignature::VariadicAny,
}
}
}

1
src/common/recordbatch/Cargo.toml
View File

@@ -11,6 +11,7 @@ workspace = true
arc-swap = "1.6"
common-error.workspace = true
common-macro.workspace = true
common-telemetry.workspace = true
datafusion.workspace = true
datafusion-common.workspace = true
datatypes.workspace = true

6
src/common/recordbatch/src/adapter.rs
View File

@@ -103,7 +103,7 @@ where
"Trying to cast a RecordBatch into an incompatible schema. RecordBatch: {}, Target: {}",
projected_column.schema(),
projected_schema,
))));
)), None));
}
let mut columns = Vec::with_capacity(projected_schema.fields.len());
@@ -218,6 +218,10 @@ impl RecordBatchStreamAdapter {
}
impl RecordBatchStream for RecordBatchStreamAdapter {
fn name(&self) -> &str {
"RecordBatchStreamAdapter"
}
fn schema(&self) -> SchemaRef {
self.schema.clone()
}

8
src/common/recordbatch/src/error.rs
View File

@@ -18,6 +18,7 @@ use std::any::Any;
use common_error::ext::{BoxedError, ErrorExt};
use common_error::status_code::StatusCode;
use common_macro::stack_trace_debug;
use datafusion_common::ScalarValue;
use datatypes::prelude::ConcreteDataType;
use snafu::{Location, Snafu};
@@ -69,8 +70,9 @@ pub enum Error {
location: Location,
},
#[snafu(display("Failed to init Recordbatch stream"))]
InitRecordbatchStream {
#[snafu(display("Failed to convert {v:?} to Arrow scalar"))]
ToArrowScalar {
v: ScalarValue,
#[snafu(source)]
error: datafusion_common::DataFusionError,
location: Location,
@@ -128,7 +130,7 @@ impl ErrorExt for Error {
| Error::CreateRecordBatches { .. }
| Error::PollStream { .. }
| Error::Format { .. }
| Error::InitRecordbatchStream { .. }
| Error::ToArrowScalar { .. }
| Error::ColumnNotExists { .. }
| Error::ProjectArrowRecordBatch { .. }
| Error::ArrowCompute { .. } => StatusCode::Internal,

10
src/common/recordbatch/src/filter.rs
View File

@@ -22,7 +22,7 @@ use datafusion_common::ScalarValue;
use datatypes::vectors::VectorRef;
use snafu::ResultExt;
use crate::error::{ArrowComputeSnafu, Result, UnsupportedOperationSnafu};
use crate::error::{ArrowComputeSnafu, Result, ToArrowScalarSnafu, UnsupportedOperationSnafu};
/// An inplace expr evaluator for simple filter. Only support
/// - `col` `op` `literal`
@@ -69,9 +69,10 @@ impl SimpleFilterEvaluator {
_ => return None,
};
let literal = rhs.to_scalar().ok()?;
Some(Self {
column_name: lhs.name.clone(),
literal: rhs.clone().to_scalar(),
literal,
op,
})
}
@@ -85,7 +86,10 @@ impl SimpleFilterEvaluator {
}
pub fn evaluate_scalar(&self, input: &ScalarValue) -> Result<bool> {
let result = self.evaluate_datum(&input.to_scalar())?;
let input = input
.to_scalar()
.with_context(|_| ToArrowScalarSnafu { v: input.clone() })?;
let result = self.evaluate_datum(&input)?;
Ok(result.value(0))
}

8
src/common/recordbatch/src/lib.rs
View File

@@ -37,6 +37,10 @@ pub use recordbatch::RecordBatch;
use snafu::{ensure, ResultExt};
pub trait RecordBatchStream: Stream<Item = Result<RecordBatch>> {
fn name(&self) -> &str {
"RecordBatchStream"
}
fn schema(&self) -> SchemaRef;
fn output_ordering(&self) -> Option<&[OrderOption]>;
@@ -243,6 +247,10 @@ impl<S> RecordBatchStreamWrapper<S> {
impl<S: Stream<Item = Result<RecordBatch>> + Unpin> RecordBatchStream
for RecordBatchStreamWrapper<S>
{
fn name(&self) -> &str {
"RecordBatchStreamWrapper"
}
fn schema(&self) -> SchemaRef {
self.schema.clone()
}

3
src/common/substrait/Cargo.toml
View File

@@ -12,7 +12,9 @@ async-trait.workspace = true
bytes.workspace = true
catalog.workspace = true
common-error.workspace = true
common-function.workspace = true
common-macro.workspace = true
common-telemetry.workspace = true
datafusion.workspace = true
datafusion-common.workspace = true
datafusion-expr.workspace = true
@@ -20,6 +22,7 @@ datafusion-substrait.workspace = true
datatypes.workspace = true
promql.workspace = true
prost.workspace = true
session.workspace = true
snafu.workspace = true
[dependencies.substrait_proto]

29
src/common/substrait/src/df_substrait.rs
View File

@@ -16,18 +16,24 @@ use std::sync::Arc;
use async_trait::async_trait;
use bytes::{Buf, Bytes, BytesMut};
use datafusion::catalog::CatalogList;
use common_function::function_registry::FUNCTION_REGISTRY;
use common_function::scalars::udf::create_udf;
use datafusion::catalog::CatalogProviderList;
use datafusion::execution::context::SessionState;
use datafusion::execution::runtime_env::RuntimeEnv;
use datafusion::execution::FunctionRegistry;
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_expr::LogicalPlan;
use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
use datafusion_substrait::logical_plan::producer::to_substrait_plan;
use datafusion_substrait::substrait::proto::Plan;
use prost::Message;
use session::context::QueryContextRef;
use snafu::ResultExt;
use substrait_proto::proto::Plan;
use crate::error::{DecodeDfPlanSnafu, DecodeRelSnafu, EncodeDfPlanSnafu, EncodeRelSnafu, Error};
use crate::error::{
DFInternalSnafu, DecodeDfPlanSnafu, DecodeRelSnafu, EncodeDfPlanSnafu, EncodeRelSnafu, Error,
};
use crate::extension_serializer::ExtensionSerializer;
use crate::SubstraitPlan;
@@ -42,17 +48,20 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor {
async fn decode<B: Buf + Send>(
&self,
message: B,
catalog_list: Arc<dyn CatalogList>,
catalog: &str,
schema: &str,
catalog_list: Arc<dyn CatalogProviderList>,
mut state: SessionState,
query_ctx: QueryContextRef,
) -> Result<Self::Plan, Self::Error> {
let state_config = SessionConfig::new().with_default_catalog_and_schema(catalog, schema);
let state = SessionState::new_with_config_rt(state_config, Arc::new(RuntimeEnv::default()))
.with_serializer_registry(Arc::new(ExtensionSerializer));
// substrait decoder will look up the UDFs in SessionState, so we need to register them
for func in FUNCTION_REGISTRY.functions() {
let udf = Arc::new(create_udf(func, query_ctx.clone(), Default::default()).into());
state.register_udf(udf).context(DFInternalSnafu)?;
}
let mut context = SessionContext::new_with_state(state);
context.register_catalog_list(catalog_list);
let plan = Plan::decode(message).context(DecodeRelSnafu)?;
let df_plan = from_substrait_plan(&mut context, &plan)
let df_plan = from_substrait_plan(&context, &plan)
.await
.context(DecodeDfPlanSnafu)?;
Ok(df_plan)

13
src/common/substrait/src/extension_serializer.rs
View File

@@ -19,7 +19,7 @@ use datafusion::execution::registry::SerializerRegistry;
use datafusion_common::DataFusionError;
use datafusion_expr::UserDefinedLogicalNode;
use promql::extension_plan::{
EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize,
EmptyMetric, InstantManipulate, RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize,
};
pub struct ExtensionSerializer;
@@ -50,6 +50,13 @@ impl SerializerRegistry for ExtensionSerializer {
.expect("Failed to downcast to RangeManipulate");
Ok(range_manipulate.serialize())
}
name if name == ScalarCalculate::name() => {
let scalar_calculate = node
.as_any()
.downcast_ref::<ScalarCalculate>()
.expect("Failed to downcast to ScalarCalculate");
Ok(scalar_calculate.serialize())
}
name if name == SeriesDivide::name() => {
let series_divide = node
.as_any()
@@ -92,6 +99,10 @@ impl SerializerRegistry for ExtensionSerializer {
let series_divide = SeriesDivide::deserialize(bytes)?;
Ok(Arc::new(series_divide))
}
name if name == ScalarCalculate::name() => {
let scalar_calculate = ScalarCalculate::deserialize(bytes)?;
Ok(Arc::new(scalar_calculate))
}
name if name == EmptyMetric::name() => Err(DataFusionError::Substrait(
"EmptyMetric should not be deserialized".to_string(),
)),

10
src/common/substrait/src/lib.rs
View File

@@ -21,7 +21,9 @@ use std::sync::Arc;
use async_trait::async_trait;
use bytes::{Buf, Bytes};
use datafusion::catalog::CatalogList;
use datafusion::catalog::CatalogProviderList;
use datafusion::execution::context::SessionState;
use session::context::QueryContextRef;
pub use substrait_proto;
pub use crate::df_substrait::DFLogicalSubstraitConvertor;
@@ -35,9 +37,9 @@ pub trait SubstraitPlan {
async fn decode<B: Buf + Send>(
&self,
message: B,
catalog_list: Arc<dyn CatalogList>,
catalog: &str,
schema: &str,
catalog_list: Arc<dyn CatalogProviderList>,
state: SessionState,
query_ctx: QueryContextRef,
) -> Result<Self::Plan, Self::Error>;
fn encode(&self, plan: &Self::Plan) -> Result<Bytes, Self::Error>;

14
src/common/time/src/datetime.rs
View File

@@ -35,11 +35,11 @@ pub struct DateTime(i64);
impl Display for DateTime {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
if let Some(abs_time) = NaiveDateTime::from_timestamp_millis(self.0) {
if let Some(abs_time) = chrono::DateTime::from_timestamp_millis(self.0) {
write!(
f,
"{}",
format_utc_datetime(&abs_time, DATETIME_FORMAT_WITH_TZ)
format_utc_datetime(&abs_time.naive_utc(), DATETIME_FORMAT_WITH_TZ)
)
} else {
write!(f, "DateTime({})", self.0)
@@ -55,7 +55,7 @@ impl From<DateTime> for serde_json::Value {
impl From<NaiveDateTime> for DateTime {
fn from(value: NaiveDateTime) -> Self {
DateTime::from(value.timestamp_millis())
DateTime::from(value.and_utc().timestamp_millis())
}
}
@@ -87,13 +87,15 @@ impl DateTime {
pub fn from_str(s: &str, timezone: Option<&Timezone>) -> Result<Self> {
let s = s.trim();
let timestamp_millis = if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
dt.naive_utc().timestamp_millis()
dt.naive_utc().and_utc().timestamp_millis()
} else if let Ok(d) = NaiveDateTime::parse_from_str(s, DATETIME_FORMAT) {
match datetime_to_utc(&d, get_timezone(timezone)) {
LocalResult::None => {
return InvalidDateStrSnafu { raw: s }.fail();
}
LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => utc.timestamp_millis(),
LocalResult::Single(t) | LocalResult::Ambiguous(t, _) => {
t.and_utc().timestamp_millis()
}
}
} else if let Ok(v) = chrono::DateTime::parse_from_str(s, DATETIME_FORMAT_WITH_TZ) {
v.timestamp_millis()
@@ -116,7 +118,7 @@ impl DateTime {
/// Convert to [NaiveDateTime].
pub fn to_chrono_datetime(&self) -> Option<NaiveDateTime> {
NaiveDateTime::from_timestamp_millis(self.0)
chrono::DateTime::from_timestamp_millis(self.0).map(|x| x.naive_utc())
}
/// Format DateTime for given format and timezone.

20
src/common/time/src/timestamp.rs
View File

@@ -357,7 +357,7 @@ impl Timestamp {
pub fn to_chrono_datetime(&self) -> Option<NaiveDateTime> {
let (sec, nsec) = self.split();
NaiveDateTime::from_timestamp_opt(sec, nsec)
chrono::DateTime::from_timestamp(sec, nsec).map(|x| x.naive_utc())
}
pub fn to_chrono_datetime_with_timezone(&self, tz: Option<&Timezone>) -> Option<NaiveDateTime> {
@@ -380,8 +380,8 @@ impl Timestamp {
}
pub fn from_chrono_datetime(ndt: NaiveDateTime) -> Option<Self> {
let sec = ndt.timestamp();
let nsec = ndt.timestamp_subsec_nanos();
let sec = ndt.and_utc().timestamp();
let nsec = ndt.and_utc().timestamp_subsec_nanos();
Timestamp::from_splits(sec, nsec)
}
@@ -1063,9 +1063,9 @@ mod tests {
let _ = Timestamp::new(i64::MAX, TimeUnit::Nanosecond).split();
let _ = Timestamp::new(i64::MIN, TimeUnit::Nanosecond).split();
let (sec, nsec) = Timestamp::new(i64::MIN, TimeUnit::Nanosecond).split();
let time = NaiveDateTime::from_timestamp_opt(sec, nsec).unwrap();
assert_eq!(sec, time.timestamp());
assert_eq!(nsec, time.timestamp_subsec_nanos());
let time = DateTime::from_timestamp(sec, nsec).unwrap().naive_utc();
assert_eq!(sec, time.and_utc().timestamp());
assert_eq!(nsec, time.and_utc().timestamp_subsec_nanos());
}
#[test]
@@ -1159,12 +1159,12 @@ mod tests {
#[test]
fn test_subtract_timestamp() {
assert_eq!(
Some(chrono::Duration::milliseconds(42)),
chrono::Duration::try_milliseconds(42),
Timestamp::new_millisecond(100).sub(&Timestamp::new_millisecond(58))
);
assert_eq!(
Some(chrono::Duration::milliseconds(-42)),
chrono::Duration::try_milliseconds(-42),
Timestamp::new_millisecond(58).sub(&Timestamp::new_millisecond(100))
);
}
@@ -1286,8 +1286,8 @@ mod tests {
#[test]
fn test_from_naive_date_time() {
let naive_date_time_min = NaiveDateTime::MIN;
let naive_date_time_max = NaiveDateTime::MAX;
let naive_date_time_min = NaiveDateTime::MIN.and_utc();
let naive_date_time_max = NaiveDateTime::MAX.and_utc();
let min_sec = Timestamp::new_second(naive_date_time_min.timestamp());
let max_sec = Timestamp::new_second(naive_date_time_max.timestamp());

6
src/common/time/src/timezone.rs
View File

@@ -15,7 +15,7 @@
use std::fmt::Display;
use std::str::FromStr;
use chrono::{FixedOffset, NaiveDateTime, TimeZone};
use chrono::{FixedOffset, TimeZone};
use chrono_tz::{OffsetComponents, Tz};
use once_cell::sync::OnceCell;
use snafu::{OptionExt, ResultExt};
@@ -114,7 +114,9 @@ impl Timezone {
match self {
Self::Offset(offset) => offset.local_minus_utc().into(),
Self::Named(tz) => {
let datetime = NaiveDateTime::from_timestamp_opt(0, 0).unwrap();
let datetime = chrono::DateTime::from_timestamp(0, 0)
.map(|x| x.naive_utc())
.expect("invalid timestamp");
let datetime = tz.from_utc_datetime(&datetime);
let utc_offset = datetime.offset().base_utc_offset();
let dst_offset = datetime.offset().dst_offset();

5
src/common/time/src/util.rs
View File

@@ -69,7 +69,10 @@ pub fn current_time_rfc3339() -> String {
/// Returns the yesterday time in rfc3339 format.
pub fn yesterday_rfc3339() -> String {
let now = chrono::Utc::now();
let day_before = now - chrono::Duration::days(1);
let day_before = now
- chrono::Duration::try_days(1).unwrap_or_else(|| {
panic!("now time ('{now}') is too early to calculate the day before")
});
day_before.to_rfc3339()
}

38
src/datanode/src/region_server.rs
View File

@@ -18,14 +18,14 @@ use std::fmt::Debug;
use std::ops::Deref;
use std::sync::{Arc, Mutex, RwLock};
use api::v1::region::{region_request, QueryRequest, RegionResponse};
use api::region::RegionResponse;
use api::v1::region::{region_request, QueryRequest, RegionResponse as RegionResponseV1};
use api::v1::{ResponseHeader, Status};
use arrow_flight::{FlightData, Ticket};
use async_trait::async_trait;
use bytes::Bytes;
use common_error::ext::BoxedError;
use common_error::status_code::StatusCode;
use common_meta::datanode_manager::HandleResponse;
use common_query::logical_plan::Expr;
use common_query::physical_plan::DfPhysicalPlanAdapter;
use common_query::{DfPhysicalPlan, OutputData};
@@ -36,7 +36,7 @@ use common_telemetry::tracing_context::{FutureExt, TracingContext};
use common_telemetry::{info, warn};
use dashmap::DashMap;
use datafusion::catalog::schema::SchemaProvider;
use datafusion::catalog::{CatalogList, CatalogProvider};
use datafusion::catalog::{CatalogProvider, CatalogProviderList};
use datafusion::datasource::TableProvider;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::SessionState;
@@ -129,7 +129,7 @@ impl RegionServer {
&self,
region_id: RegionId,
request: RegionRequest,
) -> Result<HandleResponse> {
) -> Result<RegionResponse> {
self.inner.handle_request(region_id, request).await
}
@@ -218,7 +218,7 @@ impl RegionServer {
#[async_trait]
impl RegionServerHandler for RegionServer {
async fn handle(&self, request: region_request::Body) -> ServerResult<RegionResponse> {
async fn handle(&self, request: region_request::Body) -> ServerResult<RegionResponseV1> {
let is_parallel = matches!(
request,
region_request::Body::Inserts(_) | region_request::Body::Deletes(_)
@@ -276,7 +276,7 @@ impl RegionServerHandler for RegionServer {
extension.extend(result.extension);
}
Ok(RegionResponse {
Ok(RegionResponseV1 {
header: Some(ResponseHeader {
status: Some(Status {
status_code: StatusCode::Success as _,
@@ -465,7 +465,7 @@ impl RegionServerInner {
&self,
region_id: RegionId,
request: RegionRequest,
) -> Result<HandleResponse> {
) -> Result<RegionResponse> {
let request_type = request.request_type();
let _timer = crate::metrics::HANDLE_REGION_REQUEST_ELAPSED
.with_label_values(&[request_type])
@@ -490,7 +490,7 @@ impl RegionServerInner {
let engine = match self.get_engine(region_id, &region_change)? {
CurrentEngine::Engine(engine) => engine,
CurrentEngine::EarlyReturn(rows) => return Ok(HandleResponse::new(rows)),
CurrentEngine::EarlyReturn(rows) => return Ok(RegionResponse::new(rows)),
};
// Sets corresponding region status to registering/deregistering before the operation.
@@ -505,7 +505,7 @@ impl RegionServerInner {
// Sets corresponding region status to ready.
self.set_region_status_ready(region_id, engine, region_change)
.await?;
Ok(HandleResponse {
Ok(RegionResponse {
affected_rows: result.affected_rows,
extension: result.extension,
})
@@ -643,10 +643,15 @@ impl RegionServerInner {
.await?;
let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider));
let query_engine_ctx = self.query_engine.engine_context(ctx.clone());
// decode substrait plan to logical plan and execute it
let logical_plan = DFLogicalSubstraitConvertor
.decode(Bytes::from(plan), catalog_list, "", "")
.decode(
Bytes::from(plan),
catalog_list,
query_engine_ctx.state().clone(),
ctx.clone(),
)
.await
.context(DecodeLogicalPlanSnafu)?;
@@ -728,7 +733,7 @@ impl DummyCatalogList {
}
}
impl CatalogList for DummyCatalogList {
impl CatalogProviderList for DummyCatalogList {
fn as_any(&self) -> &dyn Any {
self
}
@@ -786,8 +791,8 @@ impl SchemaProvider for DummySchemaProvider {
vec![]
}
async fn table(&self, _name: &str) -> Option<Arc<dyn TableProvider>> {
Some(self.table.clone())
async fn table(&self, _name: &str) -> DfResult<Option<Arc<dyn TableProvider>>> {
Ok(Some(self.table.clone()))
}
fn table_exist(&self, _name: &str) -> bool {
@@ -827,7 +832,10 @@ impl TableProvider for DummyTableProvider {
limit: Option<usize>,
) -> DfResult<Arc<dyn DfPhysicalPlan>> {
let mut request = self.scan_request.lock().unwrap().clone();
request.projection = projection.cloned();
request.projection = match projection {
Some(x) if !x.is_empty() => Some(x.clone()),
_ => None,
};
request.filters = filters.iter().map(|e| Expr::from(e.clone())).collect();
request.limit = limit;

9
src/datanode/src/tests.rs
View File

@@ -16,6 +16,7 @@ use std::any::Any;
use std::sync::Arc;
use std::time::Duration;
use api::region::RegionResponse;
use async_trait::async_trait;
use common_error::ext::BoxedError;
use common_function::function::FunctionRef;
@@ -31,7 +32,7 @@ use query::query_engine::DescribeResult;
use query::{QueryEngine, QueryEngineContext};
use session::context::QueryContextRef;
use store_api::metadata::RegionMetadataRef;
use store_api::region_engine::{RegionEngine, RegionHandleResult, RegionRole, SetReadonlyResponse};
use store_api::region_engine::{RegionEngine, RegionRole, SetReadonlyResponse};
use store_api::region_request::{AffectedRows, RegionRequest};
use store_api::storage::{RegionId, ScanRequest};
use table::TableRef;
@@ -166,18 +167,18 @@ impl RegionEngine for MockRegionEngine {
&self,
region_id: RegionId,
request: RegionRequest,
) -> Result<RegionHandleResult, BoxedError> {
) -> Result<RegionResponse, BoxedError> {
if let Some(delay) = self.handle_request_delay {
tokio::time::sleep(delay).await;
}
if let Some(mock_fn) = &self.handle_request_mock_fn {
return mock_fn(region_id, request)
.map_err(BoxedError::new)
.map(RegionHandleResult::new);
.map(RegionResponse::new);
};
let _ = self.sender.send((region_id, request)).await;
Ok(RegionHandleResult::new(0))
Ok(RegionResponse::new(0))
}
async fn handle_query(

7
src/datatypes/src/error.rs
View File

@@ -139,6 +139,13 @@ pub enum Error {
error: arrow::error::ArrowError,
location: Location,
},
#[snafu(display("Failed to convert Arrow array to scalars"))]
ConvertArrowArrayToScalars {
#[snafu(source)]
error: datafusion_common::DataFusionError,
location: Location,
},
}
impl ErrorExt for Error {

6
src/datatypes/src/scalars.rs
View File

@@ -437,10 +437,8 @@ mod tests {
#[test]
fn test_list_value_scalar() {
let list_value = ListValue::new(
Some(Box::new(vec![Value::Int32(123)])),
ConcreteDataType::int32_datatype(),
);
let list_value =
ListValue::new(vec![Value::Int32(123)], ConcreteDataType::int32_datatype());
let list_ref = ListValueRef::Ref { val: &list_value };
assert_eq!(list_ref, list_value.as_scalar_ref());
assert_eq!(list_value, list_ref.to_owned_scalar());

4
src/datatypes/src/types/list_type.rs
View File

@@ -61,7 +61,7 @@ impl DataType for ListType {
}
fn default_value(&self) -> Value {
Value::List(ListValue::new(None, *self.item_type.clone()))
Value::List(ListValue::new(vec![], *self.item_type.clone()))
}
fn as_arrow_type(&self) -> ArrowDataType {
@@ -95,7 +95,7 @@ mod tests {
assert_eq!("List<Boolean>", t.name());
assert_eq!(LogicalTypeId::List, t.logical_type_id());
assert_eq!(
Value::List(ListValue::new(None, ConcreteDataType::boolean_datatype())),
Value::List(ListValue::new(vec![], ConcreteDataType::boolean_datatype())),
t.default_value()
);
assert_eq!(

148
src/datatypes/src/value.rs
View File

@@ -17,6 +17,7 @@ use std::fmt::{Display, Formatter};
use std::sync::Arc;
use arrow::datatypes::{DataType as ArrowDataType, Field};
use arrow_array::{Array, ListArray};
use common_base::bytes::{Bytes, StringBytes};
use common_decimal::Decimal128;
use common_telemetry::logging;
@@ -31,8 +32,7 @@ pub use ordered_float::OrderedFloat;
use serde::{Deserialize, Serialize};
use snafu::{ensure, ResultExt};
use crate::error;
use crate::error::{Error, Result, TryFromValueSnafu};
use crate::error::{self, ConvertArrowArrayToScalarsSnafu, Error, Result, TryFromValueSnafu};
use crate::prelude::*;
use crate::type_id::LogicalTypeId;
use crate::types::{IntervalType, ListType};
@@ -110,9 +110,8 @@ impl Display for Value {
Value::Interval(v) => write!(f, "{}", v.to_iso8601_string()),
Value::Duration(d) => write!(f, "{d}"),
Value::List(v) => {
let default = Box::<Vec<Value>>::default();
let items = v.items().as_ref().unwrap_or(&default);
let items = items
let items = v
.items()
.iter()
.map(|i| i.to_string())
.collect::<Vec<String>>()
@@ -424,9 +423,10 @@ pub fn to_null_scalar_value(output_type: &ConcreteDataType) -> Result<ScalarValu
IntervalType::DayTime(_) => ScalarValue::IntervalDayTime(None),
IntervalType::MonthDayNano(_) => ScalarValue::IntervalMonthDayNano(None),
},
ConcreteDataType::List(_) => {
ScalarValue::List(None, Arc::new(new_item_field(output_type.as_arrow_type())))
}
ConcreteDataType::List(_) => ScalarValue::List(Arc::new(ListArray::new_null(
Arc::new(new_item_field(output_type.as_arrow_type())),
0,
))),
ConcreteDataType::Dictionary(dict) => ScalarValue::Dictionary(
Box::new(dict.key_type().as_arrow_type()),
Box::new(to_null_scalar_value(dict.value_type())?),
@@ -715,9 +715,7 @@ impl TryFrom<Value> for serde_json::Value {
/// List value.
#[derive(Debug, Clone, PartialEq, Hash, Serialize, Deserialize)]
pub struct ListValue {
/// List of nested Values (boxed to reduce size_of(Value))
#[allow(clippy::box_collection)]
items: Option<Box<Vec<Value>>>,
items: Vec<Value>,
/// Inner values datatype, to distinguish empty lists of different datatypes.
/// Restricted by DataFusion, cannot use null datatype for empty list.
datatype: ConcreteDataType,
@@ -726,11 +724,11 @@ pub struct ListValue {
impl Eq for ListValue {}
impl ListValue {
pub fn new(items: Option<Box<Vec<Value>>>, datatype: ConcreteDataType) -> Self {
pub fn new(items: Vec<Value>, datatype: ConcreteDataType) -> Self {
Self { items, datatype }
}
pub fn items(&self) -> &Option<Box<Vec<Value>>> {
pub fn items(&self) -> &[Value] {
&self.items
}
@@ -739,38 +737,30 @@ impl ListValue {
}
fn try_to_scalar_value(&self, output_type: &ListType) -> Result<ScalarValue> {
let vs = if let Some(items) = self.items() {
Some(
items
.iter()
.map(|v| v.try_to_scalar_value(output_type.item_type()))
.collect::<Result<Vec<_>>>()?,
)
} else {
None
};
Ok(ScalarValue::List(
vs,
Arc::new(new_item_field(output_type.item_type().as_arrow_type())),
))
let vs = self
.items
.iter()
.map(|v| v.try_to_scalar_value(output_type.item_type()))
.collect::<Result<Vec<_>>>()?;
Ok(ScalarValue::List(ScalarValue::new_list(
&vs,
&self.datatype.as_arrow_type(),
)))
}
/// use 'the first item size' * 'length of items' to estimate the size.
/// it could be inaccurate.
fn estimated_size(&self) -> usize {
if let Some(items) = &self.items {
if let Some(item) = items.first() {
return item.as_value_ref().data_size() * items.len();
}
}
0
self.items
.first()
.map(|x| x.as_value_ref().data_size() * self.items.len())
.unwrap_or(0)
}
}
impl Default for ListValue {
fn default() -> ListValue {
ListValue::new(None, ConcreteDataType::null_datatype())
ListValue::new(vec![], ConcreteDataType::null_datatype())
}
}
@@ -824,17 +814,14 @@ impl TryFrom<ScalarValue> for Value {
ScalarValue::Binary(b)
| ScalarValue::LargeBinary(b)
| ScalarValue::FixedSizeBinary(_, b) => Value::from(b.map(Bytes::from)),
ScalarValue::List(vs, field) | ScalarValue::Fixedsizelist(vs, field, _) => {
let items = if let Some(vs) = vs {
let vs = vs
.into_iter()
.map(ScalarValue::try_into)
.collect::<Result<_>>()?;
Some(Box::new(vs))
} else {
None
};
let datatype = ConcreteDataType::try_from(field.data_type())?;
ScalarValue::List(array) => {
let datatype = ConcreteDataType::try_from(array.data_type())?;
let items = ScalarValue::convert_array_to_scalar_vec(array.as_ref())
.context(ConvertArrowArrayToScalarsSnafu)?
.into_iter()
.flatten()
.map(|x| x.try_into())
.collect::<Result<Vec<Value>>>()?;
Value::List(ListValue::new(items, datatype))
}
ScalarValue::Date32(d) => d.map(|x| Value::Date(Date::new(x))).unwrap_or(Value::Null),
@@ -891,8 +878,11 @@ impl TryFrom<ScalarValue> for Value {
.map(|v| Value::Decimal128(Decimal128::new(v, p, s)))
.unwrap_or(Value::Null),
ScalarValue::Decimal256(_, _, _)
| ScalarValue::Struct(_, _)
| ScalarValue::Dictionary(_, _) => {
| ScalarValue::Struct(_)
| ScalarValue::FixedSizeList(_)
| ScalarValue::LargeList(_)
| ScalarValue::Dictionary(_, _)
| ScalarValue::Union(_, _, _) => {
return error::UnsupportedArrowTypeSnafu {
arrow_type: v.data_type(),
}
@@ -1382,19 +1372,22 @@ mod tests {
assert_eq!(
Value::List(ListValue::new(
Some(Box::new(vec![Value::Int32(1), Value::Null])),
ConcreteDataType::int32_datatype()
vec![Value::Int32(1), Value::Null],
ConcreteDataType::list_datatype(ConcreteDataType::int32_datatype())
)),
ScalarValue::new_list(
Some(vec![ScalarValue::Int32(Some(1)), ScalarValue::Int32(None)]),
ArrowDataType::Int32,
)
ScalarValue::List(ScalarValue::new_list(
&[ScalarValue::Int32(Some(1)), ScalarValue::Int32(None)],
&ArrowDataType::Int32,
))
.try_into()
.unwrap()
);
assert_eq!(
Value::List(ListValue::new(None, ConcreteDataType::uint32_datatype())),
ScalarValue::new_list(None, ArrowDataType::UInt32)
Value::List(ListValue::new(
vec![],
ConcreteDataType::list_datatype(ConcreteDataType::uint32_datatype())
)),
ScalarValue::List(ScalarValue::new_list(&[], &ArrowDataType::UInt32))
.try_into()
.unwrap()
);
@@ -1664,7 +1657,7 @@ mod tests {
check_type_and_value(
&ConcreteDataType::list_datatype(ConcreteDataType::int32_datatype()),
&Value::List(ListValue::new(
Some(Box::new(vec![Value::Int32(10)])),
vec![Value::Int32(10)],
ConcreteDataType::int32_datatype(),
)),
);
@@ -1829,7 +1822,7 @@ mod tests {
assert_eq!(
json_value,
to_json(Value::List(ListValue {
items: Some(Box::new(vec![Value::Int32(123)])),
items: vec![Value::Int32(123)],
datatype: ConcreteDataType::int32_datatype(),
}))
);
@@ -1897,7 +1890,7 @@ mod tests {
check_as_value_ref!(DateTime, DateTime::new(1034));
let list = ListValue {
items: None,
items: vec![],
datatype: ConcreteDataType::int32_datatype(),
};
assert_eq!(
@@ -1935,7 +1928,7 @@ mod tests {
check_as_correct!(Time::new_second(12), Time, as_time);
check_as_correct!(Duration::new_second(12), Duration, as_duration);
let list = ListValue {
items: None,
items: vec![],
datatype: ConcreteDataType::int32_datatype(),
};
check_as_correct!(ListValueRef::Ref { val: &list }, List, as_list);
@@ -1991,7 +1984,7 @@ mod tests {
);
assert_eq!(
Value::List(ListValue::new(
Some(Box::new(vec![Value::Int8(1), Value::Int8(2)])),
vec![Value::Int8(1), Value::Int8(2)],
ConcreteDataType::int8_datatype(),
))
.to_string(),
@@ -1999,7 +1992,7 @@ mod tests {
);
assert_eq!(
Value::List(ListValue::new(
Some(Box::default()),
vec![],
ConcreteDataType::timestamp_second_datatype(),
))
.to_string(),
@@ -2007,7 +2000,7 @@ mod tests {
);
assert_eq!(
Value::List(ListValue::new(
Some(Box::default()),
vec![],
ConcreteDataType::timestamp_millisecond_datatype(),
))
.to_string(),
@@ -2015,7 +2008,7 @@ mod tests {
);
assert_eq!(
Value::List(ListValue::new(
Some(Box::default()),
vec![],
ConcreteDataType::timestamp_microsecond_datatype(),
))
.to_string(),
@@ -2023,7 +2016,7 @@ mod tests {
);
assert_eq!(
Value::List(ListValue::new(
Some(Box::default()),
vec![],
ConcreteDataType::timestamp_nanosecond_datatype(),
))
.to_string(),
@@ -2253,19 +2246,29 @@ mod tests {
#[test]
fn test_list_value_to_scalar_value() {
let items = Some(Box::new(vec![Value::Int32(-1), Value::Null]));
let items = vec![Value::Int32(-1), Value::Null];
let list = Value::List(ListValue::new(items, ConcreteDataType::int32_datatype()));
let df_list = list
.try_to_scalar_value(&ConcreteDataType::list_datatype(
ConcreteDataType::int32_datatype(),
))
.unwrap();
assert!(matches!(df_list, ScalarValue::List(_, _)));
assert!(matches!(df_list, ScalarValue::List(_)));
match df_list {
ScalarValue::List(vs, field) => {
assert_eq!(ArrowDataType::Int32, *field.data_type());
ScalarValue::List(vs) => {
assert_eq!(
ArrowDataType::List(Arc::new(Field::new_list_field(
ArrowDataType::Int32,
true
))),
*vs.data_type()
);
let vs = vs.unwrap();
let vs = ScalarValue::convert_array_to_scalar_vec(vs.as_ref())
.unwrap()
.into_iter()
.flatten()
.collect::<Vec<_>>();
assert_eq!(
vs,
vec![ScalarValue::Int32(Some(-1)), ScalarValue::Int32(None)]
@@ -2367,10 +2370,10 @@ mod tests {
check_value_ref_size_eq(
&ValueRef::List(ListValueRef::Ref {
val: &ListValue {
items: Some(Box::new(vec![
items: vec![
Value::String("hello world".into()),
Value::String("greptimedb".into()),
])),
],
datatype: ConcreteDataType::string_datatype(),
},
}),
@@ -2387,7 +2390,6 @@ mod tests {
for vec_opt in &data {
if let Some(vec) = vec_opt {
let values = vec.iter().map(|v| Value::from(*v)).collect();
let values = Some(Box::new(values));
let list_value = ListValue::new(values, ConcreteDataType::int32_datatype());
builder.push(Some(ListValueRef::Ref { val: &list_value }));

2
src/datatypes/src/vectors/constant.rs
View File

@@ -223,7 +223,7 @@ mod tests {
assert_eq!(10, c.len());
assert!(c.validity().is_all_valid());
assert!(!c.only_null());
assert_eq!(64, c.memory_size());
assert_eq!(4, c.memory_size());
for i in 0..10 {
assert!(!c.is_null(i));

57
src/datatypes/src/vectors/helper.rs
View File

@@ -26,9 +26,9 @@ use datafusion_common::ScalarValue;
use snafu::{OptionExt, ResultExt};
use crate::data_type::ConcreteDataType;
use crate::error::{self, Result};
use crate::error::{self, ConvertArrowArrayToScalarsSnafu, Result};
use crate::scalars::{Scalar, ScalarVectorBuilder};
use crate::value::{ListValue, ListValueRef};
use crate::value::{ListValue, ListValueRef, Value};
use crate::vectors::{
BinaryVector, BooleanVector, ConstantVector, DateTimeVector, DateVector, Decimal128Vector,
DurationMicrosecondVector, DurationMillisecondVector, DurationNanosecondVector,
@@ -160,19 +160,18 @@ impl Helper {
| ScalarValue::FixedSizeBinary(_, v) => {
ConstantVector::new(Arc::new(BinaryVector::from(vec![v])), length)
}
ScalarValue::List(v, field) | ScalarValue::Fixedsizelist(v, field, _) => {
let item_type = ConcreteDataType::try_from(field.data_type())?;
ScalarValue::List(array) => {
let item_type = ConcreteDataType::try_from(&array.value_type())?;
let mut builder = ListVectorBuilder::with_type_capacity(item_type.clone(), 1);
if let Some(values) = v {
let values = values
.into_iter()
.map(ScalarValue::try_into)
.collect::<Result<_>>()?;
let list_value = ListValue::new(Some(Box::new(values)), item_type);
builder.push(Some(ListValueRef::Ref { val: &list_value }));
} else {
builder.push(None);
}
let values = ScalarValue::convert_array_to_scalar_vec(array.as_ref())
.context(ConvertArrowArrayToScalarsSnafu)?
.into_iter()
.flatten()
.map(ScalarValue::try_into)
.collect::<Result<Vec<Value>>>()?;
builder.push(Some(ListValueRef::Ref {
val: &ListValue::new(values, item_type),
}));
let list_vector = builder.to_vector();
ConstantVector::new(list_vector, length)
}
@@ -236,8 +235,11 @@ impl Helper {
ConstantVector::new(Arc::new(vector), length)
}
ScalarValue::Decimal256(_, _, _)
| ScalarValue::Struct(_, _)
| ScalarValue::Dictionary(_, _) => {
| ScalarValue::Struct(_)
| ScalarValue::FixedSizeList(_)
| ScalarValue::LargeList(_)
| ScalarValue::Dictionary(_, _)
| ScalarValue::Union(_, _, _) => {
return error::ConversionSnafu {
from: format!("Unsupported scalar value: {value}"),
}
@@ -351,7 +353,11 @@ impl Helper {
| ArrowDataType::Dictionary(_, _)
| ArrowDataType::Decimal256(_, _)
| ArrowDataType::Map(_, _)
| ArrowDataType::RunEndEncoded(_, _) => {
| ArrowDataType::RunEndEncoded(_, _)
| ArrowDataType::BinaryView
| ArrowDataType::Utf8View
| ArrowDataType::ListView(_)
| ArrowDataType::LargeListView(_) => {
return error::UnsupportedArrowTypeSnafu {
arrow_type: array.as_ref().data_type().clone(),
}
@@ -396,7 +402,7 @@ mod tests {
TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
};
use arrow::datatypes::{Field, Int32Type};
use arrow::datatypes::Int32Type;
use arrow_array::DictionaryArray;
use common_decimal::Decimal128;
use common_time::time::Time;
@@ -486,13 +492,10 @@ mod tests {
#[test]
fn test_try_from_list_value() {
let value = ScalarValue::List(
Some(vec![
ScalarValue::Int32(Some(1)),
ScalarValue::Int32(Some(2)),
]),
Arc::new(Field::new("item", ArrowDataType::Int32, true)),
);
let value = ScalarValue::List(ScalarValue::new_list(
&[ScalarValue::Int32(Some(1)), ScalarValue::Int32(Some(2))],
&ArrowDataType::Int32,
));
let vector = Helper::try_from_scalar_value(value, 3).unwrap();
assert_eq!(
ConcreteDataType::list_datatype(ConcreteDataType::int32_datatype()),
@@ -501,8 +504,8 @@ mod tests {
assert_eq!(3, vector.len());
for i in 0..vector.len() {
let v = vector.get(i);
let items = v.as_list().unwrap().unwrap().items().as_ref().unwrap();
assert_eq!(vec![Value::Int32(1), Value::Int32(2)], **items);
let items = v.as_list().unwrap().unwrap().items();
assert_eq!(vec![Value::Int32(1), Value::Int32(2)], items);
}
}

43
src/datatypes/src/vectors/list.rs
View File

@@ -114,10 +114,7 @@ impl Vector for ListVector {
let values = (0..vector.len())
.map(|i| vector.get(i))
.collect::<Vec<Value>>();
Value::List(ListValue::new(
Some(Box::new(values)),
self.item_type.clone(),
))
Value::List(ListValue::new(values, self.item_type.clone()))
}
fn get_ref(&self, index: usize) -> ValueRef {
@@ -248,11 +245,8 @@ impl ListVectorBuilder {
}
fn push_list_value(&mut self, list_value: &ListValue) -> Result<()> {
if let Some(items) = list_value.items() {
for item in &**items {
self.values_builder
.try_push_value_ref(item.as_value_ref())?;
}
for v in list_value.items() {
self.values_builder.try_push_value_ref(v.as_value_ref())?;
}
self.finish_list(true);
@@ -496,7 +490,6 @@ pub mod tests {
for vec_opt in data {
if let Some(vec) = vec_opt {
let values = vec.iter().map(|v| Value::from(*v)).collect();
let values = Some(Box::new(values));
let list_value = ListValue::new(values, ConcreteDataType::int32_datatype());
builder.push(Some(ListValueRef::Ref { val: &list_value }));
@@ -576,11 +569,7 @@ pub mod tests {
assert_eq!(
Value::List(ListValue::new(
Some(Box::new(vec![
Value::Int32(1),
Value::Int32(2),
Value::Int32(3)
])),
vec![Value::Int32(1), Value::Int32(2), Value::Int32(3)],
ConcreteDataType::int32_datatype()
)),
list_vector.get(0)
@@ -599,11 +588,7 @@ pub mod tests {
assert_eq!(Value::Null, list_vector.get(1));
assert_eq!(
Value::List(ListValue::new(
Some(Box::new(vec![
Value::Int32(4),
Value::Null,
Value::Int32(6)
])),
vec![Value::Int32(4), Value::Null, Value::Int32(6)],
ConcreteDataType::int32_datatype()
)),
list_vector.get(2)
@@ -680,11 +665,7 @@ pub mod tests {
ListType::new(ConcreteDataType::int32_datatype()).create_mutable_vector(3);
builder.push_value_ref(ValueRef::List(ListValueRef::Ref {
val: &ListValue::new(
Some(Box::new(vec![
Value::Int32(4),
Value::Null,
Value::Int32(6),
])),
vec![Value::Int32(4), Value::Null, Value::Int32(6)],
ConcreteDataType::int32_datatype(),
),
}));
@@ -717,11 +698,7 @@ pub mod tests {
builder.push(None);
builder.push(Some(ListValueRef::Ref {
val: &ListValue::new(
Some(Box::new(vec![
Value::Int32(4),
Value::Null,
Value::Int32(6),
])),
vec![Value::Int32(4), Value::Null, Value::Int32(6)],
ConcreteDataType::int32_datatype(),
),
}));
@@ -772,11 +749,7 @@ pub mod tests {
builder.push(None);
builder.push(Some(ListValueRef::Ref {
val: &ListValue::new(
Some(Box::new(vec![
Value::Int32(4),
Value::Null,
Value::Int32(6),
])),
vec![Value::Int32(4), Value::Null, Value::Int32(6)],
ConcreteDataType::int32_datatype(),
),
}));

4
src/datatypes/src/vectors/primitive.rs
View File

@@ -531,9 +531,9 @@ mod tests {
#[test]
fn test_memory_size() {
let v = Int32Vector::from_slice((0..5).collect::<Vec<i32>>());
assert_eq!(64, v.memory_size());
assert_eq!(20, v.memory_size());
let v = Int64Vector::from(vec![Some(0i64), Some(1i64), Some(2i64), None, None]);
assert_eq!(128, v.memory_size());
assert_eq!(144, v.memory_size());
}
#[test]

3
src/file-engine/Cargo.toml
View File

@@ -12,7 +12,7 @@ test = ["common-test-util"]
workspace = true
[dependencies]
api = { workspace = true, optional = true }
api.workspace = true
async-trait = "0.1"
common-catalog.workspace = true
common-datasource.workspace = true
@@ -25,6 +25,7 @@ common-telemetry.workspace = true
common-test-util = { workspace = true, optional = true }
common-time.workspace = true
datafusion.workspace = true
datafusion-expr.workspace = true
datatypes.workspace = true
futures.workspace = true
object-store.workspace = true

9
src/file-engine/src/engine.rs
View File

@@ -16,6 +16,7 @@ use std::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use api::region::RegionResponse;
use async_trait::async_trait;
use common_catalog::consts::FILE_ENGINE;
use common_error::ext::BoxedError;
@@ -24,7 +25,7 @@ use common_telemetry::{error, info};
use object_store::ObjectStore;
use snafu::{ensure, OptionExt};
use store_api::metadata::RegionMetadataRef;
use store_api::region_engine::{RegionEngine, RegionHandleResult, RegionRole, SetReadonlyResponse};
use store_api::region_engine::{RegionEngine, RegionRole, SetReadonlyResponse};
use store_api::region_request::{
AffectedRows, RegionCloseRequest, RegionCreateRequest, RegionDropRequest, RegionOpenRequest,
RegionRequest,
@@ -60,7 +61,7 @@ impl RegionEngine for FileRegionEngine {
&self,
region_id: RegionId,
request: RegionRequest,
) -> Result<RegionHandleResult, BoxedError> {
) -> Result<RegionResponse, BoxedError> {
self.inner
.handle_request(region_id, request)
.await
@@ -154,7 +155,7 @@ impl EngineInner {
&self,
region_id: RegionId,
request: RegionRequest,
) -> EngineResult<RegionHandleResult> {
) -> EngineResult<RegionResponse> {
let result = match request {
RegionRequest::Create(req) => self.handle_create(region_id, req).await,
RegionRequest::Drop(req) => self.handle_drop(region_id, req).await,
@@ -165,7 +166,7 @@ impl EngineInner {
}
.fail(),
};
result.map(RegionHandleResult::new)
result.map(RegionResponse::new)
}
async fn stop(&self) -> EngineResult<()> {

16
src/file-engine/src/query/file_stream.rs
View File

@@ -23,15 +23,16 @@ use common_query::prelude::Expr;
use common_query::DfPhysicalPlan;
use common_recordbatch::adapter::RecordBatchStreamAdapter;
use common_recordbatch::SendableRecordBatchStream;
use datafusion::common::ToDFSchema;
use datafusion::common::{Statistics, ToDFSchema};
use datafusion::config::TableParquetOptions;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{FileOpener, FileScanConfig, FileStream, ParquetExec};
use datafusion::optimizer::utils::conjunction;
use datafusion::physical_expr::create_physical_expr;
use datafusion::physical_expr::execution_props::ExecutionProps;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::prelude::SessionContext;
use datafusion_expr::utils::conjunction;
use datatypes::arrow::datatypes::Schema as ArrowSchema;
use datatypes::schema::SchemaRef;
use object_store::ObjectStore;
@@ -101,6 +102,7 @@ fn build_record_batch_stream<T: FileOpener + Send + 'static>(
projection: Option<&Vec<usize>>,
limit: Option<usize>,
) -> Result<SendableRecordBatchStream> {
let statistics = Statistics::new_unknown(file_schema.as_ref());
let stream = FileStream::new(
&FileScanConfig {
object_store_url: ObjectStoreUrl::parse("empty://").unwrap(), // won't be used
@@ -109,12 +111,11 @@ fn build_record_batch_stream<T: FileOpener + Send + 'static>(
.iter()
.map(|filename| PartitionedFile::new(filename.to_string(), 0))
.collect::<Vec<_>>()],
statistics: Default::default(),
statistics,
projection: projection.cloned(),
limit,
table_partition_cols: vec![],
output_ordering: vec![],
infinite_source: false,
},
0, // partition: hard-code
opener,
@@ -173,12 +174,11 @@ fn new_parquet_stream_with_exec_plan(
.iter()
.map(|filename| PartitionedFile::new(filename.to_string(), 0))
.collect::<Vec<_>>()],
statistics: Default::default(),
statistics: Statistics::new_unknown(file_schema.as_ref()),
projection: projection.cloned(),
limit: *limit,
table_partition_cols: vec![],
output_ordering: vec![],
infinite_source: false,
};
// build predicate filter
@@ -192,7 +192,7 @@ fn new_parquet_stream_with_exec_plan(
.to_dfschema_ref()
.context(error::ParquetScanPlanSnafu)?;
let filters = create_physical_expr(&expr, &df_schema, &file_schema, &ExecutionProps::new())
let filters = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())
.context(error::ParquetScanPlanSnafu)?;
Some(filters)
} else {
@@ -201,7 +201,7 @@ fn new_parquet_stream_with_exec_plan(
// TODO(ruihang): get this from upper layer
let task_ctx = SessionContext::default().task_ctx();
let parquet_exec = ParquetExec::new(scan_config, filters, None)
let parquet_exec = ParquetExec::new(scan_config, filters, None, TableParquetOptions::default())
.with_parquet_file_reader_factory(Arc::new(DefaultParquetFileReaderFactory::new(
store.clone(),
)));

2
src/flow/clippy.toml
View File

@@ -1,3 +1,5 @@
# Whether to only check for missing documentation in items visible within the current crate. For example, pub(crate) items. (default: false)
# This is a config for clippy::missing_docs_in_private_items
missing-docs-in-crate-items = true
too-many-lines-threshold = 500

37
src/flow/src/transform/expr.rs
View File

@@ -33,6 +33,28 @@ use crate::repr::{ColumnType, RelationType};
use crate::transform::literal::{from_substrait_literal, from_substrait_type};
use crate::transform::FunctionExtensions;
// TODO: found proper place for this
/// ref to `arrow_schema::datatype` for type name
fn typename_to_cdt(name: &str) -> CDT {
match name {
"Int8" => CDT::int8_datatype(),
"Int16" => CDT::int16_datatype(),
"Int32" => CDT::int32_datatype(),
"Int64" => CDT::int64_datatype(),
"UInt8" => CDT::uint8_datatype(),
"UInt16" => CDT::uint16_datatype(),
"UInt32" => CDT::uint32_datatype(),
"UInt64" => CDT::uint64_datatype(),
"Float32" => CDT::float32_datatype(),
"Float64" => CDT::float64_datatype(),
"Boolean" => CDT::boolean_datatype(),
"String" => CDT::string_datatype(),
"Date" => CDT::date_datatype(),
"Timestamp" => CDT::timestamp_second_datatype(),
_ => CDT::null_datatype(),
}
}
impl TypedExpr {
/// Convert ScalarFunction into Flow's ScalarExpr
pub fn from_substrait_scalar_func(
@@ -87,6 +109,21 @@ impl TypedExpr {
Ok(TypedExpr::new(arg.call_unary(func), ret_type))
}
2 if fn_name == "arrow_cast" => {
let cast_to = arg_exprs[1]
.clone()
.as_literal()
.and_then(|lit| lit.as_string())
.with_context(|| InvalidQuerySnafu {
reason: "array_cast's second argument must be a literal string",
})?;
let cast_to = typename_to_cdt(&cast_to);
let func = UnaryFunc::Cast(cast_to);
let arg = arg_exprs[0].clone();
let ret_type = ColumnType::new_nullable(func.signature().output.clone());
Ok(TypedExpr::new(arg.call_unary(func), ret_type))
}
// because variadic function can also have 2 arguments, we need to check if it's a variadic function first
2 if VariadicFunc::from_str_and_types(fn_name, &arg_types).is_err() => {
let (func, signature) =

80
src/frontend/src/heartbeat/handler/tests.rs
View File

@@ -22,6 +22,7 @@ use common_meta::heartbeat::handler::{
};
use common_meta::heartbeat::mailbox::{HeartbeatMailbox, MessageMeta};
use common_meta::instruction::{CacheIdent, Instruction};
use common_meta::key::schema_name::{SchemaName, SchemaNameKey};
use common_meta::key::table_info::TableInfoKey;
use common_meta::key::TableMetaKey;
use partition::manager::TableRouteCacheInvalidator;
@@ -53,6 +54,27 @@ impl TableRouteCacheInvalidator for MockTableRouteCacheInvalidator {
}
}
pub fn test_message_meta(id: u64, subject: &str, to: &str, from: &str) -> MessageMeta {
MessageMeta {
id,
subject: subject.to_string(),
to: to.to_string(),
from: from.to_string(),
}
}
async fn handle_instruction(
executor: Arc<dyn HeartbeatResponseHandlerExecutor>,
mailbox: Arc<HeartbeatMailbox>,
instruction: Instruction,
) {
let response = HeartbeatResponse::default();
let mut ctx: HeartbeatResponseHandlerContext =
HeartbeatResponseHandlerContext::new(mailbox, response);
ctx.incoming_message = Some((test_message_meta(1, "hi", "foo", "bar"), instruction));
executor.handle(ctx).await.unwrap();
}
#[tokio::test]
async fn test_invalidate_table_cache_handler() {
let table_id = 1;
@@ -92,23 +114,45 @@ async fn test_invalidate_table_cache_handler() {
.await;
}
pub fn test_message_meta(id: u64, subject: &str, to: &str, from: &str) -> MessageMeta {
MessageMeta {
id,
subject: subject.to_string(),
to: to.to_string(),
from: from.to_string(),
}
}
#[tokio::test]
async fn test_invalidate_schema_key_handler() {
let (catalog, schema) = ("foo", "bar");
let schema_key = SchemaNameKey { catalog, schema };
let inner = HashMap::from([(schema_key.as_raw_key(), 1)]);
let backend = Arc::new(MockKvCacheInvalidator {
inner: Mutex::new(inner),
});
async fn handle_instruction(
executor: Arc<dyn HeartbeatResponseHandlerExecutor>,
mailbox: Arc<HeartbeatMailbox>,
instruction: Instruction,
) {
let response = HeartbeatResponse::default();
let mut ctx: HeartbeatResponseHandlerContext =
HeartbeatResponseHandlerContext::new(mailbox, response);
ctx.incoming_message = Some((test_message_meta(1, "hi", "foo", "bar"), instruction));
executor.handle(ctx).await.unwrap();
let executor = Arc::new(HandlerGroupExecutor::new(vec![Arc::new(
InvalidateTableCacheHandler::new(backend.clone()),
)]));
let (tx, _) = mpsc::channel(8);
let mailbox = Arc::new(HeartbeatMailbox::new(tx));
// removes a valid key
let valid_key = SchemaName {
catalog_name: catalog.to_string(),
schema_name: schema.to_string(),
};
handle_instruction(
executor.clone(),
mailbox.clone(),
Instruction::InvalidateCaches(vec![CacheIdent::SchemaName(valid_key.clone())]),
)
.await;
assert!(!backend
.inner
.lock()
.unwrap()
.contains_key(&schema_key.as_raw_key()));
// removes a invalid key
handle_instruction(
executor,
mailbox,
Instruction::InvalidateCaches(vec![CacheIdent::SchemaName(valid_key)]),
)
.await;
}

13
src/frontend/src/instance.rs
View File

@@ -491,11 +491,20 @@ pub fn check_permission(
// database ops won't be checked
Statement::CreateDatabase(_) | Statement::ShowDatabases(_) | Statement::DropDatabase(_) => {
}
// show create table and alter are not supported yet
Statement::ShowCreateTable(_) | Statement::CreateExternalTable(_) | Statement::Alter(_) => {
Statement::ShowCreateTable(stmt) => {
validate_param(&stmt.table_name, query_ctx)?;
}
Statement::CreateExternalTable(stmt) => {
validate_param(&stmt.name, query_ctx)?;
}
Statement::Alter(stmt) => {
validate_param(stmt.table_name(), query_ctx)?;
}
// set/show variable now only alter/show variable in session
Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
// show charset and show collation won't be checked
Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
Statement::Insert(insert) => {
validate_param(insert.table_name(), query_ctx)?;

11
src/frontend/src/instance/standalone.rs
View File

@@ -14,11 +14,12 @@
use std::sync::Arc;
use api::v1::region::{QueryRequest, RegionRequest, RegionResponse};
use api::region::RegionResponse;
use api::v1::region::{QueryRequest, RegionRequest, RegionResponse as RegionResponseV1};
use async_trait::async_trait;
use client::region::check_response_header;
use common_error::ext::BoxedError;
use common_meta::datanode_manager::{Datanode, DatanodeManager, DatanodeRef, HandleResponse};
use common_meta::datanode_manager::{Datanode, DatanodeManager, DatanodeRef};
use common_meta::error::{self as meta_error, Result as MetaResult};
use common_meta::peer::Peer;
use common_recordbatch::SendableRecordBatchStream;
@@ -49,7 +50,7 @@ impl RegionInvoker {
Arc::new(Self { region_server })
}
async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponse> {
async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponseV1> {
let body = request.body.with_context(|| InvalidRegionRequestSnafu {
reason: "body not found",
})?;
@@ -63,7 +64,7 @@ impl RegionInvoker {
#[async_trait]
impl Datanode for RegionInvoker {
async fn handle(&self, request: RegionRequest) -> MetaResult<HandleResponse> {
async fn handle(&self, request: RegionRequest) -> MetaResult<RegionResponse> {
let span = request
.header
.as_ref()
@@ -79,7 +80,7 @@ impl Datanode for RegionInvoker {
check_response_header(&response.header)
.map_err(BoxedError::new)
.context(meta_error::ExternalSnafu)?;
Ok(HandleResponse::from_region_response(response))
Ok(RegionResponse::from_region_response(response))
}
async fn handle_query(&self, request: QueryRequest) -> MetaResult<SendableRecordBatchStream> {

2
src/index/src/inverted_index/format/writer/single.rs
View File

@@ -120,7 +120,7 @@ where
// update min/max, assume values are appended in lexicographic order
if stats.distinct_count == 1 {
stats.min_value = value.clone();
stats.min_value.clone_from(&value);
}
stats.max_value = value;
}

33
src/meta-client/src/client.rs
View File

@@ -510,9 +510,6 @@ impl MetaClient {
#[cfg(test)]
mod tests {
use api::v1::meta::{HeartbeatRequest, Peer};
use meta_srv::metasrv::SelectorContext;
use meta_srv::selector::{Namespace, Selector, SelectorOptions};
use meta_srv::Result as MetaResult;
use super::*;
use crate::{error, mocks};
@@ -662,36 +659,6 @@ mod tests {
});
}
struct MockSelector;
#[async_trait::async_trait]
impl Selector for MockSelector {
type Context = SelectorContext;
type Output = Vec<Peer>;
async fn select(
&self,
_ns: Namespace,
_ctx: &Self::Context,
_opts: SelectorOptions,
) -> MetaResult<Self::Output> {
Ok(vec![
Peer {
id: 0,
addr: "peer0".to_string(),
},
Peer {
id: 1,
addr: "peer1".to_string(),
},
Peer {
id: 2,
addr: "peer2".to_string(),
},
])
}
}
#[tokio::test]
async fn test_range_get() {
let tc = new_client("test_range_get").await;

12
src/meta-srv/src/metasrv/builder.rs
View File

@@ -65,7 +65,7 @@ use crate::procedure::region_migration::DefaultContextFactory;
use crate::pubsub::PublishRef;
use crate::selector::lease_based::LeaseBasedSelector;
use crate::service::mailbox::MailboxRef;
use crate::service::store::cached_kv::{CheckLeader, LeaderCachedKvBackend};
use crate::service::store::cached_kv::LeaderCachedKvBackend;
use crate::state::State;
use crate::table_meta_alloc::MetasrvPeerAllocator;
@@ -436,13 +436,3 @@ impl Default for MetasrvBuilder {
Self::new()
}
}
struct CheckLeaderByElection(Option<ElectionRef>);
impl CheckLeader for CheckLeaderByElection {
fn check(&self) -> bool {
self.0
.as_ref()
.map_or(false, |election| election.is_leader())
}
}

10
src/meta-srv/src/selector/load_based.rs
View File

@@ -205,11 +205,9 @@ mod tests {
let alive_stat_kvs = filter_out_expired_datanode(stat_kvs, &lease_kvs);
assert_eq!(1, alive_stat_kvs.len());
assert!(alive_stat_kvs
.get(&StatKey {
cluster_id: 1,
node_id: 1
})
.is_some());
assert!(alive_stat_kvs.contains_key(&StatKey {
cluster_id: 1,
node_id: 1
}));
}
}

2
src/meta-srv/src/service/admin.rs
View File

@@ -29,7 +29,7 @@ use std::task::{Context, Poll};
use tonic::body::BoxBody;
use tonic::codegen::{empty_body, http, BoxFuture, Service};
use tonic::transport::NamedService;
use tonic::server::NamedService;
use crate::metasrv::Metasrv;

Some files were not shown because too many files have changed in this diff Show More