better type coalescing in arithmetic operators

PastorGL · Jan 20, 2025 · 7fe98a8 · 7fe98a8
1 parent be3f4f5
commit 7fe98a8
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 17 deletions.
diff --git a/...mmons/src/main/java/io/github/pastorgl/datacooker/commons/operators/NumericOperators.java b/...mmons/src/main/java/io/github/pastorgl/datacooker/commons/operators/NumericOperators.java
@@ -283,7 +283,7 @@ protected Long op0(Deque<Object> args) {
         }
     }
 
-    public static class ADD extends Operator.Binary<Double, Double, Double> {
+    public static class ADD extends Operator.Binary<Number, Number, Number> {
         @Override
         public int prio() {
             return 125;
@@ -305,14 +305,20 @@ public String descr() {
         }
 
         @Override
-        protected Double op0(Deque<Object> args) {
-            double a = Evaluator.popDouble(args);
-            double b = Evaluator.popDouble(args);
-            return a + b;
+        protected Number op0(Deque<Object> args) {
+            Number a = Evaluator.popNumeric(args);
+            Number b = Evaluator.popNumeric(args);
+            if (a instanceof Double || b instanceof Double) {
+                return a.doubleValue() + b.doubleValue();
+            }
+            if (a instanceof Long || b instanceof Long) {
+                return a.longValue() + b.longValue();
+            }
+            return a.intValue() + b.intValue();
         }
     }
 
-    public static class SUB extends Operator.Binary<Double, Double, Double> {
+    public static class SUB extends Operator.Binary<Number, Number, Number> {
         @Override
         public int prio() {
             return 125;
@@ -334,14 +340,20 @@ public String descr() {
         }
 
         @Override
-        protected Double op0(Deque<Object> args) {
-            double a = Evaluator.popDouble(args);
-            double b = Evaluator.popDouble(args);
-            return a - b;
+        protected Number op0(Deque<Object> args) {
+            Number a = Evaluator.popNumeric(args);
+            Number b = Evaluator.popNumeric(args);
+            if (a instanceof Double || b instanceof Double) {
+                return a.doubleValue() - b.doubleValue();
+            }
+            if (a instanceof Long || b instanceof Long) {
+                return a.longValue() - b.longValue();
+            }
+            return a.intValue() - b.intValue();
         }
     }
 
-    public static class MUL extends Operator.Binary<Double, Double, Double> {
+    public static class MUL extends Operator.Binary<Number, Number, Number> {
         @Override
         public int prio() {
             return 130;
@@ -363,14 +375,20 @@ public String descr() {
         }
 
         @Override
-        protected Double op0(Deque<Object> args) {
-            double a = Evaluator.popDouble(args);
-            double b = Evaluator.popDouble(args);
-            return a * b;
+        protected Number op0(Deque<Object> args) {
+            Number a = Evaluator.popNumeric(args);
+            Number b = Evaluator.popNumeric(args);
+            if (a instanceof Double || b instanceof Double) {
+                return a.doubleValue() * b.doubleValue();
+            }
+            if (a instanceof Long || b instanceof Long) {
+                return a.longValue() * b.longValue();
+            }
+            return a.intValue() * b.intValue();
         }
     }
 
-    public static class DIV extends Operator.Binary<Double, Double, Double> {
+    public static class DIV extends Operator.Binary<Double, Number, Number> {
         @Override
         public int prio() {
             return 130;

diff --git a/datacooker-commons/src/test/resources/test.columnarSelect.tdl b/datacooker-commons/src/test/resources/test.columnarSelect.tdl
@@ -12,7 +12,7 @@ SELECT userid,num,acc INTO ret7 FROM "source" WHERE acc < 15.0 OR acc >= 100.0;
 SELECT * INTO ret8 FROM "source" WHERE NOT (trackid MATCH '.+?non.*' AND pt != 'e2e');
 SELECT * INTO ret9 FROM "source" WHERE unknown IS NOT NULL;
 SELECT num INTO ret10 FROM "source" WHERE num BETWEEN 8 AND 15.0;
-SELECT num INTO ret20 FROM "source" WHERE num + 0 IN [8.,9.,10.];
+SELECT num INTO ret20 FROM "source" WHERE num + 0 IN [8,9,10];
 SELECT "type" INTO ret21 FROM "source" WHERE "type" NOT IN ['car','bike'];
 
 COPY ret * hadoopText() INTO 'ret';