Skip to content

Commit

Permalink
Cast the values read in the EXTERN function to the type in the row si…
Browse files Browse the repository at this point in the history 
…gnature (#15183)
  • Loading branch information
@LakshSingla
LakshSingla authored Oct 19, 2023
1 parent 7802078 commit fa311dd
Show file tree
Hide file tree
Showing 4 changed files with 124 additions and 2 deletions.
10 changes: 10 additions & 0 deletions .../multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public class ExternalSegment extends RowBasedSegment<InputRow>
{

private final InputSource inputSource;
private final RowSignature signature;
public static final String SEGMENT_ID = "__external";

/**
Expand Down Expand Up @@ -145,6 +146,7 @@ public void cleanup(CloseableIterator<InputRow> iterFromMake)
signature
);
this.inputSource = inputSource;
this.signature = signature;
}

/**
Expand All @@ -154,4 +156,12 @@ public InputSource externalInputSource()
{
return inputSource;
}

/**
* Returns the signature of the external input source
*/
public RowSignature signature()
{
return signature;
}
}
20 changes: 18 additions & 2 deletions ...query/src/main/java/org/apache/druid/msq/querykit/scan/ExternalColumnSelectorFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
}

Expand All @@ -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()
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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);
Expand Down
1 change: 1 addition & 0 deletions ...stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -392,6 +392,7 @@ private ColumnSelectorFactory wrapColumnSelectorFactoryIfNeeded(final ColumnSele
return new ExternalColumnSelectorFactory(
baseColumnSelectorFactory,
((ExternalSegment) segment).externalInputSource(),
((ExternalSegment) segment).signature(),
cursorOffset
);
}
Expand Down
95 changes: 95 additions & 0 deletions extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -857,6 +858,100 @@ public void testScanWithOrderByOnDoubleArray()
.verifyResults();
}

@Test
public void testScanExternBooleanArray()
{
final List<Object[]> 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<LONG>\"}]'\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<Object[]> 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<LONG>\"}]'\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<Object[]> expectedMultiValueFooRowsToArray()
{
List<Object[]> expectedRows = new ArrayList<>();
Expand Down

0 comments on commit fa311dd

Please sign in to comment.