summa_core/index/
helpers.rs

1//! Indexing helper functions
2//!
3//! This module provides high-level helper functions for creating indexes
4//! and indexing documents, used by `summa-tool` and `summa-server`.
5
6use std::io::BufRead;
7use std::num::NonZeroU32;
8use std::path::Path;
9
10use crate::directories::{Directory, DirectoryWriter, FsDirectory};
11use crate::dsl::{Document, Schema, SchemaBuilder, parse_single_index};
12use crate::error::{Error, Result};
13use crate::index::{IndexConfig, IndexWriter};
14
15/// Schema configuration from JSON format
16#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
17pub struct SchemaFieldConfig {
18    /// Field name
19    pub name: String,
20    /// Field type: text, u64, i64, f64, bytes, json, sparse_vector, dense_vector
21    #[serde(rename = "type")]
22    pub field_type: String,
23    /// Whether field is indexed (default: true)
24    #[serde(default = "default_true")]
25    pub indexed: bool,
26    /// Whether field is stored (default: true)
27    #[serde(default = "default_true")]
28    pub stored: bool,
29    /// Dimension for dense_vector fields
30    #[serde(default)]
31    pub dimension: usize,
32    /// Text primary key, matching the SDL `primary` attribute.
33    #[serde(default, alias = "primary", skip_serializing_if = "std::ops::Not::not")]
34    pub primary_key: bool,
35    /// Stored content fingerprint for unchanged upserts.
36    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
37    pub content_hash: bool,
38}
39
40fn default_true() -> bool {
41    true
42}
43
44/// JSON schema configuration
45#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
46pub struct SchemaConfig {
47    /// List of field definitions
48    pub fields: Vec<SchemaFieldConfig>,
49    /// Creation-time cap on retained tokens per L1 phrase (default: 64).
50    #[serde(default, skip_serializing_if = "Option::is_none")]
51    pub max_l1_phrase_terms: Option<NonZeroU32>,
52}
53
54impl SchemaConfig {
55    /// Build a Schema from this configuration
56    pub fn build(&self) -> Result<Schema> {
57        let mut builder = SchemaBuilder::default();
58        if let Some(limit) = self.max_l1_phrase_terms {
59            builder.set_max_l1_phrase_terms(limit);
60        }
61
62        if self.fields.iter().filter(|field| field.primary_key).count() > 1 {
63            return Err(Error::Schema("at most one primary key is allowed".into()));
64        }
65        for field in &self.fields {
66            if field.primary_key && field.field_type != "text" {
67                return Err(Error::Schema("primary key must be text".into()));
68            }
69            let id = match field.field_type.as_str() {
70                "text" => builder.add_text_field(&field.name, field.indexed, field.stored),
71                "u64" => builder.add_u64_field(&field.name, field.indexed, field.stored),
72                "i64" => builder.add_i64_field(&field.name, field.indexed, field.stored),
73                "f64" => builder.add_f64_field(&field.name, field.indexed, field.stored),
74                "bytes" => builder.add_bytes_field(&field.name, field.stored),
75                "json" => builder.add_json_field(&field.name, field.stored),
76                "sparse_vector" => {
77                    builder.add_sparse_vector_field(&field.name, field.indexed, field.stored)
78                }
79                "dense_vector" => builder.add_dense_vector_field(
80                    &field.name,
81                    field.dimension,
82                    field.indexed,
83                    field.stored,
84                ),
85                other => return Err(Error::Schema(format!("Unknown field type: {}", other))),
86            };
87            if field.primary_key {
88                builder.set_primary_key(id);
89            }
90            if field.content_hash {
91                builder.set_content_hash(id);
92            }
93        }
94
95        let schema = builder.build();
96        schema.validate()?;
97        Ok(schema)
98    }
99}
100
101/// Parse schema from a string (auto-detects JSON or SDL format)
102pub fn parse_schema(content: &str) -> Result<Schema> {
103    let trimmed = content.trim();
104
105    // Detect SDL format (starts with "index " or "#" for comments)
106    if trimmed.starts_with("index ") || trimmed.starts_with('#') {
107        let index_def = parse_single_index(content)
108            .map_err(|e| Error::Schema(format!("Failed to parse SDL: {}", e)))?;
109        Ok(index_def.to_schema())
110    } else {
111        // Try JSON format
112        let config: SchemaConfig = serde_json::from_str(content)
113            .map_err(|e| Error::Schema(format!("Failed to parse JSON schema: {}", e)))?;
114        config.build()
115    }
116}
117
118/// Create a new index at the given path with the provided schema
119pub async fn create_index_at_path(
120    path: impl AsRef<Path>,
121    schema: Schema,
122    config: IndexConfig,
123) -> Result<IndexWriter<FsDirectory>> {
124    let path = path.as_ref();
125
126    std::fs::create_dir_all(path).map_err(|e| {
127        Error::Io(std::io::Error::new(
128            e.kind(),
129            format!("Failed to create index directory {:?}: {}", path, e),
130        ))
131    })?;
132
133    let dir = FsDirectory::new(path);
134    IndexWriter::create(dir, schema, config).await
135}
136
137/// Create a new index from an SDL schema string
138pub async fn create_index_from_sdl(
139    path: impl AsRef<Path>,
140    sdl: &str,
141    config: IndexConfig,
142) -> Result<IndexWriter<FsDirectory>> {
143    let schema = parse_schema(sdl)?;
144    create_index_at_path(path, schema, config).await
145}
146
147/// Indexing statistics
148#[derive(Debug, Clone, Default)]
149pub struct IndexingStats {
150    /// Number of documents indexed
151    pub indexed: usize,
152    /// Number of documents that failed to parse
153    pub errors: usize,
154    /// Total time in seconds
155    pub elapsed_secs: f64,
156}
157
158impl IndexingStats {
159    /// Documents per second rate
160    pub fn docs_per_sec(&self) -> f64 {
161        if self.elapsed_secs > 0.0 {
162            self.indexed as f64 / self.elapsed_secs
163        } else {
164            0.0
165        }
166    }
167}
168
169/// Index documents from a JSONL reader
170///
171/// Each line should be a valid JSON object. Documents are parsed according
172/// to the schema and indexed. Returns statistics about the indexing operation.
173pub async fn index_documents_from_reader<D, R>(
174    writer: &mut IndexWriter<D>,
175    reader: R,
176    progress_callback: Option<&dyn Fn(usize)>,
177) -> Result<IndexingStats>
178where
179    D: Directory + DirectoryWriter,
180    R: BufRead,
181{
182    let schema = writer.schema();
183    let mut stats = IndexingStats::default();
184    let start_time = std::time::Instant::now();
185
186    for line in reader.lines() {
187        let line = line.map_err(Error::Io)?;
188        if line.trim().is_empty() {
189            continue;
190        }
191
192        let json: serde_json::Value = match serde_json::from_str(&line) {
193            Ok(v) => v,
194            Err(_) => {
195                stats.errors += 1;
196                continue;
197            }
198        };
199
200        let doc = match Document::from_json(&json, &schema) {
201            Some(d) => d,
202            None => {
203                stats.errors += 1;
204                continue;
205            }
206        };
207
208        writer.add_document(doc)?;
209        stats.indexed += 1;
210
211        if let Some(callback) = progress_callback {
212            callback(stats.indexed);
213        }
214    }
215
216    writer.commit().await?;
217    stats.elapsed_secs = start_time.elapsed().as_secs_f64();
218
219    Ok(stats)
220}
221
222/// Index a single document from JSON
223pub async fn index_json_document<D>(writer: &IndexWriter<D>, json: &serde_json::Value) -> Result<()>
224where
225    D: Directory + DirectoryWriter,
226{
227    let schema = writer.schema();
228    let doc = Document::from_json(json, &schema)
229        .ok_or_else(|| Error::Document("Failed to parse JSON document".to_string()))?;
230    writer.add_document(doc)?;
231    Ok(())
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237    use crate::directories::RamDirectory;
238
239    #[test]
240    fn test_schema_config_json() {
241        let json = r#"{
242            "fields": [
243                {"name": "title", "type": "text", "indexed": true, "stored": true},
244                {"name": "body", "type": "text"},
245                {"name": "score", "type": "f64", "indexed": false}
246            ]
247        }"#;
248
249        let config: SchemaConfig = serde_json::from_str(json).unwrap();
250        assert_eq!(config.fields.len(), 3);
251
252        let schema = config.build().unwrap();
253        assert!(schema.get_field("title").is_some());
254        assert!(schema.get_field("body").is_some());
255        assert!(schema.get_field("score").is_some());
256    }
257
258    #[test]
259    fn test_parse_schema_json() {
260        let json = r#"{"fields": [{"name": "text", "type": "text"}]}"#;
261        let schema = parse_schema(json).unwrap();
262        assert!(schema.get_field("text").is_some());
263    }
264
265    #[test]
266    fn json_schema_configures_content_hash_and_rejects_invalid_combinations() {
267        let json = r#"{"fields":[{"name":"id","type":"text","primary_key":true},{"name":"digest","type":"bytes","stored":true,"content_hash":true}]}"#;
268        let schema = parse_schema(json).unwrap();
269        assert_eq!(schema.primary_field(), schema.get_field("id"));
270        assert_eq!(schema.content_hash_field(), schema.get_field("digest"));
271        let config: SchemaConfig = serde_json::from_str(json).unwrap();
272        assert_eq!(
273            parse_schema(&serde_json::to_string(&config).unwrap())
274                .unwrap()
275                .content_hash_field(),
276            schema.content_hash_field()
277        );
278        assert!(
279            parse_schema(&json.replace("\"primary_key\":true", "\"primary_key\":false")).is_err()
280        );
281        assert!(parse_schema(&json.replace("\"stored\":true", "\"stored\":false")).is_err());
282    }
283
284    #[test]
285    fn test_parse_schema_sdl() {
286        let sdl = r#"
287            index test {
288                field text: text [indexed, stored]
289            }
290        "#;
291        let schema = parse_schema(sdl).unwrap();
292        assert!(schema.get_field("text").is_some());
293    }
294
295    #[test]
296    fn creation_schemas_preserve_positive_phrase_limits_and_reject_invalid_values() {
297        assert_eq!(Schema::default().max_l1_phrase_terms(), 64);
298        assert_eq!(SchemaBuilder::default().build().max_l1_phrase_terms(), 64);
299        for value in [None, Some(1), Some(64), Some(65), Some(300), Some(u32::MAX)] {
300            let option = value
301                .map(|value| format!("max_l1_phrase_terms: {value}"))
302                .unwrap_or_default();
303            let mut json = serde_json::json!({"fields": [{"name": "body", "type": "text"}]});
304            if let Some(value) = value {
305                json["max_l1_phrase_terms"] = value.into();
306            }
307            for input in [
308                format!("index documents {{ {option} field body: text }}"),
309                json.to_string(),
310            ] {
311                let schema = parse_schema(&input).unwrap();
312                assert_eq!(schema.max_l1_phrase_terms(), value.unwrap_or(64) as usize);
313                let serialized = serde_json::to_value(&schema).unwrap();
314                assert_eq!(
315                    serialized.get("max_l1_phrase_terms").is_some(),
316                    value.is_some()
317                );
318                let restored: Schema = serde_json::from_value(serialized).unwrap();
319                assert_eq!(restored.max_l1_phrase_terms(), schema.max_l1_phrase_terms());
320            }
321        }
322        for invalid in ["0", "-1", "1.5", "4294967296", "\"64\"", "true"] {
323            for input in [
324                format!("index documents {{ max_l1_phrase_terms: {invalid} field body: text }}"),
325                format!(r#"{{"max_l1_phrase_terms": {invalid}, "fields": []}}"#),
326            ] {
327                assert!(parse_schema(&input).is_err(), "accepted {input}");
328            }
329        }
330        assert!(
331            parse_schema("index documents { max_l1_phrase_terms: 64 max_l1_phrase_terms: 256 }")
332                .is_err()
333        );
334    }
335
336    #[tokio::test]
337    async fn test_index_documents_from_reader() {
338        let mut builder = SchemaBuilder::default();
339        let _title = builder.add_text_field("title", true, true);
340        let schema = builder.build();
341
342        let dir = RamDirectory::new();
343        let config = IndexConfig::default();
344        let mut writer = IndexWriter::create(dir, schema, config).await.unwrap();
345
346        let jsonl = r#"{"title": "Doc 1"}
347{"title": "Doc 2"}
348{"title": "Doc 3"}"#;
349
350        let reader = std::io::Cursor::new(jsonl);
351        let stats = index_documents_from_reader(&mut writer, reader, None)
352            .await
353            .unwrap();
354
355        assert_eq!(stats.indexed, 3);
356        assert_eq!(stats.errors, 0);
357    }
358}
summa_core/index/helpers.rs

summa_core/index/
helpers.rs
