scripts/asterixdb_query_generator.py - main - GraphTreeHeap/museumdb2_experiments

scripts/asterixdb_query_generator.py
128 строк6 KB
FedMam
Added AsterixDB experiment
03 июн 2026, 00:48
03 июн 2026, 00:484c65824
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 + '.adb.sh'
out_queries_file = queries_file[:queries_file.find('.')] + '.adb.sh'

rng = random.Random(hash(dataset_file))

id_of_point = dict()

dataset = []
with open(dataset_file, 'r', encoding='utf-8') as df, \
     open(out_dataset_file, 'w', encoding='utf-8') as odf:
  odf.write(f'''#!/bin/bash

curl -s -X POST http://localhost:19002/query/service --data-urlencode "statement=DROP DATAVERSE {table_name} IF EXISTS;" > /dev/null;
curl -s -X POST http://localhost:19002/query/service --data-urlencode "statement=CREATE DATAVERSE {table_name};";

function execute_query() {{
  Q_RESP=$(curl -s -X POST http://localhost:19002/query/service --data-urlencode "dataverse={table_name}" --data-urlencode "statement=$1")
  STATUS=$(echo "$Q_RESP" | jq -r \'.status\')
  
  if [ "$STATUS" != "success" ]; then
    echo "=== QUERY ERROR ===";
    echo $Q_RESP;
    exit 1;
  fi
}}

execute_query \'CREATE TYPE entry AS CLOSED {{id: int32, location: point, val_str: string}};\'
execute_query \'CREATE DATASET Entries (entry) PRIMARY KEY id;\'

''')

  row_i = 0
  for row in df:
    if row_i % 100 == 0:
      if row_i != 0:
        odf.write('];\'\n')
      odf.write('execute_query \'UPSERT INTO Entries [')
    else:
      odf.write(', ')

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

    odf.write(f'{{"id": {row_i}, "location": create_point({x}.0, {y}.0), "val_str": "{value}"}}')

    id_of_point[(x, y)] = row_i
    row_i += 1
  odf.write(f''']\';

execute_query \'CREATE INDEX rtree_location ON Entries (location) TYPE RTREE;\'
''')

with open(queries_file, 'r', encoding='utf-8') as qf, \
     open(out_queries_file, 'w', encoding='utf-8') as oqf:
  oqf.write(f'''#!/usr/bin/bash

CURRENT_FILE=${{BASH_SOURCE[0]}}
OUTPUT_FILE="${{CURRENT_FILE}}.output"
            
: > $OUTPUT_FILE;

function execute_query() {{
  Q_RESP=$(curl -s -X POST http://localhost:19002/query/service --data-urlencode "dataverse={table_name}" --data-urlencode "statement=$1")
  
  STATUS=$(echo "$Q_RESP" | jq -r \'.status\')
  
  if [ "$STATUS" != "success" ]; then
    echo "=== QUERY ERROR ===";
    echo $Q_RESP;
    exit 1;
  fi
  
  echo -n "$Q_RESP" | jq \'.metrics\' | tr \'\\n\' \' \' >> $OUTPUT_FILE;
  echo "" >> $OUTPUT_FILE;
}}

''')

  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'execute_query \'SELECT VALUE en FROM Entries en WHERE spatial_intersect(en.location, create_point({x}.0, {y}.0));\'\n')
    elif query_type == 'write':
      n_writes = int(query_args[2])
      oqf.write(f'execute_query \'UPSERT INTO Entries [')
      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(', ')
        oqf.write(f'{{"id": {id_of_point[(x, y)]}, "location": create_point({x}.0, {y}.0), "val_str": "{value}"}}')
      oqf.write(f'];\'\n')
    elif query_type == 'delete':
      n_deletes = int(query_args[2])
      oqf.write(f'execute_query \'DELETE FROM Entries AS en WHERE en.id 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'{id_of_point[(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'execute_query \'SELECT VALUE en FROM Entries en WHERE spatial_intersect(en.location, create_rectangle(create_point({left}.0, {top}.0), create_point({right}.0, {bottom}.0)));\'\n')

    row_i += 1