Apache NiFi ExecuteStreamCommand执行Python脚本报错求助

问题场景 #

在Apache NiFi中通过ExecuteStreamCommand处理器运行Python脚本，读取ListFile+FetchFile传递的Excel文件，意图合并并分割为指定大小的Parquet文件，但执行时报错：

2024-11-25 07:46:07,540 ERROR [Timer-Driven Process Thread-5] o.a.n.p.standard.ExecuteStreamCommand ExecuteStreamCommand[id=4e7bd5aa-0193-1000-6154-675d13b6c0e5] Transferring StandardFlowFileRecord[uuid=c534f453-59d0-4958-819e-01295652d375,claim=StandardContentClaim [resourceClaim=StandardResourceClaim[id=1732502757533-52, container=default, section=52], offset=1680, length=112],offset=0,name=3GNANKANA.xlsx,size=112] to nonzero status. Executable command python3 ended in an error:

原脚本代码：

import os
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import sys
import io


def merge_and_split_parquet(data, output_dir, max_size=1 * 1024**3):
    """
    Merges and splits data into Parquet files based on size.

    Args:
        data (pd.DataFrame): Data to process.
        output_dir (str): Directory to save Parquet files.
        max_size (int): Maximum size in bytes for each Parquet file.
    """
    merged_data = pd.DataFrame()
    parquet_index = 1

    # Process the data row by row
    for _, row in data.iterrows():
        merged_data = pd.concat([merged_data, pd.DataFrame([row])])

        # Check the size of the DataFrame
        table = pa.Table.from_pandas(merged_data)
        buffer = pa.BufferOutputStream()
        pq.write_table(table, buffer)
        if buffer.size() >= max_size:
            # Write to a new Parquet file
            output_file = os.path.join(output_dir, f"merged_{parquet_index}.parquet")
            pq.write_table(table, output_file)
            print(f"Generated {output_file}")
            merged_data = pd.DataFrame()  # Reset the DataFrame
            parquet_index += 1

    # Write any remaining data
    if not merged_data.empty:
        output_file = os.path.join(output_dir, f"merged_{parquet_index}.parquet")
        pq.write_table(pa.Table.from_pandas(merged_data), output_file)
        print(f"Generated {output_file}")


if __name__ == "__main__":
    # Ensure proper arguments are passed
    if len(sys.argv) < 2:
        print("Usage: python script.py <output_dir>")
        sys.exit(1)

    output_dir = sys.argv[1]

    # Ensure the output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Read data from STDIN
    try:
        input_stream = sys.stdin.read()
        excel_data = pd.read_excel(io.BytesIO(input_stream.encode()))
    except Exception as e:
        print(f"Error reading input data: {e}")
        sys.exit(1)

    # Process and write Parquet files
    merge_and_split_parquet(excel_data, output_dir)

排查与解决步骤 #

1. 修复脚本核心错误 #

原脚本读取STDIN的逻辑存在致命问题：Excel是二进制文件，sys.stdin.read()获取的是字节流，调用encode()会破坏二进制结构，导致pd.read_excel解析失败。同时逐行concat的方式性能极低，一并优化。

修改后的脚本：

import os
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import sys
import io


def merge_and_split_parquet(data, output_dir, max_size=1 * 1024**3):
    """
    Merges and splits data into Parquet files based on size.

    Args:
        data (pd.DataFrame): Data to process.
        output_dir (str): Directory to save Parquet files.
        max_size (int): Maximum size in bytes for each Parquet file.
    """
    batch_rows = []
    parquet_index = 1

    # 改用列表收集行，批量合并提升性能
    for _, row in data.iterrows():
        batch_rows.append(row)
        # 计算当前批次的Parquet文件大小
        temp_df = pd.DataFrame(batch_rows)
        temp_table = pa.Table.from_pandas(temp_df)
        buffer = pa.BufferOutputStream()
        pq.write_table(temp_table, buffer)
        
        if buffer.size() >= max_size:
            # 写入文件并重置批次
            output_file = os.path.join(output_dir, f"merged_{parquet_index}.parquet")
            pq.write_table(temp_table, output_file)
            print(f"Generated {output_file}")
            batch_rows = []
            parquet_index += 1

    # 处理剩余未达大小阈值的数据
    if batch_rows:
        final_df = pd.DataFrame(batch_rows)
        final_table = pa.Table.from_pandas(final_df)
        output_file = os.path.join(output_dir, f"merged_{parquet_index}.parquet")
        pq.write_table(final_table, output_file)
        print(f"Generated {output_file}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python script.py <output_dir>", file=sys.stderr)
        sys.exit(1)

    output_dir = sys.argv[1]
    os.makedirs(output_dir, exist_ok=True)

    # 修正二进制读取逻辑：直接读取原始字节流
    try:
        input_stream = sys.stdin.buffer.read()
        excel_data = pd.read_excel(io.BytesIO(input_stream))
    except Exception as e:
        print(f"Error reading input data: {str(e)}", file=sys.stderr)
        sys.exit(1)

    merge_and_split_parquet(excel_data, output_dir)

关键修改点：

• 读取STDIN时使用sys.stdin.buffer.read()获取原始二进制字节流，取消错误的encode()调用
• 优化数据拼接逻辑，用列表收集行后批量合并，避免逐行concat的性能损耗
• 将错误信息定向到sys.stderr，确保NiFi能捕获到详细报错内容

2. 验证ExecuteStreamCommand配置 #

• Command：填写python3的完整绝对路径（例如/usr/bin/python3），避免环境变量查找失败
• Command Arguments：确保传递的输出目录是NiFi进程拥有写入权限的路径，可使用绝对路径或NiFi属性（如${output.dir}）
• Redirect Error Stream to Output：勾选此选项，让Python的错误信息输出到NiFi日志，便于排查
• Working Directory：设置为脚本所在目录，避免文件路径解析错误

3. 依赖环境检查 #

确认NiFi运行环境的Python3已安装所有必需依赖：

pip3 install pandas pyarrow openpyxl

（注：openpyxl用于读取xlsx格式文件，若处理xls格式需安装xlrd）

4. 独立测试脚本 #

先脱离NiFi环境，在命令行直接测试脚本功能：

python3 your_script.py /path/to/output < test.xlsx

若命令行运行报错，根据终端输出的错误信息修复后，再部署到NiFi中运行。

内容的提问来源于stack exchange，提问作者Filbadeha