* 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.
*/
#include "substrait/VeloxSubstraitSignature.h"
#include "velox/common/base/tests/GTestUtils.h"
#include "velox/functions/prestosql/registration/RegistrationFunctions.h"
using namespace facebook::velox;
namespace gluten {
class VeloxSubstraitSignatureTest : public ::testing::Test {
protected:
void SetUp() override {
Test::SetUp();
functions::prestosql::registerAllScalarFunctions();
}
static TypePtr fromSubstraitSignature(const std::string& signature) {
return VeloxSubstraitSignature::fromSubstraitSignature(signature);
}
static std::string toSubstraitSignature(const TypePtr& type) {
return VeloxSubstraitSignature::toSubstraitSignature(type);
}
static std::string toSubstraitSignature(const std::string& functionName, const std::vector<TypePtr>& arguments) {
return VeloxSubstraitSignature::toSubstraitSignature(functionName, arguments);
}
};
TEST_F(VeloxSubstraitSignatureTest, toSubstraitSignatureWithType) {
ASSERT_EQ(toSubstraitSignature(BOOLEAN()), "bool");
ASSERT_EQ(toSubstraitSignature(TINYINT()), "i8");
ASSERT_EQ(toSubstraitSignature(SMALLINT()), "i16");
ASSERT_EQ(toSubstraitSignature(INTEGER()), "i32");
ASSERT_EQ(toSubstraitSignature(BIGINT()), "i64");
ASSERT_EQ(toSubstraitSignature(REAL()), "fp32");
ASSERT_EQ(toSubstraitSignature(DOUBLE()), "fp64");
ASSERT_EQ(toSubstraitSignature(VARCHAR()), "str");
ASSERT_EQ(toSubstraitSignature(VARBINARY()), "vbin");
ASSERT_EQ(toSubstraitSignature(TIMESTAMP()), "ts");
ASSERT_EQ(toSubstraitSignature(DATE()), "date");
ASSERT_EQ(toSubstraitSignature(ARRAY(BOOLEAN())), "list");
ASSERT_EQ(toSubstraitSignature(ARRAY(INTEGER())), "list");
ASSERT_EQ(toSubstraitSignature(MAP(INTEGER(), BIGINT())), "map");
ASSERT_EQ(toSubstraitSignature(ROW({INTEGER(), BIGINT()})), "struct<i32,i64>");
ASSERT_EQ(toSubstraitSignature(ROW({ARRAY(INTEGER())})), "struct<list>");
ASSERT_EQ(toSubstraitSignature(ROW({MAP(INTEGER(), INTEGER())})), "struct<map>");
ASSERT_EQ(toSubstraitSignature(ROW({ROW({INTEGER()})})), "struct<struct<i32>>");
ASSERT_EQ(toSubstraitSignature(UNKNOWN()), "u!name");
}
TEST_F(VeloxSubstraitSignatureTest, toSubstraitSignatureWithFunctionNameAndArguments) {
ASSERT_EQ(toSubstraitSignature("eq", {INTEGER(), INTEGER()}), "eq:i32_i32");
ASSERT_EQ(toSubstraitSignature("gt", {INTEGER(), INTEGER()}), "gt:i32_i32");
ASSERT_EQ(toSubstraitSignature("lt", {INTEGER(), INTEGER()}), "lt:i32_i32");
ASSERT_EQ(toSubstraitSignature("gte", {INTEGER(), INTEGER()}), "gte:i32_i32");
ASSERT_EQ(toSubstraitSignature("lte", {INTEGER(), INTEGER()}), "lte:i32_i32");
ASSERT_EQ(toSubstraitSignature("and", {BOOLEAN(), BOOLEAN()}), "and:bool_bool");
ASSERT_EQ(toSubstraitSignature("or", {BOOLEAN(), BOOLEAN()}), "or:bool_bool");
ASSERT_EQ(toSubstraitSignature("not", {BOOLEAN()}), "not:bool");
ASSERT_EQ(toSubstraitSignature("xor", {BOOLEAN(), BOOLEAN()}), "xor:bool_bool");
ASSERT_EQ(toSubstraitSignature("between", {INTEGER(), INTEGER(), INTEGER()}), "between:i32_i32_i32");
ASSERT_EQ(toSubstraitSignature("plus", {INTEGER(), INTEGER()}), "plus:i32_i32");
ASSERT_EQ(toSubstraitSignature("divide", {INTEGER(), INTEGER()}), "divide:i32_i32");
ASSERT_EQ(toSubstraitSignature("cardinality", {ARRAY(INTEGER())}), "cardinality:list");
ASSERT_EQ(toSubstraitSignature("array_sum", {ARRAY(INTEGER())}), "array_sum:list");
ASSERT_EQ(toSubstraitSignature("sum", {INTEGER()}), "sum:i32");
ASSERT_EQ(toSubstraitSignature("avg", {INTEGER()}), "avg:i32");
ASSERT_EQ(toSubstraitSignature("count", {INTEGER()}), "count:i32");
auto functionType = std::make_shared<const FunctionType>(std::vector<TypePtr>{INTEGER(), VARCHAR()}, BIGINT());
std::vector<TypePtr> types = {MAP(INTEGER(), VARCHAR()), functionType};
ASSERT_ANY_THROW(toSubstraitSignature("transform_keys", std::move(types)));
}
TEST_F(VeloxSubstraitSignatureTest, fromSubstraitSignature) {
ASSERT_EQ(fromSubstraitSignature("bool")->kind(), TypeKind::BOOLEAN);
ASSERT_EQ(fromSubstraitSignature("i8")->kind(), TypeKind::TINYINT);
ASSERT_EQ(fromSubstraitSignature("i16")->kind(), TypeKind::SMALLINT);
ASSERT_EQ(fromSubstraitSignature("i32")->kind(), TypeKind::INTEGER);
ASSERT_EQ(fromSubstraitSignature("i64")->kind(), TypeKind::BIGINT);
ASSERT_EQ(fromSubstraitSignature("fp32")->kind(), TypeKind::REAL);
ASSERT_EQ(fromSubstraitSignature("fp64")->kind(), TypeKind::DOUBLE);
ASSERT_EQ(fromSubstraitSignature("str")->kind(), TypeKind::VARCHAR);
ASSERT_EQ(fromSubstraitSignature("vbin")->kind(), TypeKind::VARBINARY);
ASSERT_EQ(fromSubstraitSignature("ts")->kind(), TypeKind::TIMESTAMP);
ASSERT_EQ(fromSubstraitSignature("date")->kind(), TypeKind::INTEGER);
ASSERT_EQ(fromSubstraitSignature("dec<18,2>")->kind(), TypeKind::BIGINT);
ASSERT_EQ(fromSubstraitSignature("dec<19,2>")->kind(), TypeKind::HUGEINT);
auto type = fromSubstraitSignature("struct<bool,vbin,dec<8,2>>");
ASSERT_EQ(type->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(0)->kind(), TypeKind::BOOLEAN);
ASSERT_EQ(type->childAt(1)->kind(), TypeKind::VARBINARY);
ASSERT_EQ(type->childAt(2)->kind(), TypeKind::BIGINT);
type = fromSubstraitSignature("struct<bool,struct<i8,fp32>>");
ASSERT_EQ(type->childAt(1)->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(1)->childAt(0)->kind(), TypeKind::TINYINT);
ASSERT_EQ(type->childAt(1)->childAt(1)->kind(), TypeKind::REAL);
type = fromSubstraitSignature("struct<bool,struct<i8,fp32>,vbin,ts,dec<9,2>>");
ASSERT_EQ(type->childAt(1)->kind(), TypeKind::ROW);
type = fromSubstraitSignature("struct<struct<ts>,i16>");
ASSERT_EQ(type->childAt(0)->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(0)->childAt(0)->kind(), TypeKind::TIMESTAMP);
type = fromSubstraitSignature("struct<struct<dec<12,9>>>");
ASSERT_EQ(type->childAt(0)->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(0)->childAt(0)->kind(), TypeKind::BIGINT);
type = fromSubstraitSignature("struct<struct<struct<i8>>>");
ASSERT_EQ(type->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(0)->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(0)->childAt(0)->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(0)->childAt(0)->childAt(0)->kind(), TypeKind::TINYINT);
type = fromSubstraitSignature("struct<struct<struct<i8,str>>>");
ASSERT_EQ(type->childAt(0)->childAt(0)->childAt(0)->kind(), TypeKind::TINYINT);
ASSERT_EQ(type->childAt(0)->childAt(0)->childAt(1)->kind(), TypeKind::VARCHAR);
type = fromSubstraitSignature("struct<struct<struct<i8,dec<19,2>>>>");
ASSERT_EQ(type->childAt(0)->childAt(0)->childAt(1)->kind(), TypeKind::HUGEINT);
type = fromSubstraitSignature("struct<fp32,struct<struct<i8,dec<19,2>,list<i32>,map<i16,i32>>>>");
ASSERT_EQ(type->childAt(0)->kind(), TypeKind::REAL);
ASSERT_EQ(type->childAt(1)->childAt(0)->childAt(0)->kind(), TypeKind::TINYINT);
ASSERT_EQ(type->childAt(1)->childAt(0)->childAt(1)->kind(), TypeKind::HUGEINT);
ASSERT_EQ(type->childAt(1)->childAt(0)->childAt(2)->childAt(0)->kind(), TypeKind::INTEGER);
ASSERT_EQ(type->childAt(1)->childAt(0)->childAt(3)->childAt(0)->kind(), TypeKind::SMALLINT);
ASSERT_EQ(type->childAt(1)->childAt(0)->childAt(3)->childAt(1)->kind(), TypeKind::INTEGER);
ASSERT_ANY_THROW(fromSubstraitSignature("other")->kind());
type = fromSubstraitSignature("map<bool,list<map<str,i32>>>");
ASSERT_EQ(type->kind(), TypeKind::MAP);
ASSERT_EQ(type->childAt(0)->kind(), TypeKind::BOOLEAN);
ASSERT_EQ(type->childAt(1)->kind(), TypeKind::ARRAY);
ASSERT_EQ(type->childAt(1)->childAt(0)->kind(), TypeKind::MAP);
type = fromSubstraitSignature("struct<map<bool,i8>,list<map<str,i32>>>");
ASSERT_EQ(type->kind(), TypeKind::ROW);
ASSERT_EQ(type->childAt(0)->kind(), TypeKind::MAP);
ASSERT_EQ(type->childAt(0)->childAt(0)->kind(), TypeKind::BOOLEAN);
ASSERT_EQ(type->childAt(0)->childAt(1)->kind(), TypeKind::TINYINT);
ASSERT_EQ(type->childAt(1)->kind(), TypeKind::ARRAY);
ASSERT_EQ(type->childAt(1)->childAt(0)->kind(), TypeKind::MAP);
ASSERT_EQ(type->childAt(1)->childAt(0)->childAt(0)->kind(), TypeKind::VARCHAR);
ASSERT_EQ(type->childAt(1)->childAt(0)->childAt(1)->kind(), TypeKind::INTEGER);
}
}
