1use crate::SourceRow;
2use crate::index::{KeyType, SourceIndex};
3use crate::{SourceDefinition, open_source_reader};
4use anyhow::{Context, Result};
5use apif_utils::FileUtils;
6use std::path::{Path, PathBuf};
7use std::sync::LazyLock;
8use std::sync::atomic::{AtomicU64, Ordering};
9
10const DEFAULT_MEMORY_LIMIT: u64 = 256 * 1024 * 1024; #[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum BuildPhase {
14 Scan,
15 Write,
16}
17
18#[derive(Debug, Default)]
19pub struct IndexMetrics {
20 pub builds_total: AtomicU64,
21 pub builds_failed: AtomicU64,
22}
23
24impl IndexMetrics {
25 pub fn record_build_success(&self) {
26 self.builds_total.fetch_add(1, Ordering::Relaxed);
27 }
28
29 pub fn record_build_failure(&self) {
30 self.builds_failed.fetch_add(1, Ordering::Relaxed);
31 }
32
33 pub fn builds_total(&self) -> u64 {
34 self.builds_total.load(Ordering::Relaxed)
35 }
36
37 pub fn builds_failed(&self) -> u64 {
38 self.builds_failed.load(Ordering::Relaxed)
39 }
40}
41
42pub static INDEX_METRICS: LazyLock<IndexMetrics, fn() -> IndexMetrics> =
43 LazyLock::new(IndexMetrics::default);
44
45pub fn index_path_for_source(source_path: &Path, key_column: &str) -> PathBuf {
46 let dir = source_path.parent().unwrap_or(Path::new("."));
47 let file_name = source_path
48 .file_name()
49 .map(|s| s.to_string_lossy().to_string())
50 .or_else(|| {
51 source_path
52 .file_stem()
53 .map(|s| s.to_string_lossy().to_string())
54 })
55 .unwrap_or_else(|| "source".to_string());
56 dir.join(format!("{file_name}.{key_column}.gcti"))
57}
58
59pub fn build_index_for_source(
60 definition: &SourceDefinition,
61 document_path: &Path,
62) -> Result<PathBuf> {
63 build_index_for_source_with_progress(definition, document_path, |_phase, _done, _total| {})
64}
65
66pub fn build_index_for_source_with_progress<F>(
67 definition: &SourceDefinition,
68 document_path: &Path,
69 mut on_progress: F,
70) -> Result<PathBuf>
71where
72 F: FnMut(BuildPhase, u64, u64),
73{
74 let result =
75 build_index_for_source_with_progress_impl(definition, document_path, &mut on_progress);
76 match result {
77 Ok(path) => {
78 INDEX_METRICS.record_build_success();
79 Ok(path)
80 }
81 Err(e) => {
82 INDEX_METRICS.record_build_failure();
83 Err(e)
84 }
85 }
86}
87
88fn build_index_for_source_with_progress_impl<F>(
89 definition: &SourceDefinition,
90 document_path: &Path,
91 on_progress: &mut F,
92) -> Result<PathBuf>
93where
94 F: FnMut(BuildPhase, u64, u64),
95{
96 let source_path = FileUtils::resolve_relative_path(document_path, &definition.file);
97 let key_columns = definition.indexed_columns();
98
99 if key_columns.is_empty() {
100 anyhow::bail!(
101 "no indexed_by column specified for source '{}'",
102 definition.file
103 );
104 }
105
106 let key_column = &key_columns[0];
107 let idx_path = index_path_for_source(&source_path, key_column);
108 let source_size = std::fs::metadata(&source_path)
109 .map(|m| m.len())
110 .unwrap_or(0);
111
112 let key_type = infer_key_type_for_column(&source_path, definition, key_column, source_size)?;
113
114 let mut reader = open_source_reader(definition, document_path)
115 .with_context(|| format!("failed to open source for indexing: {}", definition.file))?;
116
117 let mut index = SourceIndex::with_key_type(key_column, key_type);
118 let header_line = read_first_line(&source_path)?;
119 let mut byte_offset = header_line.len() as u64 + 1;
120
121 let mut row_count = 0u64;
122 on_progress(BuildPhase::Scan, byte_offset.min(source_size), source_size);
123 while let Some(row) = reader.next_row()? {
124 let key_val = row.get(key_column).ok_or_else(|| {
125 anyhow::anyhow!("column '{}' not found in row {}", key_column, row_count)
126 })?;
127
128 let row_bytes = estimate_row_size(&row);
129 index
130 .insert(key_val.to_string(), byte_offset, row_bytes)
131 .with_context(|| format!("failed to insert key '{}' at row {}", key_val, row_count))?;
132 byte_offset += row_bytes as u64 + 1;
133 row_count += 1;
134 if row_count.is_multiple_of(1024) {
135 on_progress(BuildPhase::Scan, byte_offset.min(source_size), source_size);
136 }
137 }
138 on_progress(BuildPhase::Scan, source_size, source_size);
139
140 let mut index_mut = index;
141 on_progress(BuildPhase::Write, 0, 1);
142 index_mut
143 .write_to_file(&idx_path)
144 .with_context(|| format!("failed to write index to {}", idx_path.display()))?;
145 on_progress(BuildPhase::Write, 1, 1);
146
147 if let Ok(meta) = std::fs::metadata(&idx_path) {
149 let size = meta.len();
150 if size > DEFAULT_MEMORY_LIMIT {
151 tracing::warn!(
152 "Index file {} is {} MB — exceeds {} MB limit. Consider increasing memory budget or reducing dataset size.",
153 idx_path.display(),
154 size / (1024 * 1024),
155 DEFAULT_MEMORY_LIMIT / (1024 * 1024)
156 );
157 }
158 }
159
160 Ok(idx_path)
161}
162
163fn infer_key_type_for_column(
164 source_path: &Path,
165 definition: &SourceDefinition,
166 key_column: &str,
167 source_size: u64,
168) -> Result<KeyType> {
169 let file = std::fs::File::open(source_path).with_context(|| {
170 format!(
171 "failed to open source for type inference: {}",
172 source_path.display()
173 )
174 })?;
175 let mut reader = std::io::BufReader::new(file);
176
177 let key_column_idx = if definition.format == Some(super::detect::SourceFormat::Ndjson) {
178 infer_ndjson_column_index(&mut reader, key_column)?
179 } else {
180 find_column_index(&mut reader, key_column)?
181 };
182
183 let max_bytes_scan = source_size.min(1024 * 1024);
184 let (key_type, _stats) = if definition.format == Some(super::detect::SourceFormat::Ndjson) {
185 super::index::infer_key_type_from_ndjson_stream(
186 &mut reader,
187 key_column,
188 1000,
189 max_bytes_scan,
190 )?
191 } else {
192 super::index::infer_key_type_from_stream(&mut reader, key_column_idx, 1000, max_bytes_scan)?
193 };
194
195 Ok(key_type)
196}
197
198fn infer_ndjson_column_index<R: std::io::BufRead>(
199 reader: &mut R,
200 target_column: &str,
201) -> Result<usize> {
202 let mut line = String::new();
203 loop {
204 line.clear();
205 match reader.read_line(&mut line) {
206 Ok(0) => anyhow::bail!("empty NDJSON file, cannot infer column index"),
207 Ok(_) => {}
208 Err(e) => anyhow::bail!("failed to read NDJSON for column inference: {}", e),
209 }
210 let trimmed = line.trim_ascii();
211 if trimmed.is_empty() || trimmed.starts_with('#') {
212 continue;
213 }
214 let obj: serde_json::Map<String, serde_json::Value> = serde_json::from_str(trimmed)
215 .map_err(|e| anyhow::anyhow!("invalid JSON in NDJSON: {}", e))?;
216 let mut keys: Vec<String> = obj.keys().cloned().collect();
217 keys.sort();
218 let idx = keys
219 .iter()
220 .position(|k| k == target_column)
221 .with_context(|| format!("column '{}' not found in NDJSON object", target_column))?;
222 return Ok(idx);
223 }
224}
225
226pub fn find_column_index<R: std::io::BufRead + std::io::Seek>(
227 reader: &mut R,
228 target_column: &str,
229) -> Result<usize> {
230 reader.seek(std::io::SeekFrom::Start(0))?;
231 let mut header = String::new();
232 reader.read_line(&mut header)?;
233
234 let delimiter = if header.contains('\t') { b'\t' } else { b',' };
235 let columns: Vec<&str> = header.trim_ascii().split(delimiter as char).collect();
236
237 let idx = columns
238 .iter()
239 .position(|&c| c == target_column)
240 .with_context(|| format!("column '{}' not found in source header", target_column))?;
241
242 Ok(idx)
243}
244
245pub fn load_or_build_index(
246 definition: &SourceDefinition,
247 document_path: &Path,
248) -> Result<SourceIndex> {
249 let source_path = FileUtils::resolve_relative_path(document_path, &definition.file);
250 let key_columns = definition.indexed_columns();
251
252 if key_columns.is_empty() {
253 anyhow::bail!("no indexed_by column for source '{}'", definition.file);
254 }
255
256 let key_column = &key_columns[0];
257 let idx_path = index_path_for_source(&source_path, key_column);
258
259 if idx_path.exists()
260 && let Ok(index) = SourceIndex::read_from_file(&idx_path)
261 && is_index_fresh(&idx_path, &source_path)
262 {
263 return Ok(index);
264 }
265
266 build_index_for_source(definition, document_path)?;
267 SourceIndex::read_from_file(&idx_path)
268}
269
270fn is_index_fresh(idx_path: &Path, source_path: &Path) -> bool {
271 let idx_meta = match std::fs::metadata(idx_path) {
272 Ok(m) => m,
273 Err(_) => return false,
274 };
275 let src_meta = match std::fs::metadata(source_path) {
276 Ok(m) => m,
277 Err(_) => return false,
278 };
279
280 if let (Ok(idx_time), Ok(src_time)) = (idx_meta.modified(), src_meta.modified()) {
281 return idx_time >= src_time;
282 }
283
284 true
285}
286
287fn read_first_line(path: &Path) -> Result<String> {
288 use std::io::{BufRead, BufReader};
289 let file = std::fs::File::open(path)?;
290 let mut reader = BufReader::new(file);
291 let mut line = String::new();
292 reader.read_line(&mut line)?;
293 Ok(line)
294}
295
296fn estimate_row_size(row: &SourceRow) -> u32 {
297 let mut size = 0u32;
298 for col in row.columns() {
299 size += col.len() as u32 + 1;
300 }
301 for val in row.values() {
302 size += val.len() as u32;
303 }
304 size + row.columns().len().saturating_sub(1) as u32
305}
306
307#[cfg(test)]
308mod tests {
309 use super::*;
310 #[cfg(not(miri))]
311 use std::io::Write;
312
313 #[cfg(not(miri))]
314 fn create_temp_csv(dir: &Path, name: &str, content: &str) -> PathBuf {
315 let path = dir.join(name);
316 let mut f = std::fs::File::create(&path).unwrap();
317 f.write_all(content.as_bytes()).unwrap();
318 path
319 }
320
321 #[test]
322 fn index_path_naming() {
323 let path = Path::new("data/pvz.csv");
324 let idx = index_path_for_source(path, "region_id");
325 assert_eq!(idx, PathBuf::from("data/pvz.csv.region_id.gcti"));
326 }
327
328 #[cfg(not(miri))]
329 #[test]
330 fn build_and_load_index() {
331 let dir = std::env::temp_dir().join("gctf_idx_build_test");
332 std::fs::create_dir_all(&dir).unwrap();
333 create_temp_csv(&dir, "data.csv", "id,name\n1,Alice\n2,Bob\n3,Charlie\n");
334
335 let defs: Vec<SourceDefinition> =
336 serde_yaml_ng::from_str("- file: data.csv\n name: data\n indexed_by: [id]\n")
337 .unwrap();
338
339 let doc_path = dir.join("test.gctf");
340 std::fs::write(&doc_path, "").unwrap();
341
342 let idx_path = build_index_for_source(&defs[0], &doc_path).unwrap();
343 assert!(idx_path.exists());
344
345 let index = SourceIndex::read_from_file(&idx_path).unwrap();
346 assert_eq!(index.len(), 3);
347 assert_eq!(index.key_column(), "id");
348 assert!(index.contains("1"));
349 assert!(index.contains("2"));
350 assert!(index.contains("3"));
351
352 std::fs::remove_dir_all(&dir).ok();
353 }
354
355 #[cfg(not(miri))]
356 #[test]
357 fn load_or_build_creates_on_first_call() {
358 let dir = std::env::temp_dir().join("gctf_idx_auto_test");
359 let _ = std::fs::remove_dir_all(&dir);
360 std::fs::create_dir_all(&dir).unwrap();
361 create_temp_csv(&dir, "items.csv", "code,label\nA,Alpha\nB,Bravo\n");
362
363 let defs: Vec<SourceDefinition> =
364 serde_yaml_ng::from_str("- file: items.csv\n name: items\n indexed_by: [code]\n")
365 .unwrap();
366
367 let doc_path = dir.join("test.gctf");
368 std::fs::write(&doc_path, "").unwrap();
369
370 let expected_idx = dir.join("items.csv.code.gcti");
371 assert!(
372 !expected_idx.exists(),
373 "stale index file should not exist: {}",
374 expected_idx.display()
375 );
376
377 let index = load_or_build_index(&defs[0], &doc_path).unwrap();
378 assert!(expected_idx.exists());
379 assert_eq!(index.len(), 2);
380
381 std::fs::remove_dir_all(&dir).ok();
382 }
383
384 #[cfg(not(miri))]
385 #[test]
386 fn load_or_build_reuses_existing() {
387 let dir = std::env::temp_dir().join("gctf_idx_reuse_test");
388 std::fs::create_dir_all(&dir).unwrap();
389 create_temp_csv(&dir, "data.csv", "id,val\n1,hello\n");
390
391 let defs: Vec<SourceDefinition> =
392 serde_yaml_ng::from_str("- file: data.csv\n name: d\n indexed_by: [id]\n").unwrap();
393
394 let doc_path = dir.join("test.gctf");
395 std::fs::write(&doc_path, "").unwrap();
396
397 let _idx1 = load_or_build_index(&defs[0], &doc_path).unwrap();
398 let _idx2 = load_or_build_index(&defs[0], &doc_path).unwrap();
399
400 std::fs::remove_dir_all(&dir).ok();
401 }
402
403 #[cfg(not(miri))]
404 #[test]
405 fn build_index_no_key_column_errors() {
406 let dir = std::env::temp_dir().join("gctf_idx_nokey_test");
407 std::fs::create_dir_all(&dir).unwrap();
408 create_temp_csv(&dir, "data.csv", "id,val\n1,hello\n");
409
410 let defs: Vec<SourceDefinition> =
411 serde_yaml_ng::from_str("- file: data.csv\n name: d\n").unwrap();
412
413 let doc_path = dir.join("test.gctf");
414 std::fs::write(&doc_path, "").unwrap();
415
416 let result = build_index_for_source(&defs[0], &doc_path);
417 assert!(result.is_err());
418
419 std::fs::remove_dir_all(&dir).ok();
420 }
421}