From aa69ca0455c63caac6de9826973b3aa2914b25c3 Mon Sep 17 00:00:00 2001
From: LJ <linghua@cocoindex.io>
Date: Thu, 20 Mar 2025 18:47:04 -0700
Subject: [PATCH] `StructSchema` take a optional description and put in JSON
 schema

---
 src/base/json_schema.rs                |  6 +++++-
 src/base/schema.rs                     | 11 ++++++++++-
 src/builder/analyzer.rs                |  9 +++++++++
 src/builder/flow_builder.rs            |  1 +
 src/ops/functions/split_recursively.rs |  1 +
 src/ops/sources/google_drive.rs        |  1 +
 src/ops/sources/local_file.rs          |  1 +
 7 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/base/json_schema.rs b/src/base/json_schema.rs
index ba952798..4a7ee125 100644
--- a/src/base/json_schema.rs
+++ b/src/base/json_schema.rs
@@ -1,6 +1,6 @@
 use super::schema;
 use schemars::schema::{
-    ArrayValidation, InstanceType, ObjectValidation, Schema, SchemaObject, SingleOrVec,
+    ArrayValidation, InstanceType, Metadata, ObjectValidation, Schema, SchemaObject, SingleOrVec,
 };
 
 pub trait ToJsonSchema {
@@ -70,6 +70,10 @@ impl ToJsonSchema for schema::BasicValueType {
 impl ToJsonSchema for schema::StructSchema {
     fn to_json_schema(&self) -> SchemaObject {
         SchemaObject {
+            metadata: Some(Box::new(Metadata {
+                description: self.description.as_ref().map(|s| s.to_string()),
+                ..Default::default()
+            })),
             instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::Object))),
             object: Some(Box::new(ObjectValidation {
                 properties: self
diff --git a/src/base/schema.rs b/src/base/schema.rs
index 7477f5c5..af35957f 100644
--- a/src/base/schema.rs
+++ b/src/base/schema.rs
@@ -68,12 +68,16 @@ impl std::fmt::Display for BasicValueType {
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct StructSchema {
     pub fields: Arc<Vec<FieldSchema>>,
+
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub description: Option<Arc<str>>,
 }
 
 impl StructSchema {
     pub fn without_attrs(&self) -> Self {
         Self {
             fields: Arc::new(self.fields.iter().map(|f| f.without_attrs()).collect()),
+            description: None,
         }
     }
 }
@@ -168,11 +172,16 @@ impl std::fmt::Display for CollectionSchema {
 }
 
 impl CollectionSchema {
-    pub fn new(kind: CollectionKind, fields: Vec<FieldSchema>) -> Self {
+    pub fn new(
+        kind: CollectionKind,
+        fields: Vec<FieldSchema>,
+        description: Option<Arc<str>>,
+    ) -> Self {
         Self {
             kind,
             row: StructSchema {
                 fields: Arc::new(fields),
+                description,
             },
             collectors: Default::default(),
         }
diff --git a/src/builder/analyzer.rs b/src/builder/analyzer.rs
index ab244a3b..a180ec37 100644
--- a/src/builder/analyzer.rs
+++ b/src/builder/analyzer.rs
@@ -60,6 +60,7 @@ impl TryInto<ValueType> for &ValueTypeBuilder {
 pub(super) struct StructSchemaBuilder {
     fields: Vec<FieldSchema<ValueTypeBuilder>>,
     field_name_idx: HashMap<FieldName, u32>,
+    description: Option<Arc<str>>,
 }
 
 impl StructSchemaBuilder {
@@ -91,6 +92,7 @@ impl TryFrom<&StructSchema> for StructSchemaBuilder {
         let mut result = StructSchemaBuilder {
             fields: Vec::with_capacity(schema.fields.len()),
             field_name_idx: HashMap::with_capacity(schema.fields.len()),
+            description: schema.description.clone(),
         };
         for field in schema.fields.iter() {
             result.add_field(FieldSchema::<ValueTypeBuilder>::from_alternative(field)?)?;
@@ -110,6 +112,7 @@ impl TryInto<StructSchema> for &StructSchemaBuilder {
                     .map(FieldSchema::<ValueType>::from_alternative)
                     .collect::<Result<Vec<_>>>()?,
             ),
+            description: self.description.clone(),
         })
     }
 }
@@ -284,6 +287,10 @@ fn try_make_common_struct_schemas(
     }
     Ok(StructSchema {
         fields: Arc::new(result_fields),
+        description: schema1
+            .description
+            .clone()
+            .or_else(|| schema2.description.clone()),
     })
 }
 
@@ -484,6 +491,7 @@ fn analyze_struct_mapping(
         },
         StructSchema {
             fields: Arc::new(field_schemas),
+            description: None,
         },
     ))
 }
@@ -829,6 +837,7 @@ impl AnalyzerContext<'_> {
                 } else {
                     ValueType::Struct(StructSchema {
                         fields: Arc::from(key_fields_schema.clone()),
+                        description: None,
                     })
                 };
                 let mut value_fields_schema: Vec<FieldSchema> = vec![];
diff --git a/src/builder/flow_builder.rs b/src/builder/flow_builder.rs
index 054b58b1..f984d760 100644
--- a/src/builder/flow_builder.rs
+++ b/src/builder/flow_builder.rs
@@ -563,6 +563,7 @@ impl FlowBuilder {
                     })
                     .collect(),
             ),
+            description: None,
         };
 
         {
diff --git a/src/ops/functions/split_recursively.rs b/src/ops/functions/split_recursively.rs
index 960945de..69306ead 100644
--- a/src/ops/functions/split_recursively.rs
+++ b/src/ops/functions/split_recursively.rs
@@ -591,6 +591,7 @@ impl SimpleFunctionFactoryBase for Factory {
                 FieldSchema::new("location", make_output_type(BasicValueType::Range)),
                 FieldSchema::new("text", make_output_type(BasicValueType::Str)),
             ],
+            None,
         ))
         .with_attr(
             field_attrs::CHUNK_BASE_TEXT,
diff --git a/src/ops/sources/google_drive.rs b/src/ops/sources/google_drive.rs
index 74b4d865..37decd1a 100644
--- a/src/ops/sources/google_drive.rs
+++ b/src/ops/sources/google_drive.rs
@@ -274,6 +274,7 @@ impl SourceFactoryBase for Factory {
                     }),
                 ),
             ],
+            None,
         )))
     }
 
diff --git a/src/ops/sources/local_file.rs b/src/ops/sources/local_file.rs
index 06a3c3fd..3e5b5bef 100644
--- a/src/ops/sources/local_file.rs
+++ b/src/ops/sources/local_file.rs
@@ -112,6 +112,7 @@ impl SourceFactoryBase for Factory {
                     }),
                 ),
             ],
+            None,
         )))
     }