From e9a575f90b5927e90aa539cc419d6bd52fcd6365 Mon Sep 17 00:00:00 2001
From: Yash Mayya <yash.mayya@gmail.com>
Date: Sat, 5 Oct 2024 23:47:47 +0530
Subject: [PATCH] Cleanup conditional logic in
 NumericalFilterOptimizer::optimizeChild

---
 .../filter/NumericalFilterOptimizer.java      | 47 +++++++++----------
 1 file changed, 23 insertions(+), 24 deletions(-)

diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/filter/NumericalFilterOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/filter/NumericalFilterOptimizer.java
index fcd38f01b955..c434897379cb 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/filter/NumericalFilterOptimizer.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/filter/NumericalFilterOptimizer.java
@@ -85,41 +85,40 @@ Expression optimizeChild(Expression filterExpression, @Nullable Schema schema) {
     Function function = filterExpression.getFunctionCall();
     FilterKind kind = FilterKind.valueOf(function.getOperator());
     switch (kind) {
-      case IS_NULL:
-      case IS_NOT_NULL:
-        // No need to try to optimize IS_NULL and IS_NOT_NULL operations on numerical columns.
+      case BETWEEN: {
+        // Verify that value is a numeric column before rewriting.
+        List<Expression> operands = function.getOperands();
+        Expression value = operands.get(0);
+        DataType dataType = getDataType(value, schema);
+        if (dataType != null && dataType.isNumeric()) {
+          return rewriteBetweenExpression(filterExpression, dataType);
+        }
         break;
-      default:
+      }
+      case EQUALS:
+      case NOT_EQUALS:
+      case GREATER_THAN:
+      case GREATER_THAN_OR_EQUAL:
+      case LESS_THAN:
+      case LESS_THAN_OR_EQUAL: {
         List<Expression> operands = function.getOperands();
         // Verify that LHS is a numeric column and RHS is a numeric literal before rewriting.
         Expression lhs = operands.get(0);
         Expression rhs = operands.get(1);
+
         if (isNumericLiteral(rhs)) {
           DataType dataType = getDataType(lhs, schema);
           if (dataType != null && dataType.isNumeric()) {
-            switch (kind) {
-              case EQUALS:
-              case NOT_EQUALS:
-                return rewriteEqualsExpression(filterExpression, kind, dataType, rhs);
-              case GREATER_THAN:
-              case GREATER_THAN_OR_EQUAL:
-              case LESS_THAN:
-              case LESS_THAN_OR_EQUAL:
-                return rewriteRangeExpression(filterExpression, kind, dataType, rhs);
-              default:
-                break;
+            if (kind.isRange()) {
+              return rewriteRangeExpression(filterExpression, kind, dataType, rhs);
+            } else {
+              return rewriteEqualsExpression(filterExpression, kind, dataType, rhs);
             }
           }
         }
-
-        if (kind == FilterKind.BETWEEN) {
-          // Verify that value is a numeric column before rewriting.
-          Expression value = operands.get(0);
-          DataType dataType = getDataType(value, schema);
-          if (dataType != null && dataType.isNumeric()) {
-            return rewriteBetweenExpression(filterExpression, dataType);
-          }
-        }
+        break;
+      }
+      default:
         break;
     }
     return filterExpression;