support `SET` variable #4069

waitingkuo · 2022-11-01T21:41:02Z

Which issue does this PR close?

Rationale for this change

to support

SET [VARIABLE] [TO | =] [VALUE | 'VALUE']

What changes are included in this PR?

support u64 value

❯ set datafusion.execution.batch_size = 1;
0 rows in set. Query took 0.001 seconds.
❯ show datafusion.execution.batch_size;
+---------------------------------+---------+
| name                            | setting |
+---------------------------------+---------+
| datafusion.execution.batch_size | 1       |
+---------------------------------+---------+
1 row in set. Query took 0.013 seconds.

support bool

❯ set datafusion.execution.coalesce_batches to false;
0 rows in set. Query took 0.000 seconds.
❯ show datafusion.execution.coalesce_batches;
+---------------------------------------+---------+
| name                                  | setting |
+---------------------------------------+---------+
| datafusion.execution.coalesce_batches | false   |
+---------------------------------------+---------+
1 row in set. Query took 0.007 seconds.

support single quoted string

❯ set datafusion.execution.coalesce_batches to 'false';
0 rows in set. Query took 0.000 seconds.
❯ show datafusion.execution.coalesce_batches;
+---------------------------------------+---------+
| name                                  | setting |
+---------------------------------------+---------+
| datafusion.execution.coalesce_batches | false   |
+---------------------------------------+---------+
1 row in set. Query took 0.007 seconds.

support alias TIME ZONE and TIMEZONE

as discussed in #3148 (comment)
we disallow set timezone for now until timezone integration is full completed

❯ set time zone = 1;
Plan("Changing Time Zone isn't supported yet")

throw error for unknown variable

❯ set abc = 1;
Execution("Unknown Variable abc")

throw error for incorrect data type

❯ set datafusion.execution.batch_size = a;
Execution("Failed to parse a as u64")

❯ set datafusion.execution.batch_size = -1;
Execution("Failed to parse -1 as u64")

Are there any user-facing changes?

waitingkuo · 2022-11-01T21:45:14Z

datafusion/core/src/execution/context.rs

+            LogicalPlan::SetVariable(SetVariable {
+                variable, value, ..
+            }) => {
+                let config_options = &self.state.write().config.config_options;
+


better refactor this part in the follow on pr once #3909 merged

waitingkuo · 2022-11-01T21:46:09Z

datafusion/core/src/execution/context.rs

+                let old_value =
+                    config_options.read().get(&variable).ok_or_else(|| {
+                        DataFusionError::Execution(format!(
+                            "Unknown Variable {}",
+                            variable
+                        ))
+                    })?;


return error if variable isn't in config_options

waitingkuo · 2022-11-01T21:49:24Z

datafusion/core/src/execution/context.rs

+                match old_value {
+                    ScalarValue::Boolean(_) => {
+                        let new_value = value.parse::<bool>().map_err(|_| {
+                            DataFusionError::Execution(format!(
+                                "Failed to parse {} as bool",
+                                value,
+                            ))
+                        })?;
+                        config_options.write().set_bool(&variable, new_value);
+                    }
+
+                    ScalarValue::UInt64(_) => {
+                        let new_value = value.parse::<u64>().map_err(|_| {
+                            DataFusionError::Execution(format!(
+                                "Failed to parse {} as u64",
+                                value,
+                            ))
+                        })?;
+                        config_options.write().set_u64(&variable, new_value);
+                    }
+
+                    ScalarValue::Utf8(_) => {
+                        let new_value = value.parse::<String>().map_err(|_| {
+                            DataFusionError::Execution(format!(
+                                "Failed to parse {} as String",
+                                value,
+                            ))
+                        })?;
+                        config_options.write().set_string(&variable, new_value);
+                    }
+
+                    _ => {
+                        return Err(DataFusionError::Execution(
+                            "Unsupported Scalar Value Type".to_string(),
+                        ))
+                    }
+                }


not sure whether it's a good way to check the data type. ideally we should use DataType defined in ConfigDefinition https://github.com/apache/arrow-datafusion/blob/97f2e4fd5517c762b0862d22b81f957db511e22e/datafusion/core/src/config.rs#L72-L80

but this information is gone once it's been converted to ConfigOptions
https://github.com/apache/arrow-datafusion/blob/97f2e4fd5517c762b0862d22b81f957db511e22e/datafusion/core/src/config.rs#L279-L281

I think checking the existing type is a reasonable approach

waitingkuo · 2022-11-01T21:50:31Z

datafusion/core/tests/sql/set_variable.rs

+#[tokio::test]
+async fn set_time_zone() {
+    // we don't support changing time zone for now until all time zone issues fixed and related function completed
+
+    let ctx = SessionContext::new();
+
+    // for full variable name
+    let err = plan_and_collect(&ctx, "set datafusion.execution.time_zone = '8'")
+        .await
+        .unwrap_err();
+
+    assert_eq!(
+        err.to_string(),
+        "Error during planning: Changing Time Zone isn't supported yet"
+    );


should get back to this once time zone is fully integrated

waitingkuo · 2022-11-01T21:51:28Z

datafusion/sql/src/planner.rs

+        if local {
+            return Err(DataFusionError::NotImplemented(
+                "LOCAL is not supported".to_string(),
+            ));
+        }


we don't support SET LOCAL VARIABLE TO VALUE

waitingkuo · 2022-11-01T21:51:59Z

datafusion/sql/src/planner.rs

+        if hivevar {
+            return Err(DataFusionError::NotImplemented(
+                "HIVEVAR is not supported".to_string(),
+            ));
+        }


we don't support HIVEVAR. I'm not sure what it is

waitingkuo · 2022-11-01T21:52:45Z

datafusion/sql/src/planner.rs

+        if variable_lower == "timezone" || variable_lower == "time.zone" {
+            // we could introduce alias in OptionDefinition if this string matching thing grows
+            variable_lower = "datafusion.execution.time_zone".to_string();
+        }


alias timezone and time zone
note that time zone is converted to time.zone during sql parsing

waitingkuo · 2022-11-01T21:53:30Z

datafusion/sql/src/planner.rs

+
+        // parse value string from Expr
+        let value_string = match &value[0] {
+            SQLExpr::Identifier(i) => i.to_string(),


to support non-whitespace-separated string

waitingkuo · 2022-11-01T21:53:54Z

datafusion/sql/src/planner.rs

+                Value::SingleQuotedString(s) => s.to_string(),
+                Value::Number(_, _) | Value::Boolean(_) => v.to_string(),


to support SingleQuotedString, number and Bool

waitingkuo · 2022-11-01T21:54:10Z

datafusion/sql/src/planner.rs

+                UnaryOperator::Plus => format!("+{}", expr),
+                UnaryOperator::Minus => format!("-{}", expr),


to support signed number e.g. +8, -8

An alternative (for the future) could be to parse the value as an Expr and then call ExprSimplifier::simplify(expr) on it during evaluation in SessionContext.

That would then support sql statements like

set batch_size = 8*1024

I think this way is fine, too

Example of how to do it
https://github.com/apache/arrow-datafusion/blob/10e64dc013ba210ab1f6c2a3c02c66aef4a0e802/datafusion-examples/examples/expr_api.rs#L77-L89

i'm not familiar with this yet, will try it and send a ticket or pr

alamb

This is beautiful @waitingkuo -- thank you 🏆

datafusion/core/src/config.rs

datafusion/core/src/execution/context.rs

alamb · 2022-11-02T16:16:33Z

datafusion/core/src/execution/context.rs

+                match old_value {
+                    ScalarValue::Boolean(_) => {
+                        let new_value = value.parse::<bool>().map_err(|_| {
+                            DataFusionError::Execution(format!(
+                                "Failed to parse {} as bool",
+                                value,
+                            ))
+                        })?;
+                        config_options.write().set_bool(&variable, new_value);
+                    }
+
+                    ScalarValue::UInt64(_) => {
+                        let new_value = value.parse::<u64>().map_err(|_| {
+                            DataFusionError::Execution(format!(
+                                "Failed to parse {} as u64",
+                                value,
+                            ))
+                        })?;
+                        config_options.write().set_u64(&variable, new_value);
+                    }
+
+                    ScalarValue::Utf8(_) => {
+                        let new_value = value.parse::<String>().map_err(|_| {
+                            DataFusionError::Execution(format!(
+                                "Failed to parse {} as String",
+                                value,
+                            ))
+                        })?;
+                        config_options.write().set_string(&variable, new_value);
+                    }
+
+                    _ => {
+                        return Err(DataFusionError::Execution(
+                            "Unsupported Scalar Value Type".to_string(),
+                        ))
+                    }
+                }


I think checking the existing type is a reasonable approach

alamb · 2022-11-02T16:18:52Z

datafusion/core/tests/sql/set_variable.rs

+#[tokio::test]
+async fn set_time_zone() {
+    // we don't support changing time zone for now until all time zone issues fixed and related function completed
+
+    let ctx = SessionContext::new();
+
+    // for full variable name
+    let err = plan_and_collect(&ctx, "set datafusion.execution.time_zone = '8'")
+        .await
+        .unwrap_err();
+
+    assert_eq!(
+        err.to_string(),
+        "Error during planning: Changing Time Zone isn't supported yet"
+    );


datafusion/expr/src/logical_plan/plan.rs

alamb · 2022-11-02T16:23:05Z

datafusion/sql/src/planner.rs

+                UnaryOperator::Plus => format!("+{}", expr),
+                UnaryOperator::Minus => format!("-{}", expr),


An alternative (for the future) could be to parse the value as an Expr and then call ExprSimplifier::simplify(expr) on it during evaluation in SessionContext.

That would then support sql statements like

set batch_size = 8*1024

I think this way is fine, too

Example of how to do it
https://github.com/apache/arrow-datafusion/blob/10e64dc013ba210ab1f6c2a3c02c66aef4a0e802/datafusion-examples/examples/expr_api.rs#L77-L89

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

waitingkuo · 2022-11-02T22:19:22Z

@alamb thank you for the review, learned a lot from it.
updated the test cases according to the suggested changes

ursabot · 2022-11-03T14:42:52Z

Benchmark runs are scheduled for baseline = 761e167 and contender = c9442ce. c9442ce is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ec2-t3-xlarge-us-east-2] ec2-t3-xlarge-us-east-2
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on test-mac-arm] test-mac-arm
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ursa-i9-9960x] ursa-i9-9960x
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ursa-thinkcentre-m75q] ursa-thinkcentre-m75q
Buildkite builds:
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java

* support SET * remove useless comment * add test cases * Update datafusion/core/src/config.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> * Update datafusion/core/src/execution/context.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> * fix test cases * fmt * Update datafusion/expr/src/logical_plan/plan.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

waitingkuo added 2 commits November 2, 2022 05:32

support SET

4cd9aba

remove useless comment

798998c

github-actions bot added core Core datafusion crate logical-expr Logical plan and expressions optimizer Optimizer rules sql labels Nov 1, 2022

add test cases

1f26272

waitingkuo commented Nov 1, 2022

View reviewed changes

alamb approved these changes Nov 2, 2022

View reviewed changes

waitingkuo and others added 5 commits November 3, 2022 02:56

Update datafusion/core/src/config.rs

2fe8fc2

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Update datafusion/core/src/execution/context.rs

ecd4b0f

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

fix test cases

4bc89a2

fmt

955c3e5

Update datafusion/expr/src/logical_plan/plan.rs

6704dde

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

alamb approved these changes Nov 3, 2022

View reviewed changes

alamb merged commit c9442ce into apache:master Nov 3, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

support `SET` variable #4069

support `SET` variable #4069

waitingkuo commented Nov 1, 2022

waitingkuo Nov 1, 2022

waitingkuo Nov 1, 2022

waitingkuo Nov 1, 2022

alamb Nov 2, 2022

waitingkuo Nov 1, 2022

alamb Nov 2, 2022

waitingkuo Nov 1, 2022

waitingkuo Nov 1, 2022

waitingkuo Nov 1, 2022

waitingkuo Nov 1, 2022

waitingkuo Nov 1, 2022

waitingkuo Nov 1, 2022

alamb Nov 2, 2022

waitingkuo Nov 2, 2022

alamb left a comment

alamb Nov 2, 2022

alamb Nov 2, 2022

alamb Nov 2, 2022

waitingkuo commented Nov 2, 2022

ursabot commented Nov 3, 2022

		Value::SingleQuotedString(s) => s.to_string(),
		Value::Number(_, _) \| Value::Boolean(_) => v.to_string(),

		UnaryOperator::Plus => format!("+{}", expr),
		UnaryOperator::Minus => format!("-{}", expr),

support SET variable #4069

support SET variable #4069

Conversation

waitingkuo commented Nov 1, 2022

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

waitingkuo commented Nov 2, 2022

ursabot commented Nov 3, 2022

support `SET` variable #4069

support `SET` variable #4069