Array.java
// 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.doris.nereids.trees.expressions.functions.scalar;
import org.apache.doris.catalog.FunctionSignature;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable;
import org.apache.doris.nereids.trees.expressions.functions.ChildDerivedSignature;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.types.ArrayType;
import org.apache.doris.nereids.types.DataType;
import org.apache.doris.nereids.types.DateTimeV2Type;
import org.apache.doris.nereids.types.DecimalV3Type;
import org.apache.doris.nereids.types.MapType;
import org.apache.doris.nereids.types.StructField;
import org.apache.doris.nereids.types.StructType;
import org.apache.doris.nereids.types.TimeV2Type;
import org.apache.doris.nereids.types.coercion.AnyDataType;
import org.apache.doris.nereids.types.coercion.FollowToAnyDataType;
import org.apache.doris.nereids.types.coercion.FollowToArgumentType;
import org.apache.doris.nereids.util.TypeCoercionUtils;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
/**
* ScalarFunction 'array'. This class is generated by GenerateFunction.
*/
public class Array extends ScalarFunction
implements ExplicitlyCastableSignature, AlwaysNotNullable, ChildDerivedSignature {
public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
FunctionSignature.ret(ArrayType.SYSTEM_DEFAULT).args()
);
/**
* constructor with 0 or more arguments.
*/
public Array(Expression... varArgs) {
super("array", varArgs);
}
public Array(List<Expression> varArgs) {
super("array", varArgs);
}
/** constructor for withChildren and reuse signature */
private Array(ScalarFunctionParams functionParams) {
super(functionParams);
}
@Override
public void checkLegalityBeforeTypeCoercion() {
if (children.isEmpty()) {
return;
}
DataType firstChildType = getArgument(0).getDataType();
if (firstChildType.isJsonType() || firstChildType.isVariantType()) {
throw new AnalysisException("array does not support jsonb/variant type");
}
}
/**
* withChildren.
*/
@Override
public Array withChildren(List<Expression> children) {
return new Array(getFunctionParams(children));
}
@Override
public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
return visitor.visitArray(this, context);
}
@Override
public List<FunctionSignature> getSignatures() {
if (arity() == 0) {
return SIGNATURES;
}
Optional<DataType> commonDataType = TypeCoercionUtils.findWiderCommonTypeByVariable(
children.stream().map(ExpressionTrait::getDataType).collect(Collectors.toList()), true, true);
if (commonDataType.isPresent()) {
return ImmutableList.of(
FunctionSignature.ret(ArrayType.of(commonDataType.get())).varArgs(commonDataType.get()));
}
Map<Boolean, List<DataType>> partitioned = children.stream()
.map(ExpressionTrait::getDataType)
.collect(Collectors.partitioningBy(TypeCoercionUtils::hasCharacterType));
List<DataType> needTypeCoercion = Lists.newArrayList(Sets.newHashSet(partitioned.get(true)));
if (needTypeCoercion.size() > 1 || !partitioned.get(false).isEmpty()) {
needTypeCoercion = needTypeCoercion.stream()
.map(TypeCoercionUtils::replaceCharacterToString)
.collect(Collectors.toList());
}
partitioned = partitioned.get(false).stream()
.collect(Collectors.partitioningBy(TypeCoercionUtils::hasDecimalV2Type));
if (!partitioned.get(true).isEmpty()) {
needTypeCoercion.addAll(partitioned.get(true).stream()
.map(TypeCoercionUtils::replaceDecimalV2WithDefault).collect(Collectors.toList()));
}
partitioned = partitioned.get(false).stream()
.collect(Collectors.partitioningBy(TypeCoercionUtils::hasDecimalV3Type));
if (!partitioned.get(true).isEmpty()) {
needTypeCoercion.addAll(partitioned.get(true).stream()
.map(TypeCoercionUtils::replaceDecimalV3WithWildcard).collect(Collectors.toList()));
}
partitioned = partitioned.get(false).stream()
.collect(Collectors.partitioningBy(TypeCoercionUtils::hasDateTimeV2Type));
if (!partitioned.get(true).isEmpty()) {
needTypeCoercion.addAll(partitioned.get(true).stream()
.map(TypeCoercionUtils::replaceDateTimeV2WithMax).collect(Collectors.toList()));
}
needTypeCoercion.addAll(partitioned.get(false));
return needTypeCoercion.stream()
.map(dataType -> FunctionSignature.ret(ArrayType.of(new FollowToArgumentType(0))).varArgs(dataType))
.collect(ImmutableList.toImmutableList());
}
@Override
public FunctionSignature computeSignature(FunctionSignature signature) {
if (isFullyResolved(signature.returnType)
&& signature.argumentsTypes.stream().allMatch(Array::isFullyResolved)) {
// findWiderCommonTypeByVariable already produced the exact common item type. Running generic precision
// promotion again could use one visible nested leaf to overwrite independent fields inside that type.
return signature;
}
return ExplicitlyCastableSignature.super.computeSignature(signature);
}
@Override
public FunctionSignature deriveSignatureFromChildren(
FunctionSignature resolvedSignature, List<Expression> immediateOriginArguments) {
if (children.isEmpty()) {
if (!resolvedSignature.hasVarArgs && resolvedSignature.argumentsTypes.isEmpty()) {
return resolvedSignature;
}
throw new AnalysisException(
"Cannot safely reuse a non-empty ARRAY signature for an empty ARRAY");
}
if (!resolvedSignature.hasVarArgs || resolvedSignature.argumentsTypes.isEmpty()
|| !(resolvedSignature.returnType instanceof ArrayType)) {
throw new AnalysisException(
"Cannot safely reuse an empty or fixed-arity ARRAY signature for a non-empty ARRAY");
}
List<DataType> currentTypes = children.stream()
.map(ExpressionTrait::getDataType)
.collect(Collectors.toList());
List<DataType> originTypes = new ArrayList<>(currentTypes.size());
for (int i = 0; i < currentTypes.size(); i++) {
DataType resolvedType = resolvedSignature.getArgType(i);
DataType originType = i < immediateOriginArguments.size()
? immediateOriginArguments.get(i).getDataType() : resolvedType;
ChildDerivedSignature.refreshNestedTypeMetadata(
resolvedType, currentTypes.get(i), originType);
originTypes.add(originType);
}
DataType itemType = ChildDerivedSignature.mergeNestedTypeMetadata(
resolvedSignature.getArgType(0), currentTypes, originTypes)
.orElseThrow(() -> new AnalysisException(
"Cannot safely reuse ARRAY signature with incompatible nested argument metadata"));
int refreshedArity = resolvedSignature.argumentsTypes.size() == 1
? 1 : currentTypes.size();
return resolvedSignature.withArgumentTypes(true,
Collections.nCopies(refreshedArity, itemType))
.withReturnType(ArrayType.of(itemType));
}
private static boolean isFullyResolved(DataType dataType) {
if (dataType instanceof AnyDataType
|| dataType instanceof FollowToAnyDataType
|| dataType instanceof FollowToArgumentType
|| DecimalV3Type.WILDCARD.equals(dataType)
|| DateTimeV2Type.WILDCARD.equals(dataType)
|| TimeV2Type.WILDCARD.equals(dataType)) {
return false;
}
if (dataType instanceof ArrayType) {
return isFullyResolved(((ArrayType) dataType).getItemType());
}
if (dataType instanceof MapType) {
return isFullyResolved(((MapType) dataType).getKeyType())
&& isFullyResolved(((MapType) dataType).getValueType());
}
if (dataType instanceof StructType) {
for (StructField field : ((StructType) dataType).getFields()) {
if (!isFullyResolved(field.getDataType())) {
return false;
}
}
}
return true;
}
}