[oap-native-sql] Gandiva: Add shift_left, shift_right (#85)

oap-project · Feb 3, 2021 · c45b759 · c45b759
1 parent 73005fc
commit c45b759
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 0 deletions.
diff --git a/cpp/src/gandiva/function_registry_arithmetic.cc b/cpp/src/gandiva/function_registry_arithmetic.cc
@@ -105,6 +105,11 @@ std::vector<NativeFunction> GetArithmeticFunctionRegistry() {
       UNARY_SAFE_NULL_IF_NULL(round, {}, int64, int64),
       BINARY_GENERIC_SAFE_NULL_IF_NULL(round, {}, int32, int32, int32),
       BINARY_GENERIC_SAFE_NULL_IF_NULL(round, {}, int64, int32, int64),
+      // bitwise functions
+      BINARY_GENERIC_SAFE_NULL_IF_NULL(shift_left, {}, int32, int32, int32),
+      BINARY_GENERIC_SAFE_NULL_IF_NULL(shift_left, {}, int64, int32, int64),
+      BINARY_GENERIC_SAFE_NULL_IF_NULL(shift_right, {}, int32, int32, int32),
+      BINARY_GENERIC_SAFE_NULL_IF_NULL(shift_right, {}, int64, int32, int64),
 
       // compare functions
       BINARY_RELATIONAL_BOOL_FN(equal, ({"eq", "same"})),

diff --git a/cpp/src/gandiva/precompiled/arithmetic_ops.cc b/cpp/src/gandiva/precompiled/arithmetic_ops.cc
@@ -300,6 +300,27 @@ BITWISE_NOT(int32)
 BITWISE_NOT(int64)
 
 #undef BITWISE_NOT
+#define SHIFT_LEFT_INT(LTYPE, RTYPE)                                               \
+  FORCE_INLINE                                                                      \
+  gdv_##LTYPE shift_left_##LTYPE##_##RTYPE(gdv_##LTYPE in1, gdv_##RTYPE in2) {     \
+    return static_cast<gdv_##LTYPE>(in1 << in2);                                    \
+  }
+
+SHIFT_LEFT_INT(int32, int32)
+SHIFT_LEFT_INT(int64, int32)
+
+#undef SHIFT_RIGHT_INT
+
+#define SHIFT_RIGHT_INT(LTYPE, RTYPE)                                               \
+  FORCE_INLINE                                                                      \
+  gdv_##LTYPE shift_right_##LTYPE##_##RTYPE(gdv_##LTYPE in1, gdv_##RTYPE in2) {     \
+    return static_cast<gdv_##LTYPE>(in1 >> in2);                                    \
+  }
+
+SHIFT_RIGHT_INT(int32, int32)
+SHIFT_RIGHT_INT(int64, int32)
+
+#undef SHIFT_RIGHT_INT
 
 #undef DATE_FUNCTION
 #undef DATE_TYPES

diff --git a/cpp/src/gandiva/tests/projector_test.cc b/cpp/src/gandiva/tests/projector_test.cc
@@ -210,6 +210,93 @@ TEST_F(TestProjector, TestProjectCacheDecimalCast) {
   EXPECT_EQ(projector0.get(), projector2.get());
 }
 
+TEST_F(TestProjector, TestShiftRight) {
+  // schema for input fields
+  auto field0 = field("f0", int32());
+  auto field1 = field("f1", int32());
+  auto schema = arrow::schema({field0, field1});
+
+  // output fields
+  auto field_shift_right = field("shift_right", int32());
+
+  // Build expression
+  auto shift_right_expr = TreeExprBuilder::MakeExpression("shift_right", {field0, field1},
+      field_shift_right);
+
+  std::shared_ptr<Projector> projector;
+
+  auto status =
+      Projector::Make(schema, {shift_right_expr}, TestConfiguration(), &projector);
+  EXPECT_TRUE(status.ok());
+
+  if (!status.ok()) {
+    std::cout << status.message() << std::endl;
+  }
+
+  // Create a row-batch with some sample data
+  int num_records = 4;
+  auto array0 = MakeArrowArrayInt32({4, 8, 16, 32}, {true, true, true, true});
+  auto array1 = MakeArrowArrayInt32({1, 2, 3, 4}, {true, true, true, true});
+
+  // expected output
+  auto exp_shift_right = MakeArrowArrayInt32({2, 2, 2, 2}, {true, true, true, true});
+
+  // prepare input record batch
+  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0, array1});
+
+  // Evaluate expression
+  arrow::ArrayVector outputs;
+  status = projector->Evaluate(*in_batch, pool_, &outputs);
+  EXPECT_TRUE(status.ok());
+
+  // Validate results
+  EXPECT_ARROW_ARRAY_EQUALS(exp_shift_right, outputs.at(0));
+}
+
+
+TEST_F(TestProjector, TestShiftLeft) {
+  // schema for input fields
+  auto field0 = field("f0", int32());
+  auto field1 = field("f1", int32());
+  auto schema = arrow::schema({field0, field1});
+
+  // output fields
+  auto field_shift_left = field("shift_left", int32());
+
+  // Build expression
+  auto shift_left_expr = TreeExprBuilder::MakeExpression("shift_left", {field0, field1},
+                                                          field_shift_left);
+
+  std::shared_ptr<Projector> projector;
+
+  auto status =
+      Projector::Make(schema, {shift_left_expr}, TestConfiguration(), &projector);
+  EXPECT_TRUE(status.ok());
+
+  if (!status.ok()) {
+    std::cout << status.message() << std::endl;
+  }
+
+  // Create a row-batch with some sample data
+  int num_records = 4;
+  auto array0 = MakeArrowArrayInt32({4, 8, 16, 32}, {true, true, true, true});
+  auto array1 = MakeArrowArrayInt32({4, 3, 2, 1}, {true, true, true, true});
+
+  // expected output
+  auto exp_shift_left = MakeArrowArrayInt32({64, 64, 64, 64}, {true, true, true, true});
+
+  // prepare input record batch
+  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0, array1});
+
+  // Evaluate expression
+  arrow::ArrayVector outputs;
+  status = projector->Evaluate(*in_batch, pool_, &outputs);
+  EXPECT_TRUE(status.ok());
+
+  // Validate results
+  EXPECT_ARROW_ARRAY_EQUALS(exp_shift_left, outputs.at(0));
+}
+
 TEST_F(TestProjector, TestIntSumSub) {
   // schema for input fields
   auto field0 = field("f0", int32());