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
+2 -1
View File
@@ -1,4 +1,5 @@
__pycache__ __pycache__
.vscode .vscode
checkpoints checkpoints
spark-warehouse spark-warehouse
scratchpad.py
+9 -5
View File
@@ -8,13 +8,17 @@ TODO
- Python3 - Python3
- Apache spark 3.2 (https://spark.apache.org/downloads.html) - 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) - graphframes (https://github.com/eejbyfeldt/graphframes/tree/spark-3.3)
## Setting up ## 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. - 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. - 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.
- Start the spark workload by either running `submit.sh` (slow) or `submit_graph.sh` (faster)
# 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( CREATE TABLE clusters(
address TEXT, address TEXT,
parent TEXT, id TEXT,
PRIMARY KEY (address) PRIMARY KEY (address)
); );
+9 -11
View File
@@ -42,13 +42,13 @@ class Master:
return self.spark \ return self.spark \
.read \ .read \
.table(self.CLUSTERS_TABLE) \ .table(self.CLUSTERS_TABLE) \
.groupBy("parent") \ .groupBy("id") \
.agg(F.collect_set('address').alias('addresses')) .agg(F.collect_set('address').alias('addresses'))
def insertNewCluster (self, addrs: Iterable[str], root: str | None = None) -> str: def insertNewCluster (self, addrs: Iterable[str], root: str | None = None) -> str:
if(root == None): if(root == None):
root = addrs[0] 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() df.writeTo(self.CLUSTERS_TABLE).append()
return root return root
@@ -58,14 +58,14 @@ class Master:
.zipWithIndex() \ .zipWithIndex() \
.toDF(["tx_group", "index"]) .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 \ cluster_rewrite = self.spark \
.table(self.CLUSTERS_TABLE) \ .table(self.CLUSTERS_TABLE) \
.where(F.col('parent').isin(cluster_roots)) \ .where(F.col('id').isin(cluster_roots)) \
.select('address') \ .select('address') \
.rdd \ .rdd \
.map(lambda addr: (addr['address'], new_cluster_root)) \ .map(lambda addr: (addr['address'], new_cluster_root)) \
.toDF(['address', 'parent']) \ .toDF(['address', 'id']) \
if(debug): if(debug):
print("REWRITE JOB") print("REWRITE JOB")
@@ -73,24 +73,22 @@ class Master:
print() print()
cluster_rewrite.writeTo(self.CLUSTERS_TABLE).append() cluster_rewrite.writeTo(self.CLUSTERS_TABLE).append()
# end class Master # end class Master
""" """
tuple structure: 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 Iterable[str] => list[addr] | the transaction addresses
""" """
def find(data: tuple[Row, Iterable[str]]) -> str | None: def find(data: tuple[Row, Iterable[str]]) -> str | None:
cluster = data[0] cluster = data[0]
tx = data[1] tx = data[1]
clusteraddresses = cluster['addresses'] + [cluster['parent']] clusteraddresses = cluster['addresses'] + [cluster['id']]
if any(x in tx for x in clusteraddresses): if any(x in tx for x in clusteraddresses):
return cluster['parent'] return cluster['id']
else: else:
return None return None
@@ -149,7 +147,7 @@ for i in range(0, tx_addr_groups.count()):
elif(len(matched_roots) == 1): elif(len(matched_roots) == 1):
master.insertNewCluster(tx_addrs, matched_roots[0]) master.insertNewCluster(tx_addrs, matched_roots[0])
else: 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]) master.insertNewCluster(tx_addrs, matched_roots[0])
if(debug): 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") \ .config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
.getOrCreate() .getOrCreate()
def empty_dataframe(self, schema) -> DataFrame:
return self.spark.createDataFrame(self.spark.sparkContext.emptyRDD(), schema)
def get_tx_dataframe(self) -> DataFrame: def get_tx_dataframe(self) -> DataFrame:
return self.spark.table(self.TX_TABLE) return self.spark.table(self.TX_TABLE)
def get_cluster_dataframe(self) -> DataFrame: def get_cluster_dataframe(self) -> DataFrame:
return self.spark.table(self.CLUSTERS_TABLE) 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 # end class Master
master = Master(config) master = Master(config)
master.spark.sparkContext.setCheckpointDir( master.spark.sparkContext.setCheckpointDir('./checkpoints') # spark is really adamant it needs this even if the algorithm is set to the non-checkpointed version
'./checkpoints') # spark is really adamant it needs this
# Vertex DataFrame tx_df = master.get_tx_dataframe()
transaction_as_vertices = master.get_tx_dataframe() \
transaction_as_vertices = tx_df \
.select('address') \ .select('address') \
.withColumnRenamed('address', 'id') \ .withColumnRenamed('address', 'id') \
.distinct() .distinct()
def explode_row(row: Row) -> List[Row]: def explode_row(row: Row) -> List[Row]:
addresses = row['addresses'] addresses = row['addresses']
if(len(addresses) == 1):
return []
return list(map(lambda addr: (addr, addresses[0]), addresses[1:])) return list(map(lambda addr: (addr, addresses[0]), addresses[1:]))
transactions_as_edges = tx_df \
tx_groups = master.get_tx_dataframe() \
.groupBy("tx_id") \ .groupBy("tx_id") \
.agg(F.collect_set('address').alias('addresses')) .agg(F.collect_set('address').alias('addresses')) \
transactions_as_edges = tx_groups \
.rdd \ .rdd \
.flatMap(explode_row) \ .flatMap(explode_row) \
.toDF(['src', 'dst']) .toDF(['src', 'dst'])
# Create a GraphFrame
g = GraphFrame(transaction_as_vertices, transactions_as_edges) 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: master.write_connected_components_as_clusters(components)
print(sorted(row['collect_list(id)']))
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() end = time.time()
print("ELAPSED TIME:", end-start) print("ELAPSED TIME:", end-start)