scripts/psql_script_generator.py - main - GraphTreeHeap/museumdb2_experiments

scripts/psql_script_generator.py
83 строки5 KB
FedMam
Add all scripts
25 май 2026, 19:18
25 май 2026, 19:18a73447a
import sys
import random

ANIMALS = ["Lion", "Tiger", "Elephant", "Giraffe", "Zebra", "Rhino", "Hippo", "Gorilla", "Orangutan", "Chimpanzee", "Kangaroo", "Koala", "Sloth", "Panda", "Moose", "Deer", "Fox", "Wolf", "Coyote", "Bear", "Otter", "Beaver", "Squirrel", "Hare", "Rabbit", "Hedgehog", "Badger", "Skunk", "Raccoon", "Seal", "Walrus", "Dolphin", "Whale", "Porpoise", "Shark", "Tuna", "Salmon", "Trout", "Carp", "Catfish", "Pike", "Mackerel", "Marlin", "Seahorse", "Jellyfish", "Octopus", "Squid", "Crab", "Lobster", "Shrimp", "Butterfly", "Moth", "Bee", "Wasp", "Ant", "Dragonfly", "Grasshopper", "Locust", "Beetle", "Ladybug", "Firefly", "Cicada", "Spider", "Scorpion", "Snake", "Cobra", "Viper", "Python", "Anaconda", "Turtle", "Tortoise", "Iguana", "Gecko", "Chameleon", "Alligator", "Crocodile", "Falcon", "Hawk", "Eagle", "Vulture", "Owl", "Sparrow", "Robin", "Finch", "Pigeon", "Dove", "Parrot", "Cockatoo", "Macaw", "Flamingo", "Pelican", "Heron", "Stork", "Crane", "Swan", "Goose", "Duck", "Turkey", "Chicken", "Peacock", "Emu", "Ostrich", "Penguin", "Platypus"]
ADJECTIVES = ["Agile", "Alert", "Ancient", "Armored", "Astonishing", "Attentive", "Audacious", "Austere", "Baleful", "Bashful", "Bittersweet", "Blazing", "Blithe", "Bold", "Brazen", "Bristling", "Buoyant", "Cunning", "Capricious", "Calm", "Carnivorous", "Celestial", "Charming", "Chromatic", "Clever", "Cloaked", "Cold-blooded", "Colossal", "Cozy", "Cryptic", "Dainty", "Dauntless", "Deafening", "Deft", "Delicate", "Diminutive", "Dire", "Dapper", "Dashing", "Draconic", "Dreamlike", "Durable", "Eerie", "Elusive", "Elegant", "Endearing", "Enigmatic", "Ethereal", "Ferocious", "Fierce", "Fleet", "Fluffy", "Formidable", "Forsaken", "Frosty", "Furtive", "Gallant", "Gargantuan", "Gleaming", "Golden", "Graceful", "Grizzled", "Harmonious", "Hardy", "Hypnotic", "Imposing", "Incandescent", "Industrious", "Ingenious", "Intrepid", "Iridescent", "Jagged", "Jovial", "Keen", "Lithe", "Luminous", "Majestic", "Merciless", "Mottled", "Mystical", "Nocturnal", "Noble", "Omnivorous", "Ornate", "Outsized", "Pensive", "Placid", "Primal", "Prismatic", "Prowling", "Puny", "Radiant", "Regal", "Rugged", "Sinuous", "Sleek", "Shocked", "Stealthy", "Striking", "Tenacious"]

dataset_file = sys.argv[1]
queries_file = sys.argv[2]

table_name = dataset_file[:dataset_file.find('.')]

out_dataset_file = table_name + '.sql'
out_queries_file = queries_file[:queries_file.find('.')] + '.sql'

rng = random.Random(hash(dataset_file))

dataset = []
with open(dataset_file, 'r', encoding='utf-8') as df, \
     open(out_dataset_file, 'w', encoding='utf-8') as odf:
  odf.write(f'''CREATE TABLE {table_name} (
  point GEOMETRY NOT NULL PRIMARY KEY,
  value TEXT)''')

  row_i = 0
  for row in df:
    if row_i % 1000 == 0:
      odf.write(f'''{"""
ON CONFLICT DO NOTHING""" if row_i != 0 else ""};

INSERT INTO {table_name} VALUES\n''')
    else:
      odf.write(',\n')

    x, y = map(int, row.strip().split('\t'))
    value = rng.choice(ADJECTIVES) + ' ' + rng.choice(ANIMALS)

    odf.write(f'  (ST_MakePoint({x}, {y}), \'{value}\')')

    row_i += 1
  odf.write(f'''
ON CONFLICT DO NOTHING;

CREATE INDEX {table_name}_spatial_index ON {table_name} USING GIST (point);\n''')

with open(queries_file, 'r', encoding='utf-8') as qf, \
     open(out_queries_file, 'w', encoding='utf-8') as oqf:
  for row in qf:
    query_args = row.strip().split('\t')
    query_type = query_args[0]

    if query_type == 'read':
      _, _, x, y = query_args
      x, y = int(x), int(y)

      oqf.write(f'EXPLAIN ANALYZE SELECT * FROM {table_name} WHERE point = ST_MakePoint({x}, {y});\n')
    elif query_type == 'write':
      n_writes = int(query_args[2])
      oqf.write(f'EXPLAIN ANALYZE INSERT INTO {table_name} VALUES\n')
      already_been = set()
      for i in range(n_writes):
        x, y, value = query_args[3 + 3 * i: 3 + 3 * i + 3]
        x, y = int(x), int(y)
        if (x, y) in already_been:
          continue
        already_been.add((x, y))
        if i > 0:
          oqf.write(',\n')
        oqf.write(f'  (ST_MakePoint({x}, {y}), \'{value}\')')
      oqf.write(f'\nON CONFLICT (point) DO UPDATE SET value = EXCLUDED.value;\n')
    elif query_type == 'delete':
      n_deletes = int(query_args[2])
      oqf.write(f'EXPLAIN ANALYZE DELETE FROM {table_name} WHERE point IN (')
      for i in range(n_deletes):
        x, y = map(int, query_args[3 + 2 * i: 3 + 2 * i + 2])
        if i > 0:
          oqf.write(', ')
        oqf.write(f'ST_MakePoint({x}, {y})')
      oqf.write(');\n')
    elif query_type == 'range':
      _, _, left, top, right, bottom, _ = query_args
      left, top, right, bottom = map(int, (left, top, right, bottom))

      oqf.write(f'EXPLAIN ANALYZE SELECT * FROM {table_name} WHERE ST_Contains(ST_MakeEnvelope({left}, {top}, {right}, {bottom}, 0), point);\n')