#!/usr/bin/env python3 """ Delta data generator - Generate random statement log data and write to Delta Lake table """ import os import json import random import time import string import pandas as pd import argparse import pyarrow as pa from pathlib import Path from datetime import datetime, timezone, timedelta from deltalake import DeltaTable, write_deltalake class DeltaDataGenerator: def __init__(self, output_dir="./statement_delta", num_hosts=500, records_per_host=100, interval_seconds=15, duplicate_rate=0.3): """ Initialize data generator Args: output_dir: Output Delta Lake table directory num_hosts: Number of hosts to generate data for records_per_host: Number of records per host interval_seconds: Generation interval (seconds) duplicate_rate: Duplicate rate (0.0-1.0) """ self.output_dir = Path(output_dir) self.output_dir.mkdir(exist_ok=True) self.num_hosts = num_hosts self.records_per_host = records_per_host self.interval_seconds = interval_seconds self.duplicate_rate = duplicate_rate # Predefined option sets self.schema_names = ["prod_betwallet", "test_db", "staging_env", "dev_schema", "analytics_db"] self.stmt_types = ["Select", "Insert", "Update", "Delete", "Create", "Drop"] self.charsets = ["utf8mb4", "utf8", "latin1", "ascii"] self.collations = ["utf8mb4_0900_ai_ci", "utf8mb4_general_ci", "utf8_general_ci", "latin1_swedish_ci"] self.resource_groups = ["default", "high_priority", "low_priority", "batch_jobs"] # Store generated records for duplicate rate control self.generated_records = [] self.duplicate_pool = [] # Statistics self.stats = { 'total_records': 0, 'unique_records': 0, 'duplicate_records': 0, 'current_batch': 0 } def generate_random_string(self, length=10): """Generate random string""" return ''.join(random.choices(string.ascii_letters + string.digits, k=length)) def generate_random_sql(self, stmt_type="Select"): """Generate random SQL statement""" if stmt_type == "Select": table_name = f"table_{random.randint(1000, 9999)}" columns = ["id", "name", "value", "created_at", "updated_at"] selected_cols = random.sample(columns, random.randint(1, len(columns))) conditions = ["id = ?", "name LIKE ?", "value > ?", "created_at BETWEEN ? AND ?"] sql = f"SELECT {', '.join(selected_cols)} FROM {table_name}" if random.choice([True, False]): condition = random.choice(conditions) sql += f" WHERE {condition}" if random.choice([True, False]): sql += " ORDER BY id" if random.choice([True, False]): sql += " LIMIT " + str(random.randint(1, 1000)) elif stmt_type == "Insert": table_name = f"table_{random.randint(1000, 9999)}" columns = ["id", "name", "value", "created_at"] values = [random.randint(1, 1000000), f"'{self.generate_random_string(8)}'", random.randint(1, 1000), "NOW()"] sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(map(str, values))})" elif stmt_type == "Update": table_name = f"table_{random.randint(1000, 9999)}" sql = f"UPDATE {table_name} SET name = '{self.generate_random_string(8)}' WHERE id = {random.randint(1, 1000000)}" elif stmt_type == "Delete": table_name = f"table_{random.randint(1000, 9999)}" sql = f"DELETE FROM {table_name} WHERE id = {random.randint(1, 1000000)}" else: sql = f"-- {stmt_type} statement placeholder" return sql def generate_normalized_sql(self, stmt_type="Select"): """Generate normalized SQL""" base_sql = self.generate_random_sql(stmt_type) # Normalize: remove specific values, keep structure import re normalized = re.sub(r'\d+', '?', base_sql) normalized = re.sub(r"'[^']*'", '?', normalized) return normalized def generate_sample_plan(self): """Generate sample execution plan (simulated)""" plans = [ "TableScan(table=users, range=[1, 1000])", "IndexScan(index=idx_name, range=['a', 'z'])", "HashJoin(left=users, right=orders)", "Sort(order=id)", "Limit(count=100)" ] return " -> ".join(random.sample(plans, random.randint(2, 4))) def generate_base64_like_string(self, length=100): """Generate base64-like random string""" chars = string.ascii_letters + string.digits + "+/=" return ''.join(random.choices(chars, k=length)) def generate_single_record(self, host_id, record_index): """ Generate a single record Args: host_id: Host identifier record_index: Record index for time calculation Returns: dict: Generated record data """ # Base timestamp (incremental) base_time = int(time.time()) begin_time = base_time - random.randint(1, 3600) + record_index end_time = begin_time + random.randint(1, 1000) # Randomly select statement type stmt_type = random.choice(self.stmt_types) record = { "begin": begin_time, "end": end_time, "schema_name": random.choice(self.schema_names), "digest": self.generate_random_string(64), "plan_digest": self.generate_random_string(64), "stmt_type": stmt_type, "normalized_sql": self.generate_normalized_sql(stmt_type), "table_names": f"{random.choice(self.schema_names)}.table_{random.randint(1000, 9999)}", "is_internal": random.choice([True, False]), "sample_sql": self.generate_random_sql(stmt_type), "charset": random.choice(self.charsets), "collation": random.choice(self.collations), "prev_sql": self.generate_random_sql(random.choice(self.stmt_types)) if random.choice([True, False]) else "", "sample_plan": self.generate_sample_plan(), "sample_binary_plan": self.generate_base64_like_string(200), "plan_hint": f"use_index(@`sel_1` `{random.choice(self.schema_names)}`.`table_{random.randint(1000, 9999)}` `idx_{self.generate_random_string(10)}`)" if random.choice([True, False]) else "", "index_names": [f"table_{random.randint(1000, 9999)}:idx_{self.generate_random_string(8)}" for _ in range(random.randint(0, 3))], "exec_count": random.randint(1, 1000), "sum_errors": random.randint(0, 10), "sum_warnings": random.randint(0, 50), "sum_latency": random.randint(1000, 10000000), "max_latency": random.randint(100, 100000), "min_latency": random.randint(1, 1000), "sum_parse_latency": random.randint(100, 100000), "max_parse_latency": random.randint(10, 10000), "sum_compile_latency": random.randint(1000, 1000000), "max_compile_latency": random.randint(100, 100000), "sum_num_cop_tasks": random.randint(1, 100), "max_cop_process_time": random.randint(1000, 100000), "max_cop_process_address": f"db-tikv-{random.randint(1, 10)}.db-tikv-peer.tidb{random.randint(1000000000000000000, 9999999999999999999)}.svc:20160", "sum_process_time": random.randint(1000, 1000000), "max_process_time": random.randint(100, 100000), "sum_wait_time": random.randint(100, 100000), "max_wait_time": random.randint(10, 10000), "sum_backoff_time": random.randint(0, 10000), "max_backoff_time": random.randint(0, 1000), "sum_total_keys": random.randint(1, 10000), "max_total_keys": random.randint(1, 1000), "sum_processed_keys": random.randint(1, 1000), "max_processed_keys": random.randint(1, 100), "sum_rocksdb_delete_skipped_count": random.randint(0, 100), "max_rocksdb_delete_skipped_count": random.randint(0, 10), "sum_rocksdb_key_skipped_count": random.randint(0, 1000), "max_rocksdb_key_skipped_count": random.randint(0, 100), "sum_rocksdb_block_cache_hit_count": random.randint(0, 10000), "max_rocksdb_block_cache_hit_count": random.randint(0, 1000), "sum_rocksdb_block_read_count": random.randint(0, 1000), "max_rocksdb_block_read_count": random.randint(0, 100), "sum_rocksdb_block_read_byte": random.randint(0, 1000000), "max_rocksdb_block_read_byte": random.randint(0, 100000), "commit_count": random.randint(0, 100), "sum_get_commit_ts_time": random.randint(0, 10000), "max_get_commit_ts_time": random.randint(0, 1000), "sum_prewrite_time": random.randint(0, 10000), "max_prewrite_time": random.randint(0, 1000), "sum_commit_time": random.randint(0, 10000), "max_commit_time": random.randint(0, 1000), "sum_local_latch_time": random.randint(0, 1000), "max_local_latch_time": random.randint(0, 100), "sum_commit_backoff_time": random.randint(0, 1000), "max_commit_backoff_time": random.randint(0, 100), "sum_resolve_lock_time": random.randint(0, 1000), "max_resolve_lock_time": random.randint(0, 100), "sum_write_keys": random.randint(0, 1000), "max_write_keys": random.randint(0, 100), "sum_write_size": random.randint(0, 100000), "max_write_size": random.randint(0, 10000), "sum_prewrite_region_num": random.randint(0, 50), "max_prewrite_region_num": random.randint(0, 10), "sum_txn_retry": random.randint(0, 10), "max_txn_retry": random.randint(0, 5), "sum_backoff_times": random.randint(0, 100), "backoff_types": {}, "auth_users": {f"user_{random.randint(1, 100)}": {}}, "sum_mem": random.randint(1000, 1000000), "max_mem": random.randint(100, 100000), "sum_disk": random.randint(0, 1000000), "max_disk": random.randint(0, 100000), "sum_affected_rows": random.randint(0, 1000), "sum_kv_total": random.randint(1000, 10000000), "sum_pd_total": random.randint(0, 10000), "sum_backoff_total": random.randint(0, 10000), "sum_write_sql_resp_total": random.randint(0, 100000), "sum_result_rows": random.randint(0, 10000), "max_result_rows": random.randint(0, 1000), "min_result_rows": random.randint(0, 100), "prepared": random.choice([True, False]), "first_seen": (datetime.now(timezone.utc) - timedelta(hours=random.randint(1, 24))).isoformat(), "last_seen": datetime.now(timezone.utc).isoformat(), "plan_in_cache": random.choice([True, False]), "plan_cache_hits": random.randint(0, 1000), "plan_in_binding": random.choice([True, False]), "exec_retry_count": random.randint(0, 10), "exec_retry_time": random.randint(0, 1000), "resource_group_name": random.choice(self.resource_groups), "sum_rru": round(random.uniform(0.1, 100.0), 6), "sum_wru": random.randint(0, 1000), "sum_ru_wait_duration": random.randint(0, 10000), "max_rru": round(random.uniform(0.01, 10.0), 6), "max_wru": random.randint(0, 100), "max_ru_wait_duration": random.randint(0, 1000), "host_id": host_id } return record def generate_batch_data(self, batch_size=1000): """ Generate a batch of data with duplicate control Args: batch_size: Batch size Returns: list: List of generated records """ batch_data = [] records_needed = batch_size # Calculate number of new records and duplicate records needed unique_count = int(records_needed * (1 - self.duplicate_rate)) duplicate_count = records_needed - unique_count # Generate new records for _ in range(unique_count): host_id = f"host_{random.randint(1, self.num_hosts)}" record = self.generate_single_record(host_id, len(self.generated_records)) batch_data.append(record) self.generated_records.append(record) self.duplicate_pool.append(record) self.stats['unique_records'] += 1 # Add duplicate records for _ in range(duplicate_count): if self.duplicate_pool: # Deep copy record to avoid modifying original duplicate_record = json.loads(json.dumps(random.choice(self.duplicate_pool))) # Make slight modifications to duplicate record to avoid exact duplicates duplicate_record['end'] += random.randint(1, 100) duplicate_record['sum_latency'] += random.randint(-1000, 1000) batch_data.append(duplicate_record) self.stats['duplicate_records'] += 1 self.stats['total_records'] += len(batch_data) return batch_data def write_batch_to_delta(self, batch_data, batch_num): """ Write a batch of data to Delta Lake table Args: batch_data: Batch data batch_num: Batch number """ try: # Preprocess data, handle None values and empty structures processed_data = [] for item in batch_data: if item is None: continue # Handle None values and complex types in nested structures processed_item = {} for key, value in item.items(): if value is None: processed_item[key] = None elif isinstance(value, dict): # Recursively process nested dictionaries cleaned_dict = {k: v for k, v in value.items() if v is not None} processed_item[key] = json.dumps(cleaned_dict, ensure_ascii=False, default=str) if cleaned_dict else None elif isinstance(value, list): # Handle None values in lists and convert to JSON string cleaned_list = [v for v in value if v is not None] processed_item[key] = json.dumps(cleaned_list, ensure_ascii=False, default=str) if cleaned_list else None else: processed_item[key] = value processed_data.append(processed_item) if not processed_data: print(f"Warning: Batch {batch_num} has no valid data, skipping") return # Convert to DataFrame df = pd.DataFrame(processed_data) # Add processing timestamp and batch number df['processed_at'] = datetime.now(timezone.utc) df['batch_num'] = batch_num # Generate Delta Lake schema delta_schema = self._get_delta_schema(df) # Save to Delta Lake table_path = self.output_dir / "rand_statement_logs" table_path_str = str(table_path.absolute()) print(f"Writing batch {batch_num} with {len(df)} records...") # Ensure output directory exists table_path.parent.mkdir(parents=True, exist_ok=True) # Convert to PyArrow Table and write pa_table = pa.Table.from_pandas(df, schema=delta_schema, preserve_index=False) write_deltalake( table_path_str, pa_table, mode="append", ) print(f"Batch {batch_num} write completed") except Exception as e: print(f"Error writing batch {batch_num}: {e}") raise def _get_delta_schema(self, df): """ Generate Delta Lake schema from DataFrame Args: df: Input DataFrame Returns: pa.Schema: PyArrow schema """ schema_fields = [] for column in df.columns: if column in ['processed_at', 'batch_num']: # Timestamp columns schema_fields.append(pa.field(column, pa.timestamp('us', tz='UTC'), nullable=True)) elif column in ['begin', 'end']: # Timestamp columns schema_fields.append(pa.field(column, pa.timestamp('s'), nullable=True)) elif column == 'first_seen': # ISO format time string schema_fields.append(pa.field(column, pa.string(), nullable=True)) elif column == 'last_seen': # ISO format time string schema_fields.append(pa.field(column, pa.string(), nullable=True)) elif column in ['is_internal', 'prepared', 'plan_in_cache', 'plan_in_binding']: # Boolean values schema_fields.append(pa.field(column, pa.bool_(), nullable=True)) elif column in ['index_names', 'backoff_types', 'auth_users']: # Complex types, convert to string schema_fields.append(pa.field(column, pa.string(), nullable=True)) elif df[column].dtype in ['int64', 'int32']: # Integer types schema_fields.append(pa.field(column, pa.int64(), nullable=True)) elif df[column].dtype in ['float64', 'float32']: # Float types schema_fields.append(pa.field(column, pa.float64(), nullable=True)) else: # Default string type schema_fields.append(pa.field(column, pa.string(), nullable=True)) return pa.schema(schema_fields) def run_continuous_generation(self, duration_minutes=60): """ Continuously generate data Args: duration_minutes: Running duration (minutes) """ print(f"Starting continuous data generation for {duration_minutes} minutes...") print(f"Parameters: hosts={self.num_hosts}, records_per_host={self.records_per_host}, interval={self.interval_seconds}s, duplicate_rate={self.duplicate_rate}") start_time = time.time() end_time = start_time + (duration_minutes * 60) batch_num = 1 while time.time() < end_time: # Calculate batch size (based on records_per_host and num_hosts, but reasonably allocated) batch_size = min(1000, self.records_per_host * min(10, self.num_hosts // 50)) # Generate a batch of data batch_data = self.generate_batch_data(batch_size) # Write to Delta Lake self.write_batch_to_delta(batch_data, batch_num) self.stats['current_batch'] = batch_num # Wait for next generation if time.time() + self.interval_seconds < end_time: print(f"Waiting {self.interval_seconds} seconds...") time.sleep(self.interval_seconds) else: break batch_num += 1 self.print_stats() def run_single_batch(self): """Generate single batch of data""" print("Generating single batch of data...") # Generate a batch of data batch_size = self.records_per_host * min(10, self.num_hosts // 50) batch_data = self.generate_batch_data(batch_size) # Write to Delta Lake self.write_batch_to_delta(batch_data, 1) self.print_stats() def print_stats(self): """Print statistics""" print("\n" + "="*50) print("Generation Statistics:") print(f"Total records: {self.stats['total_records']}") print(f"Unique records: {self.stats['unique_records']}") print(f"Duplicate records: {self.stats['duplicate_records']}") print(f"Duplicate rate: {self.stats['duplicate_records']/max(self.stats['total_records'], 1)*100:.2f}%") print(f"Total batches: {self.stats['current_batch']}") print("="*50) def main(): parser = argparse.ArgumentParser(description='Generate random statement log data and write to Delta Lake table') parser.add_argument('--output-dir', default='./statement_delta', help='Output Delta Lake table directory') parser.add_argument('--num-hosts', type=int, default=500, help='Number of hosts to generate data for') parser.add_argument('--records-per-host', type=int, default=100, help='Number of records per host') parser.add_argument('--interval-seconds', type=int, default=15, help='Generation interval (seconds)') parser.add_argument('--duplicate-rate', type=float, default=0.3, help='Duplicate rate (0.0-1.0)') parser.add_argument('--duration-minutes', type=int, default=60, help='Continuous running duration (minutes)') parser.add_argument('--single-batch', action='store_true', help='Generate only a single batch of data') args = parser.parse_args() # Create generator generator = DeltaDataGenerator( output_dir=args.output_dir, num_hosts=args.num_hosts, records_per_host=args.records_per_host, interval_seconds=args.interval_seconds, duplicate_rate=args.duplicate_rate ) if args.single_batch: generator.run_single_batch() else: generator.run_continuous_generation(args.duration_minutes) print(f"\nāœ… Data generation completed") print(f"šŸ“ Results saved in: {generator.output_dir / 'statement_logs'}") if __name__ == "__main__": main()