Add CoalesceExpression support

This commit is contained in:
Beyyes 2024-09-25 11:40:31 +08:00 committed by GitHub
parent 905c2e2af0
commit 9dda5d93ac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 131 additions and 2 deletions

View File

@ -389,6 +389,25 @@ public class IoTDBMultiIDsWithAttributesTableIT {
DATABASE_NAME);
}
@Test
public void coalesceTest() {
String[] expectedHeader = new String[] {"time", "level", "coalesce_attr2", "str"};
String[] retArray =
new String[] {
"1970-01-01T00:00:00.040Z,l3,a,apricot,",
"1970-01-01T00:00:00.040Z,l3,CCC,apricot,",
"1970-01-01T00:00:00.020Z,l2,CCC,pineapple,",
"1970-01-01T00:00:00.020Z,l2,zz,pineapple,",
"1970-01-01T00:00:00.000Z,l1,d,coconut,",
"1970-01-01T00:00:00.000Z,l1,c,coconut,",
};
tableResultSetEqualTest(
"select time,level,coalesce(attr2, 'CCC', 'DDD') as coalesce_attr2,str from table0 order by num+1,attr1 limit 6",
expectedHeader,
retArray,
DATABASE_NAME);
}
// ========== SubQuery Test =========
@Test
public void subQueryTest1() {

View File

@ -79,6 +79,7 @@ import org.apache.iotdb.db.queryengine.transformation.dag.column.leaf.IdentityCo
import org.apache.iotdb.db.queryengine.transformation.dag.column.leaf.LeafColumnTransformer;
import org.apache.iotdb.db.queryengine.transformation.dag.column.leaf.NullColumnTransformer;
import org.apache.iotdb.db.queryengine.transformation.dag.column.leaf.TimeColumnTransformer;
import org.apache.iotdb.db.queryengine.transformation.dag.column.multi.CoalesceColumnTransformer;
import org.apache.iotdb.db.queryengine.transformation.dag.column.multi.InBinaryMultiColumnTransformer;
import org.apache.iotdb.db.queryengine.transformation.dag.column.multi.InBooleanMultiColumnTransformer;
import org.apache.iotdb.db.queryengine.transformation.dag.column.multi.InDoubleMultiColumnTransformer;
@ -1283,7 +1284,27 @@ public class ColumnTransformerBuilder
@Override
protected ColumnTransformer visitCoalesceExpression(CoalesceExpression node, Context context) {
throw new UnsupportedOperationException(String.format(UNSUPPORTED_EXPRESSION, node));
if (!context.cache.containsKey(node)) {
if (context.hasSeen.containsKey(node)) {
ColumnTransformer columnTransformer = context.hasSeen.get(node);
IdentityColumnTransformer identity =
new IdentityColumnTransformer(
columnTransformer.getType(),
context.originSize + context.commonTransformerList.size());
columnTransformer.addReferenceCount();
context.commonTransformerList.add(columnTransformer);
context.leafList.add(identity);
context.inputDataTypes.add(getTSDataType(columnTransformer.getType()));
context.cache.put(node, identity);
} else {
List<ColumnTransformer> children =
node.getChildren().stream().map(c -> process(c, context)).collect(Collectors.toList());
context.cache.put(node, new CoalesceColumnTransformer(children.get(0).getType(), children));
}
}
ColumnTransformer res = context.cache.get(node);
res.addReferenceCount();
return res;
}
@Override

View File

@ -20,9 +20,14 @@
package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
import com.google.common.collect.ImmutableList;
import org.apache.tsfile.utils.ReadWriteIOUtils;
import javax.annotation.Nonnull;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
@ -89,4 +94,27 @@ public final class CoalesceExpression extends Expression {
public boolean shallowEquals(Node other) {
return sameClass(this, other);
}
// =============== serialize =================
@Override
public TableExpressionType getExpressionType() {
return TableExpressionType.COALESCE;
}
@Override
protected void serialize(DataOutputStream stream) throws IOException {
ReadWriteIOUtils.write(operands.size(), stream);
for (Expression operand : operands) {
serialize(operand, stream);
}
}
public CoalesceExpression(ByteBuffer byteBuffer) {
super(null);
int size = ReadWriteIOUtils.readInt(byteBuffer);
this.operands = new ArrayList<>(size);
while (size-- > 0) {
operands.add(deserialize(byteBuffer));
}
}
}

View File

@ -146,6 +146,9 @@ public abstract class Expression extends Node {
case 24:
expression = new SymbolReference(byteBuffer);
break;
case 25:
expression = new CoalesceExpression(byteBuffer);
break;
default:
throw new IllegalArgumentException("Invalid expression type: " + type);
}

View File

@ -43,7 +43,8 @@ public enum TableExpressionType {
LONG_LITERAL((short) 21),
NULL_LITERAL((short) 22),
STRING_LITERAL((short) 23),
SYMBOL_REFERENCE((short) 24);
SYMBOL_REFERENCE((short) 24),
COALESCE((short) 25);
TableExpressionType(short type) {
this.type = type;

View File

@ -0,0 +1,57 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.apache.iotdb.db.queryengine.transformation.dag.column.multi;
import org.apache.iotdb.db.queryengine.transformation.dag.column.ColumnTransformer;
import org.apache.tsfile.block.column.Column;
import org.apache.tsfile.block.column.ColumnBuilder;
import org.apache.tsfile.read.common.type.Type;
import java.util.List;
public class CoalesceColumnTransformer extends MultiColumnTransformer {
public CoalesceColumnTransformer(Type returnType, List<ColumnTransformer> columnTransformerList) {
super(returnType, columnTransformerList);
}
@Override
protected void doTransform(
List<Column> childrenColumns, ColumnBuilder builder, int positionCount) {
for (int i = 0; i < positionCount; i++) {
boolean allNull = true;
for (Column column : childrenColumns) {
if (!column.isNull(i)) {
allNull = false;
builder.write(column, i);
}
}
if (allNull) {
builder.appendNull();
}
}
}
@Override
protected void checkType() {
// do nothing, has checked in FE
}
}