feat: casting arrays to array<string> for array_combine

WIP, string allocation need fix
4paradigm · Jun 3, 2024 · 6ebe49b · 6ebe49b
1 parent 3c180bf
commit 6ebe49b
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 5 deletions.
diff --git a/cases/query/udf_query.yaml b/cases/query/udf_query.yaml
@@ -581,6 +581,18 @@ cases:
       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,
+  #   expect:
+  #     columns:
+  #       - c0 string
+  #     rows:
+  #       - ["1-3,1-4,2-3,2-4"]
+
   # ================================================================
   # Map data type
   # FIXME: request mode tests disabled, because TestRequestEngineForLastRow cause SEG FAULT

diff --git a/hybridse/src/codegen/array_ir_builder.cc b/hybridse/src/codegen/array_ir_builder.cc
@@ -18,6 +18,7 @@
 
 #include <string>
 
+#include "codegen/cast_expr_ir_builder.h"
 #include "codegen/context.h"
 #include "codegen/ir_base_builder.h"
 
@@ -136,5 +137,60 @@ absl::StatusOr<llvm::Value*> ArrayIRBuilder::NumElements(CodeGenContextBase* ctx
     return out;
 }
 
+absl::StatusOr<llvm::Value*> ArrayIRBuilder::CastFrom(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_;
+    auto fields = src_builder->Load(ctx, src);
+    CHECK_ABSL_STATUSOR(fields);
+    llvm::Value* src_raws = fields.value().at(RAW_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");
+    }
+
+    auto builder = ctx->GetBuilder();
+    auto* raw_array_ptr = builder->CreateAlloca(element_type_, num_elements);
+    auto* nullables_ptr = builder->CreateAlloca(builder->getInt1Ty(), num_elements);
+
+    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));
+
+            NativeValue out;
+            CHECK_STATUS(cast_builder.Cast(NativeValue::Create(src_ele_value), element_type_, &out));
+
+            builder->CreateStore(out.GetRaw(), builder->CreateGEP(element_type_, raw_array_ptr, idx_val));
+            builder->CreateStore(out.GetIsNull(builder),
+                                 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;
+}
+
 }  // namespace codegen
 }  // namespace hybridse
diff --git a/hybridse/src/codegen/array_ir_builder.h b/hybridse/src/codegen/array_ir_builder.h
@@ -42,6 +42,8 @@ class ArrayIRBuilder : public StructTypeIRBuilder {
         CHECK_TRUE(false, common::kCodegenError, "casting to array un-implemented");
     };
 
+    absl::StatusOr<llvm::Value*> CastFrom(CodeGenContextBase* ctx, llvm::Value* src);
+
     absl::StatusOr<NativeValue> ExtractElement(CodeGenContextBase* ctx, const NativeValue& arr,
                                                const NativeValue& key) const override;
 

diff --git a/hybridse/src/codegen/struct_ir_builder.cc b/hybridse/src/codegen/struct_ir_builder.cc
@@ -16,6 +16,8 @@
 
 #include "codegen/struct_ir_builder.h"
 
+#include <utility>
+
 #include "absl/status/status.h"
 #include "absl/strings/substitute.h"
 #include "codegen/array_ir_builder.h"
@@ -299,15 +301,25 @@ absl::StatusOr<NativeValue> Combine(CodeGenContextBase* ctx, const NativeValue d
     }
     llvm::Value* input_arrays = builder->CreateAlloca(input_arr_type, builder->getInt32(args.size()), "array_data");
     node::NodeManager nm;
+    std::vector<NativeValue> casted_args(args.size());
     for (int i = 0; i < args.size(); ++i) {
         const node::TypeNode* tp = nullptr;
         if (!GetFullType(&nm, args.at(i).GetType(), &tp)) {
             return absl::InternalError("codegen error: fail to get valid type from llvm value");
         }
-        if (!tp->IsArray() || tp->GetGenericSize() != 1 || !tp->GetGenericType(0)->IsString()) {
-            return absl::InternalError("codegen error: arguments to array_combine is not ARRAY<STRING>");
+        if (!tp->IsArray() || tp->GetGenericSize() != 1) {
+            return absl::InternalError("codegen error: arguments to array_combine is not ARRAY");
+        }
+        if (!tp->GetGenericType(0)->IsString()) {
+            auto s = arr_builder.CastFrom(ctx, args.at(i).GetRaw());
+            CHECK_ABSL_STATUSOR(s);
+            casted_args.at(i) = NativeValue::Create(s.value());
+        } else {
+            casted_args.at(i) = args.at(i);
         }
-        auto safe_str_arr = builder->CreateSelect(args.at(i).GetIsNull(builder), empty_arr, args.at(i).GetRaw());
+
+        auto safe_str_arr =
+            builder->CreateSelect(casted_args.at(i).GetIsNull(builder), empty_arr, casted_args.at(i).GetRaw());
         builder->CreateStore(safe_str_arr, builder->CreateGEP(input_arr_type, input_arrays, builder->getInt32(i)));
     }
 

diff --git a/hybridse/src/udf/default_defs/array_def.cc b/hybridse/src/udf/default_defs/array_def.cc
@@ -162,8 +162,6 @@ void DefaultUdfLibrary::InitArrayUdfs() {
                 CHECK_TRUE(delimit.type()->IsString(), common::kCodegenError, "delimiter must be string");
                 for (auto & val : arg_attrs) {
                     CHECK_TRUE(val.type()->IsArray(), common::kCodegenError, "argument to array_combine must be array");
-                    CHECK_TRUE(val.type()->GetGenericType(0)->IsString(), common::kCodegenError,
-                               "argument to array_combine must be array of string");
                 }
                 auto nm = ctx->node_manager();
                 out->SetType(nm->MakeNode<node::TypeNode>(node::kArray, nm->MakeNode<node::TypeNode>(node::kVarchar)));