diff --git a/processing/src/main/java/org/apache/druid/math/expr/Expr.java b/processing/src/main/java/org/apache/druid/math/expr/Expr.java index 50eabf3dac5..11bc0c922a6 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/Expr.java +++ b/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(); } /** diff --git a/processing/src/main/java/org/apache/druid/math/expr/Exprs.java b/processing/src/main/java/org/apache/druid/math/expr/Exprs.java index 618afa5b4f0..e8ad020fe70 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/Exprs.java +++ b/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) { diff --git a/processing/src/main/java/org/apache/druid/query/expression/LookupExprMacro.java b/processing/src/main/java/org/apache/druid/query/expression/LookupExprMacro.java index 07fd7631d88..f824038586c 100644 --- a/processing/src/main/java/org/apache/druid/query/expression/LookupExprMacro.java +++ b/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()) - .appendCacheable(extractionFn) - .build(); + builder.appendCacheable(extractionFn); } } diff --git a/server/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java b/server/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java index a4bb1ca0d3c..d0097bcc810 100644 --- a/server/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java +++ b/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) {