Optimizing Parallel Relationship Loading in Graph Databases: The Mix and Batch Technique – 3 of 4

# Thread 1 is creating a relationship from Alice to Bob
def thread_1_operation(session):
    # This locks the 'Alice' node first
    session.run("""
        MATCH (a:Person {name: 'Alice'})
        MATCH (b:Person {name: 'Bob'})
        CREATE (a)-[:KNOWS]->(b)
    """)

# Thread 2 is creating a relationship from Bob to Alice
def thread_2_operation(session):
    # This locks the 'Bob' node first
    session.run("""
        MATCH (b:Person {name: 'Bob'})
        MATCH (a:Person {name: 'Alice'})
        CREATE (b)-[:KNOWS]->(a)
    """)

# Result: Thread 1 locks Alice, waits for Bob
#         Thread 2 locks Bob, waits for Alice
#         DEADLOCK!

# Safe but painfully slow
def load_relationships_sequential(relationships, session):
    for source, target, rel_type in relationships:
        session.run("""
            MATCH (s {id: $source})
            MATCH (t {id: $target})
            CREATE (s)-[r:$type]->(t)
        """, source=source, target=target, type=rel_type)

def create_relationship_with_retry(source, target, rel_type, session, max_retries=5):
    for attempt in range(max_retries):
        try:
            session.run("""
                MATCH (s {id: $source})
                MATCH (t {id: $target})
                CREATE (s)-[r:$type]->(t)
            """, source=source, target=target, type=rel_type)
            return True
        except DeadlockException:
            time.sleep(2 ** attempt)  # Exponential backoff
    return False

def batch_create_relationships(relationships, batch_size=1000):
    for i in range(0, len(relationships), batch_size):
        batch = relationships[i:i+batch_size]
        # This can still deadlock with other batches!
        create_batch(batch)

def partition_nodes(relationships, num_partitions=10):
    """
    Assign each node to a partition using a deterministic function.
    """
    node_partitions = {}
    
    # Extract all unique nodes
    nodes = set()
    for source, target, _ in relationships:
        nodes.add(source)
        nodes.add(target)
    
    # Assign partitions
    for node_id in nodes:
        # Use modulo for numeric IDs, hash for strings
        if isinstance(node_id, (int, float)):
            partition = int(node_id) % num_partitions
        else:
            partition = hash(str(node_id)) % num_partitions
        
        node_partitions[node_id] = partition
    
    return node_partitions

def create_partition_codes(relationships, node_partitions):
    """
    Assign a partition code to each relationship.
    """
    partition_codes = {}
    
    for idx, (source, target, _) in enumerate(relationships):
        source_partition = node_partitions[source]
        target_partition = node_partitions[target]
        
        # Create partition code
        partition_code = f"{source_partition}-{target_partition}"
        partition_codes[idx] = partition_code
    
    return partition_codes

def organize_batches(partition_codes, num_partitions=10):
    """
    Organize relationships into non-conflicting batches.
    """
    # Group relationships by partition code
    code_to_indices = defaultdict(list)
    for idx, code in partition_codes.items():
        code_to_indices[code].append(idx)
    
    batches = []
    
    # Create batches using diagonal pattern
    for offset in range(num_partitions):
        batch = []
        for i in range(num_partitions):
            j = (i + offset) % num_partitions
            code = f"{i}-{j}"
            if code in code_to_indices:
                batch.extend(code_to_indices[code])
        
        if batch:
            batches.append(batch)
    
    return batches

def process_batches(batches, relationships, neo4j_driver, num_workers=8):
    """
    Process batches with guaranteed deadlock-free parallelism.
    """
    total_created = 0
    
    for batch_num, batch in enumerate(batches):
        print(f"Processing batch {batch_num + 1}/{len(batches)}")
        
        # Within this batch, we can parallelize safely!
        with ThreadPoolExecutor(max_workers=num_workers) as executor:
            futures = []
            
            # Split batch into chunks for workers
            chunk_size = max(1, len(batch) // num_workers)
            for i in range(0, len(batch), chunk_size):
                chunk = batch[i:i + chunk_size]
                chunk_rels = [relationships[idx] for idx in chunk]
                
                future = executor.submit(create_relationships_chunk, 
                                       chunk_rels, neo4j_driver)
                futures.append(future)
            
            # Collect results
            for future in as_completed(futures):
                total_created += future.result()
    
    return total_created

import hashlib
import logging
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Tuple, Dict, Any

class MixAndBatchLoader:
    """
    Production-ready Mix and Batch implementation for Neo4j.
    """
    
    def __init__(self, driver, num_partitions=10, concurrency=4):
        """
        Initialize the Mix and Batch loader.
        
        Args:
            driver: Neo4j driver instance
            num_partitions: Number of partitions (affects parallelism)
            concurrency: Number of concurrent workers per batch
        """
        self.driver = driver
        self.num_partitions = num_partitions
        self.concurrency = concurrency
        self.logger = logging.getLogger(__name__)
        
        # Performance metrics
        self.partitioning_time = 0
        self.batching_time = 0
        self.execution_time = 0
    
    def load_relationships(self, relationships: List[Tuple[Any, Any, str, Dict]]):
        """
        Load relationships using Mix and Batch technique.
        
        Args:
            relationships: List of (source_id, target_id, type, properties)
            
        Returns:
            Tuple of (relationships_created, performance_metrics)
        """
        import time
        start_time = time.time()
        
        # Phase 1: Partition nodes
        phase1_start = time.time()
        node_ids = self._extract_node_ids(relationships)
        node_partitions = self._partition_node_ids(node_ids)
        self.partitioning_time = time.time() - phase1_start
        
        self.logger.info(f"Partitioned {len(node_ids)} nodes in {self.partitioning_time:.2f}s")
        
        # Phase 2: Create partition codes
        phase2_start = time.time()
        partition_codes = self._create_partition_codes(relationships, node_partitions)
        
        # Phase 3: Organize batches
        batches = self._organize_batches(partition_codes)
        self.batching_time = time.time() - phase2_start
        
        self.logger.info(f"Organized {len(relationships)} relationships into "
                        f"{len(batches)} batches in {self.batching_time:.2f}s")
        
        # Phase 4: Execute batches
        phase4_start = time.time()
        total_created = self._process_batches(batches, relationships)
        self.execution_time = time.time() - phase4_start
        
        # Calculate metrics
        total_time = time.time() - start_time
        metrics = {
            "partitioning_time": self.partitioning_time,
            "batching_time": self.batching_time,
            "execution_time": self.execution_time,
            "total_time": total_time,
            "relationships_per_second": total_created / total_time if total_time > 0 else 0,
            "batch_count": len(batches),
            "average_batch_size": len(relationships) / len(batches) if batches else 0
        }
        
        return total_created, metrics
    
    def _extract_node_ids(self, relationships):
        """Extract all unique node IDs from relationships."""
        node_ids = set()
        for source, target, _, _ in relationships:
            node_ids.add(source)
            node_ids.add(target)
        return node_ids
    
    def _partition_node_ids(self, node_ids):
        """Assign each node ID to a partition."""
        partitions = {}
        
        for node_id in node_ids:
            # Use consistent hashing for string IDs
            if isinstance(node_id, str):
                hash_value = int(hashlib.md5(node_id.encode()).hexdigest(), 16)
                partition = hash_value % self.num_partitions
            else:
                # Direct modulo for numeric IDs
                partition = int(node_id) % self.num_partitions
            
            partitions[node_id] = partition
        
        return partitions
    
    def _create_partition_codes(self, relationships, node_partitions):
        """Create partition codes for relationships."""
        partition_codes = {}
        
        for idx, (source, target, _, _) in enumerate(relationships):
            source_partition = node_partitions[source]
            target_partition = node_partitions[target]
            
            # Create partition code
            code = f"{source_partition}-{target_partition}"
            partition_codes[idx] = code
        
        return partition_codes
    
    def _organize_batches(self, partition_codes):
        """Organize relationships into non-conflicting batches."""
        # Group by partition code
        code_to_indices = defaultdict(list)
        for idx, code in partition_codes.items():
            code_to_indices[code].append(idx)
        
        batches = []
        
        # Create batches using diagonal pattern
        for offset in range(self.num_partitions):
            batch = []
            
            for i in range(self.num_partitions):
                j = (i + offset) % self.num_partitions
                code = f"{i}-{j}"
                
                if code in code_to_indices:
                    batch.extend(code_to_indices[code])
            
            if batch:
                batches.append(batch)
        
        return batches
    
    def _process_batches(self, batches, relationships):
        """Process batches with parallel execution within each batch."""
        total_created = 0
        
        for batch_idx, batch in enumerate(batches):
            batch_start = time.time()
            
            # Process this batch in parallel
            created = self._process_single_batch(batch, relationships)
            total_created += created
            
            batch_time = time.time() - batch_start
            self.logger.info(f"Batch {batch_idx + 1}/{len(batches)}: "
                           f"{created} relationships in {batch_time:.2f}s "
                           f"({created/batch_time:.0f} rel/s)")
        
        return total_created
    
    def _process_single_batch(self, batch_indices, relationships):
        """Process a single batch with parallel workers."""
        # Divide batch into chunks for workers
        chunk_size = max(1, len(batch_indices) // self.concurrency)
        chunks = []
        
        for i in range(0, len(batch_indices), chunk_size):
            chunk = batch_indices[i:i + chunk_size]
            chunk_rels = [relationships[idx] for idx in chunk]
            chunks.append(chunk_rels)
        
        # Process chunks in parallel
        created = 0
        with ThreadPoolExecutor(max_workers=self.concurrency) as executor:
            futures = [
                executor.submit(self._create_relationships_chunk, chunk)
                for chunk in chunks
            ]
            
            for future in as_completed(futures):
                try:
                    created += future.result()
                except Exception as e:
                    self.logger.error(f"Error in chunk processing: {e}")
        
        return created
    
    def _create_relationships_chunk(self, chunk_relationships):
        """Create a chunk of relationships in a single transaction."""
        with self.driver.session() as session:
            # Prepare batch data
            batch_data = []
            for source, target, rel_type, properties in chunk_relationships:
                batch_data.append({
                    'source': source,
                    'target': target,
                    'type': rel_type,
                    'props': properties or {}
                })
            
            # Execute batch creation
            result = session.run("""
                UNWIND $batch AS rel
                MATCH (source {id: rel.source})
                MATCH (target {id: rel.target})
                CREATE (source)-[r:REL]->(target)
                SET r = rel.props
                SET r.type = rel.type
                RETURN count(r) as created
            """, batch=batch_data)
            
            return result.single()['created']

# Initialize Neo4j driver
from neo4j import GraphDatabase

driver = GraphDatabase.driver("bolt://localhost:7687", 
                            auth=("neo4j", "password"))

# Prepare your relationships
relationships = [
    ("user_1", "product_100", "PURCHASED", {"date": "2024-01-01"}),
    ("user_2", "product_101", "VIEWED", {"timestamp": 1234567890}),
    # ... millions more
]

# Create loader
loader = MixAndBatchLoader(driver, num_partitions=10, concurrency=8)

# Load relationships
created, metrics = loader.load_relationships(relationships)

print(f"Created {created} relationships")
print(f"Performance: {metrics['relationships_per_second']:.0f} rel/s")
print(f"Partitioning: {metrics['partitioning_time']:.2f}s")
print(f"Batching: {metrics['batching_time']:.2f}s")
print(f"Execution: {metrics['execution_time']:.2f}s")

def organize_bipartite_batches(self, partition_codes, set_a_partitions, set_b_partitions):
    """
    Optimized batching for bipartite graphs.
    """
    # We know relationships only go from Set A to Set B
    # This allows for more efficient batching
    
    batches = []
    num_a = len(set_a_partitions)
    num_b = len(set_b_partitions)
    
    # Create batches that maximize parallelism
    for offset in range(max(num_a, num_b)):
        batch = []
        
        for i in range(num_a):
            j = (i + offset) % num_b
            code = f"A{i}-B{j}"
            
            if code in code_to_indices:
                batch.extend(code_to_indices[code])
        
        if batch:
            batches.append(batch)
    
    return batches

def organize_monopartite_batches(self, partition_codes, num_partitions):
    """
    Optimized batching for monopartite graphs with bidirectional relationships.
    """
    # Group relationships by normalized partition codes
    normalized_codes = defaultdict(list)
    
    for idx, code in partition_codes.items():
        parts = code.split('-')
        source_p, target_p = int(parts[0]), int(parts[1])
        
        # Normalize code to handle bidirectional relationships
        normalized = f"{min(source_p, target_p)}-{max(source_p, target_p)}"
        normalized_codes[normalized].append(idx)
    
    # Create batches ensuring no conflicts
    batches = []
    for k in range(num_partitions):
        batch = []
        
        for i in range(num_partitions):
            j = (i + k) % num_partitions
            code = f"{min(i, j)}-{max(i, j)}"
            
            if code in normalized_codes:
                batch.extend(normalized_codes[code])
        
        if batch:
            batches.append(batch)
    
    return batches

def handle_supernodes(self, relationships, threshold=1000):
    """
    Special handling for highly connected nodes.
    """
    # Count connections per node
    node_degree = defaultdict(int)
    for source, target, _, _ in relationships:
        node_degree[source] += 1
        node_degree[target] += 1
    
    # Identify supernodes
    supernodes = {node for node, degree in node_degree.items() 
                  if degree > threshold}
    
    # Separate supernode relationships
    supernode_rels = []
    regular_rels = []
    
    for rel in relationships:
        if rel[0] in supernodes or rel[1] in supernodes:
            supernode_rels.append(rel)
        else:
            regular_rels.append(rel)
    
    # Process with different strategies
    return regular_rels, supernode_rels

def calculate_optimal_partitions(self, relationships):
    """
    Dynamically determine optimal partition count.
    """
    num_nodes = len(self._extract_node_ids(relationships))
    num_relationships = len(relationships)
    
    # Estimate relationship density
    density = num_relationships / (num_nodes ** 2) if num_nodes > 0 else 0
    
    # More partitions for denser graphs
    if density > 0.1:
        return min(32, max(16, int(num_nodes ** 0.25)))
    elif density > 0.01:
        return min(16, max(8, int(num_nodes ** 0.25)))
    else:
        return min(10, max(4, int(num_nodes ** 0.25)))

def process_relationships_streaming(self, relationship_iterator, batch_size=100000):
    """
    Process relationships in a streaming fashion for memory efficiency.
    """
    buffer = []
    total_created = 0
    
    for rel in relationship_iterator:
        buffer.append(rel)
        
        if len(buffer) >= batch_size:
            # Process this chunk
            created, _ = self.load_relationships(buffer)
            total_created += created
            buffer = []
    
    # Don't forget the last chunk
    if buffer:
        created, _ = self.load_relationships(buffer)
        total_created += created
    
    return total_created

def get_diagnostics(self):
    """
    Provide detailed diagnostics for performance tuning.
    """
    return {
        "partition_distribution": self._analyze_partition_distribution(),
        "batch_efficiency": self._calculate_batch_efficiency(),
        "deadlock_count": 0,  # Always zero with Mix and Batch!
        "average_batch_size": sum(len(b) for b in self.batches) / len(self.batches),
        "parallelism_factor": self.concurrency * len(self.batches),
        "theoretical_speedup": self._calculate_theoretical_speedup()
    }