// 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.

#pragma once

#include <arrow/record_batch.h>

#include <cctz/time_zone.h>

#include "common/status.h"

#include "core/column/column.h"

#include "core/data_type/data_type.h"

#include "exprs/function/array/function_array_utils.h"

#include "exprs/table_function/table_function.h"

#include "udf/python/python_udtf_client.h"

namespace doris {

#include "common/compile_check_begin.h"

/**

 * PythonUDTFFunction - Python User-Defined Table Function

 * 

 * Execution Flow:

 * 1. open() - Create Python UDTF client and establish RPC connection

 * 2. process_init(block) - Batch evaluate all rows:

 *    - Convert input block to Arrow RecordBatch

 *    - Call Python UDTF server via RPC (evaluates all rows in one call)

 *    - Receive Arrow ListArray (one list per input row)

 *    - Convert to array<T> column using DataTypeArraySerDe

 * 3. process_row(row_idx) - Set array offset for current row

 * 4. get_value()/get_same_many_values() - Extract values from array column

 * 5. process_close() - Clean up batch state (array column, offsets)

 * 6. close() - Close Python UDTF client and RPC connection

 */

class PythonUDTFFunction final : public TableFunction {

    ENABLE_FACTORY_CREATOR(PythonUDTFFunction);

public:

    PythonUDTFFunction(const TFunction& t_fn);

    ~PythonUDTFFunction() override = default;

    Status open() override;

    Status process_init(Block* block, RuntimeState* state) override;

    void process_row(size_t row_idx) override;

    void process_close() override;

    void get_same_many_values(MutableColumnPtr& column, int length) override;

    int get_value(MutableColumnPtr& column, int max_step) override;

    Status close() override;

private:

    /**

     * Convert Python UDTF output (Arrow ListArray) to Doris array column

     * 

     * Input from Python server (via Arrow RPC):

     * - list_array: Arrow ListArray where each element corresponds to one input row's outputs

     *

     * Format:

     * - Single-column output: List<T> (e.g., List<int>, List<string>)

     * - Multi-column output: List<Struct<col1, col2, ...>>

     *

     * Example: 3 input rows producing variable output rows

     *   ListArray structure:

     *     [0]: [val1, val2, val3]       (3 elements)

     *     [1]: []                       (0 elements - empty array)

     *     [2]: [val4, val5, val6, val7] (4 elements)

     *

     * @param list_array Arrow ListArray containing UDTF output (length = num_input_rows)

     * @return Status indicating success or validation/conversion errors

     */

    Status _convert_list_array_to_array_column(const std::shared_ptr<arrow::ListArray>& list_array);

    const TFunction& _t_fn;

    DataTypePtr _return_type;

    PythonUDTFClientPtr _udtf_client;

    cctz::time_zone _timezone_obj;

    // Result storage (similar to Java UDTF)

    ColumnPtr _array_result_column;                // Array column storing all results

    ColumnArrayExecutionData _array_column_detail; // Array metadata for efficient access

    int64_t _array_offset = 0;                     // Offset into array for current row

};

#include "common/compile_check_end.h"

} // namespace doris