Coverage for src/lilbee/data/store/shard_merge.py: 100%

1"""Fold the per-worker shard stores of a multi-GPU ingest into one index.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6from typing import TYPE_CHECKING 

7 

8from lilbee.core.config import CHUNKS_TABLE, INGEST_SOURCE_COLUMNS, META_TABLE, SOURCES_TABLE 

9from lilbee.data.store.lance_helpers import escape_sql_string, table_names 

10 

11if TYPE_CHECKING: 

12 from pathlib import Path 

13 

14 import lancedb 

15 

16 from lilbee.data.store.core import Store 

17 

18log = logging.getLogger(__name__) 

19 

20# Rows read from a shard per append. A chunks row carries its vector, so a whole 

21# shard table does not fit in memory. 

22_MERGE_BATCH_ROWS = 10_000 

23 

24# Source names per ``IN`` predicate when only part of a shard is merged. LanceDB 

25# parses the predicate as one SQL string, so the names are chunked rather than 

26# joined into a single clause per sync. 

27_NAMES_PER_PREDICATE = 500 

28 

29 

30def merge_shards( 

31 store: Store, shard_dirs: list[Path], *, sources: set[str] | None = None 

32) -> dict[str, int]: 

33 """Append every shard's rows into *store*, returning the rows merged per table. 

34 

35 Shards own disjoint sources, so their union is an append with no dedup. A 

36 whole-shard copy (*sources* None) is the fresh-index case. Naming the touched 

37 sources is the re-sync case: the store's own rows for those keys are dropped 

38 first, so re-merging replaces them instead of doubling them. 

39 """ 

40 import lancedb 

41 

42 if sources is not None: 

43 store.remove_documents(sorted(sources)) 

44 merged: dict[str, int] = {} 

45 adopted = _adopt_chunks(store, shard_dirs, sources) 

46 if adopted is not None: 

47 merged[CHUNKS_TABLE] = adopted 

48 for shard_dir in shard_dirs: 

49 database = lancedb.connect(str(shard_dir)) 

50 for name in table_names(database): 

51 # The merged store writes its own meta row from the running config; 

52 # a shard's copy would land beside it as a second row. 

53 if name == META_TABLE: 

54 continue 

55 if name == CHUNKS_TABLE and adopted is not None: 

56 continue # already taken over whole, without reading a row 

57 rows = _copy_table(database.open_table(name), store, name, sources) 

58 merged[name] = merged.get(name, 0) + rows 

59 log.info("Merged %d shard(s): %s", len(shard_dirs), merged) 

60 _reconcile_sources(store, shard_dirs) 

61 return merged 

62 

63 

64def _adopt_chunks(store: Store, shard_dirs: list[Path], sources: set[str] | None) -> int | None: 

65 """Take over every shard's chunk fragments whole; None when that cannot apply. 

66 

67 The chunks table carries the vectors, so it is the whole cost of the merge: 

68 at 8.8M rows by 4096 dims the row copy rewrites about 144GB, and since the 

69 shard stores stay as resume state the corpus then sits on disk twice. 

70 Adopting the fragments is metadata only, and the hard links mean one physical 

71 copy with two names. 

72 

73 Whole-fragment, so only a full merge qualifies: a scoped re-sync names its 

74 sources, and a fragment there holds touched and untouched rows together. 

75 Returns None when the caller should copy rows instead, which also covers a 

76 shard on another filesystem or a data file whose name is already taken; the 

77 merge is correct either way, only slower. 

78 """ 

79 if sources is not None: 

80 return None 

81 tables = [shard_dir / f"{CHUNKS_TABLE}.lance" for shard_dir in shard_dirs] 

82 present = [table for table in tables if table.exists()] 

83 if not present: 

84 return None 

85 try: 

86 return store.adopt_fragments(CHUNKS_TABLE, present) 

87 except OSError as exc: 

88 log.warning("Adopting shard fragments failed (%s); copying rows instead", exc) 

89 return None 

90 

91 

92def _reconcile_sources(store: Store, shard_dirs: list[Path]) -> None: 

93 """Say so when the merged index tracks fewer sources than the workers hold. 

94 

95 A scoped merge only takes what the run touched, so a source a worker holds and 

96 the index does not (an earlier merge that failed, a removal against the index 

97 alone) would otherwise stay missing with nothing to show for it. 

98 """ 

99 import lancedb 

100 

101 held = sum(_source_count(lancedb.connect(str(shard_dir))) for shard_dir in shard_dirs) 

102 merged = _source_count(store.get_db()) 

103 if merged < held: 

104 log.warning( 

105 "The index tracks %d source(s) against %d across the ingest workers. " 

106 "Re-run with --force to fold every worker's shard back in.", 

107 merged, 

108 held, 

109 ) 

110 

111 

112def _source_count(database: lancedb.DBConnection) -> int: 

113 """Rows in a store's source table, zero when it has none.""" 

114 if SOURCES_TABLE not in table_names(database): 

115 return 0 

116 return int(database.open_table(SOURCES_TABLE).count_rows()) 

117 

118 

119def _copy_table( 

120 table: lancedb.table.Table, store: Store, name: str, sources: set[str] | None 

121) -> int: 

122 """Append the rows of *table* that this merge wants into *store*.""" 

123 return sum( 

124 _copy_rows(table, store, name, predicate) for predicate in _predicates(name, sources) 

125 ) 

126 

127 

128def _predicates(name: str, sources: set[str] | None) -> list[str | None]: 

129 """The where-clauses selecting the rows to merge from table *name*. 

130 

131 ``None`` is the whole table. A table with no source column holds corpus-level 

132 aggregates that the post-merge passes rebuild, so a scoped merge skips it. 

133 """ 

134 if sources is None: 

135 return [None] 

136 column = INGEST_SOURCE_COLUMNS.get(name) 

137 if column is None: 

138 return [] 

139 names = sorted(sources) 

140 return [ 

141 _in_predicate(column, names[start : start + _NAMES_PER_PREDICATE]) 

142 for start in range(0, len(names), _NAMES_PER_PREDICATE) 

143 ] 

144 

145 

146def _in_predicate(column: str, names: list[str]) -> str: 

147 """``column IN (...)`` over *names*.""" 

148 quoted = ", ".join(f"'{escape_sql_string(name)}'" for name in names) 

149 return f"{column} IN ({quoted})" 

150 

151 

152def _copy_rows(table: lancedb.table.Table, store: Store, name: str, predicate: str | None) -> int: 

153 """Stream the rows *predicate* selects from *table* into *store*.""" 

154 import pyarrow as pa 

155 

156 query = table.search() 

157 if predicate is not None: 

158 query = query.where(predicate) 

159 reader = query.limit(0).to_batches(_MERGE_BATCH_ROWS) 

160 copied = 0 

161 for batch in reader: 

162 copied += store.absorb_rows(name, pa.Table.from_batches([batch], schema=reader.schema)) 

163 return copied