Support more data types for IN filter pushdown (oap-project#139)

Support tiny int, small int and date for IN filter pushdown.
zhejiangxiaomai · Mar 29, 2023 · a25daf5 · a25daf5
1 parent 48788db
commit a25daf5
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 0 deletions.
diff --git a/velox/substrait/SubstraitToVeloxPlan.cpp b/velox/substrait/SubstraitToVeloxPlan.cpp
@@ -1851,6 +1851,60 @@ void SubstraitVeloxPlanConverter::setInFilter<TypeKind::INTEGER>(
       common::createBigintValues(values, nullAllowed);
 }
 
+template <>
+void SubstraitVeloxPlanConverter::setInFilter<TypeKind::SMALLINT>(
+    const std::vector<variant>& variants,
+    bool nullAllowed,
+    const std::string& inputName,
+    connector::hive::SubfieldFilters& filters) {
+  // Use bigint values for small int type.
+  std::vector<int64_t> values;
+  values.reserve(variants.size());
+  for (const auto& variant : variants) {
+    // Use the matched type to get value from variant.
+    int64_t value = variant.value<int16_t>();
+    values.emplace_back(value);
+  }
+  filters[common::Subfield(inputName)] =
+      common::createBigintValues(values, nullAllowed);
+}
+
+template <>
+void SubstraitVeloxPlanConverter::setInFilter<TypeKind::TINYINT>(
+    const std::vector<variant>& variants,
+    bool nullAllowed,
+    const std::string& inputName,
+    connector::hive::SubfieldFilters& filters) {
+  // Use bigint values for tiny int type.
+  std::vector<int64_t> values;
+  values.reserve(variants.size());
+  for (const auto& variant : variants) {
+    // Use the matched type to get value from variant.
+    int64_t value = variant.value<int8_t>();
+    values.emplace_back(value);
+  }
+  filters[common::Subfield(inputName)] =
+      common::createBigintValues(values, nullAllowed);
+}
+
+template <>
+void SubstraitVeloxPlanConverter::setInFilter<TypeKind::DATE>(
+    const std::vector<variant>& variants,
+    bool nullAllowed,
+    const std::string& inputName,
+    connector::hive::SubfieldFilters& filters) {
+  // Use bigint values for int type.
+  std::vector<int64_t> values;
+  values.reserve(variants.size());
+  for (const auto& variant : variants) {
+    // Use int32 to get value from date variant.
+    int64_t value = variant.value<int32_t>();
+    values.emplace_back(value);
+  }
+  filters[common::Subfield(inputName)] =
+      common::createBigintValues(values, nullAllowed);
+}
+
 template <>
 void SubstraitVeloxPlanConverter::setInFilter<TypeKind::VARCHAR>(
     const std::vector<variant>& variants,
@@ -2017,6 +2071,14 @@ connector::hive::SubfieldFilters SubstraitVeloxPlanConverter::mapToFilters(
   for (uint32_t colIdx = 0; colIdx < inputNameList.size(); colIdx++) {
     auto inputType = inputTypeList[colIdx];
     switch (inputType->kind()) {
+      case TypeKind::TINYINT:
+        constructSubfieldFilters<TypeKind::TINYINT, common::BigintRange>(
+            colIdx, inputNameList[colIdx], colInfoMap[colIdx], filters);
+        break;
+      case TypeKind::SMALLINT:
+        constructSubfieldFilters<TypeKind::SMALLINT, common::BigintRange>(
+            colIdx, inputNameList[colIdx], colInfoMap[colIdx], filters);
+        break;
       case TypeKind::INTEGER:
         constructSubfieldFilters<TypeKind::INTEGER, common::BigintRange>(
             colIdx, inputNameList[colIdx], colInfoMap[colIdx], filters);

diff --git a/velox/substrait/TypeUtils.h b/velox/substrait/TypeUtils.h
@@ -32,6 +32,20 @@ std::string_view getNameBeforeDelimiter(
 template <TypeKind KIND>
 struct RangeTraits {};
 
+template <>
+struct RangeTraits<TypeKind::TINYINT> {
+  using RangeType = common::BigintRange;
+  using MultiRangeType = common::BigintMultiRange;
+  using NativeType = int8_t;
+};
+
+template <>
+struct RangeTraits<TypeKind::SMALLINT> {
+  using RangeType = common::BigintRange;
+  using MultiRangeType = common::BigintMultiRange;
+  using NativeType = int16_t;
+};
+
 template <>
 struct RangeTraits<TypeKind::INTEGER> {
   using RangeType = common::BigintRange;