Add Option to Optimize NOT Predicates with Push Down

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -142,6 +142,10 @@ public static boolean isSkipScanFilterReorder(Map<String, String> queryOptions)
     return "false".equalsIgnoreCase(queryOptions.get(QueryOptionKey.USE_SCAN_REORDER_OPTIMIZATION));
   }
 
+  public static boolean isPushDownNot(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.PUSH_DOWN_NOT));
+  }
+
   @Nullable
   public static Map<String, Set<FieldConfig.IndexType>> getSkipIndexes(Map<String, String> queryOptions) {
     // Example config:  skipIndexes='col1=inverted,range&col2=inverted'

pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/QueryOptimizer.java

@@ -24,12 +24,14 @@
 import javax.annotation.Nullable;
 import org.apache.pinot.common.request.Expression;
 import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.common.utils.config.QueryOptionsUtils;
 import org.apache.pinot.core.query.optimizer.filter.FilterOptimizer;
 import org.apache.pinot.core.query.optimizer.filter.FlattenAndOrFilterOptimizer;
 import org.apache.pinot.core.query.optimizer.filter.IdenticalPredicateFilterOptimizer;
 import org.apache.pinot.core.query.optimizer.filter.MergeEqInFilterOptimizer;
 import org.apache.pinot.core.query.optimizer.filter.MergeRangeFilterOptimizer;
 import org.apache.pinot.core.query.optimizer.filter.NumericalFilterOptimizer;
+import org.apache.pinot.core.query.optimizer.filter.PushDownNotFilterOptimizer;
 import org.apache.pinot.core.query.optimizer.filter.TextMatchFilterOptimizer;
 import org.apache.pinot.core.query.optimizer.filter.TimePredicateFilterOptimizer;
 import org.apache.pinot.core.query.optimizer.statement.StatementOptimizer;
@@ -69,6 +71,9 @@ public void optimize(PinotQuery pinotQuery, @Nullable TableConfig tableConfig, @
       for (FilterOptimizer filterOptimizer : FILTER_OPTIMIZERS) {
         filterExpression = filterOptimizer.optimize(filterExpression, schema);
       }
+      if (pinotQuery.getQueryOptions() != null && QueryOptionsUtils.isPushDownNot(pinotQuery.getQueryOptions())) {
+        filterExpression = new PushDownNotFilterOptimizer().optimize(filterExpression, schema);
+      }
       pinotQuery.setFilterExpression(filterExpression);
     }
 

pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/filter/PushDownNotFilterOptimizer.java

@@ -0,0 +1,128 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.optimizer.filter;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.pinot.common.request.Expression;
+import org.apache.pinot.common.request.ExpressionType;
+import org.apache.pinot.common.request.Function;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.sql.FilterKind;
+
+/**
+ * The {@code PushdownNotFilterOptimizer} pushes down NOT expressions to simplify and flatten filter expressions,
+ * applying DeMorgan's law recursively. For example:
+ * <ul>
+ *   <li>NOT(X AND Y) becomes NOT(X) OR NOT(Y)</li>
+ *   <li>NOT(X OR Y) becomes NOT(X) AND NOT(Y)</li>
+ *   <li>Handles nested cases like NOT(X AND (Y OR Z))</li>
+ * </ul>
+ */
+public class PushDownNotFilterOptimizer implements FilterOptimizer {
+
+  @Override
+  public Expression optimize(Expression filterExpression, Schema schema) {
+    return filterExpression.getType() == ExpressionType.FUNCTION ? optimize(filterExpression) : filterExpression;
+  }
+
+  @VisibleForTesting
+  Expression optimize(Expression expression) {
+    Function function = expression.getFunctionCall();
+    FilterKind filterKind;
+    try {
+      filterKind = FilterKind.valueOf(function.getOperator());
+    } catch (Exception e) {
+      return expression;
+    }
+    List<Expression> operands = function.getOperands();
+
+    if (filterKind == FilterKind.NOT) {
+      Expression operand = operands.get(0);
+      if (operand.getType() == ExpressionType.FUNCTION) {
+        Function innerFunction = operand.getFunctionCall();
+        FilterKind innerFilterKind = FilterKind.valueOf(innerFunction.getOperator());
+
+        if (innerFilterKind == FilterKind.AND) {
+          return optimize(createOrExpression(innerFunction.getOperands()));
+        } else if (innerFilterKind == FilterKind.OR) {
+          return optimize(createAndExpression(innerFunction.getOperands()));
+        }
+      }
+      return expression;
+    }
+
+    for (int i = 0; i < operands.size(); i++) {
+      operands.set(i, optimize(operands.get(i)));
+    }
+
+    expression.setType(ExpressionType.FUNCTION);
+    return expression;
+  }
+
+  /**
+   * Creates an OR expression by applying NOT to each operand and connecting them with OR.
+   */
+  private Expression createOrExpression(List<Expression> operands) {
+    Function orFunction = new Function();
+    orFunction.setOperator(FilterKind.OR.name());
+    List<Expression> orOperands = new ArrayList<>();
+    for (Expression operand : operands) {
+      orOperands.add(createNotExpression(operand));
+    }
+    orFunction.setOperands(orOperands);
+    Expression orExpression = new Expression();
+    orExpression.setFunctionCall(orFunction);
+    orExpression.setType(ExpressionType.FUNCTION);
+    return orExpression;
+  }
+
+  /**
+   * Creates an AND expression by applying NOT to each operand and connecting them with AND.
+   */
+  private Expression createAndExpression(List<Expression> operands) {
+    Function andFunction = new Function();
+    andFunction.setOperator(FilterKind.AND.name());
+    List<Expression> andOperands = new ArrayList<>();
+    for (Expression operand : operands) {
+      andOperands.add(createNotExpression(operand));
+    }
+    andFunction.setOperands(andOperands);
+    Expression andExpression = new Expression();
+    andExpression.setFunctionCall(andFunction);
+    andExpression.setType(ExpressionType.FUNCTION);
+    return andExpression;
+  }
+
+  /**
+   * Wraps an operand in a NOT expression.
+   */
+  private Expression createNotExpression(Expression operand) {
+    Function notFunction = new Function();
+    notFunction.setOperator(FilterKind.NOT.name());
+    List<Expression> notOperands = new ArrayList<>();
+    notOperands.add(operand);
+    notFunction.setOperands(notOperands);
+    Expression notExpression = new Expression();
+    notExpression.setFunctionCall(notFunction);
+    notExpression.setType(ExpressionType.FUNCTION);
+    return notExpression;
+  }
+}

pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/filter/PushDownNotFilterOptimizerTest.java

@@ -0,0 +1,53 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.optimizer.filter;
+
+import org.apache.pinot.common.request.Expression;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+public class PushDownNotFilterOptimizerTest {
+  private static final PushDownNotFilterOptimizer OPTIMIZER = new PushDownNotFilterOptimizer();
+
+  @Test
+  public void testPushDownNotFilterOptimizer() {
+    testPushDownNot("NOT(AirTime = 141 AND ArrDelay > 1)", "NOT(AirTime = 141) OR NOT(ArrDelay > 1)");
+    testPushDownNot("NOT(AirTime = 141 AND ArrDelay > 1 OR ArrTime > 808)", "(NOT(AirTime = 141) OR NOT(ArrDelay > 1)) "
+        + "AND NOT(ArrTime > 808)");
+    testPushDownNot("NOT(AirTime = 141 OR ArrDelay > 1)", "NOT(AirTime = 141) AND NOT(ArrDelay > 1)");
+    testPushDownNot("NOT(AirTime = 141 OR ArrDelay > 1 OR ArrTime > 808)", "NOT(AirTime = 141) AND NOT(ArrDelay > 1) "
+        + "AND NOT(ArrTime > 808)");
+  }
+
+  @Test
+  public void testNoOptimizerChange() {
+    testPushDownNot("NOT(AirTime = 141)", "NOT(AirTime = 141)");
+    testPushDownNot("AirTime = 141", "AirTime = 141");
+    testPushDownNot("true", "true");
+  }
+
+  private void testPushDownNot(String filterString, String expectedOptimizedFilterString) {
+    Expression expectedExpression = CalciteSqlParser.compileToExpression(expectedOptimizedFilterString);
+    Expression optimizedFilterExpression = OPTIMIZER.optimize(CalciteSqlParser.compileToExpression(filterString));
+    assertEquals(optimizedFilterExpression, expectedExpression);
+  }
+}

pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java

@@ -489,6 +489,8 @@ public static class QueryOptionKey {
         // possible.
         public static final String OPTIMIZE_MAX_INITIAL_RESULT_HOLDER_CAPACITY =
             "optimizeMaxInitialResultHolderCapacity";
+
+        public static final String PUSH_DOWN_NOT = "false";
       }
 
       public static class QueryOptionValue {