Support bool type for range filter (facebookincubator#202)

rui-mo · Apr 25, 2023 · f4bb7f0 · f4bb7f0
1 parent 011716b
commit f4bb7f0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 11 deletions.
diff --git a/velox/substrait/SubstraitToVeloxPlan.cpp b/velox/substrait/SubstraitToVeloxPlan.cpp
@@ -506,23 +506,22 @@ core::PlanNodePtr SubstraitVeloxPlanConverter::toVeloxPlan(
 
   std::vector<std::vector<core::TypedExprPtr>> projectSetExprs;
   projectSetExprs.reserve(expandRel.fields_size());
-  
+
   for (const auto& projections : expandRel.fields()) {
     std::vector<core::TypedExprPtr> projectExprs;
     projectExprs.reserve(projections.switching_field().duplicates_size());
 
-    for (
-      const auto& projectExpr : projections.switching_field().duplicates()) {
+    for (const auto& projectExpr : projections.switching_field().duplicates()) {
       if (projectExpr.has_selection()) {
         auto expression =
-          exprConverter_->toVeloxExpr(projectExpr.selection(), inputType);
+            exprConverter_->toVeloxExpr(projectExpr.selection(), inputType);
         projectExprs.emplace_back(expression);
       } else if (projectExpr.has_literal()) {
-        auto expression =
-          exprConverter_->toVeloxExpr(projectExpr.literal());
+        auto expression = exprConverter_->toVeloxExpr(projectExpr.literal());
         projectExprs.emplace_back(expression);
       } else {
-        VELOX_FAIL("The project in Expand Operator only support field or literal.");
+        VELOX_FAIL(
+            "The project in Expand Operator only support field or literal.");
       }
     }
     projectSetExprs.emplace_back(projectExprs);
@@ -536,10 +535,7 @@ core::PlanNodePtr SubstraitVeloxPlanConverter::toVeloxPlan(
   }
 
   return std::make_shared<core::ExpandNode>(
-      nextPlanNodeId(),
-      projectSetExprs,
-      std::move(names),
-      childNode);
+      nextPlanNodeId(), projectSetExprs, std::move(names), childNode);
 }
 
 const core::WindowNode::Frame createWindowFrame(
@@ -1722,6 +1718,12 @@ void SubstraitVeloxPlanConverter::setFilterMap(
       }
       setColInfoMap<double>(functionName, colIdxVal, val, reverse, colInfoMap);
       break;
+    case TypeKind::BOOLEAN:
+      if (substraitLit) {
+        val = variant(substraitLit.value().boolean());
+      }
+      setColInfoMap<bool>(functionName, colIdxVal, val, reverse, colInfoMap);
+      break;
     case TypeKind::VARCHAR:
       if (substraitLit) {
         val = variant(substraitLit.value().string());
@@ -2131,6 +2133,10 @@ connector::hive::SubfieldFilters SubstraitVeloxPlanConverter::mapToFilters(
         constructSubfieldFilters<TypeKind::DOUBLE, common::Filter>(
             colIdx, inputNameList[colIdx], colInfoMap[colIdx], filters);
         break;
+      case TypeKind::BOOLEAN:
+        constructSubfieldFilters<TypeKind::BOOLEAN, common::BigintRange>(
+            colIdx, inputNameList[colIdx], colInfoMap[colIdx], filters);
+        break;
       case TypeKind::VARCHAR:
         constructSubfieldFilters<TypeKind::VARCHAR, common::Filter>(
             colIdx, inputNameList[colIdx], colInfoMap[colIdx], filters);

diff --git a/velox/substrait/TypeUtils.h b/velox/substrait/TypeUtils.h
@@ -67,6 +67,13 @@ struct RangeTraits<TypeKind::DOUBLE> {
   using NativeType = double;
 };
 
+template <>
+struct RangeTraits<TypeKind::BOOLEAN> {
+  using RangeType = common::BigintRange;
+  using MultiRangeType = common::BigintMultiRange;
+  using NativeType = bool;
+};
+
 template <>
 struct RangeTraits<TypeKind::VARCHAR> {
   using RangeType = common::BytesRange;