Add interface method for returning canonical lookup name (#16557)

* Add interface method for returning canonical lookup name * Address review comment * Add test in LookupReferencesManagerTest for coverage check * Add test in LookupSerdeModuleTest for coverage check
apache · Jun 5, 2024 · 6d7d2ff · 6d7d2ff
1 parent 7aecc09
commit 6d7d2ff
Show file tree

Hide file tree

Showing 14 changed files with 81 additions and 1 deletion.
diff --git a/benchmarks/src/test/java/org/apache/druid/benchmark/JoinAndLookupBenchmark.java b/benchmarks/src/test/java/org/apache/druid/benchmark/JoinAndLookupBenchmark.java
@@ -305,6 +305,12 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
                       return Optional.empty();
                     }
                   }
+
+                  @Override
+                  public String getCanonicalLookupName(String lookupName)
+                  {
+                    return lookupName;
+                  }
                 }
             )
         )

diff --git a/.../src/main/java/org/apache/druid/query/lookup/LookupExtractorFactoryContainerProvider.java b/.../src/main/java/org/apache/druid/query/lookup/LookupExtractorFactoryContainerProvider.java
@@ -40,4 +40,9 @@ public interface LookupExtractorFactoryContainerProvider
    * Returns a lookup container for the provided lookupName, if it exists.
    */
   Optional<LookupExtractorFactoryContainer> get(String lookupName);
+
+  /**
+   * Returns the canonical lookup name from a given lookup name, if special syntax exists.
+   */
+  String getCanonicalLookupName(String lookupName);
 }
diff --git a/processing/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java b/processing/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java
@@ -53,6 +53,12 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
           {
             return Optional.empty();
           }
+
+          @Override
+          public String getCanonicalLookupName(String lookupName)
+          {
+            return lookupName;
+          }
         })
     );
   }

diff --git a/server/src/main/java/org/apache/druid/query/lookup/LookupReferencesManager.java b/server/src/main/java/org/apache/druid/query/lookup/LookupReferencesManager.java
@@ -332,6 +332,12 @@ public Set<String> getAllLookupNames()
     return stateRef.get().lookupMap.keySet();
   }
 
+  @Override
+  public String getCanonicalLookupName(String lookupName)
+  {
+    return lookupName;
+  }
+
   // Note that this should ensure that "toLoad" and "toDrop" are disjoint.
   LookupsState<LookupExtractorFactoryContainer> getAllLookupsState()
   {

diff --git a/server/src/main/java/org/apache/druid/query/lookup/LookupSerdeModule.java b/server/src/main/java/org/apache/druid/query/lookup/LookupSerdeModule.java
@@ -79,5 +79,11 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
     {
       return Optional.empty();
     }
+
+    @Override
+    public String getCanonicalLookupName(String lookupName)
+    {
+      return lookupName;
+    }
   }
 }
diff --git a/server/src/test/java/org/apache/druid/query/expression/LookupEnabledTestExprMacroTable.java b/server/src/test/java/org/apache/druid/query/expression/LookupEnabledTestExprMacroTable.java
@@ -95,6 +95,12 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
           return Optional.empty();
         }
       }
+
+      @Override
+      public String getCanonicalLookupName(String lookupName)
+      {
+        return lookupName;
+      }
     };
   }
 

diff --git a/server/src/test/java/org/apache/druid/query/lookup/LookupReferencesManagerTest.java b/server/src/test/java/org/apache/druid/query/lookup/LookupReferencesManagerTest.java
@@ -579,6 +579,13 @@ public void testGetAllLookupNames() throws Exception
     );
   }
 
+  @Test
+  public void testGetCanonicalLookupName()
+  {
+    String lookupName = "lookupName1";
+    Assert.assertEquals(lookupName, lookupReferencesManager.getCanonicalLookupName(lookupName));
+  }
+
   @Test
   public void testGetAllLookupsState() throws Exception
   {

diff --git a/server/src/test/java/org/apache/druid/query/lookup/LookupSerdeModuleTest.java b/server/src/test/java/org/apache/druid/query/lookup/LookupSerdeModuleTest.java
@@ -117,4 +117,12 @@ public void testExpressionTransformSerde() throws Exception
         objectMapper.readValue(objectMapper.writeValueAsBytes(transform), ExpressionTransform.class)
     );
   }
+
+  @Test
+  public void testGetCanonicalLookupName()
+  {
+    LookupExtractorFactoryContainerProvider instance = injector.getInstance(LookupExtractorFactoryContainerProvider.class);
+    String lookupName = "lookupName1";
+    Assert.assertEquals(lookupName, instance.getCanonicalLookupName(lookupName));
+  }
 }
diff --git a/server/src/test/java/org/apache/druid/segment/LookupSegmentWranglerTest.java b/server/src/test/java/org/apache/druid/segment/LookupSegmentWranglerTest.java
@@ -67,6 +67,12 @@ public Optional<LookupExtractorFactoryContainer> get(final String lookupName)
             return Optional.empty();
           }
         }
+
+        @Override
+        public String getCanonicalLookupName(String lookupName)
+        {
+          return lookupName;
+        }
       }
   );
 

diff --git a/server/src/test/java/org/apache/druid/segment/join/LookupJoinableFactoryTest.java b/server/src/test/java/org/apache/druid/segment/join/LookupJoinableFactoryTest.java
@@ -79,6 +79,12 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
                 return Optional.empty();
               }
             }
+
+            @Override
+            public String getCanonicalLookupName(String lookupName)
+            {
+              return lookupName;
+            }
           }
       );
     }

diff --git a/server/src/test/java/org/apache/druid/server/SpecificSegmentsQuerySegmentWalker.java b/server/src/test/java/org/apache/druid/server/SpecificSegmentsQuerySegmentWalker.java
@@ -89,6 +89,12 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
         {
           return Optional.empty();
         }
+
+        @Override
+        public String getCanonicalLookupName(String lookupName)
+        {
+          return lookupName;
+        }
       };
 
   public static SpecificSegmentsQuerySegmentWalker createWalker(

diff --git a/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java b/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java
@@ -481,6 +481,12 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
       {
         return Optional.empty();
       }
+
+      @Override
+      public String getCanonicalLookupName(String lookupName)
+      {
+        return lookupName;
+      }
     }
 
     final TimeseriesQuery query =

diff --git a/...n/java/org/apache/druid/sql/calcite/expression/builtin/QueryLookupOperatorConversion.java b/...n/java/org/apache/druid/sql/calcite/expression/builtin/QueryLookupOperatorConversion.java
@@ -84,7 +84,7 @@ public DruidExpression toDruidExpression(
           final String lookupName = (String) lookupNameExpr.getLiteralValue();
 
           // Add the lookup name to the set of lookups to selectively load.
-          plannerContext.addLookupToLoad(lookupName);
+          plannerContext.addLookupToLoad(lookupExtractorFactoryContainerProvider.getCanonicalLookupName(lookupName));
 
           if (arg.isSimpleExtraction() && lookupNameExpr.isLiteral()) {
             return arg.getSimpleExtraction().cascade(

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestLookupProvider.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestLookupProvider.java
@@ -55,4 +55,10 @@ public Optional<LookupExtractorFactoryContainer> get(String lookupName)
       return Optional.empty();
     }
   }
+
+  @Override
+  public String getCanonicalLookupName(String lookupName)
+  {
+    return lookupName;
+  }
 }