Add cardinality API to QueryBuilder.

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

Author: ienkovich

omniscidb/QueryBuilder/QueryBuilder.cpp

@@ -1587,6 +1587,30 @@ BuilderExpr BuilderExpr::at(int64_t idx) const {
   return at(builder_->cst(idx, builder_->ctx_.int64(false)));
 }
 
+BuilderExpr BuilderExpr::cardinality() const {
+  if (!expr_->type()->isArray()) {
+    throw InvalidQueryError() << "Non-array type used for CARDINALITY operator: "
+                              << expr_->type()->toString();
+  }
+
+  if (expr_->is<Constant>()) {
+    if (expr_->as<Constant>()->isNull()) {
+      return builder_->nullCst(builder_->ctx_.int32());
+    }
+    return builder_->cst(static_cast<int32_t>(expr_->as<Constant>()->valueList().size()),
+                         builder_->ctx_.int32(false));
+  }
+
+  if (expr_->type()->size() > 0 && !expr_->type()->nullable()) {
+    auto elem_type = expr_->type()->as<hdk::ir::ArrayBaseType>()->elemType();
+    return builder_->cst(expr_->type()->size() / elem_type->size(),
+                         builder_->ctx_.int32(false));
+  }
+
+  auto cardinality_expr = makeExpr<CardinalityExpr>(expr_);
+  return {builder_, cardinality_expr, "", true};
+}
+
 BuilderExpr BuilderExpr::over() const {
   return over(std::vector<BuilderExpr>());
 }

omniscidb/QueryBuilder/QueryBuilder.h

@@ -216,6 +216,8 @@ class BuilderExpr {
   BuilderExpr at(int idx) const;
   BuilderExpr at(int64_t idx) const;
 
+  BuilderExpr cardinality() const;
+
   BuilderExpr over() const;
   BuilderExpr over(const BuilderExpr& key) const;
   BuilderExpr over(const std::vector<BuilderExpr>& keys) const;

omniscidb/Tests/QueryBuilderTest.cpp

@@ -302,6 +302,12 @@ void checkWindowCollation(const BuilderExpr& expr,
   ASSERT_EQ(wnd_fn->collation()[idx].nulls_first, null_pos == NullSortedPosition::First);
 }
 
+void checkCardinality(const BuilderExpr& expr, const BuilderExpr& op) {
+  auto cardinality_expr = expr.expr()->as<CardinalityExpr>();
+  ASSERT_TRUE(cardinality_expr);
+  ASSERT_TRUE(cardinality_expr->arg()->equal(op.expr().get()));
+}
+
 }  // anonymous namespace
 
 class QueryBuilderTest : public TestSuite {
@@ -463,6 +469,18 @@ class QueryBuilderTest : public TestSuite {
                     "2,3,2,3,2,3,2,3,2,3,2,3,2,3,2,3,,,,,,,,\n"
                     "3,2,3,2,3,2,3,2,3,2,3,2,3,2,3,2,,,,,,,,\n"
                     ",1,4,1,,1,4,1,,1,4,1,,1,4,1,,,,,,,,\n");
+
+    createTable("test_cardinality",
+                {{"col_arr_i32", ctx().arrayVarLen(ctx().int32())},
+                 {"col_arr_i32x2", ctx().arrayFixed(2, ctx().int32())},
+                 {"col_arr_i32nn", ctx().arrayVarLen(ctx().int32(), 4, false)},
+                 {"col_arr_i32x2nn", ctx().arrayFixed(2, ctx().int32(), false)}});
+    insertJsonValues(
+        "test_cardinality",
+        R"___({"col_arr_i32": [1], "col_arr_i32x2": [0, 1], "col_arr_i32nn": [1, 2], "col_arr_i32x2nn": [3, 4]}
+              {"col_arr_i32": null, "col_arr_i32x2": null, "col_arr_i32nn": [1], "col_arr_i32x2nn": [3, 4]}
+              {"col_arr_i32": [], "col_arr_i32x2": [0, 1], "col_arr_i32nn": [], "col_arr_i32x2nn": [3, 4]}
+              {"col_arr_i32": [1, 2], "col_arr_i32x2": null, "col_arr_i32nn": [null], "col_arr_i32x2nn": [3, 4]})___");
   }
 
   static void TearDownTestSuite() {
@@ -479,6 +497,8 @@ class QueryBuilderTest : public TestSuite {
     dropTable("withNull");
     dropTable("test_tmstmp");
     dropTable("test_tmstmp_nullable");
+    dropTable("test_bitwise");
+    dropTable("test_cardinality");
   }
 
   void compare_res_fields(const ExecutionResult& res,
@@ -5196,6 +5216,42 @@ TEST_F(QueryBuilderTest, BwNot_Exec) {
                             inline_null_value<int64_t>()}));
 }
 
+TEST_F(QueryBuilderTest, Cardinality) {
+  QueryBuilder builder(ctx(), schema_mgr_, configPtr());
+  auto scan = builder.scan("test_cardinality");
+  checkCardinality(scan.ref("col_arr_i32").cardinality(), scan.ref("col_arr_i32"));
+  checkCardinality(scan.ref("col_arr_i32x2").cardinality(), scan.ref("col_arr_i32x2"));
+  checkCardinality(scan.ref("col_arr_i32nn").cardinality(), scan.ref("col_arr_i32nn"));
+  checkCst(scan.ref("col_arr_i32x2nn").cardinality(), 2, ctx().int32(false));
+  checkCst(builder.cst({1, 2, 3, 4}, "array(int8)").cardinality(), 4, ctx().int32(false));
+  checkNullCst(builder.nullCst("array(int8)").cardinality(), ctx().int32());
+  auto scan2 = builder.scan("test3");
+  EXPECT_THROW(scan2.ref("col_i").cardinality(), InvalidQueryError);
+  EXPECT_THROW(scan2.ref("col_d").cardinality(), InvalidQueryError);
+  EXPECT_THROW(scan2.ref("col_dec").cardinality(), InvalidQueryError);
+  EXPECT_THROW(scan2.ref("col_time").cardinality(), InvalidQueryError);
+  EXPECT_THROW(scan2.ref("col_str").cardinality(), InvalidQueryError);
+  EXPECT_THROW(builder.cst(1).cardinality(), InvalidQueryError);
+}
+
+TEST_F(QueryBuilderTest, Cardinality_Exec) {
+  QueryBuilder builder(ctx(), schema_mgr_, configPtr());
+  auto scan = builder.scan("test_cardinality");
+  auto dag = scan.proj({scan.ref("col_arr_i32").cardinality(),
+                        scan.ref("col_arr_i32x2").cardinality(),
+                        scan.ref("col_arr_i32nn").cardinality(),
+                        scan.ref("col_arr_i32x2nn").cardinality()})
+                 .finalize();
+  auto res = runQuery(std::move(dag));
+  compare_res_data(
+      res,
+      std::vector<int32_t>({1, inline_null_value<int32_t>(), 0, 2}),
+      std::vector<int32_t>(
+          {2, inline_null_value<int32_t>(), 2, inline_null_value<int32_t>()}),
+      std::vector<int32_t>({2, 1, 0, 1}),
+      std::vector<int32_t>({2, 2, 2, 2}));
+}
+
 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

@@ -88,6 +88,7 @@ cdef extern from "omniscidb/QueryBuilder/QueryBuilder.h":
     CBuilderExpr ge(const CBuilderExpr&) except +
 
     CBuilderExpr at(const CBuilderExpr&) except +
+    CBuilderExpr cardinality() except +
 
     CBuilderExpr over(const vector[CBuilderExpr]&) except +
     CBuilderExpr orderBy(const vector[CBuilderOrderByKey]&) except +

python/pyhdk/_builder.pyx

@@ -308,6 +308,11 @@ cdef class QueryExpr:
     res.c_expr = self.c_expr.at((<QueryExpr>value).c_expr)
     return res
 
+  def cardinality(self):
+    res = QueryExpr()
+    res.c_expr = self.c_expr.cardinality()
+    return res
+
   def over(self, *args):
     cdef vector[CBuilderExpr] keys
     for arg in args:

python/pyhdk/hdk.py

@@ -1347,6 +1347,29 @@ def at(self, index):
         """
         pass
 
+    def cardinality(self):
+        """
+        Create cardinality expression to access array size.
+
+        -------
+        QueryExpr
+
+        Examples
+        --------
+        >>> hdk = pyhdk.init()
+        >>> ht = hdk.create_table("test_arr", [("a", "array(int)")])
+        >>> hdk.import_pydict({"a": [[1, 2], [2, 3, 4], None, []]}, ht)
+        >>> ht.proj(len=ht["a"].cardinality()).run()
+        Schema:
+          len: INT32
+        Data:
+        2
+        3
+        null
+        0
+        """
+        pass
+
     def over(self, *args):
         """
         Add partition keys to window function expression. Also, transforms min, max, sum, avg, and count

python/tests/test_pyhdk_api.py

@@ -917,6 +917,18 @@ def test_at(self, exe_cfg):
 
         hdk.drop_table(ht)
 
+    def test_cardinality(self, exe_cfg):
+        hdk = pyhdk.init()
+        ht = hdk.create_table("test_arr", [("a", "array(int)")])
+        hdk.import_pydict({"a": [[1, 2], [2, 3, 4], None, []]}, ht)
+
+        check_res(
+            ht.proj(len=ht["a"].cardinality()).run(device_type=exe_cfg.device_type),
+            {"len": [2, 3, "null", 0]},
+        )
+
+        hdk.drop_table(ht)
+
     def test_run_on_res(self, exe_cfg):
         hdk = pyhdk.init()
         ht = hdk.import_pydict({"a": [1, 2, 3, 4, 5], "b": [10, 20, 30, 40, 50]})