Skip to content

Improve filter predicates with Utf8View literals #11043

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alamb merged 6 commits into from
Jun 26, 2024
Merged

Improve filter predicates with Utf8View literals #11043

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c69dc7e
refactor: Improve type coercion logic in TypeCoercionRewriter
Weijun-H Jun 21, 2024
da77901
refactor: Improve type coercion logic in TypeCoercionRewriter
Weijun-H Jun 21, 2024
690f8d0
chore
Weijun-H Jun 22, 2024
7cb3d39
chore: Update test
Weijun-H Jun 22, 2024
f8db675
refactor: Improve type coercion logic in TypeCoercionRewriter
Weijun-H Jun 25, 2024
f06c505
refactor: Remove unused import and update code formatting in unwrap_c…
Weijun-H Jun 25, 2024
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
26 changes: 10 additions & 16 deletions datafusion/optimizer/src/unwrap_cast_in_comparison.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
use datafusion_common::{internal_err, DFSchema, DFSchemaRef, Result, ScalarValue};
use datafusion_expr::expr::{BinaryExpr, Cast, InList, TryCast};
use datafusion_expr::utils::merge_schema;
use datafusion_expr::{lit, Expr, ExprSchemable, LogicalPlan, Operator};
use datafusion_expr::{lit, Expr, ExprSchemable, LogicalPlan};

/// [`UnwrapCastInComparison`] attempts to remove casts from
/// comparisons to literals ([`ScalarValue`]s) by applying the casts
Expand Down Expand Up @@ -154,7 +154,7 @@ impl TreeNodeRewriter for UnwrapCastExprRewriter {
};
is_supported_type(&left_type)
&& is_supported_type(&right_type)
&& is_comparison_op(op)
&& op.is_comparison_operator()
} =>
{
match (left.as_mut(), right.as_mut()) {
Expand Down Expand Up @@ -270,18 +270,6 @@ impl TreeNodeRewriter for UnwrapCastExprRewriter {
}
}

fn is_comparison_op(op: &Operator) -> bool {
Copy link
Contributor

@alamb alamb Jun 25, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice cleanup

matches!(
op,
Operator::Eq
| Operator::NotEq
| Operator::Gt
| Operator::GtEq
| Operator::Lt
| Operator::LtEq
)
}

/// Returns true if [UnwrapCastExprRewriter] supports this data type
fn is_supported_type(data_type: &DataType) -> bool {
is_supported_numeric_type(data_type)
Expand All @@ -308,7 +296,10 @@ fn is_supported_numeric_type(data_type: &DataType) -> bool {

/// Returns true if [UnwrapCastExprRewriter] supports casting this value as a string
fn is_supported_string_type(data_type: &DataType) -> bool {
matches!(data_type, DataType::Utf8 | DataType::LargeUtf8)
matches!(
data_type,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
)
}

/// Returns true if [UnwrapCastExprRewriter] supports casting this value as a dictionary
Expand Down Expand Up @@ -481,12 +472,15 @@ fn try_cast_string_literal(
target_type: &DataType,
) -> Option<ScalarValue> {
let string_value = match lit_value {
ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) => s.clone(),
ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s) => {
s.clone()
}
_ => return None,
};
let scalar_value = match target_type {
DataType::Utf8 => ScalarValue::Utf8(string_value),
DataType::LargeUtf8 => ScalarValue::LargeUtf8(string_value),
DataType::Utf8View => ScalarValue::Utf8View(string_value),
_ => return None,
};
Some(scalar_value)
Expand Down
18 changes: 16 additions & 2 deletions datafusion/sqllogictest/test_files/string_view.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,12 +226,18 @@ logical_plan
02)--Filter: test.column1_utf8view = Utf8View("Andrew")
03)----TableScan: test projection=[column1_utf8, column1_utf8view]

# should not be casting the column: https://github.com/apache/datafusion/issues/10998
query TT
explain SELECT column1_utf8 from test where column1_utf8 = arrow_cast('Andrew', 'Utf8View');
----
logical_plan
01)Filter: CAST(test.column1_utf8 AS Utf8View) = Utf8View("Andrew")
01)Filter: test.column1_utf8 = Utf8("Andrew")
02)--TableScan: test projection=[column1_utf8]

query TT
explain SELECT column1_utf8 from test where arrow_cast('Andrew', 'Utf8View') = column1_utf8;
----
logical_plan
01)Filter: test.column1_utf8 = Utf8("Andrew")
02)--TableScan: test projection=[column1_utf8]

query TT
Expand All @@ -242,6 +248,14 @@ logical_plan
02)--Filter: test.column1_utf8view = Utf8View("Andrew")
03)----TableScan: test projection=[column1_utf8, column1_utf8view]

query TT
explain SELECT column1_utf8 from test where arrow_cast('Andrew', 'Dictionary(Int32, Utf8)') = column1_utf8view;
----
logical_plan
01)Projection: test.column1_utf8
02)--Filter: test.column1_utf8view = Utf8View("Andrew")
03)----TableScan: test projection=[column1_utf8, column1_utf8view]

# compare string / stringview
# Should cast string -> stringview (which is cheap), not stringview -> string (which is not)
query TT
Expand Down