[Backport] add VirtualColumns.findEquivalent and VirtualColumn.EquivalenceKey

processing/src/main/java/org/apache/druid/segment/VirtualColumn.java

@@ -21,6 +21,7 @@
 
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import org.apache.druid.annotations.SubclassesMustOverrideEqualsAndHashCode;
 import org.apache.druid.java.util.common.Cacheable;
 import org.apache.druid.query.dimension.DimensionSpec;
 import org.apache.druid.query.filter.ColumnIndexSelector;
@@ -336,4 +337,22 @@ default ColumnIndexSupplier getIndexSupplier(
   {
     return NoIndexesColumnIndexSupplier.getInstance();
   }
+
+  /**
+   * Returns a key used for "equivalence" comparisons, for checking if some virtual column is equivalent to some other
+   * virtual column, regardless of the output name. If this method returns null, it does not participate in equivalence
+   * comparisons.
+   * 
+   * @see VirtualColumns#findEquivalent(VirtualColumn) 
+   */
+  @Nullable
+  default EquivalenceKey getEquivalanceKey()
+  {
+    return null;
+  }
+
+  @SubclassesMustOverrideEqualsAndHashCode
+  interface EquivalenceKey
+  {
+  }
 }

processing/src/main/java/org/apache/druid/segment/VirtualColumns.java

@@ -26,6 +26,7 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import org.apache.druid.java.util.common.Cacheable;
 import org.apache.druid.java.util.common.IAE;
@@ -131,10 +132,13 @@ public static VirtualColumns nullToEmpty(@Nullable VirtualColumns virtualColumns
   // For equals, hashCode, toString, and serialization:
   private final List<VirtualColumn> virtualColumns;
   private final List<String> virtualColumnNames;
+  // For equivalence
+  private final Map<VirtualColumn.EquivalenceKey, VirtualColumn> equivalence;
 
   // For getVirtualColumn:
   private final Map<String, VirtualColumn> withDotSupport;
   private final Map<String, VirtualColumn> withoutDotSupport;
+  private final boolean hasNoDotColumns;
 
   private VirtualColumns(
       List<VirtualColumn> virtualColumns,
@@ -146,10 +150,15 @@ private VirtualColumns(
     this.withDotSupport = withDotSupport;
     this.withoutDotSupport = withoutDotSupport;
     this.virtualColumnNames = new ArrayList<>(virtualColumns.size());
-
+    this.hasNoDotColumns = withDotSupport.isEmpty();
+    this.equivalence = Maps.newHashMapWithExpectedSize(virtualColumns.size());
     for (VirtualColumn virtualColumn : virtualColumns) {
       detectCycles(virtualColumn, null);
       virtualColumnNames.add(virtualColumn.getOutputName());
+      VirtualColumn.EquivalenceKey key = virtualColumn.getEquivalanceKey();
+      if (key != null) {
+        equivalence.put(key, virtualColumn);
+      }
     }
   }
 
@@ -172,10 +181,23 @@ public VirtualColumn getVirtualColumn(String columnName)
     if (vc != null) {
       return vc;
     }
+    if (hasNoDotColumns) {
+      return null;
+    }
     final String baseColumnName = splitColumnName(columnName).lhs;
     return withDotSupport.get(baseColumnName);
   }
 
+  /**
+   * Check if a virtual column is already defined which is the same as some other virtual column, ignoring output name,
+   * returning that virtual column if it exists, or null if there is no equivalent virtual column.
+   */
+  @Nullable
+  public VirtualColumn findEquivalent(VirtualColumn virtualColumn)
+  {
+    return equivalence.get(virtualColumn.getEquivalanceKey());
+  }
+
   /**
    * Get the {@link ColumnIndexSupplier} of the specified virtual column, with the assistance of a
    * {@link ColumnSelector} to allow reading things from segments. If the column does not have indexes this method

processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVirtualColumn.java

@@ -61,9 +61,7 @@ public class ExpressionVirtualColumn implements VirtualColumn
   private static final Logger log = new Logger(ExpressionVirtualColumn.class);
 
   private final String name;
-  private final String expression;
-  @Nullable
-  private final ColumnType outputType;
+  private final Expression expression;
   private final Supplier<Expr> parsedExpression;
   private final Supplier<byte[]> cacheKey;
 
@@ -126,8 +124,7 @@ private ExpressionVirtualColumn(
   )
   {
     this.name = Preconditions.checkNotNull(name, "name");
-    this.expression = Preconditions.checkNotNull(expression, "expression");
-    this.outputType = outputType;
+    this.expression = new Expression(Preconditions.checkNotNull(expression, "expression"), outputType);
     this.parsedExpression = parsedExpression;
     this.cacheKey = makeCacheKeySupplier();
   }
@@ -142,14 +139,14 @@ public String getOutputName()
   @JsonProperty
   public String getExpression()
   {
-    return expression;
+    return expression.expressionString;
   }
 
   @Nullable
   @JsonProperty
   public ColumnType getOutputType()
   {
-    return outputType;
+    return expression.outputType;
   }
 
   @JsonIgnore
@@ -273,7 +270,7 @@ public ColumnIndexSupplier getIndexSupplier(
       ColumnIndexSelector columnIndexSelector
   )
   {
-    return getParsedExpression().get().asColumnIndexSupplier(columnIndexSelector, outputType);
+    return getParsedExpression().get().asColumnIndexSupplier(columnIndexSelector, expression.outputType);
   }
 
   @Override
@@ -283,7 +280,7 @@ public ColumnCapabilities capabilities(String columnName)
     // are unable to compute the output type of the expression, either due to incomplete type information of the
     // inputs or because of unimplemented methods on expression implementations themselves, or, because a
     // ColumnInspector is not available
-
+    final ColumnType outputType = expression.outputType;
     if (ExpressionProcessing.processArraysAsMultiValueStrings() && outputType != null && outputType.isArray()) {
       return new ColumnCapabilitiesImpl().setType(ColumnType.STRING).setHasMultipleValues(true);
     }
@@ -299,6 +296,8 @@ public ColumnCapabilities capabilities(ColumnInspector inspector, String columnN
       return inspector.getColumnCapabilities(parsedExpression.get().getBindingIfIdentifier());
     }
 
+    final ColumnType outputType = expression.outputType;
+
     final ExpressionPlan plan = ExpressionPlanner.plan(inspector, parsedExpression.get());
     final ColumnCapabilities inferred = plan.inferColumnCapabilities(outputType);
     // if we can infer the column capabilities from the expression plan, then use that
@@ -311,14 +310,14 @@ public ColumnCapabilities capabilities(ColumnInspector inspector, String columnN
           log.warn(
               "Projected output type %s of expression %s does not match provided type %s",
               inferred.asTypeString(),
-              expression,
+              expression.expressionString,
               outputType
           );
         } else {
           log.debug(
               "Projected output type %s of expression %s does not match provided type %s",
               inferred.asTypeString(),
-              expression,
+              expression.expressionString,
               outputType
           );
         }
@@ -348,6 +347,13 @@ public byte[] getCacheKey()
     return cacheKey.get();
   }
 
+  @Nullable
+  @Override
+  public EquivalenceKey getEquivalanceKey()
+  {
+    return expression;
+  }
+
   @Override
   public boolean equals(final Object o)
   {
@@ -359,23 +365,21 @@ public boolean equals(final Object o)
     }
     final ExpressionVirtualColumn that = (ExpressionVirtualColumn) o;
     return Objects.equals(name, that.name) &&
-           Objects.equals(expression, that.expression) &&
-           Objects.equals(outputType, that.outputType);
+           Objects.equals(expression, that.expression);
   }
 
   @Override
   public int hashCode()
   {
-    return Objects.hash(name, expression, outputType);
+    return Objects.hash(name, expression);
   }
 
   @Override
   public String toString()
   {
     return "ExpressionVirtualColumn{" +
            "name='" + name + '\'' +
-           ", expression='" + expression + '\'' +
-           ", outputType=" + outputType +
+           ", expression=" + expression +
            '}';
   }
 
@@ -389,10 +393,10 @@ private boolean isDirectAccess(final ColumnInspector inspector)
       final ColumnCapabilities baseCapabilities =
           inspector.getColumnCapabilities(parsedExpression.get().getBindingIfIdentifier());
 
-      if (outputType == null) {
+      if (expression.outputType == null) {
         // No desired output type. Anything from the source is fine.
         return true;
-      } else if (baseCapabilities != null && outputType.equals(baseCapabilities.toColumnType())) {
+      } else if (baseCapabilities != null && expression.outputType.equals(baseCapabilities.toColumnType())) {
         // Desired output type matches the type from the source.
         return true;
       }
@@ -408,10 +412,57 @@ private Supplier<byte[]> makeCacheKeySupplier()
           .appendString(name)
           .appendCacheable(parsedExpression.get());
 
-      if (outputType != null) {
-        builder.appendString(outputType.toString());
+      if (expression.outputType != null) {
+        builder.appendString(expression.outputType.toString());
       }
       return builder.build();
     });
   }
+
+  /**
+   * {@link VirtualColumn.EquivalenceKey} for expressions. Note that this does not check true equivalence of
+   * expressions, for example it will not currently consider something like 'a + b' equivalent to 'b + a'. This is ok
+   * for current uses of this functionality, but in theory we could push down equivalence to the parsed expression
+   * instead of checking for an identical string expression, it would just be a lot more expensive.
+   */
+  private static final class Expression implements EquivalenceKey
+  {
+    private final String expressionString;
+    @Nullable
+    private final ColumnType outputType;
+
+    private Expression(String expression, @Nullable ColumnType outputType)
+    {
+      this.expressionString = expression;
+      this.outputType = outputType;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      Expression that = (Expression) o;
+      return Objects.equals(expressionString, that.expressionString) && Objects.equals(outputType, that.outputType);
+    }
+
+    @Override
+    public int hashCode()
+    {
+      return Objects.hash(expressionString, outputType);
+    }
+
+    @Override
+    public String toString()
+    {
+      return "Expression{" +
+             "expression='" + expressionString + '\'' +
+             ", outputType=" + outputType +
+             '}';
+    }
+  }
 }