Skip to content

Commit 7334481

Browse files
liukun4515liukun4515
authored
support array with scalar arithmetic operation for decimal data type (apache#2233)
* support array with scalar arithmetic operation for decimal data type * add sql level test for arithmetic op * add sql level test for decimal for arithmetic operation
1 parent 4c31f78 commit 7334481

File tree

2 files changed

+382
-0
lines changed

2 files changed

+382
-0
lines changed

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

+300
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,306 @@ async fn decimal_logic_op() -> Result<()> {
363363
Ok(())
364364
}
365365

366+
#[tokio::test]
367+
async fn decimal_arithmetic_op() -> Result<()> {
368+
let ctx = SessionContext::new();
369+
register_decimal_csv_table_by_sql(&ctx).await;
370+
// add
371+
let sql = "select c1+1 from decimal_simple"; // add scalar
372+
let actual = execute_to_batches(&ctx, sql).await;
373+
// array decimal(10,6) + scalar decimal(20,0) => decimal(21,6)
374+
assert_eq!(
375+
&DataType::Decimal(27, 6),
376+
actual[0].schema().field(0).data_type()
377+
);
378+
let expected = vec![
379+
"+---------------------------------+",
380+
"| decimal_simple.c1 Plus Int64(1) |",
381+
"+---------------------------------+",
382+
"| 1.000010 |",
383+
"| 1.000020 |",
384+
"| 1.000020 |",
385+
"| 1.000030 |",
386+
"| 1.000030 |",
387+
"| 1.000030 |",
388+
"| 1.000040 |",
389+
"| 1.000040 |",
390+
"| 1.000040 |",
391+
"| 1.000040 |",
392+
"| 1.000050 |",
393+
"| 1.000050 |",
394+
"| 1.000050 |",
395+
"| 1.000050 |",
396+
"| 1.000050 |",
397+
"+---------------------------------+",
398+
];
399+
assert_batches_eq!(expected, &actual);
400+
// array decimal(10,6) + array decimal(12,7) => decimal(13,7)
401+
let sql = "select c1+c5 from decimal_simple";
402+
let actual = execute_to_batches(&ctx, sql).await;
403+
assert_eq!(
404+
&DataType::Decimal(13, 7),
405+
actual[0].schema().field(0).data_type()
406+
);
407+
let expected = vec![
408+
"+------------------------------------------+",
409+
"| decimal_simple.c1 Plus decimal_simple.c5 |",
410+
"+------------------------------------------+",
411+
"| 0.0000240 |",
412+
"| 0.0000450 |",
413+
"| 0.0000390 |",
414+
"| 0.0000620 |",
415+
"| 0.0000650 |",
416+
"| 0.0000410 |",
417+
"| 0.0000840 |",
418+
"| 0.0000800 |",
419+
"| 0.0000800 |",
420+
"| 0.0000840 |",
421+
"| 0.0001020 |",
422+
"| 0.0001280 |",
423+
"| 0.0000830 |",
424+
"| 0.0001180 |",
425+
"| 0.0001500 |",
426+
"+------------------------------------------+",
427+
];
428+
assert_batches_eq!(expected, &actual);
429+
// subtract
430+
let sql = "select c1-1 from decimal_simple";
431+
let actual = execute_to_batches(&ctx, sql).await;
432+
assert_eq!(
433+
&DataType::Decimal(27, 6),
434+
actual[0].schema().field(0).data_type()
435+
);
436+
let expected = vec![
437+
"+----------------------------------+",
438+
"| decimal_simple.c1 Minus Int64(1) |",
439+
"+----------------------------------+",
440+
"| -0.999990 |",
441+
"| -0.999980 |",
442+
"| -0.999980 |",
443+
"| -0.999970 |",
444+
"| -0.999970 |",
445+
"| -0.999970 |",
446+
"| -0.999960 |",
447+
"| -0.999960 |",
448+
"| -0.999960 |",
449+
"| -0.999960 |",
450+
"| -0.999950 |",
451+
"| -0.999950 |",
452+
"| -0.999950 |",
453+
"| -0.999950 |",
454+
"| -0.999950 |",
455+
"+----------------------------------+",
456+
];
457+
assert_batches_eq!(expected, &actual);
458+
459+
let sql = "select c1-c5 from decimal_simple";
460+
let actual = execute_to_batches(&ctx, sql).await;
461+
assert_eq!(
462+
&DataType::Decimal(13, 7),
463+
actual[0].schema().field(0).data_type()
464+
);
465+
let expected = vec![
466+
"+-------------------------------------------+",
467+
"| decimal_simple.c1 Minus decimal_simple.c5 |",
468+
"+-------------------------------------------+",
469+
"| -0.0000040 |",
470+
"| -0.0000050 |",
471+
"| 0.0000010 |",
472+
"| -0.0000020 |",
473+
"| -0.0000050 |",
474+
"| 0.0000190 |",
475+
"| -0.0000040 |",
476+
"| 0.0000000 |",
477+
"| 0.0000000 |",
478+
"| -0.0000040 |",
479+
"| -0.0000020 |",
480+
"| -0.0000280 |",
481+
"| 0.0000170 |",
482+
"| -0.0000180 |",
483+
"| -0.0000500 |",
484+
"+-------------------------------------------+",
485+
];
486+
assert_batches_eq!(expected, &actual);
487+
// multiply
488+
let sql = "select c1*20 from decimal_simple";
489+
let actual = execute_to_batches(&ctx, sql).await;
490+
assert_eq!(
491+
&DataType::Decimal(31, 6),
492+
actual[0].schema().field(0).data_type()
493+
);
494+
let expected = vec![
495+
"+--------------------------------------+",
496+
"| decimal_simple.c1 Multiply Int64(20) |",
497+
"+--------------------------------------+",
498+
"| 0.000200 |",
499+
"| 0.000400 |",
500+
"| 0.000400 |",
501+
"| 0.000600 |",
502+
"| 0.000600 |",
503+
"| 0.000600 |",
504+
"| 0.000800 |",
505+
"| 0.000800 |",
506+
"| 0.000800 |",
507+
"| 0.000800 |",
508+
"| 0.001000 |",
509+
"| 0.001000 |",
510+
"| 0.001000 |",
511+
"| 0.001000 |",
512+
"| 0.001000 |",
513+
"+--------------------------------------+",
514+
];
515+
assert_batches_eq!(expected, &actual);
516+
517+
let sql = "select c1*c5 from decimal_simple";
518+
let actual = execute_to_batches(&ctx, sql).await;
519+
assert_eq!(
520+
&DataType::Decimal(23, 13),
521+
actual[0].schema().field(0).data_type()
522+
);
523+
let expected = vec![
524+
"+----------------------------------------------+",
525+
"| decimal_simple.c1 Multiply decimal_simple.c5 |",
526+
"+----------------------------------------------+",
527+
"| 0.0000000001400 |",
528+
"| 0.0000000005000 |",
529+
"| 0.0000000003800 |",
530+
"| 0.0000000009600 |",
531+
"| 0.0000000010500 |",
532+
"| 0.0000000003300 |",
533+
"| 0.0000000017600 |",
534+
"| 0.0000000016000 |",
535+
"| 0.0000000016000 |",
536+
"| 0.0000000017600 |",
537+
"| 0.0000000026000 |",
538+
"| 0.0000000039000 |",
539+
"| 0.0000000016500 |",
540+
"| 0.0000000034000 |",
541+
"| 0.0000000050000 |",
542+
"+----------------------------------------------+",
543+
];
544+
assert_batches_eq!(expected, &actual);
545+
// divide
546+
let sql = "select c1/cast(0.00001 as decimal(5,5)) from decimal_simple";
547+
let actual = execute_to_batches(&ctx, sql).await;
548+
assert_eq!(
549+
&DataType::Decimal(21, 12),
550+
actual[0].schema().field(0).data_type()
551+
);
552+
let expected = vec![
553+
"+-------------------------------------------------------------+",
554+
"| decimal_simple.c1 / CAST(Float64(0.00001) AS Decimal(5, 5)) |",
555+
"+-------------------------------------------------------------+",
556+
"| 1.000000000000 |",
557+
"| 2.000000000000 |",
558+
"| 2.000000000000 |",
559+
"| 3.000000000000 |",
560+
"| 3.000000000000 |",
561+
"| 3.000000000000 |",
562+
"| 4.000000000000 |",
563+
"| 4.000000000000 |",
564+
"| 4.000000000000 |",
565+
"| 4.000000000000 |",
566+
"| 5.000000000000 |",
567+
"| 5.000000000000 |",
568+
"| 5.000000000000 |",
569+
"| 5.000000000000 |",
570+
"| 5.000000000000 |",
571+
"+-------------------------------------------------------------+",
572+
];
573+
assert_batches_eq!(expected, &actual);
574+
575+
let sql = "select c1/c5 from decimal_simple";
576+
let actual = execute_to_batches(&ctx, sql).await;
577+
assert_eq!(
578+
&DataType::Decimal(30, 19),
579+
actual[0].schema().field(0).data_type()
580+
);
581+
let expected = vec![
582+
"+--------------------------------------------+",
583+
"| decimal_simple.c1 Divide decimal_simple.c5 |",
584+
"+--------------------------------------------+",
585+
"| 0.7142857142857143296 |",
586+
"| 0.8000000000000000000 |",
587+
"| 1.0526315789473683456 |",
588+
"| 0.9375000000000000000 |",
589+
"| 0.8571428571428571136 |",
590+
"| 2.7272727272727269376 |",
591+
"| 0.9090909090909090816 |",
592+
"| 1.0000000000000000000 |",
593+
"| 1.0000000000000000000 |",
594+
"| 0.9090909090909090816 |",
595+
"| 0.9615384615384614912 |",
596+
"| 0.6410256410256410624 |",
597+
"| 1.5151515151515152384 |",
598+
"| 0.7352941176470588416 |",
599+
"| 0.5000000000000000000 |",
600+
"+--------------------------------------------+",
601+
];
602+
assert_batches_eq!(expected, &actual);
603+
604+
// modulo
605+
let sql = "select c5%cast(0.00001 as decimal(5,5)) from decimal_simple";
606+
let actual = execute_to_batches(&ctx, sql).await;
607+
assert_eq!(
608+
&DataType::Decimal(7, 7),
609+
actual[0].schema().field(0).data_type()
610+
);
611+
let expected = vec![
612+
"+-------------------------------------------------------------+",
613+
"| decimal_simple.c5 % CAST(Float64(0.00001) AS Decimal(5, 5)) |",
614+
"+-------------------------------------------------------------+",
615+
"| 0.0000040 |",
616+
"| 0.0000050 |",
617+
"| 0.0000090 |",
618+
"| 0.0000020 |",
619+
"| 0.0000050 |",
620+
"| 0.0000010 |",
621+
"| 0.0000040 |",
622+
"| 0.0000000 |",
623+
"| 0.0000000 |",
624+
"| 0.0000040 |",
625+
"| 0.0000020 |",
626+
"| 0.0000080 |",
627+
"| 0.0000030 |",
628+
"| 0.0000080 |",
629+
"| 0.0000000 |",
630+
"+-------------------------------------------------------------+",
631+
];
632+
assert_batches_eq!(expected, &actual);
633+
634+
let sql = "select c1%c5 from decimal_simple";
635+
let actual = execute_to_batches(&ctx, sql).await;
636+
assert_eq!(
637+
&DataType::Decimal(11, 7),
638+
actual[0].schema().field(0).data_type()
639+
);
640+
let expected = vec![
641+
"+--------------------------------------------+",
642+
"| decimal_simple.c1 Modulo decimal_simple.c5 |",
643+
"+--------------------------------------------+",
644+
"| 0.0000100 |",
645+
"| 0.0000200 |",
646+
"| 0.0000010 |",
647+
"| 0.0000300 |",
648+
"| 0.0000300 |",
649+
"| 0.0000080 |",
650+
"| 0.0000400 |",
651+
"| 0.0000000 |",
652+
"| 0.0000000 |",
653+
"| 0.0000400 |",
654+
"| 0.0000500 |",
655+
"| 0.0000500 |",
656+
"| 0.0000170 |",
657+
"| 0.0000500 |",
658+
"| 0.0000500 |",
659+
"+--------------------------------------------+",
660+
];
661+
assert_batches_eq!(expected, &actual);
662+
663+
Ok(())
664+
}
665+
366666
#[tokio::test]
367667
async fn decimal_sort() -> Result<()> {
368668
let ctx = SessionContext::new();

0 commit comments

Comments
 (0)