Expr getCacheKey now delegates to children (#14287)

* Expr getCacheKey now delegates to children

* Removed the LOOKUP_EXPR_CACHE_KEY as we do not need it

* Adding an unit test

* Update processing/src/main/java/org/apache/druid/math/expr/Expr.java

Co-authored-by: Clint Wylie <cjwylie@gmail.com>

---------

Co-authored-by: Clint Wylie <cjwylie@gmail.com>

processing/src/main/java/org/apache/druid/math/expr/Expr.java

@@ -185,10 +185,30 @@ public interface Expr extends Cacheable
     throw Exprs.cannotVectorize(this);
   }
 
+
+  /**
+   * Decorates the {@link CacheKeyBuilder} for the default implementation of {@link #getCacheKey()}. The default cache
+   * key implementation includes the output of {@link #stringify()} and then uses a {@link Shuttle} to call this method
+   * on all children. The stringified representation is sufficient for most expressions, but for any which rely on
+   * external state that might change, this method allows the cache key to change when the state does, even if the
+   * expression itself is otherwise the same.
+   */
+  default void decorateCacheKeyBuilder(CacheKeyBuilder builder)
+  {
+    // no op
+  }
+
   @Override
   default byte[] getCacheKey()
   {
-    return new CacheKeyBuilder(Exprs.EXPR_CACHE_KEY).appendString(stringify()).build();
+    CacheKeyBuilder builder = new CacheKeyBuilder(Exprs.EXPR_CACHE_KEY).appendString(stringify());
+    // delegate to the child expressions through shuttle
+    Shuttle keyShuttle = expr -> {
+      expr.decorateCacheKeyBuilder(builder);
+      return expr;
+    };
+    this.visit(keyShuttle);
+    return builder.build();
   }
 
   /**

processing/src/main/java/org/apache/druid/math/expr/Exprs.java

@@ -30,7 +30,6 @@ import java.util.Stack;
 public class Exprs
 {
   public static final byte EXPR_CACHE_KEY = 0x00;
-  public static final byte LOOKUP_EXPR_CACHE_KEY = 0x01;
 
   public static UnsupportedOperationException cannotVectorize(Expr expr)
   {

processing/src/main/java/org/apache/druid/query/expression/LookupExprMacro.java

@@ -26,7 +26,6 @@ import org.apache.druid.math.expr.Expr;
 import org.apache.druid.math.expr.ExprEval;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.math.expr.ExpressionType;
-import org.apache.druid.math.expr.Exprs;
 import org.apache.druid.query.cache.CacheKeyBuilder;
 import org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider;
 import org.apache.druid.query.lookup.RegisteredLookupExtractionFn;
@@ -109,11 +108,9 @@ public class LookupExprMacro implements ExprMacroTable.ExprMacro
       }
 
       @Override
-      public byte[] getCacheKey()
+      public void decorateCacheKeyBuilder(CacheKeyBuilder builder)
       {
-        return new CacheKeyBuilder(Exprs.LOOKUP_EXPR_CACHE_KEY).appendString(stringify())
+        builder.appendCacheable(extractionFn);
-                                                               .appendCacheable(extractionFn)
-                                                               .build();
       }
     }
 

server/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java

@@ -81,6 +81,30 @@ public class LookupExprMacroTest extends InitializedNullHandlingTest
     }
   }
 
+  @Test
+  public void testCacheKeyChangesWhenLookupChangesSubExpr()
+  {
+    final String expression = "concat(lookup(x, 'lookyloo'))";
+    final Expr expr = Parser.parse(expression, LookupEnabledTestExprMacroTable.INSTANCE);
+    final Expr exprSameLookup = Parser.parse(expression, LookupEnabledTestExprMacroTable.INSTANCE);
+    final Expr exprChangedLookup = Parser.parse(
+        expression,
+        new ExprMacroTable(LookupEnabledTestExprMacroTable.makeTestMacros(ImmutableMap.of("x", "y", "a", "b")))
+    );
+    // same should have same cache key
+    Assert.assertArrayEquals(expr.getCacheKey(), exprSameLookup.getCacheKey());
+    // different should not have same key
+    final byte[] exprBytes = expr.getCacheKey();
+    final byte[] expr2Bytes = exprChangedLookup.getCacheKey();
+    if (exprBytes.length == expr2Bytes.length) {
+      // only check for equality if lengths are equal
+      boolean allEqual = true;
+      for (int i = 0; i < exprBytes.length; i++) {
+        allEqual = allEqual && (exprBytes[i] == expr2Bytes[i]);
+      }
+      Assert.assertFalse(allEqual);
+    }
+  }
 
   private void assertExpr(final String expression, final Object expectedResult)
   {