scripts/tpch_deeplake_ingest.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""
TPC-H Ingestion Script for pg_deeplake
Each table gets its own connection through a load balancer (HAProxy),
which distributes tables across backend instances via round-robin.
Usage:
# Via HAProxy (parallel, one connection per table):
python tpch_deeplake_ingest.py
# Direct to single instance:
python tpch_deeplake_ingest.py --port 5433 --sequential
"""
import argparse
import io
import sys
import time
import psycopg2
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
def log(msg):
print(msg, flush=True)
# TPC-H table definitions
TPCH_TABLES = {
'region': {
'columns': [
('r_regionkey', 'INTEGER'),
('r_name', 'VARCHAR(25)'),
('r_comment', 'VARCHAR(152)')
],
},
'nation': {
'columns': [
('n_nationkey', 'INTEGER'),
('n_name', 'VARCHAR(25)'),
('n_regionkey', 'INTEGER'),
('n_comment', 'VARCHAR(152)')
],
},
'supplier': {
'columns': [
('s_suppkey', 'INTEGER'),
('s_name', 'VARCHAR(25)'),
('s_address', 'VARCHAR(40)'),
('s_nationkey', 'INTEGER'),
('s_phone', 'VARCHAR(15)'),
('s_acctbal', 'DECIMAL(15,2)'),
('s_comment', 'VARCHAR(101)')
],
},
'customer': {
'columns': [
('c_custkey', 'INTEGER'),
('c_name', 'VARCHAR(25)'),
('c_address', 'VARCHAR(40)'),
('c_nationkey', 'INTEGER'),
('c_phone', 'VARCHAR(15)'),
('c_acctbal', 'DECIMAL(15,2)'),
('c_mktsegment', 'VARCHAR(10)'),
('c_comment', 'VARCHAR(117)')
],
},
'part': {
'columns': [
('p_partkey', 'INTEGER'),
('p_name', 'VARCHAR(55)'),
('p_mfgr', 'VARCHAR(25)'),
('p_brand', 'VARCHAR(10)'),
('p_type', 'VARCHAR(25)'),
('p_size', 'INTEGER'),
('p_container', 'VARCHAR(10)'),
('p_retailprice', 'DECIMAL(15,2)'),
('p_comment', 'VARCHAR(23)')
],
},
'partsupp': {
'columns': [
('ps_partkey', 'INTEGER'),
('ps_suppkey', 'INTEGER'),
('ps_availqty', 'INTEGER'),
('ps_supplycost', 'DECIMAL(15,2)'),
('ps_comment', 'VARCHAR(199)')
],
},
'orders': {
'columns': [
('o_orderkey', 'INTEGER'),
('o_custkey', 'INTEGER'),
('o_orderstatus', 'VARCHAR(1)'),
('o_totalprice', 'DECIMAL(15,2)'),
('o_orderdate', 'DATE'),
('o_orderpriority', 'VARCHAR(15)'),
('o_clerk', 'VARCHAR(15)'),
('o_shippriority', 'INTEGER'),
('o_comment', 'VARCHAR(79)')
],
},
'lineitem': {
'columns': [
('l_orderkey', 'INTEGER'),
('l_partkey', 'INTEGER'),
('l_suppkey', 'INTEGER'),
('l_linenumber', 'INTEGER'),
('l_quantity', 'DECIMAL(15,2)'),
('l_extendedprice', 'DECIMAL(15,2)'),
('l_discount', 'DECIMAL(15,2)'),
('l_tax', 'DECIMAL(15,2)'),
('l_returnflag', 'VARCHAR(1)'),
('l_linestatus', 'VARCHAR(1)'),
('l_shipdate', 'DATE'),
('l_commitdate', 'DATE'),
('l_receiptdate', 'DATE'),
('l_shipinstruct', 'VARCHAR(25)'),
('l_shipmode', 'VARCHAR(10)'),
('l_comment', 'VARCHAR(44)')
],
},
}
TABLE_LOAD_ORDER = ['region', 'nation', 'supplier', 'customer', 'part', 'partsupp', 'orders', 'lineitem']
def get_connection(host, port, database, user, password):
return psycopg2.connect(host=host, port=port, database=database, user=user, password=password)
def disable_autovacuum(conn):
old_autocommit = conn.autocommit
conn.autocommit = True
try:
with conn.cursor() as cur:
cur.execute("ALTER SYSTEM SET autovacuum = off;")
cur.execute("SELECT pg_reload_conf();")
finally:
conn.autocommit = old_autocommit
def enable_autovacuum(conn):
old_autocommit = conn.autocommit
conn.autocommit = True
try:
with conn.cursor() as cur:
cur.execute("ALTER SYSTEM SET autovacuum = on;")
cur.execute("SELECT pg_reload_conf();")
finally:
conn.autocommit = old_autocommit
def drop_table(conn, table_name):
with conn.cursor() as cur:
cur.execute(f"DROP TABLE IF EXISTS {table_name} CASCADE;")
conn.commit()
def create_table(conn, table_name, table_def):
columns = table_def['columns']
col_defs = ', '.join([f"{name} {dtype}" for name, dtype, *_ in columns])
sql = f"CREATE TABLE {table_name} ({col_defs}) USING deeplake;"
with conn.cursor() as cur:
cur.execute(sql)
conn.commit()
def table_exists(conn, table_name):
with conn.cursor() as cur:
cur.execute(
"""
SELECT 1
FROM pg_tables
WHERE schemaname = 'public' AND tablename = %s
""",
(table_name,),
)
return cur.fetchone() is not None
CHUNK_SIZE = 5_000_000 # rows per COPY batch to avoid OOM on large tables
def load_data(conn, table_name, data_file):
"""Load data in chunks to avoid OOM on large tables like lineitem (60M rows).
Each chunk is a separate COPY + COMMIT so deeplake flushes to S3
before accumulating the next chunk.
"""
total_lines = 0
chunk_num = 0
with open(data_file, 'r') as f:
while True:
chunk = io.StringIO()
lines_in_chunk = 0
for line in f:
if line.endswith('|\n'):
chunk.write(line[:-2] + '\n')
elif line.endswith('|'):
chunk.write(line[:-1])
else:
chunk.write(line)
lines_in_chunk += 1
if lines_in_chunk >= CHUNK_SIZE:
break
if lines_in_chunk == 0:
break
chunk.seek(0)
with conn.cursor() as cur:
cur.copy_expert(
f"COPY {table_name} FROM STDIN WITH (FORMAT csv, DELIMITER '|')",
chunk
)
conn.commit()
total_lines += lines_in_chunk
chunk_num += 1
if chunk_num > 1 or lines_in_chunk >= CHUNK_SIZE:
log(f" chunk {chunk_num}: {lines_in_chunk:,} rows committed ({total_lines:,} total)")
return total_lines
def get_row_count(conn, table_name):
with conn.cursor() as cur:
cur.execute(f"SELECT COUNT(*) FROM {table_name};")
return cur.fetchone()[0]
def run_vacuum(conn, table_name):
old_autocommit = conn.autocommit
conn.autocommit = True
try:
with conn.cursor() as cur:
cur.execute(f"VACUUM ANALYZE {table_name};")
finally:
conn.autocommit = old_autocommit
def ingest_one_table(table_name, args):
"""Worker: open a new connection (via LB), load one table. Returns (table, rows, seconds)."""
conn = get_connection(args.host, args.port, args.database, args.user, args.password)
try:
conn.autocommit = True
with conn.cursor() as cur:
cur.execute("SELECT inet_server_addr()")
backend_addr = cur.fetchone()[0]
conn.autocommit = False
data_file = Path(args.data_dir) / f"{table_name}.tbl"
if not data_file.exists():
log(f"[{table_name}] Data file not found: {data_file}, skipping")
return (table_name, 0, 0.0)
log(f"[{table_name}] Loading on backend {backend_addr}...")
start_time = time.time()
load_data(conn, table_name, data_file)
elapsed = time.time() - start_time
row_count = get_row_count(conn, table_name)
log(f"[{table_name}] {row_count:,} rows in {elapsed:.1f}s (backend {backend_addr})")
if args.vacuum_after_each:
run_vacuum(conn, table_name)
return (table_name, row_count, elapsed)
except Exception as e:
log(f"[{table_name}] ERROR: {e}")
raise
finally:
try:
conn.rollback()
except Exception:
pass
conn.close()
def main():
parser = argparse.ArgumentParser(description='TPC-H pg_deeplake ingestion')
# Connection
parser.add_argument('--host', default='localhost')
parser.add_argument('--port', type=int, default=5432,
help='Port (default: 5432 for HAProxy)')
parser.add_argument('--database', default='postgres')
parser.add_argument('--user', default='postgres')
parser.add_argument('--password', default='postgres')
# Ingestion
parser.add_argument('--data-dir', default='/home/admin/sasun/work/tpch/tpch_data')
parser.add_argument('--tables', nargs='*', default=None,
help='Specific tables to load (default: all 8)')
parser.add_argument('--skip-create', action='store_true')
parser.add_argument('--replace-existing', action='store_true',
help='DROP+CREATE existing tables before load')
parser.add_argument('--sequential', action='store_true',
help='Load tables sequentially on one connection (no parallel)')
parser.add_argument('--disable-autovacuum', action='store_true')
parser.add_argument('--vacuum-after-each', action='store_true')
args = parser.parse_args()
tables_to_load = args.tables if args.tables else TABLE_LOAD_ORDER
for t in tables_to_load:
if t not in TPCH_TABLES:
parser.error(f"Unknown table: {t}")
log(f"TPC-H pg_deeplake Ingestion")
log(f" Endpoint: {args.host}:{args.port}")
log(f" Database: {args.database}")
log(f" Tables: {', '.join(tables_to_load)}")
log(f" Parallel: {not args.sequential} ({len(tables_to_load)} connections)")
log("")
# Phase 1: Create tables serially on one connection
if not args.skip_create:
log("Phase 1: Creating tables...")
conn = get_connection(args.host, args.port, args.database, args.user, args.password)
try:
if args.disable_autovacuum:
disable_autovacuum(conn)
for table_name in tables_to_load:
table_def = TPCH_TABLES[table_name]
if args.replace_existing:
drop_table(conn, table_name)
create_table(conn, table_name, table_def)
log(f" {table_name}: recreated")
elif not table_exists(conn, table_name):
create_table(conn, table_name, table_def)
log(f" {table_name}: created")
else:
log(f" {table_name}: exists")
finally:
try:
conn.rollback()
except Exception:
pass
conn.close()
log("")
# Phase 2: Load data
overall_start = time.time()
if args.sequential:
# Sequential: one connection, all tables
log("Loading data sequentially...")
results = []
for table_name in tables_to_load:
result = ingest_one_table(table_name, args)
results.append(result)
else:
# Parallel: one connection per table, HAProxy distributes
log(f"Phase 2: Loading data in parallel ({len(tables_to_load)} tables, 1 connection each)...")
results = []
with ThreadPoolExecutor(max_workers=len(tables_to_load)) as executor:
futures = {
executor.submit(ingest_one_table, table_name, args): table_name
for table_name in tables_to_load
}
for future in as_completed(futures):
table_name = futures[future]
try:
results.append(future.result())
except Exception as e:
log(f"[{table_name}] FAILED: {e}")
overall_elapsed = time.time() - overall_start
total_rows = sum(r[1] for r in results)
log(f"\nDone. {total_rows:,} total rows in {overall_elapsed:.1f}s (wall time)")
for table_name, row_count, elapsed in sorted(results, key=lambda x: x[0]):
log(f" {table_name:<12} {row_count:>12,} rows {elapsed:>7.1f}s")
if __name__ == '__main__':
main()