Skip to main content

ragfs_store/
schema.rs

1//! Arrow schema definitions for `LanceDB` tables.
2
3use arrow_schema::{DataType, Field, Schema, TimeUnit};
4use serde::{Deserialize, Serialize};
5use std::sync::Arc;
6
7/// Current chunks table schema version.
8///
9/// Version 1: pre-scope columns (no `dir_path` / `dir_depth` / `path_components`).
10/// Version 2: directory-scope columns for `ragfs query --scope`.
11pub const CHUNKS_SCHEMA_VERSION: u32 = 2;
12
13/// Sidecar written beside the Lance database (`…/indices/<hash>/ragfs-schema.json`).
14pub const SCHEMA_SIDECAR_FILENAME: &str = "ragfs-schema.json";
15
16/// Persisted schema version for an existing Lance index.
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
18pub struct SchemaSidecar {
19    pub chunks_schema_version: u32,
20    /// Embedding width of the `vector` `FixedSizeList`, when known.
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub embedding_dim: Option<usize>,
23    /// Embedder model id that wrote this index, when known.
24    #[serde(default, skip_serializing_if = "Option::is_none")]
25    pub embedding_model: Option<String>,
26}
27
28/// Scope columns missing from a pre-v2 chunks table (nullable so `add_columns` can fill nulls).
29#[must_use]
30pub fn missing_scope_columns(schema: &Schema) -> Vec<Field> {
31    let mut missing = Vec::new();
32    if schema.field_with_name("dir_path").is_err() {
33        missing.push(Field::new("dir_path", DataType::Utf8, true));
34    }
35    if schema.field_with_name("dir_depth").is_err() {
36        missing.push(Field::new("dir_depth", DataType::UInt16, true));
37    }
38    if schema.field_with_name("path_components").is_err() {
39        missing.push(Field::new("path_components", DataType::Utf8, true));
40    }
41    missing
42}
43
44/// Schema for the chunks table.
45#[must_use]
46pub fn chunks_schema(embedding_dim: usize) -> Schema {
47    Schema::new(vec![
48        // Identity
49        Field::new("chunk_id", DataType::Utf8, false),
50        Field::new("file_id", DataType::Utf8, false),
51        Field::new("file_path", DataType::Utf8, false),
52        // Content
53        Field::new("content", DataType::Utf8, false),
54        Field::new("content_type", DataType::Utf8, false),
55        // Position
56        Field::new("chunk_index", DataType::UInt32, false),
57        Field::new("start_byte", DataType::UInt64, false),
58        Field::new("end_byte", DataType::UInt64, false),
59        Field::new("start_line", DataType::UInt32, true),
60        Field::new("end_line", DataType::UInt32, true),
61        // Hierarchy
62        Field::new("parent_chunk_id", DataType::Utf8, true),
63        Field::new("depth", DataType::UInt8, false),
64        // Embedding
65        Field::new(
66            "embedding",
67            DataType::FixedSizeList(
68                Arc::new(Field::new("item", DataType::Float32, true)),
69                embedding_dim as i32,
70            ),
71            false,
72        ),
73        // Metadata
74        Field::new("embedding_model", DataType::Utf8, true),
75        Field::new(
76            "indexed_at",
77            DataType::Timestamp(TimeUnit::Millisecond, None),
78            false,
79        ),
80        // File metadata (denormalized)
81        Field::new("file_mime_type", DataType::Utf8, true),
82        Field::new("file_size_bytes", DataType::UInt64, true),
83        // Code-specific
84        Field::new("language", DataType::Utf8, true),
85        Field::new("symbol_type", DataType::Utf8, true),
86        Field::new("symbol_name", DataType::Utf8, true),
87        Field::new("dir_path", DataType::Utf8, false),
88        Field::new("dir_depth", DataType::UInt16, false),
89        Field::new("path_components", DataType::Utf8, false),
90    ])
91}
92
93/// Schema for the files metadata table.
94#[must_use]
95pub fn files_schema() -> Schema {
96    Schema::new(vec![
97        Field::new("file_id", DataType::Utf8, false),
98        Field::new("path", DataType::Utf8, false),
99        Field::new("size_bytes", DataType::UInt64, false),
100        Field::new("mime_type", DataType::Utf8, false),
101        Field::new("content_hash", DataType::Utf8, false),
102        Field::new(
103            "modified_at",
104            DataType::Timestamp(TimeUnit::Millisecond, None),
105            false,
106        ),
107        Field::new(
108            "indexed_at",
109            DataType::Timestamp(TimeUnit::Millisecond, None),
110            true,
111        ),
112        Field::new("chunk_count", DataType::UInt32, false),
113        Field::new("status", DataType::Utf8, false),
114        Field::new("error_message", DataType::Utf8, true),
115    ])
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn test_missing_scope_columns_on_v1_schema() {
124        let v1 = Schema::new(vec![Field::new("file_path", DataType::Utf8, false)]);
125        let missing = missing_scope_columns(&v1);
126        assert_eq!(missing.len(), 3);
127        assert_eq!(missing[0].name(), "dir_path");
128        assert_eq!(missing[1].name(), "dir_depth");
129        assert_eq!(missing[2].name(), "path_components");
130        assert!(missing.iter().all(Field::is_nullable));
131    }
132
133    #[test]
134    fn test_missing_scope_columns_on_v2_schema() {
135        let v2 = chunks_schema(384);
136        assert!(missing_scope_columns(&v2).is_empty());
137    }
138
139    #[test]
140    fn old_sidecar_deserializes_without_embedding_fields() {
141        let s: SchemaSidecar = serde_json::from_str(r#"{"chunks_schema_version":2}"#).unwrap();
142        assert_eq!(s.chunks_schema_version, 2);
143        assert_eq!(s.embedding_dim, None);
144        assert_eq!(s.embedding_model, None);
145    }
146}