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

build(deps): update datafusion to latest and arrow to 51.0 (#3661)

Browse Source 
* chore: update datafusion

* update sqlness case of time.sql

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix: adjust range query partition

* fix: hisogram incorrect result

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix: ignore filter pushdown temporarily

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix: update limit sqlness result

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix: histogram with wrong distribution

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix: update negative ordinal sqlness case

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* feat: bump df to cd7a00b

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* resolve conflicts

* ignore test_range_filter

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix promql exec panic

* fix "select count(*)" exec error

* re-enable the "test_range_filter" test since the filter push down seems not necessary to be removed

* fix: range query schema error

* update sqlness results

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* resolve conflicts

* update datafusion, again

* fix pyo3 compile error, and update some sqlness results

* update decimal sqlness cases

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix: promql literal

* fix udaf tests

* fix filter pushdown sqlness tests

* fix?: test_cast

* fix: rspy test fail due to datafusion `sin` signature change

* rebase main to see if there are any failed tests

* debug ci

* debug ci

* debug ci

* enforce input partition

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* debug ci

* fix ci

* fix ci

* debug ci

* debug ci

* debug ci

* fix sqlness

* feat: do not return error while creating a filter

* chore: remove array from error

* chore: replace todo with unimplemented

* Update src/flow/clippy.toml

Co-authored-by: Yingwen <realevenyag@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
Co-authored-by: Ruihang Xia <waynestxia@gmail.com>
Co-authored-by: WUJingdi <taylor-lagrange@qq.com>
Co-authored-by: discord9 <discord9@163.com>
Co-authored-by: evenyag <realevenyag@gmail.com>
Co-authored-by: tison <wander4096@gmail.com>
This commit is contained in:
LFC
2024-04-18 12:07:18 +00:00
committed by GitHub
co-authored by Yingwen Ruihang Xia WUJingdi discord9 tison
parent 510782261d
commit 314f2704d4
174 changed files with 2869 additions and 2263 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+588 -352
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+28 -21
View File
@@ -73,13 +73,18 @@ clippy.implicit_clone = "warn"
rust.unknown_lints = "deny"
[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 +96,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/MichaelScofield/greptime-proto.git", rev = "bdbd4cfa871ec8d192d3dbabf11debcb2cb67748" }
humantime = "2.1"
humantime-serde = "1.1"
itertools = "0.10"
@@ -115,12 +122,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 +151,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
benchmarks/src/bin/nyc-taxi.rs
+2 -32
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};"),
),
(
src/catalog/src/table_source.rs
+11 -47
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());
}
}
src/common/datasource/Cargo.toml
+1 -1
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"
src/common/datasource/src/file_format/csv.rs
+1 -1
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());
src/common/datasource/src/file_format/tests.rs
+2 -1
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();
src/common/datasource/src/test_util.rs
+3 -3
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,
}
}
src/common/function/src/scalars/aggregate/diff.rs
+5 -17
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()
);
}
src/common/function/src/scalars/aggregate/percentile.rs
+1 -4
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(),
])
}
src/common/function/src/scalars/aggregate/polyval.rs
+1 -4
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(),
])
}
src/common/function/src/scalars/aggregate/scipy_stats_norm_cdf.rs
+1 -4
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(),
])
}
src/common/function/src/scalars/aggregate/scipy_stats_norm_pdf.rs
+1 -4
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(),
])
}
src/common/function/src/scalars/math.rs
+1 -1
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> {
src/common/macro/src/range_fn.rs
+8 -6
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> {
src/common/query/src/columnar_value.rs
+4 -2
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 })?
}
src/common/query/src/lib.rs
+2 -2
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())
}
}
}
src/common/query/src/logical_plan.rs
+8 -5
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) => {
src/common/query/src/logical_plan/accumulator.rs
+2 -2
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
src/common/query/src/logical_plan/expr.rs
+2 -2
View File
@@ -94,10 +94,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());
}
}
src/common/query/src/logical_plan/udaf.rs
+72 -26
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)
}
src/common/query/src/logical_plan/udf.rs
+2
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(),
src/common/query/src/physical_plan.rs
+26 -33
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();
src/common/query/src/signature.rs
+12
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,
}
}
}
src/common/recordbatch/Cargo.toml
+1
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
src/common/recordbatch/src/adapter.rs
+5 -1
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()
}
src/common/recordbatch/src/error.rs
+5 -3
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,
src/common/recordbatch/src/filter.rs
+7 -3
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))
}
src/common/recordbatch/src/lib.rs
+8
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()
}
src/common/substrait/Cargo.toml
+3
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]
src/common/substrait/src/df_substrait.rs
+19 -10
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)
src/common/substrait/src/lib.rs
+6 -4
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>;
src/common/time/src/datetime.rs
+8 -6
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.
src/common/time/src/timestamp.rs
+8 -8
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,7 +380,7 @@ impl Timestamp {
}
pub fn from_chrono_datetime(ndt: NaiveDateTime) -> Option<Self> {
let sec = ndt.timestamp();
let sec = ndt.and_utc().timestamp();
let nsec = ndt.timestamp_subsec_nanos();
Timestamp::from_splits(sec, nsec)
}
@@ -1063,8 +1063,8 @@ 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());
let time = DateTime::from_timestamp(sec, nsec).unwrap().naive_utc();
assert_eq!(sec, time.and_utc().timestamp());
assert_eq!(nsec, time.timestamp_subsec_nanos());
}
@@ -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());
src/common/time/src/timezone.rs
+4 -2
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();
src/common/time/src/util.rs
+4 -1
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()
}
src/datanode/src/region_server.rs
+15 -7
View File
@@ -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;
@@ -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;
src/datatypes/src/error.rs
+7
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 {
src/datatypes/src/scalars.rs
+2 -4
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());
src/datatypes/src/types/list_type.rs
+2 -2
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!(
src/datatypes/src/value.rs
+75 -73
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 }));
src/datatypes/src/vectors/constant.rs
+1 -1
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));
src/datatypes/src/vectors/helper.rs
+30 -27
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);
}
}
src/datatypes/src/vectors/list.rs
+8 -35
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(),
),
}));
src/datatypes/src/vectors/primitive.rs
+2 -2
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]
src/file-engine/Cargo.toml
+1
View File
@@ -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
src/file-engine/src/query/file_stream.rs
+8 -8
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(),
)));
src/flow/clippy.toml
+2
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
src/flow/src/transform/expr.rs
+37
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) =
src/meta-srv/src/service/admin.rs
+1 -1
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;
src/mito2/src/cache/cache_size.rs
Vendored
+1 -1
View File
@@ -137,6 +137,6 @@ mod tests {
fn test_parquet_meta_size() {
let metadata = parquet_meta();
assert_eq!(948, parquet_meta_size(&metadata));
assert_eq!(956, parquet_meta_size(&metadata));
}
}
src/mito2/src/engine/basic_test.rs
+1 -1
View File
@@ -550,7 +550,7 @@ async fn test_region_usage() {
flush_region(&engine, region_id, None).await;
let region_stat = region.region_usage().await;
assert_eq!(region_stat.sst_usage, 2962);
assert_eq!(region_stat.sst_usage, 3010);
// region total usage
// Some memtables may share items.
src/mito2/src/error.rs
+2
View File
@@ -685,7 +685,9 @@ impl ErrorExt for Error {
CleanDir { .. } => StatusCode::Unexpected,
InvalidConfig { .. } => StatusCode::InvalidArguments,
StaleLogEntry { .. } => StatusCode::Unexpected,
FilterRecordBatch { source, .. } => source.status_code(),
Upload { .. } => StatusCode::StorageUnavailable,
BiError { .. } => StatusCode::Internal,
EncodeMemtable { .. } | ReadDataPart { .. } => StatusCode::Internal,
src/mito2/src/memtable/partition_tree/tree.rs
+3 -2
View File
@@ -211,8 +211,9 @@ impl PartitionTree {
};
let filters = predicate
.map(|p| {
p.exprs()
.map(|predicate| {
predicate
.exprs()
.iter()
.filter_map(|f| SimpleFilterEvaluator::try_new(f.df_expr()))
.collect::<Vec<_>>()
src/mito2/src/memtable/time_series.rs
+14 -11
View File
@@ -238,7 +238,9 @@ impl Memtable for TimeSeriesMemtable {
.collect()
};
let iter = self.series_set.iter_series(projection, filters, self.dedup);
let iter = self
.series_set
.iter_series(projection, filters, self.dedup)?;
Ok(Box::new(iter))
}
@@ -348,7 +350,7 @@ impl SeriesSet {
projection: HashSet<ColumnId>,
predicate: Option<Predicate>,
dedup: bool,
) -> Iter {
) -> Result<Iter> {
let primary_key_schema = primary_key_schema(&self.region_metadata);
let primary_key_datatypes = self
.region_metadata
@@ -356,7 +358,7 @@ impl SeriesSet {
.map(|pk| pk.column_schema.data_type.clone())
.collect();
Iter::new(
Iter::try_new(
self.region_metadata.clone(),
self.series.clone(),
projection,
@@ -417,7 +419,7 @@ struct Iter {
impl Iter {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
pub(crate) fn try_new(
metadata: RegionMetadataRef,
series: Arc<SeriesRwLockMap>,
projection: HashSet<ColumnId>,
@@ -426,27 +428,28 @@ impl Iter {
pk_datatypes: Vec<ConcreteDataType>,
codec: Arc<McmpRowCodec>,
dedup: bool,
) -> Self {
let simple_filters = predicate
.map(|p| {
p.exprs()
) -> Result<Self> {
let predicate = predicate
.map(|predicate| {
predicate
.exprs()
.iter()
.filter_map(|f| SimpleFilterEvaluator::try_new(f.df_expr()))
.collect::<Vec<_>>()
})
.unwrap_or_default();
Self {
Ok(Self {
metadata,
series,
projection,
last_key: None,
predicate: simple_filters,
predicate,
pk_schema,
pk_datatypes,
codec,
dedup,
metrics: Metrics::default(),
}
})
}
}
src/mito2/src/sst/parquet/reader.rs
+4 -3
View File
@@ -180,11 +180,12 @@ impl ParquetReaderBuilder {
metrics.build_cost = start.elapsed();
let predicate = if let Some(p) = &self.predicate {
p.exprs()
let predicate = if let Some(predicate) = &self.predicate {
predicate
.exprs()
.iter()
.filter_map(|expr| SimpleFilterEvaluator::try_new(expr.df_expr()))
.collect()
.collect::<Vec<_>>()
} else {
vec![]
};
src/mito2/src/sst/parquet/stats.rs
+12 -2
View File
@@ -18,8 +18,8 @@ use std::borrow::Borrow;
use std::collections::HashSet;
use datafusion::physical_optimizer::pruning::PruningStatistics;
use datafusion_common::Column;
use datatypes::arrow::array::ArrayRef;
use datafusion_common::{Column, ScalarValue};
use datatypes::arrow::array::{ArrayRef, BooleanArray};
use parquet::file::metadata::RowGroupMetaData;
use store_api::storage::ColumnId;
@@ -81,4 +81,14 @@ impl<'a, T: Borrow<RowGroupMetaData>> PruningStatistics for RowGroupPruningStats
let column_id = self.column_id_to_prune(&column.name)?;
self.read_format.null_counts(self.row_groups, column_id)
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
// TODO(LFC): Impl it.
None
}
fn contained(&self, _column: &Column, _values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
// TODO(LFC): Impl it.
None
}
}
src/operator/src/delete.rs
+1
View File
@@ -124,6 +124,7 @@ impl Deleter {
let request_factory = RegionRequestFactory::new(RegionRequestHeader {
tracing_context: TracingContext::from_current_span().to_w3c(),
dbname: ctx.get_db_string(),
..Default::default()
});
let tasks = self
src/operator/src/expr_factory.rs
+11 -4
View File
@@ -224,10 +224,15 @@ fn find_primary_keys(
let columns_pk = columns
.iter()
.filter_map(|x| {
if x.options
.iter()
.any(|o| matches!(o.option, ColumnOption::Unique { is_primary: true }))
{
if x.options.iter().any(|o| {
matches!(
o.option,
ColumnOption::Unique {
is_primary: true,
..
}
)
}) {
Some(x.name.value.clone())
} else {
None
@@ -249,6 +254,7 @@ fn find_primary_keys(
name: _,
columns,
is_primary: true,
..
} => Some(columns.iter().map(|ident| ident.value.clone())),
_ => None,
})
@@ -276,6 +282,7 @@ pub fn find_time_index(constraints: &[TableConstraint]) -> Result<String> {
name: Some(name),
columns,
is_primary: false,
..
} => {
if name.value == TIME_INDEX {
Some(columns.iter().map(|ident| &ident.value))
src/operator/src/insert.rs
+1
View File
@@ -196,6 +196,7 @@ impl Inserter {
let request_factory = RegionRequestFactory::new(RegionRequestHeader {
tracing_context: TracingContext::from_current_span().to_w3c(),
dbname: ctx.get_db_string(),
..Default::default()
});
let tasks = self
src/operator/src/request.rs
+1
View File
@@ -162,6 +162,7 @@ impl Requester {
let request_factory = RegionRequestFactory::new(RegionRequestHeader {
tracing_context: TracingContext::from_current_span().to_w3c(),
dbname: db_string.unwrap_or_else(|| ctx.get_db_string()),
..Default::default()
});
let tasks = requests.into_iter().map(|req_body| {
src/operator/src/statement/copy_table_from.rs
+3 -2
View File
@@ -36,6 +36,7 @@ use datafusion::datasource::physical_plan::{FileOpener, FileScanConfig, FileStre
use datafusion::parquet::arrow::arrow_reader::ArrowReaderMetadata;
use datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion_common::Statistics;
use datatypes::arrow::compute::can_cast_types;
use datatypes::arrow::datatypes::{Schema, SchemaRef};
use datatypes::vectors::Helper;
@@ -184,17 +185,17 @@ impl StatementExecutor {
filename: &str,
file_schema: SchemaRef,
) -> Result<DfSendableRecordBatchStream> {
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
file_schema,
file_groups: vec![vec![PartitionedFile::new(filename.to_string(), 10)]],
statistics: Default::default(),
statistics,
projection: None,
limit: None,
table_partition_cols: vec![],
output_ordering: vec![],
infinite_source: false,
},
0,
opener,
src/operator/src/statement/copy_table_to.rs
+1 -1
View File
@@ -118,7 +118,7 @@ impl StatementExecutor {
let table_source = Arc::new(DefaultTableSource::new(table_provider));
let mut builder = LogicalPlanBuilder::scan_with_filters(
df_table_ref.to_owned_reference(),
df_table_ref,
table_source,
None,
filters.clone(),
src/promql/Cargo.toml
+3
View File
@@ -20,9 +20,12 @@ common-query.workspace = true
common-recordbatch.workspace = true
common-telemetry.workspace = true
datafusion.workspace = true
datafusion-expr.workspace = true
datafusion-functions.workspace = true
datatypes.workspace = true
futures = "0.3"
greptime-proto.workspace = true
itertools.workspace = true
lazy_static.workspace = true
prometheus.workspace = true
promql-parser = "0.1.1"
src/promql/src/extension_plan/empty_metric.rs
+60 -43
View File
@@ -20,16 +20,20 @@ use std::sync::Arc;
use std::task::{Context, Poll};
use datafusion::arrow::array::ArrayRef;
use datafusion::arrow::datatypes::{DataType, Schema as ArrowSchema, TimeUnit};
use datafusion::common::{DFField, DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics};
use datafusion::arrow::datatypes::{DataType, TimeUnit};
use datafusion::common::arrow::datatypes::Field;
use datafusion::common::stats::Precision;
use datafusion::common::{
DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics, TableReference,
};
use datafusion::error::DataFusionError;
use datafusion::execution::context::{SessionState, TaskContext};
use datafusion::logical_expr::{ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::{PhysicalExprRef, PhysicalSortExpr};
use datafusion::physical_expr::{EquivalenceProperties, PhysicalExprRef};
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
SendableRecordBatchStream,
DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, Partitioning, PlanProperties,
RecordBatchStream, SendableRecordBatchStream,
};
use datafusion::physical_planner::PhysicalPlanner;
use datafusion::prelude::{col, lit, Expr};
@@ -66,15 +70,14 @@ impl EmptyMetric {
field_column_name: String,
field_expr: Option<Expr>,
) -> DataFusionResult<Self> {
let qualifier = Some(TableReference::bare(""));
let ts_only_schema = build_ts_only_schema(&time_index_column_name);
let mut fields = vec![ts_only_schema.field(0).clone()];
let mut fields = vec![(qualifier.clone(), Arc::new(ts_only_schema.field(0).clone()))];
if let Some(field_expr) = &field_expr {
let field_data_type = field_expr.get_type(&ts_only_schema)?;
fields.push(DFField::new(
Some(""),
&field_column_name,
field_data_type,
true,
fields.push((
qualifier.clone(),
Arc::new(Field::new(field_column_name, field_data_type, true)),
));
}
let schema = Arc::new(DFSchema::new_with_metadata(fields, HashMap::new())?);
@@ -102,22 +105,23 @@ impl EmptyMetric {
.expr
.as_ref()
.map(|expr| {
physical_planner.create_physical_expr(
expr,
&self.time_index_schema,
&ArrowSchema::from(self.time_index_schema.as_ref()),
session_state,
)
physical_planner.create_physical_expr(expr, &self.time_index_schema, session_state)
})
.transpose()?;
let result_schema: SchemaRef = Arc::new(self.result_schema.as_ref().into());
let properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(result_schema.clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
));
Ok(Arc::new(EmptyMetricExec {
start: self.start,
end: self.end,
interval: self.interval,
time_index_schema: Arc::new(self.time_index_schema.as_ref().into()),
result_schema: Arc::new(self.result_schema.as_ref().into()),
result_schema,
expr: physical_expr,
properties,
metric: ExecutionPlanMetricsSet::new(),
}))
}
@@ -137,7 +141,11 @@ impl UserDefinedLogicalNodeCore for EmptyMetric {
}
fn expressions(&self) -> Vec<Expr> {
vec![]
if let Some(expr) = &self.expr {
vec![expr.clone()]
} else {
vec![]
}
}
fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
@@ -148,8 +156,19 @@ impl UserDefinedLogicalNodeCore for EmptyMetric {
)
}
fn from_template(&self, _expr: &[Expr], _inputs: &[LogicalPlan]) -> Self {
self.clone()
fn from_template(&self, expr: &[Expr], _inputs: &[LogicalPlan]) -> Self {
Self {
start: self.start,
end: self.end,
interval: self.interval,
expr: if !expr.is_empty() {
Some(expr[0].clone())
} else {
None
},
time_index_schema: self.time_index_schema.clone(),
result_schema: self.result_schema.clone(),
}
}
}
@@ -164,7 +183,7 @@ pub struct EmptyMetricExec {
/// Schema of the output record batch
result_schema: SchemaRef,
expr: Option<PhysicalExprRef>,
properties: Arc<PlanProperties>,
metric: ExecutionPlanMetricsSet,
}
@@ -177,12 +196,8 @@ impl ExecutionPlan for EmptyMetricExec {
self.result_schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
Partitioning::UnknownPartitioning(1)
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
None
fn properties(&self) -> &PlanProperties {
self.properties.as_ref()
}
fn maintains_input_order(&self) -> Vec<bool> {
@@ -222,16 +237,15 @@ impl ExecutionPlan for EmptyMetricExec {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
fn statistics(&self) -> DataFusionResult<Statistics> {
let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
let total_byte_size = estimated_row_num * std::mem::size_of::<Millisecond>() as f64;
Statistics {
num_rows: Some(estimated_row_num.floor() as _),
total_byte_size: Some(total_byte_size.floor() as _),
column_statistics: None,
is_exact: true,
}
Ok(Statistics {
num_rows: Precision::Inexact(estimated_row_num.floor() as _),
total_byte_size: Precision::Inexact(total_byte_size.floor() as _),
column_statistics: Statistics::unknown_column(&self.schema()),
})
}
}
@@ -285,21 +299,21 @@ impl Stream for EmptyMetricStream {
let num_rows = time_array.len();
let input_record_batch =
RecordBatch::try_new(self.time_index_schema.clone(), vec![time_array.clone()])
.map_err(DataFusionError::ArrowError)?;
.map_err(|e| DataFusionError::ArrowError(e, None))?;
let mut result_arrays: Vec<ArrayRef> = vec![time_array];
// evaluate the field expr and get the result
if let Some(field_expr) = &self.expr {
result_arrays.push(
field_expr
.evaluate(&input_record_batch)?
.into_array(num_rows),
.evaluate(&input_record_batch)
.and_then(|x| x.into_array(num_rows))?,
);
}
// assemble the output record batch
let batch = RecordBatch::try_new(self.result_schema.clone(), result_arrays)
.map_err(DataFusionError::ArrowError);
.map_err(|e| DataFusionError::ArrowError(e, None));
Poll::Ready(Some(batch))
} else {
@@ -311,14 +325,17 @@ impl Stream for EmptyMetricStream {
/// Build a schema that only contains **millisecond** timestamp column
fn build_ts_only_schema(column_name: &str) -> DFSchema {
let ts_field = DFField::new(
Some(""),
let ts_field = Field::new(
column_name,
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
);
// safety: should not fail (UT covers this)
DFSchema::new_with_metadata(vec![ts_field], HashMap::new()).unwrap()
DFSchema::new_with_metadata(
vec![(Some(TableReference::bare("")), Arc::new(ts_field))],
HashMap::new(),
)
.unwrap()
}
// Convert timestamp column to UNIX epoch second:
src/promql/src/extension_plan/histogram_fold.rs
+48 -42
View File
@@ -24,16 +24,17 @@ use datafusion::arrow::array::AsArray;
use datafusion::arrow::compute::{self, concat_batches, SortOptions};
use datafusion::arrow::datatypes::{DataType, Float64Type, SchemaRef};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::common::stats::Precision;
use datafusion::common::{ColumnStatistics, DFSchema, DFSchemaRef, Statistics};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::TaskContext;
use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::{PhysicalSortExpr, PhysicalSortRequirement};
use datafusion::physical_expr::{EquivalenceProperties, PhysicalSortRequirement};
use datafusion::physical_plan::expressions::{CastExpr as PhyCast, Column as PhyColumn};
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PhysicalExpr,
RecordBatchStream, SendableRecordBatchStream, Statistics,
PlanProperties, RecordBatchStream, SendableRecordBatchStream,
};
use datafusion::prelude::{Column, Expr};
use datatypes::prelude::{ConcreteDataType, DataType as GtDataType};
@@ -141,16 +142,13 @@ impl HistogramFold {
) -> DataFusionResult<()> {
let check_column = |col| {
if !input_schema.has_column_with_unqualified_name(col) {
return Err(DataFusionError::SchemaError(
Err(DataFusionError::SchemaError(
datafusion::common::SchemaError::FieldNotFound {
field: Box::new(Column::new(None::<String>, col)),
valid_fields: input_schema
.fields()
.iter()
.map(|f| f.qualified_column())
.collect(),
valid_fields: input_schema.columns(),
},
));
Box::new(None),
))
} else {
Ok(())
}
@@ -166,25 +164,29 @@ impl HistogramFold {
// safety: those fields are checked in `check_schema()`
let le_column_index = input_schema
.index_of_column_by_name(None, &self.le_column)
.unwrap()
.unwrap();
let field_column_index = input_schema
.index_of_column_by_name(None, &self.field_column)
.unwrap()
.unwrap();
let ts_column_index = input_schema
.index_of_column_by_name(None, &self.ts_column)
.unwrap()
.unwrap();
let output_schema: SchemaRef = Arc::new(self.output_schema.as_ref().into());
let properties = PlanProperties::new(
EquivalenceProperties::new(output_schema.clone()),
Partitioning::UnknownPartitioning(1),
exec_input.properties().execution_mode(),
);
Arc::new(HistogramFoldExec {
le_column_index,
field_column_index,
ts_column_index,
input: exec_input,
quantile: self.quantile.into(),
output_schema: Arc::new(self.output_schema.as_ref().into()),
output_schema,
metric: ExecutionPlanMetricsSet::new(),
properties,
})
}
@@ -195,15 +197,16 @@ impl HistogramFold {
input_schema: &DFSchemaRef,
le_column: &str,
) -> DataFusionResult<DFSchemaRef> {
let mut fields = input_schema.fields().clone();
let fields = input_schema.fields();
// safety: those fields are checked in `check_schema()`
let le_column_idx = input_schema
.index_of_column_by_name(None, le_column)?
.unwrap();
fields.remove(le_column_idx);
let mut new_fields = Vec::with_capacity(fields.len() - 1);
for f in fields {
if f.name() != le_column {
new_fields.push((None, f.clone()));
}
}
Ok(Arc::new(DFSchema::new_with_metadata(
fields,
new_fields,
HashMap::new(),
)?))
}
@@ -220,6 +223,7 @@ pub struct HistogramFoldExec {
ts_column_index: usize,
quantile: f64,
metric: ExecutionPlanMetricsSet,
properties: PlanProperties,
}
impl ExecutionPlan for HistogramFoldExec {
@@ -227,16 +231,8 @@ impl ExecutionPlan for HistogramFoldExec {
self
}
fn schema(&self) -> SchemaRef {
self.output_schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
self.input.output_partitioning()
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.input.output_ordering()
fn properties(&self) -> &PlanProperties {
&self.properties
}
fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> {
@@ -276,8 +272,7 @@ impl ExecutionPlan for HistogramFoldExec {
}
fn required_input_distribution(&self) -> Vec<Distribution> {
// partition on all tag columns, i.e., non-le, non-ts and non-field columns
vec![Distribution::HashPartitioned(self.tag_col_exprs())]
vec![Distribution::SinglePartition; self.children().len()]
}
fn maintains_input_order(&self) -> Vec<bool> {
@@ -302,6 +297,7 @@ impl ExecutionPlan for HistogramFoldExec {
quantile: self.quantile,
output_schema: self.output_schema.clone(),
field_column_index: self.field_column_index,
properties: self.properties.clone(),
}))
}
@@ -343,13 +339,16 @@ impl ExecutionPlan for HistogramFoldExec {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
Statistics {
num_rows: None,
total_byte_size: None,
column_statistics: None,
is_exact: false,
}
fn statistics(&self) -> DataFusionResult<Statistics> {
Ok(Statistics {
num_rows: Precision::Absent,
total_byte_size: Precision::Absent,
column_statistics: vec![
ColumnStatistics::new_unknown();
// plus one more for the removed column by function `convert_schema`
self.schema().all_fields().len() + 1
],
})
}
}
@@ -594,7 +593,7 @@ impl HistogramFoldStream {
self.output_buffered_rows = 0;
RecordBatch::try_new(self.output_schema.clone(), columns)
.map(Some)
.map_err(DataFusionError::ArrowError)
.map_err(|e| DataFusionError::ArrowError(e, None))
}
/// Find the first `+Inf` which indicates the end of the bucket group
@@ -695,6 +694,7 @@ mod test {
use datafusion::arrow::datatypes::{Field, Schema};
use datafusion::common::ToDFSchema;
use datafusion::physical_plan::memory::MemoryExec;
use datafusion::physical_plan::ExecutionMode;
use datafusion::prelude::SessionContext;
use datatypes::arrow_array::StringArray;
@@ -759,7 +759,7 @@ mod test {
#[tokio::test]
async fn fold_overall() {
let memory_exec = Arc::new(prepare_test_data());
let output_schema = Arc::new(
let output_schema: SchemaRef = Arc::new(
(*HistogramFold::convert_schema(
&Arc::new(memory_exec.schema().to_dfschema().unwrap()),
"le",
@@ -769,6 +769,11 @@ mod test {
.clone()
.into(),
);
let properties = PlanProperties::new(
EquivalenceProperties::new(output_schema.clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
);
let fold_exec = Arc::new(HistogramFoldExec {
le_column_index: 1,
field_column_index: 2,
@@ -777,6 +782,7 @@ mod test {
input: memory_exec,
output_schema,
metric: ExecutionPlanMetricsSet::new(),
properties,
});
let session_context = SessionContext::default();
src/promql/src/extension_plan/instant_manipulate.rs
+25 -19
View File
@@ -21,14 +21,14 @@ use std::task::{Context, Poll};
use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray, UInt64Array};
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::common::stats::Precision;
use datafusion::common::{ColumnStatistics, DFSchema, DFSchemaRef};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, Statistics,
};
use datatypes::arrow::compute;
@@ -194,16 +194,17 @@ impl ExecutionPlan for InstantManipulateExec {
self.input.schema()
}
fn output_partitioning(&self) -> Partitioning {
self.input.output_partitioning()
fn properties(&self) -> &PlanProperties {
self.input.properties()
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.input.output_ordering()
fn required_input_distribution(&self) -> Vec<Distribution> {
self.input.required_input_distribution()
}
// Prevent reordering of input
fn maintains_input_order(&self) -> Vec<bool> {
vec![true; self.children().len()]
vec![false; self.children().len()]
}
fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
@@ -262,23 +263,28 @@ impl ExecutionPlan for InstantManipulateExec {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
let input_stats = self.input.statistics();
fn statistics(&self) -> DataFusionResult<Statistics> {
let input_stats = self.input.statistics()?;
let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
let estimated_total_bytes = input_stats
.total_byte_size
.zip(input_stats.num_rows)
.map(|(size, rows)| (size as f64 / rows as f64) * estimated_row_num)
.map(|size| size.floor() as _);
.get_value()
.zip(input_stats.num_rows.get_value())
.map(|(size, rows)| {
Precision::Inexact(((*size as f64 / *rows as f64) * estimated_row_num).floor() as _)
})
.unwrap_or(Precision::Absent);
Statistics {
num_rows: Some(estimated_row_num.floor() as _),
Ok(Statistics {
num_rows: Precision::Inexact(estimated_row_num.floor() as _),
total_byte_size: estimated_total_bytes,
// TODO(ruihang): support this column statistics
column_statistics: None,
is_exact: false,
}
column_statistics: vec![
ColumnStatistics::new_unknown();
self.schema().all_fields().len()
],
})
}
}
@@ -438,7 +444,7 @@ impl InstantManipulateStream {
arrays[self.time_index] = Arc::new(TimestampMillisecondArray::from(aligned_ts));
let result = RecordBatch::try_new(record_batch.schema(), arrays)
.map_err(DataFusionError::ArrowError)?;
.map_err(|e| DataFusionError::ArrowError(e, None))?;
Ok(result)
}
}
src/promql/src/extension_plan/normalize.rs
+5 -10
View File
@@ -23,10 +23,9 @@ use datafusion::common::{DFSchema, DFSchemaRef, Result as DataFusionResult, Stat
use datafusion::error::DataFusionError;
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, RecordBatchStream,
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
};
use datatypes::arrow::array::TimestampMillisecondArray;
@@ -170,12 +169,8 @@ impl ExecutionPlan for SeriesNormalizeExec {
vec![Distribution::SinglePartition]
}
fn output_partitioning(&self) -> Partitioning {
self.input.output_partitioning()
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.input.output_ordering()
fn properties(&self) -> &PlanProperties {
self.input.properties()
}
fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
@@ -223,7 +218,7 @@ impl ExecutionPlan for SeriesNormalizeExec {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
fn statistics(&self) -> DataFusionResult<Statistics> {
self.input.statistics()
}
}
@@ -299,7 +294,7 @@ impl SeriesNormalizeStream {
}
let result = compute::filter_record_batch(&ordered_batch, &BooleanArray::from(filter))
.map_err(DataFusionError::ArrowError)?;
.map_err(|e| DataFusionError::ArrowError(e, None))?;
Ok(result)
}
}
src/promql/src/extension_plan/range_manipulate.rs
+53 -28
View File
@@ -23,14 +23,15 @@ use datafusion::arrow::compute;
use datafusion::arrow::datatypes::{Field, SchemaRef};
use datafusion::arrow::error::ArrowError;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::{DFField, DFSchema, DFSchemaRef};
use datafusion::common::stats::Precision;
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, Statistics,
};
use datafusion::sql::TableReference;
@@ -110,44 +111,55 @@ impl RangeManipulate {
time_index: &str,
field_columns: &[String],
) -> DataFusionResult<DFSchemaRef> {
let mut columns = input_schema.fields().clone();
let columns = input_schema.fields();
let mut new_columns = Vec::with_capacity(columns.len() + 1);
for i in 0..columns.len() {
let x = input_schema.qualified_field(i);
new_columns.push((x.0.cloned(), Arc::new(x.1.clone())));
}
// process time index column
// the raw timestamp field is preserved. And a new timestamp_range field is appended to the last.
let Some(ts_col_index) = input_schema.index_of_column_by_name(None, time_index)? else {
let Some(ts_col_index) = input_schema.index_of_column_by_name(None, time_index) else {
return Err(datafusion::common::field_not_found(
None::<TableReference>,
time_index,
input_schema.as_ref(),
));
};
let ts_col_field = columns[ts_col_index].field();
let ts_col_field = &columns[ts_col_index];
let timestamp_range_field = Field::new(
Self::build_timestamp_range_name(time_index),
RangeArray::convert_field(ts_col_field).data_type().clone(),
ts_col_field.is_nullable(),
);
columns.push(DFField::from(timestamp_range_field));
new_columns.push((None, Arc::new(timestamp_range_field)));
// process value columns
for name in field_columns {
let Some(index) = input_schema.index_of_column_by_name(None, name)? else {
let Some(index) = input_schema.index_of_column_by_name(None, name) else {
return Err(datafusion::common::field_not_found(
None::<TableReference>,
name,
input_schema.as_ref(),
));
};
columns[index] = DFField::from(RangeArray::convert_field(columns[index].field()));
new_columns[index] = (None, Arc::new(RangeArray::convert_field(&columns[index])));
}
Ok(Arc::new(DFSchema::new_with_metadata(
columns,
new_columns,
HashMap::new(),
)?))
}
pub fn to_execution_plan(&self, exec_input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
let output_schema: SchemaRef = SchemaRef::new(self.output_schema.as_ref().into());
let properties = PlanProperties::new(
EquivalenceProperties::new(output_schema.clone()),
exec_input.properties().partitioning.clone(),
exec_input.properties().execution_mode,
);
Arc::new(RangeManipulateExec {
start: self.start,
end: self.end,
@@ -157,8 +169,9 @@ impl RangeManipulate {
time_range_column: self.range_timestamp_name(),
field_columns: self.field_columns.clone(),
input: exec_input,
output_schema: SchemaRef::new(self.output_schema.as_ref().into()),
output_schema,
metric: ExecutionPlanMetricsSet::new(),
properties,
})
}
@@ -247,6 +260,7 @@ pub struct RangeManipulateExec {
input: Arc<dyn ExecutionPlan>,
output_schema: SchemaRef,
metric: ExecutionPlanMetricsSet,
properties: PlanProperties,
}
impl ExecutionPlan for RangeManipulateExec {
@@ -258,12 +272,8 @@ impl ExecutionPlan for RangeManipulateExec {
self.output_schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
self.input.output_partitioning()
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.input.output_ordering()
fn properties(&self) -> &PlanProperties {
&self.properties
}
fn maintains_input_order(&self) -> Vec<bool> {
@@ -274,6 +284,10 @@ impl ExecutionPlan for RangeManipulateExec {
vec![self.input.clone()]
}
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![Distribution::SinglePartition]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
@@ -290,6 +304,7 @@ impl ExecutionPlan for RangeManipulateExec {
output_schema: self.output_schema.clone(),
input: children[0].clone(),
metric: self.metric.clone(),
properties: self.properties.clone(),
}))
}
@@ -333,23 +348,25 @@ impl ExecutionPlan for RangeManipulateExec {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
let input_stats = self.input.statistics();
fn statistics(&self) -> DataFusionResult<Statistics> {
let input_stats = self.input.statistics()?;
let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
let estimated_total_bytes = input_stats
.total_byte_size
.zip(input_stats.num_rows)
.map(|(size, rows)| (size as f64 / rows as f64) * estimated_row_num)
.map(|size| size.floor() as _);
.get_value()
.zip(input_stats.num_rows.get_value())
.map(|(size, rows)| {
Precision::Inexact(((*size as f64 / *rows as f64) * estimated_row_num).floor() as _)
})
.unwrap_or_default();
Statistics {
num_rows: Some(estimated_row_num.floor() as _),
Ok(Statistics {
num_rows: Precision::Inexact(estimated_row_num as _),
total_byte_size: estimated_total_bytes,
// TODO(ruihang): support this column statistics
column_statistics: None,
is_exact: false,
}
column_statistics: Statistics::unknown_column(&self.schema()),
})
}
}
@@ -452,7 +469,7 @@ impl RangeManipulateStream {
RecordBatch::try_new(self.output_schema.clone(), new_columns)
.map(Some)
.map_err(DataFusionError::ArrowError)
.map_err(|e| DataFusionError::ArrowError(e, None))
}
fn calculate_range(&self, input: &RecordBatch) -> (ArrayRef, Vec<(u32, u32)>) {
@@ -500,7 +517,9 @@ mod test {
ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimestampMillisecondType,
};
use datafusion::common::ToDFSchema;
use datafusion::physical_expr::Partitioning;
use datafusion::physical_plan::memory::MemoryExec;
use datafusion::physical_plan::ExecutionMode;
use datafusion::prelude::SessionContext;
use datatypes::arrow::array::TimestampMillisecondArray;
@@ -556,6 +575,11 @@ mod test {
.as_ref()
.into(),
);
let properties = PlanProperties::new(
EquivalenceProperties::new(manipulate_output_schema.clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
);
let normalize_exec = Arc::new(RangeManipulateExec {
start,
end,
@@ -567,6 +591,7 @@ mod test {
time_index_column: time_index,
input: memory_exec,
metric: ExecutionPlanMetricsSet::new(),
properties,
});
let session_context = SessionContext::default();
let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
src/promql/src/extension_plan/series_divide.rs
+5 -19
View File
@@ -24,12 +24,12 @@ use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::error::Result as DataFusionResult;
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::{PhysicalSortExpr, PhysicalSortRequirement};
use datafusion::physical_expr::PhysicalSortRequirement;
use datafusion::physical_plan::expressions::Column as ColumnExpr;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, RecordBatchStream,
SendableRecordBatchStream, Statistics,
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
};
use datatypes::arrow::compute;
use futures::{ready, Stream, StreamExt};
@@ -130,8 +130,8 @@ impl ExecutionPlan for SeriesDivideExec {
self.input.schema()
}
fn output_partitioning(&self) -> Partitioning {
Partitioning::UnknownPartitioning(1)
fn properties(&self) -> &PlanProperties {
self.input.properties()
}
fn required_input_distribution(&self) -> Vec<Distribution> {
@@ -156,10 +156,6 @@ impl ExecutionPlan for SeriesDivideExec {
}
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.input.output_ordering()
}
fn maintains_input_order(&self) -> Vec<bool> {
vec![true; self.children().len()]
}
@@ -212,16 +208,6 @@ impl ExecutionPlan for SeriesDivideExec {
fn metrics(&self) -> Option<MetricsSet> {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
Statistics {
num_rows: None,
total_byte_size: None,
// TODO(ruihang): support this column statistics
column_statistics: None,
is_exact: false,
}
}
}
impl DisplayAs for SeriesDivideExec {
src/promql/src/extension_plan/union_distinct_on.rs
+20 -19
View File
@@ -25,11 +25,11 @@ use datafusion::common::DFSchemaRef;
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
hash_utils, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning,
RecordBatchStream, SendableRecordBatchStream, Statistics,
hash_utils, DisplayAs, DisplayFormatType, Distribution, ExecutionMode, ExecutionPlan,
Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
};
use datatypes::arrow::compute;
use futures::future::BoxFuture;
@@ -91,13 +91,20 @@ impl UnionDistinctOn {
left_exec: Arc<dyn ExecutionPlan>,
right_exec: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
let output_schema: SchemaRef = Arc::new(self.output_schema.as_ref().into());
let properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(output_schema.clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
));
Arc::new(UnionDistinctOnExec {
left: left_exec,
right: right_exec,
compare_keys: self.compare_keys.clone(),
ts_col: self.ts_col.clone(),
output_schema: Arc::new(self.output_schema.as_ref().into()),
output_schema,
metric: ExecutionPlanMetricsSet::new(),
properties,
random_state: RandomState::new(),
})
}
@@ -151,6 +158,7 @@ pub struct UnionDistinctOnExec {
ts_col: String,
output_schema: SchemaRef,
metric: ExecutionPlanMetricsSet,
properties: Arc<PlanProperties>,
/// Shared the `RandomState` for the hashing algorithm
random_state: RandomState,
@@ -169,14 +177,8 @@ impl ExecutionPlan for UnionDistinctOnExec {
vec![Distribution::SinglePartition, Distribution::SinglePartition]
}
fn output_partitioning(&self) -> Partitioning {
Partitioning::UnknownPartitioning(1)
}
/// [UnionDistinctOnExec] will output left first, then right.
/// So the order of the output is not maintained.
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
None
fn properties(&self) -> &PlanProperties {
self.properties.as_ref()
}
fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
@@ -198,6 +200,7 @@ impl ExecutionPlan for UnionDistinctOnExec {
ts_col: self.ts_col.clone(),
output_schema: self.output_schema.clone(),
metric: self.metric.clone(),
properties: self.properties.clone(),
random_state: self.random_state.clone(),
}))
}
@@ -249,10 +252,6 @@ impl ExecutionPlan for UnionDistinctOnExec {
fn metrics(&self) -> Option<MetricsSet> {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
Statistics::default()
}
}
impl DisplayAs for UnionDistinctOnExec {
@@ -472,7 +471,8 @@ fn interleave_batches(
}
// assemble new record batch
RecordBatch::try_new(schema.clone(), interleaved_arrays).map_err(DataFusionError::ArrowError)
RecordBatch::try_new(schema.clone(), interleaved_arrays)
.map_err(|e| DataFusionError::ArrowError(e, None))
}
/// Utility function to take rows from a record batch. Based on [take](datafusion::arrow::compute::take)
@@ -490,9 +490,10 @@ fn take_batch(batch: &RecordBatch, indices: &[usize]) -> DataFusionResult<Record
.iter()
.map(|array| compute::take(array, &indices_array, None))
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(DataFusionError::ArrowError)?;
.map_err(|e| DataFusionError::ArrowError(e, None))?;
let result = RecordBatch::try_new(schema, arrays).map_err(DataFusionError::ArrowError)?;
let result =
RecordBatch::try_new(schema, arrays).map_err(|e| DataFusionError::ArrowError(e, None))?;
Ok(result)
}
src/promql/src/functions/extrapolate_rate.rs
+24 -18
View File
@@ -204,15 +204,17 @@ impl ExtrapolatedRate<false, false> {
}
pub fn scalar_udf(range_length: i64) -> 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(move |input| Self::new(range_length).calc(input)),
}
&(Arc::new(|_: &_| Ok(Arc::new(Self::return_type()))) as _),
&(Arc::new(move |input: &_| Self::new(range_length).calc(input)) as _),
)
}
}
@@ -223,15 +225,17 @@ impl ExtrapolatedRate<true, true> {
}
pub fn scalar_udf(range_length: i64) -> 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(move |input| Self::new(range_length).calc(input)),
}
&(Arc::new(|_: &_| Ok(Arc::new(Self::return_type()))) as _),
&(Arc::new(move |input: &_| Self::new(range_length).calc(input)) as _),
)
}
}
@@ -242,15 +246,17 @@ impl ExtrapolatedRate<true, false> {
}
pub fn scalar_udf(range_length: i64) -> 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(move |input| Self::new(range_length).calc(input)),
}
&(Arc::new(|_: &_| Ok(Arc::new(Self::return_type()))) as _),
&(Arc::new(move |input: &_| Self::new(range_length).calc(input)) as _),
)
}
}
src/promql/src/functions/holt_winters.rs
+8 -6
View File
@@ -68,15 +68,17 @@ impl HoltWinters {
}
pub fn scalar_udf(level: f64, trend: f64) -> 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(move |input| Self::new(level, trend).calc(input)),
}
&(Arc::new(|_: &_| Ok(Arc::new(Self::return_type()))) as _),
&(Arc::new(move |input: &_| Self::new(level, trend).calc(input)) as _),
)
}
fn calc(&self, input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
src/promql/src/functions/idelta.rs
+8 -6
View File
@@ -42,15 +42,17 @@ impl<const IS_RATE: bool> IDelta<IS_RATE> {
}
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 _),
)
}
// time index column and value column
src/promql/src/functions/predict_linear.rs
+8 -6
View File
@@ -44,15 +44,17 @@ impl PredictLinear {
}
pub fn scalar_udf(t: i64) -> 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(move |input| Self::new(t).calc(input)),
}
&(Arc::new(|_: &_| Ok(Arc::new(Self::return_type()))) as _),
&(Arc::new(move |input: &_| Self::new(t).calc(input)) as _),
)
}
// time index column and value column
src/promql/src/functions/quantile.rs
+8 -6
View File
@@ -40,15 +40,17 @@ impl QuantileOverTime {
}
pub fn scalar_udf(quantile: f64) -> 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(move |input| Self::new(quantile).calc(input)),
}
&(Arc::new(|_: &_| Ok(Arc::new(Self::return_type()))) as _),
&(Arc::new(move |input: &_| Self::new(quantile).calc(input)) as _),
)
}
// time index column and value column
src/promql/src/functions/test_util.rs
+10 -2
View File
@@ -32,12 +32,20 @@ pub fn simple_range_udf_runner(
ColumnarValue::Array(Arc::new(input_ts.into_dict())),
ColumnarValue::Array(Arc::new(input_value.into_dict())),
];
let eval_result: Vec<Option<f64>> = extract_array(&(range_fn.fun)(&input).unwrap())
let eval_result: Vec<Option<f64>> = extract_array(&(range_fn.fun())(&input).unwrap())
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.iter()
.collect();
assert_eq!(eval_result, expected)
assert_eq!(eval_result.len(), expected.len());
assert!(eval_result
.iter()
.zip(expected.iter())
.all(|(x, y)| match (*x, *y) {
(Some(x), Some(y)) => (x - y).abs() < 0.0001,
(None, None) => true,
_ => false,
}));
}
src/promql/src/planner.rs
+73 -41
View File
@@ -20,20 +20,23 @@ use std::time::UNIX_EPOCH;
use async_recursion::async_recursion;
use catalog::table_source::DfTableSourceProvider;
use common_query::prelude::GREPTIME_VALUE;
use datafusion::common::{DFSchemaRef, OwnedTableReference, Result as DfResult};
use datafusion::common::{DFSchemaRef, Result as DfResult};
use datafusion::datasource::DefaultTableSource;
use datafusion::logical_expr::expr::{AggregateFunction, Alias, ScalarFunction, ScalarUDF};
use datafusion::logical_expr::expr::{
AggregateFunction, AggregateFunctionDefinition, Alias, ScalarFunction,
};
use datafusion::logical_expr::expr_rewriter::normalize_cols;
use datafusion::logical_expr::{
AggregateFunction as AggregateFunctionEnum, BinaryExpr, BuiltinScalarFunction, Cast, Extension,
LogicalPlan, LogicalPlanBuilder, Operator, ScalarUDF as ScalarUdfDef,
LogicalPlan, LogicalPlanBuilder, Operator, ScalarFunctionDefinition, ScalarUDF as ScalarUdfDef,
};
use datafusion::optimizer::utils::{self, conjunction};
use datafusion::prelude as df_prelude;
use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
use datafusion::scalar::ScalarValue;
use datafusion::sql::TableReference;
use datafusion_expr::utils::conjunction;
use datatypes::arrow::datatypes::DataType as ArrowDataType;
use itertools::Itertools;
use promql_parser::label::{MatchOp, Matcher, Matchers, METRIC_NAME};
use promql_parser::parser::{
token, AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt,
@@ -305,14 +308,14 @@ impl PromPlanner {
let left_field_columns = self.ctx.field_columns.clone();
let mut left_table_ref = self
.table_ref()
.unwrap_or_else(|_| OwnedTableReference::bare(""));
.unwrap_or_else(|_| TableReference::bare(""));
let left_context = self.ctx.clone();
let right_input = self.prom_expr_to_plan(*rhs.clone()).await?;
let right_field_columns = self.ctx.field_columns.clone();
let mut right_table_ref = self
.table_ref()
.unwrap_or_else(|_| OwnedTableReference::bare(""));
.unwrap_or_else(|_| TableReference::bare(""));
let right_context = self.ctx.clone();
// TODO(ruihang): avoid join if left and right are the same table
@@ -332,8 +335,8 @@ impl PromPlanner {
// normal join
if left_table_ref == right_table_ref {
// rename table references to avoid ambiguity
left_table_ref = OwnedTableReference::bare("lhs");
right_table_ref = OwnedTableReference::bare("rhs");
left_table_ref = TableReference::bare("lhs");
right_table_ref = TableReference::bare("rhs");
self.ctx.table_name = Some("lhs".to_string());
}
let mut field_columns =
@@ -349,13 +352,13 @@ impl PromPlanner {
let bin_expr_builder = |_: &String| {
let (left_col_name, right_col_name) = field_columns.next().unwrap();
let left_col = join_plan_schema
.field_with_name(Some(&left_table_ref), left_col_name)
.qualified_field_with_name(Some(&left_table_ref), left_col_name)
.context(DataFusionPlanningSnafu)?
.qualified_column();
.into();
let right_col = join_plan_schema
.field_with_name(Some(&right_table_ref), right_col_name)
.qualified_field_with_name(Some(&right_table_ref), right_col_name)
.context(DataFusionPlanningSnafu)?
.qualified_column();
.into();
let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
let mut binary_expr = binary_expr_builder(
@@ -859,7 +862,7 @@ impl PromPlanner {
Ok(exprs)
}
fn table_ref(&self) -> Result<OwnedTableReference> {
fn table_ref(&self) -> Result<TableReference> {
let table_name = self
.ctx
.table_name
@@ -868,12 +871,12 @@ impl PromPlanner {
// set schema name if `__schema__` is given
let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
TableReference::partial(schema_name, &table_name)
TableReference::partial(schema_name.as_str(), table_name.as_str())
} else {
TableReference::bare(&table_name)
TableReference::bare(table_name.as_str())
};
Ok(table_ref.to_owned_reference())
Ok(table_ref)
}
/// Create a table scan plan and a filter plan with given filter.
@@ -882,7 +885,7 @@ impl PromPlanner {
/// If the filter is empty
async fn create_table_scan_plan(
&mut self,
table_ref: OwnedTableReference,
table_ref: TableReference,
filter: Vec<DfExpr>,
) -> Result<LogicalPlan> {
let provider = self
@@ -1128,7 +1131,9 @@ impl PromPlanner {
right: Box::new(interval_1day_lit_expr),
});
let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
fun: BuiltinScalarFunction::DateTrunc,
func_def: ScalarFunctionDefinition::UDF(
datafusion_functions::datetime::date_trunc(),
),
args: vec![month_lit_expr, self.create_time_index_column_expr()?],
});
let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
@@ -1137,21 +1142,30 @@ impl PromPlanner {
right: Box::new(the_1month_minus_1day_expr),
});
let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
fun: BuiltinScalarFunction::DatePart,
func_def: ScalarFunctionDefinition::UDF(
datafusion_functions::datetime::date_part(),
),
args: vec![day_lit_expr, date_trunc_plus_interval_expr],
});
exprs.push(date_part_expr);
ScalarFunc::GeneratedExpr
}
_ => ScalarFunc::DataFusionBuiltin(
BuiltinScalarFunction::from_str(func.name).map_err(|_| {
UnsupportedExprSnafu {
_ => {
if let Ok(f) = BuiltinScalarFunction::from_str(func.name) {
ScalarFunc::DataFusionBuiltin(f)
} else if let Some(f) = datafusion_functions::math::functions()
.iter()
.find(|f| f.name() == func.name)
{
ScalarFunc::DataFusionUdf(f.clone())
} else {
return UnsupportedExprSnafu {
name: func.name.to_string(),
}
.build()
})?,
),
.fail();
}
}
};
for value in &self.ctx.field_columns {
@@ -1161,12 +1175,24 @@ impl PromPlanner {
ScalarFunc::DataFusionBuiltin(fun) => {
other_input_exprs.insert(field_column_pos, col_expr);
let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
fun,
func_def: ScalarFunctionDefinition::BuiltIn(fun),
args: other_input_exprs.clone().into(),
});
exprs.push(fn_expr);
let _ = other_input_exprs.remove(field_column_pos);
}
ScalarFunc::DataFusionUdf(f) => {
let args = itertools::chain!(
other_input_exprs.iter().take(field_column_pos).cloned(),
std::iter::once(col_expr),
other_input_exprs.iter().skip(field_column_pos).cloned()
)
.collect_vec();
exprs.push(DfExpr::ScalarFunction(ScalarFunction {
func_def: ScalarFunctionDefinition::UDF(f),
args,
}))
}
ScalarFunc::Udf(fun) => {
let ts_range_expr = DfExpr::Column(Column::from_name(
RangeManipulate::build_timestamp_range_name(
@@ -1175,8 +1201,8 @@ impl PromPlanner {
));
other_input_exprs.insert(field_column_pos, ts_range_expr);
other_input_exprs.insert(field_column_pos + 1, col_expr);
let fn_expr = DfExpr::ScalarUDF(ScalarUDF {
fun: Arc::new(fun),
let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
func_def: ScalarFunctionDefinition::UDF(Arc::new(fun)),
args: other_input_exprs.clone().into(),
});
exprs.push(fn_expr);
@@ -1193,8 +1219,8 @@ impl PromPlanner {
other_input_exprs.insert(field_column_pos + 1, col_expr);
other_input_exprs
.insert(field_column_pos + 2, self.create_time_index_column_expr()?);
let fn_expr = DfExpr::ScalarUDF(ScalarUDF {
fun: Arc::new(fun),
let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
func_def: ScalarFunctionDefinition::UDF(Arc::new(fun)),
args: other_input_exprs.clone().into(),
});
exprs.push(fn_expr);
@@ -1258,7 +1284,7 @@ impl PromPlanner {
exprs.push(expr);
}
utils::conjunction(exprs).context(ValueNotFoundSnafu {
conjunction(exprs).context(ValueNotFoundSnafu {
table: self.table_ref()?.to_quoted_string(),
})
}
@@ -1299,11 +1325,12 @@ impl PromPlanner {
.iter()
.map(|col| {
DfExpr::AggregateFunction(AggregateFunction {
fun: aggr.clone(),
func_def: AggregateFunctionDefinition::BuiltIn(aggr.clone()),
args: vec![DfExpr::Column(Column::from_name(col))],
distinct: false,
filter: None,
order_by: None,
null_treatment: None,
})
})
.collect();
@@ -1515,13 +1542,16 @@ impl PromPlanner {
token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
token::T_POW => Ok(Box::new(|lhs, rhs| {
Ok(DfExpr::ScalarFunction(ScalarFunction {
fun: BuiltinScalarFunction::Power,
func_def: ScalarFunctionDefinition::UDF(datafusion_functions::math::power()),
args: vec![lhs, rhs],
}))
})),
token::T_ATAN2 => Ok(Box::new(|lhs, rhs| {
Ok(DfExpr::ScalarFunction(ScalarFunction {
fun: BuiltinScalarFunction::Atan2,
// func_def: ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::Atan2),
func_def: datafusion_expr::ScalarFunctionDefinition::UDF(
datafusion_functions::math::atan2(),
),
args: vec![lhs, rhs],
}))
})),
@@ -1557,8 +1587,8 @@ impl PromPlanner {
&self,
left: LogicalPlan,
right: LogicalPlan,
left_table_ref: OwnedTableReference,
right_table_ref: OwnedTableReference,
left_table_ref: TableReference,
right_table_ref: TableReference,
) -> Result<LogicalPlan> {
let mut tag_columns = self
.ctx
@@ -1778,8 +1808,8 @@ impl PromPlanner {
.difference(&right_tag_cols_set)
.cloned()
.collect::<Vec<_>>();
let left_qualifier = left.schema().field(0).qualifier().cloned();
let right_qualifier = right.schema().field(0).qualifier().cloned();
let left_qualifier = left.schema().qualified_field(0).0.cloned();
let right_qualifier = right.schema().qualified_field(0).0.cloned();
let left_qualifier_string = left_qualifier
.as_ref()
.map(|l| l.to_string())
@@ -1953,7 +1983,7 @@ impl PromPlanner {
let field_columns_iter = result_field_columns
.into_iter()
.zip(self.ctx.field_columns.iter())
.map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, name.to_string()))));
.map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
// chain non-field columns (unchanged) and field columns (applied computation then alias)
let project_fields = non_field_columns_iter
@@ -2007,7 +2037,7 @@ impl PromPlanner {
})?,
);
let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
fun: BuiltinScalarFunction::DatePart,
func_def: ScalarFunctionDefinition::UDF(datafusion_functions::datetime::date_part()),
args: vec![lit_expr, input_expr],
});
Ok(fn_expr)
@@ -2023,6 +2053,8 @@ struct FunctionArgs {
#[derive(Debug, Clone)]
enum ScalarFunc {
DataFusionBuiltin(BuiltinScalarFunction),
/// The UDF that is defined by Datafusion itself.
DataFusionUdf(Arc<ScalarUdfDef>),
Udf(ScalarUdfDef),
// todo(ruihang): maybe merge with Udf later
/// UDF that require extra information like range length to be evaluated.
@@ -2436,7 +2468,7 @@ mod test {
#[tokio::test]
async fn aggregate_stdvar() {
do_aggregate_expr_plan("stdvar", "VARIANCE_POP").await;
do_aggregate_expr_plan("stdvar", "VAR_POP").await;
}
#[tokio::test]
src/query/src/datafusion.rs
+13 -12
View File
@@ -161,9 +161,9 @@ impl DatafusionQueryEngine {
}
#[tracing::instrument(skip_all)]
async fn delete<'a>(
async fn delete(
&self,
table_name: &ResolvedTableReference<'a>,
table_name: &ResolvedTableReference,
table: &TableRef,
column_vectors: HashMap<String, VectorRef>,
query_ctx: QueryContextRef,
@@ -205,9 +205,9 @@ impl DatafusionQueryEngine {
}
#[tracing::instrument(skip_all)]
async fn insert<'a>(
async fn insert(
&self,
table_name: &ResolvedTableReference<'a>,
table_name: &ResolvedTableReference,
column_vectors: HashMap<String, VectorRef>,
query_ctx: QueryContextRef,
) -> Result<Output> {
@@ -226,7 +226,7 @@ impl DatafusionQueryEngine {
.context(TableMutationSnafu)
}
async fn find_table(&self, table_name: &ResolvedTableReference<'_>) -> Result<TableRef> {
async fn find_table(&self, table_name: &ResolvedTableReference) -> Result<TableRef> {
let catalog_name = table_name.catalog.as_ref();
let schema_name = table_name.schema.as_ref();
let table_name = table_name.table.as_ref();
@@ -309,7 +309,9 @@ impl QueryEngine for DatafusionQueryEngine {
}
fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext {
QueryEngineContext::new(self.state.session_state(), query_ctx)
let mut state = self.state.session_state();
state.config_mut().set_extension(query_ctx.clone());
QueryEngineContext::new(state, query_ctx)
}
}
@@ -435,7 +437,7 @@ impl QueryExecutor for DatafusionQueryEngine {
let exec_timer = metrics::EXEC_PLAN_ELAPSED.start_timer();
let task_ctx = ctx.build_task_ctx();
match plan.output_partitioning().partition_count() {
match plan.properties().output_partitioning().partition_count() {
0 => Ok(Box::pin(EmptyRecordBatchStream::new(plan.schema()))),
1 => {
let stream = plan
@@ -453,7 +455,7 @@ impl QueryExecutor for DatafusionQueryEngine {
// merge into a single partition
let plan = CoalescePartitionsExec::new(df_plan.clone());
// CoalescePartitionsExec must produce a single partition
assert_eq!(1, plan.output_partitioning().partition_count());
assert_eq!(1, plan.properties().output_partitioning().partition_count());
let df_stream = plan
.execute(0, task_ctx)
.context(error::DatafusionSnafu)
@@ -475,7 +477,6 @@ impl QueryExecutor for DatafusionQueryEngine {
#[cfg(test)]
mod tests {
use std::borrow::Cow::Borrowed;
use std::sync::Arc;
use catalog::RegisterTableRequest;
@@ -576,9 +577,9 @@ mod tests {
.unwrap();
let table = engine
.find_table(&ResolvedTableReference {
catalog: Borrowed("greptime"),
schema: Borrowed("public"),
table: Borrowed("numbers"),
catalog: "greptime".into(),
schema: "public".into(),
table: "numbers".into(),
})
.await
.unwrap();
src/query/src/datafusion/planner.rs
+21 -6
View File
@@ -23,13 +23,13 @@ use common_query::logical_plan::create_aggregate_function;
use datafusion::catalog::TableReference;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::SessionState;
use datafusion::physical_plan::udaf::AggregateUDF;
use datafusion::physical_plan::udf::ScalarUDF;
use datafusion::sql::planner::ContextProvider;
use datafusion::variable::VarType;
use datafusion_common::config::ConfigOptions;
use datafusion_common::{DataFusionError, OwnedTableReference};
use datafusion_expr::{TableSource, WindowUDF};
use datafusion_physical_expr::var_provider::{is_system_variables, VarType};
use datafusion_common::DataFusionError;
use datafusion_expr::var_provider::is_system_variables;
use datafusion_expr::{AggregateUDF, TableSource, WindowUDF};
use datafusion_sql::parser::Statement as DfStatement;
use session::context::QueryContextRef;
use snafu::ResultExt;
@@ -79,7 +79,7 @@ impl DfContextProviderAdapter {
}
async fn resolve_tables(
table_names: Vec<OwnedTableReference>,
table_names: Vec<TableReference>,
table_provider: &mut DfTableSourceProvider,
) -> Result<HashMap<String, Arc<dyn TableSource>>> {
let mut tables = HashMap::with_capacity(table_names.len());
@@ -102,7 +102,7 @@ async fn resolve_tables(
}
impl ContextProvider for DfContextProviderAdapter {
fn get_table_provider(&self, name: TableReference) -> DfResult<Arc<dyn TableSource>> {
fn get_table_source(&self, name: TableReference) -> DfResult<Arc<dyn TableSource>> {
let table_ref = self.table_provider.resolve_table_ref(name)?;
self.tables
.get(&table_ref.to_string())
@@ -159,4 +159,19 @@ impl ContextProvider for DfContextProviderAdapter {
fn options(&self) -> &ConfigOptions {
self.session_state.config_options()
}
fn udfs_names(&self) -> Vec<String> {
// TODO(LFC): Impl it.
vec![]
}
fn udafs_names(&self) -> Vec<String> {
// TODO(LFC): Impl it.
vec![]
}
fn udwfs_names(&self) -> Vec<String> {
// TODO(LFC): Impl it.
vec![]
}
}
src/query/src/dist_plan/analyzer.rs
+27 -23
View File
@@ -17,7 +17,7 @@ use std::sync::Arc;
use datafusion::datasource::DefaultTableSource;
use datafusion::error::Result as DfResult;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{RewriteRecursion, Transformed, TreeNode, TreeNodeRewriter};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
use datafusion_expr::expr::{Exists, InSubquery};
use datafusion_expr::{col, Expr, LogicalPlan, LogicalPlanBuilder, Subquery};
use datafusion_optimizer::analyzer::AnalyzerRule;
@@ -47,12 +47,12 @@ impl AnalyzerRule for DistPlannerAnalyzer {
// preprocess the input plan
let optimizer_context = OptimizerContext::new();
let plan = SimplifyExpressions::new()
.try_optimize(&plan, &optimizer_context)?
.unwrap_or(plan);
.rewrite(plan, &optimizer_context)?
.data;
let plan = plan.transform(&Self::inspect_plan_with_subquery)?;
let mut rewriter = PlanRewriter::default();
let result = plan.rewrite(&mut rewriter)?;
let result = plan.data.rewrite(&mut rewriter)?.data;
Ok(result)
}
@@ -63,35 +63,40 @@ impl DistPlannerAnalyzer {
let exprs = plan
.expressions()
.into_iter()
.map(|e| e.transform(&Self::transform_subquery))
.map(|e| e.transform(&Self::transform_subquery).map(|x| x.data))
.collect::<DfResult<Vec<_>>>()?;
let inputs = plan.inputs().into_iter().cloned().collect::<Vec<_>>();
Ok(Transformed::Yes(plan.with_new_exprs(exprs, &inputs)?))
Ok(Transformed::yes(plan.with_new_exprs(exprs, inputs)?))
}
fn transform_subquery(expr: Expr) -> DfResult<Transformed<Expr>> {
match expr {
Expr::Exists(exists) => Ok(Transformed::Yes(Expr::Exists(Exists {
Expr::Exists(exists) => Ok(Transformed::yes(Expr::Exists(Exists {
subquery: Self::handle_subquery(exists.subquery)?,
negated: exists.negated,
}))),
Expr::InSubquery(in_subquery) => Ok(Transformed::Yes(Expr::InSubquery(InSubquery {
Expr::InSubquery(in_subquery) => Ok(Transformed::yes(Expr::InSubquery(InSubquery {
expr: in_subquery.expr,
subquery: Self::handle_subquery(in_subquery.subquery)?,
negated: in_subquery.negated,
}))),
Expr::ScalarSubquery(scalar_subquery) => Ok(Transformed::Yes(Expr::ScalarSubquery(
Expr::ScalarSubquery(scalar_subquery) => Ok(Transformed::yes(Expr::ScalarSubquery(
Self::handle_subquery(scalar_subquery)?,
))),
_ => Ok(Transformed::No(expr)),
_ => Ok(Transformed::no(expr)),
}
}
fn handle_subquery(subquery: Subquery) -> DfResult<Subquery> {
let mut rewriter = PlanRewriter::default();
let mut rewrote_subquery = subquery.subquery.as_ref().clone().rewrite(&mut rewriter)?;
let mut rewrote_subquery = subquery
.subquery
.as_ref()
.clone()
.rewrite(&mut rewriter)?
.data;
// Workaround. DF doesn't support the first plan in subquery to be an Extension
if matches!(rewrote_subquery, LogicalPlan::Extension(_)) {
let output_schema = rewrote_subquery.schema().clone();
@@ -232,35 +237,34 @@ impl PlanRewriter {
}
impl TreeNodeRewriter for PlanRewriter {
type N = LogicalPlan;
type Node = LogicalPlan;
/// descend
fn pre_visit<'a>(&'a mut self, node: &'a Self::N) -> DfResult<RewriteRecursion> {
fn f_down<'a>(&mut self, node: Self::Node) -> DfResult<Transformed<Self::Node>> {
self.level += 1;
self.stack.push((node.clone(), self.level));
// decendening will clear the stage
self.stage.clear();
self.set_unexpanded();
self.partition_cols = None;
Ok(RewriteRecursion::Continue)
Ok(Transformed::no(node))
}
/// ascend
///
/// Besure to call `pop_stack` before returning
fn mutate(&mut self, node: Self::N) -> DfResult<Self::N> {
fn f_up(&mut self, node: Self::Node) -> DfResult<Transformed<Self::Node>> {
// only expand once on each ascending
if self.is_expanded() {
self.pop_stack();
return Ok(node);
return Ok(Transformed::no(node));
}
// only expand when the leaf is table scan
if node.inputs().is_empty() && !matches!(node, LogicalPlan::TableScan(_)) {
self.set_expanded();
self.pop_stack();
return Ok(node);
return Ok(Transformed::no(node));
}
self.maybe_set_partitions(&node);
@@ -270,12 +274,12 @@ impl TreeNodeRewriter for PlanRewriter {
let mut node = MergeScanLogicalPlan::new(node, false).into_logical_plan();
// expand stages
for new_stage in self.stage.drain(..) {
node = new_stage.with_new_inputs(&[node])?
node = new_stage.with_new_exprs(node.expressions(), vec![node.clone()])?
}
self.set_expanded();
self.pop_stack();
return Ok(node);
return Ok(Transformed::yes(node));
};
// TODO(ruihang): avoid this clone
@@ -285,16 +289,16 @@ impl TreeNodeRewriter for PlanRewriter {
let mut node = MergeScanLogicalPlan::new(node, false).into_logical_plan();
// expand stages
for new_stage in self.stage.drain(..) {
node = new_stage.with_new_inputs(&[node])?
node = new_stage.with_new_exprs(node.expressions(), vec![node.clone()])?
}
self.set_expanded();
self.pop_stack();
return Ok(node);
return Ok(Transformed::yes(node));
}
self.pop_stack();
Ok(node)
Ok(Transformed::no(node))
}
}
src/query/src/dist_plan/commutativity.rs
+4 -5
View File
@@ -107,6 +107,7 @@ impl Categorizer {
LogicalPlan::Dml(_) => Commutativity::Unsupported,
LogicalPlan::Ddl(_) => Commutativity::Unsupported,
LogicalPlan::Copy(_) => Commutativity::Unsupported,
LogicalPlan::RecursiveQuery(_) => Commutativity::Unsupported,
}
}
@@ -142,8 +143,7 @@ impl Categorizer {
| Expr::Between(_)
| Expr::Sort(_)
| Expr::Exists(_)
| Expr::ScalarFunction(_)
| Expr::ScalarUDF(_) => Commutativity::Commutative,
| Expr::ScalarFunction(_) => Commutativity::Commutative,
Expr::Like(_)
| Expr::SimilarTo(_)
@@ -155,14 +155,13 @@ impl Categorizer {
| Expr::TryCast(_)
| Expr::AggregateFunction(_)
| Expr::WindowFunction(_)
| Expr::AggregateUDF(_)
| Expr::InList(_)
| Expr::InSubquery(_)
| Expr::ScalarSubquery(_)
| Expr::Wildcard => Commutativity::Unimplemented,
| Expr::Wildcard { .. } => Commutativity::Unimplemented,
Expr::Alias(_)
| Expr::QualifiedWildcard { .. }
| Expr::Unnest(_)
| Expr::GroupingSet(_)
| Expr::Placeholder(_)
| Expr::OuterReferenceColumn(_, _) => Commutativity::Unimplemented,
src/query/src/dist_plan/merge_scan.rs
+20 -13
View File
@@ -33,15 +33,18 @@ use common_telemetry::tracing_context::TracingContext;
use datafusion::physical_plan::metrics::{
Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, Time,
};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning};
use datafusion_common::{Result, Statistics};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, Partitioning, PlanProperties,
};
use datafusion_common::Result;
use datafusion_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion_physical_expr::PhysicalSortExpr;
use datafusion_physical_expr::EquivalenceProperties;
use datatypes::schema::{Schema, SchemaRef};
use futures_util::StreamExt;
use greptime_proto::v1::region::{QueryRequest, RegionRequestHeader};
use meter_core::data::ReadItem;
use meter_macros::read_meter;
use session::context::QueryContextRef;
use snafu::ResultExt;
use store_api::storage::RegionId;
use tokio::time::Instant;
@@ -123,6 +126,8 @@ pub struct MergeScanExec {
arrow_schema: ArrowSchemaRef,
region_query_handler: RegionQueryHandlerRef,
metric: ExecutionPlanMetricsSet,
properties: PlanProperties,
query_ctx: QueryContextRef,
}
impl std::fmt::Debug for MergeScanExec {
@@ -142,8 +147,14 @@ impl MergeScanExec {
substrait_plan: Bytes,
arrow_schema: &ArrowSchema,
region_query_handler: RegionQueryHandlerRef,
query_ctx: QueryContextRef,
) -> Result<Self> {
let arrow_schema_without_metadata = Self::arrow_schema_without_metadata(arrow_schema);
let properties = PlanProperties::new(
EquivalenceProperties::new(arrow_schema_without_metadata.clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
);
let schema_without_metadata =
Self::arrow_schema_to_schema(arrow_schema_without_metadata.clone())?;
Ok(Self {
@@ -154,6 +165,8 @@ impl MergeScanExec {
arrow_schema: arrow_schema_without_metadata,
region_query_handler,
metric: ExecutionPlanMetricsSet::new(),
properties,
query_ctx,
})
}
@@ -166,6 +179,7 @@ impl MergeScanExec {
let dbname = context.task_id().unwrap_or_default();
let tracing_context = TracingContext::from_json(context.session_id().as_str());
let tz = self.query_ctx.timezone().to_string();
let stream = Box::pin(stream!({
MERGE_SCAN_REGIONS.observe(regions.len() as f64);
@@ -178,6 +192,7 @@ impl MergeScanExec {
header: Some(RegionRequestHeader {
tracing_context: tracing_context.to_w3c(),
dbname: dbname.clone(),
timezone: tz.clone(),
}),
region_id: region_id.into(),
plan: substrait_plan.clone(),
@@ -266,12 +281,8 @@ impl ExecutionPlan for MergeScanExec {
self.arrow_schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
Partitioning::UnknownPartitioning(1)
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
None
fn properties(&self) -> &PlanProperties {
&self.properties
}
fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
@@ -297,10 +308,6 @@ impl ExecutionPlan for MergeScanExec {
)))
}
fn statistics(&self) -> Statistics {
Statistics::default()
}
fn metrics(&self) -> Option<MetricsSet> {
Some(self.metric.clone_inner())
}
src/query/src/dist_plan/planner.rs
+24 -16
View File
@@ -25,10 +25,11 @@ use datafusion::datasource::DefaultTableSource;
use datafusion::execution::context::SessionState;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeVisitor, VisitRecursion};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor};
use datafusion_common::TableReference;
use datafusion_expr::{LogicalPlan, UserDefinedLogicalNode};
use datafusion_optimizer::analyzer::Analyzer;
use session::context::QueryContext;
use snafu::{OptionExt, ResultExt};
use store_api::storage::RegionId;
use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
@@ -103,12 +104,18 @@ impl ExtensionPlanner for DistExtensionPlanner {
.encode(&amended_plan)
.context(error::EncodeSubstraitLogicalPlanSnafu)?
.into();
let query_ctx = session_state
.config()
.get_extension()
.unwrap_or_else(QueryContext::arc);
let merge_scan_plan = MergeScanExec::new(
table_name,
regions,
substrait_plan,
&schema,
self.region_query_handler.clone(),
query_ctx,
)?;
Ok(Some(Arc::new(merge_scan_plan) as _))
}
@@ -125,6 +132,7 @@ impl DistExtensionPlanner {
/// Apply the fully resolved table name to the TableScan plan
fn plan_with_full_table_name(plan: LogicalPlan, name: &TableName) -> Result<LogicalPlan> {
plan.transform(&|plan| TableNameRewriter::rewrite_table_name(plan, name))
.map(|x| x.data)
}
async fn get_regions(&self, table_name: &TableName) -> Result<Vec<RegionId>> {
@@ -163,9 +171,9 @@ struct TableNameExtractor {
}
impl TreeNodeVisitor for TableNameExtractor {
type N = LogicalPlan;
type Node = LogicalPlan;
fn pre_visit(&mut self, node: &Self::N) -> Result<VisitRecursion> {
fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
match node {
LogicalPlan::TableScan(scan) => {
if let Some(source) = scan.source.as_any().downcast_ref::<DefaultTableSource>() {
@@ -182,7 +190,7 @@ impl TreeNodeVisitor for TableNameExtractor {
info.name.clone(),
));
}
return Ok(VisitRecursion::Stop);
return Ok(TreeNodeRecursion::Stop);
}
}
match &scan.table_name {
@@ -192,32 +200,32 @@ impl TreeNodeVisitor for TableNameExtractor {
table,
} => {
self.table_name = Some(TableName::new(
catalog.clone(),
schema.clone(),
table.clone(),
catalog.to_string(),
schema.to_string(),
table.to_string(),
));
Ok(VisitRecursion::Stop)
Ok(TreeNodeRecursion::Stop)
}
// TODO(ruihang): Maybe the following two cases should not be valid
TableReference::Partial { schema, table } => {
self.table_name = Some(TableName::new(
DEFAULT_CATALOG_NAME.to_string(),
schema.clone(),
table.clone(),
schema.to_string(),
table.to_string(),
));
Ok(VisitRecursion::Stop)
Ok(TreeNodeRecursion::Stop)
}
TableReference::Bare { table } => {
self.table_name = Some(TableName::new(
DEFAULT_CATALOG_NAME.to_string(),
DEFAULT_SCHEMA_NAME.to_string(),
table.clone(),
table.to_string(),
));
Ok(VisitRecursion::Stop)
Ok(TreeNodeRecursion::Stop)
}
}
}
_ => Ok(VisitRecursion::Continue),
_ => Ok(TreeNodeRecursion::Continue),
}
}
}
@@ -236,9 +244,9 @@ impl TableNameRewriter {
name.schema_name.clone(),
name.table_name.clone(),
);
Transformed::Yes(LogicalPlan::TableScan(table_scan))
Transformed::yes(LogicalPlan::TableScan(table_scan))
}
_ => Transformed::No(plan),
_ => Transformed::no(plan),
})
}
}
src/query/src/metrics.rs
+4
View File
@@ -76,6 +76,10 @@ impl<F> OnDone<F> {
}
impl<F: FnOnce() + Unpin> RecordBatchStream for OnDone<F> {
fn name(&self) -> &str {
self.stream.name()
}
fn schema(&self) -> SchemaRef {
self.stream.schema()
}
src/query/src/optimizer/order_hint.rs
+9 -8
View File
@@ -15,7 +15,7 @@
use arrow_schema::SortOptions;
use common_recordbatch::OrderOption;
use datafusion::datasource::DefaultTableSource;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeVisitor, VisitRecursion};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor};
use datafusion_common::Result as DataFusionResult;
use datafusion_expr::expr::Sort;
use datafusion_expr::{Expr, LogicalPlan};
@@ -48,6 +48,7 @@ impl OrderHintRule {
if let Some(order_expr) = visitor.order_expr.take() {
plan.clone()
.transform_down(&|plan| Self::set_ordering_hint(plan, &order_expr))
.map(|x| x.data)
} else {
Ok(plan.clone())
}
@@ -74,7 +75,7 @@ impl OrderHintRule {
for sort in order_expr {
let name = match sort.expr.try_into_col() {
Ok(col) => col.name,
Err(_) => return Ok(Transformed::No(plan)),
Err(_) => return Ok(Transformed::no(plan)),
};
opts.push(OrderOption {
name,
@@ -89,12 +90,12 @@ impl OrderHintRule {
}
}
if transformed {
Ok(Transformed::Yes(plan))
Ok(Transformed::yes(plan))
} else {
Ok(Transformed::No(plan))
Ok(Transformed::no(plan))
}
}
_ => Ok(Transformed::No(plan)),
_ => Ok(Transformed::no(plan)),
}
}
}
@@ -106,9 +107,9 @@ struct OrderHintVisitor {
}
impl TreeNodeVisitor for OrderHintVisitor {
type N = LogicalPlan;
type Node = LogicalPlan;
fn pre_visit(&mut self, node: &Self::N) -> DataFusionResult<VisitRecursion> {
fn f_down(&mut self, node: &Self::Node) -> DataFusionResult<TreeNodeRecursion> {
if let LogicalPlan::Sort(sort) = node {
let mut exprs = vec![];
for expr in &sort.expr {
@@ -118,7 +119,7 @@ impl TreeNodeVisitor for OrderHintVisitor {
}
self.order_expr = Some(exprs);
}
Ok(VisitRecursion::Continue)
Ok(TreeNodeRecursion::Continue)
}
}
src/query/src/optimizer/string_normalization.rs
+6 -5
View File
@@ -32,10 +32,11 @@ impl AnalyzerRule for StringNormalizationRule {
let expr = plan
.expressions()
.into_iter()
.map(|e| e.rewrite(&mut converter))
.map(|e| e.rewrite(&mut converter).map(|x| x.data))
.collect::<Result<Vec<_>>>()?;
plan.with_new_exprs(expr, &inputs).map(Transformed::Yes)
plan.with_new_exprs(expr, inputs).map(Transformed::yes)
})
.map(|x| x.data)
}
fn name(&self) -> &str {
@@ -46,12 +47,12 @@ impl AnalyzerRule for StringNormalizationRule {
struct StringNormalizationConverter;
impl TreeNodeRewriter for StringNormalizationConverter {
type N = Expr;
type Node = Expr;
/// remove extra whitespaces from the String value when
/// there is a CAST from a String to Timestamp.
/// Otherwise - no modifications applied
fn mutate(&mut self, expr: Expr) -> Result<Expr> {
fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
let new_expr = match expr {
Expr::Cast(Cast { expr, data_type }) => {
let expr = match data_type {
@@ -71,7 +72,7 @@ impl TreeNodeRewriter for StringNormalizationConverter {
}
expr => expr,
};
Ok(new_expr)
Ok(Transformed::yes(new_expr))
}
}
src/query/src/optimizer/type_conversion.rs
+32 -25
View File
@@ -48,8 +48,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
schema: filter.input.schema().clone(),
query_ctx: ctx.query_ctx(),
};
let rewritten = filter.predicate.clone().rewrite(&mut converter)?;
Ok(Transformed::Yes(LogicalPlan::Filter(Filter::try_new(
let rewritten = filter.predicate.clone().rewrite(&mut converter)?.data;
Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new(
rewritten,
filter.input,
)?)))
@@ -68,9 +68,9 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
};
let rewrite_filters = filters
.into_iter()
.map(|e| e.rewrite(&mut converter))
.map(|e| e.rewrite(&mut converter).map(|x| x.data))
.collect::<Result<Vec<_>>>()?;
Ok(Transformed::Yes(LogicalPlan::TableScan(TableScan {
Ok(Transformed::yes(LogicalPlan::TableScan(TableScan {
table_name: table_name.clone(),
source: source.clone(),
projection,
@@ -100,10 +100,10 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
let expr = plan
.expressions()
.into_iter()
.map(|e| e.rewrite(&mut converter))
.map(|e| e.rewrite(&mut converter).map(|x| x.data))
.collect::<Result<Vec<_>>>()?;
plan.with_new_exprs(expr, &inputs).map(Transformed::Yes)
plan.with_new_exprs(expr, inputs).map(Transformed::yes)
}
LogicalPlan::Subquery { .. }
@@ -116,8 +116,10 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
| LogicalPlan::Unnest(_)
| LogicalPlan::Statement(_)
| LogicalPlan::Ddl(_)
| LogicalPlan::Copy(_) => Ok(Transformed::No(plan)),
| LogicalPlan::Copy(_)
| LogicalPlan::RecursiveQuery(_) => Ok(Transformed::no(plan)),
})
.map(|x| x.data)
}
fn name(&self) -> &str {
@@ -155,9 +157,9 @@ impl TypeConverter {
_ => Ok(ScalarValue::Boolean(None)),
},
(target_type, value) => {
let value_arr = value.to_array();
let arr =
compute::cast(&value_arr, target_type).map_err(DataFusionError::ArrowError)?;
let value_arr = value.to_array()?;
let arr = compute::cast(&value_arr, target_type)
.map_err(|e| DataFusionError::ArrowError(e, None))?;
ScalarValue::try_from_array(
&arr,
@@ -207,9 +209,9 @@ impl TypeConverter {
}
impl TreeNodeRewriter for TypeConverter {
type N = Expr;
type Node = Expr;
fn mutate(&mut self, expr: Expr) -> Result<Expr> {
fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
let new_expr = match expr {
Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
Operator::Eq
@@ -275,7 +277,7 @@ impl TreeNodeRewriter for TypeConverter {
},
expr => expr,
};
Ok(new_expr)
Ok(Transformed::yes(new_expr))
}
}
@@ -310,7 +312,9 @@ mod tests {
use std::sync::Arc;
use datafusion::logical_expr::expr::AggregateFunction as AggrExpr;
use datafusion_common::{Column, DFField, DFSchema};
use datafusion_common::arrow::datatypes::Field;
use datafusion_common::{Column, DFSchema};
use datafusion_expr::expr::AggregateFunctionDefinition;
use datafusion_expr::{AggregateFunction, LogicalPlanBuilder};
use datafusion_sql::TableReference;
use session::context::QueryContext;
@@ -390,11 +394,13 @@ mod tests {
let schema = Arc::new(
DFSchema::new_with_metadata(
vec![DFField::new(
vec![(
None::<TableReference>,
"ts",
DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
true,
Arc::new(Field::new(
"ts",
DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
true,
)),
)],
HashMap::new(),
)
@@ -411,12 +417,13 @@ mod tests {
None
))),
converter
.mutate(
.f_up(
Expr::Column(Column::from_name("ts")).gt(Expr::Literal(ScalarValue::Utf8(
Some("2020-09-08T05:42:29+08:00".to_string()),
)))
)
.unwrap()
.data
);
}
@@ -425,11 +432,9 @@ mod tests {
let col_name = "is_valid";
let schema = Arc::new(
DFSchema::new_with_metadata(
vec![DFField::new(
vec![(
None::<TableReference>,
col_name,
DataType::Boolean,
false,
Arc::new(Field::new(col_name, DataType::Boolean, false)),
)],
HashMap::new(),
)
@@ -444,11 +449,12 @@ mod tests {
Expr::Column(Column::from_name(col_name))
.eq(Expr::Literal(ScalarValue::Boolean(Some(true)))),
converter
.mutate(
.f_up(
Expr::Column(Column::from_name(col_name))
.eq(Expr::Literal(ScalarValue::Utf8(Some("true".to_string()))))
)
.unwrap()
.data
);
}
@@ -475,11 +481,12 @@ mod tests {
.aggregate(
Vec::<Expr>::new(),
vec![Expr::AggregateFunction(AggrExpr {
fun: AggregateFunction::Count,
func_def: AggregateFunctionDefinition::BuiltIn(AggregateFunction::Count),
args: vec![Expr::Column(Column::from_name("column1"))],
distinct: false,
filter: None,
order_by: None,
null_treatment: None,
})],
)
.unwrap()
src/query/src/parser.rs
+3 -3
View File
@@ -305,9 +305,9 @@ mod test {
sort_by: [], \
having: None, \
named_window: [], \
qualify: None \
}), order_by: [], limit: None, offset: None, fetch: None, locks: [] } }))");
qualify: None, \
value_table_mode: None \
}), order_by: [], limit: None, limit_by: [], offset: None, fetch: None, locks: [], for_clause: None } }))");
assert_eq!(format!("{stmt:?}"), expected);
}
src/query/src/plan.rs
+2 -1
View File
@@ -16,6 +16,7 @@ use std::collections::HashMap;
use std::fmt::{Debug, Display};
use common_query::prelude::ScalarValue;
use datafusion_common::ParamValues;
use datafusion_expr::LogicalPlan as DfLogicalPlan;
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::Schema;
@@ -82,7 +83,7 @@ impl LogicalPlan {
let LogicalPlan::DfPlan(plan) = self;
plan.clone()
.replace_params_with_values(values)
.replace_params_with_values(&ParamValues::List(values.to_vec()))
.context(DataFusionSnafu)
.map(LogicalPlan::DfPlan)
}
src/query/src/query_engine/state.rs
+8 -13
View File
@@ -26,7 +26,6 @@ use common_function::state::FunctionState;
use common_query::physical_plan::SessionContext;
use common_query::prelude::ScalarUdf;
use common_telemetry::warn;
use datafusion::catalog::MemoryCatalogList;
use datafusion::dataframe::DataFrame;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::{QueryPlanner, SessionConfig, SessionState};
@@ -101,18 +100,14 @@ impl QueryEngineState {
let mut optimizer = Optimizer::new();
optimizer.rules.push(Arc::new(OrderHintRule));
let session_state = SessionState::new_with_config_rt_and_catalog_list(
session_config,
runtime_env,
Arc::new(MemoryCatalogList::default()), // pass a dummy catalog list
)
.with_serializer_registry(Arc::new(ExtensionSerializer))
.with_analyzer_rules(analyzer.rules)
.with_query_planner(Arc::new(DfQueryPlanner::new(
catalog_list.clone(),
region_query_handler,
)))
.with_optimizer_rules(optimizer.rules);
let session_state = SessionState::new_with_config_rt(session_config, runtime_env)
.with_serializer_registry(Arc::new(ExtensionSerializer))
.with_analyzer_rules(analyzer.rules)
.with_query_planner(Arc::new(DfQueryPlanner::new(
catalog_list.clone(),
region_query_handler,
)))
.with_optimizer_rules(optimizer.rules);
let df_context = SessionContext::new_with_state(session_state);
src/query/src/range_select/plan.rs
+119 -82
View File
@@ -33,12 +33,14 @@ use datafusion::execution::context::SessionState;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::udaf::create_aggregate_expr as create_aggr_udf_expr;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
};
use datafusion::physical_planner::create_physical_sort_expr;
use datafusion_common::hash_utils::create_hashes;
use datafusion_common::utils::{get_arrayref_at_indices, get_row_at_idx};
use datafusion_common::{DFField, DFSchema, DFSchemaRef, DataFusionError, ScalarValue};
use datafusion_common::{DFSchema, DFSchemaRef, DataFusionError, ScalarValue};
use datafusion_expr::expr::AggregateFunctionDefinition;
use datafusion_expr::utils::{exprlist_to_fields, COUNT_STAR_EXPANSION};
use datafusion_expr::{
lit, Accumulator, AggregateFunction, Expr, ExprSchemable, LogicalPlan,
@@ -46,9 +48,9 @@ use datafusion_expr::{
};
use datafusion_physical_expr::aggregate::utils::down_cast_any_ref;
use datafusion_physical_expr::expressions::create_aggregate_expr as create_aggr_expr;
use datafusion_physical_expr::hash_utils::create_hashes;
use datafusion_physical_expr::{
create_physical_expr, AggregateExpr, Distribution, PhysicalExpr, PhysicalSortExpr,
create_physical_expr, AggregateExpr, Distribution, EquivalenceProperties, Partitioning,
PhysicalExpr, PhysicalSortExpr,
};
use datatypes::arrow::array::{
Array, ArrayRef, TimestampMillisecondArray, TimestampMillisecondBuilder, UInt32Builder,
@@ -181,7 +183,7 @@ impl Accumulator for RangeFirstListValueAcc {
Ok(())
}
fn evaluate(&self) -> DataFusionResult<ScalarValue> {
fn evaluate(&mut self) -> DataFusionResult<ScalarValue> {
Ok(self.data.clone().unwrap_or(ScalarValue::Null))
}
@@ -189,7 +191,7 @@ impl Accumulator for RangeFirstListValueAcc {
std::mem::size_of_val(self)
}
fn state(&self) -> DataFusionResult<Vec<ScalarValue>> {
fn state(&mut self) -> DataFusionResult<Vec<ScalarValue>> {
unreachable!("Accumulator::state will not be used in range query")
}
@@ -466,12 +468,13 @@ impl RangeSelect {
fill,
..
}| {
Ok(DFField::new_unqualified(
let field = Field::new(
name,
data_type.clone(),
// Only when data fill with Const option, the data can't be null
!matches!(fill, Some(Fill::Const(..))),
))
);
Ok((None, Arc::new(field)))
},
)
.collect::<DfResult<Vec<_>>>()
@@ -480,11 +483,10 @@ impl RangeSelect {
let ts_field = time_index
.to_field(input.schema().as_ref())
.context(DataFusionSnafu)?;
let time_index_name = ts_field.name().clone();
let time_index_name = ts_field.1.name().clone();
fields.push(ts_field);
// add by
let by_fields =
exprlist_to_fields(by.iter().collect::<Vec<_>>(), &input).context(DataFusionSnafu)?;
let by_fields = exprlist_to_fields(&by, &input).context(DataFusionSnafu)?;
fields.extend(by_fields.clone());
let schema_before_project = Arc::new(
DFSchema::new_with_metadata(fields, input.schema().metadata().clone())
@@ -502,7 +504,6 @@ impl RangeSelect {
if let Expr::Column(column) = project_expr {
schema_before_project
.index_of_column_by_name(column.relation.as_ref(), &column.name)
.unwrap_or(None)
.ok_or(())
} else {
Err(())
@@ -513,7 +514,10 @@ impl RangeSelect {
let schema = if let Some(project) = &schema_project {
let project_field = project
.iter()
.map(|i| schema_before_project.fields()[*i].clone())
.map(|i| {
let f = schema_before_project.qualified_field(*i);
(f.0.cloned(), Arc::new(f.1.clone()))
})
.collect();
Arc::new(
DFSchema::new_with_metadata(project_field, input.schema().metadata().clone())
@@ -555,6 +559,8 @@ impl UserDefinedLogicalNodeCore for RangeSelect {
self.range_expr
.iter()
.map(|expr| expr.expr.clone())
.chain([self.time_expr.clone()])
.chain(self.by.clone())
.collect()
}
@@ -578,18 +584,32 @@ impl UserDefinedLogicalNodeCore for RangeSelect {
)
}
fn from_template(&self, _exprs: &[Expr], inputs: &[LogicalPlan]) -> Self {
fn from_template(&self, exprs: &[Expr], inputs: &[LogicalPlan]) -> Self {
assert!(!inputs.is_empty());
assert!(exprs.len() == self.range_expr.len() + self.by.len() + 1);
let range_expr = exprs
.iter()
.zip(self.range_expr.iter())
.map(|(e, range)| RangeFn {
name: range.name.clone(),
data_type: range.data_type.clone(),
expr: e.clone(),
range: range.range,
fill: range.fill.clone(),
need_cast: range.need_cast,
})
.collect();
let time_expr = exprs[self.range_expr.len()].clone();
let by = exprs[self.range_expr.len() + 1..].to_vec();
Self {
align: self.align,
align_to: self.align_to,
range_expr: self.range_expr.clone(),
range_expr,
input: Arc::new(inputs[0].clone()),
time_index: self.time_index.clone(),
time_expr: self.time_expr.clone(),
time_expr,
schema: self.schema.clone(),
by: self.by.clone(),
by,
by_schema: self.by_schema.clone(),
schema_project: self.schema_project.clone(),
schema_before_project: self.schema_before_project.clone(),
@@ -603,7 +623,6 @@ impl RangeSelect {
is_count_aggr: bool,
exprs: &[Expr],
df_schema: &Arc<DFSchema>,
schema: &Schema,
session_state: &SessionState,
) -> DfResult<Vec<Arc<dyn PhysicalExpr>>> {
exprs
@@ -614,13 +633,12 @@ impl RangeSelect {
// At this time, aggregate plan has been replaced by a custom range plan,
// so `CountWildcardRule` has not been applied.
// We manually modify it when creating the physical plan.
Expr::Wildcard if is_count_aggr => create_physical_expr(
Expr::Wildcard { .. } if is_count_aggr => create_physical_expr(
&lit(COUNT_STAR_EXPANSION),
df_schema,
schema,
df_schema.as_ref(),
session_state.execution_props(),
),
_ => create_physical_expr(e, df_schema, schema, session_state.execution_props()),
_ => create_physical_expr(e, df_schema.as_ref(), session_state.execution_props()),
})
.collect::<DfResult<Vec<_>>>()
}
@@ -650,10 +668,25 @@ impl RangeSelect {
.iter()
.map(|range_fn| {
let expr = match &range_fn.expr {
Expr::AggregateFunction(aggr)
if aggr.fun == AggregateFunction::FirstValue
|| aggr.fun == AggregateFunction::LastValue =>
{
Expr::AggregateFunction(
aggr @ datafusion_expr::expr::AggregateFunction {
func_def:
AggregateFunctionDefinition::BuiltIn(AggregateFunction::FirstValue),
..
},
)
| Expr::AggregateFunction(
aggr @ datafusion_expr::expr::AggregateFunction {
func_def:
AggregateFunctionDefinition::BuiltIn(AggregateFunction::LastValue),
..
},
) => {
let is_last_value_func = matches!(
aggr.func_def,
AggregateFunctionDefinition::BuiltIn(AggregateFunction::LastValue)
);
// Because we only need to find the first_value/last_value,
// the complexity of sorting the entire batch is O(nlogn).
// We can sort the batch with limit 1.
@@ -665,13 +698,12 @@ impl RangeSelect {
.map(|x| {
create_physical_sort_expr(
x,
input_dfschema,
&input_schema,
input_dfschema.as_ref(),
session_state.execution_props(),
)
.map(|expr| {
// reverse the last_value sort
if aggr.fun == AggregateFunction::LastValue {
if is_last_value_func {
PhysicalSortExpr {
expr: expr.expr,
options: SortOptions {
@@ -689,14 +721,13 @@ impl RangeSelect {
// if user not assign order by, time index is needed as default ordering
let time_index = create_physical_expr(
&self.time_expr,
input_dfschema,
&input_schema,
input_dfschema.as_ref(),
session_state.execution_props(),
)?;
vec![PhysicalSortExpr {
expr: time_index,
options: SortOptions {
descending: aggr.fun == AggregateFunction::LastValue,
descending: is_last_value_func,
nulls_first: false,
},
}]
@@ -705,7 +736,6 @@ impl RangeSelect {
false,
&aggr.args,
input_dfschema,
&input_schema,
session_state,
)?;
// first_value/last_value has only one param.
@@ -723,8 +753,7 @@ impl RangeSelect {
.map(|x| {
create_physical_sort_expr(
x,
input_dfschema,
&input_schema,
input_dfschema.as_ref(),
session_state.execution_props(),
)
})
@@ -732,36 +761,39 @@ impl RangeSelect {
} else {
vec![]
};
let expr = create_aggr_expr(
&aggr.fun,
false,
&self.create_physical_expr_list(
aggr.fun == AggregateFunction::Count,
&aggr.args,
input_dfschema,
&input_schema,
session_state,
)?,
&order_by,
&input_schema,
range_fn.expr.display_name()?,
let input_phy_exprs = self.create_physical_expr_list(
matches!(
aggr.func_def,
AggregateFunctionDefinition::BuiltIn(AggregateFunction::Count,)
),
&aggr.args,
input_dfschema,
session_state,
)?;
Ok(expr)
}
Expr::AggregateUDF(aggr_udf) => {
let expr = create_aggr_udf_expr(
&aggr_udf.fun,
&self.create_physical_expr_list(
match &aggr.func_def {
AggregateFunctionDefinition::BuiltIn(fun) => create_aggr_expr(
fun,
false,
&aggr_udf.args,
input_dfschema,
&input_phy_exprs,
&order_by,
&input_schema,
session_state,
)?,
&input_schema,
range_fn.expr.display_name()?,
)?;
Ok(expr)
range_fn.expr.display_name()?,
false,
),
AggregateFunctionDefinition::UDF(fun) => create_aggr_udf_expr(
fun,
&input_phy_exprs,
&[],
&[],
&input_schema,
range_fn.expr.display_name()?,
false,
),
f => Err(DataFusionError::NotImplemented(format!(
"Range function from {f:?}"
))),
}
}
_ => Err(DataFusionError::Plan(format!(
"Unexpected Expr:{} in RangeSelect",
@@ -788,24 +820,25 @@ impl RangeSelect {
} else {
schema_before_project.clone()
};
let by = self.create_physical_expr_list(false, &self.by, input_dfschema, session_state)?;
let cache = PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
);
Ok(Arc::new(RangeSelectExec {
input: exec_input,
range_exec,
align: self.align.as_millis() as Millisecond,
align_to: self.align_to,
by: self.create_physical_expr_list(
false,
&self.by,
input_dfschema,
&input_schema,
session_state,
)?,
by,
time_index: self.time_index.clone(),
schema,
by_schema: Arc::new(Schema::new(by_fields)),
metric: ExecutionPlanMetricsSet::new(),
schema_before_project,
schema_project: self.schema_project.clone(),
cache,
}))
}
}
@@ -848,6 +881,7 @@ pub struct RangeSelectExec {
metric: ExecutionPlanMetricsSet,
schema_project: Option<Vec<usize>>,
schema_before_project: SchemaRef,
cache: PlanProperties,
}
impl DisplayAs for RangeSelectExec {
@@ -882,16 +916,12 @@ impl ExecutionPlan for RangeSelectExec {
self.schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
Partitioning::UnknownPartitioning(1)
}
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![Distribution::SinglePartition]
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.input.output_ordering()
fn properties(&self) -> &PlanProperties {
&self.cache
}
fn children(&self) -> Vec<Arc<dyn DfPhysicalPlan>> {
@@ -915,6 +945,7 @@ impl ExecutionPlan for RangeSelectExec {
metric: self.metric.clone(),
schema_before_project: self.schema_before_project.clone(),
schema_project: self.schema_project.clone(),
cache: self.cache.clone(),
}))
}
@@ -963,8 +994,8 @@ impl ExecutionPlan for RangeSelectExec {
Some(self.metric.clone_inner())
}
fn statistics(&self) -> Statistics {
self.input.statistics()
fn statistics(&self) -> DataFusionResult<Statistics> {
Ok(Statistics::new_unknown(self.schema.as_ref()))
}
}
@@ -1054,7 +1085,7 @@ impl RangeSelectStream {
.iter()
.map(|expr| {
let value = expr.evaluate(batch)?;
Ok(value.into_array(batch.num_rows()))
value.into_array(batch.num_rows())
})
.collect::<DfResult<Vec<_>>>()
}
@@ -1168,7 +1199,7 @@ impl RangeSelectStream {
for SeriesState {
row,
align_ts_accumulator,
} in self.series_map.values()
} in self.series_map.values_mut()
{
// skip empty time series
if align_ts_accumulator.is_empty() {
@@ -1184,8 +1215,8 @@ impl RangeSelectStream {
align_ts_accumulator.keys().copied().collect::<Vec<_>>()
};
for ts in &align_ts {
if let Some(slot) = align_ts_accumulator.get(ts) {
for (column, acc) in all_scalar.iter_mut().zip(slot.iter()) {
if let Some(slot) = align_ts_accumulator.get_mut(ts) {
for (column, acc) in all_scalar.iter_mut().zip(slot.iter_mut()) {
column.push(acc.evaluate()?);
}
} else {
@@ -1415,6 +1446,11 @@ mod test {
Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true),
Field::new("host", DataType::Utf8, true),
]));
let cache = PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
);
let range_select_exec = Arc::new(RangeSelectExec {
input: memory_exec,
range_exec: vec![
@@ -1450,6 +1486,7 @@ mod test {
schema_project: None,
by_schema: Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])),
metric: ExecutionPlanMetricsSet::new(),
cache,
});
let sort_exec = SortExec::new(
vec![
src/query/src/range_select/plan_rewrite.rs
+24 -28
View File
@@ -25,9 +25,8 @@ use common_time::{Interval, Timestamp, Timezone};
use datafusion::datasource::DefaultTableSource;
use datafusion::prelude::Column;
use datafusion::scalar::ScalarValue;
use datafusion_common::tree_node::{TreeNode, TreeNodeRewriter, VisitRecursion};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
use datafusion_common::{DFSchema, DataFusionError, Result as DFResult};
use datafusion_expr::expr::ScalarUDF;
use datafusion_expr::{
Aggregate, Analyze, Explain, Expr, ExprSchemable, Extension, LogicalPlan, LogicalPlanBuilder,
Projection,
@@ -163,11 +162,7 @@ fn parse_expr_list(args: &[Expr], start: usize, len: usize) -> DFResult<Vec<Expr
for i in start..start + len {
outs.push(match &args.get(i) {
Some(
Expr::Column(_)
| Expr::Literal(_)
| Expr::BinaryExpr(_)
| Expr::ScalarFunction(_)
| Expr::ScalarUDF(_),
Expr::Column(_) | Expr::Literal(_) | Expr::BinaryExpr(_) | Expr::ScalarFunction(_),
) => args[i].clone(),
other => {
return Err(dispose_parse_error(*other));
@@ -195,11 +190,11 @@ macro_rules! inconsistent_check {
}
impl<'a> TreeNodeRewriter for RangeExprRewriter<'a> {
type N = Expr;
type Node = Expr;
fn mutate(&mut self, node: Expr) -> DFResult<Expr> {
if let Expr::ScalarUDF(func) = &node {
if func.fun.name == "range_fn" {
fn f_down(&mut self, node: Expr) -> DFResult<Transformed<Expr>> {
if let Expr::ScalarFunction(func) = &node {
if func.name() == "range_fn" {
// `range_fn(func, range, fill, byc, [byv], align, to)`
// `[byv]` are variadic arguments, byc indicate the length of arguments
let range_expr = self.get_range_expr(&func.args, 0)?;
@@ -246,10 +241,10 @@ impl<'a> TreeNodeRewriter for RangeExprRewriter<'a> {
};
let alias = Expr::Column(Column::from_name(range_fn.name.clone()));
self.range_fn.insert(range_fn);
return Ok(alias);
return Ok(Transformed::yes(alias));
}
}
Ok(node)
Ok(Transformed::no(node))
}
}
@@ -317,7 +312,7 @@ impl RangePlanRewriter {
};
let new_expr = expr
.iter()
.map(|expr| expr.clone().rewrite(&mut range_rewriter))
.map(|expr| expr.clone().rewrite(&mut range_rewriter).map(|x| x.data))
.collect::<DFResult<Vec<_>>>()
.context(DataFusionSnafu)?;
if range_rewriter.by.is_empty() {
@@ -385,7 +380,7 @@ impl RangePlanRewriter {
.context(DataFusionSnafu)?
.build()
}
_ => plan.with_new_inputs(&inputs),
_ => plan.with_new_exprs(plan.expressions(), inputs),
}
.context(DataFusionSnafu)?;
Ok(Some(plan))
@@ -401,10 +396,11 @@ impl RangePlanRewriter {
/// If the user does not explicitly use the `by` keyword to indicate time series,
/// `[row_columns]` will be use as default time series
async fn get_index_by(&mut self, schema: &Arc<DFSchema>) -> Result<(Expr, Vec<Expr>)> {
let mut time_index_expr = Expr::Wildcard;
let mut time_index_expr = Expr::Wildcard { qualifier: None };
let mut default_by = vec![];
for field in schema.fields() {
if let Some(table_ref) = field.qualifier() {
for i in 0..schema.fields().len() {
let (qualifier, _) = schema.qualified_field(i);
if let Some(table_ref) = qualifier {
let table = self
.table_provider
.resolve_table(table_ref.clone())
@@ -446,7 +442,7 @@ impl RangePlanRewriter {
}
}
}
if time_index_expr == Expr::Wildcard {
if matches!(time_index_expr, Expr::Wildcard { .. }) {
TimeIndexNotFoundSnafu {
table: schema.to_string(),
}
@@ -461,13 +457,13 @@ fn have_range_in_exprs(exprs: &[Expr]) -> bool {
exprs.iter().any(|expr| {
let mut find_range = false;
let _ = expr.apply(&mut |expr| {
if let Expr::ScalarUDF(ScalarUDF { fun, .. }) = expr {
if fun.name == "range_fn" {
Ok(match expr {
Expr::ScalarFunction(func) if func.name() == "range_fn" => {
find_range = true;
return Ok(VisitRecursion::Stop);
TreeNodeRecursion::Stop
}
}
Ok(VisitRecursion::Continue)
_ => TreeNodeRecursion::Continue,
})
});
find_range
})
@@ -581,8 +577,8 @@ mod test {
let query =
r#"SELECT (covar(field_0 + field_1, field_1)/4) RANGE '5m' FROM test ALIGN '1h';"#;
let expected = String::from(
"Projection: COVARIANCE(test.field_0 + test.field_1,test.field_1) RANGE 5m / Int64(4) [COVARIANCE(test.field_0 + test.field_1,test.field_1) RANGE 5m / Int64(4):Float64;N]\
\n RangeSelect: range_exprs=[COVARIANCE(test.field_0 + test.field_1,test.field_1) RANGE 5m], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1, test.tag_2, test.tag_3, test.tag_4], time_index=timestamp [COVARIANCE(test.field_0 + test.field_1,test.field_1) RANGE 5m:Float64;N, timestamp:Timestamp(Millisecond, None), tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8]\
"Projection: COVAR(test.field_0 + test.field_1,test.field_1) RANGE 5m / Int64(4) [COVAR(test.field_0 + test.field_1,test.field_1) RANGE 5m / Int64(4):Float64;N]\
\n RangeSelect: range_exprs=[COVAR(test.field_0 + test.field_1,test.field_1) RANGE 5m], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1, test.tag_2, test.tag_3, test.tag_4], time_index=timestamp [COVAR(test.field_0 + test.field_1,test.field_1) RANGE 5m:Float64;N, timestamp:Timestamp(Millisecond, None), tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8]\
\n TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(Millisecond, None), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]"
);
query_plan_compare(query, expected).await;
@@ -662,7 +658,7 @@ mod test {
async fn complex_range_expr() {
let query = r#"SELECT gcd(CAST(max(field_0 + 1) Range '5m' FILL NULL AS Int64), CAST(tag_0 AS Int64)) + round(max(field_2+1) Range '6m' FILL NULL + 1) + max(field_2+3) Range '10m' FILL NULL * CAST(tag_1 AS Float64) + 1 FROM test ALIGN '1h' by (tag_0, tag_1);"#;
let expected = String::from(
"Projection: gcd(CAST(MAX(test.field_0 + Int64(1)) RANGE 5m FILL NULL AS Int64), CAST(test.tag_0 AS Int64)) + round(MAX(test.field_2 + Int64(1)) RANGE 6m FILL NULL + Int64(1)) + MAX(test.field_2 + Int64(3)) RANGE 10m FILL NULL * CAST(test.tag_1 AS Float64) + Int64(1) [gcd(MAX(test.field_0 + Int64(1)) RANGE 5m FILL NULL,test.tag_0) + round(MAX(test.field_2 + Int64(1)) RANGE 6m FILL NULL + Int64(1)) + MAX(test.field_2 + Int64(3)) RANGE 10m FILL NULL * test.tag_1 + Int64(1):Float64;N]\
"Projection: gcd(arrow_cast(MAX(test.field_0 + Int64(1)) RANGE 5m FILL NULL, Utf8(\"Int64\")), arrow_cast(test.tag_0, Utf8(\"Int64\"))) + round(MAX(test.field_2 + Int64(1)) RANGE 6m FILL NULL + Int64(1)) + MAX(test.field_2 + Int64(3)) RANGE 10m FILL NULL * arrow_cast(test.tag_1, Utf8(\"Float64\")) + Int64(1) [gcd(arrow_cast(MAX(test.field_0 + Int64(1)) RANGE 5m FILL NULL,Utf8(\"Int64\")),arrow_cast(test.tag_0,Utf8(\"Int64\"))) + round(MAX(test.field_2 + Int64(1)) RANGE 6m FILL NULL + Int64(1)) + MAX(test.field_2 + Int64(3)) RANGE 10m FILL NULL * arrow_cast(test.tag_1,Utf8(\"Float64\")) + Int64(1):Float64;N]\
\n RangeSelect: range_exprs=[MAX(test.field_0 + Int64(1)) RANGE 5m FILL NULL, MAX(test.field_2 + Int64(1)) RANGE 6m FILL NULL, MAX(test.field_2 + Int64(3)) RANGE 10m FILL NULL], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [MAX(test.field_0 + Int64(1)) RANGE 5m FILL NULL:Float64;N, MAX(test.field_2 + Int64(1)) RANGE 6m FILL NULL:Float64;N, MAX(test.field_2 + Int64(3)) RANGE 10m FILL NULL:Float64;N, timestamp:Timestamp(Millisecond, None), tag_0:Utf8, tag_1:Utf8]\
\n TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(Millisecond, None), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]"
);
@@ -673,7 +669,7 @@ mod test {
async fn range_linear_on_integer() {
let query = r#"SELECT min(CAST(field_0 AS Int64) + CAST(field_1 AS Int64)) RANGE '5m' FILL LINEAR FROM test ALIGN '1h' by (tag_0,tag_1);"#;
let expected = String::from(
"RangeSelect: range_exprs=[MIN(test.field_0 + test.field_1) RANGE 5m FILL LINEAR], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [MIN(test.field_0 + test.field_1) RANGE 5m FILL LINEAR:Float64;N]\
"RangeSelect: range_exprs=[MIN(arrow_cast(test.field_0,Utf8(\"Int64\")) + arrow_cast(test.field_1,Utf8(\"Int64\"))) RANGE 5m FILL LINEAR], align=3600000ms, align_to=0ms, align_by=[test.tag_0, test.tag_1], time_index=timestamp [MIN(arrow_cast(test.field_0,Utf8(\"Int64\")) + arrow_cast(test.field_1,Utf8(\"Int64\"))) RANGE 5m FILL LINEAR:Float64;N]\
\n TableScan: test [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, tag_4:Utf8, timestamp:Timestamp(Millisecond, None), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N, field_3:Float64;N, field_4:Float64;N]"
);
query_plan_compare(query, expected).await;
src/query/src/sql.rs
+1 -1
View File
@@ -378,7 +378,7 @@ pub async fn show_index(
lit("").alias(INDEX_COMMENT_COLUMN),
lit(YES_STR).alias(INDEX_VISIBLE_COLUMN),
null().alias(INDEX_EXPRESSION_COLUMN),
Expr::Wildcard,
Expr::Wildcard { qualifier: None },
];
let projects = vec![
src/query/src/sql/show_create_table.rs
+3 -1
View File
@@ -46,7 +46,7 @@ fn string_value(s: impl Into<String>) -> SqlValue {
fn sql_option(name: &str, value: SqlValue) -> SqlOption {
SqlOption {
name: name.into(),
value,
value: Expr::Value(value),
}
}
@@ -141,6 +141,7 @@ fn create_table_constraints(
name: Some(TIME_INDEX.into()),
columns: vec![Ident::with_quote(quote_style, column_name)],
is_primary: false,
characteristics: None,
});
}
if !table_meta.primary_key_indices.is_empty() {
@@ -152,6 +153,7 @@ fn create_table_constraints(
name: None,
columns,
is_primary: true,
characteristics: None,
});
}
src/script/Cargo.toml
+3 -1
View File
@@ -11,6 +11,7 @@ python = [
"dep:datafusion",
"dep:datafusion-common",
"dep:datafusion-expr",
"dep:datafusion-functions",
"dep:datafusion-physical-expr",
"dep:rustpython-vm",
"dep:rustpython-parser",
@@ -45,6 +46,7 @@ crossbeam-utils = "0.8.14"
datafusion = { workspace = true, optional = true }
datafusion-common = { workspace = true, optional = true }
datafusion-expr = { workspace = true, optional = true }
datafusion-functions = { workspace = true, optional = true }
datafusion-physical-expr = { workspace = true, optional = true }
datatypes.workspace = true
futures.workspace = true
@@ -54,7 +56,7 @@ paste = { workspace = true, optional = true }
prometheus.workspace = true
query.workspace = true
# TODO(discord9): This is a forked and tweaked version of RustPython, please update it to newest original RustPython After RustPython support GC
pyo3 = { version = "0.19", optional = true, features = ["abi3", "abi3-py37"] }
pyo3 = { version = "0.20", optional = true, features = ["abi3", "abi3-py37"] }
rustpython-codegen = { git = "https://github.com/discord9/RustPython", optional = true, rev = "9ed5137412" }
rustpython-compiler = { git = "https://github.com/discord9/RustPython", optional = true, rev = "9ed5137412" }
rustpython-compiler-core = { git = "https://github.com/discord9/RustPython", optional = true, rev = "9ed5137412" }
src/script/src/python/engine.rs
+4
View File
@@ -126,6 +126,10 @@ impl Function for PyUDF {
}
fn signature(&self) -> common_query::prelude::Signature {
if self.copr.arg_types.is_empty() {
return Signature::any(0, Volatility::Volatile);
}
// try our best to get a type signature
let mut arg_types = Vec::with_capacity(self.copr.arg_types.len());
let mut know_all_types = true;
src/script/src/python/ffi_types/copr.rs
+24 -1
View File
@@ -21,8 +21,10 @@ use std::sync::{Arc, Weak};
use common_query::OutputData;
use common_recordbatch::{RecordBatch, RecordBatches};
use datafusion_common::ScalarValue;
use datatypes::arrow::compute;
use datatypes::data_type::{ConcreteDataType, DataType};
use datatypes::prelude::Value;
use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
use datatypes::vectors::{Helper, VectorRef};
// use crate::python::builtins::greptime_builtin;
@@ -42,7 +44,9 @@ use vm::{pyclass as rspyclass, PyObjectRef, PyPayload, PyResult, VirtualMachine}
use super::py_recordbatch::PyRecordBatch;
use crate::engine::EvalContext;
use crate::python::error::{ensure, ArrowSnafu, OtherSnafu, Result, TypeCastSnafu};
use crate::python::error::{
ensure, ArrowSnafu, DataFusionSnafu, OtherSnafu, Result, TypeCastSnafu,
};
use crate::python::ffi_types::PyVector;
#[cfg(feature = "pyo3_backend")]
use crate::python::pyo3::pyo3_exec_parsed;
@@ -179,6 +183,25 @@ impl Coprocessor {
/// check if real types and annotation types(if have) is the same, if not try cast columns to annotated type
pub(crate) fn check_and_cast_type(&self, cols: &mut [VectorRef]) -> Result<()> {
for col in cols.iter_mut() {
if let ConcreteDataType::List(x) = col.data_type() {
let values =
ScalarValue::convert_array_to_scalar_vec(col.to_arrow_array().as_ref())
.context(DataFusionSnafu)?
.into_iter()
.flatten()
.map(Value::try_from)
.collect::<std::result::Result<Vec<_>, _>>()
.context(TypeCastSnafu)?;
let mut builder = x.item_type().create_mutable_vector(values.len());
for v in values.iter() {
builder.push_value_ref(v.as_value_ref());
}
*col = builder.to_vector();
}
}
let return_types = &self.return_types;
// allow ignore Return Type Annotation
if return_types.is_empty() {
src/script/src/python/ffi_types/pair_tests.rs
+1
View File
@@ -209,6 +209,7 @@ fn eval_pyo3(case: CodeBlockTestCase) {
let res_vec = locals
.get_item("ret")
.unwrap()
.unwrap()
.extract::<PyVector>()
.map_err(|e| {
dbg!(&case.script);
src/script/src/python/ffi_types/vector.rs
+6 -11
View File
@@ -525,17 +525,12 @@ pub fn val_to_pyobj(val: value::Value, vm: &VirtualMachine) -> PyResult {
// FIXME(dennis): lose the timestamp unit here
Value::Timestamp(v) => vm.ctx.new_int(v.value()).into(),
value::Value::List(list) => {
let list = list.items().as_ref();
match list {
Some(list) => {
let list: Vec<_> = list
.iter()
.map(|v| val_to_pyobj(v.clone(), vm))
.collect::<Result<_, _>>()?;
vm.ctx.new_list(list).into()
}
None => vm.ctx.new_list(Vec::new()).into(),
}
let list: Vec<_> = list
.items()
.iter()
.map(|v| val_to_pyobj(v.clone(), vm))
.collect::<Result<_, _>>()?;
vm.ctx.new_list(list).into()
}
#[allow(unreachable_patterns)]
_ => return Err(vm.new_type_error(format!("Convert from {val:?} is not supported yet"))),

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