fix filtering nested field virtual column when used with non nested column input

processing/src/main/java/org/apache/druid/query/dimension/DefaultDimensionSpec.java

@@ -41,6 +41,11 @@ public static DefaultDimensionSpec of(String dimensionName)
     return new DefaultDimensionSpec(dimensionName, dimensionName);
   }
 
+  public static DefaultDimensionSpec of(String dimensionName, ColumnType columnType)
+  {
+    return new DefaultDimensionSpec(dimensionName, dimensionName, columnType);
+  }
+
   private static final byte CACHE_TYPE_ID = 0x0;
   private final String dimension;
   private final String outputName;

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

@@ -257,17 +257,18 @@ default VectorObjectSelector makeVectorObjectSelector(
 
   /**
    * Return the {@link ColumnCapabilities} which best describe the optimal selector to read from this virtual column.
-   *
+   * <p>
    * The {@link ColumnInspector} (most likely corresponding to an underlying {@link ColumnSelectorFactory} of a query)
    * allows the virtual column to consider this information if necessary to compute its output type details.
-   *
+   * <p>
    * Examples of this include the {@link ExpressionVirtualColumn}, which takes input from other columns and uses the
    * {@link ColumnInspector} to infer the output type of expressions based on the types of the inputs.
    *
    * @param inspector column inspector to provide additional information of other available columns
    * @param columnName the name this virtual column was referenced with
    * @return capabilities, must not be null
    */
+  @Nullable
   default ColumnCapabilities capabilities(ColumnInspector inspector, String columnName)
   {
     return capabilities(columnName);

processing/src/main/java/org/apache/druid/segment/nested/NestedDataComplexColumn.java

@@ -20,12 +20,9 @@
 
 package org.apache.druid.segment.nested;
 
-import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.query.extraction.ExtractionFn;
-import org.apache.druid.segment.ColumnSelector;
 import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.DimensionSelector;
-import org.apache.druid.segment.column.BaseColumn;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.column.ColumnIndexSupplier;
 import org.apache.druid.segment.column.ColumnType;
@@ -47,29 +44,6 @@
  */
 public abstract class NestedDataComplexColumn implements ComplexColumn
 {
-  @Nullable
-  public static NestedDataComplexColumn fromColumnSelector(
-      ColumnSelector columnSelector,
-      String columnName
-  )
-  {
-    ColumnHolder holder = columnSelector.getColumnHolder(columnName);
-    if (holder == null) {
-      return null;
-    }
-    BaseColumn theColumn = holder.getColumn();
-    if (theColumn instanceof CompressedNestedDataComplexColumn) {
-      return (CompressedNestedDataComplexColumn) theColumn;
-    }
-    throw new IAE(
-        "Column [%s] is invalid type, found [%s] instead of [%s]",
-        columnName,
-        theColumn.getClass(),
-        NestedDataComplexColumn.class.getSimpleName()
-    );
-  }
-
-
   /**
    * Make a {@link DimensionSelector} for a nested literal field column associated with this nested
    * complex column specified by a sequence of {@link NestedPathPart}.

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

@@ -194,6 +194,7 @@ public ColumnCapabilities capabilities(String columnName)
     return new ColumnCapabilitiesImpl().setType(outputType == null ? ColumnType.FLOAT : outputType);
   }
 
+  @Nullable
   @Override
   public ColumnCapabilities capabilities(ColumnInspector inspector, String columnName)
   {

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

@@ -163,6 +163,7 @@ public ColumnCapabilities capabilities(String columnName)
     return ColumnCapabilitiesImpl.createDefault();
   }
 
+  @Nullable
   @SuppressWarnings("ConstantConditions")
   @Override
   public ColumnCapabilities capabilities(ColumnInspector inspector, String columnName)

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

@@ -158,6 +158,7 @@ public ColumnCapabilities capabilities(String columnName)
                                        .setHasBitmapIndexes(true);
   }
 
+  @Nullable
   @Override
   public ColumnCapabilities capabilities(ColumnInspector inspector, String columnName)
   {

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

@@ -54,6 +54,7 @@
 import org.apache.druid.segment.column.ValueTypes;
 import org.apache.druid.segment.data.IndexedInts;
 import org.apache.druid.segment.data.ReadableOffset;
+import org.apache.druid.segment.nested.CompressedNestedDataComplexColumn;
 import org.apache.druid.segment.nested.NestedDataComplexColumn;
 import org.apache.druid.segment.nested.NestedDataComplexTypeSerde;
 import org.apache.druid.segment.nested.NestedPathArrayElement;
@@ -410,6 +411,9 @@ public VectorObjectSelector makeVectorObjectSelector(
     // not a nested column, but we can still do stuff if the path is the 'root', indicated by an empty path parts
     if (parts.isEmpty()) {
       ColumnCapabilities capabilities = holder.getCapabilities();
+      // expectedType shouldn't possibly be null if we are being asked for an object selector and the underlying column
+      // is numeric, else we would have been asked for a value selector
+      Preconditions.checkArgument(expectedType != null, "Asked for a VectorObjectSelector on a numeric column, 'expectedType' must not be null");
       if (capabilities.isNumeric()) {
         return ExpressionVectorSelectors.castValueSelectorToObject(
             offset,
@@ -600,12 +604,18 @@ public ColumnIndexSupplier getIndexSupplier(
       ColumnSelector selector
   )
   {
-    final NestedDataComplexColumn column = NestedDataComplexColumn.fromColumnSelector(selector, this.columnName);
-
-    if (column == null) {
+    ColumnHolder holder = selector.getColumnHolder(this.columnName);
+    if (holder == null) {
       return null;
     }
-    return column.getColumnIndexSupplier(parts);
+    BaseColumn theColumn = holder.getColumn();
+    if (theColumn instanceof CompressedNestedDataComplexColumn) {
+      return ((CompressedNestedDataComplexColumn<?>) theColumn).getColumnIndexSupplier(parts);
+    }
+    if (parts.isEmpty()) {
+      return holder.getIndexSupplier();
+    }
+    return null;
   }
 
   @Override
@@ -625,6 +635,7 @@ public ColumnCapabilities capabilities(String columnName)
                                  .setHasNulls(expectedType == null || !expectedType.isNumeric() || (expectedType.isNumeric() && NullHandling.sqlCompatible()));
   }
 
+  @Nullable
   @Override
   public ColumnCapabilities capabilities(ColumnInspector inspector, String columnName)
   {
@@ -637,17 +648,38 @@ public ColumnCapabilities capabilities(ColumnInspector inspector, String columnN
     }
     // ColumnInspector isn't really enough... we need the ability to read the complex column itself to examine
     // the nested fields type information to really be accurate here, so we rely on the expectedType to guide us
-    final ColumnCapabilities complexCapabilites = inspector.getColumnCapabilities(this.columnName);
-    if (complexCapabilites != null && complexCapabilites.isDictionaryEncoded().isTrue()) {
-      return ColumnCapabilitiesImpl.createDefault()
-                                   .setType(expectedType != null ? expectedType : ColumnType.STRING)
-                                   .setDictionaryEncoded(true)
-                                   .setDictionaryValuesSorted(true)
-                                   .setDictionaryValuesUnique(true)
-                                   .setHasBitmapIndexes(true)
-                                   .setHasNulls(expectedType == null || (expectedType.isNumeric()
-                                                                         && NullHandling.sqlCompatible()));
+    final ColumnCapabilities capabilities = inspector.getColumnCapabilities(this.columnName);
+
+    if (capabilities != null) {
+      // if the underlying column is a nested column (and persisted to disk, re: the dictionary encoded check)
+      if (capabilities.is(ValueType.COMPLEX) &&
+          capabilities.getComplexTypeName().equals(NestedDataComplexTypeSerde.TYPE_NAME) &&
+          capabilities.isDictionaryEncoded().isTrue()) {
+        return ColumnCapabilitiesImpl.createDefault()
+                                     .setType(expectedType != null ? expectedType : ColumnType.STRING)
+                                     .setDictionaryEncoded(true)
+                                     .setDictionaryValuesSorted(true)
+                                     .setDictionaryValuesUnique(true)
+                                     .setHasBitmapIndexes(true)
+                                     .setHasNulls(expectedType == null || (expectedType.isNumeric()
+                                                                           && NullHandling.sqlCompatible()));
+      }
+      // column is not nested, use underlying column capabilities, adjusted for expectedType as necessary
+      if (parts.isEmpty()) {
+        ColumnCapabilitiesImpl copy = ColumnCapabilitiesImpl.copyOf(capabilities);
+        if (expectedType != null) {
+          copy.setType(expectedType);
+          copy.setHasNulls(
+              copy.hasNulls().or(ColumnCapabilities.Capable.of(expectedType.getType() != capabilities.getType()))
+          );
+        }
+        return copy;
+      } else if (capabilities.isPrimitive()) {
+        // path doesn't exist and column isn't nested, so effectively column doesn't exist
+        return null;
+      }
     }
+
     return capabilities(columnName);
   }