be/src/exprs/table_function/python_udtf_function.cpp

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 "exprs/table_function/python_udtf_function.h"

#include <arrow/array.h>
#include <arrow/array/array_nested.h>
#include <arrow/record_batch.h>
#include <arrow/type_fwd.h>
#include <glog/logging.h>

#include "core/assert_cast.h"
#include "core/block/block.h"
#include "core/block/column_numbers.h"
#include "core/column/column.h"
#include "core/column/column_array.h"
#include "core/column/column_nullable.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_factory.hpp"
#include "core/data_type_serde/data_type_array_serde.h"
#include "exprs/function/array/function_array_utils.h"
#include "exprs/vexpr.h"
#include "exprs/vexpr_context.h"
#include "format/arrow/arrow_block_convertor.h"
#include "format/arrow/arrow_row_batch.h"
#include "format/arrow/arrow_utils.h"
#include "runtime/runtime_state.h"
#include "runtime/user_function_cache.h"
#include "udf/python/python_env.h"
#include "udf/python/python_server.h"
#include "udf/python/python_udf_meta.h"
#include "util/timezone_utils.h"

namespace doris {
#include "common/compile_check_begin.h"

PythonUDTFFunction::PythonUDTFFunction(const TFunction& t_fn) : TableFunction(), _t_fn(t_fn) {
    _fn_name = _t_fn.name.function_name;
    TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, _timezone_obj);

    // Like Java UDTF, FE passes the element type T, and we wrap it into array<T> here
    // This makes the behavior consistent with Java UDTF
    DataTypePtr element_type = DataTypeFactory::instance().create_data_type(t_fn.ret_type);
    _return_type = make_nullable(std::make_shared<DataTypeArray>(make_nullable(element_type)));
}

Status PythonUDTFFunction::open() {
    PythonUDFMeta python_udf_meta;
    python_udf_meta.id = _t_fn.id;
    python_udf_meta.name = _t_fn.name.function_name;
    python_udf_meta.symbol = _t_fn.scalar_fn.symbol;

    if (!_t_fn.function_code.empty()) {
        python_udf_meta.type = PythonUDFLoadType::INLINE;
        python_udf_meta.location = "inline";
        python_udf_meta.inline_code = _t_fn.function_code;
    } else if (!_t_fn.hdfs_location.empty()) {
        python_udf_meta.type = PythonUDFLoadType::MODULE;
        python_udf_meta.location = _t_fn.hdfs_location;
        python_udf_meta.checksum = _t_fn.checksum;
    } else {
        python_udf_meta.type = PythonUDFLoadType::UNKNOWN;
        python_udf_meta.location = "unknown";
    }

    python_udf_meta.client_type = PythonClientType::UDTF;

    if (python_udf_meta.type == PythonUDFLoadType::MODULE) {
        RETURN_IF_ERROR(UserFunctionCache::instance()->get_pypath(
                python_udf_meta.id, python_udf_meta.location, python_udf_meta.checksum,
                &python_udf_meta.location));
    }

    PythonVersion version;
    if (_t_fn.__isset.runtime_version && !_t_fn.runtime_version.empty()) {
        RETURN_IF_ERROR(
                PythonVersionManager::instance().get_version(_t_fn.runtime_version, &version));
        python_udf_meta.runtime_version = version.full_version;
    } else {
        return Status::InvalidArgument("Python UDTF runtime version is not set");
    }

    for (const auto& arg_type : _t_fn.arg_types) {
        DataTypePtr doris_type = DataTypeFactory::instance().create_data_type(arg_type);
        python_udf_meta.input_types.push_back(doris_type);
    }

    // For Python UDTF, FE passes the element type T (like Java UDTF)
    // Use it directly as the UDF's return type for Python metadata
    python_udf_meta.return_type = DataTypeFactory::instance().create_data_type(_t_fn.ret_type);
    python_udf_meta.always_nullable = python_udf_meta.return_type->is_nullable();
    RETURN_IF_ERROR(python_udf_meta.check());

    RETURN_IF_ERROR(
            PythonServerManager::instance().get_client(python_udf_meta, version, &_udtf_client));

    if (!_udtf_client) {
        return Status::InternalError("Failed to create Python UDTF client");
    }

    return Status::OK();
}

Status PythonUDTFFunction::process_init(Block* block, RuntimeState* state) {
    // Step 1: Extract input columns from child expressions
    auto child_size = _expr_context->root()->children().size();
    ColumnNumbers child_column_idxs;
    child_column_idxs.resize(child_size);
    for (int i = 0; i < child_size; ++i) {
        int result_id = -1;
        RETURN_IF_ERROR(_expr_context->root()->children()[i]->execute(_expr_context.get(), block,
                                                                      &result_id));
        DCHECK_NE(result_id, -1);
        child_column_idxs[i] = result_id;
    }

    // Step 2: Build input block and convert to Arrow format
    Block input_block;
    for (uint32_t i = 0; i < child_column_idxs.size(); ++i) {
        input_block.insert(block->get_by_position(child_column_idxs[i]));
    }
    std::shared_ptr<arrow::Schema> input_schema;
    std::shared_ptr<arrow::RecordBatch> input_batch;
    RETURN_IF_ERROR(get_arrow_schema_from_block(input_block, &input_schema,
                                                TimezoneUtils::default_time_zone));
    RETURN_IF_ERROR(convert_to_arrow_batch(input_block, input_schema, arrow::default_memory_pool(),
                                           &input_batch, _timezone_obj));

    // Step 3: Call Python UDTF to evaluate all rows at once (similar to Java UDTF's JNI call)
    // Python returns a ListArray where each element contains outputs for one input row
    std::shared_ptr<arrow::ListArray> list_array;
    RETURN_IF_ERROR(_udtf_client->evaluate(*input_batch, &list_array));

    // Step 4: Convert Python server output (ListArray) to Doris array column
    RETURN_IF_ERROR(_convert_list_array_to_array_column(list_array));

    // Step 5: Extract array column metadata using extract_column_array_info
    if (!extract_column_array_info(*_array_result_column, _array_column_detail)) {
        return Status::NotSupported("column type {} not supported now",
                                    _array_result_column->get_name());
    }

    return Status::OK();
}

void PythonUDTFFunction::process_row(size_t row_idx) {
    TableFunction::process_row(row_idx);

    // Check if array is null for this row
    if (!_array_column_detail.array_nullmap_data ||
        !_array_column_detail.array_nullmap_data[row_idx]) {
        _array_offset = (*_array_column_detail.offsets_ptr)[row_idx - 1];
        _cur_size = (*_array_column_detail.offsets_ptr)[row_idx] - _array_offset;
    }
    // When it's NULL at row_idx, _cur_size stays 0, meaning current_empty()
    // If outer function: will continue with insert_default
    // If not outer function: will not insert any value
}

void PythonUDTFFunction::process_close() {
    _array_result_column = nullptr;
    _array_column_detail.reset();
    _array_offset = 0;
}

void PythonUDTFFunction::get_same_many_values(MutableColumnPtr& column, int length) {
    size_t pos = _array_offset + _cur_offset;
    if (current_empty() || (_array_column_detail.nested_nullmap_data &&
                            _array_column_detail.nested_nullmap_data[pos])) {
        column->insert_many_defaults(length);
    } else {
        if (_is_nullable) {
            auto* nullable_column = assert_cast<ColumnNullable*>(column.get());
            auto nested_column = nullable_column->get_nested_column_ptr();
            auto nullmap_column = nullable_column->get_null_map_column_ptr();
            nested_column->insert_many_from(*_array_column_detail.nested_col, pos, length);
            assert_cast<ColumnUInt8*>(nullmap_column.get())->insert_many_defaults(length);
        } else {
            column->insert_many_from(*_array_column_detail.nested_col, pos, length);
        }
    }
}

int PythonUDTFFunction::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;
}

Status PythonUDTFFunction::close() {
    // Close UDTF client
    if (_udtf_client) {
        Status status = _udtf_client->close();
        if (!status.ok()) {
            LOG(WARNING) << "Failed to close UDTF client: " << status.to_string();
        }
        _udtf_client.reset();
    }

    return TableFunction::close();
}

Status PythonUDTFFunction::_convert_list_array_to_array_column(
        const std::shared_ptr<arrow::ListArray>& list_array) {
    if (!list_array) {
        return Status::InternalError("Received null ListArray from Python UDTF");
    }

    size_t num_input_rows = list_array->length();

    // Handle nullable array column
    MutableColumnPtr array_col_ptr = _return_type->create_column();
    ColumnNullable* nullable_col = nullptr;
    ColumnArray* array_col = nullptr;

    if (_return_type->is_nullable()) {
        nullable_col = assert_cast<ColumnNullable*>(array_col_ptr.get());
        array_col = assert_cast<ColumnArray*>(
                nullable_col->get_nested_column_ptr()->assume_mutable().get());
    } else {
        array_col = assert_cast<ColumnArray*>(array_col_ptr.get());
    }

    // Create DataTypeArraySerDe for direct Arrow conversion
    DataTypePtr element_type = DataTypeFactory::instance().create_data_type(_t_fn.ret_type);
    DataTypePtr array_type = std::make_shared<DataTypeArray>(make_nullable(element_type));
    auto array_serde = array_type->get_serde();

    // Use read_column_from_arrow for optimized conversion
    // This directly converts Arrow ListArray to Doris ColumnArray
    // No struct unwrapping needed - Python server sends the correct format!
    RETURN_IF_ERROR(array_serde->read_column_from_arrow(
            array_col->assume_mutable_ref(), list_array.get(), 0, num_input_rows, _timezone_obj));

    // Handle nullable wrapper: all array elements are non-null
    // (empty arrays [] are non-null, different from NULL)
    if (nullable_col) {
        auto& null_map = nullable_col->get_null_map_data();
        null_map.resize_fill(num_input_rows, 0); // All non-null
    }

    _array_result_column = std::move(array_col_ptr);
    return Status::OK();
}

#include "common/compile_check_end.h"
} // namespace doris