Sort-Merge Join: Algorithm, Implementation, and Optimization

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

def external_merge_sort(input_file: str, output_file: str, 
                          memory_size: int) -> None:
    """
    External merge sort for files larger than available memory.
    
    Args:
        input_file: Path to unsorted input file
        output_file: Path for sorted output
        memory_size: Available memory in bytes for sorting
    """
    # Phase 1: Generate sorted runs
    run_files = generate_sorted_runs(input_file, memory_size)
    
    # Phase 2: Merge runs until single sorted output
    while len(run_files) > 1:
        run_files = merge_runs(run_files, memory_size)
    
    # Final run becomes output
    rename(run_files[0], output_file)
 
 
def generate_sorted_runs(input_file: str, memory_size: int) -> List[str]:
    """
    Phase 1: Read chunks, sort in memory, write sorted runs.
    
    Each run contains approximately (memory_size) bytes of sorted records.
    Uses any efficient in-memory sort (quicksort, heapsort, etc.)
    """
    runs = []
    chunk_buffer = allocate_buffer(memory_size)
    
    with open(input_file, 'rb') as f:
        run_number = 0
        while True:
            # Read memory-sized chunk
            records = read_records_into_buffer(f, chunk_buffer)
            if not records:
                break
            
            # Sort chunk in memory using any O(n log n) algorithm
            records.sort(key=lambda r: r.sort_key)
            
            # Write sorted run to temporary file
            run_file = f"run_{run_number}.tmp"
            write_records_to_file(records, run_file)
            runs.append(run_file)
            run_number += 1
    
    return runs
 
 
def merge_runs(run_files: List[str], memory_size: int) -> List[str]:
    """
    Merge multiple sorted runs into fewer, larger sorted runs.
    
    With B buffer pages, we can merge (B-1) runs at once,
    keeping 1 page for output buffer.
    """
    # Calculate how many runs we can merge at once
    # Reserve one buffer for output, rest for input runs
    buffer_pages = memory_size // PAGE_SIZE
    merge_fan_in = buffer_pages - 1
    
    new_runs = []
    
    # Process runs in groups of (merge_fan_in)
    for i in range(0, len(run_files), merge_fan_in):
        runs_to_merge = run_files[i:i + merge_fan_in]
        merged_run = multi_way_merge(runs_to_merge)
        new_runs.append(merged_run)
        
        # Clean up merged runs
        for run in runs_to_merge:
            delete_file(run)
    
    return new_runs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

def replacement_selection(input_file: str, memory_size: int) -> List[str]:
    """
    Replacement selection produces runs ~2× memory size on average.
    
    Uses a priority queue (min-heap) to always output the smallest
    record that can extend the current run.
    """
    heap = MinHeap(capacity=memory_size)
    runs = []
    current_run_file = None
    last_output_key = None
    
    with open(input_file, 'rb') as f:
        # Initialize heap with first memory-sized chunk
        initial_records = read_records(f, memory_size)
        for record in initial_records:
            heap.insert(record, current_run=True)
        
        while not heap.is_empty():
            # Find smallest record that can extend current run
            # (must be >= last output key)
            candidate = heap.get_min_for_current_run(last_output_key)
            
            if candidate is None:
                # No record can extend current run - start new run
                finalize_run(current_run_file)
                runs.append(current_run_file)
                current_run_file = new_run_file()
                last_output_key = None
                heap.reset_run_markers()
                continue
            
            # Output candidate to current run
            write_record(current_run_file, candidate)
            last_output_key = candidate.sort_key
            heap.remove(candidate)
            
            # Replace with next record from input (if available)
            next_record = read_one_record(f)
            if next_record:
                # Mark whether it can join current run
                can_join_current = next_record.sort_key >= last_output_key
                heap.insert(next_record, current_run=can_join_current)
    
    if current_run_file:
        finalize_run(current_run_file)
        runs.append(current_run_file)
    
    return runs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

import heapq
from dataclasses import dataclass
from typing import List, Iterator, Any
 
@dataclass
class HeapEntry:
    """Entry in merge heap with comparison support."""
    key: Any           # Sort key
    record: bytes      # Full record
    run_index: int     # Which run this came from
    
    def __lt__(self, other):
        return self.key < other.key
 
 
def k_way_merge(sorted_runs: List[str], output_file: str,
                buffer_pages: int) -> None:
    """
    Merge K sorted runs using a min-heap.
    
    Args:
        sorted_runs: List of sorted run file paths
        output_file: Path for merged output
        buffer_pages: Available buffer pages (K+1 needed for K runs)
    
    Time: O(N log K) where N = total records
    I/O: 2N pages (read all + write all)
    """
    K = len(sorted_runs)
    if K > buffer_pages - 1:
        raise ValueError(f"Cannot merge {K} runs with {buffer_pages} buffers")
    
    # Open all input runs
    run_readers = [BufferedReader(run_file) for run_file in sorted_runs]
    output_writer = BufferedWriter(output_file)
    
    # Initialize min-heap with first record from each run
    heap = []
    for i, reader in enumerate(run_readers):
        record = reader.read_next()
        if record:
            entry = HeapEntry(
                key=extract_sort_key(record),
                record=record,
                run_index=i
            )
            heapq.heappush(heap, entry)
    
    # Merge until all runs exhausted
    records_merged = 0
    while heap:
        # Extract minimum record
        min_entry = heapq.heappop(heap)
        output_writer.write(min_entry.record)
        records_merged += 1
        
        # Fetch next record from same run
        next_record = run_readers[min_entry.run_index].read_next()
        if next_record:
            next_entry = HeapEntry(
                key=extract_sort_key(next_record),
                record=next_record,
                run_index=min_entry.run_index
            )
            heapq.heappush(heap, next_entry)
    
    # Flush and close
    output_writer.flush()
    for reader in run_readers:
        reader.close()
    
    return records_merged
 
 
class BufferedReader:
    """Buffered sequential reader for sorted runs."""
    
    def __init__(self, file_path: str, buffer_size: int = 64 * 1024):
        self.file = open(file_path, 'rb')
        self.buffer_size = buffer_size
        self.buffer = b''
        self.position = 0
        self._fill_buffer()
    
    def _fill_buffer(self):
        """Read next chunk from disk into buffer."""
        self.buffer = self.file.read(self.buffer_size)
        self.position = 0
    
    def read_next(self) -> bytes:
        """Read next record from buffer, refilling as needed."""
        if self.position >= len(self.buffer):
            self._fill_buffer()
            if not self.buffer:
                return None  # EOF
        
        # Assume fixed-size records for simplicity
        # Real implementations handle variable-length records
        record_size = RECORD_SIZE
        record = self.buffer[self.position:self.position + record_size]
        self.position += record_size
        return record if len(record) == record_size else None

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

class TournamentTree:
    """
    Loser tree for efficient K-way merge.
    
    Each internal node stores the "loser" of that comparison.
    The winner propagates up. This reduces comparisons per
    output element from 2*log(K) to log(K).
    """
    def __init__(self, sources: List[Iterator]):
        self.K = len(sources)
        self.sources = sources
        
        # Tree has K-1 internal nodes plus K leaf positions
        self.tree = [None] * (2 * self.K)
        self.current_values = [None] * self.K
        
        # Initialize with first value from each source
        for i in range(self.K):
            self.current_values[i] = next(self.sources[i], None)
        
        self._build_tree()
    
    def _build_tree(self):
        """Build initial tournament tree structure."""
        # Leaves are at indices K to 2K-1
        for i in range(self.K):
            self.tree[self.K + i] = i
        
        # Build internal nodes bottom-up
        for i in range(self.K - 1, 0, -1):
            left_child = self.tree[2 * i]
            right_child = self.tree[2 * i + 1]
            winner, loser = self._compete(left_child, right_child)
            self.tree[i] = loser  # Loser stays at this node
            # Winner implicitly propagates to parent
    
    def _compete(self, idx1: int, idx2: int) -> Tuple[int, int]:
        """Compare two source indices, return (winner, loser)."""
        val1 = self.current_values[idx1]
        val2 = self.current_values[idx2]
        
        if val1 is None:
            return (idx2, idx1)
        if val2 is None:
            return (idx1, idx2)
        if val1 <= val2:
            return (idx1, idx2)
        return (idx2, idx1)
    
    def pop_min(self) -> Any:
        """
        Remove and return minimum element.
        Replay tournament for affected path - O(log K).
        """
        # Root contains overall loser; track winner separately
        winner_idx = self._find_overall_winner()
        if self.current_values[winner_idx] is None:
            return None  # All sources exhausted
        
        result = self.current_values[winner_idx]
        
        # Advance winning source
        self.current_values[winner_idx] = next(
            self.sources[winner_idx], None
        )
        
        # Replay from leaf to root
        self._replay(winner_idx)
        
        return result