rosetta-sql 1.0.0__py3-none-any.whl

benchmark/generate_csv_data.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""
+生成 bench_json_mv_select 测试所需的 CSV 数据文件。
+
+生成文件:
+- bench_sel.csv: 100万行字符串tag数据 (user_id, json_array)
+- bench_sel_int.csv: 100万行整数数据 (user_id, json_array)
+
+用法:
+    python generate_csv_data.py [--output-dir DIR] [--rows N]
+"""
+
+import argparse
+import csv
+import json
+import os
+import random
+from typing import List, Tuple
+
+
+def generate_string_data(num_rows: int = 1_000_000) -> List[Tuple[int, str]]:
+    """生成字符串tag数据，数组大小在 [1, 20] 范围内随机变化。"""
+    rows = []
+    for _ in range(num_rows):
+        user_id = random.randint(1, 100000)
+        array_size = random.randint(1, 20)
+        tags = [f"tag_{random.randint(0, 999)}" for _ in range(array_size)]
+        json_data = json.dumps(tags)
+        rows.append((user_id, json_data))
+    return rows
+
+
+def generate_int_data(num_rows: int = 1_000_000) -> List[Tuple[int, str]]:
+    """生成整数数据，数组大小在 [1, 20] 范围内随机变化 (值域0-9999)。"""
+    rows = []
+    for _ in range(num_rows):
+        user_id = random.randint(1, 100000)
+        array_size = random.randint(1, 20)
+        values = [random.randint(0, 9999) for _ in range(array_size)]
+        json_data = json.dumps(values)
+        rows.append((user_id, json_data))
+    return rows
+
+
+def write_csv(filename: str, rows: List[Tuple], output_dir: str):
+    """写入CSV文件，无header，制表符分隔。JSON字段不加引号。"""
+    filepath = os.path.join(output_dir, filename)
+    with open(filepath, 'w', newline='', encoding='utf-8') as f:
+        for row in rows:
+            # 手动写入，不使用csv.writer避免自动引号
+            # 格式: user_id\tjson_data\n
+            f.write(f"{row[0]}\t{row[1]}\n")
+    print(f"Generated: {filepath} ({len(rows)} rows)")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate CSV data files for benchmark')
+    parser.add_argument('--output-dir', '-o', default='.', help='Output directory (default: current dir)')
+    parser.add_argument('--rows', '-n', type=int, default=1_000_000, help='Number of rows per table (default: 1000000)')
+    args = parser.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # 设置随机种子以便可重现
+    random.seed(42)
+
+    print(f"Generating {args.rows:,} rows per table...")
+
+    # 生成字符串数据
+    print("Generating string data...")
+    str_rows = generate_string_data(args.rows)
+    write_csv('bench_sel.csv', str_rows, args.output_dir)
+
+    # 生成整数数据
+    print("Generating integer data...")
+    int_rows = generate_int_data(args.rows)
+    write_csv('bench_sel_int.csv', int_rows, args.output_dir)
+
+    print("Done!")
+
+
+if __name__ == '__main__':
+    main()

benchmark/import_data.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+"""
+将 CSV 数据导入到目标数据库。
+
+用法:
+    python import_data.py --config ../dbms_config.json --dbms tdsql
+    python import_data.py --config ../dbms_config.json --dbms all
+"""
+
+import argparse
+import json
+import sys
+import time
+
+try:
+    import pymysql
+except ImportError:
+    print("Error: pymysql is required. Install via: pip install pymysql")
+    sys.exit(1)
+
+
+CREATE_TABLE_SQL = """
+CREATE TABLE IF NOT EXISTS bench_sel (
+    id INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
+    user_id INT NOT NULL,
+    data JSON NOT NULL,
+    KEY idx_user (user_id),
+    KEY idx_mv_data ((CAST(data->'$' AS CHAR(64) ARRAY)))
+)
+"""
+
+
+def import_data(config: dict, database: str, csv_path: str, batch_size: int = 10000):
+    """Import data for a single DBMS."""
+    name = config['name']
+    print(f"\n[{name}] Connecting to {config['host']}:{config['port']}...")
+    print(f"[{name}] Target database: {database}")
+    
+    try:
+        # Connect and create database
+        conn = pymysql.connect(
+            host=config['host'],
+            port=config['port'],
+            user=config['user'],
+            password=config['password'],
+            local_infile=True,
+            charset='utf8mb4',
+        )
+        cursor = conn.cursor()
+        cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{database}`")
+        cursor.execute(f"USE `{database}`")
+        conn.commit()
+        
+        # Drop existing table and create new one with MVI index
+        print(f"[{name}] Dropping existing table...")
+        cursor.execute("DROP TABLE IF EXISTS bench_sel")
+        conn.commit()
+        print(f"[{name}] Creating table...")
+        cursor.execute(CREATE_TABLE_SQL)
+        conn.commit()
+        
+        # Import CSV using batch INSERT
+        print(f"[{name}] Importing data...")
+        insert_sql = "INSERT INTO bench_sel (user_id, data) VALUES (%s, %s)"
+        
+        total_rows = 0
+        batch = []
+        batch_times = []  # 记录每批次耗时
+        import_start_time = time.time()
+        batch_start_time = time.time()
+        
+        with open(csv_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                parts = line.rstrip('\n').split('\t', 1)
+                if len(parts) == 2:
+                    batch.append((int(parts[0]), parts[1]))
+                    
+                    if len(batch) >= batch_size:
+                        cursor.executemany(insert_sql, batch)
+                        conn.commit()
+                        total_rows += len(batch)
+                        
+                        # 记录本批次耗时
+                        batch_elapsed = time.time() - batch_start_time
+                        batch_times.append(batch_elapsed)
+                        
+                        # 计算统计信息
+                        avg_batch_time = sum(batch_times) / len(batch_times)
+                        rows_per_sec = len(batch) / batch_elapsed if batch_elapsed > 0 else 0
+                        elapsed_total = time.time() - import_start_time
+                        
+                        print(f"[{name}] Imported {total_rows:,} rows | "
+                              f"Batch: {batch_elapsed:.2f}s | "
+                              f"Avg: {avg_batch_time:.2f}s/batch | "
+                              f"Speed: {rows_per_sec:,.0f} rows/s", end='\r')
+                        batch = []
+                        batch_start_time = time.time()
+        
+        # Insert remaining rows
+        if batch:
+            cursor.executemany(insert_sql, batch)
+            conn.commit()
+            total_rows += len(batch)
+            batch_elapsed = time.time() - batch_start_time
+            if batch_elapsed > 0:
+                batch_times.append(batch_elapsed)
+        
+        import_total_time = time.time() - import_start_time
+        avg_rows_per_sec = total_rows / import_total_time if import_total_time > 0 else 0
+        
+        print(f"\n[{name}] Import completed: {total_rows:,} rows in {import_total_time:.2f}s "
+              f"({avg_rows_per_sec:,.0f} rows/s)")
+        if batch_times:
+            print(f"[{name}] Average batch time: {sum(batch_times)/len(batch_times):.3f}s "
+              f"(min: {min(batch_times):.3f}s, max: {max(batch_times):.3f}s, batches: {len(batch_times)})")
+        
+        cursor.close()
+        conn.close()
+        print(f"[{name}] Done!")
+        return True
+        
+    except Exception as e:
+        error_msg = str(e)
+        # 静默处理连接拒绝错误
+        if "Connection refused" in error_msg or "111" in error_msg:
+            print(f"[{name}] Skipped: service unavailable")
+            return False
+        print(f"[{name}] ERROR: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Import CSV data to database')
+    parser.add_argument('--config', '-c', required=True, help='Path to dbms_config.json')
+    parser.add_argument('--dbms', '-d', default='all', help='DBMS name to import (default: all)')
+    parser.add_argument('--database', '-D', default='cross_dbms_test_db', help='Target database name')
+    parser.add_argument('--csv', default='benchmark/bench_sel.csv', help='CSV file path')
+    parser.add_argument('--batch-size', '-b', type=int, default=10000, help='Batch size for INSERT')
+    args = parser.parse_args()
+    
+    # Load config
+    with open(args.config, 'r') as f:
+        data = json.load(f)
+    configs = data['databases'] if isinstance(data, dict) else data
+    
+    # Filter target DBMS
+    if args.dbms.lower() == 'all':
+        targets = configs
+    else:
+        targets = [c for c in configs if c['name'].lower() == args.dbms.lower()]
+        if not targets:
+            print(f"Error: DBMS '{args.dbms}' not found")
+            print(f"Available: {[c['name'] for c in configs]}")
+            sys.exit(1)
+    
+    # Import to each target
+    success = 0
+    for cfg in targets:
+        if import_data(cfg, args.database, args.csv, args.batch_size):
+            success += 1
+    
+    print(f"\nCompleted: {success}/{len(targets)} DBMS targets")
+
+
+if __name__ == '__main__':
+    main()

rosetta/__init__.py

@@ -0,0 +1,3 @@
+"""Rosetta — Cross-DBMS SQL behavioral consistency verification tool."""
+
+__version__ = "1.0.0"

rosetta/__main__.py

@@ -0,0 +1,8 @@
+"""Allow running rosetta as: python -m rosetta"""
+
+import sys
+
+# Use new CLI module
+from .cli.main import main
+
+sys.exit(main())