feat(udf): array_combine & array_join

cases/query/udf_query.yaml

@@ -554,6 +554,98 @@ cases:
  - c1 bool
  data: |
  true, false
+ - id: array_join
+ mode: request-unsupport
+ sql: |
+ select
+ array_join(["1", "2"], ",") c1,
+ array_join(["1", "2"], "") c2,
+ array_join(["1", "2"], cast(null as string)) c3,
+ array_join(["1", NULL, "4", "5", NULL], "-") c4,
+ array_join(array<string>[], ",") as c5
+ expect:
+ columns:
+ - c1 string
+ - c2 string
+ - c3 string
+ - c4 string
+ - c5 string
+ rows:
+ - ["1,2", "12", "12", "1-4-5", ""]
+ - id: array_combine
+ mode: request-unsupport
+ sql: |
+ select
+ array_join(array_combine("-", ["1", "2"], ["3", "4"]), ",") c0,
+ expect:
+ columns:
+ - c0 string
+ rows:
+ - ["1-3,1-4,2-3,2-4"]
+
+ - id: array_combine_2
+ desc: array_combine casting array to array<string> first
+ mode: request-unsupport
+ sql: |
+ select
+ array_join(array_combine("-", [1, 2], [3, 4]), ",") c0,
+ array_join(array_combine("-", [1, 2], array<int64>[3], ["5", "6"]), ",") c1,
+ array_join(array_combine("|", ["1"], [timestamp(1717171200000), timestamp("2024-06-02 12:00:00")]), ",") c2,
+ array_join(array_combine("|", ["1"]), ",") c3,
+ expect:
+ columns:
+ - c0 string
+ - c1 string
+ - c2 string
+ - c3 string
+ rows:
+ - ["1-3,1-4,2-3,2-4", "1-3-5,1-3-6,2-3-5,2-3-6", "1|2024-06-01 00:00:00,1|2024-06-02 12:00:00", "1"]
+ - id: array_combine_3
+ desc: null values skipped
+ mode: request-unsupport
+ sql: |
+ select
+ array_join(array_combine("-", [1, NULL], [3, 4]), ",") c0,
+ array_join(array_combine("-", ARRAY<INT>[NULL], ["9", "8"]), ",") c1,
+ array_join(array_combine(string(NULL), ARRAY<INT>[1], ["9", "8"]), ",") c2,
+ expect:
+ columns:
+ - c0 string
+ - c1 string
+ - c2 string
+ rows:
+ - ["1-3,1-4", "", "19,18"]
+ - id: array_combine_4
+ desc: construct array from table
+ mode: request-unsupport
+ inputs:
+ - name: t1
+ columns: ["col1:int32", "std_ts:timestamp", "col2:string"]
+ indexs: ["index1:col1:std_ts"]
+ rows:
+ - [1, 1590115420001, "foo"]
+ - [2, 1590115420001, "bar"]
+ sql: |
+ select
+ col1,
+ array_join(array_combine("-", [col1, 10], [col2, "c2"]), ",") c0,
+ from t1
+ expect:
+ columns:
+ - col1 int32
+ - c0 string
+ rows:
+ - [1, "1-foo,1-c2,10-foo,10-c2"]
+ - [2, "2-bar,2-c2,10-bar,10-c2"]
+ - id: array_combine_err1
+ mode: request-unsupport
+ sql: |
+ select
+ array_join(array_combine("-"), ",") c0,
+ expect:
+ success: false
+ msg: |
+ Fail to resolve expression: array_join(array_combine(-), ,)
 
  # ================================================================
  # Map data type

hybridse/src/base/cartesian_product.cc

@@ -0,0 +1,62 @@
+/**
+ * Copyright (c) 2024 OpenMLDB authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "base/cartesian_product.h"
+
+#include <algorithm>
+
+#include "absl/types/span.h"
+
+namespace hybridse {
+namespace base {
+
+static auto cartesian_product(const std::vector<std::vector<int>>& lists) {
+ std::vector<std::vector<int>> result;
+ if (std::find_if(std::begin(lists), std::end(lists), [](auto e) -> bool { return e.size() == 0; }) !=
+ std::end(lists)) {
+ return result;
+ }
+ for (auto& e : lists[0]) {
+ result.push_back({e});
+ }
+ for (size_t i = 1; i < lists.size(); ++i) {
+ std::vector<std::vector<int>> temp;
+ for (auto& e : result) {
+ for (auto f : lists[i]) {
+ auto e_tmp = e;
+ e_tmp.push_back(f);
+ temp.push_back(e_tmp);
+ }
+ }
+ result = temp;
+ }
+ return result;
+}
+
+std::vector<std::vector<int>> cartesian_product(absl::Span<int const> vec) {
+ std::vector<std::vector<int>> input;
+ for (auto& v : vec) {
+ std::vector<int> seq(v, 0);
+ for (int i = 0; i < v; ++i) {
+ seq[i] = i;
+ }
+ input.push_back(seq);
+ }
+ return cartesian_product(input);
+}
+
+} // namespace base
+} // namespace hybridse

hybridse/src/base/cartesian_product.h

@@ -0,0 +1,32 @@
+/**
+ * Copyright (c) 2024 OpenMLDB authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef HYBRIDSE_SRC_BASE_CARTESIAN_PRODUCT_H_
+#define HYBRIDSE_SRC_BASE_CARTESIAN_PRODUCT_H_
+
+#include <vector>
+
+#include "absl/types/span.h"
+
+namespace hybridse {
+namespace base {
+
+std::vector<std::vector<int>> cartesian_product(absl::Span<int const> vec);
+
+} // namespace base
+} // namespace hybridse
+
+#endif // HYBRIDSE_SRC_BASE_CARTESIAN_PRODUCT_H_

hybridse/src/codegen/array_ir_builder.cc

@@ -18,8 +18,12 @@
 
 #include <string>
 
+#include "absl/strings/substitute.h"
+#include "base/fe_status.h"
+#include "codegen/cast_expr_ir_builder.h"
 #include "codegen/context.h"
 #include "codegen/ir_base_builder.h"
+#include "codegen/string_ir_builder.h"
 
 namespace hybridse {
 namespace codegen {
@@ -122,5 +126,116 @@
  return true;
 }
 
+absl::StatusOr<NativeValue> ArrayIRBuilder::ExtractElement(CodeGenContextBase* ctx, const NativeValue& arr,
+ const NativeValue& key) const {
+ return absl::UnimplementedError("array extract element");
+}
+
+absl::StatusOr<llvm::Value*> ArrayIRBuilder::NumElements(CodeGenContextBase* ctx, llvm::Value* arr) const {
+ llvm::Value* out = nullptr;
+ if (!Load(ctx->GetCurrentBlock(), arr, SZ_IDX, &out)) {
+ return absl::InternalError("codegen: fail to extract array size");
+ }
+
+ return out;
+}
+
+absl::StatusOr<llvm::Value*> ArrayIRBuilder::CastToArrayString(CodeGenContextBase* ctx, llvm::Value* src) {
+ auto sb = StructTypeIRBuilder::CreateStructTypeIRBuilder(ctx->GetModule(), src->getType());
+ CHECK_ABSL_STATUSOR(sb);
+
+ ArrayIRBuilder* src_builder = dynamic_cast<ArrayIRBuilder*>(sb.value().get());
+ if (!src_builder) {
+ return absl::InvalidArgumentError("input value not a array");
+ }
+
+ llvm::Type* src_ele_type = src_builder->element_type_;
+ if (IsStringPtr(src_ele_type)) {
+ // already array<string>
+ return src;
+ }
+
+ auto fields = src_builder->Load(ctx, src);
+ CHECK_ABSL_STATUSOR(fields);
+ llvm::Value* src_raws = fields.value().at(RAW_IDX);
+ llvm::Value* src_nulls = fields.value().at(NULL_IDX);
+ llvm::Value* num_elements = fields.value().at(SZ_IDX);
+
+ llvm::Value* casted = nullptr;
+ if (!CreateDefault(ctx->GetCurrentBlock(), &casted)) {
+ return absl::InternalError("codegen error: fail to construct default array");
+ }
+ // initialize each element
+ CHECK_ABSL_STATUS(Initialize(ctx, casted, {num_elements}));
+
+ auto builder = ctx->GetBuilder();
+ auto dst_fields = Load(ctx, casted);
+ CHECK_ABSL_STATUSOR(fields);
+ auto* raw_array_ptr = dst_fields.value().at(RAW_IDX);
+ auto* nullables_ptr = dst_fields.value().at(NULL_IDX);
+
+ llvm::Type* idx_type = builder->getInt64Ty();
+ llvm::Value* idx = builder->CreateAlloca(idx_type);
+ builder->CreateStore(builder->getInt64(0), idx);
+ CHECK_STATUS_TO_ABSL(ctx->CreateWhile(
+ [&](llvm::Value** cond) -> base::Status {
+ *cond = builder->CreateICmpSLT(builder->CreateLoad(idx_type, idx), num_elements);
+ return {};
+ },
+ [&]() -> base::Status {
+ llvm::Value* idx_val = builder->CreateLoad(idx_type, idx);
+ codegen::CastExprIRBuilder cast_builder(ctx->GetCurrentBlock());
+
+ llvm::Value* src_ele_value =
+ builder->CreateLoad(src_ele_type, builder->CreateGEP(src_ele_type, src_raws, idx_val));
+ llvm::Value* dst_ele =
+ builder->CreateLoad(element_type_, builder->CreateGEP(element_type_, raw_array_ptr, idx_val));
+
+ codegen::StringIRBuilder str_builder(ctx->GetModule());
+ auto s = str_builder.CastFrom(ctx->GetCurrentBlock(), src_ele_value, dst_ele);
+ CHECK_TRUE(s.ok(), common::kCodegenError, s.ToString());
+
+ builder->CreateStore(
+ builder->CreateLoad(builder->getInt1Ty(), builder->CreateGEP(builder->getInt1Ty(), src_nulls, idx_val)),
+ builder->CreateGEP(builder->getInt1Ty(), nullables_ptr, idx_val));
+
+ builder->CreateStore(builder->CreateAdd(idx_val, builder->getInt64(1)), idx);
+ return {};
+ }));
+
+ CHECK_ABSL_STATUS(Set(ctx, casted, {raw_array_ptr, nullables_ptr, num_elements}));
+ return casted;
+}
+
+absl::Status ArrayIRBuilder::Initialize(CodeGenContextBase* ctx, ::llvm::Value* alloca,
+ absl::Span<llvm::Value* const> args) const {
+ auto* builder = ctx->GetBuilder();
+ StringIRBuilder str_builder(ctx->GetModule());
+ auto ele_type = str_builder.GetType();
+ if (!alloca->getType()->isPointerTy() || alloca->getType()->getPointerElementType() != struct_type_ ||
+ ele_type->getPointerTo() != element_type_) {
+ return absl::UnimplementedError(absl::Substitute(
+ "not able to Initialize array except array<string>, got type $0", GetLlvmObjectString(alloca->getType())));
+ }
+ if (args.size() != 1) {
+ // require one argument that is array size
+ return absl::InvalidArgumentError("initialize array requries one argument which is array size");
+ }
+ if (!args[0]->getType()->isIntegerTy()) {
+ return absl::InvalidArgumentError("array size argument should be integer");
+ }
+ auto sz = args[0];
+ if (sz->getType() != builder->getInt64Ty()) {
+ CastExprIRBuilder cast_builder(ctx->GetCurrentBlock());
+ base::Status s;
+ cast_builder.SafeCastNumber(sz, builder->getInt64Ty(), &sz, s);
+ CHECK_STATUS_TO_ABSL(s);
+ }
+ auto fn = ctx->GetModule()->getOrInsertFunction("hybridse_alloc_array_string", builder->getVoidTy(),
+ struct_type_->getPointerTo(), builder->getInt64Ty());
+
+ builder->CreateCall(fn, {alloca, sz});
+ return absl::OkStatus();
+}
 } // namespace codegen
 } // namespace hybridse

hybridse/src/codegen/array_ir_builder.h

@@ -42,11 +42,21 @@ class ArrayIRBuilder : public StructTypeIRBuilder {
  CHECK_TRUE(false, common::kCodegenError, "casting to array un-implemented");
  };
 
- private:
- void InitStructType() override;
+ absl::StatusOr<llvm::Value*> CastToArrayString(CodeGenContextBase* ctx, llvm::Value* src);
+
+ absl::StatusOr<NativeValue> ExtractElement(CodeGenContextBase* ctx, const NativeValue& arr,
+ const NativeValue& key) const override;
+
+ absl::StatusOr<llvm::Value*> NumElements(CodeGenContextBase* ctx, llvm::Value* arr) const override;
 
  bool CreateDefault(::llvm::BasicBlock* block, ::llvm::Value** output) override;
 
+ absl::Status Initialize(CodeGenContextBase* ctx, ::llvm::Value* alloca,
+ absl::Span<llvm::Value* const> args) const override;
+
+ private:
+ void InitStructType() override;
+
  private:
  ::llvm::Type* element_type_ = nullptr;
 };

hybridse/src/codegen/ir_base_builder.cc

@@ -575,12 +575,12 @@
  if (type_pointee->isStructTy()) {
  auto* key_type = type_pointee->getStructElementType(1);
  const node::TypeNode* key = nullptr;
- if (key_type->isPointerTy() && !GetFullType(nm, key_type->getPointerElementType(), &key)) {
+ if (!key_type->isPointerTy() || !GetFullType(nm, key_type->getPointerElementType(), &key)) {
  return false;
  }
  const node::TypeNode* value = nullptr;
  auto* value_type = type_pointee->getStructElementType(2);
- if (value_type->isPointerTy() && !GetFullType(nm, value_type->getPointerElementType(), &value)) {
+ if (!value_type->isPointerTy() || !GetFullType(nm, value_type->getPointerElementType(), &value)) {
  return false;
  }
 
@@ -590,6 +590,22 @@
  }
  return false;
  }
+ case hybridse::node::kArray: {
+ if (type->isPointerTy()) {
+ auto type_pointee = type->getPointerElementType();
+ if (type_pointee->isStructTy()) {
+ auto* key_type = type_pointee->getStructElementType(0);
+ const node::TypeNode* key = nullptr;
+ if (!key_type->isPointerTy() || !GetFullType(nm, key_type->getPointerElementType(), &key)) {
+ return false;
+ }
+
+ *type_node = nm->MakeNode<node::TypeNode>(node::DataType::kArray, key);
+ return true;
+ }
+ }
+ return false;
+ }
  default: {
  *type_node = nm->MakeTypeNode(base);
  return true;