fix: clustered segment catalog spec complex handling

multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java

@@ -31,8 +31,6 @@
 import org.apache.druid.java.util.emitter.service.AlertEvent;
 import org.apache.druid.segment.AutoTypeColumnSchema;
 import org.apache.druid.segment.DimensionHandlerUtils;
-import org.apache.druid.segment.column.ColumnCapabilities;
-import org.apache.druid.segment.column.ColumnCapabilitiesImpl;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.ValueType;
 
@@ -81,9 +79,7 @@ public static DimensionSchema createDimensionSchema(
       // for complex types that are not COMPLEX<json>, we still want to use the handler since 'auto' typing
       // only works for the 'standard' built-in types
       if (queryType != null && queryType.is(ValueType.COMPLEX) && !ColumnType.NESTED_DATA.equals(queryType)) {
-        final ColumnCapabilities capabilities = ColumnCapabilitiesImpl.createDefault().setType(queryType);
-        return DimensionHandlerUtils.getHandlerFromCapabilities(column, capabilities, null)
-                                    .getDimensionSchema(capabilities);
+        return DimensionHandlerUtils.getComplexDimensionSchema(column, queryType);
       }
 
       if (queryType != null && (queryType.isPrimitive() || queryType.isPrimitiveArray())) {
@@ -111,9 +107,7 @@ public static DimensionSchema createDimensionSchema(
       } else if (dimensionType.getType() == ValueType.ARRAY) {
         return new AutoTypeColumnSchema(column, dimensionType, null);
       } else {
-        final ColumnCapabilities capabilities = ColumnCapabilitiesImpl.createDefault().setType(dimensionType);
-        return DimensionHandlerUtils.getHandlerFromCapabilities(column, capabilities, null)
-                                    .getDimensionSchema(capabilities);
+        return DimensionHandlerUtils.getComplexDimensionSchema(column, dimensionType);
       }
     }
   }

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

@@ -24,6 +24,7 @@
 import com.google.common.primitives.Doubles;
 import com.google.common.primitives.Floats;
 import org.apache.druid.common.guava.GuavaUtils;
+import org.apache.druid.data.input.impl.DimensionSchema;
 import org.apache.druid.data.input.impl.DimensionSchema.MultiValueHandling;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.java.util.common.IAE;
@@ -130,17 +131,44 @@ private DimensionHandlerUtils()
     }
 
     if (capabilities.is(ValueType.COMPLEX) && capabilities.getComplexTypeName() != null) {
-      DimensionHandlerProvider provider = DIMENSION_HANDLER_PROVIDERS.get(capabilities.getComplexTypeName());
-      if (provider == null) {
-        throw new ISE("Can't find DimensionHandlerProvider for typeName [%s]", capabilities.getComplexTypeName());
-      }
-      return provider.get(dimensionName);
+      return getHandlerForComplexType(dimensionName, capabilities.getComplexTypeName());
     }
 
     // Return a StringDimensionHandler by default (null columns will be treated as String typed)
     return new StringDimensionHandler(dimensionName, multiValueHandling, true, false);
   }
 
+  /**
+   * The {@link DimensionHandler} registered for a complex type. Complex columns are stored by type-specific handlers,
+   * so a type contributed by an extension becomes storable as soon as that extension registers one.
+   *
+   * @throws ISE if no handler is registered for the type, which usually means the extension defining it is not loaded
+   */
+  public static DimensionHandler<?, ?, ?> getHandlerForComplexType(String dimensionName, String complexTypeName)
+  {
+    final DimensionHandlerProvider provider = DIMENSION_HANDLER_PROVIDERS.get(complexTypeName);
+    if (provider == null) {
+      throw new ISE("Can't find DimensionHandlerProvider for typeName [%s]", complexTypeName);
+    }
+    return provider.get(dimensionName);
+  }
+
+  /**
+   * The {@link DimensionSchema} to use when storing a complex column of the given type, for callers that have a
+   * declared type rather than an existing column. Handlers are free to consult the {@link ColumnCapabilities} they
+   * are given, so a default set describing the type is supplied on the caller's behalf.
+   *
+   * @throws ISE if no handler is registered for the type, which usually means the extension defining it is not loaded
+   */
+  public static DimensionSchema getComplexDimensionSchema(String dimensionName, ColumnType type)
+  {
+    if (!type.is(ValueType.COMPLEX) || type.getComplexTypeName() == null) {
+      throw new IAE("Type [%s] is not a named complex type", type);
+    }
+    return getHandlerForComplexType(dimensionName, type.getComplexTypeName())
+        .getDimensionSchema(ColumnCapabilitiesImpl.createDefault().setType(type));
+  }
+
   public static List<ColumnType> getValueTypesFromDimensionSpecs(List<DimensionSpec> dimSpecs)
   {
     List<ColumnType> types = new ArrayList<>(dimSpecs.size());

processing/src/test/java/org/apache/druid/segment/DimensionHandlerUtilsTest.java

@@ -26,6 +26,7 @@
 import org.apache.druid.data.input.impl.LongDimensionSchema;
 import org.apache.druid.data.input.impl.NewSpatialDimensionSchema;
 import org.apache.druid.data.input.impl.StringDimensionSchema;
+import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnCapabilitiesImpl;
@@ -349,4 +350,42 @@ public ColumnType getColumnType()
       return ColumnType.ofComplex(TYPE);
     }
   }
+
+  @Test
+  public void testGetComplexDimensionSchema()
+  {
+    Assert.assertEquals(
+        new TestDimensionSchema("x", null, false),
+        DimensionHandlerUtils.getComplexDimensionSchema("x", ColumnType.ofComplex(TYPE))
+    );
+  }
+
+  @Test
+  public void testGetComplexDimensionSchemaUnregisteredType()
+  {
+    Assert.assertThrows(
+        ISE.class,
+        () -> DimensionHandlerUtils.getComplexDimensionSchema("x", ColumnType.ofComplex("noSuchType"))
+    );
+  }
+
+  @Test
+  public void testGetComplexDimensionSchemaRejectsNonComplexType()
+  {
+    Assert.assertThrows(
+        IAE.class,
+        () -> DimensionHandlerUtils.getComplexDimensionSchema("x", ColumnType.STRING)
+    );
+  }
+
+  @Test
+  public void testGetHandlerForComplexType()
+  {
+    Assert.assertNotNull(DimensionHandlerUtils.getHandlerForComplexType("x", TYPE));
+    Assert.assertThrows(
+        ISE.class,
+        () -> DimensionHandlerUtils.getHandlerForComplexType("x", "noSuchType")
+    );
+  }
+
 }

server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java

@@ -26,11 +26,14 @@
 import org.apache.druid.data.input.impl.ClusteredValueGroupsBaseTableProjectionSpec;
 import org.apache.druid.data.input.impl.DimensionSchema;
 import org.apache.druid.error.InvalidInput;
+import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.segment.AutoTypeColumnSchema;
+import org.apache.druid.segment.DimensionHandlerUtils;
 import org.apache.druid.segment.NestedDataColumnSchema;
 import org.apache.druid.segment.VirtualColumns;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.ValueType;
 import org.apache.druid.utils.CollectionUtils;
 
 import javax.annotation.Nullable;
@@ -212,16 +215,54 @@ private DimensionSchema toDimensionSchema(ColumnSpec column, @Nullable Dimension
     if (ColumnType.NESTED_DATA.equals(druidType)) {
       return new NestedDataColumnSchema(column.name(), NestedDataColumnSchema.DEFAULT_FORMAT_VERSION);
     }
-    // Other complex types cannot be ingested into a clustered base table: there is no dimension handler for them,
-    // and clustered base tables have no aggregators to produce them.
+    if (druidType.is(ValueType.COMPLEX)) {
+      return complexDimensionSchema(column.name(), druidType);
+    }
     throw InvalidInput.exception(
-        "column [%s] has unsupported type [%s] for a clustered base table; supported types are primitive, primitive"
-        + " array, and COMPLEX<json> columns",
+        "column [%s] has unsupported type [%s] for a clustered base table",
         column.name(),
         druidType
     );
   }
 
+  /**
+   * Resolve a complex column through its registered {@link org.apache.druid.segment.DimensionHandler}, so that any
+   * complex type which can be stored as a dimension may be declared, including types contributed by extensions. The
+   * handler is looked up by the complex type name, so the schema it returns is specific to the declared type.
+   * <p>
+   * The returned schema is checked against the declared type before it is accepted. A schema selects its own handler
+   * at ingest time (via {@link DimensionSchema#getDimensionHandler()}, which reads
+   * {@link DimensionSchema#getColumnType()}), so a schema of some other type would quietly store the column as that
+   * type instead, contradicting the declared schema that queries are validated and coerced against.
+   */
+  private static DimensionSchema complexDimensionSchema(String name, ColumnType druidType)
+  {
+    final DimensionSchema schema;
+    try {
+      schema = DimensionHandlerUtils.getComplexDimensionSchema(name, druidType);
+    }
+    catch (ISE e) {
+      // No handler is registered for this complex type, which usually means the extension that defines it is not
+      // loaded on whichever service is validating the spec.
+      throw InvalidInput.exception(
+          "column [%s] has type [%s], which cannot be stored as a dimension of a clustered base table; if this type"
+          + " comes from an extension, check that the extension is loaded",
+          name,
+          druidType
+      );
+    }
+    if (!druidType.equals(schema.getColumnType())) {
+      throw InvalidInput.exception(
+          "column [%s] has type [%s], but the dimension handler registered for that type produced a schema of type"
+          + " [%s]; a column cannot be stored as a type other than the one it declares",
+          name,
+          druidType,
+          schema.getColumnType()
+      );
+    }
+    return schema;
+  }
+
   private void validateColumnSchemaCustomization(
       ColumnSpec column,
       DimensionSchema customSchema,

server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java

@@ -30,11 +30,15 @@
 import org.apache.druid.data.input.impl.StringDimensionSchema;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.jackson.DefaultObjectMapper;
+import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.segment.AutoTypeColumnSchema;
 import org.apache.druid.segment.DefaultColumnFormatConfig;
+import org.apache.druid.segment.DimensionHandlerUtils;
+import org.apache.druid.segment.DoubleDimensionHandler;
 import org.apache.druid.segment.NestedDataColumnSchema;
 import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
 import org.apache.druid.testing.InitializedNullHandlingTest;
@@ -447,8 +451,12 @@ public void testCreateSpecRetainsDeclaredArrayAndNestedTypes()
     );
   }
 
+  /**
+   * A complex type with no registered dimension handler cannot be stored, and the message says so rather than
+   * claiming the type is unsupported in general: the handler may simply belong to an extension that is not loaded.
+   */
   @Test
-  public void testCreateSpecUnsupportedComplexTypeFails()
+  public void testCreateSpecComplexTypeWithoutHandlerFails()
   {
     final DatasourceBaseTableMetadata metadata = new ClusteredValueGroupsBaseTableMetadata(
         Collections.singletonList("tenant"),
@@ -461,7 +469,145 @@ public void testCreateSpecUnsupportedComplexTypeFails()
         new ColumnSpec("unique_things", "COMPLEX<hyperUnique>", null)
     );
     final DruidException e = Assert.assertThrows(DruidException.class, () -> metadata.createSpec(columns));
-    Assert.assertTrue(e.getMessage().contains("column [unique_things] has unsupported type [COMPLEX<hyperUnique>]"));
+    Assert.assertTrue(
+        e.getMessage(),
+        e.getMessage().contains("column [unique_things] has type [COMPLEX<hyperUnique>], which cannot be stored")
+    );
+  }
+
+  /**
+   * A complex type that does have a registered handler resolves through it, which is how types contributed by
+   * extensions become declarable.
+   */
+  @Test
+  public void testCreateSpecComplexTypeWithRegisteredHandler()
+  {
+    final String typeName = "clusteredBaseTableTestType";
+    // Only getDimensionSchema is exercised; the handler's storage behavior is irrelevant to building a spec.
+    DimensionHandlerUtils.registerDimensionHandlerProvider(
+        typeName,
+        name -> new DoubleDimensionHandler(name)
+        {
+          @Override
+          public DimensionSchema getDimensionSchema(ColumnCapabilities capabilities)
+          {
+            return new TestComplexDimensionSchema(name, typeName);
+          }
+        }
+    );
+
+    final DatasourceBaseTableMetadata metadata = new ClusteredValueGroupsBaseTableMetadata(
+        Collections.singletonList("tenant"),
+        null,
+        null
+    );
+    final List<ColumnSpec> columns = Arrays.asList(
+        new ColumnSpec("tenant", Columns.SQL_VARCHAR, null),
+        new ColumnSpec(Columns.TIME_COLUMN, null, null),
+        new ColumnSpec("sketch", StringUtils.format("COMPLEX<%s>", typeName), null)
+    );
+
+    final List<DimensionSchema> specColumns = metadata.createSpec(columns).getDimensionsSpec().getDimensions();
+    final DimensionSchema stored = specColumns.get(specColumns.size() - 1);
+    Assert.assertEquals("sketch", stored.getName());
+    Assert.assertEquals(ColumnType.ofComplex(typeName), stored.getColumnType());
+  }
+
+  /**
+   * A handler that hands back a schema of some other type is rejected. The schema, not the declared type, selects the
+   * handler used at ingest time, so accepting it would store the column as that other type and contradict the declared
+   * schema that queries are validated and coerced against.
+   */
+  @Test
+  public void testCreateSpecComplexTypeHandlerSchemaOfOtherTypeFails()
+  {
+    final String typeName = "clusteredBaseTableMismatchedType";
+    DimensionHandlerUtils.registerDimensionHandlerProvider(
+        typeName,
+        name -> new DoubleDimensionHandler(name)
+        {
+          @Override
+          public DimensionSchema getDimensionSchema(ColumnCapabilities capabilities)
+          {
+            return new DoubleDimensionSchema(name);
+          }
+        }
+    );
+
+    final DatasourceBaseTableMetadata metadata = new ClusteredValueGroupsBaseTableMetadata(
+        Collections.singletonList("tenant"),
+        null,
+        null
+    );
+    final List<ColumnSpec> columns = Arrays.asList(
+        new ColumnSpec("tenant", Columns.SQL_VARCHAR, null),
+        new ColumnSpec(Columns.TIME_COLUMN, null, null),
+        new ColumnSpec("sketch", StringUtils.format("COMPLEX<%s>", typeName), null)
+    );
+
+    final DruidException e = Assert.assertThrows(DruidException.class, () -> metadata.createSpec(columns));
+    Assert.assertTrue(
+        e.getMessage(),
+        e.getMessage().contains(
+            StringUtils.format(
+                "column [sketch] has type [COMPLEX<%s>], but the dimension handler registered for that type produced"
+                + " a schema of type [DOUBLE]",
+                typeName
+            )
+        )
+    );
+  }
+
+  /**
+   * Minimal complex {@link DimensionSchema}, the shape an honest handler for a complex type returns: the column type
+   * it reports is the type it was registered for.
+   */
+  private static class TestComplexDimensionSchema extends DimensionSchema
+  {
+    private final String typeName;
+
+    TestComplexDimensionSchema(String name, String typeName)
+    {
+      super(name, null, false);
+      this.typeName = typeName;
+    }
+
+    @Override
+    public String getTypeName()
+    {
+      return typeName;
+    }
+
+    @Override
+    public ColumnType getColumnType()
+    {
+      return ColumnType.ofComplex(typeName);
+    }
+  }
+
+  /**
+   * A column declared COMPLEX<json> is nested by declaration, so it keeps a nested schema rather than the 'auto'
+   * schema its dimension handler would produce, which would infer the type from the ingested values instead.
+   */
+  @Test
+  public void testCreateSpecNestedTypeStaysNested()
+  {
+    final DatasourceBaseTableMetadata metadata = new ClusteredValueGroupsBaseTableMetadata(
+        Collections.singletonList("tenant"),
+        null,
+        null
+    );
+    final List<ColumnSpec> columns = Arrays.asList(
+        new ColumnSpec("tenant", Columns.SQL_VARCHAR, null),
+        new ColumnSpec(Columns.TIME_COLUMN, null, null),
+        new ColumnSpec("payload", ColumnType.NESTED_DATA.asTypeString(), null)
+    );
+
+    final List<DimensionSchema> specColumns = metadata.createSpec(columns).getDimensionsSpec().getDimensions();
+    Assert.assertEquals(
+        new NestedDataColumnSchema("payload", NestedDataColumnSchema.DEFAULT_FORMAT_VERSION),
+        specColumns.get(specColumns.size() - 1)
+    );
   }
 
   @Test