T1: Deduplicator compile fix + unit tests [checkpoint]
This commit is contained in:
Artur Kruszewski
@@ -0,0 +1,216 @@
|
||||
package com.developx.szpitale.load;
|
||||
|
||||
import com.developx.szpitale.model.*;
|
||||
import org.neo4j.driver.Driver;
|
||||
import org.neo4j.driver.Session;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import java.util.*;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
public class GraphLoader {
|
||||
|
||||
private static final Logger log = LoggerFactory.getLogger(GraphLoader.class);
|
||||
private static final int BATCH_SIZE = 1000;
|
||||
|
||||
private final Driver driver;
|
||||
|
||||
public GraphLoader(Driver driver) {
|
||||
this.driver = driver;
|
||||
}
|
||||
|
||||
public LoadStats load(List<CanonicalDataset> datasets) {
|
||||
LoadStats stats = new LoadStats();
|
||||
try (Session session = driver.session()) {
|
||||
for (CanonicalDataset dataset : datasets) {
|
||||
stats.addVoivodeship(dataset.voivodeship());
|
||||
loadHospitals(session, dataset.hospitals(), stats);
|
||||
loadPeople(session, dataset.people(), stats);
|
||||
loadRoles(session, dataset.roles(), stats);
|
||||
loadAffiliations(session, dataset.affiliations(), stats);
|
||||
loadMandates(session, dataset.mandates(), stats);
|
||||
}
|
||||
}
|
||||
return stats;
|
||||
}
|
||||
|
||||
private void loadHospitals(Session session, List<Hospital> hospitals, LoadStats stats) {
|
||||
List<Map<String, Object>> batch = new ArrayList<>();
|
||||
for (Hospital h : hospitals) {
|
||||
batch.add(Map.of(
|
||||
"id", h.id(),
|
||||
"name", h.name(),
|
||||
"shortName", h.shortName(),
|
||||
"city", h.city(),
|
||||
"voivodeship", h.voivodeship(),
|
||||
"legalForm", h.legalForm() != null ? h.legalForm().name() : null,
|
||||
"foundingBody", h.foundingBody(),
|
||||
"supervisoryBodyType", h.supervisoryBodyType() != null ? h.supervisoryBodyType().name() : null,
|
||||
"website", h.website(),
|
||||
"sourceUrls", h.sourceUrls()
|
||||
));
|
||||
}
|
||||
|
||||
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
|
||||
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
|
||||
String cypher = "UNWIND $batch AS h " +
|
||||
"MERGE (hospital:Hospital {id: h.id}) " +
|
||||
"SET hospital += {name:h.name, shortName:h.shortName, city:h.city, " +
|
||||
" voivodeship:h.voivodeship, legalForm:h.legalForm, " +
|
||||
" foundingBody:h.foundingBody, supervisoryBodyType:h.supervisoryBodyType, " +
|
||||
" website:h.website} " +
|
||||
"MERGE (v:Voivodeship {slug: h.voivodeship}) " +
|
||||
"MERGE (hospital)-[:W_WOJEWODZTWIE]->(v)";
|
||||
session.run(cypher, Map.of("batch", subList));
|
||||
stats.addHospitals(subList.size());
|
||||
}
|
||||
}
|
||||
|
||||
private void loadPeople(Session session, List<Person> people, LoadStats stats) {
|
||||
List<Map<String, Object>> batch = new ArrayList<>();
|
||||
for (Person p : people) {
|
||||
batch.add(Map.of(
|
||||
"id", p.id(),
|
||||
"fullName", p.fullName(),
|
||||
"displayName", p.displayName(),
|
||||
"titles", p.titles() != null ? p.titles() : Collections.emptyList(),
|
||||
"sourceUrls", p.sourceUrls()
|
||||
));
|
||||
}
|
||||
|
||||
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
|
||||
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
|
||||
String cypher = "UNWIND $batch AS p " +
|
||||
"MERGE (person:Person {id: p.id}) " +
|
||||
"SET person += {fullName:p.fullName, displayName:p.displayName, " +
|
||||
" titles:p.titles, sourceUrls:p.sourceUrls}";
|
||||
session.run(cypher, Map.of("batch", subList));
|
||||
stats.addPeople(subList.size());
|
||||
}
|
||||
}
|
||||
|
||||
private void loadRoles(Session session, List<Role> roles, LoadStats stats) {
|
||||
Map<String, List<Role>> byRelType = new LinkedHashMap<>();
|
||||
for (Role r : roles) {
|
||||
String relType = roleTypeToRelName(r);
|
||||
byRelType.computeIfAbsent(relType, k -> new ArrayList<>()).add(r);
|
||||
}
|
||||
|
||||
for (Map.Entry<String, List<Role>> entry : byRelType.entrySet()) {
|
||||
String relType = entry.getKey();
|
||||
List<Role> items = entry.getValue();
|
||||
|
||||
for (int i = 0; i < items.size(); i += BATCH_SIZE) {
|
||||
List<Role> subList = items.subList(i, Math.min(i + BATCH_SIZE, items.size()));
|
||||
String cypher = "UNWIND $batch AS r " +
|
||||
"MATCH (h:Hospital {id: r.hospitalId}) " +
|
||||
"MATCH (p:Person {id: r.personId}) " +
|
||||
"MERGE (h)-[rrel:`" + escapeBacktick(relType) + "`]->(p) " +
|
||||
"SET rrel += {roleType: r.roleType, roleLabel: r.roleLabel, " +
|
||||
" organ: r.organ, status: r.status, " +
|
||||
" validFrom: r.validFrom, validTo: r.validTo, note: r.note}";
|
||||
session.run(cypher, Map.of("batch", subList.stream()
|
||||
.map(r -> Map.of(
|
||||
"hospitalId", r.hospitalId(),
|
||||
"personId", r.personId(),
|
||||
"roleType", r.roleType().name(),
|
||||
"roleLabel", r.roleLabel(),
|
||||
"organ", r.organ().name(),
|
||||
"status", r.status().name(),
|
||||
"validFrom", r.validFrom(),
|
||||
"validTo", r.validTo(),
|
||||
"note", r.note()
|
||||
))
|
||||
.collect(Collectors.toList())));
|
||||
stats.addRoles(subList.size());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private void loadAffiliations(Session session, List<Affiliation> affiliations, LoadStats stats) {
|
||||
for (int i = 0; i < affiliations.size(); i += BATCH_SIZE) {
|
||||
List<Affiliation> subList = affiliations.subList(
|
||||
i, Math.min(i + BATCH_SIZE, affiliations.size()));
|
||||
String cypher = "UNWIND $batch AS a " +
|
||||
"MATCH (p:Person {id: a.personId}) " +
|
||||
"MERGE (party:Party {name: a.name}) SET party.type = a.type " +
|
||||
"MERGE (p)-[m:CZLONEK_PARTII]->(party) " +
|
||||
"SET m += {confidence: a.confidence, note: a.note, sourceUrls: a.sourceUrls}";
|
||||
session.run(cypher, Map.of("batch", subList.stream()
|
||||
.map(a -> Map.of(
|
||||
"personId", a.personId(),
|
||||
"name", a.name(),
|
||||
"type", a.type().name(),
|
||||
"confidence", a.confidence().name(),
|
||||
"note", a.note(),
|
||||
"sourceUrls", a.sourceUrls()
|
||||
))
|
||||
.collect(Collectors.toList())));
|
||||
stats.addAffiliations(subList.size());
|
||||
}
|
||||
}
|
||||
|
||||
private void loadMandates(Session session, List<Mandate> mandates, LoadStats stats) {
|
||||
if (mandates.isEmpty()) return;
|
||||
|
||||
for (int i = 0; i < mandates.size(); i += BATCH_SIZE) {
|
||||
List<Mandate> subList = mandates.subList(
|
||||
i, Math.min(i + BATCH_SIZE, mandates.size()));
|
||||
String cypher = "UNWIND $batch AS m " +
|
||||
"MATCH (p:Person {id: m.personId}) " +
|
||||
"MERGE (g:GovBody {name: m.body}) " +
|
||||
" SET g.type = m.mandateType, g.voivodeship = m.voivodeship, g.term = m.term " +
|
||||
"MERGE (p)-[:PELNI_MANDAT {mandateType: m.mandateType, term: m.term}]->(g)";
|
||||
session.run(cypher, Map.of("batch", subList.stream()
|
||||
.map(m -> Map.of(
|
||||
"personId", m.personId(),
|
||||
"body", m.body(),
|
||||
"mandateType", m.mandateType().name(),
|
||||
"term", m.term(),
|
||||
"voivodeship", m.voivodeship(),
|
||||
"sourceUrls", m.sourceUrls()
|
||||
))
|
||||
.collect(Collectors.toList())));
|
||||
stats.addMandates(subList.size());
|
||||
}
|
||||
}
|
||||
|
||||
private String roleTypeToRelName(Role r) {
|
||||
String organ = r.organ().name();
|
||||
String role = r.roleType().name();
|
||||
|
||||
if ("DYREKCJA".equals(organ) && "DYREKTOR".equals(role)) return "DYREKTOR";
|
||||
if ("DYREKCJA".equals(organ) && ("ZASTEPCA_DYREKTORA".equals(role) || "GLOWNY_KSIEGOWY".equals(role) || "PRZELOZONA_PIELEGNIARKOW".equals(role))) return "ZASTEPCA_DYREKTORA";
|
||||
if ("PRZEWODNICZACY_ORGANU".equals(role)) return "PRZEWODNICZY_ORGANOWI";
|
||||
return "CZLONEK_ORGANU";
|
||||
}
|
||||
|
||||
private String escapeBacktick(String s) {
|
||||
return s.replace("`", "\\`");
|
||||
}
|
||||
|
||||
public static class LoadStats {
|
||||
public int hospitalCount;
|
||||
public int personCount;
|
||||
public int roleCount;
|
||||
public int affiliationCount;
|
||||
public int mandateCount;
|
||||
public final List<String> voivodeships = new ArrayList<>();
|
||||
|
||||
public void addVoivodeship(String v) { voivodeships.add(v); }
|
||||
public void addHospitals(int count) { hospitalCount += count; }
|
||||
public void addPeople(int count) { personCount += count; }
|
||||
public void addRoles(int count) { roleCount += count; }
|
||||
public void addAffiliations(int count) { affiliationCount += count; }
|
||||
public void addMandates(int count) { mandateCount += count; }
|
||||
|
||||
@Override
|
||||
public String toString() {
|
||||
return String.format(
|
||||
"Load stats - voivodeships: %s, hospitals: %d, people: %d, roles: %d, affiliations: %d, mandates: %d",
|
||||
String.join(", ", voivodeships), hospitalCount, personCount, roleCount, affiliationCount, mandateCount
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user