Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@ use crate::compile::{
/// - binary operations between a literal string and an expression
/// of a different type to a string casted to that type
/// - binary operations between a timestamp and a date to a timestamp and timestamp operation
/// - comparisons of a timestamp with `DATE +/- INTERVAL` arithmetic to an explicit
/// `TIMESTAMP` cast of that arithmetic
/// - IN list expressions where expression being tested is `TIMESTAMP`
/// and values might be `DATE` to values casted to `TIMESTAMP`
/// - BETWEEN expressions where expression being tested is `TIMESTAMP`
Expand Down Expand Up @@ -1342,6 +1344,8 @@ fn grouping_set_normalize(
/// - binary operations between a literal string and an expression
/// of a different type to a string casted to that type
/// - binary operations between a timestamp and a date to a timestamp and timestamp operation
/// - comparisons of a timestamp with `DATE +/- INTERVAL` arithmetic to an explicit
/// `TIMESTAMP` cast of that arithmetic
#[inline(never)]
fn binary_expr_normalize(
optimizer: &PlanNormalize,
Expand Down Expand Up @@ -1380,6 +1384,37 @@ fn binary_expr_normalize(
return Ok(Box::new(Expr::ScalarUDF { fun, args }));
}

// DataFusion types `DATE +/- INTERVAL` as `TIMESTAMP` with no cast node. Strict dialects
// (BigQuery) type that arithmetic as `DATETIME` and reject comparing it to a `TIMESTAMP`,
// so the implicit cast is made explicit here. Unlike a BETWEEN bound it is not folded:
// the arithmetic may hold `now()`-like placeholders that only the rewrite rules resolve.
if matches!(
op,
Operator::Eq
| Operator::NotEq
| Operator::Lt
| Operator::LtEq
| Operator::Gt
| Operator::GtEq
| Operator::IsDistinctFrom
| Operator::IsNotDistinctFrom
) {
let target_type = match (&left_type, &right_type) {
(DataType::Timestamp(_, _), DataType::Timestamp(_, _) | DataType::Date32) => {
Some(&left_type)
}
(DataType::Date32, DataType::Timestamp(_, _)) => Some(&right_type),
_ => None,
};
if let Some(target_type) = target_type {
let left =
normalize_temporal_operand(optimizer, left, &left_type, target_type, schema)?;
let right =
normalize_temporal_operand(optimizer, right, &right_type, target_type, schema)?;
return Ok(Box::new(Expr::BinaryExpr { left, op, right }));
}
}

// Check if the expression is `TIMESTAMP <op> DATE` or `DATE <op> TIMESTAMP`
// and cast the `DATE` to `TIMESTAMP` to match the types.
match (&left_type, &right_type) {
Expand Down Expand Up @@ -1426,6 +1461,55 @@ fn binary_expr_normalize(
Ok(Box::new(Expr::BinaryExpr { left, op, right }))
}

/// Normalizes one side of a temporal comparison to the `TIMESTAMP` type of the other side:
/// `DATE +/- INTERVAL` arithmetic gets an explicit cast, a `DATE` side is casted and folded.
fn normalize_temporal_operand(
optimizer: &PlanNormalize,
expr: Box<Expr>,
expr_type: &DataType,
target_type: &DataType,
schema: &DFSchema,
) -> Result<Box<Expr>> {
if is_date_interval_arithmetic(&expr, schema)? {
return Ok(Box::new(Expr::Cast {
expr,
data_type: target_type.clone(),
}));
}
if matches!(expr_type, DataType::Date32) {
Comment thread
claude[bot] marked this conversation as resolved.
return evaluate_expr(optimizer, expr.cast_to(target_type, schema)?);
}
Ok(expr)
}

/// Checks if the expression is `DATE +/- INTERVAL` arithmetic, possibly offset by more
/// intervals (`CURRENT_DATE - INTERVAL '1 month' + INTERVAL '1 day'`). DataFusion types
/// such an expression as `TIMESTAMP` without an explicit cast.
fn is_date_interval_arithmetic(expr: &Expr, schema: &DFSchema) -> Result<bool> {
let is_interval = |data_type: &DataType| matches!(data_type, DataType::Interval(_));
// Walk down the chain of interval offsets to the expression they apply to.
let mut expr = expr;
loop {
let Expr::BinaryExpr { left, op, right } = expr else {
return Ok(false);
};
if !matches!(op, Operator::Plus | Operator::Minus) {
return Ok(false);
}
let base = if is_interval(&right.get_type(schema)?) {
left
} else if *op == Operator::Plus && is_interval(&left.get_type(schema)?) {
right
} else {
return Ok(false);
};
if base.get_type(schema)? == DataType::Date32 {
return Ok(true);
}
expr = base;
}
}

Comment thread
claude[bot] marked this conversation as resolved.
/// Casts a string literal expression to the given type, evaluating it to a constant.
/// Timestamp targets are parsed with Cube's date parser, which accepts date-only
/// strings (e.g. `'2026-06-01'`) that the Arrow cast kernel rejects.
Expand Down Expand Up @@ -1782,6 +1866,55 @@ mod tests {
Ok(())
}

// `DATE - INTERVAL` is typed as TIMESTAMP by DataFusion without a cast; strict dialects
// produce a DATETIME there, so the implicit cast is made explicit when compared against
// a TIMESTAMP.
#[test]
fn test_binary_expr_timestamp_computed_date_bound() -> Result<()> {
run_async_test(async move {
let meta = get_test_tenant_ctx();
let cube_ctx = create_test_postgresql_cube_context(meta)
.await
.expect("Failed to create cube context");

let schema = Schema::new(vec![Field::new(
"ts",
DataType::Timestamp(TimeUnit::Nanosecond, None),
true,
)]);

let table_scan = LogicalPlanBuilder::scan_empty(Some("test_table"), &schema, None)
.expect("Failed to create table scan")
.build()
.expect("Failed to build plan");

// 2026-06-01 minus 28 days
let date = Expr::Literal(ScalarValue::Date32(Some(20605)));
let interval = Expr::Literal(ScalarValue::IntervalDayTime(Some(28i64 << 32)));
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("ts").gt_eq(date.clone() - interval.clone()))
.expect("Failed to add filter")
.build()
.expect("Failed to build plan");

let optimizer = PlanNormalize::new(&cube_ctx);
let optimized = optimizer.optimize(&plan, &OptimizerConfig::new()).unwrap();

let LogicalPlan::Filter(Filter { predicate, .. }) = &optimized else {
panic!("Expected Filter plan, got: {:?}", optimized);
};
assert_eq!(
*predicate,
col("test_table.ts").gt_eq(Expr::Cast {
expr: Box::new(date - interval),
data_type: DataType::Timestamp(TimeUnit::Nanosecond, None),
})
);
});

Ok(())
}

// A string literal that can't be casted must not fail the whole plan
// normalization; the expression is kept as is.
#[test]
Expand Down
74 changes: 68 additions & 6 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cast.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,15 @@
use crate::compile::rewrite::{
cast_expr, rewrite, rewriter::CubeRewrite, rules::wrapper::WrapperRules,
wrapper_pullup_replacer, wrapper_pushdown_replacer,
use crate::{
compile::rewrite::{
cast_expr, rewrite,
rewriter::{CubeEGraph, CubeRewrite},
rules::wrapper::WrapperRules,
transforming_rewrite, wrapper_pullup_replacer, wrapper_pushdown_replacer,
wrapper_replacer_context, CastExprDataType,
},
var, var_iter,
};
use egg::Subst;
use std::ops::ControlFlow;

impl WrapperRules {
pub fn cast_rules(&self, rules: &mut Vec<CubeRewrite>) {
Expand All @@ -11,11 +19,65 @@ impl WrapperRules {
wrapper_pushdown_replacer(cast_expr("?expr", "?data_type"), "?context"),
cast_expr(wrapper_pushdown_replacer("?expr", "?context"), "?data_type"),
),
rewrite(
transforming_rewrite(
"wrapper-pull-up-cast",
cast_expr(wrapper_pullup_replacer("?expr", "?context"), "?data_type"),
wrapper_pullup_replacer(cast_expr("?expr", "?data_type"), "?context"),
cast_expr(
wrapper_pullup_replacer(
"?expr",
wrapper_replacer_context(
"?alias_to_cube",
"?push_to_cube",
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
"?input_data_source",
),
),
"?data_type",
),
wrapper_pullup_replacer(
cast_expr("?expr", "?data_type"),
wrapper_replacer_context(
"?alias_to_cube",
"?push_to_cube",
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
"?input_data_source",
),
),
self.transform_cast_expr("?data_type", "?input_data_source"),
),
]);
}

fn transform_cast_expr(
&self,
data_type_var: &'static str,
input_data_source_var: &'static str,
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
let data_type_var = var!(data_type_var);
let input_data_source_var = var!(input_data_source_var);
let meta = self.meta_context.clone();
move |egraph, subst| {
let Ok(data_source) = Self::get_data_source(egraph, subst, input_data_source_var)
else {
return false;
};

// Rendering a cast needs both the cast template and the template of its type
Comment thread
claude[bot] marked this conversation as resolved.
let sql_generator = match Self::template_sql_generator(&data_source, &meta) {
ControlFlow::Continue(sql_generator) => sql_generator,
ControlFlow::Break(verdict) => return verdict,
};
let templates = sql_generator.get_sql_templates();
if !templates.contains_template("expressions/cast") {
return false;
}
var_iter!(egraph[subst[data_type_var]], CastExprDataType)
.any(|data_type| templates.contains_sql_type(data_type))
}
Comment thread
claude[bot] marked this conversation as resolved.
}
}
32 changes: 22 additions & 10 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,10 @@ use crate::{
},
config::ConfigObj,
singular_eclass,
transport::{DataSource, MetaContext},
transport::{DataSource, MetaContext, SqlGenerator},
};
use egg::{Subst, Var};
use std::{fmt::Display, sync::Arc};
use std::{fmt::Display, ops::ControlFlow, sync::Arc};

pub struct WrapperRules {
meta_context: Arc<MetaContext>,
Expand Down Expand Up @@ -253,17 +253,29 @@ impl WrapperRules {
}
}

fn can_rewrite_template(data_source: &DataSource, meta: &MetaContext, template: &str) -> bool {
let sql_generator = match data_source {
/// The SQL generator whose templates a wrapper context renders with. `Break` carries the
/// verdict of a template check when there is none to consult: an unrestricted context may
/// render anything, while nothing renders for a data source `meta` does not know.
fn template_sql_generator<'meta>(
data_source: &DataSource,
meta: &'meta MetaContext,
) -> ControlFlow<bool, &'meta Arc<dyn SqlGenerator + Send + Sync>> {
match data_source {
DataSource::Specific(data_source) => {
let Some(sql_generator) = meta.data_source_to_sql_generator.get(*data_source)
else {
return false;
};
sql_generator
match meta.data_source_to_sql_generator.get(*data_source) {
Some(sql_generator) => ControlFlow::Continue(sql_generator),
None => ControlFlow::Break(false),
}
}
// TODO is it correct?
DataSource::Unrestricted => return true,
DataSource::Unrestricted => ControlFlow::Break(true),
}
}

fn can_rewrite_template(data_source: &DataSource, meta: &MetaContext, template: &str) -> bool {
let sql_generator = match Self::template_sql_generator(data_source, meta) {
ControlFlow::Continue(sql_generator) => sql_generator,
ControlFlow::Break(verdict) => return verdict,
};

sql_generator
Expand Down
Loading
Loading