rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-23 21:55:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat!: upgrade DataFusion to 55 (#8555)

Browse Source 
* feat!: upgrade DataFusion dependencies to 55

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor: migrate DataFusion 55 APIs

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: preserve table function planning behavior

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: preserve PostgreSQL query compatibility

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: preserve distributed execution plan behavior

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: cover DataFusion 55 behavior regressions

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: update DataFusion 55 SQLness expectations

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: complete DataFusion 55 test API migration

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: address DataFusion 55 CI regressions

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: address remaining DataFusion 55 regressions

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: adapt latest base code to DataFusion 55

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: normalize environment-specific DataFusion 55 plans

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: align final DataFusion 55 expectations

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: isolate DataFusion 55 regression cases

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: preserve empty result schema in timestamp widening

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: preserve JSON source column order

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* chore: use released DataFusion 55 integrations

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: adapt latest execution plan mock to DataFusion 55

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: pin DataFusion recursive schema and date repairs

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(promql): align dictionary temporality match keys

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: retain Greptime DataFusion fork behaviors on version 55

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: restore ordinary function error expectations

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: refresh distributed count compatibility plan

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(query): adapt last-row cast hint to DataFusion 55

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: refresh instant last-row empty results for Arrow 59

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* style: simplify DataFusion expression visitor imports

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: restore sorting and PostgreSQL column-order assertions

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(function): restore primitive numeric coercion signatures

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor(function): share geo integer signature types

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: cover timestamp widening overflow boundaries

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: fix decimal coercion regression imports

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(function): preserve scalar count_hash NULL state semantics

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: simplify decimal clamp case type inference

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: retain historical count_hash wrapper result

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: restore timestamp widening equality and IN pruning

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: carry upstream aggregate dynamic filter correctness fix

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: carry upstream null and predicate simplification fixes

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: restore baseline JSON ordering expectations

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: restore histogram JSON ordering expectations

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: refresh empty PromQL range result schemas

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: align native timestamp plan with DF55 decimal display

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: refresh native timestamp SQLness results for DF55

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: regenerate NULL sample empty result headers for DF55

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: use DF55 child replacement API in timestamp regressions

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: expose pushed scan dynamic filters to DF55 producers

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: encode string-backed PostgreSQL OID aliases in binary results

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: verify REGPROC binary and text over PostgreSQL protocol

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: register real PostgreSQL catalogs in server fixtures

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: complete DF55 expression inventories for custom query plans

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: correct RangeSelect expression fixture and column identities

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* ci: wait for Kafka WAL helper deployment rollout

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test: update custom storage empty result headers

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: require exact row counts in scan statistics

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: suppress deprecated partition_statistics warning in test

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

---------

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
Co-authored-by: Ning Sun <sunng@protonmail.com>
This commit is contained in:
discord9
2026-09-15 11:42:38 +00:00
committed by GitHub
co-authored by Ning Sun
parent c65e8d0962
commit 94d7e2c7fc
284 changed files with 5948 additions and 2687 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/integration.yml
+3 -4
View File
@@ -781,10 +781,9 @@ jobs:
kubectl -n kafka-cluster apply \
-f .github/actions/setup-greptimedb-cluster/kafka-wal-helper.yaml
kubectl -n kafka-cluster wait \
--for=condition=Ready \
pod -l app=kafka-wal-helper \
--timeout=120s
kubectl rollout status deployment/kafka-wal-helper \
--timeout=120s \
-n kafka-cluster
- name: Print etcd info
shell: bash
run: kubectl get all --show-labels -n etcd-cluster
Cargo.lock
Generated
+458 -414
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+47 -44
View File
@@ -100,13 +100,13 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tokio_unstable)'] }
# See for more detaiils: https://github.com/rust-lang/cargo/issues/11329
ahash = { version = "0.8", features = ["compile-time-rng"] }
aquamarine = "0.6"
arrow = { version = "58.3", features = ["prettyprint"] }
arrow-array = { version = "58.3", default-features = false, features = ["chrono-tz"] }
arrow-buffer = "58.3"
arrow-cast = "58.3"
arrow-flight = "58.3"
arrow-ipc = { version = "58.3", default-features = false, features = ["lz4", "zstd"] }
arrow-schema = { version = "58.3", features = ["serde"] }
arrow = { version = "=59.2.0", features = ["prettyprint"] }
arrow-array = { version = "=59.2.0", default-features = false, features = ["chrono-tz"] }
arrow-buffer = "=59.2.0"
arrow-cast = "=59.2.0"
arrow-flight = "=59.2.0"
arrow-ipc = { version = "=59.2.0", default-features = false, features = ["lz4", "zstd"] }
arrow-schema = { version = "=59.2.0", features = ["serde"] }
async-stream = "0.3"
async-trait = "0.1"
# Remember to update axum-extra, axum-macros when updating axum
@@ -128,22 +128,22 @@ const_format = "0.2"
criterion = "0.7"
crossbeam-utils = "0.8"
dashmap = "6.1"
datafusion = "=53.1.0"
datafusion-common = "=53.1.0"
datafusion-datasource = "=53.1.0"
datafusion-expr = "=53.1.0"
datafusion-expr-common = "=53.1.0"
datafusion-functions = "=53.1.0"
datafusion-functions-aggregate-common = "=53.1.0"
datafusion-functions-window-common = "=53.1.0"
datafusion-optimizer = "=53.1.0"
datafusion-orc = { git = "https://github.com/datafusion-contrib/datafusion-orc.git", rev = "6c07fa282dc8d62db2aa4ded06ab55485efc811a" }
datafusion-pg-catalog = "0.17.3"
datafusion-physical-expr = "=53.1.0"
datafusion-physical-plan = "=53.1.0"
datafusion-proto = "=53.1.0"
datafusion-sql = "=53.1.0"
datafusion-substrait = "=53.1.0"
datafusion = "=55.0.0"
datafusion-common = "=55.0.0"
datafusion-datasource = "=55.0.0"
datafusion-expr = "=55.0.0"
datafusion-expr-common = "=55.0.0"
datafusion-functions = "=55.0.0"
datafusion-functions-aggregate-common = "=55.0.0"
datafusion-functions-window-common = "=55.0.0"
datafusion-optimizer = "=55.0.0"
datafusion-orc = "0.10.0"
datafusion-pg-catalog = "0.18.3"
datafusion-physical-expr = "=55.0.0"
datafusion-physical-plan = "=55.0.0"
datafusion-proto = "=55.0.0"
datafusion-sql = "=55.0.0"
datafusion-substrait = "=55.0.0"
datafusion_object_store = { package = "object_store", version = "0.13.2" }
deadpool = "0.12"
deadpool-postgres = "0.14"
@@ -195,10 +195,10 @@ otel-arrow-rust = { git = "https://github.com/GreptimeTeam/otel-arrow", rev = "5
"server",
] }
parking_lot = "0.12"
parquet = { version = "58.3", default-features = false, features = ["arrow", "async", "object_store"] }
parquet-variant = "58.3"
parquet-variant-compute = "58.3"
parquet-variant-json = "58.3"
parquet = { version = "=59.2.0", default-features = false, features = ["arrow", "async", "object_store"] }
parquet-variant = "=59.2.0"
parquet-variant-compute = "=59.2.0"
parquet-variant-json = "=59.2.0"
paste = "1.0"
pin-project = "1.0"
pretty_assertions = "1.4.0"
@@ -239,7 +239,7 @@ simd-json = "0.15"
similar-asserts = "1.6.0"
smallvec = { version = "1", features = ["serde"] }
snafu = "0.8"
sqlparser = { version = "0.61.0", default-features = false, features = ["std", "visitor", "serde"] }
sqlparser = { version = "0.62.0", default-features = false, features = ["std", "visitor", "serde"] }
sqlx = { version = "0.8", default-features = false, features = [
"any",
"macros",
@@ -353,22 +353,25 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git"
rev = "5618e779cf2bb4755b499c630fba4c35e91898cb"
[patch.crates-io]
datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x
substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" }
datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-pg-catalog = { git = "https://github.com/GreptimeTeam/datafusion-postgres.git", rev = "3c77e6c32b8db80635a0d2f4b318a36b31170bc3" }
arrow-pg = { git = "https://github.com/GreptimeTeam/datafusion-postgres.git", rev = "3c77e6c32b8db80635a0d2f4b318a36b31170bc3" }
sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" }
# Temporary: use the GreptimeTeam fork of tikv-jemalloc-sys embedding
# jemalloc 5.3.1 + backport of 54f22c83 ("Initialize TSD tcache before
src/catalog/src/information_extension.rs
+10 -5
View File
@@ -27,8 +27,10 @@ use common_query::request::QueryRequest;
use common_recordbatch::adapter::{AsyncRecordBatchStreamAdapter, DfRecordBatchStreamAdapter};
use common_recordbatch::util::{ChainedRecordBatchStream, LimitedRecordBatchStream};
use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream};
use datafusion::common::Result;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::execution::TaskContext;
use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
@@ -157,10 +159,6 @@ impl ExecutionPlan for DistributedInspectExec {
"DistributedInspectExec"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema(&self) -> ArrowSchemaRef {
self.arrow_schema.clone()
}
@@ -173,6 +171,13 @@ impl ExecutionPlan for DistributedInspectExec {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
src/catalog/src/system_schema/information_schema/ssts.rs
+1 -1
View File
@@ -336,7 +336,7 @@ mod tests {
};
let plan = table.scan_to_plan(request).unwrap().unwrap();
assert!(plan.as_any().is::<EmptyExec>());
assert!(plan.as_ref().is::<EmptyExec>());
assert_eq!(1, plan.schema().fields().len());
}
}
src/catalog/src/table_source/dummy_catalog.rs
-13
View File
@@ -14,7 +14,6 @@
//! Dummy catalog for region server.
use std::any::Any;
use std::fmt;
use std::sync::Arc;
@@ -64,10 +63,6 @@ impl fmt::Debug for DummyCatalogList {
}
impl CatalogProviderList for DummyCatalogList {
fn as_any(&self) -> &dyn Any {
self
}
fn register_catalog(
&self,
_name: String,
@@ -98,10 +93,6 @@ struct DummyCatalogProvider {
}
impl CatalogProvider for DummyCatalogProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn schema_names(&self) -> Vec<String> {
vec![]
}
@@ -135,10 +126,6 @@ struct DummySchemaProvider {
#[async_trait]
impl SchemaProvider for DummySchemaProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn table_names(&self) -> Vec<String> {
vec![]
}
src/cmd/src/bin/query_perf_fixture/inspect_footer.rs
+73 -9
View File
@@ -14,18 +14,21 @@
use std::collections::BTreeSet;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use clap::Args as ClapArgs;
use datafusion_object_store::path::Path as StorePath;
use datafusion_object_store::{ObjectMeta, ObjectStore};
use futures::{StreamExt, TryStreamExt};
use datafusion_object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
use futures::future::BoxFuture;
use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt};
use object_store::config::ObjectStoreConfig;
use object_store::factory::new_raw_object_store;
use object_store::services::Fs;
use parquet::arrow::async_reader::ParquetObjectReader;
use parquet::file::metadata::ParquetMetaDataReader;
use parquet::arrow::async_reader::AsyncFileReader;
use parquet::errors::{ParquetError, Result as ParquetResult};
use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
use serde::{Deserialize, Serialize};
/// Same shape as `query_regression_runner::model::DestinationConfig`. The two
@@ -100,6 +103,70 @@ struct ListedFile {
relative_path: String,
}
/// An asynchronous Parquet reader backed directly by an object store.
///
/// The file size is retained from the listing so footer reads use bounded
/// ranges without an additional stat/head request.
#[derive(Clone, Debug)]
struct ObjectStoreReader {
store: Arc<dyn ObjectStore>,
path: StorePath,
file_size: u64,
}
impl ObjectStoreReader {
fn new(store: Arc<dyn ObjectStore>, path: StorePath, file_size: u64) -> Self {
Self {
store,
path,
file_size,
}
}
}
fn to_parquet_error(error: datafusion_object_store::Error) -> ParquetError {
ParquetError::External(Box::new(error))
}
impl AsyncFileReader for ObjectStoreReader {
fn get_bytes(
&mut self,
range: Range<u64>,
) -> BoxFuture<'_, ParquetResult<prost::bytes::Bytes>> {
self.store
.get_range(&self.path, range)
.map_err(to_parquet_error)
.boxed()
}
fn get_byte_ranges(
&mut self,
ranges: Vec<Range<u64>>,
) -> BoxFuture<'_, ParquetResult<Vec<prost::bytes::Bytes>>> {
async move {
self.store
.get_ranges(&self.path, &ranges)
.await
.map_err(to_parquet_error)
}
.boxed()
}
fn get_metadata<'a>(
&'a mut self,
_options: Option<&'a parquet::arrow::arrow_reader::ArrowReaderOptions>,
) -> BoxFuture<'a, ParquetResult<Arc<ParquetMetaData>>> {
let file_size = self.file_size;
async move {
let metadata = ParquetMetaDataReader::new()
.load_and_finish(self, file_size)
.await?;
Ok(Arc::new(metadata))
}
.boxed()
}
}
pub(super) async fn run_inspect_footer(
args: InspectFooterArgs,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
@@ -217,11 +284,8 @@ async fn inspect_file(
file: &ListedFile,
column: &str,
) -> Result<FooterFileReport, Box<dyn std::error::Error + Send + Sync>> {
let mut reader =
ParquetObjectReader::new(store, file.location.clone()).with_file_size(file.size);
let metadata = ParquetMetaDataReader::new()
.load_and_finish(&mut reader, file.size)
.await?;
let mut reader = ObjectStoreReader::new(store, file.location.clone(), file.size);
let metadata = reader.get_metadata(None).await?;
let file_metadata = metadata.file_metadata();
let row_groups = metadata.row_groups();
let mut columns = Vec::new();
src/common/datasource/Cargo.toml
+1 -1
View File
@@ -34,7 +34,7 @@ futures.workspace = true
lazy_static.workspace = true
object-store.workspace = true
object_store_opendal.workspace = true
orc-rust = { version = "0.8", default-features = false, features = ["async"] }
orc-rust = { version = "0.9", default-features = false, features = ["async"] }
parquet.workspace = true
paste.workspace = true
regex.workspace = true
src/common/datasource/src/file_format.rs
+6 -2
View File
@@ -35,7 +35,7 @@ use datafusion::datasource::file_format::file_compression_type::FileCompressionT
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
FileGroup, FileOpenFuture, FileScanConfigBuilder, FileSource, FileStream,
FileGroup, FileOpenFuture, FileScanConfigBuilder, FileSource, FileStreamBuilder,
};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::physical_plan::SendableRecordBatchStream;
@@ -321,7 +321,11 @@ pub async fn file_to_stream(
let store = Arc::new(object_store_opendal::OpendalStore::new(store.clone()));
let file_opener = config.file_source().create_file_opener(store, &config, 0)?;
let stream = FileStream::new(&config, 0, file_opener, &ExecutionPlanMetricsSet::new())?;
let stream = FileStreamBuilder::new(&config)
.with_partition(0)
.with_file_opener(file_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()?;
Ok(Box::pin(stream))
}
src/common/datasource/src/file_format/tests.rs
+10 -11
View File
@@ -19,7 +19,7 @@ use std::{assert_matches, vec};
use common_test_util::find_workspace_path;
use datafusion::assert_batches_eq;
use datafusion::datasource::physical_plan::{
CsvSource, FileScanConfig, FileSource, FileStream, JsonSource, ParquetSource,
CsvSource, FileScanConfig, FileSource, FileStreamBuilder, JsonSource, ParquetSource,
};
use datafusion::datasource::source::DataSourceExec;
use datafusion::execution::context::TaskContext;
@@ -50,16 +50,15 @@ impl Test<'_> {
.create_file_opener(store, &self.config, 0)
.unwrap();
let result = FileStream::new(
&self.config,
0,
file_opener,
&ExecutionPlanMetricsSet::new(),
)
.unwrap()
.map(|b| b.unwrap())
.collect::<Vec<_>>()
.await;
let result = FileStreamBuilder::new(&self.config)
.with_partition(0)
.with_file_opener(file_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.unwrap()
.map(|b| b.unwrap())
.collect::<Vec<_>>()
.await;
assert_batches_eq!(self.expected, &result);
}
src/common/datasource/src/test_util.rs
+13 -9
View File
@@ -20,7 +20,7 @@ use datafusion::datasource::file_format::file_compression_type::FileCompressionT
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
CsvSource, FileGroup, FileScanConfig, FileScanConfigBuilder, FileSource, FileStream,
CsvSource, FileGroup, FileScanConfig, FileScanConfigBuilder, FileSource, FileStreamBuilder,
JsonOpener, JsonSource,
};
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
@@ -110,13 +110,12 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi
let size = store.read(origin_path).await.unwrap().len();
let config = scan_config(None, origin_path, Arc::new(JsonSource::new(schema)));
let stream = FileStream::new(
&config,
0,
Arc::new(json_opener),
&ExecutionPlanMetricsSet::new(),
)
.unwrap();
let stream = FileStreamBuilder::new(&config)
.with_partition(0)
.with_file_opener(Arc::new(json_opener))
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.unwrap();
let (tmp_store, dir) = test_tmp_store("test_stream_to_json");
@@ -162,7 +161,12 @@ pub async fn setup_stream_to_csv_test(
0,
)
.unwrap();
let stream = FileStream::new(&config, 0, csv_opener, &ExecutionPlanMetricsSet::new()).unwrap();
let stream = FileStreamBuilder::new(&config)
.with_partition(0)
.with_file_opener(csv_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.unwrap();
let (tmp_store, dir) = test_tmp_store("test_stream_to_csv");
src/common/function/src/aggrs/aggr_wrapper.rs
+19 -21
View File
@@ -33,15 +33,15 @@ use datafusion::functions_aggregate::count::Count;
use datafusion::functions_aggregate::min_max::{Max, Min};
use datafusion::optimizer::AnalyzerRule;
use datafusion::optimizer::analyzer::type_coercion::TypeCoercion;
use datafusion::physical_planner::create_aggregate_expr_and_maybe_filter;
use datafusion_common::{Column, ScalarValue};
use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
use datafusion_expr::function::StateFieldsArgs;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{
Accumulator, Aggregate, AggregateUDF, AggregateUDFImpl, EmitTo, Expr, ExprSchemable,
GroupsAccumulator, LogicalPlan, Signature,
};
use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
use datafusion_physical_expr::aggregate::{AggregateFunctionExpr, LoweredAggregateBuilder};
use datatypes::arrow::datatypes::{DataType, Field};
use crate::aggrs::aggr_wrapper::fix_order::FixStateUdafOrderingAnalyzer;
@@ -210,12 +210,25 @@ impl StateMergeHelper {
lower_aggr_exprs.push(expr);
// then create the merge function using the physical expression of the original aggregate function
let (original_phy_expr, _filter, _ordering) = create_aggregate_expr_and_maybe_filter(
let (name, human_display) = match aggr_expr {
Expr::Alias(alias) => (alias.name.clone(), aggr_expr.human_display().to_string()),
Expr::AggregateFunction(_) => (
aggr_expr.schema_name().to_string(),
aggr_expr.human_display().to_string(),
),
_ => unreachable!("aggregate expression was validated above"),
};
let original_phy_expr = LoweredAggregateBuilder::new(
aggr_expr,
aggr.input.schema(),
aggr.input.schema().as_arrow(),
&Default::default(),
)?;
&PhysicalPlanningContext::default(),
)
.with_name(name)
.with_human_display(human_display)
.build()?
.aggregate;
let merge_func = MergeWrapper::new(
(*aggr_func.func).clone(),
@@ -371,9 +384,6 @@ impl AggregateUDFImpl for StateWrapper {
Ok(Box::new(StateGroupsAccum::new(inner, state_type)?))
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
self.name.as_str()
}
@@ -443,7 +453,7 @@ impl AggregateUDFImpl for StateWrapper {
&self,
statistics_args: &datafusion_expr::StatisticsArgs,
) -> Option<ScalarValue> {
let inner = self.inner().inner().as_any();
let inner = self.inner().inner();
// only count/min/max need special handling here, for getting result from statistics
// the result of count/min/max is also the result of count_state so can return directly
let can_use_stat = inner.is::<Count>() || inner.is::<Max>() || inner.is::<Min>();
@@ -572,11 +582,10 @@ impl GroupsAccumulator for StateGroupsAccum {
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> datafusion_common::Result<()> {
self.inner
.merge_batch(values, group_indices, opt_filter, total_num_groups)
.merge_batch(values, group_indices, total_num_groups)
}
fn evaluate(&mut self, emit_to: EmitTo) -> datafusion_common::Result<ArrayRef> {
@@ -596,10 +605,6 @@ impl GroupsAccumulator for StateGroupsAccum {
self.inner.convert_to_state(values, opt_filter)
}
fn supports_convert_to_state(&self) -> bool {
self.inner.supports_convert_to_state()
}
fn size(&self) -> usize {
self.inner.size()
}
@@ -800,10 +805,6 @@ impl AggregateUDFImpl for DeltaMergeWrapper {
}))
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
&self.name
}
@@ -973,9 +974,6 @@ impl AggregateUDFImpl for MergeWrapper {
Ok(Box::new(MergeAccum::new(inner_accum, &fields)))
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
self.name.as_str()
}
src/common/function/src/aggrs/aggr_wrapper/fix_order.rs
-1
View File
@@ -144,7 +144,6 @@ fn rewrite_expr(
let Some(old_state_wrapper) = aggregate_function
.func
.inner()
.as_any()
.downcast_ref::<StateWrapper>()
else {
return Ok(Transformed::no(Expr::AggregateFunction(aggregate_function)));
src/common/function/src/aggrs/aggr_wrapper/tests.rs
+36 -61
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
@@ -25,6 +24,7 @@ use arrow::record_batch::RecordBatch;
use arrow_schema::SchemaRef;
use common_telemetry::init_default_ut_logging;
use datafusion::catalog::{Session, TableProvider};
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::datasource::DefaultTableSource;
use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
use datafusion::functions_aggregate::average::avg_udaf;
@@ -39,16 +39,16 @@ use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
use datafusion::prelude::SessionContext;
use datafusion_common::arrow::array::AsArray;
use datafusion_common::arrow::datatypes::{Float64Type, UInt64Type};
use datafusion_common::{Column, TableReference};
use datafusion_common::{Column, Result, TableReference};
use datafusion_expr::expr::{AggregateFunction, NullTreatment};
use datafusion_expr::function::AccumulatorArgs;
use datafusion_expr::{
Aggregate, AggregateUDFImpl, ColumnarValue, Expr, LogicalPlan, ScalarFunctionArgs, SortExpr,
TableScan, TypeSignature, lit,
TableScanBuilder, TypeSignature, lit,
};
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
use datafusion_physical_expr::expressions::{Column as PhysicalColumn, col, lit as physical_lit};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use futures::{Stream, StreamExt as _};
use hyperloglogplus::HyperLogLog;
use pretty_assertions::assert_eq;
@@ -97,10 +97,6 @@ impl ExecutionPlan for MockInputExec {
"MockInputExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
@@ -109,6 +105,13 @@ impl ExecutionPlan for MockInputExec {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
@@ -203,10 +206,6 @@ impl Default for DummyTableProvider {
#[async_trait::async_trait]
impl TableProvider for DummyTableProvider {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema(&self) -> Arc<arrow_schema::Schema> {
self.schema.clone()
}
@@ -237,14 +236,12 @@ fn dummy_table_scan() -> LogicalPlan {
let table_provider = Arc::new(DummyTableProvider::default());
let table_source = DefaultTableSource::new(table_provider);
LogicalPlan::TableScan(
TableScan::try_new(
TableReference::bare("Number"),
Arc::new(table_source),
None,
vec![],
None,
)
.unwrap(),
TableScanBuilder::new(TableReference::bare("Number"), Arc::new(table_source))
.with_projection(None)
.with_filters(vec![])
.with_fetch(None)
.build()
.unwrap(),
)
}
@@ -252,14 +249,12 @@ fn dummy_table_scan_with_ts() -> LogicalPlan {
let table_provider = Arc::new(DummyTableProvider::with_ts(None));
let table_source = DefaultTableSource::new(table_provider);
LogicalPlan::TableScan(
TableScan::try_new(
TableReference::bare("Number"),
Arc::new(table_source),
None,
vec![],
None,
)
.unwrap(),
TableScanBuilder::new(TableReference::bare("Number"), Arc::new(table_source))
.with_projection(None)
.with_filters(vec![])
.with_fetch(None)
.build()
.unwrap(),
)
}
@@ -381,10 +376,7 @@ async fn test_sum_udaf() {
.create_physical_plan(&res.lower_state, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_state_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_state_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut state_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -414,10 +406,7 @@ async fn test_sum_udaf() {
.create_physical_plan(&res.upper_merge, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_merge_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_merge_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut merge_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -543,10 +532,7 @@ async fn test_avg_udaf() {
.create_physical_plan(&coerced_aggr_state_plan, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_state_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_state_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut state_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -582,10 +568,7 @@ async fn test_avg_udaf() {
.create_physical_plan(&res.upper_merge, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_merge_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_merge_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut merge_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -703,10 +686,7 @@ async fn test_last_value_order_by_udaf() {
.create_physical_plan(&fixed_aggr_state_plan, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_state_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_state_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let merge_input_fields = vec![Arc::new(Field::new(
@@ -796,10 +776,7 @@ async fn test_last_value_order_by_udaf() {
.create_physical_plan(&res.upper_merge, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_merge_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_merge_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut merge_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -900,7 +877,7 @@ fn test_avg_state_groups_accumulator_state_merge_evaluate() {
.update_batch(&merged_values, &merged_group_indices, None, 3)
.unwrap();
merged_accum
.merge_batch(&source_state, &[1, 2, 0], None, 3)
.merge_batch(&source_state, &[1, 2, 0], 3)
.unwrap();
let result = merged_accum.evaluate(EmitTo::All).unwrap();
@@ -1324,14 +1301,12 @@ async fn test_udaf_correct_eval_result() {
);
let table_source = DefaultTableSource::new(Arc::new(table_provider));
let logical_plan = LogicalPlan::TableScan(
TableScan::try_new(
test_table_ref.clone(),
Arc::new(table_source),
None,
vec![],
None,
)
.unwrap(),
TableScanBuilder::new(test_table_ref.clone(), Arc::new(table_source))
.with_projection(None)
.with_filters(vec![])
.with_fetch(None)
.build()
.unwrap(),
);
let args = case.args;
src/common/function/src/aggrs/approximate/uddsketch.rs
-2
View File
@@ -123,7 +123,6 @@ impl UddSketchState {
fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> {
let bucket_size = match args.exprs[0]
.as_any()
.downcast_ref::<Literal>()
.map(|lit| lit.value())
{
@@ -140,7 +139,6 @@ fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> {
};
let error_rate = match args.exprs[1]
.as_any()
.downcast_ref::<Literal>()
.map(|lit| lit.value())
{
src/common/function/src/aggrs/count_hash.rs
+26 -13
View File
@@ -26,7 +26,7 @@ use std::sync::Arc;
use ahash::RandomState;
use datafusion_common::cast::as_list_array;
use datafusion_common::error::Result;
use datafusion_common::hash_utils::create_hashes;
use datafusion_common::hash_utils::create_hashes_with_hasher;
use datafusion_common::utils::SingleRowListArrayBuilder;
use datafusion_common::{ScalarValue, internal_err, not_impl_err};
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
@@ -74,10 +74,6 @@ pub struct CountHash {
}
impl AggregateUDFImpl for CountHash {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
"count_hash"
}
@@ -208,7 +204,7 @@ impl GroupsAccumulator for CountHashGroupAccumulator {
let array = &values[0];
self.batch_hashes.clear();
self.batch_hashes.resize(array.len(), 0);
let hashes = create_hashes(
let hashes = create_hashes_with_hasher(
&[ArrayRef::clone(array)],
&self.random_state,
&mut self.batch_hashes,
@@ -279,7 +275,6 @@ impl GroupsAccumulator for CountHashGroupAccumulator {
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
_opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> Result<()> {
assert_eq!(
@@ -364,10 +359,6 @@ impl GroupsAccumulator for CountHashGroupAccumulator {
Ok(vec![Arc::new(list_array)])
}
fn supports_convert_to_state(&self) -> bool {
true
}
fn size(&self) -> usize {
// Base size of the struct
let mut size = size_of::<Self>();
@@ -423,7 +414,7 @@ impl Accumulator for CountHashAccumulator {
self.batch_hashes.clear();
self.batch_hashes.resize(arr.len(), 0);
let hashes = create_hashes(
let hashes = create_hashes_with_hasher(
&[ArrayRef::clone(arr)],
&self.random_state,
&mut self.batch_hashes,
@@ -522,6 +513,28 @@ mod tests {
Ok(())
}
#[test]
fn test_count_hash_accumulator_typed_null_state_merge() -> Result<()> {
let typed_nulls = Arc::new(Int32Array::from(vec![None, None])) as ArrayRef;
let mut fresh = create_test_accumulator();
fresh.update_batch(&[typed_nulls])?;
let fresh_state = fresh.state()?;
assert_eq!(fresh.evaluate()?, ScalarValue::Int64(Some(1)));
let persisted_state = Arc::new(
SingleRowListArrayBuilder::new(Arc::new(UInt64Array::from(vec![0])) as ArrayRef)
.build_list_array(),
) as ArrayRef;
let mut restored = create_test_accumulator();
restored.merge_batch(&[persisted_state])?;
assert_eq!(restored.evaluate()?, fresh.evaluate()?);
assert_eq!(restored.state()?, fresh_state);
Ok(())
}
#[test]
fn test_count_hash_accumulator_merge() -> Result<()> {
// Accumulator 1
@@ -622,7 +635,7 @@ mod tests {
// We will merge acc1's group 0 into acc2's group 0
// and acc1's group 1 into acc2's group 2
let merge_group_indices = vec![0, 2];
acc2.merge_batch(&state1, &merge_group_indices, None, 3)?;
acc2.merge_batch(&state1, &merge_group_indices, 3)?;
let result_array = acc2.evaluate(EmitTo::All)?;
let result = result_array.as_any().downcast_ref::<Int64Array>().unwrap();
src/common/function/src/helper.rs
+27
View File
@@ -23,6 +23,33 @@ use datatypes::types::cast::cast;
use datatypes::value::ValueRef;
use snafu::{OptionExt, ResultExt};
/// Integer types accepted by geospatial function signatures.
pub(crate) const INTEGER_TYPES: &[DataType] = &[
DataType::Int8,
DataType::Int16,
DataType::Int32,
DataType::Int64,
DataType::UInt8,
DataType::UInt16,
DataType::UInt32,
DataType::UInt64,
];
/// Legacy primitive numeric signature types; Decimal values are coerced to `Float64`.
pub(crate) const NUMERICS: &[DataType] = &[
DataType::Int8,
DataType::Int16,
DataType::Int32,
DataType::Int64,
DataType::UInt8,
DataType::UInt16,
DataType::UInt32,
DataType::UInt64,
DataType::Float16,
DataType::Float32,
DataType::Float64,
];
/// Create a function signature with oneof signatures of interleaving two arguments.
pub(crate) fn one_of_sigs2(args1: Vec<DataType>, args2: Vec<DataType>) -> Signature {
let mut sigs = Vec::with_capacity(args1.len() * args2.len());
src/common/function/src/scalars/anomaly/iqr.rs
+1 -6
View File
@@ -24,7 +24,6 @@
//! When IQR = 0 (constant quartiles), returns 0.0 if value is on the fence,
//! or +inf if value is outside.
use std::any::Any;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
@@ -32,11 +31,11 @@ use std::sync::Arc;
use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::{DataFusionError, Result, ScalarValue};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{PartitionEvaluator, Signature, Volatility, WindowUDFImpl};
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use crate::helper::NUMERICS;
use crate::scalars::anomaly::utils::{cast_to_f64, collect_window_values, percentile_sorted};
/// Minimum valid samples for IQR (linear-interpolated Q1 != Q3 is possible at n >= 3).
@@ -56,10 +55,6 @@ impl AnomalyScoreIqr {
}
impl WindowUDFImpl for AnomalyScoreIqr {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"anomaly_score_iqr"
}
src/common/function/src/scalars/anomaly/mad.rs
+1 -6
View File
@@ -20,7 +20,6 @@
//! When MAD = 0 (majority-constant window), returns 0.0 if value equals
//! median, or +inf otherwise.
use std::any::Any;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
@@ -28,11 +27,11 @@ use std::sync::Arc;
use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::{DataFusionError, Result, ScalarValue};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{PartitionEvaluator, Signature, Volatility, WindowUDFImpl};
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use crate::helper::NUMERICS;
use crate::scalars::anomaly::utils::{
anomaly_ratio, cast_to_f64, collect_window_values, median_f64,
};
@@ -57,10 +56,6 @@ impl AnomalyScoreMad {
}
impl WindowUDFImpl for AnomalyScoreMad {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"anomaly_score_mad"
}
src/common/function/src/scalars/anomaly/zscore.rs
+1 -6
View File
@@ -19,7 +19,6 @@
//! When stddev = 0 (constant window), returns 0.0 if value equals mean,
//! or +inf otherwise.
use std::any::Any;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
@@ -27,11 +26,11 @@ use std::sync::Arc;
use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::{DataFusionError, Result, ScalarValue};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{PartitionEvaluator, Signature, Volatility, WindowUDFImpl};
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use crate::helper::NUMERICS;
use crate::scalars::anomaly::utils::{anomaly_ratio, cast_to_f64, collect_window_values};
/// Minimum valid samples for zscore (stddev requires n >= 2).
@@ -51,10 +50,6 @@ impl AnomalyScoreZscore {
}
impl WindowUDFImpl for AnomalyScoreZscore {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"anomaly_score_zscore"
}
src/common/function/src/scalars/geo/geohash.rs
+3 -3
View File
@@ -22,12 +22,12 @@ use datafusion::arrow::array::{Array, AsArray, ListBuilder, StringViewBuilder};
use datafusion::arrow::datatypes::{DataType, Field, Float64Type, UInt8Type};
use datafusion::logical_expr::ColumnarValue;
use datafusion_common::DataFusionError;
use datafusion_expr::type_coercion::aggregates::INTEGERS;
use datafusion_expr::{ScalarFunctionArgs, Signature, TypeSignature, Volatility};
use geohash::Coord;
use snafu::ResultExt;
use crate::function::{Function, extract_args};
use crate::helper::INTEGER_TYPES;
use crate::scalars::geo::helpers;
fn ensure_resolution_usize(v: u8) -> datafusion_common::Result<usize> {
@@ -49,7 +49,7 @@ impl Default for GeohashFunction {
fn default() -> Self {
let mut signatures = Vec::new();
for coord_type in &[DataType::Float32, DataType::Float64] {
for resolution_type in INTEGERS {
for resolution_type in INTEGER_TYPES {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
@@ -146,7 +146,7 @@ impl Default for GeohashNeighboursFunction {
fn default() -> Self {
let mut signatures = Vec::new();
for coord_type in &[DataType::Float32, DataType::Float64] {
for resolution_type in INTEGERS {
for resolution_type in INTEGER_TYPES {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
src/common/function/src/scalars/geo/h3.rs
+4 -4
View File
@@ -26,7 +26,6 @@ use datafusion::arrow::compute;
use datafusion::arrow::datatypes::{Float64Type, Int64Type, UInt8Type, UInt64Type};
use datafusion::logical_expr::ColumnarValue;
use datafusion_common::{DataFusionError, ScalarValue};
use datafusion_expr::type_coercion::aggregates::INTEGERS;
use datafusion_expr::{ScalarFunctionArgs, Signature, TypeSignature, Volatility};
use datatypes::arrow::datatypes::{DataType, Field};
use derive_more::Display;
@@ -34,6 +33,7 @@ use h3o::{CellIndex, LatLng, Resolution};
use snafu::prelude::*;
use crate::function::{Function, extract_args};
use crate::helper::INTEGER_TYPES;
use crate::scalars::geo::helpers;
static CELL_TYPES: LazyLock<Vec<DataType>> =
@@ -42,11 +42,11 @@ static CELL_TYPES: LazyLock<Vec<DataType>> =
static COORDINATE_TYPES: LazyLock<Vec<DataType>> =
LazyLock::new(|| vec![DataType::Float32, DataType::Float64]);
static RESOLUTION_TYPES: &[DataType] = INTEGERS;
static RESOLUTION_TYPES: &[DataType] = INTEGER_TYPES;
static DISTANCE_TYPES: &[DataType] = INTEGERS;
static DISTANCE_TYPES: &[DataType] = INTEGER_TYPES;
static POSITION_TYPES: &[DataType] = INTEGERS;
static POSITION_TYPES: &[DataType] = INTEGER_TYPES;
/// Function that returns [h3] encoding cellid for a given geospatial coordinate.
///
src/common/function/src/scalars/geo/s2.rs
+2 -1
View File
@@ -25,6 +25,7 @@ use s2::latlng::LatLng;
use snafu::ensure;
use crate::function::{Function, extract_args};
use crate::helper::INTEGER_TYPES;
use crate::scalars::geo::helpers;
use crate::scalars::geo::helpers::ensure_and_coerce;
@@ -34,7 +35,7 @@ static CELL_TYPES: LazyLock<Vec<DataType>> =
static COORDINATE_TYPES: LazyLock<Vec<DataType>> =
LazyLock::new(|| vec![DataType::Float32, DataType::Float64]);
static LEVEL_TYPES: &[DataType] = datafusion_expr::type_coercion::aggregates::INTEGERS;
static LEVEL_TYPES: &[DataType] = INTEGER_TYPES;
/// Function that returns [s2] encoding cellid for a given geospatial coordinate.
///
src/common/function/src/scalars/json/json_get.rs
+31 -3
View File
@@ -532,7 +532,7 @@ mod tests {
use datafusion_common::arrow::datatypes::{Float64Type, Int64Type};
use datatypes::extension::json::Json2ExtensionType;
use datatypes::types::parse_string_to_jsonb;
use serde_json::json;
use serde_json::{Value, json};
use super::*;
@@ -595,6 +595,34 @@ mod tests {
})
}
fn assert_json_or_string_eq(actual: Option<&str>, expected: Option<&str>) {
let is_json_container = |value: &str| {
matches!(
serde_json::from_str::<Value>(value),
Ok(Value::Object(_) | Value::Array(_))
)
};
match (actual, expected) {
(Some(actual), Some(expected))
if is_json_container(actual) || is_json_container(expected) =>
{
let actual_value = serde_json::from_str::<Value>(actual).unwrap_or_else(|error| {
panic!("failed to parse actual JSON result {actual:?}: {error}")
});
let expected_value =
serde_json::from_str::<Value>(expected).unwrap_or_else(|error| {
panic!("failed to parse expected JSON result {expected:?}: {error}")
});
assert_eq!(
actual_value, expected_value,
"JSON result mismatch: actual {actual:?}, expected {expected:?}"
);
}
_ => assert_eq!(actual, expected),
}
}
#[test]
fn test_json_get_int() {
let json_get_int = JsonGetInt::default();
@@ -895,7 +923,7 @@ mod tests {
let result = result.as_string_view();
assert_eq!(1, result.len());
let actual = result.is_valid(0).then(|| result.value(0));
assert_eq!(actual, expect);
assert_json_or_string_eq(actual, expect);
}
}
@@ -1033,7 +1061,7 @@ mod tests {
let result = result.as_string_view();
assert_eq!(1, result.len());
let actual = result.is_valid(0).then(|| result.value(0));
assert_eq!(actual, expect);
assert_json_or_string_eq(actual, expect);
}
let json_strings = [
src/common/function/src/scalars/json/json_get_rewriter.rs
+5 -13
View File
@@ -59,10 +59,8 @@ impl FunctionRewrite for JsonGetRewriter {
// json_get(column, path, <data_type>)
// )
fn inject_type_from_cast_expr(cast: Cast) -> Result<Transformed<Expr>> {
let Cast {
expr,
mut data_type,
} = cast;
let Cast { expr, field } = cast;
let mut data_type = field.data_type().clone();
let mut json_get = match *expr {
Expr::ScalarFunction(f)
@@ -73,7 +71,7 @@ fn inject_type_from_cast_expr(cast: Cast) -> Result<Transformed<Expr>> {
expr => {
return Ok(Transformed::no(Expr::Cast(Cast {
expr: Box::new(expr),
data_type,
field,
})));
}
};
@@ -204,10 +202,7 @@ mod tests {
});
// Create a cast expression: json_get(...)::int8
let cast_expr = Expr::Cast(Cast {
expr: Box::new(json_expr),
data_type: DataType::Int8,
});
let cast_expr = Expr::Cast(Cast::new(Box::new(json_expr), DataType::Int8));
// Apply the rewriter
let result = rewriter.rewrite(cast_expr, &schema, &config).unwrap();
@@ -279,10 +274,7 @@ mod tests {
// Create an arrow cast function: cast(json_get(...), 'Int64')
// Note: ArrowCastFunc doesn't exist in this codebase, so this test uses a simple cast instead
let arrow_cast_expr = Expr::Cast(Cast {
expr: Box::new(json_get_expr),
data_type: DataType::Int64,
});
let arrow_cast_expr = Expr::Cast(Cast::new(Box::new(json_get_expr), DataType::Int64));
// Apply the rewriter
let result = rewriter.rewrite(arrow_cast_expr, &schema, &config).unwrap();
src/common/function/src/scalars/matches.rs
+7 -2
View File
@@ -24,6 +24,7 @@ use datafusion::execution::SessionStateBuilder;
use datafusion::logical_expr::{self, ColumnarValue, Expr, Volatility};
use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
use datafusion_common::DataFusionError;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{ScalarFunctionArgs, Signature};
use datatypes::arrow::array::RecordBatch;
use datatypes::arrow::datatypes::{DataType, Field};
@@ -110,8 +111,12 @@ impl MatchesFunction {
let input_schema = Self::input_schema();
let session_state = SessionStateBuilder::new().with_default_features().build();
let planner = DefaultPhysicalPlanner::default();
let physical_expr =
planner.create_physical_expr(&like_expr, &input_schema, &session_state)?;
let physical_expr = planner.create_physical_expr(
&like_expr,
&input_schema,
&session_state,
&PhysicalPlanningContext::default(),
)?;
let arrow_schema = Arc::new(input_schema.as_arrow().clone());
let input_record_batch = RecordBatch::try_new(arrow_schema, vec![data_array]).unwrap();
src/common/function/src/scalars/math/clamp.rs
+77 -2
View File
@@ -19,10 +19,10 @@ use datafusion::arrow::array::{Array, ArrayRef, AsArray, PrimitiveArray};
use datafusion::arrow::datatypes::DataType as ArrowDataType;
use datafusion::logical_expr::{ColumnarValue, Volatility};
use datafusion_common::{DataFusionError, ScalarValue, utils};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{ScalarFunctionArgs, Signature};
use crate::function::Function;
use crate::helper::NUMERICS;
#[derive(Clone, Debug)]
pub struct ClampFunction {
@@ -338,7 +338,9 @@ mod test {
use arrow_schema::Field;
use datafusion_common::config::ConfigOptions;
use datatypes::arrow::array::{ArrayRef, Float64Array, Int64Array, UInt64Array};
use datatypes::arrow::array::{
ArrayRef, Decimal128Array, Float64Array, Int64Array, UInt64Array,
};
use datatypes::arrow_array::StringArray;
use super::*;
@@ -370,6 +372,79 @@ mod test {
impl_test_eval!(ClampMinFunction);
impl_test_eval!(ClampMaxFunction);
fn decimal_array(values: Vec<i128>) -> ColumnarValue {
ColumnarValue::Array(Arc::new(
Decimal128Array::from(values)
.with_precision_and_scale(10, 2)
.unwrap(),
))
}
fn decimal_scalar(value: i128) -> ColumnarValue {
ColumnarValue::Scalar(ScalarValue::Decimal128(Some(value), 10, 2))
}
#[allow(deprecated)]
fn evaluate_decimal(
function: &dyn Function,
args: Vec<ColumnarValue>,
) -> datafusion_common::Result<ArrayRef> {
let input_types = args
.iter()
.map(ColumnarValue::data_type)
.collect::<Vec<_>>();
let planned_types = datafusion_expr::type_coercion::functions::data_types(
function.name(),
&input_types,
function.signature(),
)?;
let args = args
.into_iter()
.zip(planned_types)
.map(|(arg, planned_type)| arg.cast_to(&planned_type, None))
.collect::<datafusion_common::Result<Vec<_>>>()?;
function
.invoke_with_args(ScalarFunctionArgs {
args,
arg_fields: vec![],
number_rows: 3,
return_field: Arc::new(Field::new("x", ArrowDataType::Float64, false)),
config_options: Arc::new(ConfigOptions::new()),
})
.and_then(|value| value.to_array(3))
}
#[test]
fn clamp_decimal_coercion_executes_as_float64() {
let test_cases = [
(
Box::new(ClampFunction::default()) as Box<dyn Function>,
vec![
decimal_array(vec![100, 300, 500]),
decimal_scalar(200),
decimal_scalar(400),
],
vec![2.0, 3.0, 4.0],
),
(
Box::new(ClampMinFunction::default()),
vec![decimal_array(vec![100, 300, 500]), decimal_scalar(200)],
vec![2.0, 3.0, 5.0],
),
(
Box::new(ClampMaxFunction::default()),
vec![decimal_array(vec![100, 300, 500]), decimal_scalar(200)],
vec![1.0, 2.0, 2.0],
),
];
for (function, args, expected) in test_cases {
let result = evaluate_decimal(function.as_ref(), args).unwrap();
let expected: ArrayRef = Arc::new(Float64Array::from(expected));
assert_eq!(expected.as_ref(), result.as_ref());
}
}
#[test]
fn clamp_i64() {
let inputs = [
src/common/function/src/scalars/math/modulo.rs
+64 -2
View File
@@ -18,10 +18,10 @@ use std::fmt::Display;
use datafusion_common::arrow::compute;
use datafusion_common::arrow::compute::kernels::numeric;
use datafusion_common::arrow::datatypes::DataType;
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
use crate::function::{Function, extract_args};
use crate::helper::NUMERICS;
const NAME: &str = "mod";
@@ -90,12 +90,74 @@ mod tests {
use std::sync::Arc;
use arrow_schema::Field;
use datafusion_common::ScalarValue;
use datafusion_common::arrow::array::{
AsArray, Float64Array, Int32Array, StringViewArray, UInt32Array,
AsArray, Decimal128Array, Float64Array, Int32Array, StringViewArray, UInt32Array,
};
use datafusion_common::arrow::datatypes::{Float64Type, Int64Type, UInt64Type};
use super::*;
fn decimal_array(values: Vec<i128>) -> ColumnarValue {
ColumnarValue::Array(Arc::new(
Decimal128Array::from(values)
.with_precision_and_scale(10, 2)
.unwrap(),
))
}
fn decimal_scalar(value: i128) -> ColumnarValue {
ColumnarValue::Scalar(ScalarValue::Decimal128(Some(value), 10, 2))
}
#[test]
#[allow(deprecated)]
fn modulo_decimal_coercion_executes_as_float64() {
let function = ModuloFunction::default();
let test_cases = [
(
vec![decimal_array(vec![500, 600]), decimal_scalar(200)],
vec![1.0, 0.0],
),
(
vec![decimal_scalar(500), decimal_array(vec![200, 300])],
vec![1.0, 2.0],
),
];
for (args, expected) in test_cases {
let input_types = args
.iter()
.map(ColumnarValue::data_type)
.collect::<Vec<_>>();
let planned_types = datafusion_expr::type_coercion::functions::data_types(
function.name(),
&input_types,
function.signature(),
)
.unwrap();
assert_eq!(vec![DataType::Float64; 2], planned_types);
let args = args
.into_iter()
.zip(planned_types)
.map(|(arg, planned_type)| arg.cast_to(&planned_type, None))
.collect::<datafusion_common::Result<Vec<_>>>()
.unwrap();
let result = function
.invoke_with_args(ScalarFunctionArgs {
args,
arg_fields: vec![],
number_rows: 2,
return_field: Arc::new(Field::new("x", DataType::Float64, false)),
config_options: Arc::new(Default::default()),
})
.unwrap()
.to_array(2)
.unwrap();
let result = result.as_primitive::<Float64Type>();
assert_eq!(&Float64Array::from(expected), result);
}
}
#[test]
fn test_mod_function_signed() {
let function = ModuloFunction::default();
src/common/function/src/scalars/math/rate.rs
+8 -7
View File
@@ -18,11 +18,11 @@ use common_query::error;
use datafusion::arrow::compute::kernels::numeric;
use datafusion_common::arrow::compute::kernels::cast;
use datafusion_common::arrow::datatypes::DataType;
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
use snafu::ResultExt;
use crate::function::{Function, extract_args};
use crate::helper::NUMERICS;
/// generates rates from a sequence of adjacent data points.
#[derive(Clone, Debug)]
@@ -96,12 +96,13 @@ mod tests {
let rate = RateFunction::default();
assert_eq!("rate", rate.name());
assert_eq!(DataType::Float64, rate.return_type(&[]).unwrap());
assert!(matches!(rate.signature(),
Signature {
type_signature: TypeSignature::Uniform(2, valid_types),
volatility: Volatility::Immutable,
..
} if valid_types == NUMERICS
assert!(matches!(
rate.signature(),
Signature {
type_signature: TypeSignature::Uniform(2, valid_types),
volatility: Volatility::Immutable,
..
} if valid_types == NUMERICS
));
let values = vec![1.0, 3.0, 6.0];
let ts = vec![0, 1, 2];
src/common/function/src/scalars/udf.rs
-5
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::fmt::{Debug, Formatter};
use std::hash::{Hash, Hasher};
@@ -49,10 +48,6 @@ impl Hash for ScalarUdf {
}
impl ScalarUDFImpl for ScalarUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
self.function.name()
}
src/common/function/src/system/pg_catalog.rs
+3
View File
@@ -389,6 +389,9 @@ impl PGCatalogFunction {
registry.register(pg_catalog::create_pg_get_partition_ancestors_udf());
registry.register(pg_catalog::quote_ident_udf::create_quote_ident_udf());
registry.register(pg_catalog::quote_ident_udf::create_parse_ident_udf());
// Register array bound UDFs used by pg_catalog views.
registry.register(pg_catalog::array_bounds_udf::create_array_upper_udf());
registry.register(pg_catalog::array_bounds_udf::create_array_lower_udf());
registry.register_scalar(ObjDescriptionFunction::new());
registry.register_scalar(ColDescriptionFunction::new());
registry.register_scalar(ShobjDescriptionFunction::new());
src/common/macro/src/admin_fn.rs
-4
View File
@@ -239,10 +239,6 @@ fn build_struct(
// Implement DataFusion's ScalarUDFImpl trait
impl datafusion::logical_expr::ScalarUDFImpl for #name {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
#display_name
}
src/common/query/src/request.rs
+8 -10
View File
@@ -137,9 +137,7 @@ fn portable_remote_dyn_filter_expr(
bounds_only: bool,
) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
expr.transform_up(|node| {
if node.as_any().is::<HashTableLookupExpr>()
|| (bounds_only && node.as_any().is::<InListExpr>())
{
if node.is::<HashTableLookupExpr>() || (bounds_only && node.is::<InListExpr>()) {
Ok(Transformed::yes(lit(true)))
} else {
Ok(Transformed::no(node))
@@ -183,7 +181,7 @@ fn validate_payload_size(
fn validate_supported_payload_expr(expr: &Arc<dyn PhysicalExpr>) -> DataFusionResult<()> {
expr.apply(|node| {
if node.as_any().is::<HashTableLookupExpr>() {
if node.is::<HashTableLookupExpr>() {
return Err(DataFusionError::Plan(
"HashTableLookupExpr cannot be encoded into DynFilterPayload::Datafusion"
.to_string(),
@@ -207,7 +205,7 @@ fn validate_decoded_payload_expr(
input_schema: &datafusion::arrow::datatypes::Schema,
) -> DataFusionResult<()> {
expr.apply(|node| {
if let Some(column) = node.as_any().downcast_ref::<Column>() {
if let Some(column) = node.downcast_ref::<Column>() {
let Some(field) = input_schema.fields().get(column.index()) else {
return Err(DataFusionError::Plan(format!(
"Decoded Column '{}' references out-of-bounds index {} for input schema of size {}",
@@ -391,8 +389,8 @@ mod tests {
.decode_datafusion_expr(&TaskContext::default(), &schema, 1024)
.unwrap();
let original = expr.as_any().downcast_ref::<Column>().unwrap();
let decoded = decoded.as_any().downcast_ref::<Column>().unwrap();
let original = expr.downcast_ref::<Column>().unwrap();
let decoded = decoded.downcast_ref::<Column>().unwrap();
assert_eq!(decoded.name(), original.name());
assert_eq!(decoded.index(), original.index());
@@ -457,7 +455,7 @@ mod tests {
)) as Arc<dyn PhysicalExpr>;
let lookup = Arc::new(HashTableLookupExpr::new(
vec![Arc::clone(&device_id)],
SeededRandomState::with_seeds(0, 0, 0, 0),
SeededRandomState::with_seed(0),
Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))),
"hash_lookup".to_string(),
)) as Arc<dyn PhysicalExpr>;
@@ -539,10 +537,10 @@ mod tests {
));
}
fn contains_expr<T: 'static>(expr: &Arc<dyn PhysicalExpr>) -> bool {
fn contains_expr<T: PhysicalExpr>(expr: &Arc<dyn PhysicalExpr>) -> bool {
let mut found = false;
expr.apply(|node| {
if node.as_any().is::<T>() {
if node.is::<T>() {
found = true;
Ok(TreeNodeRecursion::Stop)
} else {
src/common/query/src/request/initial_remote_dyn_filter_reg.rs
+1 -1
View File
@@ -371,7 +371,7 @@ mod tests {
let decoded = reg
.decode_children(&TaskContext::default(), &schema, 1024)
.unwrap();
let decoded = decoded[0].as_any().downcast_ref::<Column>().unwrap();
let decoded = decoded[0].downcast_ref::<Column>().unwrap();
assert_eq!(reg.filter_id, "filter-1");
assert_eq!(decoded.name(), "host");
src/common/query/src/stream.rs
+11 -6
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::fmt::{Debug, Formatter};
use std::sync::{Arc, Mutex};
@@ -22,8 +21,11 @@ use datafusion::execution::SendableRecordBatchStream as DfSendableRecordBatchStr
use datafusion::execution::context::TaskContext;
use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalSortExpr};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties,
};
use datafusion_common::DataFusionError;
use datafusion_common::tree_node::TreeNodeRecursion;
use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
use datatypes::schema::SchemaRef;
@@ -99,10 +101,6 @@ impl DisplayAs for StreamScanAdapter {
}
impl ExecutionPlan for StreamScanAdapter {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> ArrowSchemaRef {
self.arrow_schema.clone()
}
@@ -115,6 +113,13 @@ impl ExecutionPlan for StreamScanAdapter {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> datafusion_common::Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
// DataFusion will swap children unconditionally.
// But since this node is leaf node, it's safe to just return self.
fn with_new_children(
src/common/recordbatch/src/adapter.rs
+18 -7
View File
@@ -29,6 +29,7 @@ use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::ExecutionProps;
use datafusion::logical_expr::Expr;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::logical_expr::utils::conjunction;
use datafusion::physical_expr::create_physical_expr;
use datafusion::physical_plan::metrics::{BaselineMetrics, MetricValue};
@@ -98,8 +99,13 @@ where
.to_dfschema_ref()
.context(error::PhysicalExprSnafu)?;
let filters = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())
.context(error::PhysicalExprSnafu)?;
let filters = create_physical_expr(
&expr,
&df_schema,
&ExecutionProps::new(),
&PhysicalPlanningContext::default(),
)
.context(error::PhysicalExprSnafu)?;
Some(filters)
} else {
None
@@ -931,7 +937,6 @@ fn convert_map_to_json_binary(
#[cfg(test)]
mod test {
use std::any::Any;
use std::time::Duration;
use common_error::ext::BoxedError;
@@ -942,6 +947,7 @@ mod test {
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet};
use datafusion::physical_plan::{DisplayAs, PlanProperties};
use datafusion_common::tree_node::TreeNodeRecursion;
use datatypes::arrow::array::{ArrayRef, MapArray, StringArray, StructArray};
use datatypes::arrow::buffer::OffsetBuffer;
use datatypes::arrow::datatypes::Field;
@@ -1021,10 +1027,6 @@ mod test {
REGION_SCAN_EXEC_NAME
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
@@ -1033,6 +1035,15 @@ mod test {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(
&Arc<dyn PhysicalExpr>,
) -> datafusion_common::Result<TreeNodeRecursion>,
) -> datafusion_common::Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
src/common/recordbatch/src/filter.rs
+15 -2
View File
@@ -616,6 +616,7 @@ mod test {
use std::sync::Arc;
use datafusion::execution::context::ExecutionProps;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::logical_expr::{BinaryExpr, col, lit};
use datafusion::physical_expr::create_physical_expr;
use datafusion_common::{Column, DFSchema};
@@ -737,7 +738,13 @@ mod test {
]);
let df_schema = DFSchema::try_from(schema.clone()).unwrap();
let props = ExecutionProps::new();
let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap();
let physical_expr = create_physical_expr(
&expr,
&df_schema,
&props,
&PhysicalPlanningContext::default(),
)
.unwrap();
let batch = RecordBatch::try_new(
Arc::new(schema),
vec![
@@ -783,7 +790,13 @@ mod test {
let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]);
let df_schema = DFSchema::try_from(schema.clone()).unwrap();
let props = ExecutionProps::new();
let physical_expr = create_physical_expr(&col_or_expr, &df_schema, &props).unwrap();
let physical_expr = create_physical_expr(
&col_or_expr,
&df_schema,
&props,
&PhysicalPlanningContext::default(),
)
.unwrap();
// Create test data
let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![
src/common/recordbatch/src/lib.rs
+44 -6
View File
@@ -44,7 +44,9 @@ pub use datatypes::arrow::record_batch::RecordBatch as DfRecordBatch;
use datatypes::arrow::util::display::{
ArrayFormatter, ArrayFormatterFactory, DisplayIndex, FormatOptions, FormatResult,
};
use datatypes::arrow::util::pretty::pretty_format_batches_with_options;
use datatypes::arrow::util::pretty::{
pretty_format_batches_with_options, pretty_format_batches_with_schema,
};
use datatypes::extension::json::is_any_json_extension_type;
use datatypes::prelude::{ConcreteDataType, DataType, VectorRef};
use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
@@ -396,12 +398,19 @@ impl RecordBatches {
.iter()
.map(|x| x.df_record_batch().clone())
.collect::<Vec<_>>();
let options =
FormatOptions::default().with_formatter_factory(Some(&BinaryFormatterFactory));
let result =
pretty_format_batches_with_options(df_batches, &options).context(error::FormatSnafu)?;
let result: String = if df_batches.is_empty() {
pretty_format_batches_with_schema(self.schema.arrow_schema().clone(), df_batches)
.context(error::FormatSnafu)?
.to_string()
} else {
let options =
FormatOptions::default().with_formatter_factory(Some(&BinaryFormatterFactory));
pretty_format_batches_with_options(df_batches, &options)
.context(error::FormatSnafu)?
.to_string()
};
Ok(result.to_string())
Ok(result)
}
pub fn try_new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Result<Self> {
@@ -1082,6 +1091,35 @@ mod tests {
assert_eq!(r.take(), expected);
}
#[tokio::test]
async fn test_recordbatches_pretty_print_empty_batches_preserves_schema() {
let schema = Arc::new(Schema::new(vec![
ColumnSchema::new("unit", ConcreteDataType::string_datatype(), false),
ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
ColumnSchema::new(
"lhs.degrees(val) + rhs.radians(val)",
ConcreteDataType::float64_datatype(),
false,
),
]));
let batches =
RecordBatches::try_collect(Box::pin(EmptyRecordBatchStream::new(schema.clone())))
.await
.unwrap();
assert_eq!(schema, batches.schema());
let expected = "\
+------+----+-------------------------------------+
| unit | ts | lhs.degrees(val) + rhs.radians(val) |
+------+----+-------------------------------------+
+------+----+-------------------------------------+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
#[test]
fn test_recordbatches_try_new() {
let column_a = ColumnSchema::new("a", ConcreteDataType::int32_datatype(), false);
src/datanode/src/region_server/catalog.rs
-9
View File
@@ -179,9 +179,6 @@ impl NameAwareCatalogList {
}
impl CatalogProviderList for NameAwareCatalogList {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn register_catalog(
&self,
_name: String,
@@ -203,9 +200,6 @@ struct NameAwareCatalogProvider {
}
impl CatalogProvider for NameAwareCatalogProvider {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema_names(&self) -> Vec<String> {
vec![]
}
@@ -229,9 +223,6 @@ impl std::fmt::Debug for NameAwareSchemaProvider {
#[async_trait::async_trait]
impl SchemaProvider for NameAwareSchemaProvider {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn table_names(&self) -> Vec<String> {
vec![]
}
src/datatypes/src/value.rs
+2
View File
@@ -1240,6 +1240,8 @@ impl TryFrom<ScalarValue> for Value {
| ScalarValue::Decimal256(_, _, _)
| ScalarValue::FixedSizeList(_)
| ScalarValue::LargeList(_)
| ScalarValue::ListView(_)
| ScalarValue::LargeListView(_)
| ScalarValue::Union(_, _, _)
| ScalarValue::Float16(_)
| ScalarValue::Utf8View(_)
src/datatypes/src/vectors/helper.rs
+2
View File
@@ -150,6 +150,8 @@ impl Helper {
| ScalarValue::Decimal256(_, _, _)
| ScalarValue::FixedSizeList(_)
| ScalarValue::LargeList(_)
| ScalarValue::ListView(_)
| ScalarValue::LargeListView(_)
| ScalarValue::Dictionary(_, _)
| ScalarValue::Union(_, _, _)
| ScalarValue::Utf8View(_)
src/file-engine/src/query/file_stream.rs
+16 -10
View File
@@ -23,7 +23,8 @@ use datafusion::config::CsvOptions;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
CsvSource, FileGroup, FileScanConfigBuilder, FileSource, FileStream, JsonSource, ParquetSource,
CsvSource, FileGroup, FileScanConfigBuilder, FileSource, FileStreamBuilder, JsonSource,
ParquetSource,
};
use datafusion::datasource::source::DataSourceExec;
use datafusion::physical_expr::create_physical_expr;
@@ -34,6 +35,7 @@ use datafusion::physical_plan::{
};
use datafusion::prelude::SessionContext;
use datafusion_expr::expr::Expr;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::utils::conjunction;
use datatypes::schema::SchemaRef;
use object_store::ObjectStore;
@@ -66,13 +68,12 @@ fn build_record_batch_stream(
));
let file_opener = config.file_source().create_file_opener(store, &config, 0)?;
let stream = FileStream::new(
&config,
0, // partition: hard-code
file_opener,
&ExecutionPlanMetricsSet::new(),
)
.context(error::BuildStreamSnafu)?;
let stream = FileStreamBuilder::new(&config)
.with_partition(0) // partition: hard-code
.with_file_opener(file_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.context(error::BuildStreamSnafu)?;
Ok(Box::pin(stream))
}
@@ -138,8 +139,13 @@ fn new_parquet_stream_with_exec_plan(
.to_dfschema_ref()
.context(error::ParquetScanPlanSnafu)?;
let filters = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())
.context(error::ParquetScanPlanSnafu)?;
let filters = create_physical_expr(
&expr,
&df_schema,
&ExecutionProps::new(),
&PhysicalPlanningContext::default(),
)
.context(error::ParquetScanPlanSnafu)?;
parquet_source = parquet_source.with_predicate(filters);
};
src/flow/src/batching_mode/time_window.rs
+11 -2
View File
@@ -37,6 +37,7 @@ use datafusion_common::tree_node::{
Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor,
};
use datafusion_common::{DFSchema, TableReference};
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{ColumnarValue, LogicalPlan};
use datafusion_physical_expr::PhysicalExprRef;
use datatypes::prelude::{ConcreteDataType, DataType};
@@ -830,7 +831,15 @@ fn to_phy_expr(
let phy_planner = DefaultPhysicalPlanner::default();
let phy_expr: PhysicalExprRef = phy_planner
.create_physical_expr(expr, df_schema, session)
// Time-window expressions are standalone scalar expressions over the input
// time column, so they cannot contain scalar subqueries or lambda variables
// that would require a plan-level physical planning context.
.create_physical_expr(
expr,
df_schema,
session,
&PhysicalPlanningContext::default(),
)
.with_context(|_e| DatafusionSnafu {
context: format!(
"Failed to create physical expression from {expr:?} using {df_schema:?}"
@@ -993,7 +1002,7 @@ mod test {
Some(Timestamp::new(0, TimeUnit::Millisecond)),
Some(Timestamp::new(300000, TimeUnit::Millisecond)),
),
"SELECT sum(numbers_with_ts.number), numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE ((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts <= CAST('1970-01-01 00:05:00' AS TIMESTAMP)))) GROUP BY numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts), bucket_name",
"SELECT sum(number), number, date_bin('5 minutes', ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE ((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts <= CAST('1970-01-01 00:05:00' AS TIMESTAMP)))) GROUP BY number, date_bin('5 minutes', ts), bucket_name",
),
// complex subquery alias
(
src/flow/src/batching_mode/utils.rs
+3 -5
View File
@@ -31,7 +31,7 @@ use datafusion_common::tree_node::{
use datafusion_common::{
Column, DFSchema, DataFusionError, NullEquality, ScalarValue, TableReference,
};
use datafusion_expr::logical_plan::{Aggregate, TableScan};
use datafusion_expr::logical_plan::{Aggregate, TableScanBuilder};
use datafusion_expr::{
Distinct, ExprSchemable, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, Projection, and,
binary_expr, bitwise_and, bitwise_or, bitwise_xor, is_null, or, when,
@@ -657,17 +657,15 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge(
let table_provider = Arc::new(DfTableProviderAdapter::new(sink_table));
let table_source = Arc::new(DefaultTableSource::new(table_provider));
let sink_scan = LogicalPlan::TableScan(
TableScan::try_new(
TableScanBuilder::new(
TableReference::Full {
catalog: sink_table_name[0].clone().into(),
schema: sink_table_name[1].clone().into(),
table: sink_table_name[2].clone().into(),
},
table_source,
None,
vec![],
None,
)
.build()
.with_context(|_| DatafusionSnafu {
context: "Failed to build sink table scan for incremental sink merge".to_string(),
})?,
src/flow/src/batching_mode/utils/test.rs
+4 -6
View File
@@ -18,7 +18,7 @@ use catalog::RegisterTableRequest;
use common_recordbatch::RecordBatch;
use common_time::Timestamp;
use datafusion_common::tree_node::TreeNode as _;
use datafusion_expr::GroupingSet;
use datafusion_expr::{GroupingSet, TableScanBuilder};
use datatypes::prelude::{ConcreteDataType, MutableVector, Scalar, ScalarVectorBuilder, VectorRef};
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::timestamp::TimestampMillisecond;
@@ -92,17 +92,15 @@ fn test_sink_scan(sink_table: TableRef, sink_table_name: &TableName) -> LogicalP
let table_provider = Arc::new(DfTableProviderAdapter::new(sink_table));
let table_source = Arc::new(DefaultTableSource::new(table_provider));
LogicalPlan::TableScan(
TableScan::try_new(
TableScanBuilder::new(
TableReference::Full {
catalog: sink_table_name[0].clone().into(),
schema: sink_table_name[1].clone().into(),
table: sink_table_name[2].clone().into(),
},
table_source,
None,
vec![],
None,
)
.build()
.unwrap(),
)
}
@@ -296,7 +294,7 @@ async fn test_add_filter() {
// complex subquery without alias
(
"SELECT sum(number), number, date_bin('5 minutes', ts) as time_window, bucket_name FROM (SELECT number, ts, case when number < 5 THEN 'bucket_0_5' when number >= 5 THEN 'bucket_5_inf' END as bucket_name FROM numbers_with_ts) GROUP BY number, time_window, bucket_name;",
"SELECT sum(numbers_with_ts.number), numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE (number > 4)) GROUP BY numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts), bucket_name",
"SELECT sum(number), number, date_bin('5 minutes', ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE (number > 4)) GROUP BY number, date_bin('5 minutes', ts), bucket_name",
),
// complex subquery alias
(
src/flow/src/transform/aggr.rs
+1 -1
View File
@@ -725,7 +725,7 @@ mod test {
df_scalar_fn: DfScalarFunction::try_from_raw_fn(
RawDfScalarFn {
f: BytesMut::from(
b"\x08\x02\"\x0f\x1a\r\n\x0b\xa2\x02\x08\n\0\x12\x04\x10\x1e \t\"\n\x1a\x08\x12\x06\n\x04\x12\x02\x08\x01".as_ref(),
b"\x08\x02\x1a\x07\x8a\x02\x04\x08\x03\x18\x01\"\x0f\x1a\r\n\x0b\xa2\x02\x08\n\0\x12\x04\x10\x1e \t\"\n\x1a\x08\x12\x06\n\x04\x12\x02\x08\x01".as_ref(),
),
input_schema: RelationType::new(vec![ColumnType::new(
ConcreteDataType::interval_month_day_nano_datatype(),
src/flow/src/transform/expr.rs
+10 -5
View File
@@ -20,6 +20,7 @@ use common_error::ext::BoxedError;
use common_telemetry::debug;
use datafusion::execution::SessionStateBuilder;
use datafusion::functions::all_default_functions;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_physical_expr::PhysicalExpr;
use datafusion_substrait::logical_plan::consumer::DefaultSubstraitConsumer;
use datatypes::data_type::ConcreteDataType as CDT;
@@ -101,11 +102,15 @@ pub(crate) async fn from_scalar_fn_to_df_fn_impl(
context: "Failed to convert substrait scalar function to datafusion scalar function",
}
})?;
let phy_expr =
datafusion::physical_expr::create_physical_expr(&expr, &schema, &Default::default())
.context(DatafusionSnafu {
context: "Failed to create physical expression from logical expression",
})?;
let phy_expr = datafusion::physical_expr::create_physical_expr(
&expr,
&schema,
&Default::default(),
&PhysicalPlanningContext::default(),
)
.context(DatafusionSnafu {
context: "Failed to create physical expression from logical expression",
})?;
Ok(phy_expr)
}
src/frontend/src/instance/dashboard.rs
+1 -1
View File
@@ -29,9 +29,9 @@ use common_query::OutputData;
use common_recordbatch::util as record_util;
use common_telemetry::info;
use common_time::FOREVER;
use datafusion::common::TableReference;
use datafusion::datasource::DefaultTableSource;
use datafusion::logical_expr::col;
use datafusion::sql::TableReference;
use datafusion_expr::{DmlStatement, LogicalPlan, lit};
use datatypes::arrow::array::{Array, AsArray};
use servers::error::{
src/mito2/src/engine/basic_test.rs
+1
View File
@@ -1341,6 +1341,7 @@ async fn test_all_index_metas_list_all_types_with_format(flat_format: bool, expe
if let Some(inverted) = value.get_mut("inverted").and_then(|v| v.as_object_mut()) {
inverted.insert("base_offset".to_string(), serde_json::Value::from(0));
}
value.sort_all_objects();
*meta_json = value.to_string();
}
}
src/mito2/src/engine/flush_test.rs
+4 -2
View File
@@ -1057,8 +1057,10 @@ async fn test_flush_empty_with_format(flat_format: bool) {
let stream = scanner.scan().await.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
let expected = "\
++
++";
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
src/mito2/src/engine/scan_test.rs
+4 -2
View File
@@ -1053,8 +1053,10 @@ async fn test_scan_with_min_sst_sequence_with_format(flat_format: bool) {
Some(9),
0,
"\
++
++",
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+",
)
.await;
}
src/mito2/src/engine/sync_test.rs
+10 -2
View File
@@ -148,7 +148,11 @@ async fn test_sync_after_flush_region_with_format(flat_format: bool) {
common_telemetry::info!("Scan the region on the follower engine");
// Scan the region on the follower engine
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
scan_check(&follower_engine, region_id, expected, 0, 0).await;
// Returns error since the max manifest is 1
@@ -262,7 +266,11 @@ async fn test_sync_after_alter_region_with_format(flat_format: bool) {
+-------+-------+---------+---------------------+";
scan_check(&engine, region_id, expected, 0, 1).await;
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
scan_check(&follower_engine, region_id, expected, 0, 0).await;
// Sync the region from the leader engine to the follower engine
src/mito2/src/engine/truncate_test.rs
+10 -2
View File
@@ -155,7 +155,11 @@ async fn test_engine_truncate_region_basic_with_format(flat_format: bool) {
let request = ScanRequest::default();
let stream = engine.scan_to_stream(region_id, request).await.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
@@ -401,7 +405,11 @@ async fn test_engine_truncate_reopen_with_format(flat_format: bool) {
let request = ScanRequest::default();
let stream = engine.scan_to_stream(region_id, request).await.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
src/mito2/src/memtable/bulk/part.rs
+1 -1
View File
@@ -1566,7 +1566,7 @@ impl PruningStatistics for BatchPruningStats<'_> {
None
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
None
}
src/mito2/src/read/scan_region.rs
+1 -1
View File
@@ -1817,7 +1817,7 @@ impl PruningStatistics for FileLevelPruningStats {
}
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
None
}
src/mito2/src/sst.rs
+1 -3
View File
@@ -1079,9 +1079,7 @@ mod tests {
&builder.parquet_schema().root_schema().get_fields()[0].get_fields()[0];
assert_eq!(
parquet_remainder.get_basic_info().logical_type_ref(),
Some(&LogicalType::Variant {
specification_version: None,
})
Some(&LogicalType::variant(None))
);
let ArrowDataType::Struct(children) = builder.schema().field_with_name("data")?.data_type()
src/mito2/src/sst/parquet/index_reader.rs
+1 -1
View File
@@ -172,7 +172,7 @@ impl PruningStatistics for IndexRowGroupPruningStats<'_> {
column_null_counts(self.row_groups, column_index)
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
None
}
src/mito2/src/sst/parquet/json_align/stream.rs
+1 -1
View File
@@ -247,7 +247,7 @@ fn align_array(
return Ok(array.clone());
}
cast_column(array, field.as_ref(), &DEFAULT_CAST_OPTIONS).context(CastColumnSnafu)
cast_column(array, field.data_type(), &DEFAULT_CAST_OPTIONS).context(CastColumnSnafu)
}
#[cfg(test)]
src/mito2/src/sst/parquet/read_columns.rs
+3 -3
View File
@@ -386,7 +386,7 @@ fn is_variant_leaf(leaf_col: &ColumnDescriptor) -> bool {
mod tests {
use std::sync::Arc;
use parquet::basic::{ConvertedType, LogicalType, Repetition};
use parquet::basic::{ConvertedType, LogicalType, Repetition, VariantType};
use parquet::errors::ParquetError;
use parquet::schema::types::Type;
@@ -816,9 +816,9 @@ mod tests {
let remainder = Arc::new(
Type::group_type_builder(JSON2_REMAINDER_FIELD_NAME)
.with_repetition(Repetition::OPTIONAL)
.with_logical_type(Some(LogicalType::Variant {
.with_logical_type(Some(LogicalType::Variant(VariantType {
specification_version: None,
}))
})))
.with_fields(vec![metadata, value])
.build()?,
);
src/mito2/src/sst/parquet/reader.rs
-9
View File
@@ -2556,7 +2556,6 @@ impl FlatRowGroupReader {
#[cfg(test)]
mod tests {
use std::any::Any;
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::sync::{Arc, LazyLock};
@@ -3048,10 +3047,6 @@ mod tests {
}
impl ScalarUDFImpl for PanicDebugUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"panic_debug_udf"
}
@@ -3150,10 +3145,6 @@ mod tests {
}
impl ScalarUDFImpl for TestVolatilityUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
&self.name
}
src/mito2/src/sst/parquet/stats.rs
+1 -1
View File
@@ -174,7 +174,7 @@ impl<T: Borrow<RowGroupMetaData>> PruningStatistics for RowGroupPruningStats<'_,
}
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
// TODO(LFC): Impl it.
None
}
src/partition/src/expr.rs
+8 -2
View File
@@ -20,6 +20,7 @@ use api::v1::meta::Partition;
use datafusion_common::{ScalarValue, ToDFSchema};
use datafusion_expr::Expr;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_physical_expr::{PhysicalExpr, create_physical_expr};
use datatypes::arrow;
use datatypes::value::{
@@ -391,8 +392,13 @@ impl PartitionExpr {
.context(error::ToDFSchemaSnafu)?;
let execution_props = &ExecutionProps::default();
let expr = self.try_as_logical_expr()?;
create_physical_expr(&expr, &df_schema, execution_props)
.context(error::CreatePhysicalExprSnafu)
create_physical_expr(
&expr,
&df_schema,
execution_props,
&PhysicalPlanningContext::default(),
)
.context(error::CreatePhysicalExprSnafu)
}
pub fn and(self, rhs: PartitionExpr) -> PartitionExpr {
src/promql/src/extension_plan/absent.rs
+10 -7
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::pin::Pin;
@@ -20,6 +19,7 @@ use std::sync::Arc;
use std::task::{Context, Poll};
use datafusion::arrow::array::Array;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchemaRef, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore};
@@ -30,8 +30,8 @@ use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
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, PlanProperties,
RecordBatchStream, SendableRecordBatchStream,
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements,
Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
};
use datafusion_common::DFSchema;
use datafusion_expr::{EmptyRelation, col};
@@ -325,8 +325,11 @@ pub struct AbsentExec {
}
impl ExecutionPlan for AbsentExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn schema(&self) -> SchemaRef {
@@ -337,8 +340,8 @@ impl ExecutionPlan for AbsentExec {
&self.properties
}
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![Distribution::SinglePartition]
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
InputDistributionRequirements::new(vec![Distribution::SinglePartition])
}
fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
src/promql/src/extension_plan/empty_metric.rs
+38 -14
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::collections::HashMap;
use std::ops::Div;
use std::pin::Pin;
@@ -21,21 +20,24 @@ use std::task::{Context, Poll};
use datafusion::arrow::array::ArrayRef;
use datafusion::arrow::datatypes::{DataType, TimeUnit};
use datafusion::catalog::Session;
use datafusion::common::arrow::datatypes::Field;
use datafusion::common::stats::Precision;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{
DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics, TableReference,
};
use datafusion::datasource::{MemTable, provider_as_source};
use datafusion::error::DataFusionError;
use datafusion::execution::context::{SessionState, TaskContext};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::logical_expr::{ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::{EquivalenceProperties, PhysicalExprRef};
use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr, PhysicalExprRef};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
SendableRecordBatchStream, StatisticsArgs,
};
use datafusion::physical_planner::PhysicalPlanner;
use datafusion::prelude::{Expr, col, lit};
@@ -113,14 +115,20 @@ impl EmptyMetric {
pub fn to_execution_plan(
&self,
session_state: &SessionState,
session: &dyn Session,
physical_planner: &dyn PhysicalPlanner,
planning_ctx: &PhysicalPlanningContext,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let physical_expr = self
.expr
.as_ref()
.map(|expr| {
physical_planner.create_physical_expr(expr, &self.time_index_schema, session_state)
physical_planner.create_physical_expr(
expr,
&self.time_index_schema,
session,
planning_ctx,
)
})
.transpose()?;
let result_schema: SchemaRef = self.result_schema.inner().clone();
@@ -224,8 +232,11 @@ pub struct EmptyMetricExec {
}
impl ExecutionPlan for EmptyMetricExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
datafusion::physical_plan::apply_expression_roots(self.expr.iter(), f)
}
fn schema(&self) -> SchemaRef {
@@ -273,9 +284,14 @@ impl ExecutionPlan for EmptyMetricExec {
Some(self.metric.clone_inner())
}
fn partition_statistics(&self, partition: Option<usize>) -> DataFusionResult<Statistics> {
fn statistics_from_inputs(
&self,
_input_stats: &[Arc<Statistics>],
args: &StatisticsArgs,
) -> DataFusionResult<Arc<Statistics>> {
let partition = args.partition();
if partition.is_some() {
return Ok(Statistics::new_unknown(self.schema().as_ref()));
return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref())));
}
let estimated_row_num = if self.end > self.start {
@@ -285,11 +301,11 @@ impl ExecutionPlan for EmptyMetricExec {
};
let total_byte_size = estimated_row_num * std::mem::size_of::<Millisecond>() as f64;
Ok(Statistics {
Ok(Arc::new(Statistics {
num_rows: Precision::Inexact(estimated_row_num.floor() as _),
total_byte_size: Precision::Inexact(total_byte_size.floor() as _),
column_statistics: Statistics::unknown_column(&self.schema()),
})
}))
}
fn name(&self) -> &str {
@@ -429,7 +445,11 @@ mod test {
)
.unwrap();
let empty_metric_exec = empty_metric
.to_execution_plan(&session_context.state(), &df_default_physical_planner)
.to_execution_plan(
&session_context.state(),
&df_default_physical_planner,
&PhysicalPlanningContext::default(),
)
.unwrap();
let result =
@@ -544,7 +564,11 @@ mod test {
let empty_metric =
EmptyMetric::new(0, 200, 1000, "time".to_string(), "value".to_string(), None).unwrap();
let empty_metric_exec = empty_metric
.to_execution_plan(&session_context.state(), &df_default_physical_planner)
.to_execution_plan(
&session_context.state(),
&df_default_physical_planner,
&PhysicalPlanningContext::default(),
)
.unwrap();
let result =
src/promql/src/extension_plan/histogram_fold.rs
+22 -9
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
@@ -25,6 +24,7 @@ use datafusion::arrow::compute::{SortOptions, concat_batches};
use datafusion::arrow::datatypes::{DataType, Float64Type, SchemaRef};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::stats::Precision;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchema, DFSchemaRef, Statistics};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::TaskContext;
@@ -37,7 +37,8 @@ use datafusion::physical_plan::expressions::{Column as PhyColumn, TryCastExpr as
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
InputDistributionRequirements, Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, StatisticsArgs,
};
use datafusion::prelude::{Column, Expr};
use datafusion_expr::{EmptyRelation, col};
@@ -526,8 +527,14 @@ pub struct HistogramFoldExec {
}
impl ExecutionPlan for HistogramFoldExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
datafusion::physical_plan::apply_expression_roots(
self.tag_columns.iter().chain(self.partition_exprs.iter()),
f,
)
}
fn properties(&self) -> &Arc<PlanProperties> {
@@ -574,8 +581,10 @@ impl ExecutionPlan for HistogramFoldExec {
vec![Some(OrderingRequirements::Hard(vec![requirement]))]
}
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![Distribution::HashPartitioned(self.partition_exprs.clone())]
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
self.partition_exprs.clone(),
)])
}
fn maintains_input_order(&self) -> Vec<bool> {
@@ -664,12 +673,16 @@ impl ExecutionPlan for HistogramFoldExec {
Some(self.metric.clone_inner())
}
fn partition_statistics(&self, _: Option<usize>) -> DataFusionResult<Statistics> {
Ok(Statistics {
fn statistics_from_inputs(
&self,
_input_stats: &[Arc<Statistics>],
_args: &StatisticsArgs,
) -> DataFusionResult<Arc<Statistics>> {
Ok(Arc::new(Statistics {
num_rows: Precision::Absent,
total_byte_size: Precision::Absent,
column_statistics: Statistics::unknown_column(&self.schema()),
})
}))
}
fn name(&self) -> &str {
src/promql/src/extension_plan/instant_manipulate.rs
+28 -12
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -22,6 +21,7 @@ use datafusion::arrow::array::{Array, TimestampMillisecondArray, UInt64Array};
use datafusion::arrow::datatypes::{DataType, SchemaRef};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::stats::Precision;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchema, DFSchemaRef, ScalarValue};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
@@ -33,8 +33,9 @@ use datafusion::physical_plan::metrics::{
BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, Statistics,
ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, InputDistributionRequirements,
PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics,
StatisticsArgs,
};
use datafusion_expr::col;
use datatypes::arrow::compute;
@@ -451,8 +452,11 @@ pub struct InstantManipulateExec {
}
impl ExecutionPlan for InstantManipulateExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn schema(&self) -> SchemaRef {
@@ -463,8 +467,8 @@ impl ExecutionPlan for InstantManipulateExec {
&self.properties
}
fn required_input_distribution(&self) -> Vec<Distribution> {
self.input.required_input_distribution()
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
self.input.input_distribution_requirements()
}
// Prevent reordering of input
@@ -557,8 +561,16 @@ impl ExecutionPlan for InstantManipulateExec {
Some(self.metric.clone_inner())
}
fn partition_statistics(&self, partition: Option<usize>) -> DataFusionResult<Statistics> {
let input_stats = self.input.partition_statistics(partition)?;
fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
vec![ChildStats::At(partition)]
}
fn statistics_from_inputs(
&self,
input_stats: &[Arc<Statistics>],
_args: &StatisticsArgs,
) -> DataFusionResult<Arc<Statistics>> {
let input_stats = &input_stats[0];
let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
let estimated_total_bytes = input_stats
@@ -570,12 +582,12 @@ impl ExecutionPlan for InstantManipulateExec {
})
.unwrap_or(Precision::Absent);
Ok(Statistics {
Ok(Arc::new(Statistics {
num_rows: Precision::Inexact(estimated_row_num.floor() as _),
total_byte_size: estimated_total_bytes,
// TODO(ruihang): support this column statistics
column_statistics: Statistics::unknown_column(&self.schema()),
})
}))
}
fn name(&self) -> &str {
@@ -818,6 +830,7 @@ mod test {
use datafusion::logical_expr::{
EmptyRelation, Extension, LogicalPlan, Projection, UserDefinedLogicalNodeCore,
};
use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
use datafusion::prelude::SessionContext;
use datafusion_expr::col;
@@ -1102,7 +1115,10 @@ mod test {
)));
let exec = rebuilt
.to_execution_plan(empty_exec_input)
.with_new_children(vec![exec_input])
.replace_children(
vec![exec_input],
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
.unwrap();
let output =
datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx())
src/promql/src/extension_plan/normalize.rs
+23 -11
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -20,6 +19,7 @@ use std::task::{Context, Poll};
use common_query::native_histogram::{START_TIMESTAMP_FIELD, native_histogram_arrow_type};
use datafusion::arrow::array::{Array, BooleanArray, StructArray};
use datafusion::arrow::compute;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{Column, DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics};
use datafusion::error::DataFusionError;
use datafusion::execution::context::TaskContext;
@@ -29,8 +29,9 @@ use datafusion::physical_plan::metrics::{
BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
ChildStats, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, StatisticsArgs,
};
use datafusion_expr::col;
use datatypes::arrow::array::TimestampMillisecondArray;
@@ -279,27 +280,30 @@ pub struct SeriesNormalizeExec {
}
impl ExecutionPlan for SeriesNormalizeExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn schema(&self) -> SchemaRef {
self.input.schema()
}
fn required_input_distribution(&self) -> Vec<Distribution> {
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
if self.tag_columns.is_empty() {
return vec![Distribution::SinglePartition];
return InputDistributionRequirements::new(vec![Distribution::SinglePartition]);
}
let schema = self.input.schema();
vec![Distribution::HashPartitioned(
InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
self.tag_columns
.iter()
// Safety: the tag column names is verified in the planning phase
.map(|tag| Arc::new(ColumnExpr::new_with_schema(tag, &schema).unwrap()) as _)
.collect(),
)]
)])
}
fn properties(&self) -> &Arc<PlanProperties> {
@@ -356,8 +360,16 @@ impl ExecutionPlan for SeriesNormalizeExec {
Some(self.metric.clone_inner())
}
fn partition_statistics(&self, partition: Option<usize>) -> DataFusionResult<Statistics> {
self.input.partition_statistics(partition)
fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
vec![ChildStats::At(partition)]
}
fn statistics_from_inputs(
&self,
input_stats: &[Arc<Statistics>],
_args: &StatisticsArgs,
) -> DataFusionResult<Arc<Statistics>> {
Ok(Arc::clone(&input_stats[0]))
}
fn name(&self) -> &str {
src/promql/src/extension_plan/planner.rs
+9 -3
View File
@@ -15,8 +15,9 @@
use std::sync::Arc;
use async_trait::async_trait;
use datafusion::catalog::Session;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::SessionState;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
@@ -36,7 +37,8 @@ impl ExtensionPlanner for PromExtensionPlanner {
node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
physical_inputs: &[Arc<dyn ExecutionPlan>],
session_state: &SessionState,
session: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
) -> DfResult<Option<Arc<dyn ExecutionPlan>>> {
if let Some(node) = node.as_any().downcast_ref::<SeriesNormalize>() {
Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
@@ -47,7 +49,11 @@ impl ExtensionPlanner for PromExtensionPlanner {
} else if let Some(node) = node.as_any().downcast_ref::<SeriesDivide>() {
Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
} else if let Some(node) = node.as_any().downcast_ref::<EmptyMetric>() {
Ok(Some(node.to_execution_plan(session_state, planner)?))
Ok(Some(node.to_execution_plan(
session,
planner,
planning_ctx,
)?))
} else if let Some(node) = node.as_any().downcast_ref::<ScalarCalculate>() {
Ok(Some(node.to_execution_plan(physical_inputs[0].clone())?))
} else if let Some(node) = node.as_any().downcast_ref::<HistogramFold>() {
src/promql/src/extension_plan/range_manipulate.rs
+36 -18
View File
@@ -12,8 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::collections::{HashMap, HashSet};
use std::collections::HashSet;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -25,7 +24,8 @@ use datafusion::arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit};
use datafusion::arrow::error::ArrowError;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::stats::Precision;
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchema, DFSchemaRef, TableReference};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
@@ -34,10 +34,10 @@ use datafusion::physical_plan::metrics::{
BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, Statistics,
ChildStats, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, Statistics, StatisticsArgs,
};
use datafusion::sql::TableReference;
use datafusion_expr::col;
use datatypes::timestamp::timestamp_array_to_primitive;
use futures::{Stream, StreamExt, ready};
@@ -181,7 +181,7 @@ impl RangeManipulate {
Ok(Arc::new(DFSchema::new_with_metadata(
new_columns,
HashMap::new(),
input_schema.metadata().clone(),
)?))
}
@@ -452,8 +452,11 @@ pub struct RangeManipulateExec {
}
impl ExecutionPlan for RangeManipulateExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn schema(&self) -> SchemaRef {
@@ -472,14 +475,17 @@ impl ExecutionPlan for RangeManipulateExec {
vec![&self.input]
}
fn required_input_distribution(&self) -> Vec<Distribution> {
let input_requirement = self.input.required_input_distribution();
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
let input_requirement = self
.input
.input_distribution_requirements()
.into_per_child();
if input_requirement.is_empty() {
// if the input is EmptyMetric, its required_input_distribution() is empty so we can't
// use its input distribution.
vec![Distribution::UnspecifiedDistribution]
InputDistributionRequirements::new(vec![Distribution::UnspecifiedDistribution])
} else {
input_requirement
InputDistributionRequirements::new(input_requirement)
}
}
@@ -567,8 +573,16 @@ impl ExecutionPlan for RangeManipulateExec {
Some(self.metric.clone_inner())
}
fn partition_statistics(&self, partition: Option<usize>) -> DataFusionResult<Statistics> {
let input_stats = self.input.partition_statistics(partition)?;
fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
vec![ChildStats::At(partition)]
}
fn statistics_from_inputs(
&self,
input_stats: &[Arc<Statistics>],
_args: &StatisticsArgs,
) -> DataFusionResult<Arc<Statistics>> {
let input_stats = &input_stats[0];
let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
let estimated_total_bytes = input_stats
@@ -580,12 +594,12 @@ impl ExecutionPlan for RangeManipulateExec {
})
.unwrap_or_default();
Ok(Statistics {
Ok(Arc::new(Statistics {
num_rows: Precision::Inexact(estimated_row_num as _),
total_byte_size: estimated_total_bytes,
// TODO(ruihang): support this column statistics
column_statistics: Statistics::unknown_column(&self.schema()),
})
}))
}
fn name(&self) -> &str {
@@ -836,6 +850,7 @@ mod test {
use datafusion::physical_expr::Partitioning;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::memory::MemoryStream;
use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
use datafusion::prelude::SessionContext;
use datatypes::arrow::array::TimestampMillisecondArray;
use futures::FutureExt;
@@ -1228,7 +1243,10 @@ mod test {
)));
let exec = rebuilt
.to_execution_plan(empty_exec_input)
.with_new_children(vec![exec_input])
.replace_children(
vec![exec_input],
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
.unwrap();
let output =
datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx())
src/promql/src/extension_plan/scalar_calculate.rs
+34 -15
View File
@@ -12,25 +12,27 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use datafusion::common::stats::Precision;
use datafusion::common::{DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics};
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{
DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics, TableReference,
};
use datafusion::error::DataFusionError;
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties,
RecordBatchStream, SendableRecordBatchStream,
ChildStats, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
InputDistributionRequirements, Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, StatisticsArgs,
};
use datafusion::prelude::Expr;
use datafusion::sql::TableReference;
use datafusion_expr::col;
use datatypes::arrow::array::{Array, ArrayRef, Float64Array, TimestampMillisecondArray};
use datatypes::arrow::compute::{CastOptions, cast_with_options, concat_batches};
@@ -129,7 +131,10 @@ impl ScalarCalculate {
.output_schema
.fields()
.iter()
.map(|field| Field::new(field.name(), field.data_type().clone(), field.is_nullable()))
.map(|field| {
Field::new(field.name(), field.data_type().clone(), field.is_nullable())
.with_metadata(field.metadata().clone())
})
.collect();
let input_schema = exec_input.schema();
let ts_index = input_schema
@@ -138,7 +143,10 @@ impl ScalarCalculate {
let val_index = input_schema
.index_of(&self.field_column)
.map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
let schema = Arc::new(Schema::new(fields));
let schema = Arc::new(Schema::new_with_metadata(
fields,
input_schema.metadata().clone(),
));
let properties = exec_input.properties();
let properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
@@ -389,8 +397,11 @@ struct ScalarCalculateExec {
}
impl ExecutionPlan for ScalarCalculateExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn schema(&self) -> SchemaRef {
@@ -405,8 +416,8 @@ impl ExecutionPlan for ScalarCalculateExec {
vec![true; self.children().len()]
}
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![Distribution::SinglePartition]
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
InputDistributionRequirements::new(vec![Distribution::SinglePartition])
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
@@ -469,8 +480,16 @@ impl ExecutionPlan for ScalarCalculateExec {
Some(self.metric.clone_inner())
}
fn partition_statistics(&self, partition: Option<usize>) -> DataFusionResult<Statistics> {
let input_stats = self.input.partition_statistics(partition)?;
fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
vec![ChildStats::At(partition)]
}
fn statistics_from_inputs(
&self,
input_stats: &[Arc<Statistics>],
_args: &StatisticsArgs,
) -> DataFusionResult<Arc<Statistics>> {
let input_stats = &input_stats[0];
let estimated_row_num = (self.end - self.start) as f64 / self.interval as f64;
let estimated_total_bytes = input_stats
@@ -482,12 +501,12 @@ impl ExecutionPlan for ScalarCalculateExec {
})
.unwrap_or_default();
Ok(Statistics {
Ok(Arc::new(Statistics {
num_rows: Precision::Inexact(estimated_row_num as _),
total_byte_size: estimated_total_bytes,
// TODO(ruihang): support this column statistics
column_statistics: Statistics::unknown_column(&self.schema()),
})
}))
}
fn name(&self) -> &str {
src/promql/src/extension_plan/series_divide.rs
+12 -9
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -20,6 +19,7 @@ use std::task::{Context, Poll};
use datafusion::arrow::array::{Array, ArrayRef, UInt64Array};
use datafusion::arrow::datatypes::{DataType, SchemaRef};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::error::Result as DataFusionResult;
use datafusion::execution::context::TaskContext;
@@ -30,8 +30,8 @@ use datafusion::physical_plan::metrics::{
BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements,
PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
};
use datafusion_expr::col;
use datatypes::arrow::compute;
@@ -334,8 +334,11 @@ pub struct SeriesDivideExec {
}
impl ExecutionPlan for SeriesDivideExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn schema(&self) -> SchemaRef {
@@ -346,18 +349,18 @@ impl ExecutionPlan for SeriesDivideExec {
self.input.properties()
}
fn required_input_distribution(&self) -> Vec<Distribution> {
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
if self.tag_columns.is_empty() {
return vec![Distribution::SinglePartition];
return InputDistributionRequirements::new(vec![Distribution::SinglePartition]);
}
let schema = self.input.schema();
vec![Distribution::HashPartitioned(
InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
self.tag_columns
.iter()
// Safety: the tag column names is verified in the planning phase
.map(|tag| Arc::new(ColumnExpr::new_with_schema(tag, &schema).unwrap()) as _)
.collect(),
)]
)])
}
fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
src/promql/src/extension_plan/union_distinct_on.rs
+30 -17
View File
@@ -12,15 +12,16 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use ahash::{HashSet, RandomState};
use ahash::HashSet;
use datafusion::arrow::array::UInt64Array;
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::hash_utils::RandomState as FixedState;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::context::TaskContext;
@@ -29,8 +30,9 @@ use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties,
RecordBatchStream, SendableRecordBatchStream, hash_utils,
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements,
Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
hash_utils,
};
use datafusion_expr::col;
use datatypes::arrow::compute;
@@ -179,7 +181,7 @@ impl UnionDistinctOn {
output_schema,
metric: ExecutionPlanMetricsSet::new(),
properties,
random_state: RandomState::new(),
random_state: FixedState::with_seed(0),
})
}
@@ -347,21 +349,27 @@ pub struct UnionDistinctOnExec {
metric: ExecutionPlanMetricsSet,
properties: Arc<PlanProperties>,
/// Shared the `RandomState` for the hashing algorithm
random_state: RandomState,
/// Shared deterministic hash state for the hashing algorithm.
random_state: FixedState,
}
impl ExecutionPlan for UnionDistinctOnExec {
fn as_any(&self) -> &dyn Any {
self
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn schema(&self) -> SchemaRef {
self.output_schema.clone()
}
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![Distribution::SinglePartition, Distribution::SinglePartition]
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
InputDistributionRequirements::new(vec![
Distribution::SinglePartition,
Distribution::SinglePartition,
])
}
fn properties(&self) -> &Arc<PlanProperties> {
@@ -452,7 +460,7 @@ pub struct UnionDistinctOnStream {
/// Include time index
compare_keys: Vec<usize>,
output_schema: SchemaRef,
random_state: RandomState,
random_state: FixedState,
lhs_signatures: HashSet<u64>,
hashes: Vec<u64>,
phase: StreamPhase,
@@ -911,12 +919,17 @@ mod test {
}
impl ExecutionPlan for TestExec {
fn name(&self) -> &str {
"TestExec"
fn apply_expressions(
&self,
_f: &mut dyn FnMut(
&Arc<dyn PhysicalExpr>,
) -> datafusion_common::Result<TreeNodeRecursion>,
) -> DataFusionResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn as_any(&self) -> &dyn Any {
self
fn name(&self) -> &str {
"TestExec"
}
fn properties(&self) -> &Arc<PlanProperties> {
@@ -965,7 +978,7 @@ mod test {
right: None,
compare_keys: vec![1, 0],
output_schema,
random_state: RandomState::new(),
random_state: FixedState::with_seed(0),
lhs_signatures: HashSet::default(),
hashes: Vec::new(),
phase: StreamPhase::Left,
src/promql/src/functions/native_histogram.rs
-9
View File
@@ -14,7 +14,6 @@
//! Native histogram PromQL helpers.
use std::any::Any;
use std::hash::{Hash, Hasher};
use std::mem::size_of;
use std::sync::Arc;
@@ -170,10 +169,6 @@ impl Hash for NativeHistogramAnnotationUdf {
}
impl ScalarUDFImpl for NativeHistogramAnnotationUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
self.name
}
@@ -1517,10 +1512,6 @@ impl Hash for MixedRangeUdf {
}
impl ScalarUDFImpl for MixedRangeUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
self.output.name()
}
src/promql/src/functions/quantile_aggr.rs
-1
View File
@@ -82,7 +82,6 @@ impl QuantileAccumulator {
}
let q = match &args.exprs[0]
.as_any()
.downcast_ref::<Literal>()
.map(|lit| lit.value())
{
src/query/Cargo.toml
+1
View File
@@ -43,6 +43,7 @@ datafusion-expr.workspace = true
datafusion-expr-common.workspace = true
datafusion-functions.workspace = true
datafusion-optimizer.workspace = true
datafusion-pg-catalog.workspace = true
datafusion-physical-expr.workspace = true
datafusion-proto.workspace = true
datafusion-sql.workspace = true
src/query/src/analyze.rs
+48 -20
View File
@@ -16,7 +16,6 @@
//!
//! The code skeleton is taken from `datafusion/physical-plan/src/analyze.rs`
use std::any::Any;
use std::fmt::Display;
use std::sync::Arc;
@@ -30,11 +29,12 @@ use datafusion::execution::TaskContext;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, accept,
ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan,
InputDistributionRequirements, PlanProperties, ReplaceChildrenOptions, accept,
};
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion_common::{DataFusionError, assert_eq_or_internal_err, internal_err};
use datafusion_physical_expr::{Distribution, EquivalenceProperties, Partitioning};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use futures::StreamExt;
use serde::Serialize;
use serde_json::{Value, json};
@@ -110,7 +110,6 @@ pub fn analyze_plan_metrics_to_json_value(
verbose: bool,
) -> serde_json::Result<Value> {
let input = plan
.as_any()
.downcast_ref::<DistAnalyzeExec>()
.map(|exec| exec.input().clone())
.unwrap_or_else(|| plan.clone());
@@ -125,7 +124,7 @@ pub fn analyze_plan_metrics_to_json_value(
}));
let _ = input.apply(|plan| {
if let Some(merge_scan) = plan.as_any().downcast_ref::<MergeScanExec>() {
if let Some(merge_scan) = plan.downcast_ref::<MergeScanExec>() {
for (node, metric) in merge_scan.sub_stage_metrics().into_iter().enumerate() {
stages.push(json!({
"stage": 1,
@@ -157,11 +156,6 @@ impl ExecutionPlan for DistAnalyzeExec {
"DistAnalyzeExec"
}
/// Return a reference to Any that can be used for downcasting
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
@@ -170,14 +164,25 @@ impl ExecutionPlan for DistAnalyzeExec {
vec![&self.input]
}
/// AnalyzeExec is handled specially so this value is ignored
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![]
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> DfResult<TreeNodeRecursion>,
) -> DfResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
/// AnalyzeExec is handled specially so this value is ignored
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
// AnalyzeExec is handled specially so this value is ignored.
InputDistributionRequirements::new(vec![
datafusion_physical_expr::Distribution::UnspecifiedDistribution,
])
}
fn replace_children(
self: Arc<Self>,
mut children: Vec<Arc<dyn ExecutionPlan>>,
_options: ReplaceChildrenOptions,
) -> DfResult<Arc<dyn ExecutionPlan>> {
assert_eq_or_internal_err!(
children.len(),
@@ -191,6 +196,17 @@ impl ExecutionPlan for DistAnalyzeExec {
)))
}
#[allow(deprecated)]
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
self.replace_children(
children,
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
}
fn execute(
&self,
partition: usize,
@@ -293,7 +309,7 @@ fn create_output_batch(
// Find merge scan and append its sub_stage_metrics
input.apply(|plan| {
if let Some(merge_scan) = plan.as_any().downcast_ref::<MergeScanExec>() {
if let Some(merge_scan) = plan.downcast_ref::<MergeScanExec>() {
let sub_stage_metrics = merge_scan.sub_stage_metrics();
for (node, metric) in sub_stage_metrics.into_iter().enumerate() {
builder.append_metric(1, node as _, metrics_to_string(metric, format)?);
@@ -421,7 +437,14 @@ mod tests {
AnalyzeFormat::TEXT,
));
assert!(ExecutionPlan::with_new_children(analyze, vec![]).is_err());
assert!(
ExecutionPlan::replace_children(
analyze,
vec![],
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
.is_err()
);
}
#[test]
@@ -432,14 +455,14 @@ mod tests {
AnalyzeFormat::TEXT,
));
let result = ExecutionPlan::with_new_children(
let result = ExecutionPlan::replace_children(
analyze,
vec![empty_plan("first"), empty_plan("second")],
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
);
if let Ok(plan) = result {
let retained = plan
.as_any()
.downcast_ref::<DistAnalyzeExec>()
.unwrap()
.input()
@@ -460,8 +483,13 @@ mod tests {
));
let replacement = empty_plan("replacement");
let rebuilt = ExecutionPlan::with_new_children(analyze, vec![replacement]).unwrap();
let rebuilt = rebuilt.as_any().downcast_ref::<DistAnalyzeExec>().unwrap();
let rebuilt = ExecutionPlan::replace_children(
analyze,
vec![replacement],
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
.unwrap();
let rebuilt = rebuilt.downcast_ref::<DistAnalyzeExec>().unwrap();
assert_eq!(rebuilt.input().schema().field(0).name(), "replacement");
}
src/query/src/datafusion.rs
+4 -4
View File
@@ -16,6 +16,7 @@
mod error;
mod json_expr_planner;
mod pg_oid_alias_expr_planner;
mod planner;
use std::any::Any;
@@ -110,7 +111,7 @@ fn query_load_region_id(plan: &Arc<dyn ExecutionPlan>) -> Option<u64> {
while let Some(plan) = stack.pop() {
if plan.name() == REGION_SCAN_EXEC_NAME
&& let Some(scan) = plan.as_any().downcast_ref::<RegionScanExec>()
&& let Some(scan) = plan.downcast_ref::<RegionScanExec>()
&& let Some(scan_region_id) = scan.query_load_region_id()
{
match region_id {
@@ -139,7 +140,7 @@ fn query_stat_counters(plan: &Arc<dyn ExecutionPlan>) -> Option<RegionQueryStatC
while let Some(plan) = stack.pop() {
if plan.name() == REGION_SCAN_EXEC_NAME
&& let Some(scan) = plan.as_any().downcast_ref::<RegionScanExec>()
&& let Some(scan) = plan.downcast_ref::<RegionScanExec>()
&& let Some(scan_counters) = scan.query_stat_counters()
{
match &counters {
@@ -519,8 +520,7 @@ impl DatafusionQueryEngine {
// let config = state.config_options();
// skip optimize AnalyzeExec plan
let optimized_plan = if let Some(analyze_plan) = plan.as_any().downcast_ref::<AnalyzeExec>()
{
let optimized_plan = if let Some(analyze_plan) = plan.downcast_ref::<AnalyzeExec>() {
let format = if let Some(format) = ctx.query_ctx().explain_format()
&& format.to_lowercase() == "json"
{
src/query/src/datafusion/pg_oid_alias_expr_planner.rs
+340
View File
@@ -0,0 +1,340 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use arrow_schema::DataType;
use datafusion_common::{DFSchema, ExprSchema, Result, ScalarValue};
use datafusion_expr::expr::BinaryExpr;
use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
use datafusion_expr::{Expr, Operator};
use datafusion_pg_catalog::pg_catalog::oid_field::{OID_ALIAS_KEY, kind};
use sqlparser::ast::BinaryOperator;
/// Rewrites PostgreSQL's regproc zero sentinel before DataFusion type coercion.
#[derive(Debug)]
pub(crate) struct PgOidAliasExprPlanner;
impl ExprPlanner for PgOidAliasExprPlanner {
fn plan_binary_op(
&self,
expr: RawBinaryExpr,
schema: &DFSchema,
) -> Result<PlannerResult<RawBinaryExpr>> {
let RawBinaryExpr {
op,
mut left,
mut right,
} = expr;
let operator = match op {
BinaryOperator::Eq => Operator::Eq,
BinaryOperator::NotEq => Operator::NotEq,
_ => return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right })),
};
let (column, zero_on_left) = match (&left, &right) {
(Expr::Literal(value, _), Expr::Column(column)) if is_integral_zero(value) => {
(column, true)
}
(Expr::Column(column), Expr::Literal(value, _)) if is_integral_zero(value) => {
(column, false)
}
_ => return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right })),
};
// A raw SQL column is resolved against the schema before the default
// coercion planner runs. Do not infer alias semantics from casts or any
// other expression shape.
let Ok(field) = schema.field_from_column(column) else {
return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
};
if field.metadata().get(OID_ALIAS_KEY).map(String::as_str) != Some(kind::REGPROC) {
return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
}
let Some(sentinel) = regproc_zero_sentinel(field.data_type()) else {
return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
};
let sentinel = Expr::Literal(sentinel, None);
if zero_on_left {
left = sentinel;
} else {
right = sentinel;
}
Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr::new(
Box::new(left),
operator,
Box::new(right),
))))
}
}
fn is_integral_zero(value: &ScalarValue) -> bool {
matches!(
value,
ScalarValue::Int8(Some(0))
| ScalarValue::Int16(Some(0))
| ScalarValue::Int32(Some(0))
| ScalarValue::Int64(Some(0))
| ScalarValue::UInt8(Some(0))
| ScalarValue::UInt16(Some(0))
| ScalarValue::UInt32(Some(0))
| ScalarValue::UInt64(Some(0))
)
}
fn regproc_zero_sentinel(data_type: &DataType) -> Option<ScalarValue> {
match data_type {
DataType::Utf8 => Some(ScalarValue::Utf8(Some("-".to_string()))),
DataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(Some("-".to_string()))),
DataType::Utf8View => Some(ScalarValue::Utf8View(Some("-".to_string()))),
_ => None,
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use arrow_schema::{Field, Fields};
use datafusion_common::Column;
use datafusion_expr::ExprSchemable;
use datafusion_expr::expr::Cast;
use datafusion_expr::simplify::SimplifyContext;
use datafusion_optimizer::simplify_expressions::ExprSimplifier;
use super::*;
fn schema(data_type: DataType, alias: Option<&str>) -> DFSchema {
let mut field = Field::new("typreceive", data_type, true);
if let Some(alias) = alias {
field = field.with_metadata(HashMap::from([(
OID_ALIAS_KEY.to_string(),
alias.to_string(),
)]));
}
DFSchema::from_unqualified_fields(Fields::from(vec![field]), HashMap::new()).unwrap()
}
fn column() -> Expr {
Expr::Column(Column::new_unqualified("typreceive"))
}
fn plan(expr: RawBinaryExpr, schema: &DFSchema) -> PlannerResult<RawBinaryExpr> {
PgOidAliasExprPlanner.plan_binary_op(expr, schema).unwrap()
}
fn assert_planned_sentinel(
planned: PlannerResult<RawBinaryExpr>,
operator: Operator,
zero_on_left: bool,
sentinel: ScalarValue,
) -> Expr {
let PlannerResult::Planned(Expr::BinaryExpr(expr)) = planned else {
panic!("expected a planned binary expression");
};
assert_eq!(expr.op, operator);
let literal = Expr::Literal(sentinel, None);
if zero_on_left {
assert_eq!(expr.left.as_ref(), &literal);
assert_eq!(expr.right.as_ref(), &column());
} else {
assert_eq!(expr.left.as_ref(), &column());
assert_eq!(expr.right.as_ref(), &literal);
}
Expr::BinaryExpr(expr)
}
#[test]
fn rewrites_zero_regproc_comparisons_in_both_operand_orders() {
let schema = schema(DataType::Utf8, Some(kind::REGPROC));
for (sql_operator, operator) in [
(BinaryOperator::Eq, Operator::Eq),
(BinaryOperator::NotEq, Operator::NotEq),
] {
for zero_on_left in [true, false] {
let zero = Expr::Literal(ScalarValue::Int64(Some(0)), None);
let (left, right) = if zero_on_left {
(zero, column())
} else {
(column(), zero)
};
let planned = assert_planned_sentinel(
plan(
RawBinaryExpr {
op: sql_operator.clone(),
left,
right,
},
&schema,
),
operator,
zero_on_left,
ScalarValue::Utf8(Some("-".to_string())),
);
assert!(planned.nullable(&schema).unwrap());
}
}
}
#[test]
fn rewrites_every_integral_zero_with_the_column_string_storage_type() {
let zero_literals = [
ScalarValue::Int8(Some(0)),
ScalarValue::Int16(Some(0)),
ScalarValue::Int32(Some(0)),
ScalarValue::Int64(Some(0)),
ScalarValue::UInt8(Some(0)),
ScalarValue::UInt16(Some(0)),
ScalarValue::UInt32(Some(0)),
ScalarValue::UInt64(Some(0)),
];
let string_types = [
(DataType::Utf8, ScalarValue::Utf8(Some("-".to_string()))),
(
DataType::LargeUtf8,
ScalarValue::LargeUtf8(Some("-".to_string())),
),
(
DataType::Utf8View,
ScalarValue::Utf8View(Some("-".to_string())),
),
];
for (data_type, sentinel) in string_types {
let schema = schema(data_type, Some(kind::REGPROC));
for zero in &zero_literals {
assert_planned_sentinel(
plan(
RawBinaryExpr {
op: BinaryOperator::Eq,
left: column(),
right: Expr::Literal(zero.clone(), None),
},
&schema,
),
Operator::Eq,
false,
sentinel.clone(),
);
}
}
}
#[test]
fn leaves_non_matching_comparisons_untouched() {
let regproc = schema(DataType::Utf8, Some(kind::REGPROC));
let int32_regproc = schema(DataType::Int32, Some(kind::REGPROC));
let untagged = schema(DataType::Utf8, None);
let regtype = schema(DataType::Utf8, Some(kind::REGTYPE));
let cases = [
(
RawBinaryExpr {
op: BinaryOperator::Eq,
left: column(),
right: Expr::Literal(ScalarValue::Int64(Some(1)), None),
},
&regproc,
),
(
RawBinaryExpr {
op: BinaryOperator::Eq,
left: column(),
right: Expr::Literal(ScalarValue::Int64(None), None),
},
&regproc,
),
(
RawBinaryExpr {
op: BinaryOperator::Lt,
left: column(),
right: Expr::Literal(ScalarValue::Int64(Some(0)), None),
},
&regproc,
),
(
RawBinaryExpr {
op: BinaryOperator::Eq,
left: column(),
right: Expr::Literal(ScalarValue::Int64(Some(0)), None),
},
&int32_regproc,
),
(
RawBinaryExpr {
op: BinaryOperator::Eq,
left: column(),
right: Expr::Literal(ScalarValue::Int64(Some(0)), None),
},
&untagged,
),
(
RawBinaryExpr {
op: BinaryOperator::Eq,
left: column(),
right: Expr::Literal(ScalarValue::Int64(Some(0)), None),
},
&regtype,
),
];
for (expr, schema) in cases {
assert!(matches!(plan(expr, schema), PlannerResult::Original(_)));
}
}
#[test]
fn leaves_casts_and_the_adbc_array_receiver_predicate_untouched() {
let schema = schema(DataType::Utf8, Some(kind::REGPROC));
let cast = Expr::Cast(Cast::new(Box::new(column()), DataType::Utf8));
let expr = RawBinaryExpr {
op: BinaryOperator::NotEq,
left: cast,
right: Expr::Literal(ScalarValue::Utf8(Some("array_recv".to_string())), None),
};
assert!(matches!(plan(expr, &schema), PlannerResult::Original(_)));
}
#[test]
fn type_coercion_keeps_regproc_as_a_string_after_the_rewrite() {
let schema = Arc::new(schema(DataType::Utf8, Some(kind::REGPROC)));
let planned = assert_planned_sentinel(
plan(
RawBinaryExpr {
op: BinaryOperator::NotEq,
left: column(),
right: Expr::Literal(ScalarValue::Int64(Some(0)), None),
},
&schema,
),
Operator::NotEq,
false,
ScalarValue::Utf8(Some("-".to_string())),
);
let simplifier = ExprSimplifier::new(
SimplifyContext::builder()
.with_schema(schema.clone())
.build(),
);
let coerced = simplifier.coerce(planned, &schema).unwrap();
assert!(!format!("{coerced}").contains("CAST"));
assert!(!format!("{coerced:?}").contains("Int64"));
}
}
src/query/src/datafusion/planner.rs
+203 -12
View File
@@ -19,7 +19,8 @@ use std::sync::Arc;
use arrow_schema::DataType;
use catalog::table_source::DfTableSourceProvider;
use common_function::function::FunctionContext;
use datafusion::common::TableReference;
use datafusion::catalog::TableFunctionArgs;
use datafusion::common::{DFSchema, TableReference};
use datafusion::datasource::cte_worktable::CteWorkTable;
use datafusion::datasource::file_format::{FileFormatFactory, format_as_file_type};
use datafusion::datasource::provider_as_source;
@@ -33,12 +34,13 @@ use datafusion_common::config::ConfigOptions;
use datafusion_common::file_options::file_type::FileType;
use datafusion_expr::planner::{ExprPlanner, TypePlanner};
use datafusion_expr::var_provider::is_system_variables;
use datafusion_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF};
use datafusion_sql::parser::Statement as DfStatement;
use session::context::QueryContextRef;
use snafu::{Location, ResultExt};
use crate::datafusion::json_expr_planner::JsonExprPlanner;
use crate::datafusion::pg_oid_alias_expr_planner::PgOidAliasExprPlanner;
use crate::error::{CatalogSnafu, Result};
use crate::query_engine::{DefaultPlanDecoder, QueryEngineState};
@@ -90,6 +92,7 @@ impl DfContextProviderAdapter {
let mut expr_planners = SessionStateDefaults::default_expr_planners();
expr_planners.insert(0, Arc::new(JsonExprPlanner));
expr_planners.insert(0, Arc::new(PgOidAliasExprPlanner));
Ok(Self {
engine_state,
@@ -161,6 +164,13 @@ impl ContextProvider for DfContextProviderAdapter {
)
}
fn get_higher_order_meta(&self, name: &str) -> Option<Arc<HigherOrderUDF>> {
self.session_state
.higher_order_functions()
.get(name)
.cloned()
}
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
self.engine_state.aggr_function(name).map_or_else(
|| self.session_state.aggregate_functions().get(name).cloned(),
@@ -200,6 +210,14 @@ impl ContextProvider for DfContextProviderAdapter {
names
}
fn higher_order_function_names(&self) -> Vec<String> {
self.session_state
.higher_order_functions()
.keys()
.cloned()
.collect()
}
fn udaf_names(&self) -> Vec<String> {
let mut names = self.engine_state.aggr_names();
names.extend(self.session_state.aggregate_functions().keys().cloned());
@@ -228,22 +246,47 @@ impl ContextProvider for DfContextProviderAdapter {
name: &str,
args: Vec<datafusion_expr::Expr>,
) -> DfResult<Arc<dyn TableSource>> {
if let Some(tbl_func) = self.engine_state.table_function(name) {
let provider = tbl_func.create_table_provider(&args)?;
Ok(provider_as_source(provider))
// Constant-fold the args before resolving the table function. DataFusion's
// SQL planner does not fold table-function arguments (constant folding
// happens later, in the analyzer), but table functions such as
// `generate_series`/`range` are resolved during planning and require
// literal bounds. Folding here lets immutable-UDF bounds like
// `array_upper(ARRAY[...], 1)` reach them as concrete literals.
// Non-constant args are returned unchanged by the simplifier.
let simplify_info = datafusion_expr::simplify::SimplifyContext::builder()
.with_config_options(Arc::clone(self.session_state.config_options()))
.with_query_execution_start_time(
self.session_state
.execution_props()
.query_execution_start_time,
)
.build();
let simplifier =
datafusion_optimizer::simplify_expressions::ExprSimplifier::new(simplify_info);
let schema = DFSchema::empty();
let args = args
.into_iter()
.map(|arg| {
simplifier
.coerce(arg, &schema)
.and_then(|arg| simplifier.simplify(arg))
})
.collect::<DfResult<Vec<_>>>()?;
let table_args = TableFunctionArgs::new(&args, &self.session_state);
let tbl_func = if let Some(tbl_func) = self.engine_state.table_function(name) {
tbl_func
} else {
let tbl_func = self
.session_state
self.session_state
.table_functions()
.get(name)
.cloned()
.ok_or_else(|| {
DataFusionError::Plan(format!("table function '{name}' not found"))
})?;
let provider = tbl_func.create_table_provider(&args)?;
})?
};
let provider = tbl_func.create_table_provider_with_args(table_args)?;
Ok(provider_as_source(provider))
}
Ok(provider_as_source(provider))
}
fn create_cte_work_table(
@@ -260,6 +303,154 @@ impl ContextProvider for DfContextProviderAdapter {
}
fn get_type_planner(&self) -> Option<Arc<dyn TypePlanner>> {
None
// Provide the SQL planner with Postgres oid-alias type names
// (`regclass`, `regproc`, `regtype`, `regnamespace`, `oid`, ...) and
// `pg_catalog.`-qualified builtins. DataFusion rejects these as
// "Unsupported SQL type" otherwise. The planner maps each to its Arrow
// type so reverse / column-operand casts like `prorettype::regtype::text`
// parse. Forward name->oid casts (`'x'::regclass`) are resolved earlier,
// at SQL-parse time, by the `PostgresCompatibilityParser`'s built-in
// `RewriteRegCastToSubquery` rule.
// Stateless, so a fresh instance per query is cheap.
Some(Arc::new(
datafusion_pg_catalog::pg_catalog::oid_type_planner::PgOidTypePlanner,
))
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use common_base::Plugins;
use datafusion::catalog::{TableFunction, TableFunctionArgs, TableFunctionImpl, TableProvider};
use datafusion::datasource::MemTable;
use datafusion::execution::SessionStateBuilder;
use datafusion::execution::context::{SessionConfig, SessionContext, SessionState};
use datafusion_common::ScalarValue;
use datafusion_expr::expr::BinaryExpr;
use datafusion_expr::{Expr, Operator, lit};
use session::context::QueryContext;
use super::*;
use crate::options::QueryOptions;
#[derive(Debug, Default)]
struct RecordingTableFunction {
called: AtomicBool,
args: Mutex<Vec<Expr>>,
target_partitions: Mutex<Option<usize>>,
}
impl TableFunctionImpl for RecordingTableFunction {
fn call_with_args(&self, args: TableFunctionArgs) -> DfResult<Arc<dyn TableProvider>> {
let session_state = args
.session()
.as_any()
.downcast_ref::<SessionState>()
.expect("table function must receive the SessionState");
*self.args.lock().unwrap() = args.exprs().to_vec();
*self.target_partitions.lock().unwrap() =
Some(session_state.config().target_partitions());
self.called.store(true, Ordering::SeqCst);
Ok(Arc::new(MemTable::try_new(
Arc::new(arrow_schema::Schema::empty()),
vec![vec![]],
)?))
}
}
fn query_engine_state() -> Arc<QueryEngineState> {
Arc::new(QueryEngineState::new(
catalog::memory::new_memory_catalog_manager().unwrap(),
None,
None,
None,
None,
None,
false,
Plugins::default(),
QueryOptions::default(),
))
}
async fn context_provider(
engine_state: Arc<QueryEngineState>,
session_state: SessionState,
) -> DfContextProviderAdapter {
DfContextProviderAdapter::try_new(engine_state, session_state, None, QueryContext::arc())
.await
.unwrap()
}
fn plus(left: Expr, right: Expr) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left),
op: Operator::Plus,
right: Box::new(right),
})
}
#[tokio::test]
async fn table_function_arguments_are_folded_before_engine_function_creation() {
let engine_state = query_engine_state();
let function = Arc::new(RecordingTableFunction::default());
engine_state.register_table_function(Arc::new(TableFunction::new(
"capture_engine_args".to_string(),
function.clone(),
)));
let provider = context_provider(engine_state.clone(), engine_state.session_state()).await;
provider
.get_table_function_source("capture_engine_args", vec![plus(lit(1_i64), lit(2_i64))])
.unwrap();
assert!(function.called.load(Ordering::SeqCst));
assert_eq!(
*function.args.lock().unwrap(),
vec![Expr::Literal(ScalarValue::Int64(Some(3)), None)]
);
}
#[tokio::test]
async fn table_function_argument_simplification_errors_are_propagated() {
let engine_state = query_engine_state();
let function = Arc::new(RecordingTableFunction::default());
engine_state.register_table_function(Arc::new(TableFunction::new(
"reject_invalid_args".to_string(),
function.clone(),
)));
let provider = context_provider(engine_state.clone(), engine_state.session_state()).await;
let error = match provider
.get_table_function_source("reject_invalid_args", vec![plus(lit(true), lit(1_i64))])
{
Ok(_) => panic!("invalid table-function argument must fail planning"),
Err(error) => error,
};
assert!(!error.to_string().is_empty());
assert!(!function.called.load(Ordering::SeqCst));
}
#[tokio::test]
async fn session_table_function_receives_table_function_args_session() {
let engine_state = query_engine_state();
let session_state = SessionStateBuilder::new_from_existing(engine_state.session_state())
.with_config(SessionConfig::new().with_target_partitions(7))
.build();
let session_context = SessionContext::new_with_state(session_state);
let function = Arc::new(RecordingTableFunction::default());
session_context.register_udtf("capture_session_args", function.clone());
let provider = context_provider(engine_state, session_context.state()).await;
provider
.get_table_function_source("capture_session_args", vec![lit(1_i64)])
.unwrap();
assert!(function.called.load(Ordering::SeqCst));
assert_eq!(*function.target_partitions.lock().unwrap(), Some(7));
}
}
src/query/src/dist_plan/analyzer.rs
+1 -5
View File
@@ -668,13 +668,9 @@ impl PlanRewriter {
}
if let LogicalPlan::TableScan(table_scan) = plan
&& let Some(source) = table_scan
.source
.as_any()
.downcast_ref::<DefaultTableSource>()
&& let Some(source) = table_scan.source.downcast_ref::<DefaultTableSource>()
&& let Some(provider) = source
.table_provider
.as_any()
.downcast_ref::<DfTableProviderAdapter>()
{
let table = provider.table();
src/query/src/dist_plan/analyzer/fallback.rs
+2 -5
View File
@@ -45,14 +45,11 @@ impl TreeNodeRewriter for FallbackPlanRewriter {
plan: Self::Node,
) -> DfResult<datafusion_common::tree_node::Transformed<Self::Node>> {
if let LogicalPlan::TableScan(table_scan) = &plan {
let partition_cols = if let Some(source) = table_scan
.source
.as_any()
.downcast_ref::<DefaultTableSource>()
let partition_cols = if let Some(source) =
table_scan.source.downcast_ref::<DefaultTableSource>()
{
if let Some(provider) = source
.table_provider
.as_any()
.downcast_ref::<DfTableProviderAdapter>()
{
if provider.table().table_type() == TableType::Base {
src/query/src/dist_plan/analyzer/test.rs
+76 -26
View File
@@ -31,15 +31,15 @@ use datafusion::functions_aggregate::min_max::{max, min};
use datafusion::functions_nested::expr_fn::make_array;
use datafusion::prelude::SessionContext;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{ExprSchema, JoinType, ScalarValue};
use datafusion_common::{ExprSchema, JoinType, ScalarValue, TableReference};
use datafusion_expr::expr::{Exists, ScalarFunction};
use datafusion_expr::utils::split_conjunction;
use datafusion_expr::{
AggregateUDF, Expr, ExprSchemable as _, Extension, LogicalPlanBuilder, Operator, Subquery,
binary_expr, col, lit,
};
use datafusion_functions::datetime::date_bin;
use datafusion_functions::datetime::expr_fn::now;
use datafusion_sql::TableReference;
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef};
use futures::Stream;
@@ -302,6 +302,53 @@ fn find_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> {
plan.inputs().into_iter().find_map(find_merge_scan)
}
fn find_table_scan<'a>(
plan: &'a LogicalPlan,
table_name: &str,
) -> Option<&'a datafusion_expr::logical_plan::TableScan> {
if let LogicalPlan::TableScan(table_scan) = plan
&& table_scan.table_name.to_string() == table_name
{
return Some(table_scan);
}
plan.inputs()
.into_iter()
.find_map(|input| find_table_scan(input, table_name))
}
fn find_merge_scan_for_table<'a>(
plan: &'a LogicalPlan,
table_name: &str,
) -> Option<&'a MergeScanLogicalPlan> {
if let LogicalPlan::Extension(extension) = plan
&& let Some(merge_scan) = extension
.node
.as_any()
.downcast_ref::<MergeScanLogicalPlan>()
&& find_table_scan(merge_scan.input(), table_name).is_some()
{
return Some(merge_scan);
}
plan.inputs()
.into_iter()
.find_map(|input| find_merge_scan_for_table(input, table_name))
}
fn has_filter_above_table_scan(plan: &LogicalPlan, table_name: &str, predicate: &Expr) -> bool {
if let LogicalPlan::Filter(filter) = plan
&& split_conjunction(&filter.predicate).contains(&predicate)
&& find_table_scan(filter.input.as_ref(), table_name).is_some()
{
return true;
}
plan.inputs()
.into_iter()
.any(|input| has_filter_above_table_scan(input, table_name, predicate))
}
#[test]
fn frontend_only_histogram_folds_stay_above_merge_scan() {
let table = TestTable::table_with_name(0, "t".to_string());
@@ -2586,34 +2633,37 @@ fn test_join_side_local_filter_pushdown_into_merge_scan() {
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
assert_remote_table_scan_filters_are_safe(&result);
let plan_str = result.to_string();
// After PushDownFilter runs, the predicate `t1.pk1 = Utf8("v")` should appear
// inside the left MergeScan's remote_input. The pre-MergeScan optimizer may
// combine it with join-derived IS NOT NULL pushdowns, so it may not appear as
// a standalone Filter: line. It must still be in TableScan partial_filters
// and below the Inner Join.
let predicate = col("t1.pk1").eq(lit("v"));
let t1_remote_input = find_merge_scan_for_table(&result, "t1")
.expect("expected MergeScan for t1")
.input();
let t1_scan = find_table_scan(t1_remote_input, "t1").expect("expected t1 TableScan");
assert!(
plan_str.contains("t1.pk1 = Utf8(\"v\")"),
"Expected predicate t1.pk1 = Utf8(\"v\") in plan, got:\n{plan_str}"
);
assert!(
plan_str.contains(
"TableScan: t1, partial_filters=[t1.pk1 = Utf8(\"v\"), t1.number IS NOT NULL]"
),
"Expected t1 TableScan partial_filters to contain pushed predicate, got:\n{plan_str}"
t1_scan
.filters
.iter()
.flat_map(|filter| split_conjunction(filter))
.any(|filter| filter == &predicate),
"expected t1 TableScan to contain the pushed predicate: {t1_remote_input}"
);
// Find the position of the filter and verify it appears after a MergeScan
// opening (i.e., inside remote_input) rather than before the Join.
let filter_pos = plan_str
.find("TableScan: t1, partial_filters=[t1.pk1 = Utf8(\"v\"), t1.number IS NOT NULL]")
.unwrap();
let join_pos = plan_str.find("Inner Join").unwrap();
// The filter should be after the Join (meaning it was pushed down below the Join,
// into a MergeScan's remote_input)
// Inexact provider pushdown must retain the predicate in an ancestor Filter.
assert!(
filter_pos > join_pos,
"Filter should be pushed below Join (into MergeScan remote_input), but found before Join"
has_filter_above_table_scan(t1_remote_input, "t1", &predicate),
"expected an ancestor Filter for t1 to retain the pushed predicate: {t1_remote_input}"
);
let t2_remote_input = find_merge_scan_for_table(&result, "t2")
.expect("expected MergeScan for t2")
.input();
assert!(
!find_table_scan(t2_remote_input, "t2")
.expect("expected t2 TableScan")
.filters
.iter()
.flat_map(|filter| split_conjunction(filter))
.any(|filter| filter == &predicate),
"t2 TableScan must not contain t1's predicate: {t2_remote_input}"
);
}
src/query/src/dist_plan/commutativity.rs
+3
View File
@@ -328,6 +328,9 @@ impl Categorizer {
| Expr::WindowFunction(_)
| Expr::InSubquery(_)
| Expr::ScalarSubquery(_)
| Expr::HigherOrderFunction(_)
| Expr::Lambda(_)
| Expr::LambdaVariable(_)
| Expr::Wildcard { .. } => Commutativity::Unimplemented,
Expr::Alias(alias) => Self::check_expr(&alias.expr),
src/query/src/dist_plan/dyn_filter_bridge.rs
+1 -5
View File
@@ -346,10 +346,6 @@ mod tests {
impl Eq for UnserializableExpr {}
impl datafusion_physical_expr::PhysicalExpr for UnserializableExpr {
fn as_any(&self) -> &dyn Any {
self
}
fn data_type(
&self,
_input_schema: &arrow_schema::Schema,
@@ -766,7 +762,7 @@ mod tests {
captured_dyn_filters[0].filter_id.to_string()
);
assert_eq!(decoded_children.len(), 1);
assert!(decoded_children[0].as_any().is::<Column>());
assert!(decoded_children[0].is::<Column>());
}
#[test]
src/query/src/dist_plan/merge_scan.rs
+74 -39
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
#[cfg(test)]
use std::cell::Cell;
use std::sync::{Arc, Mutex};
@@ -40,9 +39,10 @@ use datafusion::physical_plan::metrics::{
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
SendableRecordBatchStream,
SendableRecordBatchStream, apply_expression_roots,
};
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{Column as ColumnExpr, DFSchemaRef, DataFusionError, Result, Statistics};
use datafusion_expr::{Expr, Extension, FetchType, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion_physical_expr::expressions::Column;
@@ -911,7 +911,7 @@ impl MergeScanExec {
}
pub fn try_with_new_distribution(&self, distribution: Distribution) -> Option<Self> {
let Distribution::HashPartitioned(hash_exprs) = distribution else {
let Distribution::KeyPartitioned(hash_exprs) = distribution else {
// not applicable
return None;
};
@@ -926,8 +926,7 @@ impl MergeScanExec {
let hash_expr_col_names: HashSet<_> = hash_exprs
.iter()
.filter_map(|expr| {
expr.as_any()
.downcast_ref::<Column>()
expr.downcast_ref::<Column>()
.map(|col_expr| col_expr.name())
})
.collect();
@@ -949,8 +948,7 @@ impl MergeScanExec {
let overlaps: Vec<_> = hash_exprs
.iter()
.filter(|expr| {
expr.as_any()
.downcast_ref::<Column>()
expr.downcast_ref::<Column>()
.is_some_and(|col_expr| all_partition_col_aliases.contains(col_expr.name()))
})
.cloned()
@@ -1153,10 +1151,6 @@ impl Drop for PartitionMetrics {
}
impl ExecutionPlan for MergeScanExec {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> ArrowSchemaRef {
self.arrow_schema.clone()
}
@@ -1169,6 +1163,24 @@ impl ExecutionPlan for MergeScanExec {
vec![]
}
fn apply_expressions(
&self,
f: &mut dyn FnMut(
&Arc<dyn datafusion_physical_expr::PhysicalExpr>,
) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
let captured_remote_dyn_filters = self.captured_remote_dyn_filters();
apply_expression_roots(
captured_remote_dyn_filters
.into_iter()
.map(|captured_dyn_filter| {
captured_dyn_filter.alive_dyn_filter
as Arc<dyn datafusion_physical_expr::PhysicalExpr>
}),
f,
)
}
// DataFusion will swap children unconditionally.
// But since this node is leaf node, it's safe to just return self.
fn with_new_children(
@@ -1250,14 +1262,14 @@ impl ExecutionPlan for MergeScanExec {
Some(self.metric.clone_inner())
}
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
if partition.is_some() {
return Ok(Statistics::new_unknown(&self.arrow_schema));
return Ok(Arc::new(Statistics::new_unknown(&self.arrow_schema)));
}
let mut statistics = Statistics::new_unknown(&self.arrow_schema);
statistics.num_rows = self.estimated_num_rows();
Ok(statistics)
Ok(Arc::new(statistics))
}
fn name(&self) -> &str {
@@ -1422,12 +1434,13 @@ mod tests {
use datafusion::config::ConfigOptions;
use datafusion::execution::SessionStateBuilder;
use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult;
use datafusion::physical_plan::{StatisticsArgs, StatisticsContext};
use datafusion_common::TableReference;
use datafusion_expr::{LogicalPlanBuilder, col, lit};
use datafusion_physical_expr::Distribution;
use datafusion_physical_expr::expressions::{
Column, DynamicFilterPhysicalExpr, lit as physical_lit,
};
use datafusion_physical_expr::{Distribution, PhysicalExpr};
use datatypes::prelude::{ConcreteDataType, VectorRef};
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector};
@@ -1501,6 +1514,12 @@ mod tests {
.unwrap()
}
fn merge_scan_statistics(exec: &MergeScanExec) -> Arc<Statistics> {
StatisticsContext::new()
.compute(exec, &StatisticsArgs::new())
.unwrap()
}
fn task_context_with_engine_state(
state: Arc<QueryEngineState>,
query_ctx: QueryContextRef,
@@ -1649,9 +1668,7 @@ mod tests {
.build()
.unwrap();
assert_eq!(
merge_scan_exec_with_plan(regions.clone(), limited, 10)
.partition_statistics(None)
.unwrap()
merge_scan_statistics(&merge_scan_exec_with_plan(regions.clone(), limited, 10))
.num_rows,
Precision::Inexact(100)
);
@@ -1665,10 +1682,12 @@ mod tests {
.build()
.unwrap();
assert_eq!(
merge_scan_exec_with_plan(vec![RegionId::new(1024, 1)], large_limit, 10)
.partition_statistics(None)
.unwrap()
.num_rows,
merge_scan_statistics(&merge_scan_exec_with_plan(
vec![RegionId::new(1024, 1)],
large_limit,
10,
))
.num_rows,
Precision::Inexact(large_bound)
);
@@ -1678,17 +1697,16 @@ mod tests {
.build()
.unwrap();
assert_eq!(
merge_scan_exec_with_plan(regions.clone(), uncapped.clone(), 10)
.partition_statistics(None)
.unwrap()
.num_rows,
merge_scan_statistics(&merge_scan_exec_with_plan(
regions.clone(),
uncapped.clone(),
10,
))
.num_rows,
Precision::Absent
);
assert_eq!(
merge_scan_exec_with_plan(Vec::new(), uncapped, 10)
.partition_statistics(None)
.unwrap()
.num_rows,
merge_scan_statistics(&merge_scan_exec_with_plan(Vec::new(), uncapped, 10)).num_rows,
Precision::Inexact(0)
);
@@ -1702,10 +1720,12 @@ mod tests {
.build()
.unwrap();
assert_eq!(
merge_scan_exec_with_plan(regions.clone(), global_aggregate, 10)
.partition_statistics(None)
.unwrap()
.num_rows,
merge_scan_statistics(&merge_scan_exec_with_plan(
regions.clone(),
global_aggregate,
10,
))
.num_rows,
Precision::Inexact(2)
);
@@ -1725,10 +1745,7 @@ mod tests {
.build()
.unwrap();
assert_eq!(
merge_scan_exec_with_plan(regions, grouping_sets, 10)
.partition_statistics(None)
.unwrap()
.num_rows,
merge_scan_statistics(&merge_scan_exec_with_plan(regions, grouping_sets, 10)).num_rows,
Precision::Absent
);
}
@@ -3251,7 +3268,7 @@ mod tests {
// A distribution that differs from the current partitioning but shares a
// column name present in partition_cols, so try_with_new_distribution
// produces a clone instead of returning None.
let new_dist = Distribution::HashPartitioned(vec![
let new_dist = Distribution::KeyPartitioned(vec![
Arc::new(Column::new("col1", 0)),
Arc::new(Column::new("col2", 1)),
]);
@@ -3267,6 +3284,24 @@ mod tests {
);
}
#[test]
fn merge_scan_apply_expressions_exposes_remote_dyn_filter_id() {
let query_ctx = QueryContext::arc();
let exec =
remote_dyn_filter_test_exec(Arc::new(TestRegionQueryHandler::default()), query_ctx);
let dyn_filter = install_remote_dyn_filter(&exec);
let expected_expression_id = dyn_filter.expression_id();
let mut expression_ids = Vec::new();
exec.apply_expressions(&mut |expr| {
expression_ids.push(expr.expression_id());
Ok(TreeNodeRecursion::Continue)
})
.unwrap();
assert_eq!(expression_ids, vec![expected_expression_id]);
}
#[test]
fn remote_dyn_filter_preflight_removes_parent_filter_after_dn_runtime_is_ready() {
let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
src/query/src/dist_plan/merge_sort.rs
+65 -47
View File
@@ -16,7 +16,6 @@
//! `SortPreservingMergeExec` operator in datafusion
//!
use std::any::Any;
use std::fmt;
use std::sync::Arc;
@@ -27,12 +26,14 @@ use datafusion::physical_plan::projection::{ProjectionExec, make_with_child, upd
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
Statistics,
ChildStats, ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan,
InputDistributionRequirements, PlanProperties, ReplaceChildrenOptions,
SendableRecordBatchStream, Statistics, StatisticsArgs, apply_expression_roots,
};
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{DataFusionError, Result};
use datafusion_expr::{Extension, LogicalPlan, SortExpr, UserDefinedLogicalNodeCore};
use datafusion_physical_expr::{Distribution, LexOrdering, OrderingRequirements};
use datafusion_physical_expr::{LexOrdering, OrderingRequirements, PhysicalExpr};
/// MergeSort Logical Plan, have same field as `Sort`, but indicate it is a merge sort,
/// which assume each input partition is a sorted stream, and will use `SortPreserveingMergeExec`
@@ -94,7 +95,7 @@ impl MergeSortExec {
fn input_with_fetch(&self, fetch: Option<usize>) -> Arc<dyn ExecutionPlan> {
let input = Arc::clone(self.inner.input());
if let Some(sort) = input.as_any().downcast_ref::<SortExec>()
if let Some(sort) = input.downcast_ref::<SortExec>()
&& sort.preserve_partitioning()
&& sort.expr() == self.inner.expr()
{
@@ -149,7 +150,7 @@ impl ExecutionPlan for MergeSortExec {
/// `MergeSortExec` delegates most behavior to DataFusion's
/// `SortPreservingMergeExec`, but it must not expose itself as that type.
/// DataFusion's `EnforceSorting` optimizer recognizes a bare
/// `SortPreservingMergeExec` via `as_any().downcast_ref::<...>()` and may
/// `SortPreservingMergeExec` via `downcast_ref::<...>()` and may
/// replace it with an unordered `CoalescePartitionsExec(fetch)` when the
/// parent does not require sorted output.
///
@@ -165,10 +166,6 @@ impl ExecutionPlan for MergeSortExec {
/// below `MergeSortExec` when `MergeScanExec` cannot preserve per-partition
/// ordering. This opacity is specifically about protecting the merge stage
/// itself from the `EnforceSorting` rewrite above.
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
self.inner.properties()
}
@@ -193,8 +190,8 @@ impl ExecutionPlan for MergeSortExec {
})
}
fn required_input_distribution(&self) -> Vec<Distribution> {
self.inner.required_input_distribution()
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
self.inner.input_distribution_requirements()
}
fn benefits_from_input_partitioning(&self) -> Vec<bool> {
@@ -205,7 +202,7 @@ impl ExecutionPlan for MergeSortExec {
/// ordered. This is the contract that makes `EnforceSorting` insert a
/// `SortExec` below `MergeSortExec` when the input cannot preserve ordering.
///
/// The opacity of `MergeSortExec::as_any`, not this requirement, is what
/// The opacity of `MergeSortExec`'s downcast identity, not this requirement, is what
/// prevents DataFusion from rewriting the merge stage itself as a bare
/// `SortPreservingMergeExec`.
fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
@@ -220,9 +217,17 @@ impl ExecutionPlan for MergeSortExec {
self.inner.children()
}
fn with_new_children(
fn apply_expressions(
&self,
f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
apply_expression_roots(self.inner.expr().iter().map(|sort_expr| &sort_expr.expr), f)
}
fn replace_children(
self: Arc<Self>,
mut children: Vec<Arc<dyn ExecutionPlan>>,
options: ReplaceChildrenOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
if children.len() != 1 {
return Err(DataFusionError::Internal(format!(
@@ -231,11 +236,31 @@ impl ExecutionPlan for MergeSortExec {
)));
}
Ok(Arc::new(Self::new(
self.inner.expr().clone(),
children.swap_remove(0),
self.inner.fetch(),
)))
match options.children_properties {
ChildrenPropertiesMode::Keep => Ok(Arc::new(Self {
inner: SortPreservingMergeExec::new(
self.inner.expr().clone(),
children.swap_remove(0),
)
.with_fetch(self.inner.fetch()),
})),
ChildrenPropertiesMode::Recompute => Ok(Arc::new(Self::new(
self.inner.expr().clone(),
children.swap_remove(0),
self.inner.fetch(),
))),
}
}
#[allow(deprecated)]
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> Result<Arc<dyn ExecutionPlan>> {
self.replace_children(
children,
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
}
fn execute(
@@ -250,8 +275,16 @@ impl ExecutionPlan for MergeSortExec {
self.inner.metrics()
}
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
self.inner.partition_statistics(partition)
fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
self.inner.child_stats_requests(partition)
}
fn statistics_from_inputs(
&self,
input_stats: &[Arc<Statistics>],
args: &StatisticsArgs,
) -> Result<Arc<Statistics>> {
self.inner.statistics_from_inputs(input_stats, args)
}
fn cardinality_effect(&self) -> CardinalityEffect {
@@ -421,10 +454,6 @@ mod tests {
"PreserveOrderProbeExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
self.inner.properties()
}
@@ -433,6 +462,13 @@ mod tests {
vec![&self.inner]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
self: Arc<Self>,
mut children: Vec<Arc<dyn ExecutionPlan>>,
@@ -489,7 +525,6 @@ mod tests {
assert_eq!(merge_sort.name(), "MergeSortExec");
assert!(
merge_sort
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.is_none(),
"MergeSortExec must stay opaque to EnforceSorting's bare SortPreservingMerge rewrite"
@@ -503,7 +538,7 @@ mod tests {
assert!(!tree.contains("SortPreservingMergeExec"));
let fetched = merge_sort.with_fetch(Some(2)).unwrap();
assert!(fetched.as_any().downcast_ref::<MergeSortExec>().is_some());
assert!(fetched.downcast_ref::<MergeSortExec>().is_some());
assert_eq!(fetched.fetch(), Some(2));
}
@@ -545,12 +580,9 @@ mod tests {
let fetched = merge_sort.with_fetch(Some(2)).unwrap();
assert!(fetched.as_any().downcast_ref::<MergeSortExec>().is_some());
assert!(fetched.downcast_ref::<MergeSortExec>().is_some());
assert_eq!(fetched.fetch(), Some(2));
let child_sort = fetched.children()[0]
.as_any()
.downcast_ref::<SortExec>()
.unwrap();
let child_sort = fetched.children()[0].downcast_ref::<SortExec>().unwrap();
assert_eq!(child_sort.fetch(), Some(2));
assert!(child_sort.preserve_partitioning());
}
@@ -568,14 +600,12 @@ mod tests {
let preserved_spm = bare_spm.with_preserve_order(true).unwrap();
assert!(
preserved_spm
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.is_some(),
"bare SPM should rebuild as bare SPM"
);
assert!(
preserved_spm.children()[0]
.as_any()
.downcast_ref::<PreserveOrderProbeExec>()
.unwrap()
.preserve_order
@@ -585,14 +615,12 @@ mod tests {
let preserved_merge_sort = merge_sort.with_preserve_order(true).unwrap();
assert!(
preserved_merge_sort
.as_any()
.downcast_ref::<MergeSortExec>()
.is_some(),
"MergeSortExec must rewrap the preserve-order child as MergeSortExec"
);
assert!(
preserved_merge_sort
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.is_none(),
"MergeSortExec must not expose a bare SPM after with_preserve_order"
@@ -605,7 +633,6 @@ mod tests {
);
assert!(
preserved_merge_sort.children()[0]
.as_any()
.downcast_ref::<PreserveOrderProbeExec>()
.unwrap()
.preserve_order
@@ -637,7 +664,6 @@ mod tests {
.expect("SPM should accept a narrowing projection that preserves the sort key");
assert!(
swapped_spm
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.is_some(),
"bare SPM should rebuild as bare SPM"
@@ -657,15 +683,11 @@ mod tests {
.expect("MergeSortExec should accept the same projection swap as SPM");
assert!(
swapped_merge_sort
.as_any()
.downcast_ref::<MergeSortExec>()
.is_some(),
swapped_merge_sort.downcast_ref::<MergeSortExec>().is_some(),
"MergeSortExec must rewrap projection swaps as MergeSortExec"
);
assert!(
swapped_merge_sort
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.is_none(),
"MergeSortExec must not expose a bare SPM after projection swap"
@@ -673,7 +695,6 @@ mod tests {
assert_eq!(swapped_merge_sort.fetch(), Some(1));
assert!(
swapped_merge_sort.children()[0]
.as_any()
.downcast_ref::<ProjectionExec>()
.is_some(),
"the projection should move below MergeSortExec"
@@ -754,7 +775,6 @@ mod tests {
.plan;
assert!(
optimized_spm
.as_any()
.downcast_ref::<CoalescePartitionsExec>()
.is_some(),
"this regression test must exercise EnforceSorting's bare SPM -> CoalescePartitionsExec rewrite"
@@ -776,14 +796,12 @@ mod tests {
.plan;
assert!(
optimized_merge_sort
.as_any()
.downcast_ref::<MergeSortExec>()
.is_some(),
"MergeSortExec must stay opaque to the bare SPM rewrite"
);
assert!(
optimized_merge_sort
.as_any()
.downcast_ref::<CoalescePartitionsExec>()
.is_none(),
"MergeSortExec(fetch) is the required distributed TopK merge stage, not an unordered coalesce"
src/query/src/dist_plan/planner.rs
+43 -16
View File
@@ -22,9 +22,11 @@ use async_trait::async_trait;
use catalog::CatalogManagerRef;
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use common_telemetry::debug;
use datafusion::catalog::Session;
use datafusion::common::Result;
use datafusion::datasource::DefaultTableSource;
use datafusion::execution::context::SessionState;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
@@ -60,15 +62,21 @@ pub struct MergeSortExtensionPlanner {}
impl MergeSortExtensionPlanner {
fn ordering(
session_state: &SessionState,
planner: &dyn PhysicalPlanner,
session: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
merge_sort: &MergeSortLogicalPlan,
) -> Result<LexOrdering> {
let ordering = merge_sort
.expr
.iter()
.map(|sort_expr| {
let physical_expr = session_state
.create_physical_expr(sort_expr.expr.clone(), merge_sort.input.schema())?;
let physical_expr = planner.create_physical_expr(
&sort_expr.expr,
merge_sort.input.schema(),
session,
planning_ctx,
)?;
Ok(PhysicalSortExpr::new(
physical_expr,
SortOptions {
@@ -91,11 +99,12 @@ impl MergeSortExtensionPlanner {
impl ExtensionPlanner for MergeSortExtensionPlanner {
async fn plan_extension(
&self,
_planner: &dyn PhysicalPlanner,
planner: &dyn PhysicalPlanner,
node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
physical_inputs: &[Arc<dyn ExecutionPlan>],
session_state: &SessionState,
session: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
if let Some(merge_sort) = node.as_any().downcast_ref::<MergeSortLogicalPlan>() {
if let LogicalPlan::Extension(ext) = &merge_sort.input.as_ref()
@@ -110,14 +119,14 @@ impl ExtensionPlanner for MergeSortExtensionPlanner {
"Expect MergeSort to have one physical input".to_string(),
)
})?;
if input.as_any().downcast_ref::<MergeScanExec>().is_none() {
if input.downcast_ref::<MergeScanExec>().is_none() {
return Err(DataFusionError::Internal(format!(
"Expect MergeSort's input is a MergeScanExec, found {:?}",
physical_inputs
)));
}
let ordering = Self::ordering(session_state, merge_sort)?;
let ordering = Self::ordering(planner, session, planning_ctx, merge_sort)?;
Ok(Some(Arc::new(MergeSortExec::new(
ordering,
input,
@@ -163,7 +172,8 @@ impl ExtensionPlanner for DistExtensionPlanner {
node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
_physical_inputs: &[Arc<dyn ExecutionPlan>],
session_state: &SessionState,
session: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
let Some(merge_scan) = node.as_any().downcast_ref::<MergeScanLogicalPlan>() else {
return Ok(None);
@@ -171,9 +181,9 @@ impl ExtensionPlanner for DistExtensionPlanner {
let input_plan = merge_scan.input();
let fallback = |logical_plan| async move {
let optimized_plan = self.optimize_input_logical_plan(session_state, logical_plan)?;
let optimized_plan = self.optimize_input_logical_plan(session, logical_plan)?;
planner
.create_physical_plan(&optimized_plan, session_state)
.create_physical_plan(&optimized_plan, session)
.await
.map(Some)
};
@@ -196,6 +206,14 @@ impl ExtensionPlanner for DistExtensionPlanner {
// TODO(ruihang): generate different execution plans for different variant merge operation
let schema = merge_scan.schema().as_arrow();
let session_state = session
.as_any()
.downcast_ref::<SessionState>()
.ok_or_else(|| {
DataFusionError::Internal(
"MergeScan requires a SessionState for physical planning".to_string(),
)
})?;
let query_ctx = session_state
.config()
.get_extension()
@@ -208,7 +226,7 @@ impl ExtensionPlanner for DistExtensionPlanner {
schema,
self.region_query_handler.clone(),
query_ctx,
session_state.config().target_partitions(),
session.config().target_partitions(),
merge_scan.partition_cols().clone(),
merge_scan.remote_dyn_filter_producer_id(),
self.enable_per_region_metrics,
@@ -420,11 +438,21 @@ impl DistExtensionPlanner {
/// Input logical plan is analyzed. Thus only call logical optimizer to optimize it.
fn optimize_input_logical_plan(
&self,
session_state: &SessionState,
session: &dyn Session,
plan: &LogicalPlan,
) -> Result<LogicalPlan> {
let state = session_state.clone();
state.optimizer().optimize(plan.clone(), &state, |_, _| {})
let session_state = session
.as_any()
.downcast_ref::<SessionState>()
.ok_or_else(|| {
DataFusionError::Internal(
"MergeScan requires a SessionState for logical optimization".to_string(),
)
})?;
session_state
.optimizer()
.optimize(plan.clone(), session_state, |_, _| {})
}
}
@@ -448,10 +476,9 @@ impl TreeNodeVisitor<'_> for TableNameExtractor {
fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
match node {
LogicalPlan::TableScan(scan) => {
if let Some(source) = scan.source.as_any().downcast_ref::<DefaultTableSource>()
if let Some(source) = scan.source.downcast_ref::<DefaultTableSource>()
&& let Some(provider) = source
.table_provider
.as_any()
.downcast_ref::<DfTableProviderAdapter>()
{
if provider.table().table_type() == TableType::Base {
src/query/src/dist_plan/predicate_extractor.rs
+23 -23
View File
@@ -433,16 +433,16 @@ impl DataFusionExprConverter {
Expr::Cast(cast_expr) => {
// For safe casts, unwrap to the inner expression
// For unsafe casts, skip with debug logging
if Self::is_safe_cast_for_partition_pruning(&cast_expr.data_type) {
if Self::is_safe_cast_for_partition_pruning(cast_expr.field.data_type()) {
Self::convert_to_operand(&cast_expr.expr)
} else {
debug!(
"Skipping unsafe cast for partition pruning: {:?}",
cast_expr.data_type
cast_expr.field.data_type()
);
Err(datafusion_common::DataFusionError::Plan(format!(
"Cast to {:?} not supported for partition pruning",
cast_expr.data_type
cast_expr.field.data_type()
)))
}
}
@@ -638,10 +638,10 @@ mod tests {
fn test_dictionary_cast_preserves_partition_constraint() {
let dictionary_type =
DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8));
let filter = col("tag").eq(Expr::Cast(datafusion_expr::expr::Cast {
expr: Box::new(lit("b")),
data_type: dictionary_type,
}));
let filter = col("tag").eq(Expr::Cast(datafusion_expr::Cast::new(
Box::new(lit("b")),
dictionary_type,
)));
let partition_expr = DataFusionExprConverter::convert(&filter).unwrap();
assert_eq!(
@@ -1011,10 +1011,10 @@ mod tests {
let cases = vec![
FilterTestCase::new(
"safe_cast",
Expr::Cast(datafusion_expr::Cast {
expr: Box::new(col("user_id")),
data_type: DataType::Int64,
})
Expr::Cast(datafusion_expr::Cast::new(
Box::new(col("user_id")),
DataType::Int64,
))
.eq(lit(100i64)),
vec![PartitionExpr::new(
Operand::Column("user_id".to_string()),
@@ -1025,10 +1025,10 @@ mod tests {
),
FilterTestCase::new(
"cast_with_alias",
Expr::Cast(datafusion_expr::Cast {
expr: Box::new(col("user_id").alias("uid")),
data_type: DataType::Int64,
})
Expr::Cast(datafusion_expr::Cast::new(
Box::new(col("user_id").alias("uid")),
DataType::Int64,
))
.eq(lit(100i64)),
vec![PartitionExpr::new(
Operand::Column("user_id".to_string()),
@@ -1039,12 +1039,12 @@ mod tests {
),
FilterTestCase::new(
"unsafe_cast",
Expr::Cast(datafusion_expr::Cast {
expr: Box::new(col("user_id")),
data_type: DataType::List(std::sync::Arc::new(
Expr::Cast(datafusion_expr::Cast::new(
Box::new(col("user_id")),
DataType::List(std::sync::Arc::new(
datafusion::arrow::datatypes::Field::new("item", DataType::Int32, true),
)),
})
))
.eq(lit(100i64)),
vec![],
vec!["user_id"],
@@ -1122,10 +1122,10 @@ mod tests {
let in_expr = col("user_id")
.alias("uid")
.in_list(vec![lit(100i64), lit(200i64)], false);
let cast_expr = Expr::Cast(datafusion_expr::Cast {
expr: Box::new(col("user_id")),
data_type: DataType::Int64,
});
let cast_expr = Expr::Cast(datafusion_expr::Cast::new(
Box::new(col("user_id")),
DataType::Int64,
));
let between_expr = cast_expr.between(lit(300i64), lit(400i64));
in_expr.or(between_expr)
},
src/query/src/dist_plan/remote_dyn_filter_receiver.rs
+5 -3
View File
@@ -18,8 +18,9 @@ use std::hash::{Hash, Hasher};
use std::sync::Arc;
use async_trait::async_trait;
use datafusion::catalog::Session;
use datafusion::common::Result;
use datafusion::execution::context::SessionState;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::physical_expr::utils::conjunction;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::expressions::Column;
@@ -176,7 +177,7 @@ fn remap_physical_expr_columns(
expr: Arc<dyn PhysicalExpr>,
input_schema: &datafusion::arrow::datatypes::Schema,
) -> Result<Arc<dyn PhysicalExpr>> {
if let Some(column) = expr.as_any().downcast_ref::<Column>() {
if let Some(column) = expr.downcast_ref::<Column>() {
return Ok(Arc::new(Column::new_with_schema(
column.name(),
input_schema,
@@ -205,7 +206,8 @@ impl ExtensionPlanner for RemoteDynFilterReceiverExtensionPlanner {
node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
physical_inputs: &[Arc<dyn ExecutionPlan>],
_session_state: &SessionState,
_session: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
let Some(receiver) = node
.as_any()
src/query/src/dummy_catalog.rs
-16
View File
@@ -70,10 +70,6 @@ impl DummyCatalogList {
}
impl CatalogProviderList for DummyCatalogList {
fn as_any(&self) -> &dyn Any {
self
}
fn register_catalog(
&self,
_name: String,
@@ -98,10 +94,6 @@ struct DummyCatalogProvider {
}
impl CatalogProvider for DummyCatalogProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn schema_names(&self) -> Vec<String> {
vec![]
}
@@ -119,10 +111,6 @@ struct DummySchemaProvider {
#[async_trait]
impl SchemaProvider for DummySchemaProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn table_names(&self) -> Vec<String> {
vec![]
}
@@ -162,10 +150,6 @@ impl fmt::Debug for DummyTableProvider {
#[async_trait]
impl TableProvider for DummyTableProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
let schema = self.metadata.schema.arrow_schema();
if !supports_pk_dictionary_encoding(self.engine.name()) {
src/query/src/log_query/planner.rs
+1 -4
View File
@@ -17,12 +17,11 @@ use catalog::table_source::DfTableSourceProvider;
use common_function::utils::escape_like_pattern;
use datafusion::datasource::DefaultTableSource;
use datafusion::execution::SessionState;
use datafusion_common::{DFSchema, ScalarValue};
use datafusion_common::{DFSchema, ScalarValue, TableReference};
use datafusion_expr::utils::{conjunction, disjunction};
use datafusion_expr::{
BinaryExpr, Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder, Operator, col, lit, not,
};
use datafusion_sql::TableReference;
use datatypes::schema::Schema;
use log_query::{AggFunc, BinaryOperator, EqualValue, LogExpr, LogQuery, TimeFilter};
use snafu::{OptionExt, ResultExt};
@@ -56,11 +55,9 @@ impl LogQueryPlanner {
.await
.context(CatalogSnafu)?;
let schema = table_source
.as_any()
.downcast_ref::<DefaultTableSource>()
.context(UnknownTableSnafu)?
.table_provider
.as_any()
.downcast_ref::<DfTableProviderAdapter>()
.context(UnknownTableSnafu)?
.table()
src/query/src/metrics.rs
+1 -1
View File
@@ -280,7 +280,7 @@ fn collect_region_watermarks(plan: Arc<dyn ExecutionPlan>) -> Vec<RegionWatermar
let mut stack = vec![plan];
while let Some(plan) = stack.pop() {
if let Some(merge_scan) = plan.as_any().downcast_ref::<MergeScanExec>()
if let Some(merge_scan) = plan.downcast_ref::<MergeScanExec>()
&& !merge_scan.is_flow_sink_scan()
{
merge_merge_scan_region_watermarks(
src/query/src/optimizer.rs
+1
View File
@@ -16,6 +16,7 @@ pub mod const_normalization;
pub mod constant_term;
pub mod count_nest_aggr;
pub mod count_wildcard;
pub mod enforce_sorting;
pub mod global_limit;
pub(crate) mod insert_assignment;
pub(crate) mod json_schema_concretize;
src/query/src/optimizer/const_normalization.rs
+22 -9
View File
@@ -561,11 +561,11 @@ enum CastInputKind {
/// Returns the input expression and target type for `CAST` and `TRY_CAST` expressions.
fn extract_cast_input(expr: &Expr) -> Option<(CastInputKind, &Expr, &DataType)> {
match expr {
Expr::Cast(Cast { expr, data_type }) => {
Some((CastInputKind::Cast, expr.as_ref(), data_type))
Expr::Cast(Cast { expr, field }) => {
Some((CastInputKind::Cast, expr.as_ref(), field.data_type()))
}
Expr::TryCast(TryCast { expr, data_type }) => {
Some((CastInputKind::TryCast, expr.as_ref(), data_type))
Expr::TryCast(TryCast { expr, field }) => {
Some((CastInputKind::TryCast, expr.as_ref(), field.data_type()))
}
_ => None,
}
@@ -921,7 +921,6 @@ mod tests {
.await
.unwrap();
let filter = physical_plan
.as_any()
.downcast_ref::<FilterExec>()
.expect("regex residual must remain a FilterExec");
assert!(matches!(
@@ -1332,6 +1331,24 @@ mod tests {
expected_greptime: "Filter: CAST(t.ts_ms AS Timestamp(ns)) = TimestampNanosecond(5000000000, None)\n TableScan: t",
expected_datafusion: "Filter: t.ts_ms = TimestampMillisecond(5000, None)\n TableScan: t",
},
Case {
name: "timestamp widening try_cast exact",
fields: vec![Field::new(
"ts_ms",
DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
false,
)],
predicate: try_cast(
col("ts_ms"),
DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
)
.eq(lit(ScalarValue::TimestampNanosecond(
Some(5_000_000_000),
None,
))),
expected_greptime: "Filter: TRY_CAST(t.ts_ms AS Timestamp(ns)) = TimestampNanosecond(5000000000, None)\n TableScan: t",
expected_datafusion: "Filter: t.ts_ms = TimestampMillisecond(5000, None)\n TableScan: t",
},
];
for case in cases {
@@ -1506,10 +1523,6 @@ mod tests {
#[async_trait]
impl TableProvider for ExactPushdownProvider {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema(&self) -> arrow_schema::SchemaRef {
self.schema.clone()
}
src/query/src/optimizer/constant_term.rs
+10 -19
View File
@@ -76,10 +76,6 @@ impl PartialEq for PreCompiledMatchesTermExpr {
impl Eq for PreCompiledMatchesTermExpr {}
impl PhysicalExpr for PreCompiledMatchesTermExpr {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn data_type(
&self,
_input_schema: &arrow_schema::Schema,
@@ -166,10 +162,10 @@ impl PhysicalOptimizerRule for MatchesConstantTermOptimizer {
) -> DfResult<Arc<dyn ExecutionPlan>> {
let res = plan
.transform_down(&|plan: Arc<dyn ExecutionPlan>| {
if let Some(filter) = plan.as_any().downcast_ref::<FilterExec>() {
if let Some(filter) = plan.downcast_ref::<FilterExec>() {
let pred = filter.predicate().clone();
let new_pred = pred.transform_down(&|expr: Arc<dyn PhysicalExpr>| {
if let Some(func) = expr.as_any().downcast_ref::<ScalarFunctionExpr>() {
if let Some(func) = expr.downcast_ref::<ScalarFunctionExpr>() {
if !func.name().eq_ignore_ascii_case("matches_term") {
return Ok(Transformed::no(expr));
}
@@ -178,7 +174,7 @@ impl PhysicalOptimizerRule for MatchesConstantTermOptimizer {
return Ok(Transformed::no(expr));
}
if let Some(lit) = args[1].as_any().downcast_ref::<Literal>()
if let Some(lit) = args[1].downcast_ref::<Literal>()
&& let ScalarValue::Utf8(Some(term)) = lit.value()
{
let finder = MatchesTermFinder::new(term);
@@ -248,6 +244,7 @@ mod tests {
use datafusion::physical_plan::get_plan_string;
use datafusion_common::{Column, DFSchema};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{Expr, Literal, ScalarUDF};
use datafusion_physical_expr::{ScalarFunctionExpr, create_physical_expr};
use datatypes::prelude::ConcreteDataType;
@@ -343,6 +340,7 @@ mod tests {
)),
&DFSchema::try_from(batch.schema().clone()).unwrap(),
&Default::default(),
&PhysicalPlanningContext::default(),
)
.unwrap();
@@ -357,16 +355,11 @@ mod tests {
.optimize(Arc::new(filter), &Default::default())
.unwrap();
let optimized_filter = optimized_plan
.as_any()
.downcast_ref::<FilterExec>()
.unwrap();
let optimized_filter = optimized_plan.downcast_ref::<FilterExec>().unwrap();
let predicate = optimized_filter.predicate();
// The predicate should be a PreCompiledMatchesTermExpr
assert!(
std::any::TypeId::of::<PreCompiledMatchesTermExpr>() == predicate.as_any().type_id()
);
assert!(predicate.is::<PreCompiledMatchesTermExpr>());
}
#[test]
@@ -413,6 +406,7 @@ mod tests {
)),
&DFSchema::try_from(batch.schema().clone()).unwrap(),
&Default::default(),
&PhysicalPlanningContext::default(),
)
.unwrap();
@@ -426,14 +420,11 @@ mod tests {
.optimize(Arc::new(filter), &Default::default())
.unwrap();
let optimized_filter = optimized_plan
.as_any()
.downcast_ref::<FilterExec>()
.unwrap();
let optimized_filter = optimized_plan.downcast_ref::<FilterExec>().unwrap();
let predicate = optimized_filter.predicate();
// The predicate should still be a ScalarFunctionExpr
assert!(std::any::TypeId::of::<ScalarFunctionExpr>() == predicate.as_any().type_id());
assert!(predicate.is::<ScalarFunctionExpr>());
}
#[tokio::test]
src/query/src/optimizer/count_wildcard.rs
+3 -9
View File
@@ -16,13 +16,12 @@ use datafusion::datasource::DefaultTableSource;
use datafusion_common::tree_node::{
Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
};
use datafusion_common::{Column, Result as DataFusionResult, ScalarValue};
use datafusion_common::{Column, Result as DataFusionResult, ScalarValue, TableReference};
use datafusion_expr::expr::{AggregateFunction, WindowFunction};
use datafusion_expr::utils::COUNT_STAR_EXPANSION;
use datafusion_expr::{Expr, LogicalPlan, WindowFunctionDefinition, col, lit};
use datafusion_optimizer::AnalyzerRule;
use datafusion_optimizer::utils::NamePreserver;
use datafusion_sql::TableReference;
use table::table::adapter::DfTableProviderAdapter;
/// A replacement to DataFusion's [`CountWildcardRule`]. This rule
@@ -155,13 +154,9 @@ impl TreeNodeVisitor<'_> for TimeIndexFinder {
}
if let LogicalPlan::TableScan(table_scan) = &node
&& let Some(source) = table_scan
.source
.as_any()
.downcast_ref::<DefaultTableSource>()
&& let Some(source) = table_scan.source.downcast_ref::<DefaultTableSource>()
&& let Some(adapter) = source
.table_provider
.as_any()
.downcast_ref::<DfTableProviderAdapter>()
{
let table_info = adapter.table().table_info();
@@ -206,9 +201,8 @@ mod test {
use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
use datafusion::functions_aggregate::count::count_all;
use datafusion::functions_aggregate::min_max::max;
use datafusion_common::Column;
use datafusion_common::{Column, TableReference};
use datafusion_expr::LogicalPlanBuilder;
use datafusion_sql::TableReference;
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::{ColumnSchema, Schema, SchemaBuilder};
use datatypes::vectors::{Int64Vector, TimestampMillisecondVector, VectorRef};
src/query/src/optimizer/enforce_sorting.rs
+91
View File
@@ -0,0 +1,91 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! Sorting enforcement that runs after GreptimeDB's custom physical rules.
//!
//! DataFusion 55 moved the standalone `EnforceSorting` phases into
//! `EnsureRequirements`. GreptimeDB still needs to rerun those phases after
//! custom rules modify scan partitioning and distribution.
use std::sync::Arc;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{
OrderPreservationContext, replace_with_order_preserving_variants,
};
use datafusion::physical_optimizer::enforce_sorting::sort_pushdown::{
SortPushDown, assign_initial_requirements, pushdown_sorts,
};
use datafusion::physical_optimizer::enforce_sorting::{
PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
parallelize_sorts, replace_with_partial_sort,
};
use datafusion::physical_plan::ExecutionPlan;
use datafusion_common::Result;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
/// Runs the standalone sorting-enforcement pipeline removed in DataFusion 55.
#[derive(Debug)]
pub struct EnforceSorting;
impl PhysicalOptimizerRule for EnforceSorting {
fn optimize(
&self,
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
// Phase 1: ensure sorting requirements and remove redundant sorts.
let sorting = PlanWithCorrespondingSort::new_default(plan);
let sorting = sorting.transform_up(ensure_sorting)?.data;
// Phase 2: optionally turn CoalescePartitions + Sort into parallel
// sorts followed by a SortPreservingMerge.
let plan = if config.optimizer.repartition_sorts {
let parallel = PlanWithCorrespondingCoalescePartitions::new_default(sorting.plan)
.transform_up(parallelize_sorts)
.data()?;
parallel.plan
} else {
sorting.plan
};
// Phase 3: use order-preserving executor variants where appropriate.
let variants = OrderPreservationContext::new_default(plan);
let variants = variants
.transform_up(|context| {
replace_with_order_preserving_variants(context, false, true, config)
})
.data()?;
// Phase 4: push sorts down through order-preserving operators.
let mut pushdown = SortPushDown::new_default(variants.plan);
assign_initial_requirements(&mut pushdown);
let pushed = pushdown_sorts(pushdown)?;
// Phase 5: exploit an already-satisfied prefix on unbounded inputs.
pushed
.plan
.transform_up(|plan| Ok(Transformed::yes(replace_with_partial_sort(plan)?)))
.data()
}
fn name(&self) -> &str {
"EnforceSorting"
}
fn schema_check(&self) -> bool {
true
}
}
src/query/src/optimizer/global_limit.rs
+40 -40
View File
@@ -21,7 +21,9 @@ use datafusion::physical_plan::filter::FilterExec;
use datafusion::physical_plan::limit::GlobalLimitExec;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
use datafusion::physical_plan::{
ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions,
};
use datafusion_common::Result as DfResult;
use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning};
@@ -57,7 +59,7 @@ impl EnsureGlobalLimitForFetch {
let plan = if children.is_empty() {
plan
} else {
let required_input_distribution = plan.required_input_distribution();
let required_input_distribution = plan.input_distribution_requirements();
let required_input_ordering = plan.required_input_ordering();
let maintains_input_order = plan.maintains_input_order();
let child_parent = ParentContext {
@@ -72,7 +74,7 @@ impl EnsureGlobalLimitForFetch {
.enumerate()
.map(|(idx, child)| {
let required_distribution = required_input_distribution
.get(idx)
.child_distribution(idx)
.cloned()
.unwrap_or(Distribution::UnspecifiedDistribution);
let partitioning_to_restore =
@@ -101,7 +103,10 @@ impl EnsureGlobalLimitForFetch {
Self::optimize_plan(Arc::clone(child), parent)
})
.collect::<DfResult<Vec<_>>>()?;
plan.with_new_children(children)?
plan.replace_children(
children,
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)?
};
let Some(fetch) = plan.fetch() else {
@@ -111,7 +116,7 @@ impl EnsureGlobalLimitForFetch {
if parent
.global_fetch
.is_some_and(|parent_fetch| parent_fetch <= fetch)
|| !plan.as_any().is::<FilterExec>()
|| !plan.is::<FilterExec>()
|| plan.output_partitioning().partition_count() <= 1
{
return Ok(plan);
@@ -149,10 +154,10 @@ impl Default for ParentContext {
fn provided_global_fetch(plan: &Arc<dyn ExecutionPlan>) -> Option<usize> {
let fetch = plan.fetch()?;
(plan.as_any().is::<GlobalLimitExec>()
|| plan.as_any().is::<CoalescePartitionsExec>()
|| plan.as_any().is::<SortPreservingMergeExec>()
|| plan.as_any().is::<MergeSortExec>())
(plan.is::<GlobalLimitExec>()
|| plan.is::<CoalescePartitionsExec>()
|| plan.is::<SortPreservingMergeExec>()
|| plan.is::<MergeSortExec>())
.then_some(fetch)
}
@@ -196,7 +201,7 @@ fn partitioning_to_restore_for(
child: &Arc<dyn ExecutionPlan>,
required_distribution: &Distribution,
) -> Option<Partitioning> {
if !matches!(required_distribution, Distribution::HashPartitioned(_))
if !matches!(required_distribution, Distribution::KeyPartitioned(_))
|| child.output_partitioning().partition_count() <= 1
{
return None;
@@ -233,7 +238,7 @@ fn inherited_partitioning_to_restore(
let satisfies_parent_distribution = matches!(
parent.required_distribution,
Distribution::HashPartitioned(_)
Distribution::KeyPartitioned(_)
) && plan
.output_partitioning()
.satisfaction(
@@ -272,7 +277,7 @@ mod tests {
let optimized =
EnsureGlobalLimitForFetch::optimize_plan(filter, ParentContext::default()).unwrap();
assert!(optimized.as_any().is::<CoalescePartitionsExec>());
assert!(optimized.is::<CoalescePartitionsExec>());
assert_eq!(optimized.fetch(), Some(1));
assert_eq!(optimized.output_partitioning().partition_count(), 1);
}
@@ -295,7 +300,7 @@ mod tests {
let projection = optimized.children()[0];
let coalesce = projection.children()[0];
assert!(coalesce.as_any().is::<CoalescePartitionsExec>());
assert!(coalesce.is::<CoalescePartitionsExec>());
assert_eq!(coalesce.fetch(), Some(5));
}
@@ -310,8 +315,8 @@ mod tests {
EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
let child = optimized.children()[0];
assert!(optimized.as_any().is::<SortPreservingMergeExec>());
assert!(child.as_any().is::<FilterExec>());
assert!(optimized.is::<SortPreservingMergeExec>());
assert!(child.is::<FilterExec>());
}
#[test]
@@ -325,8 +330,8 @@ mod tests {
EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
let child = optimized.children()[0];
assert!(optimized.as_any().is::<SortPreservingMergeExec>());
assert!(child.as_any().is::<SortPreservingMergeExec>());
assert!(optimized.is::<SortPreservingMergeExec>());
assert!(child.is::<SortPreservingMergeExec>());
assert_eq!(child.fetch(), Some(5));
}
@@ -340,8 +345,8 @@ mod tests {
EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
let child = optimized.children()[0];
assert!(optimized.as_any().is::<MergeSortExec>());
assert!(child.as_any().is::<FilterExec>());
assert!(optimized.is::<MergeSortExec>());
assert!(child.is::<FilterExec>());
}
#[test]
@@ -354,10 +359,10 @@ mod tests {
EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
let child = optimized.children()[0];
assert!(optimized.as_any().is::<MergeSortExec>());
assert!(child.as_any().is::<SortPreservingMergeExec>());
assert!(optimized.is::<MergeSortExec>());
assert!(child.is::<SortPreservingMergeExec>());
assert_eq!(child.fetch(), Some(5));
assert!(child.children()[0].as_any().is::<FilterExec>());
assert!(child.children()[0].is::<FilterExec>());
}
#[test]
@@ -371,8 +376,8 @@ mod tests {
EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
let child = optimized.children()[0];
assert!(optimized.as_any().is::<SortPreservingMergeExec>());
assert!(child.as_any().is::<SortPreservingMergeExec>());
assert!(optimized.is::<SortPreservingMergeExec>());
assert!(child.is::<SortPreservingMergeExec>());
assert_eq!(child.fetch(), Some(1));
}
@@ -396,10 +401,7 @@ mod tests {
None,
)
.unwrap();
let merge = optimized
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.unwrap();
let merge = optimized.downcast_ref::<SortPreservingMergeExec>().unwrap();
assert_eq!(merge.expr(), &actual_ordering);
}
@@ -423,9 +425,9 @@ mod tests {
let projection = optimized.children()[0];
let child = projection.children()[0];
assert!(optimized.as_any().is::<SortPreservingMergeExec>());
assert!(projection.as_any().is::<ProjectionExec>());
assert!(child.as_any().is::<SortPreservingMergeExec>());
assert!(optimized.is::<SortPreservingMergeExec>());
assert!(projection.is::<ProjectionExec>());
assert!(child.is::<SortPreservingMergeExec>());
assert_eq!(child.fetch(), Some(1));
}
@@ -455,13 +457,13 @@ mod tests {
let optimized =
EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
let left = optimized.children()[0];
let repartition = left.as_any().downcast_ref::<RepartitionExec>().unwrap();
let repartition = left.downcast_ref::<RepartitionExec>().unwrap();
assert!(matches!(
repartition.partitioning(),
Partitioning::Hash(_, 3)
));
assert!(repartition.input().as_any().is::<CoalescePartitionsExec>());
assert!(repartition.input().is::<CoalescePartitionsExec>());
assert_eq!(repartition.input().fetch(), Some(1));
}
@@ -499,16 +501,15 @@ mod tests {
EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
let projection = optimized.children()[0];
let repartition = projection.children()[0]
.as_any()
.downcast_ref::<RepartitionExec>()
.unwrap();
assert!(projection.as_any().is::<ProjectionExec>());
assert!(projection.is::<ProjectionExec>());
assert!(matches!(
repartition.partitioning(),
Partitioning::Hash(_, 3)
));
assert!(repartition.input().as_any().is::<CoalescePartitionsExec>());
assert!(repartition.input().is::<CoalescePartitionsExec>());
assert_eq!(repartition.input().fetch(), Some(1));
}
@@ -542,17 +543,16 @@ mod tests {
let outer_projection = optimized.children()[0];
let inner_projection = outer_projection.children()[0];
let repartition = inner_projection.children()[0]
.as_any()
.downcast_ref::<RepartitionExec>()
.unwrap();
assert!(outer_projection.as_any().is::<ProjectionExec>());
assert!(inner_projection.as_any().is::<ProjectionExec>());
assert!(outer_projection.is::<ProjectionExec>());
assert!(inner_projection.is::<ProjectionExec>());
assert!(matches!(
repartition.partitioning(),
Partitioning::Hash(_, 3)
));
assert!(repartition.input().as_any().is::<CoalescePartitionsExec>());
assert!(repartition.input().is::<CoalescePartitionsExec>());
assert_eq!(repartition.input().fetch(), Some(1));
}
src/query/src/optimizer/insert_assignment.rs
+4 -1
View File
@@ -163,11 +163,14 @@ fn retarget_assignment_cast(
let expr = unalias_mut(expr);
let Expr::Cast(Cast {
expr: source,
data_type: DataType::Timestamp(unit, None),
field,
}) = expr
else {
return Ok(false);
};
let DataType::Timestamp(unit, None) = field.data_type() else {
return Ok(false);
};
let unit = *unit;
if !matches!(
src/query/src/optimizer/json_type_concretize.rs
+4 -7
View File
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::collections::HashMap;
use arrow_schema::DataType;
@@ -50,11 +51,7 @@ impl OptimizerRule for JsonTypeConcretizeRule {
plan.transform_down(|plan| match &plan {
LogicalPlan::TableScan(table_scan) => {
let Some(source) = table_scan
.source
.as_any()
.downcast_ref::<DefaultTableSource>()
else {
let Some(source) = table_scan.source.downcast_ref::<DefaultTableSource>() else {
return Ok(Transformed::no(plan));
};
@@ -95,12 +92,12 @@ fn apply_json_type_hint(
return false;
}
if let Some(adapter) = provider.as_any().downcast_ref::<DummyTableProvider>() {
if let Some(adapter) = (provider as &dyn Any).downcast_ref::<DummyTableProvider>() {
adapter.with_json_type_hint(json_types);
return true;
}
if let Some(adapter) = provider.as_any().downcast_ref::<DfTableProviderAdapter>() {
if let Some(adapter) = (provider as &dyn Any).downcast_ref::<DfTableProviderAdapter>() {
adapter.with_json_type_hint(json_types);
return true;
}

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