add db write to graph impl

This commit is contained in:
nitowa
2022-08-25 13:42:34 -04:00
parent a614991ff0
commit d42f70d33c
5 changed files with 45 additions and 39 deletions
+1
View File
@@ -2,3 +2,4 @@ __pycache__
.vscode
checkpoints
spark-warehouse
scratchpad.py
+8 -4
View File
@@ -8,13 +8,17 @@ TODO
- Python3
- Apache spark 3.2 (https://spark.apache.org/downloads.html)
- Cassandra DB (https://cassandra.apache.org/_/index.html, locally the docker build is recommended: https://hub.docker.com/_/cassandra)
- Cassandra DB (https://cassandra.apache.org/\_/index.html, locally the docker build is recommended: https://hub.docker.com/\_/cassandra)
For the graph implementation specifically you need to install `graphframes` manually since the official release is incompatible with `spark 3.x` (pull request pending). A prebuilt copy is supplied in the `spark-packages` directory.
For the graph implementation specifically you need to install `graphframes` manually from a third party since the official release is incompatible with `spark 3.x` ([pull request pending](https://github.com/graphframes/graphframes/pull/415)). A prebuilt copy is supplied in the `spark-packages` directory.
- graphframes (https://github.com/eejbyfeldt/graphframes/tree/spark-3.3)
## Setting up
- Modify `settings.json` to reflect your setup. If you are running everything locally you can use `start_services.sh` to turn everything on in one swoop.
- Load the development database by running `python3 setup.py` from the project root.
- Modify `settings.json` to reflect your setup. If you are running everything locally you can use `start_services.sh` to turn everything on in one swoop. It might take a few minutes for Cassandra to become available.
- Load the development database by running `python3 setup.py` from the project root. Per default this will move `small_test_data.csv` into the transactions table.
# Deploying:
- Start the spark workload by either running `submit.sh` (slow) or `submit_graph.sh` (faster)
- If you need to clean out the Database you can run `python3 clean.py`. Be wary that this wipes all data.
+1 -1
View File
@@ -1,5 +1,5 @@
CREATE TABLE clusters(
address TEXT,
parent TEXT,
id TEXT,
PRIMARY KEY (address)
);
+9 -11
View File
@@ -42,13 +42,13 @@ class Master:
return self.spark \
.read \
.table(self.CLUSTERS_TABLE) \
.groupBy("parent") \
.groupBy("id") \
.agg(F.collect_set('address').alias('addresses'))
def insertNewCluster (self, addrs: Iterable[str], root: str | None = None) -> str:
if(root == None):
root = addrs[0]
df = self.spark.createDataFrame(map(lambda addr: (addr, root), addrs), schema=['address', 'parent'])
df = self.spark.createDataFrame(map(lambda addr: (addr, root), addrs), schema=['address', 'id'])
df.writeTo(self.CLUSTERS_TABLE).append()
return root
@@ -58,14 +58,14 @@ class Master:
.zipWithIndex() \
.toDF(["tx_group", "index"])
def rewrite_cluster_parent(self, cluster_roots: Iterable[str], new_cluster_root: str) -> None:
def rewrite_cluster_id(self, cluster_roots: Iterable[str], new_cluster_root: str) -> None:
cluster_rewrite = self.spark \
.table(self.CLUSTERS_TABLE) \
.where(F.col('parent').isin(cluster_roots)) \
.where(F.col('id').isin(cluster_roots)) \
.select('address') \
.rdd \
.map(lambda addr: (addr['address'], new_cluster_root)) \
.toDF(['address', 'parent']) \
.toDF(['address', 'id']) \
if(debug):
print("REWRITE JOB")
@@ -73,24 +73,22 @@ class Master:
print()
cluster_rewrite.writeTo(self.CLUSTERS_TABLE).append()
# end class Master
"""
tuple structure:
Row => Row(parent=addr, addresses=list[addr] | the cluster
Row => Row(id=addr, addresses=list[addr] | the cluster
Iterable[str] => list[addr] | the transaction addresses
"""
def find(data: tuple[Row, Iterable[str]]) -> str | None:
cluster = data[0]
tx = data[1]
clusteraddresses = cluster['addresses'] + [cluster['parent']]
clusteraddresses = cluster['addresses'] + [cluster['id']]
if any(x in tx for x in clusteraddresses):
return cluster['parent']
return cluster['id']
else:
return None
@@ -149,7 +147,7 @@ for i in range(0, tx_addr_groups.count()):
elif(len(matched_roots) == 1):
master.insertNewCluster(tx_addrs, matched_roots[0])
else:
master.rewrite_cluster_parent(matched_roots[1:], matched_roots[0])
master.rewrite_cluster_id(matched_roots[1:], matched_roots[0])
master.insertNewCluster(tx_addrs, matched_roots[0])
if(debug):
+24 -21
View File
@@ -29,52 +29,55 @@ class Master:
.config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
.getOrCreate()
def empty_dataframe(self, schema) -> DataFrame:
return self.spark.createDataFrame(self.spark.sparkContext.emptyRDD(), schema)
def get_tx_dataframe(self) -> DataFrame:
return self.spark.table(self.TX_TABLE)
def get_cluster_dataframe(self) -> DataFrame:
return self.spark.table(self.CLUSTERS_TABLE)
def write_connected_components_as_clusters(self, conn_comp: DataFrame) -> None:
conn_comp \
.withColumnRenamed('id', 'address') \
.withColumnRenamed('component', 'id') \
.writeTo(self.CLUSTERS_TABLE) \
.append()
# end class Master
master = Master(config)
master.spark.sparkContext.setCheckpointDir(
'./checkpoints') # spark is really adamant it needs this
master.spark.sparkContext.setCheckpointDir('./checkpoints') # spark is really adamant it needs this even if the algorithm is set to the non-checkpointed version
# Vertex DataFrame
transaction_as_vertices = master.get_tx_dataframe() \
tx_df = master.get_tx_dataframe()
transaction_as_vertices = tx_df \
.select('address') \
.withColumnRenamed('address', 'id') \
.distinct()
def explode_row(row: Row) -> List[Row]:
addresses = row['addresses']
if(len(addresses) == 1):
return []
return list(map(lambda addr: (addr, addresses[0]), addresses[1:]))
tx_groups = master.get_tx_dataframe() \
transactions_as_edges = tx_df \
.groupBy("tx_id") \
.agg(F.collect_set('address').alias('addresses'))
transactions_as_edges = tx_groups \
.agg(F.collect_set('address').alias('addresses')) \
.rdd \
.flatMap(explode_row) \
.toDF(['src', 'dst'])
# Create a GraphFrame
g = GraphFrame(transaction_as_vertices, transactions_as_edges)
res = g.connectedComponents().groupBy('component').agg(F.collect_list('id')).collect()
components = g.connectedComponents(algorithm='graphframes')
for row in res:
print(sorted(row['collect_list(id)']))
master.write_connected_components_as_clusters(components)
if(debug):
clusters = components \
.groupBy('component') \
.agg(F.collect_list('id')) \
.collect()
for cluster in clusters:
print(sorted(cluster['collect_list(id)']))
end = time.time()
print("ELAPSED TIME:", end-start)