working graph implementation and improved shell scripts

.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+.vscode
+checkpoints
+spark-warehouse

.vscode/settings.json

@@ -1,5 +0,0 @@
-{
-    "python.analysis.extraPaths": [
-        "./config/db"
-    ]
-}

README.md

@@ -0,0 +1,20 @@
+# Project Description
+
+TODO
+
+# Installation
+
+## Prerequisites:
+
+- 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)
+
+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. 
+- 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.
+- Start the spark workload by either running `submit.sh` (slow) or `submit_graph.sh` (faster)

scratchpad.py

@@ -1,29 +0,0 @@
-import time
-import sys
-import json
-from typing import Dict
-from cassandra.cluster import Cluster
-
-sys.path.append("config/db")
-
-config = json.load(open("./settings.json"))
-
-print(
-    f"Attempting Cassandra connection @ {config['cassandra_addresses']}:{config['cassandra_port']}")
-cluster = Cluster(config['cassandra_addresses'],
-                    port=config['cassandra_port'])
-session = cluster.connect(config['cassandra_keyspace'])
-print(f"Connection OK")
-
-result = session.execute("SELECT * FROM clusters")
-
-map = dict()
-
-for e in result.all():
-    if(e[1] not in map):
-        map[e[1]] = []
-    
-    map[e[1]].append(e[0])
-
-for key in map:
-    print(sorted(map[key]))

settings.json

@@ -3,6 +3,7 @@
     "cassandra_port": 9042,
     "cassandra_keyspace": "distributedunionfind",
     "cassandra_catalog": "DUFCatalog",
+    "cassandra_output_consistency": "ONE",
 
     "setup_db_dir": "config/db",
     "setup_tables_dir": "config/db/tables",
@@ -11,7 +12,10 @@
     "tx_table_name": "transactions",
     "clusters_table_name": "clusters",
 
+    "spark_home": "/home/osboxes/Downloads/spark-3.2.2-bin-hadoop3.2",
     "spark_master": "spark://osboxes:7077",
+    "spark_worker_memory": "1g",
+    "spark_event_logging": "true",
 
-    "debug": true
+    "debug": false
 }

src/spark/main.py

@@ -28,13 +28,7 @@ class Master:
     def makeSparkContext(self,config) -> SparkSession:
         return SparkSession.builder \
         .appName('SparkCassandraApp') \
-        .config('spark.cassandra.connection.host', ','.join(config['cassandra_addresses'])) \
-        .config('spark.cassandra.connection.port', config["cassandra_port"]) \
-        .config('spark.cassandra.output.consistency.level', 'ONE') \
-        .config("spark.sql.extensions",  "com.datastax.spark.connector.CassandraSparkExtensions") \
         .config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
-        .config('directJoinSetting', 'on') \
-        .master(config['spark_master']) \
         .getOrCreate()
 
     def group_tx_addrs(self) -> DataFrame:
@@ -103,7 +97,7 @@ def find(data: tuple[Row, Iterable[str]]) -> str | None:
 master = Master(config)
 
 tx_addr_groups = master.group_tx_addrs()
-tx_groups_indexed = master.enumerate(tx_addr_groups)
+tx_groups_indexed = master.enumerate(tx_addr_groups).cache()
 
 for i in range(0, tx_addr_groups.count()):
     cluster_addr_groups = master.group_cluster_addrs()
@@ -129,7 +123,7 @@ for i in range(0, tx_addr_groups.count()):
 
     cluster_tx_mapping = cluster_addr_groups \
         .rdd \
-        .map(lambda cluster: (cluster, tx_addrs)) 
+        .map(lambda cluster: (cluster, tx_addrs))
 
     if(debug):
         print("cluster_tx_mapping")

src/spark/main_graphs.py

@@ -0,0 +1,80 @@
+from typing import List
+from graphframes import GraphFrame
+import json
+from pyspark.sql import SparkSession, DataFrame, Row
+from pyspark.sql import functions as F
+
+import time
+start = time.time()
+
+
+config = json.load(open("./settings.json"))
+debug = config['debug']
+
+
+class Master:
+    spark: SparkSession
+    CLUSTERS_TABLE: str
+    TX_TABLE: str
+
+    def __init__(self, config):
+        self.spark = self.makeSparkContext(config)
+        self.config = config
+        self.CLUSTERS_TABLE = f"{config['cassandra_catalog']}.{config['cassandra_keyspace']}.{config['clusters_table_name']}"
+        self.TX_TABLE = f"{config['cassandra_catalog']}.{config['cassandra_keyspace']}.{config['tx_table_name']}"
+
+    def makeSparkContext(self, config) -> SparkSession:
+        return SparkSession.builder \
+            .appName('DistributedUnionFindWithGraphs') \
+            .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)
+
+# end class Master
+
+
+master = Master(config)
+master.spark.sparkContext.setCheckpointDir(
+    './checkpoints')  # spark is really adamant it needs this
+
+# Vertex DataFrame
+transaction_as_vertices = master.get_tx_dataframe() \
+    .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() \
+    .groupBy("tx_id") \
+    .agg(F.collect_set('address').alias('addresses'))
+
+transactions_as_edges = tx_groups \
+    .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()
+
+for row in res:
+    print(sorted(row['collect_list(id)']))
+
+end = time.time()
+print("ELAPSED TIME:", end-start)

start_services.sh

@@ -1,9 +1,12 @@
-SPARK_HOME="/home/osboxes/Downloads/spark-3.2.2-bin-hadoop3.2"
-SPARK_MASTER="spark://osboxes:7077"
+SPARK_HOME=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_home"])')
+SPARK_MASTER=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_master"])')
 
 echo "Starting spark master..."
 "$SPARK_HOME"/sbin/start-master.sh
 echo "Starting spark workers..."
 SPARK_WORKER_INSTANCES=5 "$SPARK_HOME"/sbin/start-worker.sh "$SPARK_MASTER"
 echo "Starting cassandra container..."
-docker run -d -p 9042:9042 cassandra
+docker run -d -p 9042:9042 cassandra
+echo "Starting spark history server..."
+mkdir -p /tmp/spark-events
+"$SPARK_HOME"/sbin/start-history-server.sh

submit.sh

@@ -1,14 +1,19 @@
-SPARK_HOME="/home/osboxes/Downloads/spark-3.2.2-bin-hadoop3.2"
-MEMORY="1g"
-SPARK_MASTER="spark://osboxes:7077"
-CASSANDRA_HOST="localhost"
-
+SPARK_HOME=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_home"])')
+MEMORY=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_worker_memory"])')
+SPARK_MASTER=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_master"])')
+CASSANDRA_HOST=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(",".join(config["cassandra_addresses"]))')
+CASSANDRA_PORT=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["cassandra_port"])')
+CASSANDRA_OUT_CONSISTENCY=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["cassandra_output_consistency"])')
+EVENT_LOGGING=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_event_logging"])')
 
 "$SPARK_HOME"/bin/spark-submit \
 --master "$SPARK_MASTER" \
 --conf spark.executor.memory="$MEMORY" \
 --conf spark.cassandra.connection.host="$CASSANDRA_HOST" \
+--conf spark.cassandra.connection.port="$CASSANDRA_PORT" \
+--conf spark.cassandra.output.consistency.level="$CASSANDRA_OUT_CONSISTENCY" \
+--conf spark.eventLog.enabled="$EVENT_LOGGING" \
 --conf spark.sql.session.timeZone=UTC \
 --conf spark.sql.extensions=com.datastax.spark.connector.CassandraSparkExtensions \
 --packages com.datastax.spark:spark-cassandra-connector_2.12:3.2.0 \
-./src/spark/main.py
+./src/spark/main.py

submit_graph.sh

@@ -0,0 +1,20 @@
+SPARK_HOME=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_home"])')
+MEMORY=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_worker_memory"])')
+SPARK_MASTER=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_master"])')
+CASSANDRA_HOST=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(",".join(config["cassandra_addresses"]))')
+CASSANDRA_PORT=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["cassandra_port"])')
+CASSANDRA_OUT_CONSISTENCY=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["cassandra_output_consistency"])')
+EVENT_LOGGING=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_event_logging"])')
+
+"$SPARK_HOME"/bin/spark-submit \
+--master "$SPARK_MASTER" \
+--conf spark.executor.memory="$MEMORY" \
+--conf spark.cassandra.connection.host="$CASSANDRA_HOST" \
+--conf spark.cassandra.connection.port="$CASSANDRA_PORT" \
+--conf spark.cassandra.output.consistency.level="$CASSANDRA_OUT_CONSISTENCY" \
+--conf spark.eventLog.enabled="$EVENT_LOGGING" \
+--conf spark.sql.session.timeZone=UTC \
+--conf spark.sql.extensions=com.datastax.spark.connector.CassandraSparkExtensions \
+--packages com.datastax.spark:spark-cassandra-connector_2.12:3.2.0 \
+--jars ./spark-packages/graphframe_3.3.jar \
+./src/spark/main_graphs.py