From 53e8c51d22cef1fc3b51ac1c8ea5a4215677086f Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 30 Jul 2026 18:13:45 -0700 Subject: [PATCH 1/2] fix: clustered segment catalog spec complex handling --- .../druid/msq/util/DimensionSchemaUtils.java | 10 +-- .../druid/segment/DimensionHandlerUtils.java | 38 +++++++-- .../segment/DimensionHandlerUtilsTest.java | 39 ++++++++++ ...ClusteredValueGroupsBaseTableMetadata.java | 33 +++++++- ...teredValueGroupsBaseTableMetadataTest.java | 78 ++++++++++++++++++- 5 files changed, 179 insertions(+), 19 deletions(-) diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java b/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java index 424700c30f01..4b7e6e7eb2ba 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java +++ b/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, 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); } } } diff --git a/processing/src/main/java/org/apache/druid/segment/DimensionHandlerUtils.java b/processing/src/main/java/org/apache/druid/segment/DimensionHandlerUtils.java index e129ceb41778..e8ee7d83634f 100644 --- a/processing/src/main/java/org/apache/druid/segment/DimensionHandlerUtils.java +++ b/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 getValueTypesFromDimensionSpecs(List dimSpecs) { List types = new ArrayList<>(dimSpecs.size()); diff --git a/processing/src/test/java/org/apache/druid/segment/DimensionHandlerUtilsTest.java b/processing/src/test/java/org/apache/druid/segment/DimensionHandlerUtilsTest.java index 1e7d5c371016..1d2087fe64c5 100644 --- a/processing/src/test/java/org/apache/druid/segment/DimensionHandlerUtilsTest.java +++ b/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") + ); + } + } diff --git a/server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java b/server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java index feaf647e402c..b9a38b4db5a6 100644 --- a/server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java +++ b/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,38 @@ 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 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. + */ + private static DimensionSchema complexDimensionSchema(String name, ColumnType druidType) + { + try { + return 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 + ); + } + } + private void validateColumnSchemaCustomization( ColumnSpec column, DimensionSchema customSchema, diff --git a/server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java b/server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java index a31a8e4dad8b..0e7d922ec88c 100644 --- a/server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java +++ b/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,73 @@ public void testCreateSpecUnsupportedComplexTypeFails() new ColumnSpec("unique_things", "COMPLEX", null) ); final DruidException e = Assert.assertThrows(DruidException.class, () -> metadata.createSpec(columns)); - Assert.assertTrue(e.getMessage().contains("column [unique_things] has unsupported type [COMPLEX]")); + Assert.assertTrue( + e.getMessage(), + e.getMessage().contains("column [unique_things] has type [COMPLEX], 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"; + 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 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 specColumns = metadata.createSpec(columns).getDimensionsSpec().getDimensions(); + Assert.assertEquals( + new DoubleDimensionSchema("sketch"), + specColumns.get(specColumns.size() - 1) + ); + } + + /** + * A column declared COMPLEX 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 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 specColumns = metadata.createSpec(columns).getDimensionsSpec().getDimensions(); + Assert.assertEquals( + new NestedDataColumnSchema("payload", NestedDataColumnSchema.DEFAULT_FORMAT_VERSION), + specColumns.get(specColumns.size() - 1) + ); } @Test From 47a3e0502bcb7ba34088c08e0c753a901a0aa1f9 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Fri, 31 Jul 2026 11:59:58 -0700 Subject: [PATCH 2/2] stronger validation for complex types from provider registry --- ...ClusteredValueGroupsBaseTableMetadata.java | 18 ++++- ...teredValueGroupsBaseTableMetadataTest.java | 80 ++++++++++++++++++- 2 files changed, 93 insertions(+), 5 deletions(-) diff --git a/server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java b/server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java index b9a38b4db5a6..cf0f3367997f 100644 --- a/server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java +++ b/server/src/main/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadata.java @@ -229,11 +229,17 @@ private DimensionSchema toDimensionSchema(ColumnSpec column, @Nullable Dimension * 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. + *

+ * 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 { - return DimensionHandlerUtils.getComplexDimensionSchema(name, druidType); + 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 @@ -245,6 +251,16 @@ private static DimensionSchema complexDimensionSchema(String name, ColumnType dr 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( diff --git a/server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java b/server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java index 0e7d922ec88c..8690ea5b415c 100644 --- a/server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java +++ b/server/src/test/java/org/apache/druid/catalog/model/ClusteredValueGroupsBaseTableMetadataTest.java @@ -483,6 +483,7 @@ public void testCreateSpecComplexTypeWithoutHandlerFails() 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) @@ -490,7 +491,7 @@ public void testCreateSpecComplexTypeWithRegisteredHandler() @Override public DimensionSchema getDimensionSchema(ColumnCapabilities capabilities) { - return new DoubleDimensionSchema(name); + return new TestComplexDimensionSchema(name, typeName); } } ); @@ -507,12 +508,83 @@ public DimensionSchema getDimensionSchema(ColumnCapabilities capabilities) ); final List specColumns = metadata.createSpec(columns).getDimensionsSpec().getDimensions(); - Assert.assertEquals( - new DoubleDimensionSchema("sketch"), - specColumns.get(specColumns.size() - 1) + 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 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 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.