Browse Source

create load data on databases scripts.

master
mehrabi 4 months ago
commit
193b0f72bd
  1. 368
      load_postgres.py
  2. 173
      load_redis.py

368
load_postgres.py

@ -0,0 +1,368 @@
import csv
import os
import psycopg2
from psycopg2 import sql
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed
import time
import sys
import random
import tempfile
# ═══════════════════════════════════════════════════════════════
# تنظیمات
# ═══════════════════════════════════════════════════════════════
DB_CONFIG = {
"host": "localhost",
"port": 5432,
"database": "numbers_db",
"user": "myuser",
"password": "mypassword"
}
CSV_FILE = "numbers.csv"
TABLE_NAME = "phone_numbers"
COLUMN_NAME = "phone_number"
RANDOM_COLUMN_NAME = "random_code"
NUM_WORKERS = 12
BATCH_SIZE = 2_000_000
# ═══════════════════════════════════════════════════════════════
# شروع توالی عددی
# ═══════════════════════════════════════════════════════════════
PHONE_START = 100_000_000 # کمترین عدد ۹ رقمی
def optimize_connection(conn):
"""تنظیمات بهینه برای سرعت حداکثری"""
with conn.cursor() as cur:
cur.execute("SET work_mem = '512MB'")
cur.execute("SET synchronous_commit = OFF")
cur.execute("SET maintenance_work_mem = '2GB'")
cur.execute("SET max_parallel_workers_per_gather = 4")
cur.execute("SET random_page_cost = 1.1")
conn.commit()
def get_file_chunks(filename, num_chunks):
"""تقسیم فایل به بخش‌های مساوی"""
print(f"Calculating chunks for {filename}...")
file_size = os.path.getsize(filename)
print(f"File size: {file_size / (1024 ** 3):.2f} GB")
chunk_size = file_size // num_chunks
chunks = []
with open(filename, 'rb') as f:
f.readline()
start = f.tell()
for i in range(num_chunks):
if i < num_chunks - 1:
f.seek(chunk_size * (i + 1))
f.readline()
end = f.tell()
chunks.append((start, end))
start = end
else:
chunks.append((start, file_size))
print(f"Created {len(chunks)} chunks")
return chunks
def count_total_rows():
"""شمارش کل رکوردها"""
print("Counting total rows...")
count = 0
with open(CSV_FILE, 'r') as f:
next(f)
for _ in f:
count += 1
return count
def load_chunk(args):
"""هر Worker یک فایل temp می‌سازد و با COPY لود می‌کند"""
chunk_start, chunk_end, worker_id, rows_in_chunk = args
print(f"[Worker {worker_id}] Starting... chunk: {chunk_start} to {chunk_end}")
conn = psycopg2.connect(**DB_CONFIG)
temp_file = None
try:
optimize_connection(conn)
print(f"[Worker {worker_id}] Connected to DB")
start_time = time.time()
loaded = 0
# ═══════════════════════════════════════════════════════
# محاسبه شروع عدد توالی برای این Worker
# ═══════════════════════════════════════════════════════
phone_start = PHONE_START + (worker_id * rows_in_chunk)
current_phone = phone_start
print(f"[Worker {worker_id}] Phone range: {phone_start:,} to {phone_start + rows_in_chunk - 1:,}")
temp_fd, temp_file = tempfile.mkstemp(suffix='.csv', prefix=f'worker_{worker_id}_')
os.close(temp_fd)
with open(temp_file, 'w', buffering=8192 * 1024) as temp_f:
with open(CSV_FILE, 'r', buffering=8192 * 1024) as f:
f.seek(chunk_start)
# رد کردن خط header اگر در ابتدای chunk باشد
if chunk_start > 0:
f.readline()
while f.tell() < chunk_end:
line = f.readline()
if line.strip():
# ═══════════════════════════════════════
# استفاده از عدد توالی به جای خواندن از CSV
# ═══════════════════════════════════════
random_code = random.randint(1000, 9999)
temp_f.write(f"{current_phone},{random_code}\n")
current_phone += 1
loaded += 1
if loaded % 5_000_000 == 0:
print(f"[Worker {worker_id}] Wrote {loaded:,} rows to temp file")
print(f"[Worker {worker_id}] Temp file created: {loaded:,} rows")
with open(temp_file, 'r') as f:
with conn.cursor() as cur:
cur.copy_expert(
f"COPY {TABLE_NAME} ({COLUMN_NAME}, {RANDOM_COLUMN_NAME}) FROM STDIN WITH CSV",
f
)
conn.commit()
elapsed = time.time() - start_time
rate = loaded / elapsed if elapsed > 0 else 0
print(f"[Worker {worker_id}] Done! {loaded:,} rows in {elapsed:.1f}s ({rate:,.0f} rows/sec)")
return loaded
except Exception as e:
print(f"[Worker {worker_id}] ERROR: {e}")
raise
finally:
conn.close()
if temp_file and os.path.exists(temp_file):
os.remove(temp_file)
def create_table(conn):
"""ساخت جدول Unlogged برای سرعت بیشتر"""
with conn.cursor() as cur:
cur.execute(f"""
DROP TABLE IF EXISTS {TABLE_NAME} CASCADE;
CREATE UNLOGGED TABLE {TABLE_NAME} (
{COLUMN_NAME} BIGINT NOT NULL,
{RANDOM_COLUMN_NAME} INTEGER NOT NULL
) WITH (fillfactor = 100);
""")
conn.commit()
def index_exists(conn, index_name):
"""بررسی وجود ایندکس"""
with conn.cursor() as cur:
cur.execute("""
SELECT EXISTS (
SELECT 1 FROM pg_indexes
WHERE indexname = %s
)
""", (index_name,))
return cur.fetchone()[0]
def create_index(conn, column_name):
"""ساخت ایندکس با بررسی و تأیید"""
index_name = f"idx_{TABLE_NAME}_{column_name}"
print(f"\n>>> Creating index on '{column_name}'...")
if index_exists(conn, index_name):
print(f" Index '{index_name}' already exists, dropping...")
with conn.cursor() as cur:
cur.execute(f"DROP INDEX IF EXISTS {index_name}")
conn.commit()
with conn.cursor() as cur:
start = time.time()
cur.execute(f"CREATE INDEX {index_name} ON {TABLE_NAME} ({column_name})")
conn.commit()
elapsed = time.time() - start
print(f" ✓ Index '{index_name}' created in {elapsed:.1f}s")
if index_exists(conn, index_name):
print(f" ✓ Verified: Index '{index_name}' exists")
return True
else:
print(f" ✗ ERROR: Index '{index_name}' was not created!")
return False
def verify_indexes():
"""تأیید نهایی وجود همه ایندکس‌ها"""
conn = psycopg2.connect(**DB_CONFIG)
try:
print("\n" + "=" * 50)
print("INDEX VERIFICATION")
print("=" * 50)
with conn.cursor() as cur:
cur.execute("""
SELECT indexname, indexdef
FROM pg_indexes
WHERE tablename = %s
ORDER BY indexname
""", (TABLE_NAME,))
indexes = cur.fetchall()
if not indexes:
print("✗ NO INDEXES FOUND!")
return False
print(f"\nFound {len(indexes)} index(es):\n")
for idx_name, idx_def in indexes:
print(f" ✓ {idx_name}")
print(f" Definition: {idx_def}\n")
phone_idx_exists = any('phone_number' in idx[1] for idx in indexes)
random_idx_exists = any('random_code' in idx[1] for idx in indexes)
print("-" * 50)
print(f"phone_number indexed: {'✓ YES' if phone_idx_exists else '✗ NO'}")
print(f"random_code indexed: {'✓ YES' if random_idx_exists else '✗ NO'}")
print("-" * 50)
return phone_idx_exists and random_idx_exists
finally:
conn.close()
def verify_data(start_check=100_000_000, sample_points=None):
"""بررسی صحت داده‌های لود شده"""
if sample_points is None:
sample_points = [0, 10_000_000, 50_000_000, 100_000_000, 124_999_999]
print("\n" + "=" * 50)
print("DATA VERIFICATION")
print("=" * 50)
conn = psycopg2.connect(**DB_CONFIG)
try:
with conn.cursor() as cur:
for point in sample_points:
expected_phone = PHONE_START + point
cur.execute(
f"SELECT {COLUMN_NAME}, {RANDOM_COLUMN_NAME} FROM {TABLE_NAME} LIMIT 1 OFFSET %s",
(point,)
)
row = cur.fetchone()
if row:
status = "" if row[0] == expected_phone else ""
print(f" {status} Offset {point:>12,}: Expected {expected_phone:,}, Got {row[0]:,}")
finally:
conn.close()
def main():
print("=" * 60)
print("PostgreSQL Bulk Loader - Sequential Phone Numbers")
print("=" * 60)
print(f"Phone number range: {PHONE_START:,} to {PHONE_START + 125_000_000 - 1:,}")
print(f"Total records: 125,000,000\n")
num_workers = min(NUM_WORKERS, mp.cpu_count())
total_rows = count_total_rows()
print(f"CSV rows: {total_rows:,}")
print(f"Using {num_workers} workers")
print(f"Batch size: {BATCH_SIZE:,}\n")
start_time = time.time()
try:
# اتصال اولیه و ساخت جدول
conn = psycopg2.connect(**DB_CONFIG)
print("Connected to DB, creating UNLOGGED table...")
optimize_connection(conn)
create_table(conn)
conn.close()
print("Table ready\n")
# تقسیم فایل
print("Calculating file chunks...")
chunks = get_file_chunks(CSV_FILE, num_workers)
# ═══════════════════════════════════════════════════════
# محاسبه تعداد رکوردها در هر chunk
# ═══════════════════════════════════════════════════════
rows_per_chunk = total_rows // num_workers
args = [
(start, end, i, rows_per_chunk)
for i, (start, end) in enumerate(chunks)
]
print("\nStarting workers...\n")
# اجرای موازی
with ProcessPoolExecutor(max_workers=num_workers) as executor:
futures = [executor.submit(load_chunk, arg) for arg in args]
total = 0
for future in as_completed(futures):
try:
total += future.result()
except Exception as e:
print(f"Worker failed: {e}")
elapsed = time.time() - start_time
rate = total / elapsed if elapsed > 0 else 0
remaining = (total_rows - total) / rate if rate > 0 else 0
percent = (total * 100) // total_rows
print(f"\n>>> Progress: {total:,}/{total_rows:,} ({percent}%) | "
f"Speed: {rate:,.0f} rows/sec | "
f"ETA: {remaining / 60:.1f} min")
print(f"\n✓ Total rows loaded: {total:,}")
# ═══════════════════════════════════════════════════════
# ساخت ایندکس‌ها با تأیید
# ═══════════════════════════════════════════════════════
print("\n" + "=" * 60)
print("CREATING INDEXES")
print("=" * 60)
conn = psycopg2.connect(**DB_CONFIG)
success1 = create_index(conn, COLUMN_NAME)
conn.close()
# ═══════════════════════════════════════════════════════
# تأیید نهایی
# ═══════════════════════════════════════════════════════
all_ok = verify_indexes()
# بررسی صحت داده‌ها
verify_data()
total_time = time.time() - start_time
print(f"\n✓ Done! Total time: {total_time / 60:.1f} minutes")
print(f"✓ Average speed: {total_rows / total_time:,.0f} rows/sec")
if all_ok:
print("\n✓✓✓ ALL INDEXES CREATED SUCCESSFULLY ✓✓✓")
else:
print("\n✗✗✗ SOME INDEXES MISSING! ✗✗✗")
except Exception as e:
print(f"\n✗ Error: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()

173
load_redis.py

@ -0,0 +1,173 @@
import os
import redis
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed
import time
import random
# ─────────────────────────────────────────────
# کانفیگ Redis
# ─────────────────────────────────────────────
REDIS_CONFIG = {
"host": "localhost",
"port": 6379,
"db": 0,
"decode_responses": True,
"socket_connect_timeout": 5,
"socket_keepalive": True,
}
NUM_WORKERS = 8
BATCH_SIZE = 100_000
# ═══════════════════════════════════════════════════════════════
# بازه اعداد - مطابق با داده‌های لود شده
# ═══════════════════════════════════════════════════════════════
TOTAL_RECORDS = 125_000_000
PHONE_START = 100_000_000
PHONE_END = PHONE_START + TOTAL_RECORDS - 1 # 224,999,999
def optimize_redis(r):
"""تنظیمات بهینه Redis"""
r.config_set("maxmemory", "16gb")
r.config_set("maxmemory-policy", "allkeys-lru")
r.config_set("save", "")
print("✓ Redis optimized")
def load_chunk_to_redis(args):
"""لود یک بخش از بازه اعداد به Redis"""
worker_id, start_phone, end_phone = args
r = redis.Redis(**REDIS_CONFIG)
try:
print(f"[Worker {worker_id}] Starting... phone: {start_phone:,} to {end_phone:,}")
loaded = 0
batch_count = 0
start_time = time.time()
pipe = r.pipeline()
current_phone = start_phone
while current_phone <= end_phone:
random_value = random.randint(1000, 9999)
pipe.set(str(current_phone), random_value)
loaded += 1
current_phone += 1
if loaded >= BATCH_SIZE * (batch_count + 1):
pipe.execute()
batch_count += 1
elapsed = time.time() - start_time
rate = loaded / elapsed if elapsed > 0 else 0
print(f"[Worker {worker_id}] Batch {batch_count} done ({loaded:,} rows, {rate:,.0f} rows/sec)")
if len(pipe.command_stack) > 0:
pipe.execute()
elapsed = time.time() - start_time
rate = loaded / elapsed if elapsed > 0 else 0
print(f"[Worker {worker_id}] Done! {loaded:,} rows in {elapsed:.1f}s ({rate:,.0f} rows/sec)")
return loaded
except Exception as e:
print(f"[Worker {worker_id}] ERROR: {e}")
raise
finally:
r.close()
def main():
print("=" * 60)
print("Redis Bulk Loader - Sequential Phone Numbers")
print("=" * 60)
print(f"Phone Range: {PHONE_START:,} to {PHONE_END:,}")
print(f"Total Records: {TOTAL_RECORDS:,}\n")
num_workers = min(NUM_WORKERS, mp.cpu_count())
print(f"Using {num_workers} workers\n")
r = redis.Redis(**REDIS_CONFIG)
try:
optimize_redis(r)
print("Clearing existing keys...")
r.flushdb()
finally:
r.close()
start_time = time.time()
try:
# ═══════════════════════════════════════════════════════
# تقسیم بازه اعداد بین workers
# ═══════════════════════════════════════════════════════
records_per_worker = TOTAL_RECORDS // num_workers
args = []
for i in range(num_workers):
if i < num_workers - 1:
start_phone = PHONE_START + (i * records_per_worker)
end_phone = PHONE_START + ((i + 1) * records_per_worker) - 1
else:
# آخرین worker باقیمانده را هم می‌گیرد
start_phone = PHONE_START + (i * records_per_worker)
end_phone = PHONE_END
args.append((i, start_phone, end_phone))
print(f"Worker {i}: {start_phone:,} to {end_phone:,} ({end_phone - start_phone + 1:,} records)")
print("\nStarting workers...\n")
with ProcessPoolExecutor(max_workers=num_workers) as executor:
futures = [executor.submit(load_chunk_to_redis, arg) for arg in args]
total = 0
for future in as_completed(futures):
try:
total += future.result()
except Exception as e:
print(f"Worker failed: {e}")
elapsed = time.time() - start_time
rate = total / elapsed if elapsed > 0 else 0
remaining = (TOTAL_RECORDS - total) / rate if rate > 0 else 0
percent = (total * 100) // TOTAL_RECORDS
print(f"\n>>> Progress: {total:,}/{TOTAL_RECORDS:,} ({percent}%) | "
f"Speed: {rate:,.0f} rows/sec | "
f"ETA: {remaining / 60:.1f} min")
print(f"\n✓ Total rows loaded: {total:,}")
# بررسی نهایی
r = redis.Redis(**REDIS_CONFIG)
try:
print("\n--- Sample Results ---")
sample_phones = [
PHONE_START,
PHONE_START + 10_000_000,
PHONE_START + 50_000_000,
PHONE_END
]
for phone in sample_phones:
value = r.get(str(phone))
print(f" {phone:,} → {value}")
final_count = r.dbsize()
print(f"\n✓ Final count in Redis: {final_count:,}")
finally:
r.close()
total_time = time.time() - start_time
print(f"\n✓ Done! Total time: {total_time / 60:.1f} minutes")
print(f"✓ Average speed: {TOTAL_RECORDS / total_time:,.0f} rows/sec")
except Exception as e:
print(f"\n✗ Error: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()
Loading…
Cancel
Save