#!/usr/bin/env python3 import argparse import logging import sys import os from datetime import datetime, timedelta from deltalake import DeltaTable # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) def optimize_delta_table(s3_path, region, start_date=None, end_date=None, days_batch=3): """ Optimize and vacuum a Delta table in S3, optionally for specific date partitions. Args: s3_path (str): Full S3 path to the Delta table (e.g., 's3://bucket/prefix/') region (str): AWS region where the S3 bucket is located start_date (str, optional): Start date in 'YYYYMMDD' format (inclusive) end_date (str, optional): End date in 'YYYYMMDD' format (inclusive) days_batch (int): Number of days to process in one batch (default: 3) """ endpoint_url = os.getenv("AWS_ENDPOINT_S3") if endpoint_url: storage_options = { 'AWS_REGION': region, 'AWS_ENDPOINT_URL': endpoint_url, 'AWS_COPY_IF_NOT_EXISTS': 'header-with-status:x-oss-forbid-overwrite:true:409' } else: storage_options = { 'AWS_REGION': region, } def optimize_table(delta_table, partition_filters=None): """Optimize a Delta table with optional partition filters""" try: logger.info(f"Optimizing table: {s3_path}") if partition_filters: logger.info(f"Using partition filters: {partition_filters}") optimize_metrics = delta_table.optimize.compact( partition_filters=partition_filters, target_size=512 * 1024 * 1024 ) else: optimize_metrics = delta_table.optimize.compact( target_size=512 * 1024 * 1024 ) logger.info(f"Table compacted. Metrics: {optimize_metrics}") except Exception as e: logger.error(f"Error optimizing table: {str(e)}", exc_info=True) return False try: logger.info("Running vacuum...") # Run vacuum with a shorter retention period for batch processing delta_table.create_checkpoint() vacuum_metrics = delta_table.vacuum( retention_hours=2, enforce_retention_duration=False, dry_run=False ) delta_table.cleanup_metadata() logger.info(f"Vacuum completed. Removed {vacuum_metrics} files.") return True except Exception as e: logger.error(f"Error during vacuum: {str(e)}") logger.warning("Skipping vacuum for this batch. Continuing with next batch...") return False def process_date_range(delta_table, start, end, batch_size): """Process date range in batches using partition filters""" current_date = start success = True batch_count = 0 while current_date <= end and success: batch_count += 1 batch_end = min(current_date + timedelta(days=batch_size - 1), end) start_str = current_date.strftime('%Y%m%d') end_str = batch_end.strftime('%Y%m%d') logger.info(f"\n=== Processing batch {batch_count}: {start_str} to {end_str} ===") # Optimize current batch logger.info(f"Optimizing date range: {start_str} to {end_str}") partition_filters = [ ("date", ">=", start_str), ("date", "<=", end_str) ] if not optimize_table(delta_table, partition_filters): return False current_date = batch_end + timedelta(days=1) return True try: # Initialize Delta table logger.info(f"Connecting to Delta table at {s3_path}") delta_table = DeltaTable(s3_path, storage_options=storage_options) logger.info(f"Table version: {delta_table.version()}") if start_date and end_date: # Process specific date range in batches start = datetime.strptime(start_date, '%Y%m%d').date() end = datetime.strptime(end_date, '%Y%m%d').date() if start > end: logger.error("Start date cannot be after end date") return 1 logger.info(f"Processing date range: {start_date} to {end_date} in batches of {days_batch} days") if not process_date_range(delta_table, start, end, days_batch): return 1 else: # Process full table if not optimize_table(delta_table): return 1 except Exception as e: logger.error(f"Error during optimization: {str(e)}", exc_info=True) return 1 def parse_date(date_str): try: return datetime.strptime(date_str, '%Y%m%d').date() except ValueError: raise argparse.ArgumentTypeError(f"Invalid date format: {date_str}. Expected format: YYYYMMDD") def main(): parser = argparse.ArgumentParser(description='Optimize and vacuum a Delta table in S3 with partition support.') parser.add_argument('--s3-path', required=True, help='S3 path to the Delta table (e.g., s3://bucket/prefix/)') parser.add_argument('--region', required=True, help='AWS region where the S3 bucket is located') parser.add_argument('--start-date', type=parse_date, help='Start date for partition optimization (format: YYYYMMDD)') parser.add_argument('--end-date', type=parse_date, help='End date for partition optimization (format: YYYYMMDD, default: same as start_date)') parser.add_argument('--days-batch', type=int, default=3, help='Number of days to process in one batch (default: 3)') args = parser.parse_args() # Validate arguments if (args.start_date or args.end_date) and not (args.start_date and args.end_date): parser.error("Both --start-date and --end-date must be specified together") if args.start_date and args.end_date and args.start_date > args.end_date: parser.error("Start date cannot be after end date") if args.days_batch < 1: parser.error("Days batch must be at least 1") # Ensure s3_path ends with a slash s3_path = args.s3_path.rstrip('/') + '/' # Format dates back to strings start_date = args.start_date.strftime('%Y%m%d') if args.start_date else None end_date = args.end_date.strftime('%Y%m%d') if args.end_date else None return optimize_delta_table( s3_path=s3_path, region=args.region, start_date=start_date, end_date=end_date, days_batch=args.days_batch ) if __name__ == "__main__": sys.exit(main())