Parquet split

This commit is contained in:
Cyber MacGeddon 2024-09-30 17:54:12 +01:00
parent e0866d9f92
commit 012b6f7672
21 changed files with 55 additions and 7 deletions

View file

@ -1,45 +0,0 @@
#!/usr/bin/env python3
"""
Concatenates multiple parquet files into a single parquet output
"""
import pyarrow as pa
import pyarrow.parquet as pq
import pandas as pd
import sys
import argparse
parser = argparse.ArgumentParser(
prog="combine-parquet",
description=__doc__
)
parser.add_argument(
'-i', '--input',
nargs='*',
help=f'Input files'
)
parser.add_argument(
'-o', '--output',
help=f'Output files'
)
args = parser.parse_args()
df = None
for file in args.input:
part = pq.read_table(file).to_pandas()
if df is None:
df = part
else:
df = pd.concat([df, part], ignore_index=True)
if df is not None:
table = pa.Table.from_pandas(df)
pq.write_table(table, args.output)

View file

@ -1,24 +0,0 @@
#!/usr/bin/env python3
import pyarrow as pa
import pyarrow.csv as pc
import pyarrow.parquet as pq
import pandas as pd
import sys
df = None
for file in sys.argv[1:]:
part = pq.read_table(file).to_pandas()
if df is None:
df = part
else:
df = pd.concat([df, part], ignore_index=True)
if df is not None:
table = pa.Table.from_pandas(df)
pc.write_csv(table, sys.stdout.buffer)

View file

@ -1,6 +0,0 @@
#!/usr/bin/env python3
from trustgraph.dump.graph_embeddings.parquet import run
run()

View file

@ -1,145 +0,0 @@
#!/usr/bin/env python3
"""
Loads Graph embeddings into TrustGraph processing.
"""
import pulsar
from pulsar.schema import JsonSchema
from trustgraph.schema import GraphEmbeddings, Value
from trustgraph.schema import graph_embeddings_store_queue
import argparse
import os
import time
import pyarrow as pa
import pyarrow.parquet as pq
from trustgraph.log_level import LogLevel
class Loader:
def __init__(
self,
pulsar_host,
output_queue,
log_level,
file,
):
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.producer = self.client.create_producer(
topic=output_queue,
schema=JsonSchema(GraphEmbeddings),
chunking_enabled=True,
)
self.file = file
def run(self):
try:
path = self.file
print("Reading file...")
table = pq.read_table(path)
print("Loaded.")
names = set(table.column_names)
if "embeddings" not in names:
print("No 'embeddings' column")
if "entity" not in names:
print("No 'entity' column")
embc = table.column("embeddings")
entc = table.column("entity")
for emb, ent in zip(embc, entc):
b = emb.as_py()
n = ent.as_py()
r = GraphEmbeddings(
vectors=b,
entity=Value(
value=n,
is_uri=n.startswith("https:")
)
)
self.producer.send(r)
except Exception as e:
print(e, flush=True)
def __del__(self):
self.client.close()
def main():
parser = argparse.ArgumentParser(
prog='loader',
description=__doc__,
)
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_output_queue = graph_embeddings_store_queue
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'-o', '--output-queue',
default=default_output_queue,
help=f'Output queue (default: {default_output_queue})'
)
parser.add_argument(
'-l', '--log-level',
type=LogLevel,
default=LogLevel.ERROR,
choices=list(LogLevel),
help=f'Output queue (default: info)'
)
parser.add_argument(
'-f', '--file',
required=True,
help=f'File to load'
)
args = parser.parse_args()
while True:
try:
p = Loader(
pulsar_host=args.pulsar_host,
output_queue=args.output_queue,
log_level=args.log_level,
file=args.file,
)
p.run()
print("File loaded.")
break
except Exception as e:
print("Exception:", e, flush=True)
print("Will retry...", flush=True)
time.sleep(10)
main()

View file

@ -1,144 +0,0 @@
#!/usr/bin/env python3
"""
Loads Graph embeddings into TrustGraph processing.
"""
import pulsar
from pulsar.schema import JsonSchema
from trustgraph.schema import Triple, Value
from trustgraph.schema import triples_store_queue
import argparse
import os
import time
import pyarrow as pa
import pyarrow.parquet as pq
from trustgraph.log_level import LogLevel
class Loader:
def __init__(
self,
pulsar_host,
output_queue,
log_level,
file,
):
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.producer = self.client.create_producer(
topic=output_queue,
schema=JsonSchema(Triple),
chunking_enabled=True,
)
self.file = file
def run(self):
try:
path = self.file
print("Reading file...")
table = pq.read_table(path)
print("Loaded.")
names = set(table.column_names)
if "s" not in names:
print("No 's' column")
if "p" not in names:
print("No 'p' column")
if "o" not in names:
print("No 'o' column")
sc = table.column("s")
pc = table.column("p")
oc = table.column("o")
for s, p, o in zip(sc, pc, oc):
r = Triple(
s=Value(value=s.as_py(), is_uri=True),
p=Value(value=p.as_py(), is_uri=True),
o=Value(value=o.as_py(), is_uri=o.as_py().startswith("https:"))
)
self.producer.send(r)
except Exception as e:
print(e, flush=True)
def __del__(self):
self.client.close()
def main():
parser = argparse.ArgumentParser(
prog='loader',
description=__doc__,
)
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_output_queue = triples_store_queue
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'-o', '--output-queue',
default=default_output_queue,
help=f'Output queue (default: {default_output_queue})'
)
parser.add_argument(
'-l', '--log-level',
type=LogLevel,
default=LogLevel.ERROR,
choices=list(LogLevel),
help=f'Output queue (default: info)'
)
parser.add_argument(
'-f', '--file',
required=True,
help=f'File to load'
)
args = parser.parse_args()
while True:
try:
p = Loader(
pulsar_host=args.pulsar_host,
output_queue=args.output_queue,
log_level=args.log_level,
file=args.file,
)
p.run()
print("File loaded.")
break
except Exception as e:
print("Exception:", e, flush=True)
print("Will retry...", flush=True)
time.sleep(10)
main()

View file

@ -1,6 +0,0 @@
#!/usr/bin/env python3
from trustgraph.dump.triples.parquet import run
run()

View file

@ -51,7 +51,6 @@ setuptools.setup(
"anthropic",
"pyyaml",
"prometheus-client",
"pyarrow",
"cohere",
"openai",
"neo4j",
@ -60,16 +59,13 @@ setuptools.setup(
scripts=[
"scripts/chunker-recursive",
"scripts/chunker-token",
"scripts/concat-parquet",
"scripts/de-query-milvus",
"scripts/de-query-qdrant",
"scripts/de-write-milvus",
"scripts/de-write-qdrant",
"scripts/document-rag",
"scripts/dump-parquet",
"scripts/embeddings-ollama",
"scripts/embeddings-vectorize",
"scripts/ge-dump-parquet",
"scripts/ge-query-milvus",
"scripts/ge-query-qdrant",
"scripts/ge-write-milvus",
@ -81,10 +77,8 @@ setuptools.setup(
"scripts/kg-extract-definitions",
"scripts/kg-extract-topics",
"scripts/kg-extract-relationships",
"scripts/load-graph-embeddings",
"scripts/load-pdf",
"scripts/load-text",
"scripts/load-triples",
"scripts/metering",
"scripts/object-extract-row",
"scripts/oe-write-milvus",
@ -103,7 +97,6 @@ setuptools.setup(
"scripts/text-completion-openai",
"scripts/tg-init-pulsar",
"scripts/tg-processor-state",
"scripts/triples-dump-parquet",
"scripts/triples-query-cassandra",
"scripts/triples-query-neo4j",
"scripts/triples-write-cassandra",

View file

@ -1,3 +0,0 @@
from . processor import *

View file

@ -1,7 +0,0 @@
#!/usr/bin/env python3
from . write import run
if __name__ == '__main__':
run()

View file

@ -1,85 +0,0 @@
"""
Write graph embeddings to parquet files in a directory.
"""
import pulsar
import base64
import os
import argparse
import time
from .... schema import GraphEmbeddings
from .... schema import graph_embeddings_store_queue
from .... base import Consumer
from . writer import ParquetWriter
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = graph_embeddings_store_queue
default_subscriber = module
default_graph_host='localhost'
default_directory = "."
default_file_template = "graph-embeds-{id}.parquet"
default_rotation_time = 60
class Processor(Consumer):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
subscriber = params.get("subscriber", default_subscriber)
directory = params.get("directory", default_directory)
file_template = params.get("file_template", default_file_template)
rotation_time = params.get("rotation_time", default_rotation_time)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": GraphEmbeddings,
}
)
self.writer = ParquetWriter(directory, file_template, rotation_time)
def __del__(self):
if hasattr(self, "writer"):
del self.writer
def handle(self, msg):
v = msg.value()
self.writer.write(v.vectors, v.entity.value)
@staticmethod
def add_args(parser):
Consumer.add_args(
parser, default_input_queue, default_subscriber,
)
parser.add_argument(
'-d', '--directory',
default=default_directory,
help=f'Directory to write to (default: {default_directory})'
)
parser.add_argument(
'-f', '--file-template',
default=default_file_template,
help=f'Directory to write to (default: {default_file_template})'
)
parser.add_argument(
'-t', '--rotation-time',
type=int,
default=default_rotation_time,
help=f'Rotation time / seconds (default: {default_rotation_time})'
)
def run():
Processor.start(module, __doc__)

View file

@ -1,94 +0,0 @@
import threading
import queue
import time
import uuid
import pyarrow as pa
import pyarrow.parquet as pq
class ParquetWriter:
def __init__(self, directory, file_template, rotation_time):
self.directory = directory
self.file_template = file_template
self.rotation_time = rotation_time
self.q = queue.Queue()
self.running = True
self.thread = threading.Thread(target=(self.writer_thread))
self.thread.start()
def writer_thread(self):
items = []
timeout = None
while self.running:
try:
item = self.q.get(timeout=1)
if timeout == None:
timeout = time.time() + self.rotation_time
items.append(item)
except queue.Empty:
pass
if timeout:
if time.time() > timeout:
self.write_file(items)
timeout = None
items = []
def write_file(self, items):
try:
schema = pa.schema([
pa.field('embeddings', pa.list_(pa.list_(pa.float64()))),
pa.field('entity', pa.string()),
])
fname = self.file_template.format(id=str(uuid.uuid4()))
path = f"{self.directory}/{fname}"
writer = pq.ParquetWriter(path, schema)
batch = pa.record_batch(
[
[i[0] for i in items],
[i[1] for i in items],
],
names=['embeddings', 'entity']
)
writer.write_batch(batch)
writer.close()
print(f"Wrote {path}.")
except Exception as e:
print("Parquet write:", e)
def write(self, embeds, ent):
self.q.put((embeds, ent))
def __del__(self):
self.running = False
if hasattr(self, "q"):
self.thread.join()

View file

@ -1,3 +0,0 @@
from . processor import *

View file

@ -1,7 +0,0 @@
#!/usr/bin/env python3
from . write import run
if __name__ == '__main__':
run()

View file

@ -1,85 +0,0 @@
"""
Write graphs triples to parquet files in a directory.
"""
import pulsar
import base64
import os
import argparse
import time
from .... schema import Triple
from .... schema import triples_store_queue
from .... base import Consumer
from . writer import ParquetWriter
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = triples_store_queue
default_subscriber = module
default_graph_host='localhost'
default_directory = "."
default_file_template = "triples-{id}.parquet"
default_rotation_time = 60
class Processor(Consumer):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
subscriber = params.get("subscriber", default_subscriber)
directory = params.get("directory", default_directory)
file_template = params.get("file_template", default_file_template)
rotation_time = params.get("rotation_time", default_rotation_time)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": Triple,
}
)
self.writer = ParquetWriter(directory, file_template, rotation_time)
def __del__(self):
if hasattr(self, "writer"):
del self.writer
def handle(self, msg):
v = msg.value()
self.writer.write(v.s.value, v.p.value, v.o.value)
@staticmethod
def add_args(parser):
Consumer.add_args(
parser, default_input_queue, default_subscriber,
)
parser.add_argument(
'-d', '--directory',
default=default_directory,
help=f'Directory to write to (default: {default_directory})'
)
parser.add_argument(
'-f', '--file-template',
default=default_file_template,
help=f'Directory to write to (default: {default_file_template})'
)
parser.add_argument(
'-t', '--rotation-time',
type=int,
default=default_rotation_time,
help=f'Rotation time / seconds (default: {default_rotation_time})'
)
def run():
Processor.start(module, __doc__)

View file

@ -1,96 +0,0 @@
import threading
import queue
import time
import uuid
import pyarrow as pa
import pyarrow.parquet as pq
class ParquetWriter:
def __init__(self, directory, file_template, rotation_time):
self.directory = directory
self.file_template = file_template
self.rotation_time = rotation_time
self.q = queue.Queue()
self.running = True
self.thread = threading.Thread(target=(self.writer_thread))
self.thread.start()
def writer_thread(self):
triples = []
timeout = None
while self.running:
try:
item = self.q.get(timeout=1)
if timeout == None:
timeout = time.time() + self.rotation_time
triples.append(item)
except queue.Empty:
pass
if timeout:
if time.time() > timeout:
self.write_file(triples)
timeout = None
triples = []
def write_file(self, triples):
try:
schema = pa.schema([
pa.field('s', pa.string()),
pa.field('p', pa.string()),
pa.field('o', pa.string()),
])
fname = self.file_template.format(id=str(uuid.uuid4()))
path = f"{self.directory}/{fname}"
writer = pq.ParquetWriter(path, schema)
batch = pa.record_batch(
[
[tpl[0] for tpl in triples],
[tpl[1] for tpl in triples],
[tpl[2] for tpl in triples],
],
names=['s', 'p', 'o']
)
writer.write_batch(batch)
writer.close()
print(f"Wrote {path}.")
except Exception as e:
print("Parquet write:", e)
def write(self, s, p, o):
self.q.put((s, p, o))
def __del__(self):
self.running = False
if hasattr(self, "q"):
self.thread.join()