Apache Arrow

# PyArrow를 사용한 고성능 데이터 처리 예제

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq
import pyarrow.dataset as ds
import pyarrow.flight as flight
import pandas as pd
import numpy as np

# ==========================================
# 1. Arrow Table 생성 및 기본 연산
# ==========================================

# Python 리스트에서 Arrow Table 생성
data = {
    'user_id': [1, 2, 3, 4, 5],
    'name': ['Alice', 'Bob', 'Charlie', 'Diana', 'Eve'],
    'age': [25, 30, 35, 28, 32],
    'salary': [50000.0, 60000.0, 75000.0, 55000.0, 80000.0],
    'department': ['Engineering', 'Marketing', 'Engineering', 'Sales', 'Engineering']
}

# Arrow Table 생성
table = pa.table(data)
print(f"Schema: {table.schema}")
print(f"Rows: {table.num_rows}, Columns: {table.num_columns}")

# Arrow Compute 함수를 사용한 벡터화 연산
# 급여 10% 인상 계산
salary_increase = pc.multiply(table.column('salary'), 1.1)
print(f"Salary after 10% raise: {salary_increase}")

# 조건 필터링 (Engineering 부서)
mask = pc.equal(table.column('department'), 'Engineering')
engineering_team = pc.filter(table, mask)
print(f"Engineering team: {engineering_team.to_pandas()}")

# 집계 연산
avg_salary = pc.mean(table.column('salary'))
max_age = pc.max(table.column('age'))
print(f"Average salary: {avg_salary}, Max age: {max_age}")

# ==========================================
# 2. Pandas와의 제로카피 변환
# ==========================================

# 대용량 데이터 생성
n_rows = 10_000_000
large_data = {
    'id': np.arange(n_rows),
    'value': np.random.randn(n_rows),
    'category': np.random.choice(['A', 'B', 'C'], n_rows)
}

# Pandas DataFrame으로 생성
df = pd.DataFrame(large_data)

# Arrow Table로 변환 (제로카피 가능한 경우)
arrow_table = pa.Table.from_pandas(df, preserve_index=False)

# 다시 Pandas로 변환 (제로카피)
df_back = arrow_table.to_pandas(self_destruct=True)  # 메모리 효율적

# ==========================================
# 3. Parquet 파일 읽기/쓰기
# ==========================================

# Parquet 파일 쓰기 (압축 및 파티셔닝)
pq.write_table(
    table,
    'employees.parquet',
    compression='snappy',
    row_group_size=1000000
)

# Parquet 파일 읽기 (특정 컬럼만)
table_read = pq.read_table(
    'employees.parquet',
    columns=['name', 'salary']
)

# 파티션된 데이터셋 쓰기
pq.write_to_dataset(
    arrow_table,
    root_path='partitioned_data',
    partition_cols=['category']
)

# ==========================================
# 4. Dataset API로 대용량 데이터 처리
# ==========================================

# 여러 Parquet 파일을 하나의 Dataset으로
dataset = ds.dataset(
    'partitioned_data',
    format='parquet',
    partitioning='hive'
)

# 푸시다운 필터링과 프로젝션
scanner = dataset.scanner(
    columns=['id', 'value'],
    filter=ds.field('category') == 'A'
)

# 배치 단위로 스트리밍 처리
for batch in scanner.to_batches():
    # 각 배치를 처리 (메모리 효율적)
    result = pc.mean(batch.column('value'))
    print(f"Batch mean: {result}")

# ==========================================
# 5. Arrow Flight로 고성능 데이터 전송
# ==========================================

class FlightServer(flight.FlightServerBase):
    def __init__(self, location, data):
        super().__init__(location)
        self.data = data

    def do_get(self, context, ticket):
        """데이터 스트리밍 제공"""
        return flight.RecordBatchStream(self.data)

    def get_flight_info(self, context, descriptor):
        """데이터셋 메타데이터 제공"""
        schema = self.data.schema
        endpoints = [flight.FlightEndpoint(
            descriptor.path[0],
            [flight.Location.for_grpc_tcp("localhost", 8815)]
        )]
        return flight.FlightInfo(
            schema, descriptor, endpoints,
            self.data.num_rows, -1
        )

# 클라이언트에서 데이터 가져오기
# client = flight.connect("grpc://localhost:8815")
# reader = client.do_get(ticket)
# table = reader.read_all()