1 files changed, 28 insertions, 21 deletions
diff --git a/backend/data/genEnwikiData.py b/backend/data/genEnwikiData.py
index 48fd2c6..879ecf6 100755
--- a/backend/data/genEnwikiData.py
+++ b/backend/data/genEnwikiData.py
@@ -4,15 +4,14 @@ import sys, re
 import sqlite3
 
 usageInfo =  f"usage: {sys.argv[0]}\n"
-usageInfo += "Reads Wikimedia enwiki data from enwiki/, a list of node names,"
-usageInfo += "and node and name data from a sqlite database, and adds\n"
-usageInfo += "description data for names that don't have them\n"
+usageInfo += "Reads Wikimedia enwiki data from enwiki/, and node and name data"
+usageInfo += "from a sqlite database, and adds description data for names that\n"
+usageInfo += "don't have them.\n"
 if len(sys.argv) > 1:
 	print(usageInfo, file=sys.stderr)
 	sys.exit(1)
 
 enwikiDb = "enwiki/enwikiData.db"
-namesFile = "reducedTol/names.txt"
 dbFile = "data.db"
 
 # Open dbs
@@ -20,40 +19,48 @@ enwikiCon = sqlite3.connect(enwikiDb)
 enwikiCur = enwikiCon.cursor()
 dbCon = sqlite3.connect(dbFile)
 dbCur = dbCon.cursor()
-# Read in names to check
-print("Getting names to check")
+# Get node names without descriptions
+print("Getting node names")
 nodeNames = set()
-with open(namesFile) as file:
-	for line in file:
-		nodeNames.add(line.rstrip())
+query = "SELECT nodes.name FROM nodes LEFT JOIN descs ON nodes.name = descs.name WHERE desc IS NULL"
+for row in dbCur.execute(query):
+	nodeNames.add(row[0])
 print("Found {} names".format(len(nodeNames)))
-# Remove names that have descriptions
-print("Checking for existing name descriptions")
-namesWithDescs = set()
-for name in nodeNames:
-	row = dbCur.execute("SELECT name FROM descs where name = ?", (name,)).fetchone()
-	if row != None:
-		namesWithDescs.add(name)
-nodeNames.difference_update(namesWithDescs)
-print("Remaining nodes: {}".format(len(nodeNames)))
 # Find page id for each node name
-nodeToPageId = {}
 print("Getting node page-ids")
+nodeToPageId = {}
+iterNum = 0
 for name in nodeNames:
-	row = enwikiCur.execute("SELECT id FROM pages where pages.title = ? COLLATE nocase", (name,)).fetchone()
+	iterNum += 1
+	if iterNum % 1e4 == 0:
+		print("At iteration {}".format(iterNum))
+	#
+	row = enwikiCur.execute("SELECT id FROM pages WHERE pages.title = ? COLLATE NOCASE", (name,)).fetchone()
 	if row != None:
 		nodeToPageId[name] = row[0]
 # Resolve redirects
 print("Resolving redirects")
 redirectingNames = set()
+iterNum = 0
 for (name, pageId) in nodeToPageId.items():
-	row = enwikiCur.execute("SELECT target_id FROM redirects where redirects.id = ?", (pageId,)).fetchone()
+	iterNum += 1
+	if iterNum % 1000 == 0:
+		print("At iteration {}".format(iterNum))
+	#
+	row = enwikiCur.execute(
+		"SELECT pages.id FROM redirects INNER JOIN pages ON redirects.target = pages.title WHERE redirects.id = ?",
+		(pageId,)).fetchone()
 	if row != None:
 		nodeToPageId[name] = row[0]
 		redirectingNames.add(name)
 # Add descriptions for each node
 print("Adding description data")
+iterNum = 0
 for (name, pageId) in nodeToPageId.items():
+	iterNum += 1
+	if iterNum % 1000 == 0:
+		print("At iteration {}".format(iterNum))
+	#
 	row = enwikiCur.execute("SELECT desc FROM descs where descs.id = ?", (pageId,)).fetchone()
 	if row != None:
 		dbCur.execute("INSERT INTO descs VALUES (?, ?, ?)", (name, row[0], 1 if name in redirectingNames else 0))