jayzhan211 commented on code in PR #27143:
URL: https://github.com/apache/datafusion/pull/27143#discussion_r4538661857
##########
datafusion/expr/src/expr_schema.rs:
##########
@@ -685,11 +688,26 @@ impl ExprSchemable for Expr {
))),
}
}
+ Expr::Case(case) => {
+ let fields = case
+ .when_then_expr
+ .iter()
+ .map(|(_, then)| then)
+ .chain(case.else_expr.iter())
+ .map(|expr| expr.to_field(schema).map(|(_, field)| field))
+ .collect::<Result<Vec<_>>>()?;
+ let metadata =
+ coerced_fields_metadata(fields.iter().map(|f|
f.as_ref()))?;
+ Ok(Arc::new(metadata.add_to_field(Field::new(
+ &schema_name,
+ self.get_type(schema)?,
+ self.nullable(schema)?,
+ ))))
Review Comment:
Planning a CASE nested inside function calls got much slower: with 12 levels
of `upper(CASE WHEN a > k THEN <inner> ELSE s END)`, `EXPLAIN` takes 19–21 s
here vs 0.34–0.44 s on main (debug build). Each level now walks its branches
three times: `to_field` on each branch, then `self.get_type` and
`self.nullable`, and for function calls those two go back through `to_field`.
So the cost grows as 3^depth instead of main's 2^depth. Taking the type from
the branch fields you already have brings it back to 0.38–0.43 s with the same
plan. The expr and optimizer lib tests and the case/metadata/coalesce/nvl slt
files still pass:
```suggestion
let metadata =
coerced_fields_metadata(fields.iter().map(|f|
f.as_ref()))?;
// Same as `get_type`: the first non-NULL branch type.
Reading it from
// `fields` avoids a second walk of the branches, which
multiplies with
// the depth of nested CASE and function calls.
let data_type = fields
.iter()
.map(|f| f.data_type())
.find(|t| !t.is_null())
.cloned()
.unwrap_or(DataType::Null);
Ok(Arc::new(metadata.add_to_field(Field::new(
&schema_name,
data_type,
self.nullable(schema)?,
))))
```
<details>
<summary>Script that generates the query, and the timings</summary>
```sh
python3 - 12 > q.sql <<'PY'
import sys
d = int(sys.argv[1]); e = "s"
for k in range(d):
e = f"upper(CASE WHEN a > {k} THEN {e} ELSE s END)"
print("CREATE TABLE t AS VALUES (1, 'x');")
print("CREATE VIEW v AS SELECT column1 AS a, column2 AS s FROM t;")
print(f"EXPLAIN SELECT {e} AS r FROM v;")
PY
time datafusion-cli -q -f q.sql > /dev/null
```
| depth | main | this PR | with the suggestion |
|---|---|---|---|
| 10 | 0.15–0.31 s | 1.7–2.7 s | 0.12 s |
| 12 | 0.34–0.44 s | 18.8–21.4 s | 0.38–0.43 s |
| 14 | 1.3–1.9 s | not run | 1.5 s |
</details>
##########
datafusion/sqllogictest/test_files/metadata.slt:
##########
@@ -687,6 +687,173 @@ NULL
l_bar
baz
+# CASE merges the metadata of its value branches with the same rules as
+# coalesce: the earlier branch wins for each key, and an extension name is kept
+# together with its parameters. The operand of a simple CASE and the WHEN
+# conditions do not contribute.
+query TTT
+select
+ arrow_metadata(case when id > 1 then name else l_name end, 'metadata_key'),
+ arrow_metadata(case id when 1 then name else l_name end, 'metadata_key'),
+ arrow_metadata(
+ case when id > 1 then cast(name as json) else cast(l_name as json) end,
+ 'ARROW:extension:name')
+from table_with_metadata;
+----
+the name field the name field arrow.json
+the name field the name field arrow.json
+the name field the name field arrow.json
+
+statement error Error during planning: Cannot merge field metadata with
differing values of 'ARROW:extension:name': 'arrow\.json' and 'b\.type'
+select case when id > 1
+ then cast(name as json)
+ else with_metadata(l_name, 'ARROW:extension:name', 'b.type') end
+from table_with_metadata;
Review Comment:
Optional cleanup: the plain-text form of this error is easier to read than
the escaped regex. I ran the file with this change applied, including the blank
line that follows it:
```suggestion
statement error
select case when id > 1
then cast(name as json)
else with_metadata(l_name, 'ARROW:extension:name', 'b.type') end
from table_with_metadata;
----
DataFusion error: Error during planning: Cannot merge field metadata with
differing values of 'ARROW:extension:name': 'arrow.json' and 'b.type'
```
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -1307,6 +1325,24 @@ impl PhysicalExpr for CaseExpr {
}
}
+ fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> {
Review Comment:
`ScalarFunctionExpr::evaluate` calls `return_field` on its arguments for
every batch, and here each nested CASE level formats the name of its whole
subtree and recomputes its nullability. That makes it O(depth²): for a CASE
nested in the ELSE, the extra cost over main is 11% of the CASE's evaluation
time at depth 4 and 22% at depth 8 (release build, 8192 rows). If nested CASE
branches compute only their metadata, it drops to 1.8% and 2.8%, and the
physical-expr lib tests and the metadata/case slt files still pass:
```diff
-use datafusion_common::metadata::coerced_fields_metadata;
+use datafusion_common::metadata::{FieldMetadata, coerced_fields_metadata};
@@ impl CaseExpr {
self.body.else_expr.as_ref()
}
+
+ /// The metadata of [`PhysicalExpr::return_field`]. A nested CASE
branch only
+ /// computes its metadata, not its name and nullability, which would
walk its
+ /// whole subtree again at every level.
+ fn return_metadata(&self, input_schema: &Schema) ->
Result<FieldMetadata> {
+ let fields = self
+ .body
+ .when_then_expr
+ .iter()
+ .map(|(_, then)| then)
+ .chain(self.body.else_expr.iter())
+ .map(|expr| match expr.downcast_ref::<CaseExpr>() {
+ Some(case) => {
+ let field = Field::new("",
case.data_type(input_schema)?, true);
+ Ok(Arc::new(
+
case.return_metadata(input_schema)?.add_to_field(field),
+ ))
+ }
+ None => expr.return_field(input_schema),
+ })
+ .collect::<Result<Vec<_>>>()?;
+ coerced_fields_metadata(fields.iter().map(|f| f.as_ref()))
+ }
}
@@ impl PhysicalExpr for CaseExpr {
fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> {
- let fields = self
- .body
- .when_then_expr
- .iter()
- .map(|(_, then)| then)
- .chain(self.body.else_expr.iter())
- .map(|expr| expr.return_field(input_schema))
- .collect::<Result<Vec<_>>>()?;
- let metadata = coerced_fields_metadata(fields.iter().map(|f|
f.as_ref()))?;
let field = Field::new(
format!("{self}"),
self.data_type(input_schema)?,
self.nullable(input_schema)?,
);
- Ok(metadata.add_to_field(field).into())
+ Ok(self
+ .return_metadata(input_schema)?
+ .add_to_field(field)
+ .into())
}
```
<details>
<summary>Extra return_field cost as a share of CASE evaluation (release,
8192 rows)</summary>
| nesting | this PR | with the diff |
|---|---|---|
| ELSE, depth 1 | 2.2% | 2.3% |
| ELSE, depth 4 | 11.1% | 1.8% |
| ELSE, depth 8 | 22.2% | 2.8% |
| ELSE, depth 32 | 137% | 7.3% |
| THEN, depth 32 | 49.5% | 2.4% |
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
