Support POWER in query builder.

Signed-off-by: ienkovich <[email protected]>

Author: ienkovich

omniscidb/QueryBuilder/QueryBuilder.cpp

@@ -1153,6 +1153,36 @@ BuilderExpr BuilderExpr::floor() const {
                             << expr_->type()->toString();
 }
 
+BuilderExpr BuilderExpr::pow(const BuilderExpr& rhs) const {
+  if (!expr_->type()->isNumber() || !rhs.type()->isNumber()) {
+    throw InvalidQueryError() << "Cannot apply POW operation for operand types: "
+                              << type()->toString() << " and " << rhs.type()->toString();
+  }
+  ExprPtrVector args;
+  args.push_back(cast(ctx().fp64(type()->nullable())).expr());
+  args.push_back(rhs.cast(ctx().fp64(rhs.type()->nullable())).expr());
+  auto res_type = ctx().fp64(type()->nullable() || rhs.type()->nullable());
+  auto pow_expr =
+      hdk::ir::makeExpr<hdk::ir::FunctionOper>(res_type, "POWER", std::move(args));
+  return {builder_, pow_expr, "", true};
+}
+
+BuilderExpr BuilderExpr::pow(int val) const {
+  return pow(builder_->cst(val, builder_->ctx_.int32(false)));
+}
+
+BuilderExpr BuilderExpr::pow(int64_t val) const {
+  return pow(builder_->cst(val, builder_->ctx_.int64(false)));
+}
+
+BuilderExpr BuilderExpr::pow(float val) const {
+  return pow(builder_->cst(val, builder_->ctx_.fp32(false)));
+}
+
+BuilderExpr BuilderExpr::pow(double val) const {
+  return pow(builder_->cst(val, builder_->ctx_.fp64(false)));
+}
+
 BuilderExpr BuilderExpr::logicalAnd(const BuilderExpr& rhs) const {
   try {
     auto bin_oper = Analyzer::normalizeOperExpr(

omniscidb/QueryBuilder/QueryBuilder.h

@@ -37,6 +37,7 @@ class BuilderExpr {
   bool isAutoNamed() const { return auto_name_; }
 
   ExprPtr expr() const { return expr_; }
+  const Type* type() const { return expr_->type(); }
 
   BuilderExpr avg() const;
   BuilderExpr min() const;
@@ -108,6 +109,12 @@ class BuilderExpr {
   BuilderExpr ceil() const;
   BuilderExpr floor() const;
 
+  BuilderExpr pow(const BuilderExpr& rhs) const;
+  BuilderExpr pow(int val) const;
+  BuilderExpr pow(int64_t val) const;
+  BuilderExpr pow(float val) const;
+  BuilderExpr pow(double val) const;
+
   BuilderExpr logicalAnd(const BuilderExpr& rhs) const;
   BuilderExpr logicalOr(const BuilderExpr& rhs) const;
 

omniscidb/Tests/QueryBuilderTest.cpp

@@ -3952,6 +3952,34 @@ TEST_F(QueryBuilderTest, Floor) {
   EXPECT_THROW(scan.ref("col_arr_i32").floor(), InvalidQueryError);
 }
 
+TEST_F(QueryBuilderTest, Pow) {
+  QueryBuilder builder(ctx(), schema_mgr_, configPtr());
+  auto scan = builder.scan("test3");
+  checkFunctionOper(scan.ref("col_bi").pow(2), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_i").pow((int64_t)3), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_f").pow(2.3f), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_d").pow(2.0), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_dec").pow(scan.ref("col_d")), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_d").pow(scan.ref("col_bi")), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_d").pow(scan.ref("col_i")), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_d").pow(scan.ref("col_f")), "POWER", 2, ctx().fp64());
+  checkFunctionOper(scan.ref("col_d").pow(scan.ref("col_dec")), "POWER", 2, ctx().fp64());
+  EXPECT_THROW(scan.ref("col_b").pow(2), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_str").pow(2), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_dict").pow(2), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_date").pow(2), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_time").pow(2), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_timestamp").pow(2), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_arr_i32").pow(2), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_d").pow(scan.ref("col_b")), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_d").pow(scan.ref("col_str")), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_d").pow(scan.ref("col_dict")), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_d").pow(scan.ref("col_date")), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_d").pow(scan.ref("col_time")), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_d").pow(scan.ref("col_timestamp")), InvalidQueryError);
+  EXPECT_THROW(scan.ref("col_d").pow(scan.ref("col_arr_i32")), InvalidQueryError);
+}
+
 TEST_F(QueryBuilderTest, SimpleProjection) {
   QueryBuilder builder(ctx(), schema_mgr_, configPtr());
   compare_test1_data(builder.scan("test1").proj({0, 1, 2, 3}));

python/pyhdk/_builder.pxd

@@ -50,6 +50,8 @@ cdef extern from "omniscidb/QueryBuilder/QueryBuilder.h":
     CBuilderExpr ceil() except +
     CBuilderExpr floor() except +
 
+    CBuilderExpr pow(const CBuilderExpr&) except +
+
     CBuilderExpr add(const CBuilderExpr&) except +
     CBuilderExpr sub(const CBuilderExpr&) except +
     CBuilderExpr mul(const CBuilderExpr&) except +

python/pyhdk/_builder.pyx

@@ -241,6 +241,12 @@ cdef class QueryExpr:
     res.c_expr = self.c_expr.unnest()
     return res
 
+  def pow(self, value):
+    value = self._process_op_expr(value)
+    res = QueryExpr()
+    res.c_expr = self.c_expr.pow((<QueryExpr>value).c_expr)
+    return res
+
   def at(self, value):
     value = self._process_op_expr(value)
     res = QueryExpr()

python/tests/test_pyhdk_api.py

@@ -836,6 +836,23 @@ def test_unnest(self):
 
         hdk.drop_table(ht)
 
+    def test_pow(self):
+        hdk = pyhdk.init()
+        ht = hdk.import_pydict({"a": [1, 2, 3, 4, 5], "b": [1.0, 2.0, 3.0, 4.0, 5.0]})
+
+        self.check_res(
+            ht.proj(
+                a1=ht["a"].pow(2), a2=ht["a"].pow(ht["b"]), b=ht["b"].pow(2.0)
+            ).run(),
+            {
+                "a1": [1.0, 4.0, 9.0, 16.0, 25.0],
+                "a2": [1.0, 4.0, 27.0, 256.0, 3125.0],
+                "b": [1.0, 4.0, 9.0, 16.0, 25.0],
+            },
+        )
+
+        hdk.drop_table(ht)
+
     def test_at(self):
         hdk = pyhdk.init()
         ht = hdk.create_table("test1", [("a", "array(int)"), ("b", "int")])