From 66329dbd2ab42583b7b113b31a63747d56af99fd Mon Sep 17 00:00:00 2001
From: PHILO-HE <feilong.he@intel.com>
Date: Tue, 7 Mar 2023 16:59:54 +0800
Subject: [PATCH] Support more data types for IN filter pushdown (#139)

Support tiny int, small int and date for IN filter pushdown.
---
 velox/substrait/SubstraitToVeloxPlan.cpp | 62 ++++++++++++++++++++++++
 velox/substrait/TypeUtils.h              | 14 ++++++
 2 files changed, 76 insertions(+)

diff --git a/velox/substrait/SubstraitToVeloxPlan.cpp b/velox/substrait/SubstraitToVeloxPlan.cpp
index 2b389afa71bf..7737a85b8db5 100644
--- 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
index d5f9e7303fbc..de78c84e3a3b 100644
--- 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;