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

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

# Generate data
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)

# Optimize Parquet settings
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

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

# Read Parquet file metadata
import pyarrow.parquet as pq

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

# Check schema
print("Schema:")
print(parquet_file.schema)

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

# Statistics per 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")
    
    # Statistics per column
    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")

# Create ORC in Spark
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")

# Analyze ORC file (using PyArrow)
import pyarrow.orc as orc

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

# Schema
print("Schema:")
print(orc_file.schema)

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

# Metadata
metadata = orc_file.metadata
print(f"Compression: {metadata.compression}")
print(f"Writer Version: {metadata.writer_version}")

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

# Use Avro in Kafka
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)

# Send message
avroProducer.produce(topic='events', value={"id": 1, "name": "Alice"})
avroProducer.flush()

import time

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

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

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

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

# Schema change test
# 1. Save data with existing schema
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 with new column added
schema_v2 = StructType([
    StructField("id", IntegerType()),
    StructField("name", StringType()),
    StructField("amount", DoubleType()),
    StructField("category", StringType())  # New column
])

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

# 3. Read both versions simultaneously
df_merged = spark.read.format(format_type).load(f"data_{format_type}_*")

# Optimal settings
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/")

# Real-time: Kafka + Avro
from confluent_kafka import avro

# Save to Kafka as Avro
avro_producer.produce(topic='events', value=event_data)

# Batch: Avro → Parquet conversion
df = spark.read.format("avro").load("s3://bucket/streaming/avro/")

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

# Maximum compression ratio settings
spark.conf.set("spark.sql.parquet.compression.codec", "gzip")  # or zstd

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

# Automatic transition with 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'}
            ]
        }]
    }
)

# Nested JSON data
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}
    ]
  }
}
"""

# Handle nested structure in Spark
df = spark.read.json("s3://bucket/raw/events.json")

# Save as Parquet (preserving nested structure)
df.write.parquet("s3://bucket/processed/events.parquet")

# Efficiently read only specific fields
df = spark.read.parquet("s3://bucket/processed/events.parquet")
df.select("user.profile.name", "event.type").show()
# Parquet reads only needed columns (nested column pruning)

from pyspark.sql import SparkSession

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

# Read CSV (schema inference)
df = spark.read \
    .option("header", "true") \
    .option("inferSchema", "true") \
    .csv("s3://bucket/raw/csv/*.csv")

# Optimize data types
from pyspark.sql.functions import col

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

# Convert to 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():,}")

# Convert Avro collected from streaming to Parquet for analytics
from pyspark.sql import SparkSession
from datetime import datetime, timedelta

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

# Process yesterday's data
yesterday = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d")

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

# Data quality check
print(f"Records: {df.count():,}")
print(f"Duplicates: {df.count() - df.dropDuplicates().count():,}")

# Remove duplicates and sort
df = df.dropDuplicates(["id"]) \
    .orderBy("event_time")

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

# Validation
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")

# Performance comparison
import time

# ORC read
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 read
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. Choose compression codec
# - Snappy: Fast compression/decompression (real-time analytics)
# - GZIP: High compression ratio (long-term storage)
# - ZSTD: Balanced performance (recommended)

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

# 2. Adjust Row Group size
spark.conf.set("spark.sql.parquet.block.size", 268435456)  # 256MB

# 3. Utilize dictionary encoding
# Automatically applied to low cardinality columns
# To manually disable:
spark.conf.set("spark.sql.parquet.enableDictionaryEncoding", "false")

# 4. Enable vectorized reader
spark.conf.set("spark.sql.parquet.enableVectorizedReader", "true")

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

# 1. Adjust Stripe size
spark.conf.set("spark.sql.orc.stripe.size", 67108864)  # 64MB

# 2. Set 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. Choose compression
# - ZLIB: Best compression ratio (default)
# - SNAPPY: Fast performance
# - LZO: Balanced

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

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

# 1. Compression settings
df.write \
    .format("avro") \
    .option("compression", "snappy") \
    .save("s3://bucket/data.avro")

# 2. Schema registry integration
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):
    """Format selection helper function"""
    
    # Real-time streaming?
    if use_case["streaming"] and use_case["schema_changes"]:
        return "Avro"
    
    # Hive-centric environment?
    if use_case["hive"] and use_case["acid"]:
        return "ORC"
    
    # Maximum compression needed?
    if use_case["storage_critical"]:
        return "ORC with ZLIB"
    
    # General analytics?
    if use_case["analytics"] and use_case["athena"]:
        return "Parquet with Snappy"
    
    # Fast writes needed?
    if use_case["write_heavy"]:
        return "Avro"
    
    # Default
    return "Parquet"

# Usage example
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}")
# Output: Recommended format: Parquet with Snappy