Apache Arrow for Go を使用してparquetファイルを出力してみた

package main

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	"github.com/apache/arrow-go/v18/arrow/memory"
	"github.com/apache/arrow-go/v18/parquet"
	"github.com/apache/arrow-go/v18/parquet/compress"
	"github.com/apache/arrow-go/v18/parquet/pqarrow"
)

func main() {
	fmt.Println("=== Arrow-Go Sample Program ===")

	// Parquetファイルを作成
	filename := "sample.parquet"
	if err := writeParquetFile(filename); err != nil {
		log.Fatalf("Failed to write parquet file: %v", err)
	}
	fmt.Printf("✓ Parquetファイルを作成しました: %s\n\n", filename)
}

// writeParquetFile はサンプルデータをParquetファイルに書き込む
func writeParquetFile(filename string) error {
	// 既存ファイルが存在する場合は削除
	if _, err := os.Stat(filename); err == nil {
		if err := os.Remove(filename); err != nil {
			return fmt.Errorf("failed to remove existing file: %w", err)
		}
		fmt.Printf("既存のファイルを削除しました: %s\n", filename)
	}

	// メモリアロケータを作成
	allocator := memory.NewGoAllocator()

	// スキーマを定義（id, name, age, score, is_active, created_at）
	schema := arrow.NewSchema(
		[]arrow.Field{
			{Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false},
			{Name: "name", Type: arrow.BinaryTypes.String, Nullable: true},
			{Name: "age", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
			{Name: "score", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
			{Name: "is_active", Type: arrow.FixedWidthTypes.Boolean, Nullable: true},
			{Name: "created_at", Type: &arrow.TimestampType{Unit: arrow.Millisecond, TimeZone: "UTC"}, Nullable: true},
		},
		nil,
	)

	// ファイルを作成
	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("failed to create file: %w", err)
	}
	defer file.Close()

	// Parquetライターのプロパティを設定
	writerProps := parquet.NewWriterProperties(
		parquet.WithCompression(compress.Codecs.Snappy), // Snappy圧縮を使用
		parquet.WithDictionaryDefault(true),             // ディクショナリエンコーディングを有効化
	)

	// Arrowライタープロパティを設定
	arrowProps := pqarrow.NewArrowWriterProperties(
		pqarrow.WithStoreSchema(), // スキーマを保存
	)

	// Parquet Writerを作成
	writer, err := pqarrow.NewFileWriter(schema, file, writerProps, arrowProps)
	if err != nil {
		return fmt.Errorf("failed to create parquet writer: %w", err)
	}
	defer writer.Close()

	// サンプルデータを作成
	records := []arrow.Record{
		createSampleRecord(schema, allocator, []recordData{
			{id: 1, name: "Alice", age: ptrInt64(30), score: ptrFloat64(95.5), isActive: ptrBool(true), createdAt: time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)},
			{id: 2, name: "Bob", age: ptrInt64(25), score: ptrFloat64(87.3), isActive: ptrBool(true), createdAt: time.Date(2024, 1, 2, 11, 0, 0, 0, time.UTC)},
			{id: 3, name: "Charlie", age: nil, score: nil, isActive: ptrBool(false), createdAt: time.Date(2024, 1, 3, 12, 0, 0, 0, time.UTC)},
			{id: 4, name: "Diana", age: ptrInt64(28), score: ptrFloat64(92.8), isActive: nil, createdAt: time.Date(2024, 1, 4, 13, 0, 0, 0, time.UTC)},
		}),
	}

	// レコードをParquetファイルに書き込み
	for _, record := range records {
		if err := writer.Write(record); err != nil {
			record.Release()
			return fmt.Errorf("failed to write record: %w", err)
		}
		record.Release()
	}

	return nil
}

// recordData はサンプルデータの構造
type recordData struct {
	id        int64
	name      string
	age       *int64
	score     *float64
	isActive  *bool
	createdAt time.Time
}

// createSampleRecord はサンプルデータからArrow Recordを作成
func createSampleRecord(schema *arrow.Schema, allocator memory.Allocator, data []recordData) arrow.Record {
	builder := array.NewRecordBuilder(allocator, schema)
	defer builder.Release()

	// id列（Int64）
	idBuilder := builder.Field(0).(*array.Int64Builder)
	for _, d := range data {
		idBuilder.Append(d.id)
	}

	// name列（String）
	nameBuilder := builder.Field(1).(*array.StringBuilder)
	for _, d := range data {
		if d.name == "" {
			nameBuilder.AppendNull()
		} else {
			nameBuilder.Append(d.name)
		}
	}

	// age列（Int64, Nullable）
	ageBuilder := builder.Field(2).(*array.Int64Builder)
	for _, d := range data {
		if d.age == nil {
			ageBuilder.AppendNull()
		} else {
			ageBuilder.Append(*d.age)
		}
	}

	// score列（Float64, Nullable）
	scoreBuilder := builder.Field(3).(*array.Float64Builder)
	for _, d := range data {
		if d.score == nil {
			scoreBuilder.AppendNull()
		} else {
			scoreBuilder.Append(*d.score)
		}
	}

	// is_active列（Boolean, Nullable）
	isActiveBuilder := builder.Field(4).(*array.BooleanBuilder)
	for _, d := range data {
		if d.isActive == nil {
			isActiveBuilder.AppendNull()
		} else {
			isActiveBuilder.Append(*d.isActive)
		}
	}

	// created_at列（Timestamp）
	createdAtBuilder := builder.Field(5).(*array.TimestampBuilder)
	for _, d := range data {
		tsMillis := arrow.Timestamp(d.createdAt.UnixMilli())
		createdAtBuilder.Append(tsMillis)
	}

	return builder.NewRecord()
}

// ptrInt64 はint64のポインタを返すヘルパー関数
func ptrInt64(v int64) *int64 {
	return &v
}

// ptrFloat64 はfloat64のポインタを返すヘルパー関数
func ptrFloat64(v float64) *float64 {
	return &v
}

// ptrBool はboolのポインタを返すヘルパー関数
func ptrBool(v bool) *bool {
	return &v
}

$ go run arrow-go-sample.go

$ duckdb

D SELECT * FROM 'sample.parquet';
┌───────┬─────────┬───────┬────────┬───────────┬──────────────────────────┐
│  id   │  name   │  age  │ score  │ is_active │        created_at        │
│ int64 │ varchar │ int64 │ double │  boolean  │ timestamp with time zone │
├───────┼─────────┼───────┼────────┼───────────┼──────────────────────────┤
│     1 │ Alice   │    30 │   95.5 │ true      │ 2024-01-01 19:00:00+09   │
│     2 │ Bob     │    25 │   87.3 │ true      │ 2024-01-02 20:00:00+09   │
│     3 │ Charlie │  NULL │   NULL │ false     │ 2024-01-03 21:00:00+09   │
│     4 │ Diana   │    28 │   92.8 │ NULL      │ 2024-01-04 22:00:00+09   │
└───────┴─────────┴───────┴────────┴───────────┴──────────────────────────┘