non graph solution but extremely memory intensive
This commit is contained in:
nitowa
+135
-102
@@ -1,5 +1,6 @@
|
||||
from gc import collect
|
||||
import json
|
||||
from select import select
|
||||
|
||||
from sqlite3 import Row
|
||||
from typing import Iterable, List
|
||||
@@ -14,6 +15,7 @@ start = time.time()
|
||||
config = json.load(open("./settings.json"))
|
||||
debug = config['debug']
|
||||
|
||||
|
||||
class Master:
|
||||
spark: SparkSession
|
||||
CLUSTERS_TABLE: str
|
||||
@@ -31,127 +33,158 @@ class Master:
|
||||
.config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
|
||||
.getOrCreate()
|
||||
|
||||
def group_tx_addrs(self) -> DataFrame:
|
||||
return self.spark \
|
||||
.read \
|
||||
.table(self.TX_TABLE) \
|
||||
.groupBy("tx_id") \
|
||||
.agg(F.collect_set('address').alias('addresses'))
|
||||
def get_tx_dataframe(self) -> DataFrame:
|
||||
return self.spark.table(self.TX_TABLE)
|
||||
|
||||
def group_cluster_addrs(self) -> DataFrame:
|
||||
return self.spark \
|
||||
.read \
|
||||
.table(self.CLUSTERS_TABLE) \
|
||||
.groupBy("id") \
|
||||
.agg(F.collect_set('address').alias('addresses'))
|
||||
def union_single_col(self, df1: DataFrame, df2: DataFrame, column: str) -> DataFrame:
|
||||
return df1 \
|
||||
.select(column) \
|
||||
.union(df2.select(column))
|
||||
|
||||
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', 'id'])
|
||||
df.writeTo(self.CLUSTERS_TABLE).append()
|
||||
return root
|
||||
def reduce_concat_array_column(self, df: DataFrame, column: str, distinct:bool = False) -> DataFrame:
|
||||
df = self.explode_array_col(df.select(column), column)
|
||||
return self.collect_col_to_array(df, column, distinct)
|
||||
|
||||
def enumerate(self, data: DataFrame) -> DataFrame:
|
||||
return data \
|
||||
def collect_col_to_array(self, df: DataFrame, column: str, distinct: bool = False) -> DataFrame:
|
||||
if(distinct):
|
||||
return df.select(F.collect_set(column).alias(column))
|
||||
else:
|
||||
return df.select(F.collect_list(column).alias(column))
|
||||
|
||||
def explode_array_col(self, df: DataFrame, column: str) -> DataFrame:
|
||||
return df \
|
||||
.rdd \
|
||||
.zipWithIndex() \
|
||||
.toDF(["tx_group", "index"])
|
||||
.flatMap(lambda row: list(map(lambda elem: (elem,), row[column]))) \
|
||||
.toDF([column])
|
||||
|
||||
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('id').isin(cluster_roots)) \
|
||||
.select('address') \
|
||||
.rdd \
|
||||
.map(lambda addr: (addr['address'], new_cluster_root)) \
|
||||
.toDF(['address', 'id']) \
|
||||
def array_col_to_elements(self, df: DataFrame, column: str, distinct:bool = False) -> DataFrame:
|
||||
exploded = master.explode_array_col(
|
||||
df,
|
||||
column
|
||||
)
|
||||
|
||||
if(debug):
|
||||
print("REWRITE JOB")
|
||||
cluster_rewrite.show(truncate=False, vertical=True)
|
||||
print()
|
||||
#this is likely redundant
|
||||
collected = master.collect_col_to_array(
|
||||
exploded,
|
||||
column,
|
||||
distinct
|
||||
)
|
||||
|
||||
return self.explode_array_col(
|
||||
collected,
|
||||
column
|
||||
)
|
||||
|
||||
cluster_rewrite.writeTo(self.CLUSTERS_TABLE).append()
|
||||
# end class Master
|
||||
|
||||
|
||||
"""
|
||||
tuple structure:
|
||||
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['id']]
|
||||
|
||||
if any(x in tx for x in clusteraddresses):
|
||||
return cluster['id']
|
||||
else:
|
||||
return None
|
||||
|
||||
master = Master(config)
|
||||
tx_df = master.get_tx_dataframe()
|
||||
|
||||
tx_addr_groups = master.group_tx_addrs()
|
||||
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()
|
||||
|
||||
if(debug):
|
||||
print("KNOWN CLUSTERS")
|
||||
cluster_addr_groups.show(truncate=True)
|
||||
print()
|
||||
|
||||
tx_addrs: Iterable[str] = tx_groups_indexed \
|
||||
.where(tx_groups_indexed.index == i) \
|
||||
.select('tx_group') \
|
||||
.collect()[0]['tx_group']['addresses']
|
||||
|
||||
if(debug):
|
||||
print("CURRENT TX")
|
||||
print(tx_addrs)
|
||||
print()
|
||||
|
||||
if (cluster_addr_groups.count() == 0):
|
||||
master.insertNewCluster(tx_addrs)
|
||||
continue
|
||||
|
||||
cluster_tx_mapping = cluster_addr_groups \
|
||||
#Turn transactions into a list of ('id', [addr, addr, ...])
|
||||
tx_grouped = tx_df \
|
||||
.groupBy('tx_id') \
|
||||
.agg(F.collect_set('address').alias('addresses')) \
|
||||
.rdd \
|
||||
.map(lambda cluster: (cluster, tx_addrs))
|
||||
|
||||
if(debug):
|
||||
print("cluster_tx_mapping")
|
||||
cluster_tx_mapping \
|
||||
.toDF(['cluster', 'tx']) \
|
||||
.show(truncate=True)
|
||||
print()
|
||||
.zipWithIndex() \
|
||||
.toDF(['tx', 'index']) \
|
||||
.select(
|
||||
F.col('tx.tx_id').alias('tx_id'),
|
||||
F.col('tx.addresses').alias('addresses'),
|
||||
'index'
|
||||
) \
|
||||
.cache()
|
||||
|
||||
|
||||
matched_roots: "List[str]" = cluster_tx_mapping \
|
||||
.map(find) \
|
||||
.filter(lambda root: root != None) \
|
||||
.collect()
|
||||
# TODO: Load clusters from DB, check if any exist, if no make initial cluster, else proceed with loaded data
|
||||
|
||||
if(debug):
|
||||
print("FOUND ROOTS")
|
||||
print(matched_roots)
|
||||
print()
|
||||
# find initial cluster
|
||||
|
||||
# take the first tx
|
||||
tx_zero = tx_grouped \
|
||||
.select(tx_grouped.tx_id, tx_grouped.addresses) \
|
||||
.where(tx_grouped.index == 0)
|
||||
|
||||
# find txs with overlapping addresses
|
||||
overlapping_txs = tx_grouped \
|
||||
.where((tx_grouped.index != 0)) \
|
||||
.join(tx_zero.withColumnRenamed('addresses', 'tx_addresses')) \
|
||||
.select(
|
||||
tx_grouped.index,
|
||||
tx_grouped.addresses,
|
||||
F.arrays_overlap(tx_grouped.addresses, 'tx_addresses').alias('overlap')
|
||||
) \
|
||||
.where(F.col('overlap') == True) \
|
||||
|
||||
# overlapped txs must not be considered anymore, so remove them candidate dataframe
|
||||
tx_grouped = tx_grouped \
|
||||
.join(overlapping_txs, 'index', 'leftanti') \
|
||||
.filter(tx_grouped.index != 0)
|
||||
|
||||
# get the distinct addresses of all overlaps in a single array
|
||||
distinct_addresses = master.reduce_concat_array_column(
|
||||
master.union_single_col(
|
||||
overlapping_txs, tx_zero, column='addresses'
|
||||
),
|
||||
column='addresses',
|
||||
distinct=True,
|
||||
)
|
||||
|
||||
#pick out a random representative for this cluster and add it to every address
|
||||
cluster = distinct_addresses \
|
||||
.rdd \
|
||||
.flatMap(lambda row: list(map(lambda addr: (addr, row['addresses'][0]), row['addresses']))) \
|
||||
.toDF(['address', 'id'])
|
||||
|
||||
# done finding initial cluster
|
||||
|
||||
#group cluster by representative and transform the result into a list of shape ('id', ['addr', 'addr', ...])
|
||||
clusters_grouped = cluster \
|
||||
.groupBy('id') \
|
||||
.agg(F.collect_list('address').alias('addresses'))
|
||||
|
||||
def take_tx_and_cluster(txs: DataFrame, clusters: DataFrame):
|
||||
if (txs.count() == 0): # done!
|
||||
return clusters
|
||||
|
||||
# take a random tx
|
||||
tx = txs \
|
||||
.select('*').limit(1)
|
||||
|
||||
# find clusters with overlapping addresses from tx
|
||||
overlapping_clusters = clusters \
|
||||
.join(tx.withColumnRenamed('addresses', 'tx_addresses')) \
|
||||
.select(
|
||||
clusters.id,
|
||||
clusters.addresses,
|
||||
F.arrays_overlap(clusters.addresses,'tx_addresses').alias('overlap')
|
||||
) \
|
||||
.where(F.col('overlap') == True)
|
||||
|
||||
#collect all addresses into single array field
|
||||
new_cluster_arr = master.reduce_concat_array_column(
|
||||
master.union_single_col(tx, overlapping_clusters, 'addresses'),
|
||||
column='addresses',
|
||||
distinct=True
|
||||
)
|
||||
|
||||
#declare cluster representative
|
||||
new_cluster = new_cluster_arr \
|
||||
.rdd \
|
||||
.flatMap(lambda row: list(map(lambda addr: (addr, row['addresses'][0]), row['addresses']))) \
|
||||
.toDF(['address', 'id']) \
|
||||
.groupBy('id') \
|
||||
.agg(F.collect_list('address').alias('addresses'))
|
||||
|
||||
#start new round with txs minus the one just used, and updated clusters
|
||||
return take_tx_and_cluster(
|
||||
txs.join(tx, 'index', 'leftanti'),
|
||||
clusters.join(overlapping_clusters, 'id', 'leftanti').union(new_cluster)
|
||||
)
|
||||
|
||||
|
||||
if(len(matched_roots) == 0):
|
||||
master.insertNewCluster(tx_addrs)
|
||||
elif(len(matched_roots) == 1):
|
||||
master.insertNewCluster(tx_addrs, matched_roots[0])
|
||||
else:
|
||||
master.rewrite_cluster_id(matched_roots[1:], matched_roots[0])
|
||||
master.insertNewCluster(tx_addrs, matched_roots[0])
|
||||
|
||||
if(debug):
|
||||
print("======================================================================")
|
||||
take_tx_and_cluster(tx_grouped, clusters_grouped).show()
|
||||
|
||||
end = time.time()
|
||||
print("ELAPSED TIME:", end-start)
|
||||
@@ -0,0 +1,157 @@
|
||||
from gc import collect
|
||||
import json
|
||||
|
||||
from sqlite3 import Row
|
||||
from typing import Iterable, List
|
||||
|
||||
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('SparkCassandraApp') \
|
||||
.config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
|
||||
.getOrCreate()
|
||||
|
||||
def group_tx_addrs(self) -> DataFrame:
|
||||
return self.spark \
|
||||
.read \
|
||||
.table(self.TX_TABLE) \
|
||||
.groupBy("tx_id") \
|
||||
.agg(F.collect_set('address').alias('addresses'))
|
||||
|
||||
def group_cluster_addrs(self) -> DataFrame:
|
||||
return self.spark \
|
||||
.read \
|
||||
.table(self.CLUSTERS_TABLE) \
|
||||
.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', 'id'])
|
||||
df.writeTo(self.CLUSTERS_TABLE).append()
|
||||
return root
|
||||
|
||||
def enumerate(self, data: DataFrame) -> DataFrame:
|
||||
return data \
|
||||
.rdd \
|
||||
.zipWithIndex() \
|
||||
.toDF(["tx_group", "index"])
|
||||
|
||||
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('id').isin(cluster_roots)) \
|
||||
.select('address') \
|
||||
.rdd \
|
||||
.map(lambda addr: (addr['address'], new_cluster_root)) \
|
||||
.toDF(['address', 'id']) \
|
||||
|
||||
if(debug):
|
||||
print("REWRITE JOB")
|
||||
cluster_rewrite.show(truncate=False, vertical=True)
|
||||
print()
|
||||
|
||||
cluster_rewrite.writeTo(self.CLUSTERS_TABLE).append()
|
||||
# end class Master
|
||||
|
||||
|
||||
"""
|
||||
tuple structure:
|
||||
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['id']]
|
||||
|
||||
if any(x in tx for x in clusteraddresses):
|
||||
return cluster['id']
|
||||
else:
|
||||
return None
|
||||
|
||||
master = Master(config)
|
||||
|
||||
tx_addr_groups = master.group_tx_addrs()
|
||||
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()
|
||||
|
||||
if(debug):
|
||||
print("KNOWN CLUSTERS")
|
||||
cluster_addr_groups.show(truncate=True)
|
||||
print()
|
||||
|
||||
tx_addrs: Iterable[str] = tx_groups_indexed \
|
||||
.where(tx_groups_indexed.index == i) \
|
||||
.select('tx_group') \
|
||||
.collect()[0]['tx_group']['addresses']
|
||||
|
||||
if(debug):
|
||||
print("CURRENT TX")
|
||||
print(tx_addrs)
|
||||
print()
|
||||
|
||||
if (cluster_addr_groups.count() == 0):
|
||||
master.insertNewCluster(tx_addrs)
|
||||
continue
|
||||
|
||||
cluster_tx_mapping = cluster_addr_groups \
|
||||
.rdd \
|
||||
.map(lambda cluster: (cluster, tx_addrs))
|
||||
|
||||
if(debug):
|
||||
print("cluster_tx_mapping")
|
||||
cluster_tx_mapping \
|
||||
.toDF(['cluster', 'tx']) \
|
||||
.show(truncate=True)
|
||||
print()
|
||||
|
||||
|
||||
matched_roots: "List[str]" = cluster_tx_mapping \
|
||||
.map(find) \
|
||||
.filter(lambda root: root != None) \
|
||||
.collect()
|
||||
|
||||
if(debug):
|
||||
print("FOUND ROOTS")
|
||||
print(matched_roots)
|
||||
print()
|
||||
|
||||
|
||||
if(len(matched_roots) == 0):
|
||||
master.insertNewCluster(tx_addrs)
|
||||
elif(len(matched_roots) == 1):
|
||||
master.insertNewCluster(tx_addrs, matched_roots[0])
|
||||
else:
|
||||
master.rewrite_cluster_id(matched_roots[1:], matched_roots[0])
|
||||
master.insertNewCluster(tx_addrs, matched_roots[0])
|
||||
|
||||
if(debug):
|
||||
print("======================================================================")
|
||||
|
||||
end = time.time()
|
||||
print("ELAPSED TIME:", end-start)
|
||||
Reference in New Issue
Block a user