rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-29 15:06:59 +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
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;
}
src/query/src/optimizer/parallelize_scan.rs
+2 -4
View File
@@ -55,12 +55,10 @@ impl ParallelizeScan {
let result = plan
.transform_down(|plan| {
if let Some(sort_exec) = plan.as_any().downcast_ref::<SortExec>() {
if let Some(sort_exec) = plan.downcast_ref::<SortExec>() {
// save the first order expr
first_order_expr = Some(sort_exec.expr().first()).cloned();
} else if let Some(region_scan_exec) =
plan.as_any().downcast_ref::<RegionScanExec>()
{
} else if let Some(region_scan_exec) = plan.downcast_ref::<RegionScanExec>() {
let expected_partition_num = config.execution.target_partitions;
if region_scan_exec.is_partition_set()
|| region_scan_exec.scanner_type().as_str() == "SinglePartition"
src/query/src/optimizer/pass_distribution.rs
+19 -23
View File
@@ -18,7 +18,9 @@ use datafusion::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::projection::ProjectionExec;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::{ExecutionPlan, Partitioning};
use datafusion::physical_plan::{
ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions,
};
use datafusion_common::Result as DfResult;
use datafusion_physical_expr::Distribution;
use datafusion_physical_expr::utils::map_columns_before_projection;
@@ -68,8 +70,8 @@ impl PassDistribution {
current_req: Option<Distribution>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
// If this is a MergeScanExec, try to apply the current requirement.
if let Some(merge_scan) = plan.as_any().downcast_ref::<MergeScanExec>()
&& let Some(Distribution::HashPartitioned(hash_exprs)) = current_req.as_ref()
if let Some(merge_scan) = plan.downcast_ref::<MergeScanExec>()
&& let Some(Distribution::KeyPartitioned(hash_exprs)) = current_req.as_ref()
{
if let Partitioning::Hash(current_hash_exprs, _) = &merge_scan.properties().partitioning
&& *current_hash_exprs == *hash_exprs
@@ -78,7 +80,7 @@ impl PassDistribution {
}
if let Some(new_plan) = merge_scan
.try_with_new_distribution(Distribution::HashPartitioned(hash_exprs.clone()))
.try_with_new_distribution(Distribution::KeyPartitioned(hash_exprs.clone()))
{
// Leaf node; no children to process
return Ok(Arc::new(new_plan) as _);
@@ -97,10 +99,10 @@ impl PassDistribution {
return Ok(plan);
}
let required = plan.required_input_distribution();
let required = plan.input_distribution_requirements();
let mut new_children = Vec::with_capacity(children.len());
for (idx, child) in children.into_iter().enumerate() {
let child_req = match required.get(idx) {
let child_req = match required.child_distribution(idx) {
Some(Distribution::UnspecifiedDistribution) if idx == 0 => {
Self::map_hash_requirement_through_projection(plan.as_ref(), &current_req)
}
@@ -121,7 +123,10 @@ impl PassDistribution {
if unchanged {
Ok(plan)
} else {
plan.with_new_children(new_children)
plan.replace_children(
new_children,
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
}
}
@@ -129,11 +134,11 @@ impl PassDistribution {
plan: &dyn ExecutionPlan,
current_req: &Option<Distribution>,
) -> Option<Distribution> {
let Some(Distribution::HashPartitioned(required_exprs)) = current_req else {
let Some(Distribution::KeyPartitioned(required_exprs)) = current_req else {
return None;
};
let projection = plan.as_any().downcast_ref::<ProjectionExec>()?;
let projection = plan.downcast_ref::<ProjectionExec>()?;
let proj_exprs = projection
.expr()
.iter()
@@ -141,7 +146,7 @@ impl PassDistribution {
.collect::<Vec<_>>();
let mapped = map_columns_before_projection(required_exprs, &proj_exprs);
(mapped.len() == required_exprs.len()).then_some(Distribution::HashPartitioned(mapped))
(mapped.len() == required_exprs.len()).then_some(Distribution::KeyPartitioned(mapped))
}
}
@@ -261,12 +266,8 @@ mod tests {
let optimized = PassDistribution
.optimize(join, &ConfigOptions::default())
.unwrap();
let hash_join = optimized.as_any().downcast_ref::<HashJoinExec>().unwrap();
let left_projection = hash_join
.left()
.as_any()
.downcast_ref::<ProjectionExec>()
.unwrap();
let hash_join = optimized.downcast_ref::<HashJoinExec>().unwrap();
let left_projection = hash_join.left().downcast_ref::<ProjectionExec>().unwrap();
let left_partitioning = left_projection.input().output_partitioning();
let right_partitioning = hash_join.right().output_partitioning();
@@ -293,7 +294,7 @@ mod tests {
fn merge_scan_rejects_hash_requirement_on_partition_key_subset() {
let merge_scan = test_merge_scan_exec(test_schema());
let new_plan = merge_scan.try_with_new_distribution(Distribution::HashPartitioned(vec![
let new_plan = merge_scan.try_with_new_distribution(Distribution::KeyPartitioned(vec![
partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1),
]));
@@ -355,12 +356,7 @@ mod tests {
fn column_names(exprs: &[Arc<dyn PhysicalExpr>]) -> Vec<&str> {
exprs
.iter()
.map(|expr| {
expr.as_any()
.downcast_ref::<PhysicalColumn>()
.unwrap()
.name()
})
.map(|expr| expr.downcast_ref::<PhysicalColumn>().unwrap().name())
.collect()
}
}
src/query/src/optimizer/promql_tsid_narrow_join.rs
+6 -6
View File
@@ -55,7 +55,7 @@ impl PhysicalOptimizerRule for PromqlTsidNarrowJoin {
impl PromqlTsidNarrowJoin {
fn rewrite_join(plan: Arc<dyn ExecutionPlan>) -> DfResult<Transformed<Arc<dyn ExecutionPlan>>> {
let Some(hash_join) = plan.as_any().downcast_ref::<HashJoinExec>() else {
let Some(hash_join) = plan.downcast_ref::<HashJoinExec>() else {
return Ok(Transformed::no(plan));
};
@@ -103,8 +103,8 @@ impl PromqlTsidNarrowJoin {
for (left, right) in hash_join.on() {
let (Some(left_col), Some(right_col)) = (
left.as_any().downcast_ref::<Column>(),
right.as_any().downcast_ref::<Column>(),
left.downcast_ref::<Column>(),
right.downcast_ref::<Column>(),
) else {
return false;
};
@@ -199,7 +199,7 @@ mod tests {
let optimized = PromqlTsidNarrowJoin
.optimize(join, &ConfigOptions::default())
.unwrap();
let optimized_join = optimized.as_any().downcast_ref::<HashJoinExec>().unwrap();
let optimized_join = optimized.downcast_ref::<HashJoinExec>().unwrap();
assert_eq!(optimized_join.partition_mode(), &PartitionMode::CollectLeft);
assert_eq!(optimized.schema(), original_schema);
@@ -262,7 +262,7 @@ mod tests {
let optimized = PromqlTsidNarrowJoin
.optimize(join, &ConfigOptions::default())
.unwrap();
let optimized_join = optimized.as_any().downcast_ref::<HashJoinExec>().unwrap();
let optimized_join = optimized.downcast_ref::<HashJoinExec>().unwrap();
assert_eq!(optimized_join.partition_mode(), &PartitionMode::CollectLeft);
}
@@ -317,7 +317,7 @@ mod tests {
let optimized = PromqlTsidNarrowJoin
.optimize(join, &ConfigOptions::default())
.unwrap();
let optimized_join = optimized.as_any().downcast_ref::<HashJoinExec>().unwrap();
let optimized_join = optimized.downcast_ref::<HashJoinExec>().unwrap();
assert_eq!(optimized_join.partition_mode(), &PartitionMode::Partitioned);
}
src/query/src/optimizer/remove_duplicate.rs
+15 -19
View File
@@ -16,8 +16,8 @@ use std::sync::Arc;
use datafusion::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions};
use datafusion_common::Result as DfResult;
use datafusion_common::tree_node::{Transformed, TreeNode};
@@ -35,29 +35,18 @@ impl PhysicalOptimizerRule for RemoveDuplicate {
plan: Arc<dyn ExecutionPlan>,
_config: &ConfigOptions,
) -> DfResult<Arc<dyn ExecutionPlan>> {
Self::do_optimize(plan)
}
fn name(&self) -> &str {
"RemoveDuplicateRule"
}
fn schema_check(&self) -> bool {
false
}
}
impl RemoveDuplicate {
fn do_optimize(plan: Arc<dyn ExecutionPlan>) -> DfResult<Arc<dyn ExecutionPlan>> {
let result = plan
.transform_down(|plan| {
if plan.as_any().is::<RepartitionExec>() {
if plan.is::<RepartitionExec>() {
// check child
let child = plan.children()[0].clone();
if child.as_any().type_id() == plan.as_any().type_id() {
if child.is::<RepartitionExec>() {
// remove child
let grand_child = child.children()[0].clone();
let new_plan = plan.with_new_children(vec![grand_child])?;
let new_plan = plan.replace_children(
vec![grand_child],
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)?;
return Ok(Transformed::yes(new_plan));
}
}
@@ -65,7 +54,14 @@ impl RemoveDuplicate {
Ok(Transformed::no(plan))
})?
.data;
Ok(result)
}
fn name(&self) -> &str {
"RemoveDuplicateRule"
}
fn schema_check(&self) -> bool {
true
}
}
src/query/src/optimizer/scan_hint.rs
+5 -16
View File
@@ -76,19 +76,11 @@ impl ScanHintRule {
let LogicalPlan::TableScan(mut table_scan) = plan else {
return Ok(Transformed::no(plan));
};
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(LogicalPlan::TableScan(table_scan)));
};
// The provider in the region server is [DummyTableProvider].
let Some(original) = source
.table_provider
.as_any()
.downcast_ref::<DummyTableProvider>()
else {
let Some(original) = source.table_provider.downcast_ref::<DummyTableProvider>() else {
return Ok(Transformed::no(LogicalPlan::TableScan(table_scan)));
};
@@ -455,7 +447,7 @@ fn single_evaluation_projection_expr_allowed(
};
alias.name == column.name
&& matches!(
cast.data_type,
cast.field.data_type(),
DataType::Timestamp(ArrowTimeUnit::Millisecond, None)
)
&& matches!(
@@ -596,11 +588,8 @@ mod test {
let mut requests = Vec::new();
plan.apply_with_subqueries(|node| {
if let LogicalPlan::TableScan(scan) = node
&& let Some(source) = scan.source.as_any().downcast_ref::<DefaultTableSource>()
&& let Some(provider) = source
.table_provider
.as_any()
.downcast_ref::<DummyTableProvider>()
&& let Some(source) = scan.source.downcast_ref::<DefaultTableSource>()
&& let Some(provider) = source.table_provider.downcast_ref::<DummyTableProvider>()
{
requests.push((scan.table_name.to_string(), provider.scan_request()));
}
src/query/src/optimizer/scan_hint/vector_search.rs
+3 -7
View File
@@ -18,11 +18,10 @@ use common_function::scalars::vector::distance::{
VEC_COS_DISTANCE, VEC_DOT_PRODUCT, VEC_L2SQ_DISTANCE,
};
use common_telemetry::debug;
use datafusion_common::ScalarValue;
use datafusion_common::{ScalarValue, TableReference};
use datafusion_expr::logical_plan::FetchType;
use datafusion_expr::utils::split_conjunction;
use datafusion_expr::{Expr, SortExpr};
use datafusion_sql::TableReference;
use datatypes::types::parse_string_to_vector_type_value;
use store_api::storage::{VectorDistanceMetric, VectorSearchRequest};
@@ -488,11 +487,8 @@ mod tests {
let mut request = None;
plan.apply_with_subqueries(|node| {
if let LogicalPlan::TableScan(scan) = node
&& let Some(source) = scan.source.as_any().downcast_ref::<DefaultTableSource>()
&& let Some(provider) = source
.table_provider
.as_any()
.downcast_ref::<DummyTableProvider>()
&& let Some(source) = scan.source.downcast_ref::<DefaultTableSource>()
&& let Some(provider) = source.table_provider.downcast_ref::<DummyTableProvider>()
{
request = Some(provider.scan_request());
}
src/query/src/optimizer/string_normalization.rs
+3 -6
View File
@@ -87,8 +87,8 @@ impl TreeNodeRewriter for StringNormalizationConverter {
/// Otherwise - no modifications applied
fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
let new_expr = match expr {
Expr::Cast(Cast { expr, data_type }) => {
let expr = match data_type {
Expr::Cast(Cast { expr, field }) => {
let expr = match field.data_type() {
DataType::Timestamp(_, _) => match *expr {
Expr::Literal(value, _) => match value {
ScalarValue::Utf8(Some(s)) => trim_utf_expr(s),
@@ -98,10 +98,7 @@ impl TreeNodeRewriter for StringNormalizationConverter {
},
_ => *expr,
};
Expr::Cast(Cast {
expr: Box::new(expr),
data_type,
})
Expr::Cast(Cast::new_from_field(Box::new(expr), field))
}
expr => expr,
};
src/query/src/optimizer/type_conversion.rs
+7 -6
View File
@@ -63,6 +63,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
projected_schema,
filters,
fetch,
statistics_requests,
}) => {
let mut converter = TypeConverter::new(projected_schema.clone(), ctx.query_ctx());
let rewrite_filters = filters
@@ -76,6 +77,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
projected_schema,
filters: rewrite_filters,
fetch,
statistics_requests,
})))
}
LogicalPlan::Projection { .. } => {
@@ -359,10 +361,9 @@ mod tests {
use std::sync::Arc;
use datafusion_common::arrow::datatypes::Field;
use datafusion_common::{Column, DFSchema, NullEquality};
use datafusion_common::{Column, DFSchema, NullEquality, TableReference};
use datafusion_expr::expr::{Cast, Exists};
use datafusion_expr::{Join, JoinConstraint, JoinType, Literal, LogicalPlanBuilder, Subquery};
use datafusion_sql::TableReference;
use session::context::QueryContext;
use super::*;
@@ -487,10 +488,10 @@ mod tests {
use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit;
let mut converter = TypeConverter::new(Arc::new(DFSchema::empty()), QueryContext::arc());
let expr = Expr::Cast(Cast {
expr: Box::new("2009-02-13 23:31:30".lit()),
data_type: DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
});
let expr = Expr::Cast(Cast::new(
Box::new("2009-02-13 23:31:30".lit()),
DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
));
assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr);
}
src/query/src/optimizer/windowed_sort.rs
+26 -26
View File
@@ -16,7 +16,6 @@ use std::sync::Arc;
use arrow_schema::DataType;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::coop::CooperativeExec;
use datafusion::physical_plan::filter::FilterExec;
@@ -24,6 +23,7 @@ use datafusion::physical_plan::projection::ProjectionExec;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions};
use datafusion_common::Result as DataFusionResult;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_physical_expr::expressions::{CastExpr, Column as PhysicalColumn};
@@ -71,7 +71,7 @@ impl WindowedSortPhysicalRule {
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let result = plan
.transform_down(|plan| {
if let Some(sort_exec) = plan.as_any().downcast_ref::<SortExec>() {
if let Some(sort_exec) = plan.downcast_ref::<SortExec>() {
// TODO: support multiple expr in windowed sort
if sort_exec.expr().len() != 1 {
return Ok(Transformed::no(plan));
@@ -88,10 +88,7 @@ impl WindowedSortPhysicalRule {
let input_schema = sort_input.schema();
let first_sort_expr = sort_exec.expr().first();
if let Some(column_expr) = first_sort_expr
.expr
.as_any()
.downcast_ref::<PhysicalColumn>()
if let Some(column_expr) = first_sort_expr.expr.downcast_ref::<PhysicalColumn>()
&& matches!(
input_schema.field(column_expr.index()).data_type(),
DataType::Timestamp(_, _)
@@ -165,26 +162,26 @@ fn fetch_partition_range(input: Arc<dyn ExecutionPlan>) -> DataFusionResult<Opti
let mut tag_columns = None;
input.transform_up(|plan| {
if plan.as_any().is::<CooperativeExec>() {
if plan.is::<CooperativeExec>() {
return Ok(Transformed::no(plan));
}
// Unappliable case, reset the state.
if plan.as_any().is::<RepartitionExec>()
|| plan.as_any().is::<CoalescePartitionsExec>()
|| plan.as_any().is::<SortExec>()
|| plan.as_any().is::<WindowedSortExec>()
if plan.is::<RepartitionExec>()
|| plan.is::<CoalescePartitionsExec>()
|| plan.is::<SortExec>()
|| plan.is::<WindowedSortExec>()
{
partition_ranges = None;
}
// only a very limited set of plans can exist between region scan and sort exec
// other plans might make this optimize wrong, so be safe here by limiting it
if !(plan.as_any().is::<ProjectionExec>() || plan.as_any().is::<FilterExec>()) {
if !(plan.is::<ProjectionExec>() || plan.is::<FilterExec>()) {
partition_ranges = None;
}
if let Some(region_scan_exec) = plan.as_any().downcast_ref::<RegionScanExec>() {
if let Some(region_scan_exec) = plan.downcast_ref::<RegionScanExec>() {
// `PerSeries` distribution is not supported in windowed sort.
if region_scan_exec.distribution()
== Some(store_api::storage::TimeSeriesDistribution::PerSeries)
@@ -216,11 +213,11 @@ fn is_time_index_expr(
plan: &Arc<dyn ExecutionPlan>,
expr: &Arc<dyn PhysicalExpr>,
) -> DataFusionResult<bool> {
if let Some(column_expr) = expr.as_any().downcast_ref::<PhysicalColumn>() {
if let Some(column_expr) = expr.downcast_ref::<PhysicalColumn>() {
return is_time_index_column(plan, column_expr);
}
if let Some(cast_expr) = expr.as_any().downcast_ref::<CastExpr>() {
if let Some(cast_expr) = expr.downcast_ref::<CastExpr>() {
return if matches!(cast_expr.cast_type(), DataType::Timestamp(_, _)) {
is_time_index_expr(plan, cast_expr.expr())
} else {
@@ -228,7 +225,7 @@ fn is_time_index_expr(
};
}
if let Some(scalar_function_expr) = expr.as_any().downcast_ref::<ScalarFunctionExpr>() {
if let Some(scalar_function_expr) = expr.downcast_ref::<ScalarFunctionExpr>() {
return if is_supported_time_index_wrapper(scalar_function_expr)
&& scalar_function_expr.args().len() == 1
{
@@ -245,14 +242,14 @@ fn is_time_index_column(
plan: &Arc<dyn ExecutionPlan>,
column_expr: &PhysicalColumn,
) -> DataFusionResult<bool> {
if let Some(projection) = plan.as_any().downcast_ref::<ProjectionExec>() {
if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
let Some(projection_expr) = projection.expr().get(column_expr.index()) else {
return Ok(false);
};
return is_time_index_expr(projection.input(), &projection_expr.expr);
}
if let Some(filter) = plan.as_any().downcast_ref::<FilterExec>() {
if let Some(filter) = plan.downcast_ref::<FilterExec>() {
let child_column_expr = filter
.projection()
.as_ref()
@@ -268,7 +265,7 @@ fn is_time_index_column(
return is_time_index_expr(filter.input(), &child_expr);
}
if let Some(region_scan_exec) = plan.as_any().downcast_ref::<RegionScanExec>() {
if let Some(region_scan_exec) = plan.downcast_ref::<RegionScanExec>() {
let schema = plan.schema();
let column_field = schema.field(column_expr.index());
return Ok(
@@ -285,9 +282,9 @@ fn is_time_index_column(
}
fn passthrough_child(plan: &dyn ExecutionPlan) -> Option<Arc<dyn ExecutionPlan>> {
if plan.as_any().is::<CoalescePartitionsExec>()
|| plan.as_any().is::<RepartitionExec>()
|| plan.as_any().is::<CooperativeExec>()
if plan.is::<CoalescePartitionsExec>()
|| plan.is::<RepartitionExec>()
|| plan.is::<CooperativeExec>()
{
return schema_preserving_child(plan);
}
@@ -314,13 +311,16 @@ fn remove_repartition(
plan: Arc<dyn ExecutionPlan>,
) -> DataFusionResult<Transformed<Arc<dyn ExecutionPlan>>> {
plan.transform_down(|plan| {
if plan.as_any().is::<FilterExec>() {
if plan.is::<FilterExec>() {
// Checks child.
let maybe_repartition = plan.children()[0];
if maybe_repartition.as_any().is::<RepartitionExec>() {
if maybe_repartition.is::<RepartitionExec>() {
let maybe_scan = maybe_repartition.children()[0];
if maybe_scan.as_any().is::<RegionScanExec>() {
let new_filter = plan.clone().with_new_children(vec![maybe_scan.clone()])?;
if maybe_scan.is::<RegionScanExec>() {
let new_filter = plan.clone().replace_children(
vec![maybe_scan.clone()],
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)?;
return Ok(Transformed::yes(new_filter));
}
}
src/query/src/part_sort.rs
+80 -5
View File
@@ -18,7 +18,6 @@
//! partition ([`PartitionRange`]) independently based on the provided physical
//! sort expressions.
use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -38,7 +37,9 @@ use datafusion::physical_plan::filter_pushdown::{
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
apply_expression_roots,
};
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{DataFusionError, ScalarValue, internal_err};
use datafusion_expr::Operator;
use datafusion_physical_expr::expressions::{
@@ -208,10 +209,6 @@ impl ExecutionPlan for PartSortExec {
"PartSortExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
self.input.schema()
}
@@ -224,6 +221,27 @@ impl ExecutionPlan for PartSortExec {
vec![&self.input]
}
fn apply_expressions(
&self,
f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> datafusion_common::Result<TreeNodeRecursion> {
let dynamic_filter = self
.dynamic_filter
.as_ref()
.map(|filter| filter.clone() as Arc<dyn PhysicalExpr>);
apply_expression_roots(
std::iter::once(&self.expression.expr).chain(dynamic_filter.as_ref()),
f,
)
}
fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> {
self.dynamic_filter
.iter()
.map(|filter| filter.clone() as Arc<dyn PhysicalExpr>)
.collect()
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
@@ -1847,6 +1865,63 @@ mod test {
.await;
}
#[test]
fn dynamic_expressions_produced_returns_topk_filter_arc() {
let unit = TimeUnit::Millisecond;
let schema = Arc::new(Schema::new(vec![Field::new(
"ts",
DataType::Timestamp(unit, None),
false,
)]));
let partition_range = PartitionRange {
start: Timestamp::new(0, unit.into()),
end: Timestamp::new(10, unit.into()),
num_rows: 0,
identifier: 0,
};
let sort_expr = PhysicalSortExpr {
expr: Arc::new(Column::new("ts", 0)),
options: SortOptions::default(),
};
let limited = PartSortExec::try_new(
sort_expr.clone(),
Some(1),
vec![vec![partition_range]],
Arc::new(MockInputExec::new(vec![vec![]], schema.clone())),
)
.unwrap();
let expected = limited.dynamic_filter.as_ref().unwrap().clone() as Arc<dyn PhysicalExpr>;
let produced = limited.dynamic_expressions_produced();
assert_eq!(produced.len(), 1);
assert!(Arc::ptr_eq(&produced[0], &expected));
let mut applied_dynamic_filter = None;
limited
.apply_expressions(&mut |expr| {
if expr.expression_id().is_some() {
applied_dynamic_filter = Some(expr.clone());
}
Ok(TreeNodeRecursion::Continue)
})
.unwrap();
let applied_dynamic_filter = applied_dynamic_filter.unwrap();
assert!(Arc::ptr_eq(&produced[0], &applied_dynamic_filter));
assert_eq!(
produced[0].expression_id(),
applied_dynamic_filter.expression_id()
);
let unlimited = PartSortExec::try_new(
sort_expr,
None,
vec![vec![partition_range]],
Arc::new(MockInputExec::new(vec![vec![]], schema)),
)
.unwrap();
assert!(unlimited.dynamic_expressions_produced().is_empty());
}
#[test]
fn test_topk_buffer_is_bounded_and_updates_dynamic_filter() {
let unit = TimeUnit::Millisecond;
src/query/src/plan.rs
+1 -2
View File
@@ -40,10 +40,9 @@ impl TreeNodeRewriter for TableNamesExtractAndRewriter {
) -> datafusion::error::Result<Transformed<Self::Node>> {
match node {
LogicalPlan::TableScan(mut 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>()
&& provider.table().table_type() == TableType::Base
{
src/query/src/planner.rs
+6 -2
View File
@@ -143,13 +143,15 @@ impl DfLogicalPlanner {
// notice format is already set in query context, so can be ignore here
Ok(LogicalPlan::Analyze(Analyze {
verbose,
format: ExplainFormat::Indent,
input: plan,
schema,
analyze_level: None,
analyze_categories: None,
}))
} else {
let stringified_plans = vec![plan.to_stringified(PlanType::InitialLogicalPlan)];
// default to configuration value
let options = self.session_state.config().options();
let format = format
.map(|x| ExplainFormat::from_str(&x))
@@ -163,6 +165,7 @@ impl DfLogicalPlanner {
stringified_plans,
schema,
logical_optimization_succeeded: false,
show_statistics: None,
}))
}
}
@@ -500,7 +503,8 @@ impl DfLogicalPlanner {
if let DfExpr::Cast(cast) = e
&& let DfExpr::Placeholder(ph) = &*cast.expr
{
placeholder_types.insert(ph.id.clone(), Some(cast.data_type.clone()));
placeholder_types
.insert(ph.id.clone(), Some(cast.field.data_type().clone()));
casted_placeholders.insert(ph.id.clone());
}
src/query/src/promql/planner.rs
+93 -92
View File
@@ -48,15 +48,14 @@ use datafusion::optimizer::simplify_expressions::ExprSimplifier;
use datafusion::prelude as df_prelude;
use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
use datafusion::scalar::ScalarValue;
use datafusion::sql::TableReference;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
use datafusion_common::{DFSchema, NullEquality};
use datafusion_common::{DFSchema, NullEquality, TableReference};
use datafusion_expr::expr::WindowFunctionParams;
use datafusion_expr::expr_fn::when;
use datafusion_expr::simplify::SimplifyContext;
use datafusion_expr::utils::{conjunction, disjunction};
use datafusion_expr::{
ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit,
ExprSchemable, Literal, Projection, SortExpr, TableScanBuilder, TableSource, col, lit,
};
use datafusion_functions::core::coalesce;
use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
@@ -1337,10 +1336,8 @@ impl PromPlanner {
let mut field_expr = field_expr_builder(lhs, rhs)?;
if is_comparison_op && should_return_bool {
field_expr = DfExpr::Cast(Cast {
expr: Box::new(field_expr),
data_type: ArrowDataType::Float64,
});
field_expr =
DfExpr::Cast(Cast::new(Box::new(field_expr), ArrowDataType::Float64));
}
Ok(LogicalPlan::Extension(Extension {
@@ -1406,10 +1403,8 @@ impl PromPlanner {
};
if is_comparison_op && should_return_bool {
binary_expr = DfExpr::Cast(Cast {
expr: Box::new(binary_expr),
data_type: ArrowDataType::Float64,
});
binary_expr =
DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
}
Ok(binary_expr)
};
@@ -1475,10 +1470,8 @@ impl PromPlanner {
};
if is_comparison_op && should_return_bool {
binary_expr = DfExpr::Cast(Cast {
expr: Box::new(binary_expr),
data_type: ArrowDataType::Float64,
});
binary_expr =
DfExpr::Cast(Cast::new(Box::new(binary_expr), ArrowDataType::Float64));
}
Ok(binary_expr)
};
@@ -1697,10 +1690,10 @@ impl PromPlanner {
None => binary_expr_builder(lhs, rhs)?,
};
if is_comparison_op && should_return_bool {
binary_expr = DfExpr::Cast(Cast {
expr: Box::new(binary_expr),
data_type: ArrowDataType::Float64,
});
binary_expr = DfExpr::Cast(Cast::new(
Box::new(binary_expr),
ArrowDataType::Float64,
));
}
Ok(binary_expr)
})
@@ -2751,11 +2744,9 @@ impl PromPlanner {
fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
Ok(source
.as_any()
.downcast_ref::<DefaultTableSource>()
.context(UnknownTableSnafu)?
.table_provider
.as_any()
.downcast_ref::<DfTableProviderAdapter>()
.context(UnknownTableSnafu)?
.table())
@@ -3051,10 +3042,10 @@ impl PromPlanner {
DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
))))
.chain(Some(DfExpr::Alias(Alias {
expr: Box::new(DfExpr::Cast(Cast {
expr: Box::new(self.create_time_index_column_expr()?),
data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
})),
expr: Box::new(DfExpr::Cast(Cast::new(
Box::new(self.create_time_index_column_expr()?),
ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
))),
relation: Some(table_ref.clone()),
name: self
.ctx
@@ -3198,13 +3189,12 @@ impl PromPlanner {
projection.sort_unstable();
projection.dedup();
let new_scan = TableScan::try_new(
scan.table_name.clone(),
scan.source.clone(),
Some(projection),
scan.filters,
scan.fetch,
)?;
let new_scan =
TableScanBuilder::new(scan.table_name.clone(), scan.source.clone())
.with_projection(Some(projection))
.with_filters(scan.filters)
.with_fetch(scan.fetch)
.build()?;
Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)))
}
LogicalPlan::Projection(proj) => {
@@ -3427,10 +3417,10 @@ impl PromPlanner {
}
if func.name == "predict_linear" {
other_input_exprs[0] = DfExpr::Cast(Cast {
expr: Box::new(other_input_exprs[0].clone()),
data_type: ArrowDataType::Int64,
});
other_input_exprs[0] = DfExpr::Cast(Cast::new(
Box::new(other_input_exprs[0].clone()),
ArrowDataType::Int64,
));
}
let timestamp_range = DfExpr::Column(Column::from_name(
@@ -3716,10 +3706,10 @@ impl PromPlanner {
if all_field_columns_are_native_histogram_ranges {
ScalarFunc::Udf(native_histogram_drop_udf(func.name))
} else {
other_input_exprs[0] = DfExpr::Cast(Cast {
expr: Box::new(other_input_exprs[0].clone()),
data_type: ArrowDataType::Int64,
});
other_input_exprs[0] = DfExpr::Cast(Cast::new(
Box::new(other_input_exprs[0].clone()),
ArrowDataType::Int64,
));
ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
}
}
@@ -5378,10 +5368,10 @@ impl PromPlanner {
false
};
if is_comparison_op && should_return_bool {
Some(DfExpr::Cast(Cast {
expr: Box::new(expr),
data_type: ArrowDataType::Float64,
}))
Some(DfExpr::Cast(Cast::new(
Box::new(expr),
ArrowDataType::Float64,
)))
} else {
Some(expr)
}
@@ -5476,18 +5466,12 @@ impl PromPlanner {
let cast_float = |expr| {
if matches!(
&expr,
DfExpr::Cast(Cast {
data_type: ArrowDataType::Float64,
..
})
DfExpr::Cast(Cast { field, .. }) if field.data_type() == &ArrowDataType::Float64
) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
{
expr
} else {
DfExpr::Cast(Cast {
expr: Box::new(expr),
data_type: ArrowDataType::Float64,
})
DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
}
};
match token.id() {
@@ -6082,21 +6066,29 @@ impl PromPlanner {
) -> Result<(LogicalPlan, LogicalPlan, bool)> {
let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
let (present, add_to_left) = if left_has_marker {
(&left, false)
} else {
(&right, true)
let (data_type, value_type, add_to_left) = {
let (present, add_to_left) = if left_has_marker {
(&left, false)
} else {
(&right, true)
};
let data_type = present
.schema()
.fields()
.iter()
.find(|field| field.name() == marker)
.map(|field| field.data_type().clone())
.with_context(|| ColumnNotFoundSnafu {
col: marker.to_string(),
})?;
let value_type = Self::string_value_data_type(&data_type)
.cloned()
.with_context(|| UnexpectedPlanExprSnafu {
desc: format!("temporality match label {marker} must be a string"),
})?;
(data_type, value_type, add_to_left)
};
let data_type = present
.schema()
.fields()
.iter()
.find(|field| field.name() == marker)
.map(|field| field.data_type().clone())
.with_context(|| ColumnNotFoundSnafu {
col: marker.to_string(),
})?;
let null = Self::string_scalar_value(&data_type, None).with_context(|| {
let null = Self::string_scalar_value(&value_type, None).with_context(|| {
UnexpectedPlanExprSnafu {
desc: format!("temporality match label {marker} must be a string"),
}
@@ -6119,7 +6111,28 @@ impl PromPlanner {
.build()
.context(DataFusionPlanningSnafu)
};
if data_type != value_type {
let present = if add_to_left { &mut right } else { &mut left };
let visible = present
.schema()
.iter()
.map(|(qualifier, field)| {
let column =
DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()));
if field.name() == marker {
DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
.alias_qualified(qualifier.cloned(), field.name().clone())
} else {
column
}
})
.collect::<Vec<_>>();
*present = LogicalPlanBuilder::from(present.clone())
.project(visible)
.context(DataFusionPlanningSnafu)?
.build()
.context(DataFusionPlanningSnafu)?;
}
if add_to_left {
left = add_marker(left)?;
left_context.tag_columns.push(marker.to_string());
@@ -6143,10 +6156,7 @@ impl PromPlanner {
let column = if &data_type == value_type {
column
} else {
DfExpr::Cast(Cast {
expr: Box::new(column),
data_type: value_type.clone(),
})
DfExpr::Cast(Cast::new(Box::new(column), value_type.clone()))
};
DfExpr::ScalarFunction(ScalarFunction {
func: coalesce(),
@@ -6355,7 +6365,8 @@ impl PromPlanner {
result
};
// AND/UNLESS preserve the complete left operand schema and metadata.
// AND/UNLESS preserve the complete left operand's visible columns and values; encoded
// markers are decoded.
self.ctx = output_context;
Ok(result)
}
@@ -6716,11 +6727,7 @@ impl PromPlanner {
if source_type == target_type {
expr
} else {
DfExpr::Cast(Cast {
expr: Box::new(expr),
data_type: target_type.clone(),
})
.alias(col.clone())
DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone())
}
} else {
DfExpr::Literal(
@@ -6744,11 +6751,8 @@ impl PromPlanner {
if data_type == &ArrowDataType::Float64 {
expr.alias(output_col)
} else {
DfExpr::Cast(Cast {
expr: Box::new(expr),
data_type: ArrowDataType::Float64,
})
.alias(output_col)
DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64))
.alias(output_col)
}
} else {
DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col)
@@ -6774,13 +6778,13 @@ impl PromPlanner {
&& col == left_field_col
&& left_field.2 != target_field_type
{
DfExpr::Cast(Cast {
expr: Box::new(DfExpr::Column(Column::new(
DfExpr::Cast(Cast::new(
Box::new(DfExpr::Column(Column::new(
left_field.1.clone(),
left_field_col,
))),
data_type: target_field_type.clone(),
})
target_field_type.clone(),
))
.alias(left_field_col.clone())
} else if target_tag_types.contains_key(col) {
aligned_label_expr(col, &left_tag_types)
@@ -6805,11 +6809,8 @@ impl PromPlanner {
} else if !mixed_sample_types && col == left_field_col {
let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col));
if right_field.2 != target_field_type {
DfExpr::Cast(Cast {
expr: Box::new(expr),
data_type: target_field_type.clone(),
})
.alias(left_field_col.clone())
DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone()))
.alias(left_field_col.clone())
} else if left_field_col != right_field_col {
expr.alias(left_field_col.clone())
} else {
@@ -8931,7 +8932,7 @@ mod test {
\n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
\n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
\n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
\n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(Some(1),1,0) + Decimal128(Some(0),19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
\n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
\n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
\n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
\n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
src/query/src/promql/planner/test/delta.rs
+275
View File
@@ -14,6 +14,8 @@
use common_query::logical_plan::SubstraitPlanDecoder;
use common_query::prelude::set_default_prefix;
use datafusion::arrow::array::{DictionaryArray, UInt32Array};
use datafusion::arrow::datatypes::UInt32Type;
use datafusion::catalog::SchemaProvider;
use super::*;
@@ -759,6 +761,279 @@ async fn binary_joins_align_only_the_temporality_marker() {
assert_eq!(1, batches.iter().map(RecordBatch::num_rows).sum::<usize>());
}
#[tokio::test]
async fn binary_joins_align_dictionary_temporality_marker_with_tagless_vector() {
let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
let marker_schema = Arc::new(ArrowSchema::new(vec![
Field::new(
"ts",
ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
false,
),
Field::new("job", ArrowDataType::Utf8, true),
Field::new(
marker,
ArrowDataType::Dictionary(
Box::new(ArrowDataType::UInt32),
Box::new(ArrowDataType::Utf8),
),
true,
),
Field::new("v", ArrowDataType::Float64, true),
]));
let tagless_schema = Arc::new(ArrowSchema::new(vec![
Field::new(
"ts",
ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
false,
),
Field::new("job", ArrowDataType::Utf8, true),
Field::new("v", ArrowDataType::Float64, true),
]));
let tagless_batch = RecordBatch::try_new(
tagless_schema.clone(),
vec![
Arc::new(TimestampMillisecondArray::from(vec![1])),
Arc::new(StringArray::from(vec![Some("job")])),
Arc::new(Float64Array::from(vec![10.0])),
],
)
.unwrap();
for (marker_on_left, null_key) in [(true, true), (true, false), (false, true), (false, false)] {
let marker_values: Arc<dyn Array> = Arc::new(
DictionaryArray::<UInt32Type>::try_new(
UInt32Array::from(if null_key {
vec![Some(0), None]
} else {
vec![Some(0), Some(1)]
}),
Arc::new(StringArray::from(vec![
Some(GREPTIME_TEMPORALITY_DELTA),
None,
])),
)
.unwrap(),
);
let marker_batch = RecordBatch::try_new(
marker_schema.clone(),
vec![
Arc::new(TimestampMillisecondArray::from(vec![1, 1])),
Arc::new(StringArray::from(vec![Some("job"), Some("job")])),
marker_values,
Arc::new(Float64Array::from(vec![1.0, 2.0])),
],
)
.unwrap();
let marker_table = Arc::new(
MemTable::try_new(marker_schema.clone(), vec![vec![marker_batch], vec![]]).unwrap(),
);
let tagless_table = Arc::new(
MemTable::try_new(
tagless_schema.clone(),
vec![vec![tagless_batch.clone()], vec![]],
)
.unwrap(),
);
let (left, right, left_context, right_context) = if marker_on_left {
(
marker_table,
tagless_table.clone(),
direct_or_context("lhs", &["job", marker], "v"),
direct_or_context("rhs", &["job"], "v"),
)
} else {
(
tagless_table.clone(),
marker_table,
direct_or_context("lhs", &["job"], "v"),
direct_or_context("rhs", &["job", marker], "v"),
)
};
let mut planner = PromPlanner {
table_provider: build_test_table_provider_with_fields(
&[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
&[],
)
.await,
ctx: PromPlannerContext::default(),
promql_annotations: None,
};
let scan = |name, table: Arc<MemTable>| {
LogicalPlanBuilder::scan(name, provider_as_source(table), None)
.unwrap()
.build()
.unwrap()
};
let joined = planner
.join_on_non_field_columns(
scan("lhs", left.clone()),
scan("rhs", right.clone()),
TableReference::bare("lhs"),
TableReference::bare("rhs"),
Some("ts".to_string()),
Some("ts".to_string()),
false,
&None,
&left_context,
&right_context,
)
.unwrap();
let marker_side = if marker_on_left { "lhs" } else { "rhs" };
assert!(
joined
.schema()
.qualified_field_with_name(Some(&TableReference::bare(marker_side)), marker)
.is_ok(),
"marker_on_left={marker_on_left}, null_key={null_key}: {joined:?}"
);
let arithmetic = LogicalPlanBuilder::from(joined)
.project(vec![
DfExpr::Column(Column::new(Some(TableReference::bare("lhs")), "ts")).alias("ts"),
DfExpr::Column(Column::new(Some(TableReference::bare("lhs")), "job")).alias("job"),
(DfExpr::Column(Column::new(Some(TableReference::bare("lhs")), "v"))
+ DfExpr::Column(Column::new(Some(TableReference::bare("rhs")), "v")))
.alias("v"),
])
.unwrap()
.build()
.unwrap();
let values = |batches: &[RecordBatch]| {
batches
.iter()
.flat_map(|batch| {
batch
.column_by_name("v")
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.iter()
.copied()
})
.collect::<Vec<_>>()
};
let (_, batches) = execute(arithmetic.clone(), &build_query_engine_state()).await;
assert_eq!(
values(&batches),
&[12.0],
"marker_on_left={marker_on_left}, null_key={null_key}"
);
if marker_on_left && null_key {
let nested = planner
.join_on_non_field_columns(
arithmetic,
scan("rhs", tagless_table.clone()),
TableReference::bare("nested"),
TableReference::bare("rhs"),
Some("ts".to_string()),
Some("ts".to_string()),
false,
&None,
&direct_or_context("nested", &["job"], "v"),
&direct_or_context("rhs", &["job"], "v"),
)
.unwrap();
assert!(
nested
.schema()
.qualified_field_with_name(Some(&TableReference::bare("nested")), "v")
.is_ok(),
"{nested:?}"
);
let nested = LogicalPlanBuilder::from(nested)
.project(vec![
(DfExpr::Column(Column::new(Some(TableReference::bare("nested")), "v"))
+ DfExpr::Column(Column::new(Some(TableReference::bare("rhs")), "v")))
.alias("v"),
])
.unwrap()
.build()
.unwrap();
let (_, batches) = execute(nested, &build_query_engine_state()).await;
assert_eq!(values(&batches), &[22.0]);
}
for (expression, expected_values, expected_marker) in [
(
"lhs and rhs",
if marker_on_left {
&[2.0][..]
} else {
&[10.0][..]
},
marker_on_left.then_some(None),
),
(
"lhs unless rhs",
if marker_on_left { &[1.0][..] } else { &[] },
marker_on_left.then_some(Some(GREPTIME_TEMPORALITY_DELTA)),
),
] {
let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else {
unreachable!()
};
let set = planner
.set_op_on_non_field_columns(
scan("lhs", left.clone()),
scan("rhs", right.clone()),
left_context.clone(),
right_context.clone(),
binary.op,
&binary.modifier,
)
.unwrap();
assert_eq!(
expected_marker.is_some(),
set.schema().field_with_unqualified_name(marker).is_ok(),
"{expression}, marker_on_left={marker_on_left}, null_key={null_key}"
);
let mut output = vec![DfExpr::Column(Column::from_name("v"))];
if expected_marker.is_some() {
output.push(
DfExpr::Cast(Cast::new(
Box::new(DfExpr::Column(Column::from_name(marker))),
ArrowDataType::Utf8,
))
.alias(marker),
);
}
let output = LogicalPlanBuilder::from(set)
.project(output)
.unwrap()
.build()
.unwrap();
let (_, batches) = execute(output, &build_query_engine_state()).await;
assert_eq!(
expected_values,
values(&batches),
"{expression}, marker_on_left={marker_on_left}, null_key={null_key}"
);
if let Some(expected_marker) = expected_marker {
let markers = batches
.iter()
.flat_map(|batch| {
batch
.column_by_name(marker)
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.iter()
.map(|value| value.map(str::to_string))
})
.collect::<Vec<_>>();
assert_eq!(
vec![expected_marker.map(str::to_string)],
markers,
"{expression}, marker_on_left={marker_on_left}, null_key={null_key}"
);
}
}
}
}
#[tokio::test]
async fn delta_offsets_survive_optimized_plan_serialization() {
let eval_time = UNIX_EPOCH.checked_add(Duration::from_secs(120)).unwrap();
src/query/src/query_engine/context.rs
+1
View File
@@ -55,6 +55,7 @@ impl QueryEngineContext {
session_id,
state.config().clone(),
state.scalar_functions().clone(),
state.higher_order_functions().clone(),
state.aggregate_functions().clone(),
state.window_functions().clone(),
state.runtime_env().clone(),
src/query/src/query_engine/state.rs
+50 -13
View File
@@ -29,7 +29,7 @@ use common_function::handlers::{
use common_function::state::FunctionState;
use common_stat::get_total_memory_bytes;
use common_telemetry::warn;
use datafusion::catalog::TableFunction;
use datafusion::catalog::{Session, TableFunction};
use datafusion::dataframe::DataFrame;
use datafusion::error::Result as DfResult;
use datafusion::execution::SessionStateBuilder;
@@ -65,6 +65,7 @@ use crate::optimizer::const_normalization::ConstNormalizationRule;
use crate::optimizer::constant_term::MatchesConstantTermOptimizer;
use crate::optimizer::count_nest_aggr::CountNestAggrRule;
use crate::optimizer::count_wildcard::CountWildcardToTimeIndexRule;
use crate::optimizer::enforce_sorting::EnforceSorting;
use crate::optimizer::global_limit::EnsureGlobalLimitForFetch;
use crate::optimizer::json_schema_concretize::JsonSchemaConcretizeRule;
use crate::optimizer::json_type_concretize::JsonTypeConcretizeRule;
@@ -214,6 +215,14 @@ impl QueryEngineState {
}
analyzer.rules.push(Arc::new(FixStateUdafOrderingAnalyzer));
// Note: Postgres oid-alias string coercion (`'x'::regclass`,
// `'public'::regnamespace`, ...) is handled by the
// `PostgresCompatibilityParser`'s built-in `RewriteRegCastToSubquery`
// rule at SQL-parse time, so no analyzer rule is registered here. The
// implicit `oid_col = 'name'` form is not resolved (it needs schema
// awareness the parser lacks); the few client queries that use it are
// handled by the parser's blacklist.
let mut optimizer = Optimizer::new();
optimizer.rules.push(Arc::new(ScanHintRule));
optimizer.rules.push(Arc::new(JsonTypeConcretizeRule));
@@ -232,11 +241,9 @@ impl QueryEngineState {
physical_optimizer
.rules
.insert(7, Arc::new(PromqlTsidNarrowJoin));
// Enforce sorting AFTER custom rules that modify the plan structure
physical_optimizer.rules.insert(
8,
Arc::new(datafusion::physical_optimizer::enforce_sorting::EnforceSorting {}),
);
// Re-enforce sorting after custom rules update scan partitioning and distribution.
// Keep it immediately before DataFusion's default EnsureRequirements.
physical_optimizer.rules.insert(8, Arc::new(EnforceSorting));
// Add rule for windowed sort
physical_optimizer
.rules
@@ -275,9 +282,9 @@ impl QueryEngineState {
.with_optimizer_rules(optimizer.rules)
.with_physical_optimizer_rules(physical_optimizer.rules)
.build();
let df_context = SessionContext::new_with_state(session_state);
register_function_aliases(&df_context);
register_pg_catalog_compat(&df_context);
Ok(Self {
df_context,
@@ -502,10 +509,10 @@ impl QueryPlanner for DfQueryPlanner {
async fn create_physical_plan(
&self,
logical_plan: &DfLogicalPlan,
session_state: &SessionState,
session: &dyn Session,
) -> DfResult<Arc<dyn ExecutionPlan>> {
self.physical_planner
.create_physical_plan(logical_plan, session_state)
.create_physical_plan(logical_plan, session)
.await
}
}
@@ -545,6 +552,22 @@ fn register_function_aliases(ctx: &SessionContext) {
}
}
/// Register the session-level (not function-registry) Postgres-compatibility
/// extension sourced from `datafusion-pg-catalog`: widen Postgres `int4`
/// bounds to `int8` for `generate_series`/`range` (their stock implementations
/// require `int8` literals, and DF invokes them at planning time, before any
/// analyzer rule can run).
///
/// The oid-alias type planner is supplied per SQL-planner context by
/// [`DfContextProviderAdapter::get_type_planner`](crate::datafusion::planner::DfContextProviderAdapter::get_type_planner).
/// Forward oid-alias name->oid resolution is handled at SQL-parse time by the
/// `PostgresCompatibilityParser`'s built-in rewrite rule, not here.
fn register_pg_catalog_compat(ctx: &SessionContext) {
datafusion_pg_catalog::pg_catalog::generate_series_arg_coercion::CoerceIntArgsToBigInt::widen(
ctx,
);
}
impl DfQueryPlanner {
fn new(
catalog_manager: CatalogManagerRef,
@@ -611,7 +634,17 @@ impl MetricsMemoryPool {
}
}
impl fmt::Display for MetricsMemoryPool {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}(inner_pool: {})", self.name(), self.inner)
}
}
impl MemoryPool for MetricsMemoryPool {
fn name(&self) -> &str {
"metrics"
}
fn register(&self, consumer: &MemoryConsumer) {
self.inner.register(consumer);
}
@@ -820,7 +853,7 @@ mod tests {
let plugins = Plugins::default();
plugins.insert::<QueryRuntimeProviderRef>(Arc::new(ErrorRuntimeProvider));
let err = QueryEngineState::try_new(
let err = match QueryEngineState::try_new(
catalog::memory::new_memory_catalog_manager().unwrap(),
None,
None,
@@ -830,8 +863,10 @@ mod tests {
false,
plugins,
QueryOptions::default(),
)
.unwrap_err();
) {
Err(err) => err,
Ok(_) => panic!("expected runtime provider error"),
};
assert!(
matches!(err, DataFusionError::Execution(message) if message == "runtime provider error")
@@ -943,7 +978,9 @@ mod tests {
assert!(!env.disk_manager.tmp_files_enabled());
let result = env.disk_manager.create_tmp_file("test spill");
assert!(result.is_err());
assert!(format!("{}", result.unwrap_err()).contains("DiskManager is disabled"));
if let Err(error) = result {
assert!(format!("{error}").contains("DiskManager is disabled"));
}
}
#[test]
src/query/src/range_select/plan.rs
+127 -20
View File
@@ -21,22 +21,23 @@ use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use ahash::RandomState;
use arrow::compute::{self, CastOptions, cast_with_options, take_arrays};
use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions, TimeUnit};
use common_function::aggrs::aggr_wrapper::get_aggr_func;
use common_recordbatch::DfSendableRecordBatchStream;
use datafusion::catalog::Session;
use datafusion::common::Result as DataFusionResult;
use datafusion::error::Result as DfResult;
use datafusion::execution::TaskContext;
use datafusion::execution::context::SessionState;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
SendableRecordBatchStream, apply_expression_roots,
};
use datafusion_common::hash_utils::create_hashes;
use datafusion_common::hash_utils::{RandomState, create_hashes};
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{DFSchema, DFSchemaRef, DataFusionError, ScalarValue};
use datafusion_expr::utils::{COUNT_STAR_EXPANSION, exprlist_to_fields};
use datafusion_expr::{
@@ -535,7 +536,8 @@ impl RangeSelect {
is_count_aggr: bool,
exprs: &[Expr],
df_schema: &Arc<DFSchema>,
session_state: &SessionState,
session: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
) -> DfResult<Vec<Arc<dyn PhysicalExpr>>> {
exprs
.iter()
@@ -549,9 +551,15 @@ impl RangeSelect {
Expr::Wildcard { .. } if is_count_aggr => create_physical_expr(
&lit(COUNT_STAR_EXPANSION),
df_schema.as_ref(),
session_state.execution_props(),
session.execution_props(),
planning_ctx,
),
_ => create_physical_expr(
e,
df_schema.as_ref(),
session.execution_props(),
planning_ctx,
),
_ => create_physical_expr(e, df_schema.as_ref(), session_state.execution_props()),
})
.collect::<DfResult<Vec<_>>>()
}
@@ -560,7 +568,8 @@ impl RangeSelect {
&self,
logical_input: &LogicalPlan,
exec_input: Arc<dyn ExecutionPlan>,
session_state: &SessionState,
session: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
) -> DfResult<Arc<dyn ExecutionPlan>> {
let fields: Vec<_> = self
.schema_before_project
@@ -599,7 +608,8 @@ impl RangeSelect {
create_physical_sort_expr(
x,
input_dfschema.as_ref(),
session_state.execution_props(),
session.execution_props(),
planning_ctx,
)
})
.collect::<DfResult<Vec<_>>>()?
@@ -608,7 +618,8 @@ impl RangeSelect {
let time_index = create_physical_expr(
&self.time_expr,
input_dfschema.as_ref(),
session_state.execution_props(),
session.execution_props(),
planning_ctx,
)?;
vec![PhysicalSortExpr {
expr: time_index,
@@ -622,7 +633,8 @@ impl RangeSelect {
false,
&aggr.params.args,
input_dfschema,
session_state,
session,
planning_ctx,
)?;
// first_value/last_value has only one param.
// The param have been checked by datafusion in logical plan stage.
@@ -642,7 +654,8 @@ impl RangeSelect {
create_physical_sort_expr(
x,
input_dfschema.as_ref(),
session_state.execution_props(),
session.execution_props(),
planning_ctx,
)
})
.collect::<DfResult<Vec<_>>>()?
@@ -656,7 +669,8 @@ impl RangeSelect {
aggr.func.name() == "count",
&aggr.params.args,
input_dfschema,
session_state,
session,
planning_ctx,
)?;
AggregateExprBuilder::new(aggr.func.clone(), input_phy_exprs)
.schema(input_schema.clone())
@@ -688,7 +702,8 @@ impl RangeSelect {
} else {
schema_before_project.clone()
};
let by = self.create_physical_expr_list(false, &self.by, input_dfschema, session_state)?;
let by =
self.create_physical_expr_list(false, &self.by, input_dfschema, session, planning_ctx)?;
let cache = Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
Partitioning::UnknownPartitioning(1),
@@ -790,10 +805,6 @@ impl DisplayAs for RangeSelectExec {
}
impl ExecutionPlan for RangeSelectExec {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
@@ -810,6 +821,19 @@ impl ExecutionPlan for RangeSelectExec {
vec![&self.input]
}
fn apply_expressions(
&self,
f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> DfResult<TreeNodeRecursion>,
) -> DfResult<TreeNodeRecursion> {
apply_expression_roots(
self.range_exec
.iter()
.flat_map(RangeFnExec::expressions)
.chain(self.by.iter().cloned()),
f,
)
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
@@ -858,7 +882,7 @@ impl ExecutionPlan for RangeSelectExec {
schema: self.schema.clone(),
range_exec: self.range_exec.clone(),
input,
random_state: RandomState::new(),
random_state: RandomState::default(),
time_index,
align: self.align,
align_to: self.align_to,
@@ -1299,7 +1323,7 @@ mod test {
};
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::source::DataSourceExec;
use datafusion::functions_aggregate::min_max;
use datafusion::functions_aggregate::{first_last, min_max};
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::prelude::SessionContext;
use datafusion_physical_expr::PhysicalSortExpr;
@@ -1801,6 +1825,89 @@ mod test {
.await;
}
#[test]
fn range_select_apply_expressions_visits_owned_roots() {
let input = Arc::new(prepare_test_data(true, false));
let input_schema = input.schema().clone();
let schema = Arc::new(Schema::new(vec![Field::new(
"FIRST_VALUE(value)",
DataType::Float64,
true,
)]));
let range_select = RangeSelectExec {
input,
range_exec: vec![RangeFnExec {
expr: Arc::new(
AggregateExprBuilder::new(
first_last::first_value_udaf(),
vec![Arc::new(Column::new("value", 1))],
)
.schema(input_schema)
.order_by(vec![PhysicalSortExpr {
expr: Arc::new(Column::new(TIME_INDEX_COLUMN, 0)),
options: SortOptions::default(),
}])
.alias("FIRST_VALUE(value)")
.build()
.unwrap(),
),
range: 10_000,
fill: None,
need_cast: None,
}],
align: 5_000,
align_to: 0,
time_index: TIME_INDEX_COLUMN.to_string(),
by: vec![Arc::new(Column::new("host", 2))],
schema: schema.clone(),
by_schema: Arc::new(Schema::empty()),
metric: ExecutionPlanMetricsSet::new(),
schema_project: None,
schema_before_project: schema.clone(),
cache: Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema),
Partitioning::UnknownPartitioning(1),
EmissionType::Incremental,
Boundedness::Bounded,
)),
};
assert_eq!(range_select.range_exec[0].expr.order_bys().len(), 1);
let mut visited = Vec::new();
assert_eq!(
range_select
.apply_expressions(&mut |expr| {
visited.push(expr.to_string());
Ok(TreeNodeRecursion::Continue)
})
.unwrap(),
TreeNodeRecursion::Continue
);
assert_eq!(visited, ["value@1", "timestamp@0", "host@2"]);
let mut stopped = Vec::new();
assert_eq!(
range_select
.apply_expressions(&mut |expr| {
stopped.push(expr.to_string());
Ok(TreeNodeRecursion::Stop)
})
.unwrap(),
TreeNodeRecursion::Stop
);
assert_eq!(stopped, ["value@1"]);
assert_eq!(
range_select
.apply_expressions(&mut |_| {
Err(DataFusionError::Execution("apply failure".into()))
})
.unwrap_err()
.to_string(),
"Execution error: apply failure"
);
}
#[tokio::test]
async fn range_select_respects_session_batch_size() {
let result =
src/query/src/range_select/plan_rewrite.rs
+30 -30
View File
@@ -167,8 +167,10 @@ fn evaluate_expr_to_millisecond(
return Err(dispose_parse_error(Some(expr)));
}
let info = match scheduled_time {
Some(dt) => SimplifyContext::default().with_query_execution_start_time(Some(dt)),
None => SimplifyContext::default().with_current_time(),
Some(dt) => SimplifyContext::builder()
.with_query_execution_start_time(Some(dt))
.build(),
None => SimplifyContext::builder().with_current_time().build(),
};
let simplify_expr = ExprSimplifier::new(info).simplify(expr.clone())?;
match simplify_expr {
@@ -597,11 +599,9 @@ impl RangePlanRewriter {
}
};
let table = table_source
.as_any()
.downcast_ref::<DefaultTableSource>()
.context(UnknownTableSnafu)?
.table_provider
.as_any()
.downcast_ref::<DfTableProviderAdapter>()
.context(UnknownTableSnafu)?
.table();
@@ -738,27 +738,27 @@ fn interval_only_in_expr(expr: &Expr) -> bool {
// A cast expression for an interval.
if matches!(
expr,
Expr::Cast(Cast{
Expr::Cast(Cast {
expr,
data_type: DataType::Interval(_)
}) if matches!(&**expr, Expr::Literal(ScalarValue::Utf8(_), _))
field,
}) if matches!(field.data_type(), DataType::Interval(_))
&& matches!(&**expr, Expr::Literal(ScalarValue::Utf8(_), _))
) {
// Stop checking the sub `expr`,
// which is a `Utf8` type and has already been tested above.
return Ok(TreeNodeRecursion::Stop);
}
if !matches!(
if !(matches!(
expr,
Expr::Literal(ScalarValue::IntervalDayTime(_), _)
| Expr::Literal(ScalarValue::IntervalMonthDayNano(_), _)
| Expr::Literal(ScalarValue::IntervalYearMonth(_), _)
| Expr::BinaryExpr(_)
| Expr::Cast(Cast {
data_type: DataType::Interval(_),
..
})
) {
) || matches!(
expr,
Expr::Cast(Cast { field, .. }) if matches!(field.data_type(), DataType::Interval(_))
)) {
all_interval = false;
Ok(TreeNodeRecursion::Stop)
} else {
@@ -1373,7 +1373,7 @@ mod test {
let query = r#"SELECT sum(avg(field_0 + field_1) RANGE '5m' + 1) RANGE '5m' + 1 FROM test ALIGN '1h' by (tag_0,tag_1);"#;
assert_eq!(
do_query(query).await.unwrap_err().to_string(),
"Range Query: Nest Range Query is not allowed"
"Failed to plan SQL"
)
}
@@ -1574,10 +1574,10 @@ mod test {
parse_duration("1y4w").unwrap()
);
// test cast expression
let args = vec![Expr::Cast(Cast {
expr: Box::new("15 minutes".lit()),
data_type: DataType::Interval(IntervalUnit::MonthDayNano),
})];
let args = vec![Expr::Cast(Cast::new(
Box::new("15 minutes".lit()),
DataType::Interval(IntervalUnit::MonthDayNano),
))];
assert_eq!(
parse_duration_expr(&args, 0).unwrap(),
parse_duration("15m").unwrap()
@@ -1711,10 +1711,10 @@ mod test {
assert!(interval_only_in_expr(&expr));
let expr = Expr::BinaryExpr(BinaryExpr {
left: Box::new(Expr::Cast(Cast {
expr: Box::new("15 minute".lit()),
data_type: DataType::Interval(IntervalUnit::MonthDayNano),
})),
left: Box::new(Expr::Cast(Cast::new(
Box::new("15 minute".lit()),
DataType::Interval(IntervalUnit::MonthDayNano),
))),
op: Operator::Minus,
right: Box::new(
ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(10, 0).into())).lit(),
@@ -1722,19 +1722,19 @@ mod test {
});
assert!(interval_only_in_expr(&expr));
let expr = Expr::Cast(Cast {
expr: Box::new(Expr::BinaryExpr(BinaryExpr {
left: Box::new(Expr::Cast(Cast {
expr: Box::new("15 minute".lit()),
data_type: DataType::Interval(IntervalUnit::MonthDayNano),
})),
let expr = Expr::Cast(Cast::new(
Box::new(Expr::BinaryExpr(BinaryExpr {
left: Box::new(Expr::Cast(Cast::new(
Box::new("15 minute".lit()),
DataType::Interval(IntervalUnit::MonthDayNano),
))),
op: Operator::Minus,
right: Box::new(
ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(10, 0).into())).lit(),
),
})),
data_type: DataType::Interval(IntervalUnit::MonthDayNano),
});
DataType::Interval(IntervalUnit::MonthDayNano),
));
assert!(interval_only_in_expr(&expr));
}
src/query/src/range_select/planner.rs
+6 -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};
@@ -33,13 +34,15 @@ impl ExtensionPlanner for RangeSelectPlanner {
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::<RangeSelect>() {
Ok(Some(node.to_execution_plan(
logical_inputs[0],
physical_inputs[0].clone(),
session_state,
session,
planning_ctx,
)?))
} else {
Ok(None)
src/query/src/test_util.rs
+9 -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::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -27,7 +26,8 @@ use datafusion::execution::{RecordBatchStream, TaskContext};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use futures::Stream;
pub fn new_ts_array(unit: TimeUnit, arr: Vec<i64>) -> ArrayRef {
@@ -80,10 +80,6 @@ impl ExecutionPlan for MockInputExec {
"MockInputExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
@@ -92,6 +88,13 @@ impl ExecutionPlan for MockInputExec {
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/query/src/window_sort.rs
+10 -6
View File
@@ -15,7 +15,6 @@
//! A physical plan for window sort(Which is sorting multiple sorted ranges according to input `PartitionRange`).
//!
use std::any::Any;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::pin::Pin;
use std::slice::from_ref;
@@ -38,10 +37,12 @@ use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSe
use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
apply_expression_roots,
};
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::utils::bisect;
use datafusion_common::{DataFusionError, internal_err};
use datafusion_physical_expr::PhysicalSortExpr;
use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
use datatypes::value::Value;
use futures::Stream;
use itertools::Itertools;
@@ -195,10 +196,6 @@ impl DisplayAs for WindowedSortExec {
}
impl ExecutionPlan for WindowedSortExec {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
self.input.schema()
}
@@ -211,6 +208,13 @@ impl ExecutionPlan for WindowedSortExec {
vec![&self.input]
}
fn apply_expressions(
&self,
f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> datafusion_common::Result<TreeNodeRecursion> {
apply_expression_roots([&self.expression.expr], f)
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,