use std::io::BufRead;
use std::num::NonZeroU32;
use std::path::Path;
use crate::directories::{Directory, DirectoryWriter, FsDirectory};
use crate::dsl::{Document, Schema, SchemaBuilder, parse_single_index};
use crate::error::{Error, Result};
use crate::index::{IndexConfig, IndexWriter};
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
pub struct SchemaFieldConfig {
pub name: String,
#[serde(rename = "type")]
pub field_type: String,
#[serde(default = "default_true")]
pub indexed: bool,
#[serde(default = "default_true")]
pub stored: bool,
#[serde(default)]
pub dimension: usize,
#[serde(default, alias = "primary", skip_serializing_if = "std::ops::Not::not")]
pub primary_key: bool,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub content_hash: bool,
}
fn default_true() -> bool {
true
}
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
pub struct SchemaConfig {
pub fields: Vec<SchemaFieldConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_l1_phrase_terms: Option<NonZeroU32>,
}
impl SchemaConfig {
pub fn build(&self) -> Result<Schema> {
let mut builder = SchemaBuilder::default();
if let Some(limit) = self.max_l1_phrase_terms {
builder.set_max_l1_phrase_terms(limit);
}
if self.fields.iter().filter(|field| field.primary_key).count() > 1 {
return Err(Error::Schema("at most one primary key is allowed".into()));
}
for field in &self.fields {
if field.primary_key && field.field_type != "text" {
return Err(Error::Schema("primary key must be text".into()));
}
let id = match field.field_type.as_str() {
"text" => builder.add_text_field(&field.name, field.indexed, field.stored),
"u64" => builder.add_u64_field(&field.name, field.indexed, field.stored),
"i64" => builder.add_i64_field(&field.name, field.indexed, field.stored),
"f64" => builder.add_f64_field(&field.name, field.indexed, field.stored),
"bytes" => builder.add_bytes_field(&field.name, field.stored),
"json" => builder.add_json_field(&field.name, field.stored),
"sparse_vector" => {
builder.add_sparse_vector_field(&field.name, field.indexed, field.stored)
}
"dense_vector" => builder.add_dense_vector_field(
&field.name,
field.dimension,
field.indexed,
field.stored,
),
other => return Err(Error::Schema(format!("Unknown field type: {}", other))),
};
if field.primary_key {
builder.set_primary_key(id);
}
if field.content_hash {
builder.set_content_hash(id);
}
}
let schema = builder.build();
schema.validate()?;
Ok(schema)
}
}
pub fn parse_schema(content: &str) -> Result<Schema> {
let trimmed = content.trim();
if trimmed.starts_with("index ") || trimmed.starts_with('#') {
let index_def = parse_single_index(content)
.map_err(|e| Error::Schema(format!("Failed to parse SDL: {}", e)))?;
Ok(index_def.to_schema())
} else {
let config: SchemaConfig = serde_json::from_str(content)
.map_err(|e| Error::Schema(format!("Failed to parse JSON schema: {}", e)))?;
config.build()
}
}
pub async fn create_index_at_path(
path: impl AsRef<Path>,
schema: Schema,
config: IndexConfig,
) -> Result<IndexWriter<FsDirectory>> {
let path = path.as_ref();
std::fs::create_dir_all(path).map_err(|e| {
Error::Io(std::io::Error::new(
e.kind(),
format!("Failed to create index directory {:?}: {}", path, e),
))
})?;
let dir = FsDirectory::new(path);
IndexWriter::create(dir, schema, config).await
}
pub async fn create_index_from_sdl(
path: impl AsRef<Path>,
sdl: &str,
config: IndexConfig,
) -> Result<IndexWriter<FsDirectory>> {
let schema = parse_schema(sdl)?;
create_index_at_path(path, schema, config).await
}
#[derive(Debug, Clone, Default)]
pub struct IndexingStats {
pub indexed: usize,
pub errors: usize,
pub elapsed_secs: f64,
}
impl IndexingStats {
pub fn docs_per_sec(&self) -> f64 {
if self.elapsed_secs > 0.0 {
self.indexed as f64 / self.elapsed_secs
} else {
0.0
}
}
}
pub async fn index_documents_from_reader<D, R>(
writer: &mut IndexWriter<D>,
reader: R,
progress_callback: Option<&dyn Fn(usize)>,
) -> Result<IndexingStats>
where
D: Directory + DirectoryWriter,
R: BufRead,
{
let schema = writer.schema();
let mut stats = IndexingStats::default();
let start_time = std::time::Instant::now();
for line in reader.lines() {
let line = line.map_err(Error::Io)?;
if line.trim().is_empty() {
continue;
}
let json: serde_json::Value = match serde_json::from_str(&line) {
Ok(v) => v,
Err(_) => {
stats.errors += 1;
continue;
}
};
let doc = match Document::from_json(&json, &schema) {
Some(d) => d,
None => {
stats.errors += 1;
continue;
}
};
writer.add_document(doc)?;
stats.indexed += 1;
if let Some(callback) = progress_callback {
callback(stats.indexed);
}
}
writer.commit().await?;
stats.elapsed_secs = start_time.elapsed().as_secs_f64();
Ok(stats)
}
pub async fn index_json_document<D>(writer: &IndexWriter<D>, json: &serde_json::Value) -> Result<()>
where
D: Directory + DirectoryWriter,
{
let schema = writer.schema();
let doc = Document::from_json(json, &schema)
.ok_or_else(|| Error::Document("Failed to parse JSON document".to_string()))?;
writer.add_document(doc)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::directories::RamDirectory;
#[test]
fn test_schema_config_json() {
let json = r#"{
"fields": [
{"name": "title", "type": "text", "indexed": true, "stored": true},
{"name": "body", "type": "text"},
{"name": "score", "type": "f64", "indexed": false}
]
}"#;
let config: SchemaConfig = serde_json::from_str(json).unwrap();
assert_eq!(config.fields.len(), 3);
let schema = config.build().unwrap();
assert!(schema.get_field("title").is_some());
assert!(schema.get_field("body").is_some());
assert!(schema.get_field("score").is_some());
}
#[test]
fn test_parse_schema_json() {
let json = r#"{"fields": [{"name": "text", "type": "text"}]}"#;
let schema = parse_schema(json).unwrap();
assert!(schema.get_field("text").is_some());
}
#[test]
fn json_schema_configures_content_hash_and_rejects_invalid_combinations() {
let json = r#"{"fields":[{"name":"id","type":"text","primary_key":true},{"name":"digest","type":"bytes","stored":true,"content_hash":true}]}"#;
let schema = parse_schema(json).unwrap();
assert_eq!(schema.primary_field(), schema.get_field("id"));
assert_eq!(schema.content_hash_field(), schema.get_field("digest"));
let config: SchemaConfig = serde_json::from_str(json).unwrap();
assert_eq!(
parse_schema(&serde_json::to_string(&config).unwrap())
.unwrap()
.content_hash_field(),
schema.content_hash_field()
);
assert!(
parse_schema(&json.replace("\"primary_key\":true", "\"primary_key\":false")).is_err()
);
assert!(parse_schema(&json.replace("\"stored\":true", "\"stored\":false")).is_err());
}
#[test]
fn test_parse_schema_sdl() {
let sdl = r#"
index test {
field text: text [indexed, stored]
}
"#;
let schema = parse_schema(sdl).unwrap();
assert!(schema.get_field("text").is_some());
}
#[test]
fn creation_schemas_preserve_positive_phrase_limits_and_reject_invalid_values() {
assert_eq!(Schema::default().max_l1_phrase_terms(), 64);
assert_eq!(SchemaBuilder::default().build().max_l1_phrase_terms(), 64);
for value in [None, Some(1), Some(64), Some(65), Some(300), Some(u32::MAX)] {
let option = value
.map(|value| format!("max_l1_phrase_terms: {value}"))
.unwrap_or_default();
let mut json = serde_json::json!({"fields": [{"name": "body", "type": "text"}]});
if let Some(value) = value {
json["max_l1_phrase_terms"] = value.into();
}
for input in [
format!("index documents {{ {option} field body: text }}"),
json.to_string(),
] {
let schema = parse_schema(&input).unwrap();
assert_eq!(schema.max_l1_phrase_terms(), value.unwrap_or(64) as usize);
let serialized = serde_json::to_value(&schema).unwrap();
assert_eq!(
serialized.get("max_l1_phrase_terms").is_some(),
value.is_some()
);
let restored: Schema = serde_json::from_value(serialized).unwrap();
assert_eq!(restored.max_l1_phrase_terms(), schema.max_l1_phrase_terms());
}
}
for invalid in ["0", "-1", "1.5", "4294967296", "\"64\"", "true"] {
for input in [
format!("index documents {{ max_l1_phrase_terms: {invalid} field body: text }}"),
format!(r#"{{"max_l1_phrase_terms": {invalid}, "fields": []}}"#),
] {
assert!(parse_schema(&input).is_err(), "accepted {input}");
}
}
assert!(
parse_schema("index documents { max_l1_phrase_terms: 64 max_l1_phrase_terms: 256 }")
.is_err()
);
}
#[tokio::test]
async fn test_index_documents_from_reader() {
let mut builder = SchemaBuilder::default();
let _title = builder.add_text_field("title", true, true);
let schema = builder.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir, schema, config).await.unwrap();
let jsonl = r#"{"title": "Doc 1"}
{"title": "Doc 2"}
{"title": "Doc 3"}"#;
let reader = std::io::Cursor::new(jsonl);
let stats = index_documents_from_reader(&mut writer, reader, None)
.await
.unwrap();
assert_eq!(stats.indexed, 3);
assert_eq!(stats.errors, 0);
}
}