如何在Go中用Apache Arrow按时间间隔划分Kafka事件？

核心问题分析 #

RecordBuilder.UnmarshalJSON 依赖固定Go结构体解析JSON，但你的场景是用户自定义动态Schema，无法提前定义结构体。解决思路是先将单条JSON解析为Arrow Record，再从中提取datetime字段值，完成时间分区判断。

实现步骤 #

1. 扩展Segment结构体 #

新增字段跟踪当前批次的时间分区：

type Segment struct {
    mu             sync.Mutex
    schema         *arrow.Schema
    evtStruct      *arrow.StructType
    builder        *array.RecordBuilder
    writer         *pqarrow.FileWriter
    timestampIndex int
    currentPartition int64 // 记录当前所属的5分钟间隔起始时间戳
}

2. 重写InsertData函数 #

使用Arrow JSON解析器处理动态Schema，提取datetime值并判断分区：

import (
    "bytes"
    "fmt"
    "math"
    "time"

    "github.com/apache/arrow/go/v14/arrow"
    "github.com/apache/arrow/go/v14/arrow/array"
    "github.com/apache/arrow/go/v14/arrow/json"
    "github.com/apache/arrow/go/v14/arrow/memory"
)

func (s *Segment) InsertData(data []byte) error {
    s.mu.Lock()
    defer s.mu.Unlock()

    // 解析单条JSON为Arrow Record
    buf := bytes.NewBuffer(data)
    reader, err := json.NewReader(memory.NewGoAllocator(), buf, s.schema, json.WithBatchSize(1))
    if err != nil {
        return err
    }
    defer reader.Release()

    rec, err := reader.Read()
    if err != nil {
        return err
    }
    defer rec.Release()

    // 提取datetime字段值（适配不同类型）
    dtCol := rec.Column(s.timestampIndex)
    var dtVal int64
    switch dtCol.DataType().ID() {
    case arrow.Timestamp:
        tsArr := dtCol.(*array.Timestamp)
        // 根据Schema定义的时间单位转换为秒（示例：毫秒转秒）
        unit := tsArr.DataType().(*arrow.TimestampType).Unit
        dtVal = int64(tsArr.Value(0)) / int64(unit)
    case arrow.String:
        strArr := dtCol.(*array.String)
        t, err := time.Parse(time.RFC3339, strArr.Value(0))
        if err != nil {
            return fmt.Errorf("parse datetime string failed: %w", err)
        }
        dtVal = t.Unix()
    default:
        return fmt.Errorf("unsupported datetime type: %s", dtCol.DataType().Name())
    }

    // 计算5分钟间隔的分区起始时间戳
    interval := int64(5 * 60)
    dtPartition := dtVal - (dtVal % interval)

    // 判断是否需要刷写当前批次
    if s.currentPartition != 0 && dtPartition > s.currentPartition {
        if err := s.Flush(); err != nil {
            return err
        }
        // 重置Builder准备新批次
        s.builder.Release()
        mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
        s.builder = array.NewRecordBuilder(mem, s.schema)
    }

    // 将记录追加到Builder
    if err := s.builder.AppendRecord(rec); err != nil {
        return err
    }

    // 初始化或更新当前分区
    if s.currentPartition == 0 {
        s.currentPartition = dtPartition
    }

    return nil
}

3. 修复Flush函数 #

原函数错误关闭了Writer和Builder，调整为仅重置Builder状态：

func (s *Segment) Flush() error {
    s.mu.Lock()
    defer s.mu.Unlock()

    rec := s.builder.NewRecord()
    defer rec.Release()

    if err := s.writer.WriteBuffered(rec); err != nil {
        return err
    }

    // 重置Builder，而非释放
    s.builder.Reset()
    return nil
}

关键说明 #

• 用arrow/json.NewReader实现动态Schema的JSON解析，无需提前定义Go结构体。
• 根据Schema中datetime字段的实际类型（Timestamp/String）做对应转换，确保时间戳计算准确。
• Flush时仅重置Builder，保留Writer以便后续批次写入。

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