From 3d0ad5d8c86b0edbf5b47bb343a35e690b63b294 Mon Sep 17 00:00:00 2001 From: luofucong Date: Thu, 21 Mar 2024 11:35:47 +0800 Subject: [PATCH] chore: update datafusion --- Cargo.toml | 33 ++-- src/common/datasource/Cargo.toml | 3 +- src/common/datasource/src/file_format/csv.rs | 2 +- .../datasource/src/file_format/tests.rs | 3 +- src/common/datasource/src/test_util.rs | 6 +- .../function/src/scalars/aggregate/diff.rs | 22 +-- .../src/scalars/aggregate/percentile.rs | 5 +- .../function/src/scalars/aggregate/polyval.rs | 5 +- .../scalars/aggregate/scipy_stats_norm_cdf.rs | 5 +- .../scalars/aggregate/scipy_stats_norm_pdf.rs | 5 +- src/common/function/src/scalars/math.rs | 2 +- src/common/macro/src/range_fn.rs | 14 +- src/common/query/src/columnar_value.rs | 6 +- src/common/query/src/logical_plan.rs | 13 +- .../query/src/logical_plan/accumulator.rs | 4 +- src/common/query/src/logical_plan/expr.rs | 4 +- src/common/query/src/logical_plan/udaf.rs | 2 + src/common/query/src/logical_plan/udf.rs | 2 + src/common/query/src/physical_plan.rs | 58 +++---- src/common/query/src/signature.rs | 12 ++ src/common/recordbatch/src/adapter.rs | 2 +- src/common/recordbatch/src/error.rs | 8 +- src/common/recordbatch/src/filter.rs | 38 +++-- src/common/substrait/Cargo.toml | 4 - src/common/substrait/src/df_substrait.rs | 8 +- src/common/substrait/src/lib.rs | 4 +- src/common/time/src/datetime.rs | 14 +- src/common/time/src/timestamp.rs | 16 +- src/common/time/src/timezone.rs | 6 +- src/common/time/src/util.rs | 5 +- src/datanode/src/region_server.rs | 8 +- src/datatypes/src/error.rs | 9 ++ src/datatypes/src/scalars.rs | 6 +- src/datatypes/src/types/list_type.rs | 4 +- src/datatypes/src/value.rs | 147 +++++++++--------- src/datatypes/src/vectors/helper.rs | 50 +++--- src/datatypes/src/vectors/list.rs | 43 +---- src/file-engine/Cargo.toml | 1 + src/file-engine/src/query/file_stream.rs | 16 +- src/mito2/src/cache/cache_size.rs | 2 +- src/mito2/src/engine/basic_test.rs | 2 +- src/mito2/src/error.rs | 14 +- src/mito2/src/memtable/partition_tree/tree.rs | 15 +- src/mito2/src/memtable/time_series.rs | 36 +++-- src/mito2/src/sst/parquet/reader.rs | 16 +- src/mito2/src/sst/parquet/stats.rs | 14 +- src/operator/src/expr_factory.rs | 15 +- src/operator/src/statement/copy_table_from.rs | 5 +- src/promql/Cargo.toml | 3 + src/promql/src/extension_plan/empty_metric.rs | 57 ++++--- .../src/extension_plan/histogram_fold.rs | 26 +--- .../src/extension_plan/instant_manipulate.rs | 38 +++-- src/promql/src/extension_plan/normalize.rs | 15 +- .../src/extension_plan/range_manipulate.rs | 36 ++--- .../src/extension_plan/series_divide.rs | 24 +-- .../src/extension_plan/union_distinct_on.rs | 39 ++--- src/promql/src/functions/extrapolate_rate.rs | 42 ++--- src/promql/src/functions/holt_winters.rs | 14 +- src/promql/src/functions/idelta.rs | 14 +- src/promql/src/functions/predict_linear.rs | 14 +- src/promql/src/functions/quantile.rs | 14 +- src/promql/src/functions/test_util.rs | 12 +- src/promql/src/planner.rs | 75 ++++++--- src/query/src/datafusion.rs | 4 +- src/query/src/datafusion/planner.rs | 20 ++- src/query/src/dist_plan/analyzer.rs | 46 +++--- src/query/src/dist_plan/commutativity.rs | 9 +- src/query/src/dist_plan/merge_scan.rs | 27 ++-- src/query/src/dist_plan/planner.rs | 21 +-- src/query/src/optimizer/order_hint.rs | 17 +- .../src/optimizer/string_normalization.rs | 11 +- src/query/src/optimizer/type_conversion.rs | 38 +++-- src/query/src/parser.rs | 6 +- src/query/src/plan.rs | 3 +- src/query/src/query_engine/state.rs | 21 +-- src/query/src/range_select/plan.rs | 138 ++++++++-------- src/query/src/range_select/plan_rewrite.rs | 44 +++--- src/query/src/sql/show_create_table.rs | 4 +- src/script/Cargo.toml | 4 +- src/script/src/python/ffi_types/vector.rs | 17 +- src/script/src/python/pyo3/utils.rs | 7 +- src/script/src/python/rspython/builtins.rs | 53 ++++--- .../src/python/rspython/builtins/test.rs | 18 +-- src/script/src/python/rspython/utils.rs | 15 +- src/servers/src/mysql/helper.rs | 1 + src/servers/src/postgres/types.rs | 20 ++- src/servers/src/prom_store.rs | 4 +- src/sql/src/parser.rs | 18 +-- src/sql/src/parsers/alter_parser.rs | 14 +- src/sql/src/parsers/copy_parser.rs | 4 +- src/sql/src/parsers/create_parser.rs | 45 ++++-- src/sql/src/parsers/describe_parser.rs | 2 +- src/sql/src/parsers/drop_parser.rs | 4 +- src/sql/src/parsers/explain_parser.rs | 22 +-- src/sql/src/parsers/show_parser.rs | 10 +- src/sql/src/parsers/truncate_parser.rs | 2 +- src/sql/src/statements.rs | 36 +++-- src/sql/src/statements/insert.rs | 11 +- .../src/statements/transform/type_alias.rs | 76 +++++---- src/sql/src/util.rs | 11 +- src/table/src/predicate.rs | 9 +- src/table/src/predicate/stats.rs | 13 +- src/table/src/table/scan.rs | 19 ++- tests-integration/tests/sql.rs | 8 +- 104 files changed, 1070 insertions(+), 919 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 54354000b451..51395c3c7c55 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -75,11 +75,11 @@ rust.unknown_lints = "deny" [workspace.dependencies] 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 = "50.0" } +arrow-array = "50.0" +arrow-flight = "50.0" +arrow-ipc = { version = "50.0", features = ["lz4"] } +arrow-schema = { version = "50.0", features = ["serde"] } async-stream = "0.3" async-trait = "0.1" axum = { version = "0.6", features = ["headers"] } @@ -91,13 +91,14 @@ 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 = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } +datafusion-common = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } +datafusion-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } +datafusion-functions = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } +datafusion-optimizer = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } +datafusion-physical-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } +datafusion-sql = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } +datafusion-substrait = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b0b329ba39403b9e87156d6f9b8c5464dc6d2480" } derive_builder = "0.12" etcd-client = "0.12" fst = "0.4.7" @@ -118,7 +119,7 @@ opentelemetry-proto = { git = "https://github.com/waynexia/opentelemetry-rust.gi "metrics", "trace", ] } -parquet = "47.0" +parquet = "50.0" paste = "1.0" pin-project = "1.0" prometheus = { version = "0.13.3", features = ["process"] } @@ -140,13 +141,13 @@ 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" diff --git a/src/common/datasource/Cargo.toml b/src/common/datasource/Cargo.toml index 8f11043afd73..a246907d0836 100644 --- a/src/common/datasource/Cargo.toml +++ b/src/common/datasource/Cargo.toml @@ -30,7 +30,8 @@ derive_builder.workspace = true futures.workspace = true lazy_static.workspace = true object-store.workspace = true -orc-rust = "0.2" +# TODO(LFC): "orc-rs" has been archived, use "datafusion-orc" instead. +orc-rust = { git = "https://github.com/MichaelScofield/orc-rs.git", rev = "b6520cd68931f77b2941a7fee4643ab0339eaccc" } parquet.workspace = true paste = "1.0" regex = "1.7" diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index 0767722d9bbc..4cf2b9e1336d 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -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()); diff --git a/src/common/datasource/src/file_format/tests.rs b/src/common/datasource/src/file_format/tests.rs index be8650c9d5a1..7be8664b7231 100644 --- a/src/common/datasource/src/file_format/tests.rs +++ b/src/common/datasource/src/file_format/tests.rs @@ -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(); diff --git a/src/common/datasource/src/test_util.rs b/src/common/datasource/src/test_util.rs index 04125f161343..8f1af59c90e7 100644 --- a/src/common/datasource/src/test_util.rs +++ b/src/common/datasource/src/test_util.rs @@ -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, 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, } } diff --git a/src/common/function/src/scalars/aggregate/diff.rs b/src/common/function/src/scalars/aggregate/diff.rs index 9893d6199b71..b83ed6d00496 100644 --- a/src/common/function/src/scalars/aggregate/diff.rs +++ b/src/common/function/src/scalars/aggregate/diff.rs @@ -56,7 +56,7 @@ where .map(|&n| n.into()) .collect::>(); 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::>(); - 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() ); } diff --git a/src/common/function/src/scalars/aggregate/percentile.rs b/src/common/function/src/scalars/aggregate/percentile.rs index 49b981a7ee0e..231e0bf43a2c 100644 --- a/src/common/function/src/scalars/aggregate/percentile.rs +++ b/src/common/function/src/scalars/aggregate/percentile.rs @@ -104,10 +104,7 @@ where .map(|&n| n.into()) .collect::>(); 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(), ]) } diff --git a/src/common/function/src/scalars/aggregate/polyval.rs b/src/common/function/src/scalars/aggregate/polyval.rs index b56a692c8df7..ae6ca101c45a 100644 --- a/src/common/function/src/scalars/aggregate/polyval.rs +++ b/src/common/function/src/scalars/aggregate/polyval.rs @@ -72,10 +72,7 @@ where .map(|&n| n.into()) .collect::>(); 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(), ]) } diff --git a/src/common/function/src/scalars/aggregate/scipy_stats_norm_cdf.rs b/src/common/function/src/scalars/aggregate/scipy_stats_norm_cdf.rs index 2ec954051341..e6c92225a682 100644 --- a/src/common/function/src/scalars/aggregate/scipy_stats_norm_cdf.rs +++ b/src/common/function/src/scalars/aggregate/scipy_stats_norm_cdf.rs @@ -56,10 +56,7 @@ where .map(|&x| x.into()) .collect::>(); 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(), ]) } diff --git a/src/common/function/src/scalars/aggregate/scipy_stats_norm_pdf.rs b/src/common/function/src/scalars/aggregate/scipy_stats_norm_pdf.rs index d1bf432c993a..3045ae8665dd 100644 --- a/src/common/function/src/scalars/aggregate/scipy_stats_norm_pdf.rs +++ b/src/common/function/src/scalars/aggregate/scipy_stats_norm_pdf.rs @@ -56,10 +56,7 @@ where .map(|&x| x.into()) .collect::>(); 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(), ]) } diff --git a/src/common/function/src/scalars/math.rs b/src/common/function/src/scalars/math.rs index f7d50f881dfb..6635e70b171f 100644 --- a/src/common/function/src/scalars/math.rs +++ b/src/common/function/src/scalars/math.rs @@ -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 { diff --git a/src/common/macro/src/range_fn.rs b/src/common/macro/src/range_fn.rs index c907f1d0d165..582fff523dd4 100644 --- a/src/common/macro/src/range_fn.rs +++ b/src/common/macro/src/range_fn.rs @@ -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 { diff --git a/src/common/query/src/columnar_value.rs b/src/common/query/src/columnar_value.rs index 4e79d7819f0e..419bc5ed3a86 100644 --- a/src/common/query/src/columnar_value.rs +++ b/src/common/query/src/columnar_value.rs @@ -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 })? } diff --git a/src/common/query/src/logical_plan.rs b/src/common/query/src/logical_plan.rs index ac20c74b5b58..ab5dedc14f87 100644 --- a/src/common/query/src/logical_plan.rs +++ b/src/common/query/src/logical_plan.rs @@ -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) => { diff --git a/src/common/query/src/logical_plan/accumulator.rs b/src/common/query/src/logical_plan/accumulator.rs index f0c272ada1f5..32f1b4587c13 100644 --- a/src/common/query/src/logical_plan/accumulator.rs +++ b/src/common/query/src/logical_plan/accumulator.rs @@ -126,7 +126,7 @@ impl DfAccumulatorAdaptor { } impl DfAccumulator for DfAccumulatorAdaptor { - fn state(&self) -> DfResult> { + fn state(&mut self) -> DfResult> { 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 { + fn evaluate(&mut self) -> DfResult { let value = self.accumulator.evaluate()?; let output_type = self.creator.output_type()?; let scalar_value = value diff --git a/src/common/query/src/logical_plan/expr.rs b/src/common/query/src/logical_plan/expr.rs index e5abddc4d273..79f2363a6e12 100644 --- a/src/common/query/src/logical_plan/expr.rs +++ b/src/common/query/src/logical_plan/expr.rs @@ -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()); } } diff --git a/src/common/query/src/logical_plan/udaf.rs b/src/common/query/src/logical_plan/udaf.rs index b96de5b888f5..535bf46a4506 100644 --- a/src/common/query/src/logical_plan/udaf.rs +++ b/src/common/query/src/logical_plan/udaf.rs @@ -90,6 +90,8 @@ impl AggregateFunction { impl From for DfAggregateUdf { fn from(udaf: AggregateFunction) -> Self { + // TODO(LFC): See how to fit the new DataFusion UDAF implementation. + #[allow(deprecated)] DfAggregateUdf::new( &udaf.name, &udaf.signature.into(), diff --git a/src/common/query/src/logical_plan/udf.rs b/src/common/query/src/logical_plan/udf.rs index 31d356174502..df5cec762c6d 100644 --- a/src/common/query/src/logical_plan/udf.rs +++ b/src/common/query/src/logical_plan/udf.rs @@ -70,6 +70,8 @@ impl ScalarUdf { impl From for DfScalarUDF { fn from(udf: ScalarUdf) -> Self { + // TODO(LFC): remove deprecated + #[allow(deprecated)] DfScalarUDF::new( &udf.name, &udf.signature.into(), diff --git a/src/common/query/src/physical_plan.rs b/src/common/query/src/physical_plan.rs index ddfb1bfb6c1f..e921ada12cf9 100644 --- a/src/common/query/src/physical_plan.rs +++ b/src/common/query/src/physical_plan.rs @@ -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 { @@ -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> { 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 { self.0.metrics() } + + fn properties(&self) -> &PlanProperties { + self.0.properties() + } } impl DisplayAs for DfPhysicalPlanAdapter { @@ -234,8 +221,9 @@ mod test { use common_recordbatch::{RecordBatch, RecordBatches}; 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 +260,13 @@ mod test { _filters: &[Expr], _limit: Option, ) -> DfResult> { - 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 +280,10 @@ mod test { } } - #[derive(Debug)] + #[derive(Debug, Clone)] struct MyExecutionPlan { schema: SchemaRef, + properties: PlanProperties, } impl PhysicalPlan for MyExecutionPlan { @@ -303,8 +295,8 @@ mod test { self.schema.clone() } - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) + fn properties(&self) -> &PlanProperties { + &self.properties } fn children(&self) -> Vec { @@ -312,7 +304,7 @@ mod test { } fn with_new_children(&self, _children: Vec) -> Result { - unimplemented!() + Ok(Arc::new(self.clone())) } fn execute( @@ -381,7 +373,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::().unwrap(); diff --git a/src/common/query/src/signature.rs b/src/common/query/src/signature.rs index 9e92a10e1730..a234990bf69c 100644 --- a/src/common/query/src/signature.rs +++ b/src/common/query/src/signature.rs @@ -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 for DfTypeSignature { TypeSignature::OneOf(ts) => { DfTypeSignature::OneOf(ts.into_iter().map(Into::into).collect()) } + TypeSignature::VariadicAny => DfTypeSignature::VariadicAny, } } } diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index 12f5ecfdc9e1..e2d2baec5c12 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -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()); diff --git a/src/common/recordbatch/src/error.rs b/src/common/recordbatch/src/error.rs index 42a2754bb2c5..b22c0f488d74 100644 --- a/src/common/recordbatch/src/error.rs +++ b/src/common/recordbatch/src/error.rs @@ -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, diff --git a/src/common/recordbatch/src/filter.rs b/src/common/recordbatch/src/filter.rs index 3175ace37eaf..9470ab2ca860 100644 --- a/src/common/recordbatch/src/filter.rs +++ b/src/common/recordbatch/src/filter.rs @@ -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` @@ -43,7 +43,7 @@ pub struct SimpleFilterEvaluator { } impl SimpleFilterEvaluator { - pub fn try_new(predicate: &Expr) -> Option { + pub fn try_new(predicate: &Expr) -> Result> { match predicate { Expr::BinaryExpr(binary) => { // check if the expr is in the supported form @@ -54,7 +54,7 @@ impl SimpleFilterEvaluator { | Operator::LtEq | Operator::Gt | Operator::GtEq => {} - _ => return None, + _ => return Ok(None), } // swap the expr if it is in the form of `literal` `op` `col` @@ -66,16 +66,19 @@ impl SimpleFilterEvaluator { op = op.swap().unwrap(); (col, lit) } - _ => return None, + _ => return Ok(None), }; - Some(Self { + let literal = rhs + .to_scalar() + .with_context(|_| ToArrowScalarSnafu { v: rhs.clone() })?; + Ok(Some(Self { column_name: lhs.name.clone(), - literal: rhs.clone().to_scalar(), + literal, op, - }) + })) } - _ => None, + _ => Ok(None), } } @@ -85,7 +88,10 @@ impl SimpleFilterEvaluator { } pub fn evaluate_scalar(&self, input: &ScalarValue) -> Result { - 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)) } @@ -139,7 +145,7 @@ mod test { op: Operator::Plus, right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))), }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + assert!(SimpleFilterEvaluator::try_new(&expr).unwrap().is_none()); // two literal is not supported let expr = Expr::BinaryExpr(BinaryExpr { @@ -147,7 +153,7 @@ mod test { op: Operator::Eq, right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))), }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + assert!(SimpleFilterEvaluator::try_new(&expr).unwrap().is_none()); // two column is not supported let expr = Expr::BinaryExpr(BinaryExpr { @@ -161,7 +167,7 @@ mod test { name: "bar".to_string(), })), }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + assert!(SimpleFilterEvaluator::try_new(&expr).unwrap().is_none()); // compound expr is not supported let expr = Expr::BinaryExpr(BinaryExpr { @@ -176,7 +182,7 @@ mod test { op: Operator::Eq, right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))), }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + assert!(SimpleFilterEvaluator::try_new(&expr).unwrap().is_none()); } #[test] @@ -201,7 +207,7 @@ mod test { name: "foo".to_string(), })), }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap().unwrap(); assert_eq!(evaluator.op, Operator::Gt); assert_eq!(evaluator.column_name, "foo".to_string()); } @@ -216,7 +222,7 @@ mod test { op: Operator::Eq, right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))), }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap().unwrap(); let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _; let result = evaluator.evaluate_array(&input_1).unwrap(); @@ -241,7 +247,7 @@ mod test { op: Operator::Lt, right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))), }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap().unwrap(); let input_1 = ScalarValue::Int64(Some(1)); let result = evaluator.evaluate_scalar(&input_1).unwrap(); diff --git a/src/common/substrait/Cargo.toml b/src/common/substrait/Cargo.toml index 9ac4fc150f77..b3f328ff44ba 100644 --- a/src/common/substrait/Cargo.toml +++ b/src/common/substrait/Cargo.toml @@ -22,10 +22,6 @@ promql.workspace = true prost.workspace = true snafu.workspace = true -[dependencies.substrait_proto] -package = "substrait" -version = "0.17" - [dev-dependencies] datatypes.workspace = true tokio.workspace = true diff --git a/src/common/substrait/src/df_substrait.rs b/src/common/substrait/src/df_substrait.rs index c4e1db9a560e..d69a4913f91d 100644 --- a/src/common/substrait/src/df_substrait.rs +++ b/src/common/substrait/src/df_substrait.rs @@ -16,16 +16,16 @@ use std::sync::Arc; use async_trait::async_trait; use bytes::{Buf, Bytes, BytesMut}; -use datafusion::catalog::CatalogList; +use datafusion::catalog::CatalogProviderList; use datafusion::execution::context::SessionState; use datafusion::execution::runtime_env::RuntimeEnv; 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 snafu::ResultExt; -use substrait_proto::proto::Plan; use crate::error::{DecodeDfPlanSnafu, DecodeRelSnafu, EncodeDfPlanSnafu, EncodeRelSnafu, Error}; use crate::extension_serializer::ExtensionSerializer; @@ -42,7 +42,7 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor { async fn decode( &self, message: B, - catalog_list: Arc, + catalog_list: Arc, catalog: &str, schema: &str, ) -> Result { @@ -52,7 +52,7 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor { 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) diff --git a/src/common/substrait/src/lib.rs b/src/common/substrait/src/lib.rs index e0c3046b0868..5400482691d0 100644 --- a/src/common/substrait/src/lib.rs +++ b/src/common/substrait/src/lib.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use async_trait::async_trait; use bytes::{Buf, Bytes}; -use datafusion::catalog::CatalogList; +use datafusion::catalog::CatalogProviderList; pub use crate::df_substrait::DFLogicalSubstraitConvertor; @@ -35,7 +35,7 @@ pub trait SubstraitPlan { async fn decode( &self, message: B, - catalog_list: Arc, + catalog_list: Arc, catalog: &str, schema: &str, ) -> Result; diff --git a/src/common/time/src/datetime.rs b/src/common/time/src/datetime.rs index d12b87f7f164..f1980a38d1af 100644 --- a/src/common/time/src/datetime.rs +++ b/src/common/time/src/datetime.rs @@ -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 for serde_json::Value { impl From 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 { 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::from_timestamp_millis(self.0) + chrono::DateTime::from_timestamp_millis(self.0).map(|x| x.naive_utc()) } /// Format DateTime for given format and timezone. diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 7010ff8293a3..95dba04adcae 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -357,7 +357,7 @@ impl Timestamp { pub fn to_chrono_datetime(&self) -> Option { 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 { @@ -380,7 +380,7 @@ impl Timestamp { } pub fn from_chrono_datetime(ndt: NaiveDateTime) -> Option { - 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()); diff --git a/src/common/time/src/timezone.rs b/src/common/time/src/timezone.rs index 8d98b072a980..b1fefceb21d8 100644 --- a/src/common/time/src/timezone.rs +++ b/src/common/time/src/timezone.rs @@ -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(); diff --git a/src/common/time/src/util.rs b/src/common/time/src/util.rs index 6ce824764a2b..19fe3bc9119e 100644 --- a/src/common/time/src/util.rs +++ b/src/common/time/src/util.rs @@ -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() } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0dc9845850b2..cf38c5740a0e 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -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; @@ -730,7 +730,7 @@ impl DummyCatalogList { } } -impl CatalogList for DummyCatalogList { +impl CatalogProviderList for DummyCatalogList { fn as_any(&self) -> &dyn Any { self } @@ -788,8 +788,8 @@ impl SchemaProvider for DummySchemaProvider { vec![] } - async fn table(&self, _name: &str) -> Option> { - Some(self.table.clone()) + async fn table(&self, _name: &str) -> DfResult>> { + Ok(Some(self.table.clone())) } fn table_exist(&self, _name: &str) -> bool { diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index 316b50e3276c..16eae089099a 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -14,6 +14,7 @@ use std::any::Any; +use arrow_array::ArrayRef; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; @@ -139,6 +140,14 @@ pub enum Error { error: arrow::error::ArrowError, location: Location, }, + + #[snafu(display("Failed to convert Arrow array {:?} to scalars", array))] + ConvertArrowArrayToScalars { + array: ArrayRef, + #[snafu(source)] + error: datafusion_common::DataFusionError, + location: Location, + }, } impl ErrorExt for Error { diff --git a/src/datatypes/src/scalars.rs b/src/datatypes/src/scalars.rs index 60e6da218369..272d91eebfb0 100644 --- a/src/datatypes/src/scalars.rs +++ b/src/datatypes/src/scalars.rs @@ -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()); diff --git a/src/datatypes/src/types/list_type.rs b/src/datatypes/src/types/list_type.rs index 18115837f570..91f4c8654f5c 100644 --- a/src/datatypes/src/types/list_type.rs +++ b/src/datatypes/src/types/list_type.rs @@ -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", 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!( diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 64635eeae2d0..7b19e8f3b53d 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -17,6 +17,7 @@ use std::fmt::{Display, Formatter}; use std::sync::Arc; use arrow::datatypes::{DataType as ArrowDataType, Field}; +use arrow_array::{Array, ArrayRef, 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::>::default(); - let items = v.items().as_ref().unwrap_or(&default); - let items = items + let items = v + .items() .iter() .map(|i| i.to_string()) .collect::>() @@ -424,9 +423,10 @@ pub fn to_null_scalar_value(output_type: &ConcreteDataType) -> Result 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 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>>, + items: Vec, /// 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>>, datatype: ConcreteDataType) -> Self { + pub fn new(items: Vec, datatype: ConcreteDataType) -> Self { Self { items, datatype } } - pub fn items(&self) -> &Option>> { + pub fn items(&self) -> &[Value] { &self.items } @@ -739,38 +737,30 @@ impl ListValue { } fn try_to_scalar_value(&self, output_type: &ListType) -> Result { - let vs = if let Some(items) = self.items() { - Some( - items - .iter() - .map(|v| v.try_to_scalar_value(output_type.item_type())) - .collect::>>()?, - ) - } 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::>>()?; + 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,16 @@ impl TryFrom 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::>()?; - 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()) + .with_context(|_| ConvertArrowArrayToScalarsSnafu { + array: array as ArrayRef, + })? + .into_iter() + .flatten() + .map(|x| x.try_into()) + .collect::>>()?; Value::List(ListValue::new(items, datatype)) } ScalarValue::Date32(d) => d.map(|x| Value::Date(Date::new(x))).unwrap_or(Value::Null), @@ -891,7 +880,9 @@ impl TryFrom for Value { .map(|v| Value::Decimal128(Decimal128::new(v, p, s))) .unwrap_or(Value::Null), ScalarValue::Decimal256(_, _, _) - | ScalarValue::Struct(_, _) + | ScalarValue::Struct(_) + | ScalarValue::FixedSizeList(_) + | ScalarValue::LargeList(_) | ScalarValue::Dictionary(_, _) => { return error::UnsupportedArrowTypeSnafu { arrow_type: v.data_type(), @@ -1380,19 +1371,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() ); @@ -1662,7 +1656,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(), )), ); @@ -1827,7 +1821,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(), })) ); @@ -1895,7 +1889,7 @@ mod tests { check_as_value_ref!(DateTime, DateTime::new(1034)); let list = ListValue { - items: None, + items: vec![], datatype: ConcreteDataType::int32_datatype(), }; assert_eq!( @@ -1933,7 +1927,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); @@ -1989,7 +1983,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(), @@ -1997,7 +1991,7 @@ mod tests { ); assert_eq!( Value::List(ListValue::new( - Some(Box::default()), + vec![], ConcreteDataType::timestamp_second_datatype(), )) .to_string(), @@ -2005,7 +1999,7 @@ mod tests { ); assert_eq!( Value::List(ListValue::new( - Some(Box::default()), + vec![], ConcreteDataType::timestamp_millisecond_datatype(), )) .to_string(), @@ -2013,7 +2007,7 @@ mod tests { ); assert_eq!( Value::List(ListValue::new( - Some(Box::default()), + vec![], ConcreteDataType::timestamp_microsecond_datatype(), )) .to_string(), @@ -2021,7 +2015,7 @@ mod tests { ); assert_eq!( Value::List(ListValue::new( - Some(Box::default()), + vec![], ConcreteDataType::timestamp_nanosecond_datatype(), )) .to_string(), @@ -2251,19 +2245,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::>(); assert_eq!( vs, vec![ScalarValue::Int32(Some(-1)), ScalarValue::Int32(None)] @@ -2364,10 +2368,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(), }, }), @@ -2384,7 +2388,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 })); diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs index ef7dab842fb1..f58a9069acc9 100644 --- a/src/datatypes/src/vectors/helper.rs +++ b/src/datatypes/src/vectors/helper.rs @@ -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,20 @@ 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::>()?; - 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()) + .with_context(|_| ConvertArrowArrayToScalarsSnafu { + array: array as ArrayRef, + })? + .into_iter() + .flatten() + .map(ScalarValue::try_into) + .collect::>>()?; + builder.push(Some(ListValueRef::Ref { + val: &ListValue::new(values, item_type), + })); let list_vector = builder.to_vector(); ConstantVector::new(list_vector, length) } @@ -236,7 +237,9 @@ impl Helper { ConstantVector::new(Arc::new(vector), length) } ScalarValue::Decimal256(_, _, _) - | ScalarValue::Struct(_, _) + | ScalarValue::Struct(_) + | ScalarValue::FixedSizeList(_) + | ScalarValue::LargeList(_) | ScalarValue::Dictionary(_, _) => { return error::ConversionSnafu { from: format!("Unsupported scalar value: {value}"), @@ -396,7 +399,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 +489,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 +501,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); } } diff --git a/src/datatypes/src/vectors/list.rs b/src/datatypes/src/vectors/list.rs index aaf88bbd9866..99399edcf468 100644 --- a/src/datatypes/src/vectors/list.rs +++ b/src/datatypes/src/vectors/list.rs @@ -114,10 +114,7 @@ impl Vector for ListVector { let values = (0..vector.len()) .map(|i| vector.get(i)) .collect::>(); - 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(), ), })); diff --git a/src/file-engine/Cargo.toml b/src/file-engine/Cargo.toml index 3ce83b6b791c..c9c5865e1289 100644 --- a/src/file-engine/Cargo.toml +++ b/src/file-engine/Cargo.toml @@ -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 diff --git a/src/file-engine/src/query/file_stream.rs b/src/file-engine/src/query/file_stream.rs index 105785979eee..f3afd1fbc529 100644 --- a/src/file-engine/src/query/file_stream.rs +++ b/src/file-engine/src/query/file_stream.rs @@ -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( projection: Option<&Vec>, limit: Option, ) -> Result { + 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( .iter() .map(|filename| PartitionedFile::new(filename.to_string(), 0)) .collect::>()], - 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::>()], - 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(), ))); diff --git a/src/mito2/src/cache/cache_size.rs b/src/mito2/src/cache/cache_size.rs index 8ecd2d5e99c6..3d79bcbe5801 100644 --- a/src/mito2/src/cache/cache_size.rs +++ b/src/mito2/src/cache/cache_size.rs @@ -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)); } } diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 78cdb285450d..3c554e442ffb 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -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. diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 39c1527e0834..a8dce120ec19 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -27,6 +27,7 @@ use prost::{DecodeError, EncodeError}; use snafu::{Location, Snafu}; use store_api::manifest::ManifestVersion; use store_api::storage::RegionId; +use table::predicate::Predicate; use crate::cache::file_cache::FileType; use crate::sst::file::FileId; @@ -575,6 +576,13 @@ pub enum Error { #[snafu(display("Invalid region options, {}", reason))] InvalidRegionOptions { reason: String, location: Location }, + + #[snafu(display("Failed to create filter from predicate {:?}", predicate))] + CreateFilterFromPredicate { + predicate: Predicate, + source: common_recordbatch::error::Error, + location: Location, + }, } pub type Result = std::result::Result; @@ -676,7 +684,11 @@ impl ErrorExt for Error { CleanDir { .. } => StatusCode::Unexpected, InvalidConfig { .. } => StatusCode::InvalidArguments, StaleLogEntry { .. } => StatusCode::Unexpected, - FilterRecordBatch { source, .. } => source.status_code(), + + FilterRecordBatch { source, .. } | CreateFilterFromPredicate { source, .. } => { + source.status_code() + } + Upload { .. } => StatusCode::StorageUnavailable, BiError { .. } => StatusCode::Internal, EncodeMemtable { .. } | ReadDataPart { .. } => StatusCode::Internal, diff --git a/src/mito2/src/memtable/partition_tree/tree.rs b/src/mito2/src/memtable/partition_tree/tree.rs index ca0b478c8724..7d5091909a6b 100644 --- a/src/mito2/src/memtable/partition_tree/tree.rs +++ b/src/mito2/src/memtable/partition_tree/tree.rs @@ -30,7 +30,9 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::ColumnId; use table::predicate::Predicate; -use crate::error::{PrimaryKeyLengthMismatchSnafu, Result, SerializeFieldSnafu}; +use crate::error::{ + CreateFilterFromPredicateSnafu, PrimaryKeyLengthMismatchSnafu, Result, SerializeFieldSnafu, +}; use crate::flush::WriteBufferManagerRef; use crate::memtable::key_values::KeyValue; use crate::memtable::partition_tree::metrics::WriteMetrics; @@ -211,12 +213,15 @@ impl PartitionTree { }; let filters = predicate - .map(|p| { - p.exprs() + .map(|predicate| { + predicate + .exprs() .iter() - .filter_map(|f| SimpleFilterEvaluator::try_new(f.df_expr())) - .collect::>() + .filter_map(|f| SimpleFilterEvaluator::try_new(f.df_expr()).transpose()) + .collect::, _>>() + .context(CreateFilterFromPredicateSnafu { predicate }) }) + .transpose()? .unwrap_or_default(); let mut tree_iter_metric = TreeIterMetrics::default(); diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index 5403a8fea225..7f998c5d96c3 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -36,7 +36,10 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::ColumnId; use table::predicate::Predicate; -use crate::error::{ComputeArrowSnafu, ConvertVectorSnafu, PrimaryKeyLengthMismatchSnafu, Result}; +use crate::error::{ + ComputeArrowSnafu, ConvertVectorSnafu, CreateFilterFromPredicateSnafu, + PrimaryKeyLengthMismatchSnafu, Result, +}; use crate::flush::WriteBufferManagerRef; use crate::memtable::key_values::KeyValue; use crate::memtable::{ @@ -236,7 +239,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)) } @@ -346,7 +351,7 @@ impl SeriesSet { projection: HashSet, predicate: Option, dedup: bool, - ) -> Iter { + ) -> Result { let primary_key_schema = primary_key_schema(&self.region_metadata); let primary_key_datatypes = self .region_metadata @@ -354,7 +359,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, @@ -415,7 +420,7 @@ struct Iter { impl Iter { #[allow(clippy::too_many_arguments)] - pub(crate) fn new( + pub(crate) fn try_new( metadata: RegionMetadataRef, series: Arc, projection: HashSet, @@ -424,27 +429,30 @@ impl Iter { pk_datatypes: Vec, codec: Arc, dedup: bool, - ) -> Self { - let simple_filters = predicate - .map(|p| { - p.exprs() + ) -> Result { + let predicate = predicate + .map(|predicate| { + predicate + .exprs() .iter() - .filter_map(|f| SimpleFilterEvaluator::try_new(f.df_expr())) - .collect::>() + .filter_map(|f| SimpleFilterEvaluator::try_new(f.df_expr()).transpose()) + .collect::, _>>() + .context(CreateFilterFromPredicateSnafu { predicate }) }) + .transpose()? .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(), - } + }) } } diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 827e5d851c88..033b6449c6e0 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -40,8 +40,8 @@ use table::predicate::Predicate; use crate::cache::CacheManagerRef; use crate::error::{ - ArrowReaderSnafu, FieldTypeMismatchSnafu, FilterRecordBatchSnafu, InvalidMetadataSnafu, - InvalidParquetSnafu, ReadParquetSnafu, Result, + ArrowReaderSnafu, CreateFilterFromPredicateSnafu, FieldTypeMismatchSnafu, + FilterRecordBatchSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, ReadParquetSnafu, Result, }; use crate::metrics::{ PRECISE_FILTER_ROWS_TOTAL, READ_ROWS_IN_ROW_GROUP_TOTAL, READ_ROWS_TOTAL, @@ -180,11 +180,15 @@ 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() + .filter_map(|expr| SimpleFilterEvaluator::try_new(expr.df_expr()).transpose()) + .collect::, _>>() + .with_context(|_| CreateFilterFromPredicateSnafu { + predicate: predicate.clone(), + })? } else { vec![] }; diff --git a/src/mito2/src/sst/parquet/stats.rs b/src/mito2/src/sst/parquet/stats.rs index a17e8ace45ae..43e0e3cfb963 100644 --- a/src/mito2/src/sst/parquet/stats.rs +++ b/src/mito2/src/sst/parquet/stats.rs @@ -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> 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 { + // TODO(LFC): Impl it. + None + } + + fn contained(&self, _column: &Column, _values: &HashSet) -> Option { + // TODO(LFC): Impl it. + None + } } diff --git a/src/operator/src/expr_factory.rs b/src/operator/src/expr_factory.rs index 0975f9b17dbd..42d8ee3187c2 100644 --- a/src/operator/src/expr_factory.rs +++ b/src/operator/src/expr_factory.rs @@ -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 { name: Some(name), columns, is_primary: false, + .. } => { if name.value == TIME_INDEX { Some(columns.iter().map(|ident| &ident.value)) diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index c8c1ae368804..cb922742ff7b 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -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 { + 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, diff --git a/src/promql/Cargo.toml b/src/promql/Cargo.toml index 9e7bc01ae8ac..893b8b24c372 100644 --- a/src/promql/Cargo.toml +++ b/src/promql/Cargo.toml @@ -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" diff --git a/src/promql/src/extension_plan/empty_metric.rs b/src/promql/src/extension_plan/empty_metric.rs index 74d25d8ebb8f..2332025f78dc 100644 --- a/src/promql/src/extension_plan/empty_metric.rs +++ b/src/promql/src/extension_plan/empty_metric.rs @@ -20,16 +20,17 @@ 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::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::common::stats::Precision; use datafusion::common::{DFField, DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics}; 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}; @@ -102,22 +103,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(), })) } @@ -164,7 +166,7 @@ pub struct EmptyMetricExec { /// Schema of the output record batch result_schema: SchemaRef, expr: Option, - + properties: Arc, metric: ExecutionPlanMetricsSet, } @@ -177,12 +179,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 { @@ -222,16 +220,15 @@ impl ExecutionPlan for EmptyMetricExec { Some(self.metric.clone_inner()) } - fn statistics(&self) -> Statistics { + fn statistics(&self) -> DataFusionResult { 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::() 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: vec![], + }) } } @@ -285,21 +282,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 = 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 { diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index c5fabd4a0f84..848b1d3b31ba 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -28,12 +28,12 @@ use datafusion::common::{DFSchema, DFSchemaRef}; 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::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, + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PhysicalExpr, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, }; use datafusion::prelude::{Column, Expr}; use datatypes::prelude::{ConcreteDataType, DataType as GtDataType}; @@ -150,6 +150,7 @@ impl HistogramFold { .map(|f| f.qualified_column()) .collect(), }, + Box::new(None), )); } else { Ok(()) @@ -231,12 +232,8 @@ impl ExecutionPlan for HistogramFoldExec { 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.input.properties() } fn required_input_ordering(&self) -> Vec>> { @@ -342,15 +339,6 @@ impl ExecutionPlan for HistogramFoldExec { fn metrics(&self) -> Option { Some(self.metric.clone_inner()) } - - fn statistics(&self) -> Statistics { - Statistics { - num_rows: None, - total_byte_size: None, - column_statistics: None, - is_exact: false, - } - } } impl HistogramFoldExec { @@ -594,7 +582,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 diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index e65592bb374e..989ac3d2a705 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -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::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_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream, + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; use datatypes::arrow::compute; @@ -194,12 +194,8 @@ impl ExecutionPlan for InstantManipulateExec { self.input.schema() } - 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 maintains_input_order(&self) -> Vec { @@ -262,23 +258,25 @@ impl ExecutionPlan for InstantManipulateExec { Some(self.metric.clone_inner()) } - fn statistics(&self) -> Statistics { - let input_stats = self.input.statistics(); + fn statistics(&self) -> DataFusionResult { + 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 _); - - Statistics { - num_rows: Some(estimated_row_num.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); + + 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![], + }) } } @@ -438,7 +436,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) } } diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index 968e8a6d3782..957c55fade54 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -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> { @@ -223,7 +218,7 @@ impl ExecutionPlan for SeriesNormalizeExec { Some(self.metric.clone_inner()) } - fn statistics(&self) -> Statistics { + fn statistics(&self) -> DataFusionResult { 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) } } diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index e6e93d8b50d2..dd4816db2cd7 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -23,14 +23,14 @@ use datafusion::arrow::compute; use datafusion::arrow::datatypes::{Field, SchemaRef}; use datafusion::arrow::error::ArrowError; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::stats::Precision; use datafusion::common::{DFField, 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, ExecutionPlan, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; use datafusion::sql::TableReference; @@ -258,12 +258,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.input.properties() } fn maintains_input_order(&self) -> Vec { @@ -333,23 +329,25 @@ impl ExecutionPlan for RangeManipulateExec { Some(self.metric.clone_inner()) } - fn statistics(&self) -> Statistics { - let input_stats = self.input.statistics(); + fn statistics(&self) -> DataFusionResult { + 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: vec![], + }) } } @@ -452,7 +450,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)>) { diff --git a/src/promql/src/extension_plan/series_divide.rs b/src/promql/src/extension_plan/series_divide.rs index 4abf8d87885a..d0524b5d5391 100644 --- a/src/promql/src/extension_plan/series_divide.rs +++ b/src/promql/src/extension_plan/series_divide.rs @@ -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 { @@ -156,10 +156,6 @@ impl ExecutionPlan for SeriesDivideExec { } } - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - self.input.output_ordering() - } - fn maintains_input_order(&self) -> Vec { vec![true; self.children().len()] } @@ -212,16 +208,6 @@ impl ExecutionPlan for SeriesDivideExec { fn metrics(&self) -> Option { 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 { diff --git a/src/promql/src/extension_plan/union_distinct_on.rs b/src/promql/src/extension_plan/union_distinct_on.rs index 20e40415c99c..4624544645da 100644 --- a/src/promql/src/extension_plan/union_distinct_on.rs +++ b/src/promql/src/extension_plan/union_distinct_on.rs @@ -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, right_exec: Arc, ) -> Arc { + 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, /// 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> { @@ -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 { 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, _>>() - .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) } diff --git a/src/promql/src/functions/extrapolate_rate.rs b/src/promql/src/functions/extrapolate_rate.rs index 7cba367145f3..7a57efae04a6 100644 --- a/src/promql/src/functions/extrapolate_rate.rs +++ b/src/promql/src/functions/extrapolate_rate.rs @@ -204,15 +204,17 @@ impl ExtrapolatedRate { } 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 { } 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 { } 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 _), + ) } } diff --git a/src/promql/src/functions/holt_winters.rs b/src/promql/src/functions/holt_winters.rs index 4480ccc52dc4..c047c1883d04 100644 --- a/src/promql/src/functions/holt_winters.rs +++ b/src/promql/src/functions/holt_winters.rs @@ -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 { diff --git a/src/promql/src/functions/idelta.rs b/src/promql/src/functions/idelta.rs index cf8ad2d0d63b..9a74b65fec27 100644 --- a/src/promql/src/functions/idelta.rs +++ b/src/promql/src/functions/idelta.rs @@ -42,15 +42,17 @@ impl IDelta { } 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 diff --git a/src/promql/src/functions/predict_linear.rs b/src/promql/src/functions/predict_linear.rs index 965fa28afce0..c9b24a76a896 100644 --- a/src/promql/src/functions/predict_linear.rs +++ b/src/promql/src/functions/predict_linear.rs @@ -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 diff --git a/src/promql/src/functions/quantile.rs b/src/promql/src/functions/quantile.rs index 62ff2f2126e0..d055ad122734 100644 --- a/src/promql/src/functions/quantile.rs +++ b/src/promql/src/functions/quantile.rs @@ -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 diff --git a/src/promql/src/functions/test_util.rs b/src/promql/src/functions/test_util.rs index 5b9d4adef310..8f9558e59046 100644 --- a/src/promql/src/functions/test_util.rs +++ b/src/promql/src/functions/test_util.rs @@ -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> = extract_array(&(range_fn.fun)(&input).unwrap()) + let eval_result: Vec> = extract_array(&(range_fn.fun())(&input).unwrap()) .unwrap() .as_any() .downcast_ref::() .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, + })); } diff --git a/src/promql/src/planner.rs b/src/promql/src/planner.rs index a2d0c2d79a63..f244ca3bc3f6 100644 --- a/src/promql/src/planner.rs +++ b/src/promql/src/planner.rs @@ -22,18 +22,21 @@ use catalog::table_source::DfTableSourceProvider; use common_query::prelude::GREPTIME_VALUE; use datafusion::common::{DFSchemaRef, OwnedTableReference, 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, @@ -1093,7 +1096,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 { @@ -1102,21 +1107,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 { @@ -1126,12 +1140,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( @@ -1140,8 +1166,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); @@ -1158,8 +1184,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); @@ -1223,7 +1249,7 @@ impl PromPlanner { exprs.push(expr); } - utils::conjunction(exprs).context(ValueNotFoundSnafu { + conjunction(exprs).context(ValueNotFoundSnafu { table: self.ctx.table_name.clone().unwrap(), }) } @@ -1264,11 +1290,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(); @@ -1480,13 +1507,13 @@ 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::BuiltIn(BuiltinScalarFunction::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), args: vec![lhs, rhs], })) })), @@ -1918,7 +1945,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::, name)))); // chain non-field columns (unchanged) and field columns (applied computation then alias) let project_fields = non_field_columns_iter @@ -1972,7 +1999,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) @@ -1988,6 +2015,8 @@ struct FunctionArgs { #[derive(Debug, Clone)] enum ScalarFunc { DataFusionBuiltin(BuiltinScalarFunction), + /// The UDF that is defined by Datafusion itself. + DataFusionUdf(Arc), Udf(ScalarUdfDef), // todo(ruihang): maybe merge with Udf later /// UDF that require extra information like range length to be evaluated. @@ -2382,7 +2411,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] diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index b6f0e9543f97..ead9e2d41361 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -435,7 +435,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 +453,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) diff --git a/src/query/src/datafusion/planner.rs b/src/query/src/datafusion/planner.rs index 67ec4c8a3c10..320fbb962270 100644 --- a/src/query/src/datafusion/planner.rs +++ b/src/query/src/datafusion/planner.rs @@ -26,10 +26,11 @@ 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::var_provider::is_system_variables; use datafusion_expr::{TableSource, WindowUDF}; -use datafusion_physical_expr::var_provider::{is_system_variables, VarType}; use datafusion_sql::parser::Statement as DfStatement; use session::context::QueryContextRef; use snafu::ResultExt; @@ -102,7 +103,7 @@ async fn resolve_tables( } impl ContextProvider for DfContextProviderAdapter { - fn get_table_provider(&self, name: TableReference) -> DfResult> { + fn get_table_source(&self, name: TableReference) -> DfResult> { let table_ref = self.table_provider.resolve_table_ref(name)?; self.tables .get(&table_ref.to_string()) @@ -159,4 +160,19 @@ impl ContextProvider for DfContextProviderAdapter { fn options(&self) -> &ConfigOptions { self.session_state.config_options() } + + fn udfs_names(&self) -> Vec { + // TODO(LFC): Impl it. + vec![] + } + + fn udafs_names(&self) -> Vec { + // TODO(LFC): Impl it. + vec![] + } + + fn udwfs_names(&self) -> Vec { + // TODO(LFC): Impl it. + vec![] + } } diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 10c2c2e9e25b..2846eef992fb 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -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; @@ -52,7 +52,7 @@ impl AnalyzerRule for DistPlannerAnalyzer { 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::>>()?; let inputs = plan.inputs().into_iter().cloned().collect::>(); - Ok(Transformed::Yes(plan.with_new_exprs(exprs, &inputs)?)) + Ok(Transformed::yes(plan.with_new_exprs(exprs, inputs)?)) } fn transform_subquery(expr: Expr) -> DfResult> { 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 { 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 { + fn f_down<'a>(&mut self, node: Self::Node) -> DfResult> { 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 { + fn f_up(&mut self, node: Self::Node) -> DfResult> { // only expand once on each ascending if self.is_expanded() { self.pop_stack(); - return Ok(node); + return Ok(Transformed::yes(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::yes(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::yes(node)) } } diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index 5bb125e1ba18..f59dc57268cb 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -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, diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index ed31ecb0c87c..a08687eb9edc 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -33,10 +33,12 @@ 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}; @@ -123,6 +125,7 @@ pub struct MergeScanExec { arrow_schema: ArrowSchemaRef, region_query_handler: RegionQueryHandlerRef, metric: ExecutionPlanMetricsSet, + properties: PlanProperties, } impl std::fmt::Debug for MergeScanExec { @@ -143,6 +146,11 @@ impl MergeScanExec { arrow_schema: &ArrowSchema, region_query_handler: RegionQueryHandlerRef, ) -> Result { + let properties = PlanProperties::new( + EquivalenceProperties::new(Arc::new(arrow_schema.clone())), + Partitioning::UnknownPartitioning(1), + ExecutionMode::Bounded, + ); let arrow_schema_without_metadata = Self::arrow_schema_without_metadata(arrow_schema); let schema_without_metadata = Self::arrow_schema_to_schema(arrow_schema_without_metadata.clone())?; @@ -154,6 +162,7 @@ impl MergeScanExec { arrow_schema: arrow_schema_without_metadata, region_query_handler, metric: ExecutionPlanMetricsSet::new(), + properties, }) } @@ -266,12 +275,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> { @@ -297,10 +302,6 @@ impl ExecutionPlan for MergeScanExec { ))) } - fn statistics(&self) -> Statistics { - Statistics::default() - } - fn metrics(&self) -> Option { Some(self.metric.clone_inner()) } diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index c3e0cca94e52..5c1af19bb1e4 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -25,7 +25,7 @@ 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; @@ -125,6 +125,7 @@ impl DistExtensionPlanner { /// Apply the fully resolved table name to the TableScan plan fn plan_with_full_table_name(plan: LogicalPlan, name: &TableName) -> Result { plan.transform(&|plan| TableNameRewriter::rewrite_table_name(plan, name)) + .map(|x| x.data) } async fn get_regions(&self, table_name: &TableName) -> Result> { @@ -163,9 +164,9 @@ struct TableNameExtractor { } impl TreeNodeVisitor for TableNameExtractor { - type N = LogicalPlan; + type Node = LogicalPlan; - fn pre_visit(&mut self, node: &Self::N) -> Result { + fn f_down(&mut self, node: &Self::Node) -> Result { match node { LogicalPlan::TableScan(scan) => { if let Some(source) = scan.source.as_any().downcast_ref::() { @@ -182,7 +183,7 @@ impl TreeNodeVisitor for TableNameExtractor { info.name.clone(), )); } - return Ok(VisitRecursion::Stop); + return Ok(TreeNodeRecursion::Stop); } } match &scan.table_name { @@ -196,7 +197,7 @@ impl TreeNodeVisitor for TableNameExtractor { schema.clone(), table.clone(), )); - Ok(VisitRecursion::Stop) + Ok(TreeNodeRecursion::Stop) } // TODO(ruihang): Maybe the following two cases should not be valid TableReference::Partial { schema, table } => { @@ -205,7 +206,7 @@ impl TreeNodeVisitor for TableNameExtractor { schema.clone(), table.clone(), )); - Ok(VisitRecursion::Stop) + Ok(TreeNodeRecursion::Stop) } TableReference::Bare { table } => { self.table_name = Some(TableName::new( @@ -213,11 +214,11 @@ impl TreeNodeVisitor for TableNameExtractor { DEFAULT_SCHEMA_NAME.to_string(), table.clone(), )); - Ok(VisitRecursion::Stop) + Ok(TreeNodeRecursion::Stop) } } } - _ => Ok(VisitRecursion::Continue), + _ => Ok(TreeNodeRecursion::Continue), } } } @@ -236,9 +237,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), }) } } diff --git a/src/query/src/optimizer/order_hint.rs b/src/query/src/optimizer/order_hint.rs index b9a34df0a8bd..6c027568a244 100644 --- a/src/query/src/optimizer/order_hint.rs +++ b/src/query/src/optimizer/order_hint.rs @@ -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 { + fn f_down(&mut self, node: &Self::Node) -> DataFusionResult { 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) } } diff --git a/src/query/src/optimizer/string_normalization.rs b/src/query/src/optimizer/string_normalization.rs index 1dfcc4d518d0..56fb36c3df14 100644 --- a/src/query/src/optimizer/string_normalization.rs +++ b/src/query/src/optimizer/string_normalization.rs @@ -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::>>()?; - 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 { + fn f_up(&mut self, expr: Expr) -> Result> { 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)) } } diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index aa1fc73d4c1a..1589971060cc 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -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::>>()?; - 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::>>()?; - 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 { + fn f_up(&mut self, expr: Expr) -> Result> { 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)) } } @@ -311,6 +313,7 @@ mod tests { use datafusion::logical_expr::expr::AggregateFunction as AggrExpr; use datafusion_common::{Column, DFField, DFSchema}; + use datafusion_expr::expr::AggregateFunctionDefinition; use datafusion_expr::{AggregateFunction, LogicalPlanBuilder}; use datafusion_sql::TableReference; use session::context::QueryContext; @@ -411,12 +414,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 ); } @@ -444,11 +448,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 +480,12 @@ mod tests { .aggregate( Vec::::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() diff --git a/src/query/src/parser.rs b/src/query/src/parser.rs index 92cab91481b3..6db1d2aa15cc 100644 --- a/src/query/src/parser.rs +++ b/src/query/src/parser.rs @@ -295,9 +295,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); } diff --git a/src/query/src/plan.rs b/src/query/src/plan.rs index 4462302d98e8..0e2dd710e7a4 100644 --- a/src/query/src/plan.rs +++ b/src/query/src/plan.rs @@ -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) } diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 18af09973e57..96e88255188f 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -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); diff --git a/src/query/src/range_select/plan.rs b/src/query/src/range_select/plan.rs index bff30d47891d..5ac3509873fe 100644 --- a/src/query/src/range_select/plan.rs +++ b/src/query/src/range_select/plan.rs @@ -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, 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_expr::expr::AggregateFunctionDefinition; use datafusion_expr::utils::{exprlist_to_fields, COUNT_STAR_EXPANSION}; use datafusion_expr::{ lit, Accumulator, AggregateFunction, Expr, ExprSchemable, LogicalPlan, @@ -46,7 +48,6 @@ 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, }; @@ -181,7 +182,7 @@ impl Accumulator for RangeFirstListValueAcc { Ok(()) } - fn evaluate(&self) -> DataFusionResult { + fn evaluate(&mut self) -> DataFusionResult { Ok(self.data.clone().unwrap_or(ScalarValue::Null)) } @@ -189,7 +190,7 @@ impl Accumulator for RangeFirstListValueAcc { std::mem::size_of_val(self) } - fn state(&self) -> DataFusionResult> { + fn state(&mut self) -> DataFusionResult> { unreachable!("Accumulator::state will not be used in range query") } @@ -483,8 +484,7 @@ impl RangeSelect { let time_index_name = ts_field.name().clone(); fields.push(ts_field); // add by - let by_fields = - exprlist_to_fields(by.iter().collect::>(), &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()) @@ -603,7 +603,6 @@ impl RangeSelect { is_count_aggr: bool, exprs: &[Expr], df_schema: &Arc, - schema: &Schema, session_state: &SessionState, ) -> DfResult>> { exprs @@ -614,13 +613,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::>>() } @@ -650,10 +648,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 +678,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 +701,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 +716,6 @@ impl RangeSelect { false, &aggr.args, input_dfschema, - &input_schema, session_state, )?; // first_value/last_value has only one param. @@ -723,8 +733,7 @@ impl RangeSelect { .map(|x| { create_physical_sort_expr( x, - input_dfschema, - &input_schema, + input_dfschema.as_ref(), session_state.execution_props(), ) }) @@ -732,36 +741,36 @@ 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()?, + ), + f => Err(DataFusionError::NotImplemented(format!( + "Range function from {f:?}" + ))), + } } _ => Err(DataFusionError::Plan(format!( "Unexpected Expr:{} in RangeSelect", @@ -788,18 +797,13 @@ impl RangeSelect { } else { schema_before_project.clone() }; + let by = self.create_physical_expr_list(false, &self.by, input_dfschema, session_state)?; 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)), @@ -882,16 +886,12 @@ impl ExecutionPlan for RangeSelectExec { self.schema.clone() } - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) - } - fn required_input_distribution(&self) -> Vec { vec![Distribution::SinglePartition] } - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - self.input.output_ordering() + fn properties(&self) -> &PlanProperties { + self.input.properties() } fn children(&self) -> Vec> { @@ -963,7 +963,7 @@ impl ExecutionPlan for RangeSelectExec { Some(self.metric.clone_inner()) } - fn statistics(&self) -> Statistics { + fn statistics(&self) -> DataFusionResult { self.input.statistics() } } @@ -1054,7 +1054,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::>>() } @@ -1168,7 +1168,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 +1184,8 @@ impl RangeSelectStream { align_ts_accumulator.keys().copied().collect::>() }; 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 { diff --git a/src/query/src/range_select/plan_rewrite.rs b/src/query/src/range_select/plan_rewrite.rs index bab81279e7a1..d27bae359074 100644 --- a/src/query/src/range_select/plan_rewrite.rs +++ b/src/query/src/range_select/plan_rewrite.rs @@ -25,9 +25,9 @@ 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::expr::ScalarFunction; use datafusion_expr::{ Aggregate, Analyze, Explain, Expr, ExprSchemable, Extension, LogicalPlan, LogicalPlanBuilder, Projection, @@ -163,11 +163,7 @@ fn parse_expr_list(args: &[Expr], start: usize, len: usize) -> DFResult args[i].clone(), other => { return Err(dispose_parse_error(*other)); @@ -195,11 +191,11 @@ macro_rules! inconsistent_check { } impl<'a> TreeNodeRewriter for RangeExprRewriter<'a> { - type N = Expr; + type Node = Expr; - fn mutate(&mut self, node: Expr) -> DFResult { - if let Expr::ScalarUDF(func) = &node { - if func.fun.name == "range_fn" { + fn f_down(&mut self, node: Expr) -> DFResult> { + if let Expr::ScalarFunction(func) = &node { + if func.func_def.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 +242,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 +313,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::>>() .context(DataFusionSnafu)?; if range_rewriter.by.is_empty() { @@ -385,7 +381,7 @@ impl RangePlanRewriter { .context(DataFusionSnafu)? .build() } - _ => plan.with_new_inputs(&inputs), + _ => plan.with_new_exprs(plan.expressions(), inputs), } .context(DataFusionSnafu)?; Ok(Some(plan)) @@ -401,7 +397,7 @@ 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) -> Result<(Expr, Vec)> { - 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() { @@ -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,15 @@ 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(ScalarFunction { func_def, .. }) + if func_def.name() == "range_fn" => + { find_range = true; - return Ok(VisitRecursion::Stop); + TreeNodeRecursion::Stop } - } - Ok(VisitRecursion::Continue) + _ => TreeNodeRecursion::Continue, + }) }); find_range }) @@ -662,7 +660,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]" ); diff --git a/src/query/src/sql/show_create_table.rs b/src/query/src/sql/show_create_table.rs index 97d8aba4fbbf..2004590aa561 100644 --- a/src/query/src/sql/show_create_table.rs +++ b/src/query/src/sql/show_create_table.rs @@ -46,7 +46,7 @@ fn string_value(s: impl Into) -> 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, }); } diff --git a/src/script/Cargo.toml b/src/script/Cargo.toml index 50dca7ccbb76..4602a334c9ef 100644 --- a/src/script/Cargo.toml +++ b/src/script/Cargo.toml @@ -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" } diff --git a/src/script/src/python/ffi_types/vector.rs b/src/script/src/python/ffi_types/vector.rs index 1f75efdecf69..bbede552cf15 100644 --- a/src/script/src/python/ffi_types/vector.rs +++ b/src/script/src/python/ffi_types/vector.rs @@ -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::>()?; - 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::>()?; + vm.ctx.new_list(list).into() } #[allow(unreachable_patterns)] _ => return Err(vm.new_type_error(format!("Convert from {val:?} is not supported yet"))), diff --git a/src/script/src/python/pyo3/utils.rs b/src/script/src/python/pyo3/utils.rs index 04146a6a6f69..43d917c951e4 100644 --- a/src/script/src/python/pyo3/utils.rs +++ b/src/script/src/python/pyo3/utils.rs @@ -75,9 +75,10 @@ pub fn val_to_py_any(py: Python<'_>, val: Value) -> PyResult { Value::DateTime(val) => val.val().to_object(py), Value::Timestamp(val) => val.value().to_object(py), Value::List(val) => { - let list = val.items().clone().unwrap_or(Default::default()); - let list = list - .into_iter() + let list = val + .items() + .iter() + .cloned() .map(|v| val_to_py_any(py, v)) .collect::>>()?; list.to_object(py) diff --git a/src/script/src/python/rspython/builtins.rs b/src/script/src/python/rspython/builtins.rs index da8c3d6a51a8..20e73d67cf60 100644 --- a/src/script/src/python/rspython/builtins.rs +++ b/src/script/src/python/rspython/builtins.rs @@ -17,14 +17,12 @@ #[cfg(test)] mod test; -use std::sync::Arc; - use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_expr::ColumnarValue as DFColValue; use datafusion_physical_expr::AggregateExpr; use datatypes::arrow::array::ArrayRef; use datatypes::arrow::compute; -use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; +use datatypes::arrow::datatypes::DataType as ArrowDataType; use datatypes::vectors::Helper as HelperVec; use rustpython_vm::builtins::{PyBaseExceptionRef, PyBool, PyFloat, PyInt, PyList, PyStr}; use rustpython_vm::{pymodule, AsObject, PyObjectRef, PyPayload, PyResult, VirtualMachine}; @@ -63,10 +61,6 @@ fn collect_diff_types_string(values: &[ScalarValue], ty: &ArrowDataType) -> Stri .unwrap_or_else(|| "Nothing".to_string()) } -fn new_item_field(data_type: ArrowDataType) -> Field { - Field::new("item", data_type, false) -} - /// try to turn a Python Object into a PyVector or a scalar that can be use for calculate /// /// supported scalar are(leftside is python data type, right side is rust type): @@ -119,8 +113,7 @@ pub fn try_into_columnar_value(obj: PyObjectRef, vm: &VirtualMachine) -> PyResul if ret.is_empty() { // TODO(dennis): empty list, we set type as null. return Ok(DFColValue::Scalar(ScalarValue::List( - None, - Arc::new(new_item_field(ArrowDataType::Null)), + ScalarValue::new_list(&[], &ArrowDataType::Null), ))); } @@ -132,8 +125,7 @@ pub fn try_into_columnar_value(obj: PyObjectRef, vm: &VirtualMachine) -> PyResul ))); } Ok(DFColValue::Scalar(ScalarValue::List( - Some(ret), - Arc::new(new_item_field(ty)), + ScalarValue::new_list(&ret, &ty), ))) } else { Err(vm.new_type_error(format!( @@ -176,9 +168,11 @@ fn scalar_val_try_into_py_obj(val: ScalarValue, vm: &VirtualMachine) -> PyResult ScalarValue::Float64(Some(v)) => Ok(PyFloat::from(v).into_pyobject(vm)), ScalarValue::Int64(Some(v)) => Ok(PyInt::from(v).into_pyobject(vm)), ScalarValue::UInt64(Some(v)) => Ok(PyInt::from(v).into_pyobject(vm)), - ScalarValue::List(Some(col), _) => { - let list = col + ScalarValue::List(list) => { + let list = ScalarValue::convert_array_to_scalar_vec(list.as_ref()) + .map_err(|e| from_df_err(e, vm))? .into_iter() + .flatten() .map(|v| scalar_val_try_into_py_obj(v, vm)) .collect::>()?; let list = vm.ctx.new_list(list); @@ -512,7 +506,11 @@ pub(crate) mod greptime_builtin { /// simple math function, the backing implement is datafusion's `tan` math function #[pyfunction] fn tan(val: PyObjectRef, vm: &VirtualMachine) -> PyResult { - bind_call_unary_math_function!(tan, vm, val); + let args = &[try_into_columnar_value(val, vm).and_then(|x| all_to_f64(x, vm))?]; + datafusion_functions::math::tan() + .invoke(args) + .map_err(|e| from_df_err(e, vm)) + .and_then(|x| try_into_py_obj(x, vm)) } /// simple math function, the backing implement is datafusion's `asin` math function @@ -548,8 +546,10 @@ pub(crate) mod greptime_builtin { #[pyfunction] fn round(val: PyObjectRef, vm: &VirtualMachine) -> PyResult { let value = try_into_columnar_value(val, vm)?; - let array = value.into_array(1); - let result = math_expressions::round(&[array]).map_err(|e| from_df_err(e, vm))?; + let result = value + .into_array(1) + .and_then(|x| math_expressions::round(&[x])) + .map_err(|e| from_df_err(e, vm))?; try_into_py_obj(DFColValue::Array(result), vm) } @@ -563,7 +563,11 @@ pub(crate) mod greptime_builtin { /// simple math function, the backing implement is datafusion's `abs` math function #[pyfunction] fn abs(val: PyObjectRef, vm: &VirtualMachine) -> PyResult { - bind_call_unary_math_function!(abs, vm, val); + let args = &[try_into_columnar_value(val, vm).and_then(|x| all_to_f64(x, vm))?]; + datafusion_functions::math::abs() + .invoke(args) + .map_err(|e| from_df_err(e, vm)) + .and_then(|x| try_into_py_obj(x, vm)) } /// simple math function, the backing implement is datafusion's `signum` math function @@ -673,13 +677,16 @@ pub(crate) mod greptime_builtin { /// effectively equals to `list(vector)` #[pyfunction] fn array_agg(values: PyVectorRef, vm: &VirtualMachine) -> PyResult { - bind_aggr_fn!( - ArrayAgg, - vm, + eval_aggr_fn( + expressions::ArrayAgg::new( + Arc::new(expressions::Column::new("expr0", 0)) as _, + "ArrayAgg", + values.arrow_data_type(), + false, + ), &[values.to_arrow_array()], - values.arrow_data_type(), - expr0 - ); + vm, + ) } /// directly port from datafusion's `avg` function diff --git a/src/script/src/python/rspython/builtins/test.rs b/src/script/src/python/rspython/builtins/test.rs index 0245cd3e9d1e..636c94a64bce 100644 --- a/src/script/src/python/rspython/builtins/test.rs +++ b/src/script/src/python/rspython/builtins/test.rs @@ -18,10 +18,11 @@ use std::io::Read; use std::path::Path; use std::sync::Arc; +use arrow::array::Array; use common_telemetry::{error, info}; use datatypes::arrow::array::{Float64Array, Int64Array}; use datatypes::arrow::compute; -use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; +use datatypes::arrow::datatypes::DataType as ArrowDataType; use datatypes::vectors::{Float64Vector, Int64Vector, VectorRef}; use ron::from_str as from_ron_string; use rustpython_vm::builtins::{PyFloat, PyInt, PyList}; @@ -68,18 +69,15 @@ fn convert_scalar_to_py_obj_and_back() { } else { panic!("Convert errors, expect 1") } - let col = DFColValue::Scalar(ScalarValue::List( - Some(vec![ - ScalarValue::Int64(Some(1)), - ScalarValue::Int64(Some(2)), - ]), - Arc::new(Field::new("item", ArrowDataType::Int64, false)), - )); + let col = DFColValue::Scalar(ScalarValue::List(ScalarValue::new_list( + &[ScalarValue::Int64(Some(1)), ScalarValue::Int64(Some(2))], + &ArrowDataType::Int64, + ))); let to = try_into_py_obj(col, vm).unwrap(); let back = try_into_columnar_value(to, vm).unwrap(); - if let DFColValue::Scalar(ScalarValue::List(Some(list), field)) = back { + if let DFColValue::Scalar(ScalarValue::List(list)) = back { assert_eq!(list.len(), 2); - assert_eq!(*field.data_type(), ArrowDataType::Int64); + assert_eq!(list.data_type(), &ArrowDataType::Int64); } let list: Vec = vec![vm.ctx.new_int(1).into(), vm.ctx.new_int(2).into()]; let nested_list: Vec = diff --git a/src/script/src/python/rspython/utils.rs b/src/script/src/python/rspython/utils.rs index e24b3005c341..ccaf5454fc2b 100644 --- a/src/script/src/python/rspython/utils.rs +++ b/src/script/src/python/rspython/utils.rs @@ -14,12 +14,13 @@ use std::sync::Arc; +use arrow::array::ArrayRef; use datafusion_common::ScalarValue; use datafusion_expr::ColumnarValue as DFColValue; use datatypes::prelude::ScalarVector; use datatypes::value::Value; use datatypes::vectors::{ - BooleanVector, Float64Vector, Helper, Int64Vector, NullVector, StringVector, VectorRef, + BooleanVector, Float64Vector, Helper, Int64Vector, StringVector, VectorRef, }; use rustpython_vm::builtins::{PyBaseExceptionRef, PyBool, PyFloat, PyInt, PyList, PyStr}; use rustpython_vm::object::PyObjectPayload; @@ -134,15 +135,9 @@ pub fn py_obj_to_vec( try_into_columnar_value(obj.clone(), vm).map_err(|e| format_py_error(e, vm))?; match columnar_value { - DFColValue::Scalar(ScalarValue::List(scalars, _datatype)) => match scalars { - Some(scalars) => { - let array = - ScalarValue::iter_to_array(scalars).context(error::DataFusionSnafu)?; - - Helper::try_into_vector(array).context(error::TypeCastSnafu) - } - None => Ok(Arc::new(NullVector::new(0))), - }, + DFColValue::Scalar(ScalarValue::List(array)) => { + Helper::try_into_vector(array as ArrayRef).context(error::TypeCastSnafu) + } _ => unreachable!(), } } else { diff --git a/src/servers/src/mysql/helper.rs b/src/servers/src/mysql/helper.rs index 6109ec706d84..c5d509016da4 100644 --- a/src/servers/src/mysql/helper.rs +++ b/src/servers/src/mysql/helper.rs @@ -179,6 +179,7 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result(idx, &client_type)?; match server_type { ConcreteDataType::Timestamp(unit) => match *unit { - TimestampType::Second(_) => { - ScalarValue::TimestampSecond(data.map(|ts| ts.timestamp()), None) - } + TimestampType::Second(_) => ScalarValue::TimestampSecond( + data.map(|ts| ts.and_utc().timestamp()), + None, + ), TimestampType::Millisecond(_) => ScalarValue::TimestampMillisecond( - data.map(|ts| ts.timestamp_millis()), + data.map(|ts| ts.and_utc().timestamp_millis()), None, ), TimestampType::Microsecond(_) => ScalarValue::TimestampMicrosecond( - data.map(|ts| ts.timestamp_micros()), + data.map(|ts| ts.and_utc().timestamp_micros()), None, ), TimestampType::Nanosecond(_) => ScalarValue::TimestampNanosecond( - data.map(|ts| ts.timestamp_micros()), + data.map(|ts| ts.and_utc().timestamp_micros()), None, ), }, ConcreteDataType::DateTime(_) => { - ScalarValue::Date64(data.map(|d| d.timestamp_millis())) + ScalarValue::Date64(data.map(|d| d.and_utc().timestamp_millis())) } _ => { return Err(invalid_parameter_error( @@ -790,10 +791,7 @@ mod test { } let err = encode_value( - &Value::List(ListValue::new( - Some(Box::default()), - ConcreteDataType::int16_datatype(), - )), + &Value::List(ListValue::new(vec![], ConcreteDataType::int16_datatype())), &mut builder, ) .unwrap_err(); diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 7553d9791225..8d4cb2540186 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -103,11 +103,11 @@ pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result { } // Case sensitive regexp match MatcherType::Re => { - conditions.push(regexp_match(vec![col(name), lit(value)]).is_not_null()); + conditions.push(regexp_match(col(name), lit(value)).is_not_null()); } // Case sensitive regexp not match MatcherType::Nre => { - conditions.push(regexp_match(vec![col(name), lit(value)]).is_null()); + conditions.push(regexp_match(col(name), lit(value)).is_null()); } } } diff --git a/src/sql/src/parser.rs b/src/sql/src/parser.rs index 9010106d850a..0b35f3e86a4f 100644 --- a/src/sql/src/parser.rs +++ b/src/sql/src/parser.rs @@ -83,14 +83,14 @@ impl<'a> ParserContext<'a> { } pub(crate) fn intern_parse_table_name(&mut self) -> Result { - let raw_table_name = self - .parser - .parse_object_name() - .context(error::UnexpectedSnafu { - sql: self.sql, - expected: "a table name", - actual: self.parser.peek_token().to_string(), - })?; + let raw_table_name = + self.parser + .parse_object_name(false) + .context(error::UnexpectedSnafu { + sql: self.sql, + expected: "a table name", + actual: self.parser.peek_token().to_string(), + })?; Ok(Self::canonicalize_object_name(raw_table_name)) } @@ -100,7 +100,7 @@ impl<'a> ParserContext<'a> { .try_with_sql(sql) .context(SyntaxSnafu)?; - let function_name = parser.parse_identifier().context(SyntaxSnafu)?; + let function_name = parser.parse_identifier(false).context(SyntaxSnafu)?; parser .parse_function(ObjectName(vec![function_name])) .context(SyntaxSnafu) diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index 483c432636c6..05695cbbe69f 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -33,7 +33,7 @@ impl<'a> ParserContext<'a> { let parser = &mut self.parser; parser.expect_keywords(&[Keyword::ALTER, Keyword::TABLE])?; - let raw_table_name = parser.parse_object_name()?; + let raw_table_name = parser.parse_object_name(false)?; let table_name = Self::canonicalize_object_name(raw_table_name); let alter_operation = if parser.parse_keyword(Keyword::ADD) { @@ -48,7 +48,7 @@ impl<'a> ParserContext<'a> { } else if let Token::Word(word) = parser.peek_token().token { if word.value.to_ascii_uppercase() == "AFTER" { let _ = parser.next_token(); - let name = Self::canonicalize_identifier(parser.parse_identifier()?); + let name = Self::canonicalize_identifier(parser.parse_identifier(false)?); Some(AddColumnLocation::After { column_name: name.value, }) @@ -65,7 +65,7 @@ impl<'a> ParserContext<'a> { } } else if parser.parse_keyword(Keyword::DROP) { if parser.parse_keyword(Keyword::COLUMN) { - let name = Self::canonicalize_identifier(self.parser.parse_identifier()?); + let name = Self::canonicalize_identifier(self.parser.parse_identifier(false)?); AlterTableOperation::DropColumn { name } } else { return Err(ParserError::ParserError(format!( @@ -74,7 +74,7 @@ impl<'a> ParserContext<'a> { ))); } } else if parser.parse_keyword(Keyword::RENAME) { - let new_table_name_obj_raw = parser.parse_object_name()?; + let new_table_name_obj_raw = parser.parse_object_name(false)?; let new_table_name_obj = Self::canonicalize_object_name(new_table_name_obj_raw); let new_table_name = match &new_table_name_obj.0[..] { [table] => table.value.clone(), @@ -128,7 +128,7 @@ mod tests { location, } => { assert_eq!("tagk_i", column_def.name.value); - assert_eq!(DataType::String, column_def.data_type); + assert_eq!(DataType::String(None), column_def.data_type); assert!(column_def .options .iter() @@ -164,7 +164,7 @@ mod tests { location, } => { assert_eq!("tagk_i", column_def.name.value); - assert_eq!(DataType::String, column_def.data_type); + assert_eq!(DataType::String(None), column_def.data_type); assert!(column_def .options .iter() @@ -200,7 +200,7 @@ mod tests { location, } => { assert_eq!("tagk_i", column_def.name.value); - assert_eq!(DataType::String, column_def.data_type); + assert_eq!(DataType::String(None), column_def.data_type); assert!(column_def .options .iter() diff --git a/src/sql/src/parsers/copy_parser.rs b/src/sql/src/parsers/copy_parser.rs index bd365e51ef3a..952449f8ae49 100644 --- a/src/sql/src/parsers/copy_parser.rs +++ b/src/sql/src/parsers/copy_parser.rs @@ -49,7 +49,7 @@ impl<'a> ParserContext<'a> { fn parser_copy_database(&mut self) -> Result { let database_name = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a database name", @@ -84,7 +84,7 @@ impl<'a> ParserContext<'a> { fn parse_copy_table(&mut self) -> Result { let raw_table_name = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a table name", diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index ec61a20d5378..c3c922efde1b 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -111,14 +111,14 @@ impl<'a> ParserContext<'a> { self.parser .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); - let database_name = self - .parser - .parse_object_name() - .context(error::UnexpectedSnafu { - sql: self.sql, - expected: "a database name", - actual: self.peek_token_as_string(), - })?; + let database_name = + self.parser + .parse_object_name(false) + .context(error::UnexpectedSnafu { + sql: self.sql, + expected: "a database name", + actual: self.peek_token_as_string(), + })?; let database_name = Self::canonicalize_object_name(database_name); Ok(Statement::CreateDatabase(CreateDatabase { name: database_name, @@ -319,6 +319,7 @@ impl<'a> ParserContext<'a> { quote_style: None, }], is_primary: false, + characteristics: None, }; constraints.push(constraint); } @@ -367,7 +368,7 @@ impl<'a> ParserContext<'a> { pub fn parse_column_def(&mut self) -> std::result::Result { let parser = &mut self.parser; - let name = parser.parse_identifier()?; + let name = parser.parse_identifier(false)?; if name.quote_style.is_none() && // "ALL_KEYWORDS" are sorted. ALL_KEYWORDS.binary_search(&name.value.to_uppercase().as_str()).is_ok() @@ -380,14 +381,14 @@ impl<'a> ParserContext<'a> { let data_type = parser.parse_data_type()?; let collation = if parser.parse_keyword(Keyword::COLLATE) { - Some(parser.parse_object_name()?) + Some(parser.parse_object_name(false)?) } else { None }; let mut options = vec![]; loop { if parser.parse_keyword(Keyword::CONSTRAINT) { - let name = Some(parser.parse_identifier()?); + let name = Some(parser.parse_identifier(false)?); if let Some(option) = Self::parse_optional_column_option(parser)? { options.push(ColumnOptionDef { name, option }); } else { @@ -415,7 +416,7 @@ impl<'a> ParserContext<'a> { ) -> std::result::Result, ParserError> { if parser.parse_keywords(&[Keyword::CHARACTER, Keyword::SET]) { Ok(Some(ColumnOption::CharacterSet( - parser.parse_object_name()?, + parser.parse_object_name(false)?, ))) } else if parser.parse_keywords(&[Keyword::NOT, Keyword::NULL]) { Ok(Some(ColumnOption::NotNull)) @@ -432,9 +433,15 @@ impl<'a> ParserContext<'a> { } else if parser.parse_keyword(Keyword::DEFAULT) { Ok(Some(ColumnOption::Default(parser.parse_expr()?))) } else if parser.parse_keywords(&[Keyword::PRIMARY, Keyword::KEY]) { - Ok(Some(ColumnOption::Unique { is_primary: true })) + Ok(Some(ColumnOption::Unique { + is_primary: true, + characteristics: None, + })) } else if parser.parse_keyword(Keyword::UNIQUE) { - Ok(Some(ColumnOption::Unique { is_primary: false })) + Ok(Some(ColumnOption::Unique { + is_primary: false, + characteristics: None, + })) } else if parser.parse_keywords(&[Keyword::TIME, Keyword::INDEX]) { // Use a DialectSpecific option for time index Ok(Some(ColumnOption::DialectSpecific(vec![ @@ -456,7 +463,10 @@ impl<'a> ParserContext<'a> { fn parse_optional_table_constraint(&mut self) -> Result> { let name = if self.parser.parse_keyword(Keyword::CONSTRAINT) { - let raw_name = self.parser.parse_identifier().context(error::SyntaxSnafu)?; + let raw_name = self + .parser + .parse_identifier(false) + .context(error::SyntaxSnafu)?; Some(Self::canonicalize_identifier(raw_name)) } else { None @@ -485,6 +495,7 @@ impl<'a> ParserContext<'a> { name, columns, is_primary: true, + characteristics: None, })) } TokenWithLocation { @@ -524,6 +535,7 @@ impl<'a> ParserContext<'a> { }), columns, is_primary: false, + characteristics: None, })) } unexpected => { @@ -568,6 +580,7 @@ fn validate_time_index(columns: &[ColumnDef], constraints: &[TableConstraint]) - name: Some(ident), columns, is_primary: false, + .. } = c { if ident.value == TIME_INDEX { @@ -1108,6 +1121,7 @@ ENGINE=mito"; name, columns, is_primary, + .. } => { assert_eq!(name.unwrap().to_string(), "__time_index"); assert_eq!(columns.len(), 1); @@ -1314,6 +1328,7 @@ ENGINE=mito"; name, columns, is_primary, + .. } => { assert_eq!(name.unwrap().to_string(), "__time_index"); assert_eq!(columns.len(), 1); diff --git a/src/sql/src/parsers/describe_parser.rs b/src/sql/src/parsers/describe_parser.rs index 7ba1ffc00366..81bfd36e9df7 100644 --- a/src/sql/src/parsers/describe_parser.rs +++ b/src/sql/src/parsers/describe_parser.rs @@ -32,7 +32,7 @@ impl<'a> ParserContext<'a> { fn parse_describe_table(&mut self) -> Result { let raw_table_idents = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a table name", diff --git a/src/sql/src/parsers/drop_parser.rs b/src/sql/src/parsers/drop_parser.rs index d5d872ee1689..2b327c4c41ce 100644 --- a/src/sql/src/parsers/drop_parser.rs +++ b/src/sql/src/parsers/drop_parser.rs @@ -41,7 +41,7 @@ impl<'a> ParserContext<'a> { let if_exists = self.parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]); let raw_table_ident = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a table name", @@ -64,7 +64,7 @@ impl<'a> ParserContext<'a> { let if_exists = self.parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]); let database_name = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a database name", diff --git a/src/sql/src/parsers/explain_parser.rs b/src/sql/src/parsers/explain_parser.rs index 906bd792eceb..57349220fd4b 100644 --- a/src/sql/src/parsers/explain_parser.rs +++ b/src/sql/src/parsers/explain_parser.rs @@ -13,6 +13,7 @@ // limitations under the License. use snafu::ResultExt; +use sqlparser::ast::DescribeAlias; use crate::error::{self, Result}; use crate::parser::ParserContext; @@ -22,14 +23,14 @@ use crate::statements::statement::Statement; /// EXPLAIN statement parser implementation impl<'a> ParserContext<'a> { pub(crate) fn parse_explain(&mut self) -> Result { - let explain_statement = - self.parser - .parse_explain(false) - .with_context(|_| error::UnexpectedSnafu { - sql: self.sql, - expected: "a query statement", - actual: self.peek_token_as_string(), - })?; + let explain_statement = self + .parser + .parse_explain(DescribeAlias::Explain) + .with_context(|_| error::UnexpectedSnafu { + sql: self.sql, + expected: "a query statement", + actual: self.peek_token_as_string(), + })?; Ok(Statement::Explain(Explain::try_from(explain_statement)?)) } @@ -80,6 +81,7 @@ mod tests { having: None, qualify: None, named_window: vec![], + value_table_mode: None, }; let sp_statement = SpStatement::Query(Box::new(SpQuery { @@ -87,13 +89,15 @@ mod tests { body: Box::new(sqlparser::ast::SetExpr::Select(Box::new(select))), order_by: vec![], limit: None, + limit_by: vec![], offset: None, fetch: None, locks: vec![], + for_clause: None, })); let explain = Explain::try_from(SpStatement::Explain { - describe_alias: false, + describe_alias: DescribeAlias::Explain, analyze: false, verbose: false, statement: Box::new(sp_statement), diff --git a/src/sql/src/parsers/show_parser.rs b/src/sql/src/parsers/show_parser.rs index a71397095eeb..4a656bd33c4c 100644 --- a/src/sql/src/parsers/show_parser.rs +++ b/src/sql/src/parsers/show_parser.rs @@ -48,7 +48,7 @@ impl<'a> ParserContext<'a> { } else if self.consume_token("VARIABLES") { let variable = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a variable name", @@ -64,7 +64,7 @@ impl<'a> ParserContext<'a> { fn parse_show_create_table(&mut self) -> Result { let raw_table_name = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a table name", @@ -94,7 +94,7 @@ impl<'a> ParserContext<'a> { Token::Word(w) => match w.keyword { Keyword::IN | Keyword::FROM => { let _ = self.parser.next_token(); - let db_name = self.parser.parse_object_name().with_context(|_| { + let db_name = self.parser.parse_object_name(false).with_context(|_| { error::UnexpectedSnafu { sql: self.sql, expected: "a database name", @@ -123,7 +123,7 @@ impl<'a> ParserContext<'a> { Token::Word(w) => match w.keyword { Keyword::LIKE => { let _ = self.parser.next_token(); - ShowKind::Like(self.parser.parse_identifier().with_context(|_| { + ShowKind::Like(self.parser.parse_identifier(false).with_context(|_| { error::UnexpectedSnafu { sql: self.sql, expected: "LIKE", @@ -162,7 +162,7 @@ impl<'a> ParserContext<'a> { } Token::Word(w) => match w.keyword { Keyword::LIKE => Ok(Statement::ShowDatabases(ShowDatabases::new( - ShowKind::Like(self.parser.parse_identifier().with_context(|_| { + ShowKind::Like(self.parser.parse_identifier(false).with_context(|_| { error::UnexpectedSnafu { sql: self.sql, expected: "LIKE", diff --git a/src/sql/src/parsers/truncate_parser.rs b/src/sql/src/parsers/truncate_parser.rs index d4ac0c4d1427..efd5e3a56932 100644 --- a/src/sql/src/parsers/truncate_parser.rs +++ b/src/sql/src/parsers/truncate_parser.rs @@ -28,7 +28,7 @@ impl<'a> ParserContext<'a> { let raw_table_ident = self.parser - .parse_object_name() + .parse_object_name(false) .with_context(|_| error::UnexpectedSnafu { sql: self.sql, expected: "a table name", diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index d51378c0982f..003efe6f0719 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -345,7 +345,7 @@ pub fn has_primary_key_option(column_def: &ColumnDef) -> bool { .options .iter() .any(|options| match options.option { - ColumnOption::Unique { is_primary } => is_primary, + ColumnOption::Unique { is_primary, .. } => is_primary, _ => false, }) } @@ -414,10 +414,15 @@ pub fn sql_column_def_to_grpc_column_def( .context(ConvertToGrpcDataTypeSnafu)? .to_parts(); - let is_primary_key = col - .options - .iter() - .any(|o| matches!(o.option, ColumnOption::Unique { is_primary: true })); + let is_primary_key = col.options.iter().any(|o| { + matches!( + o.option, + ColumnOption::Unique { + is_primary: true, + .. + } + ) + }); let semantic_type = if is_primary_key { SemanticType::Tag @@ -453,7 +458,7 @@ pub fn sql_data_type_to_concrete_data_type(data_type: &SqlDataType) -> Result Ok(ConcreteDataType::string_datatype()), + | SqlDataType::String(_) => Ok(ConcreteDataType::string_datatype()), SqlDataType::Float(_) => Ok(ConcreteDataType::float32_datatype()), SqlDataType::Double => Ok(ConcreteDataType::float64_datatype()), SqlDataType::Boolean => Ok(ConcreteDataType::boolean_datatype()), @@ -502,7 +507,7 @@ pub fn concrete_data_type_to_sql_data_type(data_type: &ConcreteDataType) -> Resu ConcreteDataType::UInt16(_) => Ok(SqlDataType::UnsignedSmallInt(None)), ConcreteDataType::Int8(_) => Ok(SqlDataType::TinyInt(None)), ConcreteDataType::UInt8(_) => Ok(SqlDataType::UnsignedTinyInt(None)), - ConcreteDataType::String(_) => Ok(SqlDataType::String), + ConcreteDataType::String(_) => Ok(SqlDataType::String(None)), ConcreteDataType::Float32(_) => Ok(SqlDataType::Float(None)), ConcreteDataType::Float64(_) => Ok(SqlDataType::Double), ConcreteDataType::Boolean(_) => Ok(SqlDataType::Boolean), @@ -588,7 +593,10 @@ mod tests { ConcreteDataType::string_datatype(), ); check_type(SqlDataType::Text, ConcreteDataType::string_datatype()); - check_type(SqlDataType::String, ConcreteDataType::string_datatype()); + check_type( + SqlDataType::String(None), + ConcreteDataType::string_datatype(), + ); check_type( SqlDataType::Float(None), ConcreteDataType::float32_datatype(), @@ -966,7 +974,10 @@ mod tests { collation: None, options: vec![ColumnOptionDef { name: None, - option: ColumnOption::Unique { is_primary: true }, + option: ColumnOption::Unique { + is_primary: true, + characteristics: None, + }, }], }; @@ -1044,7 +1055,10 @@ mod tests { collation: None, options: vec![ColumnOptionDef { name: None, - option: ColumnOption::Unique { is_primary: true }, + option: ColumnOption::Unique { + is_primary: true, + characteristics: None, + }, }], }; assert!(has_primary_key_option(&column_def)); @@ -1081,7 +1095,7 @@ mod tests { let column_def = ColumnDef { name: "col2".into(), - data_type: SqlDataType::String, + data_type: SqlDataType::String(None), collation: None, options: vec![ ColumnOptionDef { diff --git a/src/sql/src/statements/insert.rs b/src/sql/src/statements/insert.rs index 81ef55734d85..4eae7f1e1874 100644 --- a/src/sql/src/statements/insert.rs +++ b/src/sql/src/statements/insert.rs @@ -55,10 +55,10 @@ impl Insert { match &self.inner { Statement::Insert { source: - box Query { + Some(box Query { body: box SetExpr::Values(Values { rows, .. }), .. - }, + }), .. } => sql_exprs_to_values(rows), _ => unreachable!(), @@ -71,10 +71,10 @@ impl Insert { match &self.inner { Statement::Insert { source: - box Query { + Some(box Query { body: box SetExpr::Values(Values { rows, .. }), .. - }, + }), .. } => rows.iter().all(|es| { es.iter().all(|expr| match expr { @@ -100,7 +100,8 @@ impl Insert { pub fn query_body(&self) -> Result> { Ok(match &self.inner { Statement::Insert { - source: box query, .. + source: Some(box query), + .. } => Some(query.clone().try_into()?), _ => None, }) diff --git a/src/sql/src/statements/transform/type_alias.rs b/src/sql/src/statements/transform/type_alias.rs index 0db61538e069..865327b0f4eb 100644 --- a/src/sql/src/statements/transform/type_alias.rs +++ b/src/sql/src/statements/transform/type_alias.rs @@ -57,28 +57,56 @@ impl TransformRule for TypeAliasTransformRule { } fn visit_expr(&self, expr: &mut Expr) -> ControlFlow<()> { + fn cast_expr_to_arrow_cast_func(expr: Expr, cast_type: String) -> Function { + Function { + name: ObjectName(vec![Ident::new("arrow_cast")]), + args: vec![ + FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)), + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value( + Value::SingleQuotedString(cast_type), + ))), + ], + filter: None, + null_treatment: None, + over: None, + distinct: false, + special: false, + order_by: vec![], + } + } + match expr { + // In new sqlparser, the "INT64" is no longer parsed to custom datatype. + // The new "Int64" is not recognizable by Datafusion, cannot directly "CAST" to it. + // We have to replace the expr to "arrow_cast" function call here. + // Same for "FLOAT64". + Expr::Cast { + expr: cast_expr, + data_type, + .. + } if matches!(data_type, DataType::Int64 | DataType::Float64) => { + if let Some(new_type) = get_data_type_by_alias_name(&data_type.to_string()) { + if let Ok(new_type) = sql_data_type_to_concrete_data_type(&new_type) { + *expr = Expr::Function(cast_expr_to_arrow_cast_func( + (**cast_expr).clone(), + new_type.as_arrow_type().to_string(), + )); + } + } + } + // Type alias Expr::Cast { data_type: DataType::Custom(name, tokens), expr: cast_expr, + .. } if name.0.len() == 1 && tokens.is_empty() => { if let Some(new_type) = get_data_type_by_alias_name(name.0[0].value.as_str()) { - if let Ok(concrete_type) = sql_data_type_to_concrete_data_type(&new_type) { - let new_type = concrete_type.as_arrow_type(); - *expr = Expr::Function(Function { - name: ObjectName(vec![Ident::new("arrow_cast")]), - args: vec![ - FunctionArg::Unnamed(FunctionArgExpr::Expr((**cast_expr).clone())), - FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value( - Value::SingleQuotedString(new_type.to_string()), - ))), - ], - over: None, - distinct: false, - special: false, - order_by: vec![], - }); + if let Ok(new_type) = sql_data_type_to_concrete_data_type(&new_type) { + *expr = Expr::Function(cast_expr_to_arrow_cast_func( + (**cast_expr).clone(), + new_type.as_arrow_type().to_string(), + )); } } } @@ -88,24 +116,16 @@ impl TransformRule for TypeAliasTransformRule { Expr::Cast { data_type: DataType::Timestamp(precision, zone), expr: cast_expr, + .. } => { if let Ok(concrete_type) = sql_data_type_to_concrete_data_type(&DataType::Timestamp(*precision, *zone)) { let new_type = concrete_type.as_arrow_type(); - *expr = Expr::Function(Function { - name: ObjectName(vec![Ident::new("arrow_cast")]), - args: vec![ - FunctionArg::Unnamed(FunctionArgExpr::Expr((**cast_expr).clone())), - FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value( - Value::SingleQuotedString(new_type.to_string()), - ))), - ], - over: None, - distinct: false, - special: false, - order_by: vec![], - }); + *expr = Expr::Function(cast_expr_to_arrow_cast_func( + (**cast_expr).clone(), + new_type.to_string(), + )); } } diff --git a/src/sql/src/util.rs b/src/sql/src/util.rs index 9a931e02cb9e..0abf92beb2de 100644 --- a/src/sql/src/util.rs +++ b/src/sql/src/util.rs @@ -17,7 +17,7 @@ use std::fmt::{Display, Formatter}; use std::sync::LazyLock; use regex::Regex; -use sqlparser::ast::{ObjectName, SqlOption, Value}; +use sqlparser::ast::{Expr, ObjectName, SqlOption, Value}; static SQL_SECRET_PATTERNS: LazyLock> = LazyLock::new(|| { vec![ @@ -47,9 +47,11 @@ pub fn format_raw_object_name(name: &ObjectName) -> String { format!("{}", Inner { name }) } -pub fn parse_option_string(value: Value) -> Option { +pub fn parse_option_string(value: Expr) -> Option { match value { - Value::SingleQuotedString(v) | Value::DoubleQuotedString(v) => Some(v), + Expr::Value(Value::SingleQuotedString(v)) | Expr::Value(Value::DoubleQuotedString(v)) => { + Some(v) + } _ => None, } } @@ -60,7 +62,8 @@ pub fn to_lowercase_options_map(opts: &[SqlOption]) -> HashMap { let mut map = HashMap::with_capacity(opts.len()); for SqlOption { name, value } in opts { let value_str = match value { - Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => s.clone(), + Expr::Value(Value::SingleQuotedString(s)) + | Expr::Value(Value::DoubleQuotedString(s)) => s.clone(), _ => value.to_string(), }; let _ = map.insert(name.value.to_lowercase().clone(), value_str); diff --git a/src/table/src/predicate.rs b/src/table/src/predicate.rs index bd5d6d2d1818..f7047926265a 100644 --- a/src/table/src/predicate.rs +++ b/src/table/src/predicate.rs @@ -89,8 +89,7 @@ impl Predicate { .exprs .iter() .filter_map(|expr| { - create_physical_expr(expr.df_expr(), df_schema.as_ref(), schema, execution_props) - .ok() + create_physical_expr(expr.df_expr(), df_schema.as_ref(), execution_props).ok() }) .collect::>()) } @@ -284,7 +283,11 @@ impl<'a> TimeRangePredicateBuilder<'a> { | Operator::BitwiseShiftLeft | Operator::StringConcat | Operator::ArrowAt - | Operator::AtArrow => None, + | Operator::AtArrow + | Operator::LikeMatch + | Operator::ILikeMatch + | Operator::NotLikeMatch + | Operator::NotILikeMatch => None, } } diff --git a/src/table/src/predicate/stats.rs b/src/table/src/predicate/stats.rs index 4d707a82e9cf..e1cc48c320b9 100644 --- a/src/table/src/predicate/stats.rs +++ b/src/table/src/predicate/stats.rs @@ -12,13 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; use std::sync::Arc; use datafusion::parquet::file::metadata::RowGroupMetaData; use datafusion::parquet::file::statistics::Statistics as ParquetStats; use datafusion::physical_optimizer::pruning::PruningStatistics; use datafusion_common::{Column, ScalarValue}; -use datatypes::arrow::array::{ArrayRef, UInt64Array}; +use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use datatypes::arrow::datatypes::DataType; use paste::paste; @@ -115,4 +116,14 @@ impl<'a> PruningStatistics for RowGroupPruningStatistics<'a> { } Some(Arc::new(UInt64Array::from(values))) } + + fn row_counts(&self, _column: &Column) -> Option { + // TODO(LFC): Impl it. + None + } + + fn contained(&self, _column: &Column, _values: &HashSet) -> Option { + // TODO(LFC): Impl it. + None + } } diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 0d5b76913646..1414575846a1 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -28,7 +28,8 @@ use common_telemetry::tracing::Span; use common_telemetry::tracing_context::TracingContext; use datafusion::execution::context::TaskContext; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_physical_expr::PhysicalSortExpr; +use datafusion::physical_plan::{ExecutionMode, PlanProperties}; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalSortExpr}; use datatypes::schema::SchemaRef; use futures::{Stream, StreamExt}; use snafu::OptionExt; @@ -41,6 +42,7 @@ pub struct StreamScanAdapter { schema: SchemaRef, output_ordering: Option>, metric: ExecutionPlanMetricsSet, + properties: PlanProperties, } impl Debug for StreamScanAdapter { @@ -55,12 +57,17 @@ impl Debug for StreamScanAdapter { impl StreamScanAdapter { pub fn new(stream: SendableRecordBatchStream) -> Self { let schema = stream.schema(); - + let properties = PlanProperties::new( + EquivalenceProperties::new(schema.arrow_schema().clone()), + Partitioning::UnknownPartitioning(1), + ExecutionMode::Bounded, + ); Self { stream: Mutex::new(Some(stream)), schema, output_ordering: None, metric: ExecutionPlanMetricsSet::new(), + properties, } } @@ -79,12 +86,8 @@ impl PhysicalPlan for StreamScanAdapter { self.schema.clone() } - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) - } - - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - self.output_ordering.as_deref() + fn properties(&self) -> &PlanProperties { + &self.properties } fn children(&self) -> Vec { diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 7cff590c39eb..245e05ed5fd0 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -143,7 +143,7 @@ pub async fn test_mysql_crud(store_type: StorageType) { .unwrap(); for i in 0..10 { let dt: DateTime = DateTime::from_naive_utc_and_offset( - NaiveDateTime::from_timestamp_opt(60, i).unwrap(), + chrono::DateTime::from_timestamp(60, i).unwrap().naive_utc(), Utc, ); let d = NaiveDate::from_yo_opt(2015, 100).unwrap(); @@ -175,7 +175,9 @@ pub async fn test_mysql_crud(store_type: StorageType) { let expected_d = NaiveDate::from_yo_opt(2015, 100).unwrap(); assert_eq!(expected_d, d); let expected_dt: DateTime = DateTime::from_naive_utc_and_offset( - NaiveDateTime::from_timestamp_opt(60, i as u32).unwrap(), + chrono::DateTime::from_timestamp(60, i as u32) + .unwrap() + .naive_utc(), Utc, ); assert_eq!( @@ -356,7 +358,7 @@ pub async fn test_postgres_crud(store_type: StorageType) { for i in 0..10 { let d = NaiveDate::from_yo_opt(2015, 100).unwrap(); - let dt = d.and_hms_opt(0, 0, 0).unwrap().timestamp_millis(); + let dt = d.and_hms_opt(0, 0, 0).unwrap().and_utc().timestamp_millis(); sqlx::query("insert into demo values($1, $2, $3, $4)") .bind(i)