diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java index f53c2cf7c540..93f24cbdff6d 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java @@ -45,6 +45,7 @@ public class ExternalSegment extends RowBasedSegment { private final InputSource inputSource; + private final RowSignature signature; public static final String SEGMENT_ID = "__external"; /** @@ -145,6 +146,7 @@ public void cleanup(CloseableIterator iterFromMake) signature ); this.inputSource = inputSource; + this.signature = signature; } /** @@ -154,4 +156,12 @@ public InputSource externalInputSource() { return inputSource; } + + /** + * Returns the signature of the external input source + */ + public RowSignature signature() + { + return signature; + } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ExternalColumnSelectorFactory.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ExternalColumnSelectorFactory.java index fbe82c240dbf..fc9f59ad32c6 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ExternalColumnSelectorFactory.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ExternalColumnSelectorFactory.java @@ -21,6 +21,8 @@ import org.apache.druid.data.input.InputSource; import org.apache.druid.java.util.common.parsers.ParseException; +import org.apache.druid.math.expr.ExprEval; +import org.apache.druid.math.expr.ExpressionType; import org.apache.druid.query.dimension.DimensionSpec; import org.apache.druid.query.filter.DruidPredicateFactory; import org.apache.druid.query.filter.ValueMatcher; @@ -32,6 +34,7 @@ import org.apache.druid.segment.RowIdSupplier; import org.apache.druid.segment.SimpleSettableOffset; import org.apache.druid.segment.column.ColumnCapabilities; +import org.apache.druid.segment.column.RowSignature; import org.apache.druid.segment.data.IndexedInts; import javax.annotation.Nullable; @@ -48,16 +51,19 @@ public class ExternalColumnSelectorFactory implements ColumnSelectorFactory private final ColumnSelectorFactory delegate; private final InputSource inputSource; + private final RowSignature rowSignature; private final SimpleSettableOffset offset; public ExternalColumnSelectorFactory( final ColumnSelectorFactory delgate, final InputSource inputSource, + final RowSignature rowSignature, final SimpleSettableOffset offset ) { this.delegate = delgate; this.inputSource = inputSource; + this.rowSignature = rowSignature; this.offset = offset; } @@ -67,6 +73,7 @@ public DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec) return new DimensionSelector() { final DimensionSelector delegateDimensionSelector = delegate.makeDimensionSelector(dimensionSpec); + final ExpressionType expressionType = ExpressionType.fromColumnType(dimensionSpec.getOutputType()); @Override public IndexedInts getRow() @@ -97,7 +104,10 @@ public void inspectRuntimeShape(RuntimeShapeInspector inspector) public Object getObject() { try { - return delegateDimensionSelector.getObject(); + if (expressionType == null) { + return delegateDimensionSelector.getObject(); + } + return ExprEval.ofType(expressionType, delegateDimensionSelector.getObject()).value(); } catch (Exception e) { throw createException(e, dimensionSpec.getDimension(), inputSource, offset); @@ -144,6 +154,9 @@ public ColumnValueSelector makeColumnValueSelector(String columnName) return new ColumnValueSelector() { final ColumnValueSelector delegateColumnValueSelector = delegate.makeColumnValueSelector(columnName); + final ExpressionType expressionType = ExpressionType.fromColumnType( + rowSignature.getColumnType(columnName).orElse(null) + ); @Override public double getDouble() @@ -195,7 +208,10 @@ public boolean isNull() public Object getObject() { try { - return delegateColumnValueSelector.getObject(); + if (expressionType == null) { + return delegateColumnValueSelector.getObject(); + } + return ExprEval.ofType(expressionType, delegateColumnValueSelector.getObject()).value(); } catch (Exception e) { throw createException(e, columnName, inputSource, offset); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java index 1541d314f215..278a9c251dea 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java @@ -392,6 +392,7 @@ private ColumnSelectorFactory wrapColumnSelectorFactoryIfNeeded(final ColumnSele return new ExternalColumnSelectorFactory( baseColumnSelectorFactory, ((ExternalSegment) segment).externalInputSource(), + ((ExternalSegment) segment).signature(), cursorOffset ); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java index 2b152cfbe1c4..282dacc115c2 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java @@ -21,6 +21,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import org.apache.druid.data.input.impl.InlineInputSource; import org.apache.druid.data.input.impl.JsonInputFormat; import org.apache.druid.data.input.impl.LocalInputSource; import org.apache.druid.java.util.common.ISE; @@ -857,6 +858,100 @@ public void testScanWithOrderByOnDoubleArray() .verifyResults(); } + @Test + public void testScanExternBooleanArray() + { + final List expectedRows = Collections.singletonList( + new Object[]{Arrays.asList(1L, 0L, null)} + ); + + RowSignature scanSignature = RowSignature.builder() + .add("a_bool", ColumnType.LONG_ARRAY) + .build(); + + Query expectedQuery = newScanQueryBuilder() + .dataSource( + new ExternalDataSource( + new InlineInputSource("{\"a_bool\":[true,false,null]}"), + new JsonInputFormat(null, null, null, null, null), + scanSignature + ) + ) + .intervals(querySegmentSpec(Filtration.eternity())) + .columns("a_bool") + .context(defaultScanQueryContext(context, scanSignature)) + .build(); + + testSelectQuery().setSql("SELECT a_bool FROM TABLE(\n" + + " EXTERN(\n" + + " '{\"type\": \"inline\", \"data\":\"{\\\"a_bool\\\":[true,false,null]}\"}',\n" + + " '{\"type\": \"json\"}',\n" + + " '[{\"name\": \"a_bool\", \"type\": \"ARRAY\"}]'\n" + + " )\n" + + ")") + .setQueryContext(context) + .setExpectedMSQSpec(MSQSpec + .builder() + .query(expectedQuery) + .columnMappings(new ColumnMappings(ImmutableList.of( + new ColumnMapping("a_bool", "a_bool") + ))) + .tuningConfig(MSQTuningConfig.defaultConfig()) + .destination(TaskReportMSQDestination.INSTANCE) + .build() + ) + .setExpectedRowSignature(scanSignature) + .setExpectedResultRows(expectedRows) + .verifyResults(); + } + + @Test + public void testScanExternArrayWithNonConvertibleType() + { + final List expectedRows = Collections.singletonList( + new Object[]{Arrays.asList(null, null)} + ); + + RowSignature scanSignature = RowSignature.builder() + .add("a_bool", ColumnType.LONG_ARRAY) + .build(); + + Query expectedQuery = newScanQueryBuilder() + .dataSource( + new ExternalDataSource( + new InlineInputSource("{\"a_bool\":[\"Test\",\"Test2\"]}"), + new JsonInputFormat(null, null, null, null, null), + scanSignature + ) + ) + .intervals(querySegmentSpec(Filtration.eternity())) + .columns("a_bool") + .context(defaultScanQueryContext(context, scanSignature)) + .build(); + + testSelectQuery().setSql("SELECT a_bool FROM TABLE(\n" + + " EXTERN(\n" + + " '{\"type\": \"inline\", \"data\":\"{\\\"a_bool\\\":[\\\"Test\\\",\\\"Test2\\\"]}\"}',\n" + + " '{\"type\": \"json\"}',\n" + + " '[{\"name\": \"a_bool\", \"type\": \"ARRAY\"}]'\n" + + " )\n" + + ")") + .setQueryContext(context) + .setExpectedMSQSpec(MSQSpec + .builder() + .query(expectedQuery) + .columnMappings(new ColumnMappings(ImmutableList.of( + new ColumnMapping("a_bool", "a_bool") + ))) + .tuningConfig(MSQTuningConfig.defaultConfig()) + .destination(TaskReportMSQDestination.INSTANCE) + .build() + ) + .setExpectedRowSignature(scanSignature) + .setExpectedResultRows(expectedRows) + .verifyResults(); + } + private List expectedMultiValueFooRowsToArray() { List expectedRows = new ArrayList<>();