Skip to content

Commit 6862860

Browse files
waitingkuowaitingkuo
committed
fix examples, add more test cases
1 parent 5ef371d commit 6862860

File tree

6 files changed

+259
-149
lines changed

6 files changed

+259
-149
lines changed

‎datafusion-examples/examples/rewrite_expr.rs

+7
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,13 @@ impl ContextProvider for MyContextProvider {
148148
fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
149149
None
150150
}
151+
152+
fn get_config_option(
153+
&self,
154+
_variable: &str,
155+
) -> Option<datafusion_common::ScalarValue> {
156+
None
157+
}
151158
}
152159

153160
struct MyTableSource {

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

-2
Original file line numberDiff line numberDiff line change
@@ -1802,8 +1802,6 @@ impl ContextProvider for SessionState {
18021802
fn get_config_option(&self, variable: &str) -> Option<ScalarValue> {
18031803
self.config.config_options.read().get(variable)
18041804
}
1805-
1806-
18071805
}
18081806

18091807
impl FunctionRegistry for SessionState {

‎datafusion/core/tests/sql/set_variable.rs

+224-23
Original file line numberDiff line numberDiff line change
@@ -260,37 +260,238 @@ async fn set_u64_variable_bad_value() {
260260

261261
#[tokio::test]
262262
async fn set_time_zone() {
263-
// we don't support changing time zone for now until all time zone issues fixed and related function completed
263+
let ctx =
264+
SessionContext::with_config(SessionConfig::new().with_information_schema(true));
264265

265-
let ctx = SessionContext::new();
266+
plan_and_collect(&ctx, "SET datafusion.execution.time_zone = '+08:00'")
267+
.await
268+
.unwrap();
266269

267-
// for full variable name
268-
let err = plan_and_collect(&ctx, "set datafusion.execution.time_zone = '8'")
270+
let result = plan_and_collect(&ctx, "SHOW datafusion.execution.time_zone")
269271
.await
270-
.unwrap_err();
272+
.unwrap();
273+
let expected = vec![
274+
"+--------------------------------+---------+",
275+
"| name | setting |",
276+
"+--------------------------------+---------+",
277+
"| datafusion.execution.time_zone | +08:00 |",
278+
"+--------------------------------+---------+",
279+
];
280+
assert_batches_eq!(expected, &result);
281+
}
271282

272-
assert_eq!(
273-
err.to_string(),
274-
"Error during planning: Changing Time Zone isn't supported yet"
275-
);
283+
#[tokio::test]
284+
async fn set_time_zone_with_alias_variable_name() {
285+
let ctx =
286+
SessionContext::with_config(SessionConfig::new().with_information_schema(true));
276287

277-
// for alias time zone
278-
let err = plan_and_collect(&ctx, "set time zone = '8'")
288+
// TIME ZONE with space
289+
plan_and_collect(&ctx, "SET TIME ZONE = '+08:00'")
279290
.await
280-
.unwrap_err();
291+
.unwrap();
281292

282-
assert_eq!(
283-
err.to_string(),
284-
"Error during planning: Changing Time Zone isn't supported yet"
285-
);
293+
let result = plan_and_collect(&ctx, "SHOW TIME ZONE").await.unwrap();
294+
let expected = vec![
295+
"+--------------------------------+---------+",
296+
"| name | setting |",
297+
"+--------------------------------+---------+",
298+
"| datafusion.execution.time_zone | +08:00 |",
299+
"+--------------------------------+---------+",
300+
];
301+
assert_batches_eq!(expected, &result);
286302

287-
// for alias timezone
288-
let err = plan_and_collect(&ctx, "set timezone = '8'")
303+
// TIMEZONE without space
304+
plan_and_collect(&ctx, "SET TIMEZONE = '+07:00'")
289305
.await
290-
.unwrap_err();
306+
.unwrap();
291307

292-
assert_eq!(
293-
err.to_string(),
294-
"Error during planning: Changing Time Zone isn't supported yet"
295-
);
308+
let result = plan_and_collect(&ctx, "SHOW TIMEZONE").await.unwrap();
309+
let expected = vec![
310+
"+--------------------------------+---------+",
311+
"| name | setting |",
312+
"+--------------------------------+---------+",
313+
"| datafusion.execution.time_zone | +07:00 |",
314+
"+--------------------------------+---------+",
315+
];
316+
assert_batches_eq!(expected, &result);
317+
}
318+
319+
#[tokio::test]
320+
async fn set_time_zone_good_time_zone_format() {
321+
let ctx =
322+
SessionContext::with_config(SessionConfig::new().with_information_schema(true));
323+
324+
plan_and_collect(&ctx, "SET TIME ZONE = '+08:00'")
325+
.await
326+
.unwrap();
327+
328+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
329+
let result =
330+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
331+
.await
332+
.unwrap();
333+
let expected = vec![
334+
"+-----------------------------+",
335+
"| Utf8(\"2000-01-01T00:00:00\") |",
336+
"+-----------------------------+",
337+
"| 2000-01-01T08:00:00+08:00 |",
338+
"+-----------------------------+",
339+
];
340+
// this might break once https://github.com/apache/arrow-rs/issues/1936 fixed
341+
assert_batches_eq!(expected, &result);
342+
343+
plan_and_collect(&ctx, "SET TIME ZONE = '-08:00'")
344+
.await
345+
.unwrap();
346+
347+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
348+
let result =
349+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
350+
.await
351+
.unwrap();
352+
let expected = vec![
353+
"+-----------------------------+",
354+
"| Utf8(\"2000-01-01T00:00:00\") |",
355+
"+-----------------------------+",
356+
"| 1999-12-31T16:00:00-08:00 |",
357+
"+-----------------------------+",
358+
];
359+
// this might break once https://github.com/apache/arrow-rs/issues/1936 fixed
360+
assert_batches_eq!(expected, &result);
361+
362+
plan_and_collect(&ctx, "SET TIME ZONE = '+0800'")
363+
.await
364+
.unwrap();
365+
366+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
367+
let result =
368+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
369+
.await
370+
.unwrap();
371+
let expected = vec![
372+
"+-----------------------------+",
373+
"| Utf8(\"2000-01-01T00:00:00\") |",
374+
"+-----------------------------+",
375+
"| 2000-01-01T08:00:00+08:00 |",
376+
"+-----------------------------+",
377+
];
378+
// this might break once https://github.com/apache/arrow-rs/issues/1936 fixed
379+
assert_batches_eq!(expected, &result);
380+
381+
plan_and_collect(&ctx, "SET TIME ZONE = '+08'")
382+
.await
383+
.unwrap();
384+
385+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
386+
let result =
387+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
388+
.await
389+
.unwrap();
390+
let expected = vec![
391+
"+-----------------------------+",
392+
"| Utf8(\"2000-01-01T00:00:00\") |",
393+
"+-----------------------------+",
394+
"| 2000-01-01T08:00:00+08:00 |",
395+
"+-----------------------------+",
396+
];
397+
// this might break once https://github.com/apache/arrow-rs/issues/1936 fixed
398+
assert_batches_eq!(expected, &result);
399+
}
400+
401+
#[tokio::test]
402+
async fn set_time_zone_bad_time_zone_format() {
403+
let ctx =
404+
SessionContext::with_config(SessionConfig::new().with_information_schema(true));
405+
406+
plan_and_collect(&ctx, "SET TIME ZONE = '+08:00:00'")
407+
.await
408+
.unwrap();
409+
410+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
411+
let result =
412+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
413+
.await
414+
.unwrap();
415+
let expected = vec![
416+
"+-----------------------------------------------------+",
417+
"| Utf8(\"2000-01-01T00:00:00\") |",
418+
"+-----------------------------------------------------+",
419+
"| 2000-01-01T00:00:00 (Unknown Time Zone '+08:00:00') |",
420+
"+-----------------------------------------------------+",
421+
];
422+
assert_batches_eq!(expected, &result);
423+
424+
plan_and_collect(&ctx, "SET TIME ZONE = '08:00'")
425+
.await
426+
.unwrap();
427+
428+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
429+
let result =
430+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
431+
.await
432+
.unwrap();
433+
let expected = vec![
434+
"+-------------------------------------------------+",
435+
"| Utf8(\"2000-01-01T00:00:00\") |",
436+
"+-------------------------------------------------+",
437+
"| 2000-01-01T00:00:00 (Unknown Time Zone '08:00') |",
438+
"+-------------------------------------------------+",
439+
];
440+
assert_batches_eq!(expected, &result);
441+
442+
plan_and_collect(&ctx, "SET TIME ZONE = '08'")
443+
.await
444+
.unwrap();
445+
446+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
447+
let result =
448+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
449+
.await
450+
.unwrap();
451+
let expected = vec![
452+
"+----------------------------------------------+",
453+
"| Utf8(\"2000-01-01T00:00:00\") |",
454+
"+----------------------------------------------+",
455+
"| 2000-01-01T00:00:00 (Unknown Time Zone '08') |",
456+
"+----------------------------------------------+",
457+
];
458+
assert_batches_eq!(expected, &result);
459+
460+
// we dont support named time zone yet
461+
plan_and_collect(&ctx, "SET TIME ZONE = 'Asia/Taipei'")
462+
.await
463+
.unwrap();
464+
465+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
466+
let result =
467+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
468+
.await
469+
.unwrap();
470+
let expected = vec![
471+
"+-------------------------------------------------------+",
472+
"| Utf8(\"2000-01-01T00:00:00\") |",
473+
"+-------------------------------------------------------+",
474+
"| 2000-01-01T00:00:00 (Unknown Time Zone 'Asia/Taipei') |",
475+
"+-------------------------------------------------------+",
476+
];
477+
assert_batches_eq!(expected, &result);
478+
479+
// this is invalid even after we support named time zone
480+
plan_and_collect(&ctx, "SET TIME ZONE = 'Asia/Taipei2'")
481+
.await
482+
.unwrap();
483+
484+
// casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
485+
let result =
486+
plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
487+
.await
488+
.unwrap();
489+
let expected = vec![
490+
"+--------------------------------------------------------+",
491+
"| Utf8(\"2000-01-01T00:00:00\") |",
492+
"+--------------------------------------------------------+",
493+
"| 2000-01-01T00:00:00 (Unknown Time Zone 'Asia/Taipei2') |",
494+
"+--------------------------------------------------------+",
495+
];
496+
assert_batches_eq!(expected, &result);
296497
}

‎datafusion/optimizer/tests/integration-test.rs

+4-1
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,10 @@ impl ContextProvider for MySchemaProvider {
331331
None
332332
}
333333

334-
fn get_config_option(&self, _variable: &str) -> Option<datafusion_common::ScalarValue> {
334+
fn get_config_option(
335+
&self,
336+
_variable: &str,
337+
) -> Option<datafusion_common::ScalarValue> {
335338
None
336339
}
337340
}

‎datafusion/sql/examples/sql.rs

+4-1
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,10 @@ impl ContextProvider for MySchemaProvider {
121121
None
122122
}
123123

124-
fn get_config_option(&self, _variable: &str) -> Option<datafusion_common::ScalarValue> {
124+
fn get_config_option(
125+
&self,
126+
_variable: &str,
127+
) -> Option<datafusion_common::ScalarValue> {
125128
None
126129
}
127130
}

0 commit comments

Comments
 (0)