from pyspark.sql import SparkSession
from pyspark.sql.types import *

spark = SparkSession.builder \
    .appName("Parquet Example") \
    .getOrCreate()

# 데이터 생성
data = [
    (1, "Alice", 100.5, "2024-01-15"),
    (2, "Bob", 200.3, "2024-01-15"),
    (3, "Charlie", 150.7, "2024-01-15")
]

schema = StructType([
    StructField("id", IntegerType(), False),
    StructField("name", StringType(), False),
    StructField("amount", DoubleType(), False),
    StructField("date", StringType(), False)
])

df = spark.createDataFrame(data, schema)

# Parquet 설정 최적화
spark.conf.set("spark.sql.parquet.compression.codec", "snappy")
spark.conf.set("spark.sql.parquet.block.size", 134217728)  # 128MB
spark.conf.set("spark.sql.parquet.page.size", 1048576)     # 1MB

# 저장
df.write \
    .mode("overwrite") \
    .parquet("s3://bucket/data/events.parquet")

# Parquet 파일 메타데이터 읽기
import pyarrow.parquet as pq

parquet_file = pq.ParquetFile('events.parquet')

# 스키마 확인
print("Schema:")
print(parquet_file.schema)

# Row Group 정보
print(f"\nRow Groups: {parquet_file.num_row_groups}")

# Row Group별 통계
for i in range(parquet_file.num_row_groups):
    rg = parquet_file.metadata.row_group(i)
    print(f"\nRow Group {i}:")
    print(f"  Rows: {rg.num_rows}")
    print(f"  Total Size: {rg.total_byte_size / 1024 / 1024:.2f} MB")
    
    # 컬럼별 통계
    for j in range(rg.num_columns):
        col = rg.column(j)
        print(f"  Column {col.path_in_schema}:")
        print(f"    Compressed: {col.total_compressed_size / 1024:.2f} KB")
        print(f"    Uncompressed: {col.total_uncompressed_size / 1024:.2f} KB")
        print(f"    Compression Ratio: {col.total_uncompressed_size / col.total_compressed_size:.2f}x")

# Spark에서 ORC 생성
df.write \
    .format("orc") \
    .option("compression", "zlib") \
    .option("orc.stripe.size", 67108864) \
    .option("orc.compress.size", 262144) \
    .option("orc.bloom.filter.columns", "user_id,product_id") \
    .mode("overwrite") \
    .save("s3://bucket/data/events.orc")

# ORC 파일 분석 (PyArrow 사용)
import pyarrow.orc as orc

orc_file = orc.ORCFile('events.orc')

# 스키마
print("Schema:")
print(orc_file.schema)

# Stripe 정보
print(f"\nStripes: {orc_file.nstripes}")
print(f"Rows: {orc_file.nrows}")

# 메타데이터
metadata = orc_file.metadata
print(f"Compression: {metadata.compression}")
print(f"Writer Version: {metadata.writer_version}")

# Spark에서 Avro 생성
df.write \
    .format("avro") \
    .option("compression", "snappy") \
    .mode("overwrite") \
    .save("s3://bucket/data/events.avro")

# Kafka에서 Avro 사용
from confluent_kafka import avro
from confluent_kafka.avro import AvroProducer

value_schema_str = """
{
   "namespace": "com.example",
   "type": "record",
   "name": "Event",
   "fields" : [
     {"name": "id", "type": "int"},
     {"name": "name", "type": "string"}
   ]
}
"""

value_schema = avro.loads(value_schema_str)

avroProducer = AvroProducer({
    'bootstrap.servers': 'localhost:9092',
    'schema.registry.url': 'http://localhost:8081'
}, default_value_schema=value_schema)

# 메시지 전송
avroProducer.produce(topic='events', value={"id": 1, "name": "Alice"})
avroProducer.flush()

import time

# Parquet 쓰기
start = time.time()
df.write.mode("overwrite").parquet("output.parquet")
parquet_time = time.time() - start

# ORC 쓰기
start = time.time()
df.write.format("orc").mode("overwrite").save("output.orc")
orc_time = time.time() - start

# Avro 쓰기
start = time.time()
df.write.format("avro").mode("overwrite").save("output.avro")
avro_time = time.time() - start

print(f"Parquet: {parquet_time:.2f}s")  # 결과: 142.3s
print(f"ORC: {orc_time:.2f}s")          # 결과: 156.8s
print(f"Avro: {avro_time:.2f}s")        # 결과: 98.4s

# 스키마 변경 테스트
# 1. 기존 스키마로 데이터 저장
schema_v1 = StructType([
    StructField("id", IntegerType()),
    StructField("name", StringType()),
    StructField("amount", DoubleType())
])

df_v1.write.format(format_type).save(f"data_{format_type}_v1")

# 2. 새 컬럼 추가된 스키마
schema_v2 = StructType([
    StructField("id", IntegerType()),
    StructField("name", StringType()),
    StructField("amount", DoubleType()),
    StructField("category", StringType())  # 새 컬럼
])

df_v2.write.format(format_type).save(f"data_{format_type}_v2")

# 3. 두 버전 동시 읽기
df_merged = spark.read.format(format_type).load(f"data_{format_type}_*")

# 최적 설정
spark.conf.set("spark.sql.parquet.compression.codec", "snappy")
spark.conf.set("spark.sql.parquet.block.size", 134217728)
spark.conf.set("spark.sql.parquet.page.size", 1048576)
spark.conf.set("spark.sql.parquet.enableVectorizedReader", "true")

df.write \
    .partitionBy("date") \
    .parquet("s3://bucket/analytics/")

# 실시간: Kafka + Avro
from confluent_kafka import avro

# Avro로 Kafka에 저장
avro_producer.produce(topic='events', value=event_data)

# 배치: Avro → Parquet 변환
df = spark.read.format("avro").load("s3://bucket/streaming/avro/")

df.write \
    .partitionBy("date") \
    .parquet("s3://bucket/analytics/parquet/")

# 최대 압축률 설정
spark.conf.set("spark.sql.parquet.compression.codec", "gzip")  # 또는 zstd

df.write \
    .partitionBy("date") \
    .parquet("s3://bucket/logs/")

# Lifecycle policy로 자동 전환
import boto3

s3 = boto3.client('s3')
s3.put_bucket_lifecycle_configuration(
    Bucket='bucket',
    LifecycleConfiguration={
        'Rules': [{
            'Id': 'TransitionLogs',
            'Status': 'Enabled',
            'Prefix': 'logs/',
            'Transitions': [
                {'Days': 30, 'StorageClass': 'STANDARD_IA'},
                {'Days': 90, 'StorageClass': 'GLACIER'}
            ]
        }]
    }
)

# 중첩 JSON 데이터
json_data = """
{
  "user": {
    "id": 123,
    "profile": {
      "name": "Alice",
      "email": "alice@example.com"
    }
  },
  "event": {
    "type": "purchase",
    "items": [
      {"id": 1, "price": 100.5},
      {"id": 2, "price": 50.3}
    ]
  }
}
"""

# Spark에서 중첩 구조 처리
df = spark.read.json("s3://bucket/raw/events.json")

# Parquet로 저장 (중첩 구조 유지)
df.write.parquet("s3://bucket/processed/events.parquet")

# 특정 필드만 효율적으로 읽기
df = spark.read.parquet("s3://bucket/processed/events.parquet")
df.select("user.profile.name", "event.type").show()
# Parquet는 필요한 컬럼만 읽음 (nested column pruning)

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("CSV to Parquet") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

# CSV 읽기 (스키마 추론)
df = spark.read \
    .option("header", "true") \
    .option("inferSchema", "true") \
    .csv("s3://bucket/raw/csv/*.csv")

# 데이터 타입 최적화
from pyspark.sql.functions import col

df = df \
    .withColumn("amount", col("amount").cast("decimal(10,2)")) \
    .withColumn("event_time", col("event_time").cast("timestamp"))

# Parquet로 변환
df.repartition(100) \
    .write \
    .mode("overwrite") \
    .option("compression", "snappy") \
    .partitionBy("date") \
    .parquet("s3://bucket/processed/parquet/")

print(f"Original CSV: {df.inputFiles()[0]}")
print(f"Rows: {df.count():,}")

# 스트리밍에서 수집된 Avro를 분석용 Parquet로 변환
from pyspark.sql import SparkSession
from datetime import datetime, timedelta

spark = SparkSession.builder \
    .appName("Avro to Parquet Batch") \
    .getOrCreate()

# 어제 날짜 데이터 처리
yesterday = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d")

# Avro 읽기
avro_path = f"s3://bucket/streaming/avro/date={yesterday}/"
df = spark.read.format("avro").load(avro_path)

# 데이터 품질 체크
print(f"Records: {df.count():,}")
print(f"Duplicates: {df.count() - df.dropDuplicates().count():,}")

# 중복 제거 및 정렬
df = df.dropDuplicates(["id"]) \
    .orderBy("event_time")

# Parquet로 저장
parquet_path = f"s3://bucket/analytics/parquet/date={yesterday}/"
df.repartition(20) \
    .write \
    .mode("overwrite") \
    .parquet(parquet_path)

# 검증
parquet_df = spark.read.parquet(parquet_path)
assert df.count() == parquet_df.count(), "Record count mismatch!"

print(f"✓ Migration completed: {yesterday}")

# ORC → Parquet
orc_df = spark.read.format("orc").load("s3://bucket/data.orc")
orc_df.write.parquet("s3://bucket/data.parquet")

# Parquet → ORC
parquet_df = spark.read.parquet("s3://bucket/data.parquet")
parquet_df.write.format("orc").save("s3://bucket/data.orc")

# 성능 비교
import time

# ORC 읽기
start = time.time()
orc_df = spark.read.format("orc").load("s3://bucket/large_data.orc")
orc_count = orc_df.count()
orc_time = time.time() - start

# Parquet 읽기
start = time.time()
parquet_df = spark.read.parquet("s3://bucket/large_data.parquet")
parquet_count = parquet_df.count()
parquet_time = time.time() - start

print(f"ORC: {orc_time:.2f}s, {orc_count:,} rows")
print(f"Parquet: {parquet_time:.2f}s, {parquet_count:,} rows")

# 1. 압축 코덱 선택
# - Snappy: 빠른 압축/해제 (실시간 분석)
# - GZIP: 높은 압축률 (장기 보관)
# - ZSTD: 균형잡힌 성능 (권장)

spark.conf.set("spark.sql.parquet.compression.codec", "zstd")

# 2. Row Group 크기 조정
spark.conf.set("spark.sql.parquet.block.size", 268435456)  # 256MB

# 3. Dictionary encoding 활용
# 카디널리티 낮은 컬럼에 자동 적용
# 수동으로 비활성화하려면:
spark.conf.set("spark.sql.parquet.enableDictionaryEncoding", "false")

# 4. Vectorized reader 활성화
spark.conf.set("spark.sql.parquet.enableVectorizedReader", "true")

# 5. Binary as string 최적화
spark.conf.set("spark.sql.parquet.binaryAsString", "false")

# 1. Stripe 크기 조정
spark.conf.set("spark.sql.orc.stripe.size", 67108864)  # 64MB

# 2. Bloom filter 설정
df.write \
    .format("orc") \
    .option("orc.bloom.filter.columns", "user_id,product_id") \
    .option("orc.bloom.filter.fpp", 0.05) \
    .save("s3://bucket/data.orc")

# 3. 압축 선택
# - ZLIB: 최고 압축률 (기본값)
# - SNAPPY: 빠른 성능
# - LZO: 균형

spark.conf.set("spark.sql.orc.compression.codec", "zlib")

# 4. Index stride (row index 간격)
spark.conf.set("orc.row.index.stride", 10000)

# 1. 압축 설정
df.write \
    .format("avro") \
    .option("compression", "snappy") \
    .save("s3://bucket/data.avro")

# 2. 스키마 레지스트리 연동
from confluent_kafka import avro
from confluent_kafka.avro import AvroProducer

producer_config = {
    'bootstrap.servers': 'localhost:9092',
    'schema.registry.url': 'http://localhost:8081'
}

producer = AvroProducer(producer_config, default_value_schema=schema)

def choose_format(use_case):
    """포맷 선택 도우미 함수"""
    
    # 실시간 스트리밍?
    if use_case["streaming"] and use_case["schema_changes"]:
        return "Avro"
    
    # Hive 중심 환경?
    if use_case["hive"] and use_case["acid"]:
        return "ORC"
    
    # 최대 압축률 필요?
    if use_case["storage_critical"]:
        return "ORC with ZLIB"
    
    # 범용 분석?
    if use_case["analytics"] and use_case["athena"]:
        return "Parquet with Snappy"
    
    # 빠른 쓰기 필요?
    if use_case["write_heavy"]:
        return "Avro"
    
    # 기본값
    return "Parquet"

# 사용 예시
use_case = {
    "streaming": False,
    "schema_changes": False,
    "hive": False,
    "acid": False,
    "storage_critical": False,
    "analytics": True,
    "athena": True,
    "write_heavy": False
}

recommended = choose_format(use_case)
print(f"Recommended format: {recommended}")
# 출력: Recommended format: Parquet with Snappy