From 43496cce48500510628cb02188df8d34de5b3c02 Mon Sep 17 00:00:00 2001
From: zhangstar333 <2561612514@qq.com>
Date: Fri, 19 Apr 2024 14:38:49 +0800
Subject: [PATCH] override get_value function

---
 .../table_function/udf_table_function.cpp     | 27 +++++++++++++++++++
 .../exprs/table_function/udf_table_function.h |  1 +
 2 files changed, 28 insertions(+)

diff --git a/be/src/vec/exprs/table_function/udf_table_function.cpp b/be/src/vec/exprs/table_function/udf_table_function.cpp
index 611d6e8f0b6f6f..bc4c815ceb15cd 100644
--- a/be/src/vec/exprs/table_function/udf_table_function.cpp
+++ b/be/src/vec/exprs/table_function/udf_table_function.cpp
@@ -24,6 +24,7 @@
 #include "vec/columns/column_nullable.h"
 #include "vec/common/assert_cast.h"
 #include "vec/core/block.h"
+#include "vec/core/types.h"
 #include "vec/data_types/data_type_array.h"
 #include "vec/data_types/data_type_factory.hpp"
 #include "vec/exec/jni_connector.h"
@@ -169,4 +170,30 @@ void UDFTableFunction::get_value(MutableColumnPtr& column) {
         }
     }
 }
+
+int UDFTableFunction::get_value(MutableColumnPtr& column, int max_step) {
+    max_step = std::min(max_step, (int)(_cur_size - _cur_offset));
+    size_t pos = _array_offset + _cur_offset;
+    if (current_empty()) {
+        column->insert_default();
+        max_step = 1;
+    } else {
+        if (_is_nullable) {
+            auto* nullable_column = assert_cast<ColumnNullable*>(column.get());
+            auto nested_column = nullable_column->get_nested_column_ptr();
+            auto* nullmap_column =
+                    assert_cast<ColumnUInt8*>(nullable_column->get_null_map_column_ptr().get());
+            nested_column->insert_range_from(*_array_column_detail.nested_col, pos, max_step);
+            size_t old_size = nullmap_column->size();
+            nullmap_column->resize(old_size + max_step);
+            memcpy(nullmap_column->get_data().data() + old_size,
+                   _array_column_detail.nested_nullmap_data + pos * sizeof(UInt8),
+                   max_step * sizeof(UInt8));
+        } else {
+            column->insert_range_from(*_array_column_detail.nested_col, pos, max_step);
+        }
+    }
+    forward(max_step);
+    return max_step;
+}
 } // namespace doris::vectorized
diff --git a/be/src/vec/exprs/table_function/udf_table_function.h b/be/src/vec/exprs/table_function/udf_table_function.h
index b268a39845913a..ae6a7c13b355a9 100644
--- a/be/src/vec/exprs/table_function/udf_table_function.h
+++ b/be/src/vec/exprs/table_function/udf_table_function.h
@@ -39,6 +39,7 @@ class UDFTableFunction final : public TableFunction {
     void process_row(size_t row_idx) override;
     void process_close() override;
     void get_value(MutableColumnPtr& column) override;
+    int get_value(MutableColumnPtr& column, int max_step) override;
     Status close() override {
         if (_jni_ctx) {
             RETURN_IF_ERROR(_jni_ctx->close());