feat: only write defined schema fields to parquet (#3454)

Co-authored-by: Hengfei Yang <hengfei.yang@gmail.com>

src/config/src/utils/json.rs

@@ -21,6 +21,13 @@ pub fn get_float_value(val: &Value) -> f64 {
     match val {
         Value::String(v) => v.parse::<f64>().unwrap_or(0.0),
         Value::Number(v) => v.as_f64().unwrap_or(0.0),
+        Value::Bool(v) => {
+            if *v {
+                1.0
+            } else {
+                0.0
+            }
+        }
         _ => 0.0,
     }
 }
@@ -29,6 +36,13 @@ pub fn get_int_value(val: &Value) -> i64 {
     match val {
         Value::String(v) => v.parse::<i64>().unwrap_or(0),
         Value::Number(v) => v.as_i64().unwrap_or(0),
+        Value::Bool(v) => {
+            if *v {
+                1
+            } else {
+                0
+            }
+        }
         _ => 0,
     }
 }
@@ -37,6 +51,13 @@ pub fn get_uint_value(val: &Value) -> u64 {
     match val {
         Value::String(v) => v.parse::<u64>().unwrap_or(0),
         Value::Number(v) => v.as_u64().unwrap_or(0),
+        Value::Bool(v) => {
+            if *v {
+                1
+            } else {
+                0
+            }
+        }
         _ => 0,
     }
 }
@@ -57,6 +78,15 @@ pub fn get_string_value(value: &Value) -> String {
     }
 }
 
+pub fn pickup_string_value(val: Value) -> String {
+    match val {
+        Value::String(v) => v,
+        Value::Number(v) => v.to_string(),
+        Value::Bool(v) => v.to_string(),
+        _ => val.to_string(),
+    }
+}
+
 pub fn estimate_json_bytes(val: &Value) -> usize {
     let mut size = 0;
     match val {

src/config/src/utils/record_batch_ext.rs

@@ -47,14 +47,20 @@ pub fn convert_json_to_record_batch(
     data: &[Arc<serde_json::Value>],
 ) -> Result<RecordBatch, ArrowError> {
     // collect all keys from the json data
-    let mut keys = HashSet::new();
+    let mut keys = HashSet::with_capacity(schema.fields().len());
+    let mut max_keys = 0;
     for record in data.iter() {
-        for (k, _) in record.as_object().unwrap() {
+        let record = record.as_object().unwrap();
+        if record.len() > max_keys {
+            max_keys = record.len();
+        }
+        for (k, _) in record {
             keys.insert(k);
         }
     }
 
     // create builders for each key
+    let records_len = data.len();
     let mut builders: FxIndexMap<&String, (&DataType, Box<dyn ArrayBuilder>)> = schema
         .fields()
         .iter()
@@ -62,14 +68,15 @@ pub fn convert_json_to_record_batch(
         .map(|f| {
             (
                 f.name(),
-                (f.data_type(), make_builder(f.data_type(), data.len())),
+                (f.data_type(), make_builder(f.data_type(), records_len)),
             )
         })
         .collect();
 
     // fill builders with data
+    let mut record_keys = HashSet::with_capacity(max_keys);
     for record in data.iter() {
-        let mut record_keys = HashSet::new();
+        record_keys.clear();
         for (k, v) in record.as_object().unwrap() {
             record_keys.insert(k);
             let res = builders.get_mut(k);
@@ -200,7 +207,7 @@ pub fn convert_json_to_record_batch(
         if let Some((_, builder)) = builders.get_mut(field.name()) {
             cols.push(builder.finish());
         } else {
-            cols.push(new_null_array(field.data_type(), data.len()))
+            cols.push(new_null_array(field.data_type(), records_len))
         }
     }
 

src/config/src/utils/schema_ext.rs

@@ -38,21 +38,19 @@ impl SchemaExt for Schema {
 
     // ensure schema is compatible
     fn cloned_from(&self, schema: &Schema) -> Schema {
-        let mut schema_latest_map = HashMap::with_capacity(schema.fields().len());
-        for field in schema.fields() {
-            schema_latest_map.insert(field.name(), field.clone());
-        }
-        let mut fields = Vec::with_capacity(self.fields().len());
-        for field in self.fields() {
-            match schema_latest_map.get(field.name()) {
-                Some(f) => {
-                    fields.push(f.clone());
-                }
-                None => {
-                    fields.push(field.clone());
-                }
-            }
-        }
+        let schema_latest_map: HashMap<_, _> = schema
+            .fields()
+            .iter()
+            .map(|field| (field.name(), field))
+            .collect();
+        let fields = self
+            .fields()
+            .iter()
+            .map(|f| match schema_latest_map.get(f.name()) {
+                Some(f) => (*f).clone(),
+                None => f.clone(),
+            })
+            .collect::<Vec<_>>();
         Schema::new(fields)
     }
 

src/handler/http/request/search/mod.rs

@@ -981,7 +981,7 @@ async fn values_v1(
     let key = format!("{org_id}/{stream_type}/{stream_name}");
     let r = STREAM_SCHEMAS_LATEST.read().await;
     let schema = if let Some(schema) = r.get(&key) {
-        schema.clone()
+        schema.schema().clone()
     } else {
         arrow_schema::Schema::empty()
     };

src/handler/http/request/status/mod.rs

@@ -331,7 +331,7 @@ async fn get_stream_schema_status() -> (usize, usize, usize) {
         stream_num += 1;
         stream_schema_num += 1;
         mem_size += key.len();
-        mem_size += schema.size();
+        mem_size += schema.schema().size();
     }
     drop(r);
     (stream_num, stream_schema_num, mem_size)

src/infra/src/schema/mod.rs

@@ -13,10 +13,12 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+use std::collections::HashMap;
+
 use chrono::Utc;
 use config::{
     meta::stream::{PartitionTimeLevel, StreamSettings, StreamType},
-    utils::json,
+    utils::{json, schema_ext::SchemaExt},
     RwAHashMap, CONFIG,
 };
 use datafusion::arrow::datatypes::{DataType, Field, Schema};
@@ -34,7 +36,8 @@ pub static STREAM_SCHEMAS: Lazy<RwAHashMap<String, Vec<(i64, Schema)>>> =
     Lazy::new(Default::default);
 pub static STREAM_SCHEMAS_COMPRESSED: Lazy<RwAHashMap<String, Vec<(i64, bytes::Bytes)>>> =
     Lazy::new(Default::default);
-pub static STREAM_SCHEMAS_LATEST: Lazy<RwAHashMap<String, Schema>> = Lazy::new(Default::default);
+pub static STREAM_SCHEMAS_LATEST: Lazy<RwAHashMap<String, SchemaCache>> =
+    Lazy::new(Default::default);
 pub static STREAM_SCHEMAS_FIELDS: Lazy<RwAHashMap<String, (i64, Vec<String>)>> =
     Lazy::new(Default::default);
 pub static STREAM_SETTINGS: Lazy<RwAHashMap<String, StreamSettings>> = Lazy::new(Default::default);
@@ -54,13 +57,31 @@ pub async fn get(org_id: &str, stream_name: &str, stream_type: StreamType) -> Re
 
     let r = STREAM_SCHEMAS_LATEST.read().await;
     if let Some(schema) = r.get(cache_key) {
-        return Ok(schema.clone());
+        return Ok(schema.schema.clone());
     }
     drop(r);
     // if not found in cache, get from db
     get_from_db(org_id, stream_name, stream_type).await
 }
 
+pub async fn get_cache(
+    org_id: &str,
+    stream_name: &str,
+    stream_type: StreamType,
+) -> Result<SchemaCache> {
+    let key = mk_key(org_id, stream_type, stream_name);
+    let cache_key = key.strip_prefix("/schema/").unwrap();
+
+    let r = STREAM_SCHEMAS_LATEST.read().await;
+    if let Some(schema) = r.get(cache_key) {
+        return Ok(schema.clone());
+    }
+    drop(r);
+    // if not found in cache, get from db
+    let schema = get_from_db(org_id, stream_name, stream_type).await?;
+    Ok(SchemaCache::new(schema))
+}
+
 pub async fn get_from_db(
     org_id: &str,
     stream_name: &str,
@@ -617,6 +638,42 @@ pub fn get_merge_schema_changes(
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct SchemaCache {
+    schema: Schema,
+    fields_map: HashMap<String, usize>,
+    hash_key: String,
+}
+
+impl SchemaCache {
+    pub fn new(schema: Schema) -> Self {
+        let hash_key = schema.hash_key();
+        let fields_map = schema
+            .fields()
+            .iter()
+            .enumerate()
+            .map(|(i, f)| (f.name().to_owned(), i))
+            .collect();
+        Self {
+            schema,
+            fields_map,
+            hash_key,
+        }
+    }
+
+    pub fn hash_key(&self) -> &str {
+        &self.hash_key
+    }
+
+    pub fn schema(&self) -> &Schema {
+        &self.schema
+    }
+
+    pub fn fields_map(&self) -> &HashMap<String, usize> {
+        &self.fields_map
+    }
+}
+
 pub fn is_widening_conversion(from: &DataType, to: &DataType) -> bool {
     let allowed_type = match from {
         DataType::Boolean => vec![DataType::Utf8],

src/ingester/src/entry.rs

@@ -38,6 +38,15 @@ pub struct Entry {
 }
 
 impl Entry {
+    pub fn new() -> Self {
+        Self {
+            stream: "".into(),
+            schema_key: "".into(),
+            partition_key: "".into(),
+            data: Vec::new(),
+            data_size: 0,
+        }
+    }
     pub fn into_bytes(&mut self) -> Result<Vec<u8>> {
         let mut buf = Vec::with_capacity(4096);
         let stream = self.stream.as_bytes();
@@ -103,6 +112,12 @@ impl Entry {
     }
 }
 
+impl Default for Entry {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 pub struct RecordBatchEntry {
     pub data: RecordBatch,
     pub data_json_size: usize,

src/ingester/src/lib.rs

@@ -49,7 +49,7 @@ pub async fn init() -> errors::Result<()> {
         loop {
             // persist immutable data to disk
             if let Err(e) = immutable::persist().await {
-                log::error!("persist error: {}", e);
+                log::error!("immutable persist error: {}", e);
             }
             // shrink metadata cache
             WAL_PARQUET_METADATA.write().await.shrink_to_fit();
@@ -57,5 +57,20 @@ pub async fn init() -> errors::Result<()> {
         }
     });
 
+    // start a job to flush memtable to immutable
+    tokio::task::spawn(async move {
+        let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(
+            config::CONFIG.limit.max_file_retention_time,
+        ));
+        interval.tick().await; // the first tick is immediate
+        loop {
+            // check memtable ttl
+            if let Err(e) = writer::check_ttl().await {
+                log::error!("memtable check ttl error: {}", e);
+            }
+            interval.tick().await;
+        }
+    });
+
     Ok(())
 }

src/ingester/src/writer.rs

@@ -107,6 +107,18 @@ pub async fn read_from_memtable(
     Ok(batches)
 }
 
+pub async fn check_ttl() -> Result<()> {
+    for w in WRITERS.iter() {
+        let w = w.read().await;
+        for r in w.values() {
+            // check writer
+            r.write(Arc::new(Schema::empty()), Entry::default(), true)
+                .await?;
+        }
+    }
+    Ok(())
+}
+
 pub async fn flush_all() -> Result<()> {
     for w in WRITERS.iter() {
         let mut w = w.write().await;
@@ -156,11 +168,21 @@ impl Writer {
         }
     }
 
-    pub async fn write(&self, schema: Arc<Schema>, mut entry: Entry) -> Result<()> {
-        if entry.data.is_empty() {
+    // check_ttl is used to check if the memtable has expired
+    pub async fn write(
+        &self,
+        schema: Arc<Schema>,
+        mut entry: Entry,
+        check_ttl: bool,
+    ) -> Result<()> {
+        if entry.data.is_empty() && !check_ttl {
             return Ok(());
         }
-        let entry_bytes = entry.into_bytes()?;
+        let entry_bytes = if !check_ttl {
+            entry.into_bytes()?
+        } else {
+            Vec::new()
+        };
         let mut wal = self.wal.lock().await;
         let mut mem = self.memtable.write().await;
         if self.check_wal_threshold(wal.size(), entry_bytes.len())
@@ -211,11 +233,13 @@ impl Writer {
             });
         }
 
-        // write into wal
-        wal.write(&entry_bytes, false).context(WalSnafu)?;
+        if !check_ttl {
+            // write into wal
+            wal.write(&entry_bytes, false).context(WalSnafu)?;
+            // write into memtable
+            mem.write(schema, entry).await?;
+        }
 
-        // write into memtable
-        mem.write(schema, entry).await?;
         Ok(())
     }
 

src/job/files/parquet.rs

@@ -37,7 +37,7 @@ use config::{
 };
 use datafusion::{arrow::json as arrow_json, datasource::MemTable, prelude::*};
 use hashbrown::HashSet;
-use infra::{cache, storage};
+use infra::{cache, schema::SchemaCache, storage};
 use ingester::WAL_PARQUET_METADATA;
 use once_cell::sync::Lazy;
 use parquet::arrow::ParquetRecordBatchStreamBuilder;
@@ -51,7 +51,6 @@ use crate::{
     job::files::idx::write_to_disk,
     service::{
         db,
-        schema::SchemaCache,
         search::datafusion::{
             exec::merge_parquet_files, string_to_array_v2_udf::STRING_TO_ARRAY_V2_UDF,
         },

src/service/compact/file_list.rs

@@ -57,7 +57,7 @@ pub async fn run_merge(offset: i64) -> Result<(), anyhow::Error> {
         offset = time_now.timestamp_micros();
         let r = STREAM_SCHEMAS_LATEST.read().await;
         for (key, val) in r.iter() {
-            if let Some(val) = val.metadata.get("created_at") {
+            if let Some(val) = val.schema().metadata.get("created_at") {
                 let time_min = val.parse().unwrap();
                 if time_min == 0 {
                     log::info!(

src/service/db/schema.rs

@@ -22,8 +22,8 @@ use hashbrown::{HashMap, HashSet};
 use infra::{
     cache,
     schema::{
-        unwrap_stream_settings, STREAM_SCHEMAS, STREAM_SCHEMAS_COMPRESSED, STREAM_SCHEMAS_LATEST,
-        STREAM_SETTINGS,
+        unwrap_stream_settings, SchemaCache, STREAM_SCHEMAS, STREAM_SCHEMAS_COMPRESSED,
+        STREAM_SCHEMAS_LATEST, STREAM_SETTINGS,
     },
 };
 #[cfg(feature = "enterprise")]
@@ -170,7 +170,7 @@ async fn list_stream_schemas(
                     stream_name,
                     stream_type,
                     schema: if fetch_schema {
-                        val.clone()
+                        val.schema().clone()
                     } else {
                         Schema::empty()
                     },
@@ -277,6 +277,7 @@ pub async fn watch() -> Result<(), anyhow::Error> {
                 let r = STREAM_SCHEMAS_LATEST.read().await;
                 let prev_schema_start_dt = if let Some(schema) = r.get(&item_key.to_owned()) {
                     schema
+                        .schema()
                         .metadata()
                         .get("start_dt")
                         .unwrap_or(&"0".to_string())
@@ -337,7 +338,10 @@ pub async fn watch() -> Result<(), anyhow::Error> {
                 w.insert(item_key.to_string(), settings);
                 drop(w);
                 let mut w = STREAM_SCHEMAS_LATEST.write().await;
-                w.insert(item_key.to_string(), latest_schema.clone());
+                w.insert(
+                    item_key.to_string(),
+                    SchemaCache::new(latest_schema.clone()),
+                );
                 drop(w);
                 if CONFIG.common.schema_cache_compress_enabled {
                     let schema_versions = schema_versions
@@ -485,7 +489,10 @@ pub async fn cache() -> Result<(), anyhow::Error> {
         w.insert(item_key.to_string(), settings);
         drop(w);
         let mut w = STREAM_SCHEMAS_LATEST.write().await;
-        w.insert(item_key.to_string(), latest_schema.clone());
+        w.insert(
+            item_key.to_string(),
+            SchemaCache::new(latest_schema.clone()),
+        );
         drop(w);
         if CONFIG.common.schema_cache_compress_enabled {
             let schema_versions = schema_versions

src/service/enrichment_table/mod.rs

@@ -34,7 +34,10 @@ use config::{
 use futures::{StreamExt, TryStreamExt};
 use infra::{
     cache::stats,
-    schema::{STREAM_SCHEMAS, STREAM_SCHEMAS_COMPRESSED, STREAM_SCHEMAS_LATEST, STREAM_SETTINGS},
+    schema::{
+        SchemaCache, STREAM_SCHEMAS, STREAM_SCHEMAS_COMPRESSED, STREAM_SCHEMAS_LATEST,
+        STREAM_SETTINGS,
+    },
 };
 
 use crate::{
@@ -43,7 +46,7 @@ use crate::{
         compact::retention,
         db, format_stream_name,
         ingestion::write_file,
-        schema::{check_for_schema, stream_schema_exists, SchemaCache},
+        schema::{check_for_schema, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };

src/service/ingestion/mod.rs

@@ -14,7 +14,7 @@
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 use std::{
-    collections::{BTreeMap, HashMap},
+    collections::{BTreeMap, HashMap, HashSet},
     sync::Arc,
 };
 
@@ -379,6 +379,7 @@ pub async fn write_file(
                     data: entry.records,
                     data_size: entry.records_size,
                 },
+                false,
             )
             .await
         {
@@ -515,7 +516,7 @@ pub async fn get_stream_routing(
 
 pub async fn get_user_defined_schema(
     streams: &[StreamParams],
-    user_defined_schema_map: &mut HashMap<String, Vec<String>>,
+    user_defined_schema_map: &mut HashMap<String, HashSet<String>>,
 ) {
     for stream in streams {
         let stream_settings =
@@ -524,7 +525,10 @@ pub async fn get_user_defined_schema(
                 .unwrap_or_default();
         if let Some(fields) = stream_settings.defined_schema_fields {
             if !fields.is_empty() {
-                user_defined_schema_map.insert(stream.stream_name.to_string(), fields);
+                user_defined_schema_map.insert(
+                    stream.stream_name.to_string(),
+                    fields.iter().cloned().collect(),
+                );
             }
         }
     }

src/service/logs/bulk.rs

@@ -14,7 +14,7 @@
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     io::{BufRead, BufReader},
     sync::Arc,
 };
@@ -32,7 +32,7 @@ use config::{
     utils::{flatten, json, schema_ext::SchemaExt, time::parse_timestamp_micro_from_value},
     BLOCKED_STREAMS, CONFIG, DISTINCT_FIELDS,
 };
-use infra::schema::unwrap_partition_time_level;
+use infra::schema::{unwrap_partition_time_level, SchemaCache};
 
 use super::{add_record, cast_to_schema_v1, StreamMeta};
 use crate::{
@@ -49,9 +49,7 @@ use crate::{
         db, format_stream_name,
         ingestion::{evaluate_trigger, write_file, TriggerAlertData},
         metadata::{distinct_values::DvItem, write, MetadataItem, MetadataType},
-        schema::{
-            get_invalid_schema_start_dt, get_upto_discard_error, stream_schema_exists, SchemaCache,
-        },
+        schema::{get_invalid_schema_start_dt, get_upto_discard_error, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };
@@ -112,7 +110,7 @@ pub async fn ingest(
 
     let mut stream_routing_map: HashMap<String, Vec<Routing>> = HashMap::new();
 
-    let mut user_defined_schema_map: HashMap<String, Vec<String>> = HashMap::new();
+    let mut user_defined_schema_map: HashMap<String, HashSet<String>> = HashMap::new();
 
     let mut next_line_is_data = false;
     let reader = BufReader::new(body.as_ref());
@@ -262,29 +260,31 @@ pub async fn ingest(
 
             // Start row based transform
             if let Some(transforms) = stream_functions_map.get(&key) {
-                let mut ret_value = value.clone();
-                ret_value = crate::service::ingestion::apply_stream_functions(
-                    transforms,
-                    ret_value,
-                    &stream_vrl_map,
-                    &stream_name,
-                    &mut runtime,
-                )?;
+                if !transforms.is_empty() {
+                    let mut ret_value = value.clone();
+                    ret_value = crate::service::ingestion::apply_stream_functions(
+                        transforms,
+                        ret_value,
+                        &stream_vrl_map,
+                        &stream_name,
+                        &mut runtime,
+                    )?;
 
-                if ret_value.is_null() || !ret_value.is_object() {
-                    bulk_res.errors = true;
-                    add_record_status(
-                        stream_name.clone(),
-                        doc_id.clone(),
-                        action.clone(),
-                        Some(value),
-                        &mut bulk_res,
-                        Some(TRANSFORM_FAILED.to_owned()),
-                        Some(TRANSFORM_FAILED.to_owned()),
-                    );
-                    continue;
-                } else {
-                    value = ret_value;
+                    if ret_value.is_null() || !ret_value.is_object() {
+                        bulk_res.errors = true;
+                        add_record_status(
+                            stream_name.clone(),
+                            doc_id.clone(),
+                            action.clone(),
+                            Some(value),
+                            &mut bulk_res,
+                            Some(TRANSFORM_FAILED.to_owned()),
+                            Some(TRANSFORM_FAILED.to_owned()),
+                        );
+                        continue;
+                    } else {
+                        value = ret_value;
+                    }
                 }
             }
             // End row based transform
@@ -296,7 +296,7 @@ pub async fn ingest(
             };
 
             if let Some(fields) = user_defined_schema_map.get(&stream_name) {
-                crate::service::logs::refactor_map(&mut local_val, fields);
+                local_val = crate::service::logs::refactor_map(local_val, fields);
             }
 
             // set _id

src/service/logs/ingest.rs

@@ -1,4 +1,4 @@
-// Copyright 2023 Zinc Labs Inc.
+// Copyright 2024 Zinc Labs Inc.
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -28,6 +28,7 @@ use config::{
     CONFIG, DISTINCT_FIELDS,
 };
 use flate2::read::GzDecoder;
+use infra::schema::SchemaCache;
 use vrl::compiler::runtime::Runtime;
 
 use crate::{
@@ -45,7 +46,7 @@ use crate::{
         ingestion::{check_ingestion_allowed, evaluate_trigger, write_file, TriggerAlertData},
         logs::StreamMeta,
         metadata::{distinct_values::DvItem, write, MetadataItem, MetadataType},
-        schema::{get_upto_discard_error, SchemaCache},
+        schema::get_upto_discard_error,
         usage::report_request_usage_stats,
     },
 };

src/service/logs/mod.rs

@@ -13,27 +13,28 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    io::Write,
+    sync::Arc,
+};
 
 use anyhow::Result;
 use arrow_schema::{DataType, Field, Schema};
 use config::{
     meta::stream::{PartitionTimeLevel, StreamPartition, StreamType},
     utils::{
-        json::{estimate_json_bytes, get_string_value, Map, Value},
+        json::{estimate_json_bytes, get_string_value, pickup_string_value, Map, Number, Value},
         schema_ext::SchemaExt,
     },
     CONFIG,
 };
-use infra::schema::unwrap_partition_time_level;
+use infra::schema::{unwrap_partition_time_level, SchemaCache};
 
 use super::{ingestion::TriggerAlertData, schema::get_invalid_schema_start_dt};
 use crate::{
     common::meta::{alerts::Alert, ingestion::RecordStatus, stream::SchemaRecords},
-    service::{
-        ingestion::get_wal_time_key,
-        schema::{check_for_schema, SchemaCache},
-    },
+    service::{ingestion::get_wal_time_key, schema::check_for_schema},
 };
 
 pub mod bulk;
@@ -150,101 +151,127 @@ pub fn cast_to_schema_v1(
     value: &mut Map<String, Value>,
     schema_map: &HashMap<&String, &DataType>,
 ) -> Result<(), anyhow::Error> {
-    let mut parse_error = String::new();
-    for (field_name, data_type) in schema_map {
-        let field_name_str = field_name.to_owned().to_owned();
-        // let field_name = field_name_str.to_owned();
-        let Some(val) = value.get(&field_name_str) else {
-            continue;
-        };
+    let mut errors = Vec::new();
+    for (key, val) in value.iter_mut() {
         if val.is_null() {
-            value.insert(field_name_str, Value::Null);
             continue;
         }
+        let Some(data_type) = schema_map.get(key) else {
+            continue;
+        };
         match data_type {
             DataType::Utf8 => {
                 if val.is_string() {
                     continue;
                 }
-                value.insert(field_name_str, Value::String(get_string_value(val)));
+                *val = Value::String(get_string_value(val));
             }
             DataType::Int64 | DataType::Int32 | DataType::Int16 | DataType::Int8 => {
                 if val.is_i64() {
                     continue;
                 }
-                if val.is_boolean() {
-                    value.insert(
-                        field_name_str,
-                        Value::Number(bool_to_serde_json_number(val.as_bool().unwrap())),
-                    );
-                    continue;
-                }
-                let val = get_string_value(val);
-                match val.parse::<i64>() {
-                    Ok(val) => {
-                        value.insert(field_name_str, Value::Number(val.into()));
-                    }
-                    Err(_) => set_parsing_error_v1(&mut parse_error, &field_name_str, data_type),
-                };
-            }
-            DataType::UInt64 | DataType::UInt32 | DataType::UInt16 | DataType::UInt8 => {
                 if val.is_u64() {
                     continue;
                 }
-                if val.is_boolean() {
-                    value.insert(
-                        field_name_str,
-                        Value::Number(bool_to_serde_json_number(val.as_bool().unwrap())),
-                    );
+                if val.is_f64() {
+                    *val = Value::Number((val.as_f64().unwrap() as i64).into());
                     continue;
                 }
-                let val = get_string_value(val);
-                match val.parse::<u64>() {
-                    Ok(val) => {
-                        value.insert(field_name_str, Value::Number(val.into()));
+                if val.is_boolean() {
+                    *val = Value::Number((val.as_bool().unwrap() as i64).into());
+                    continue;
+                }
+                let local_val = get_string_value(val);
+                match local_val.parse::<i64>() {
+                    Ok(v) => {
+                        *val = Value::Number(v.into());
                     }
-                    Err(_) => set_parsing_error_v1(&mut parse_error, &field_name_str, data_type),
+                    Err(_) => errors.push((key, *data_type)),
+                };
+            }
+            DataType::UInt64 | DataType::UInt32 | DataType::UInt16 | DataType::UInt8 => {
+                if val.is_i64() {
+                    continue;
+                }
+                if val.is_u64() {
+                    continue;
+                }
+                if val.is_f64() {
+                    *val = Value::Number((val.as_f64().unwrap() as u64).into());
+                    continue;
+                }
+                if val.is_boolean() {
+                    *val = Value::Number((val.as_bool().unwrap() as u64).into());
+                    continue;
+                }
+                let local_val = get_string_value(val);
+                match local_val.parse::<u64>() {
+                    Ok(v) => {
+                        *val = Value::Number(v.into());
+                    }
+                    Err(_) => errors.push((key, *data_type)),
                 };
             }
             DataType::Float64 | DataType::Float32 | DataType::Float16 => {
                 if val.is_f64() {
                     continue;
                 }
+                if val.is_i64() {
+                    *val = Value::Number(Number::from_f64(val.as_i64().unwrap() as f64).unwrap());
+                    continue;
+                }
+                if val.is_u64() {
+                    *val = Value::Number(Number::from_f64(val.as_u64().unwrap() as f64).unwrap());
+                    continue;
+                }
                 if val.is_boolean() {
-                    value.insert(
-                        field_name_str,
-                        Value::Number(bool_to_serde_json_number(val.as_bool().unwrap())),
+                    *val = Value::Number(
+                        Number::from_f64((val.as_bool().unwrap() as i64) as f64).unwrap(),
                     );
                     continue;
                 }
-                let val = get_string_value(val);
-                match val.parse::<f64>() {
-                    Ok(val) => {
-                        value.insert(
-                            field_name_str,
-                            Value::Number(serde_json::Number::from_f64(val).unwrap()),
-                        );
+                let local_val = get_string_value(val);
+                match local_val.parse::<f64>() {
+                    Ok(local_val) => {
+                        *val = Value::Number(serde_json::Number::from_f64(local_val).unwrap());
                     }
-                    Err(_) => set_parsing_error_v1(&mut parse_error, &field_name_str, data_type),
+                    Err(_) => errors.push((key, *data_type)),
                 };
             }
             DataType::Boolean => {
                 if val.is_boolean() {
                     continue;
                 }
-                let val = get_string_value(val);
-                match val.parse::<bool>() {
-                    Ok(val) => {
-                        value.insert(field_name_str, Value::Bool(val));
+                if val.is_i64() {
+                    *val = Value::Bool(val.as_i64().unwrap() > 0);
+                    continue;
+                }
+                if val.is_u64() {
+                    *val = Value::Bool(val.as_u64().unwrap() > 0);
+                    continue;
+                }
+                if val.is_f64() {
+                    *val = Value::Bool(val.as_f64().unwrap() > 0.0);
+                    continue;
+                }
+                let local_val: String = get_string_value(val);
+                match local_val.parse::<bool>() {
+                    Ok(local_val) => {
+                        *val = Value::Bool(local_val);
                     }
-                    Err(_) => set_parsing_error_v1(&mut parse_error, &field_name_str, data_type),
+                    Err(_) => errors.push((key, *data_type)),
                 };
             }
-            _ => set_parsing_error_v1(&mut parse_error, &field_name_str, data_type),
+            _ => errors.push((key, *data_type)),
         };
     }
-    if !parse_error.is_empty() {
-        Err(anyhow::Error::msg(parse_error))
+    if !errors.is_empty() {
+        let error_message = errors
+            .iter()
+            .map(|(field, dt)| format!("Failed to cast Field: {}, DataType: {:?}", field, dt))
+            .collect::<Vec<_>>()
+            .join(", ");
+        Err(anyhow::Error::msg(error_message))
     } else {
         Ok(())
     }
@@ -266,7 +293,7 @@ async fn add_valid_record(
         .unwrap();
 
     // check schema
-    let (schema_evolution, _) = match check_for_schema(
+    let schema_evolution = match check_for_schema(
         &stream_meta.org_id,
         &stream_meta.stream_name,
         StreamType::Logs,
@@ -390,13 +417,6 @@ fn set_parsing_error(parse_error: &mut String, field: &Field) {
     ));
 }
 
-fn set_parsing_error_v1(parse_error: &mut String, field_name: &str, field_data_type: &DataType) {
-    parse_error.push_str(&format!(
-        "Failed to cast {} to type {} ",
-        field_name, field_data_type
-    ));
-}
-
 async fn add_record(
     stream_meta: &StreamMeta<'_>,
     write_buf: &mut HashMap<String, SchemaRecords>,
@@ -431,11 +451,6 @@ async fn add_record(
     Ok(())
 }
 
-fn bool_to_serde_json_number(value: bool) -> serde_json::Number {
-    let num = value as i64;
-    serde_json::Number::from(num)
-}
-
 struct StreamMeta<'a> {
     org_id: String,
     stream_name: String,
@@ -444,24 +459,43 @@ struct StreamMeta<'a> {
     stream_alerts_map: &'a HashMap<String, Vec<Alert>>,
 }
 
-pub fn refactor_map(original_map: &mut Map<String, Value>, defined_schema_keys: &[String]) {
-    let mut non_schema_map = HashMap::new();
+pub fn refactor_map(
+    original_map: Map<String, Value>,
+    defined_schema_keys: &HashSet<String>,
+) -> Map<String, Value> {
+    let mut new_map = Map::with_capacity(defined_schema_keys.len() + 2);
+    let mut non_schema_map = Vec::with_capacity(1024); // 1KB
 
-    for (key, value) in original_map.iter() {
-        if !defined_schema_keys.contains(key) {
-            non_schema_map.insert(key.clone(), get_string_value(value));
+    let mut has_elements = false;
+    non_schema_map.write_all(b"{").unwrap();
+    for (key, value) in original_map {
+        if defined_schema_keys.contains(&key) {
+            new_map.insert(key, value);
+        } else {
+            if has_elements {
+                non_schema_map.write_all(b",").unwrap();
+            } else {
+                has_elements = true;
+            }
+            non_schema_map.write_all(b"\"").unwrap();
+            non_schema_map.write_all(key.as_bytes()).unwrap();
+            non_schema_map.write_all(b"\":\"").unwrap();
+            non_schema_map
+                .write_all(pickup_string_value(value).as_bytes())
+                .unwrap();
+            non_schema_map.write_all(b"\"").unwrap();
         }
     }
+    non_schema_map.write_all(b"}").unwrap();
 
-    original_map.retain(|key, _| defined_schema_keys.contains(key));
-
-    if !non_schema_map.is_empty() {
-        let non_schema_json = serde_json::to_string(&non_schema_map).unwrap_or_default();
-        original_map.insert(
+    if has_elements {
+        new_map.insert(
             CONFIG.common.all_fields_name.to_string(),
-            Value::String(non_schema_json),
+            Value::String(String::from_utf8(non_schema_map).unwrap()),
         );
     }
+
+    new_map
 }
 
 #[cfg(test)]

src/service/logs/multi.rs

@@ -1,4 +1,4 @@
-// Copyright 2023 Zinc Labs Inc.
+// Copyright 2024 Zinc Labs Inc.
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -27,6 +27,7 @@ use config::{
     utils::{flatten, json, time::parse_timestamp_micro_from_value},
     CONFIG, DISTINCT_FIELDS,
 };
+use infra::schema::SchemaCache;
 
 use crate::{
     common::meta::{
@@ -39,7 +40,7 @@ use crate::{
         ingestion::{check_ingestion_allowed, evaluate_trigger, write_file, TriggerAlertData},
         logs::StreamMeta,
         metadata::{distinct_values::DvItem, write, MetadataItem, MetadataType},
-        schema::{get_upto_discard_error, SchemaCache},
+        schema::get_upto_discard_error,
         usage::report_request_usage_stats,
     },
 };

src/service/logs/otlp_grpc.rs

@@ -1,4 +1,4 @@
-// Copyright 2023 Zinc Labs Inc.
+// Copyright 2024 Zinc Labs Inc.
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -26,6 +26,7 @@ use config::{
     utils::{flatten, json, time::parse_timestamp_micro_from_value},
     CONFIG, DISTINCT_FIELDS,
 };
+use infra::schema::SchemaCache;
 use opentelemetry::trace::{SpanId, TraceId};
 use opentelemetry_proto::tonic::collector::logs::v1::{
     ExportLogsServiceRequest, ExportLogsServiceResponse,
@@ -48,7 +49,7 @@ use crate::{
             write_file, TriggerAlertData,
         },
         metadata::{distinct_values::DvItem, write, MetadataItem, MetadataType},
-        schema::{get_upto_discard_error, stream_schema_exists, SchemaCache},
+        schema::{get_upto_discard_error, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };

src/service/logs/otlp_http.rs

@@ -25,6 +25,7 @@ use config::{
     utils::{flatten, json},
     CONFIG, DISTINCT_FIELDS,
 };
+use infra::schema::SchemaCache;
 use opentelemetry::trace::{SpanId, TraceId};
 use opentelemetry_proto::tonic::collector::logs::v1::{
     ExportLogsPartialSuccess, ExportLogsServiceRequest, ExportLogsServiceResponse,
@@ -44,7 +45,7 @@ use crate::{
         db, get_formatted_stream_name,
         ingestion::{evaluate_trigger, get_val_for_attr, write_file, TriggerAlertData},
         metadata::{distinct_values::DvItem, write, MetadataItem, MetadataType},
-        schema::{get_upto_discard_error, stream_schema_exists, SchemaCache},
+        schema::{get_upto_discard_error, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };

src/service/logs/syslog.rs

@@ -1,4 +1,4 @@
-// Copyright 2023 Zinc Labs Inc.
+// Copyright 2024 Zinc Labs Inc.
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -25,6 +25,7 @@ use config::{
     utils::{flatten, json, time::parse_timestamp_micro_from_value},
     CONFIG, DISTINCT_FIELDS,
 };
+use infra::schema::SchemaCache;
 use syslog_loose::{Message, ProcId, Protocol};
 
 use super::StreamMeta;
@@ -43,7 +44,7 @@ use crate::{
         db, get_formatted_stream_name,
         ingestion::{evaluate_trigger, write_file, TriggerAlertData},
         metadata::{distinct_values::DvItem, write, MetadataItem, MetadataType},
-        schema::{get_upto_discard_error, SchemaCache},
+        schema::get_upto_discard_error,
     },
 };
 

src/service/metrics/json.rs

@@ -28,7 +28,7 @@ use config::{
     CONFIG,
 };
 use datafusion::arrow::datatypes::Schema;
-use infra::schema::unwrap_partition_time_level;
+use infra::schema::{unwrap_partition_time_level, SchemaCache};
 use vrl::compiler::runtime::Runtime;
 
 use super::get_exclude_labels;
@@ -42,7 +42,7 @@ use crate::{
     service::{
         db, format_stream_name,
         ingestion::{get_wal_time_key, write_file},
-        schema::{check_for_schema, SchemaCache},
+        schema::check_for_schema,
         usage::report_request_usage_stats,
     },
 };
@@ -123,13 +123,7 @@ pub async fn ingest(org_id: &str, body: web::Bytes, thread_id: usize) -> Result<
                     )
                     .await?;
                 }
-                let fields_map = schema
-                    .fields()
-                    .iter()
-                    .enumerate()
-                    .map(|(i, f)| (f.name().to_owned(), i))
-                    .collect();
-                stream_schema_map.insert(stream_name.clone(), SchemaCache::new(schema, fields_map));
+                stream_schema_map.insert(stream_name.clone(), SchemaCache::new(schema));
             }
             continue;
         }
@@ -223,13 +217,7 @@ pub async fn ingest(org_id: &str, body: web::Bytes, thread_id: usize) -> Result<
                 )
                 .await;
             }
-            let fields_map = schema
-                .fields()
-                .iter()
-                .enumerate()
-                .map(|(i, f)| (f.name().to_owned(), i))
-                .collect();
-            stream_schema_map.insert(stream_name.clone(), SchemaCache::new(schema, fields_map));
+            stream_schema_map.insert(stream_name.clone(), SchemaCache::new(schema));
         }
 
         // check for schema evolution

src/service/metrics/otlp_grpc.rs

@@ -29,7 +29,7 @@ use config::{
     CONFIG,
 };
 use hashbrown::HashSet;
-use infra::schema::{unwrap_partition_time_level, update_setting};
+use infra::schema::{unwrap_partition_time_level, update_setting, SchemaCache};
 use opentelemetry::trace::{SpanId, TraceId};
 use opentelemetry_proto::tonic::{
     collector::metrics::v1::{ExportMetricsServiceRequest, ExportMetricsServiceResponse},
@@ -52,7 +52,7 @@ use crate::{
             write_file, TriggerAlertData,
         },
         metrics::{format_label_name, get_exclude_labels},
-        schema::{check_for_schema, stream_schema_exists, SchemaCache},
+        schema::{check_for_schema, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };

src/service/metrics/otlp_http.rs

@@ -29,7 +29,7 @@ use config::{
     CONFIG,
 };
 use hashbrown::HashSet;
-use infra::schema::{unwrap_partition_time_level, update_setting};
+use infra::schema::{unwrap_partition_time_level, update_setting, SchemaCache};
 use opentelemetry::trace::{SpanId, TraceId};
 use opentelemetry_proto::tonic::{
     collector::metrics::v1::{ExportMetricsServiceRequest, ExportMetricsServiceResponse},
@@ -50,7 +50,7 @@ use crate::{
         db, format_stream_name,
         ingestion::{evaluate_trigger, get_val_for_attr, write_file, TriggerAlertData},
         metrics::{format_label_name, get_exclude_labels, otlp_grpc::handle_grpc_request},
-        schema::{check_for_schema, stream_schema_exists, SchemaCache},
+        schema::{check_for_schema, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };

src/service/metrics/prom.rs

@@ -31,7 +31,7 @@ use datafusion::arrow::datatypes::Schema;
 use infra::{
     cache::stats,
     errors::{Error, Result},
-    schema::{unwrap_partition_time_level, update_setting},
+    schema::{unwrap_partition_time_level, update_setting, SchemaCache},
 };
 use promql_parser::{label::MatchOp, parser};
 use prost::Message;
@@ -51,7 +51,7 @@ use crate::{
         db, format_stream_name,
         ingestion::{evaluate_trigger, write_file, TriggerAlertData},
         metrics::format_label_name,
-        schema::{check_for_schema, stream_schema_exists, SchemaCache},
+        schema::{check_for_schema, stream_schema_exists},
         search as search_service,
         usage::report_request_usage_stats,
     },

src/service/mod.rs

@@ -66,35 +66,26 @@ pub fn format_partition_key(input: &str) -> String {
 // format stream name
 pub async fn get_formatted_stream_name(
     params: &mut StreamParams,
-    schema_map: &mut HashMap<String, schema::SchemaCache>,
+    schema_map: &mut HashMap<String, infra::schema::SchemaCache>,
 ) -> String {
     let mut stream_name = params.stream_name.to_string();
 
-    let schema = infra::schema::get(&params.org_id, &stream_name, params.stream_type)
+    let schema = infra::schema::get_cache(&params.org_id, &stream_name, params.stream_type)
         .await
         .unwrap();
 
-    let schema = if schema.fields().is_empty() {
+    let schema = if schema.fields_map().is_empty() {
         stream_name = RE_CORRECT_STREAM_NAME
             .replace_all(&stream_name, "_")
             .to_string();
-        infra::schema::get(&params.org_id, &stream_name, params.stream_type)
+        infra::schema::get_cache(&params.org_id, &stream_name, params.stream_type)
             .await
             .unwrap()
     } else {
         schema
     };
 
-    let fields_map = schema
-        .fields()
-        .iter()
-        .enumerate()
-        .map(|(i, f)| (f.name().to_owned(), i))
-        .collect();
-    schema_map.insert(
-        stream_name.to_owned(),
-        schema::SchemaCache::new(schema, fields_map),
-    );
+    schema_map.insert(stream_name.to_owned(), schema);
 
     params.stream_name = stream_name.to_owned().into();
 

src/service/schema.rs

@@ -22,6 +22,8 @@ use config::{
     CONFIG,
 };
 use datafusion::arrow::datatypes::{Field, Schema};
+use hashbrown::HashSet;
+use infra::schema::{get_settings, unwrap_stream_settings, SchemaCache};
 use serde_json::{Map, Value};
 
 use crate::{
@@ -52,35 +54,6 @@ pub(crate) fn get_request_columns_limit_error(
     )
 }
 
-pub struct SchemaCache {
-    schema: Schema,
-    fields_map: HashMap<String, usize>,
-    hash_key: String,
-}
-
-impl SchemaCache {
-    pub fn new(schema: Schema, fields_map: HashMap<String, usize>) -> Self {
-        let hash_key = schema.hash_key();
-        Self {
-            schema,
-            fields_map,
-            hash_key,
-        }
-    }
-
-    pub fn hash_key(&self) -> &str {
-        &self.hash_key
-    }
-
-    pub fn schema(&self) -> &Schema {
-        &self.schema
-    }
-
-    pub fn fields_map(&self) -> &HashMap<String, usize> {
-        &self.fields_map
-    }
-}
-
 pub async fn check_for_schema(
     org_id: &str,
     stream_name: &str,
@@ -88,48 +61,31 @@ pub async fn check_for_schema(
     stream_schema_map: &mut HashMap<String, SchemaCache>,
     record_val: Vec<&Map<String, Value>>,
     record_ts: i64,
-) -> Result<(SchemaEvolution, Option<Schema>)> {
+) -> Result<SchemaEvolution> {
     if !stream_schema_map.contains_key(stream_name) {
-        let schema = infra::schema::get(org_id, stream_name, stream_type)
-            .await
-            .unwrap();
-        let fields_map = schema
-            .fields()
-            .iter()
-            .enumerate()
-            .map(|(i, f)| (f.name().to_owned(), i))
-            .collect();
-        stream_schema_map.insert(
-            stream_name.to_string(),
-            SchemaCache::new(schema, fields_map),
-        );
+        let schema = infra::schema::get_cache(org_id, stream_name, stream_type).await?;
+        stream_schema_map.insert(stream_name.to_string(), schema);
     }
     let schema = stream_schema_map.get(stream_name).unwrap();
     if !schema.schema().fields().is_empty() && CONFIG.common.skip_schema_validation {
-        return Ok((
-            SchemaEvolution {
-                schema_compatible: true,
-                is_schema_changed: false,
-                types_delta: None,
-            },
-            None,
-        ));
+        return Ok(SchemaEvolution {
+            schema_compatible: true,
+            is_schema_changed: false,
+            types_delta: None,
+        });
     }
 
     // get infer schema
     let value_iter = record_val.into_iter();
-    let inferred_schema = infer_json_schema_from_map(value_iter, stream_type).unwrap();
+    let inferred_schema = infer_json_schema_from_map(value_iter, stream_type)?;
 
     // fast path
     if schema.schema().fields.eq(&inferred_schema.fields) {
-        return Ok((
-            SchemaEvolution {
-                schema_compatible: true,
-                is_schema_changed: false,
-                types_delta: None,
-            },
-            None,
-        ));
+        return Ok(SchemaEvolution {
+            schema_compatible: true,
+            is_schema_changed: false,
+            types_delta: None,
+        });
     }
 
     if inferred_schema.fields.len() > CONFIG.limit.req_cols_per_record_limit {
@@ -148,21 +104,28 @@ pub async fn check_for_schema(
             let inferred_schema = if field_datatype_delta.is_empty() {
                 inferred_schema
             } else {
-                let schema_latest = stream_schema_map.get(stream_name).unwrap();
-                inferred_schema.cloned_from(schema_latest.schema())
+                inferred_schema.cloned_from(schema.schema())
             };
-            return Ok((
-                SchemaEvolution {
-                    schema_compatible: true,
-                    is_schema_changed: false,
-                    types_delta: Some(field_datatype_delta),
-                },
-                Some(inferred_schema),
-            ));
+            // check defined_schema_fields
+            let stream_setting = get_settings(org_id, stream_name, stream_type).await;
+            let defined_schema_fields = stream_setting
+                .and_then(|s| s.defined_schema_fields)
+                .unwrap_or_default();
+            if !defined_schema_fields.is_empty() {
+                let schema = SchemaCache::new(inferred_schema);
+                let schema =
+                    generate_schema_for_defined_schema_fields(schema, &defined_schema_fields);
+                stream_schema_map.insert(stream_name.to_string(), schema);
+            }
+            return Ok(SchemaEvolution {
+                schema_compatible: true,
+                is_schema_changed: false,
+                types_delta: Some(field_datatype_delta),
+            });
         }
         if !field_datatype_delta.is_empty() {
             // check if the min_ts < current_version_created_at, if yes, discard the data
-            let schema_metadata = schema.schema.metadata();
+            let schema_metadata = schema.schema().metadata();
             if let Some(start_dt) = schema_metadata.get("start_dt") {
                 let created_at = start_dt.parse().unwrap_or_default();
                 if record_ts <= created_at {
@@ -188,12 +151,7 @@ pub async fn check_for_schema(
         is_schema_changed: false,
         types_delta: None,
     });
-
-    // generate new schema
-    let schema_latest = stream_schema_map.get(stream_name).unwrap();
-    let inferred_schema = inferred_schema.cloned_from(schema_latest.schema());
-
-    Ok((ret, Some(inferred_schema)))
+    Ok(ret)
 }
 
 pub async fn get_merged_schema(
@@ -239,15 +197,9 @@ async fn handle_diff_schema(
     if is_new {
         let mut metadata = HashMap::with_capacity(1);
         metadata.insert("created_at".to_string(), record_ts.to_string());
-        let fields_map = inferred_schema
-            .fields()
-            .iter()
-            .enumerate()
-            .map(|(i, f)| (f.name().to_owned(), i))
-            .collect();
         stream_schema_map.insert(
             stream_name.to_string(),
-            SchemaCache::new(inferred_schema.clone().with_metadata(metadata), fields_map),
+            SchemaCache::new(inferred_schema.clone().with_metadata(metadata)),
         );
     }
 
@@ -313,17 +265,17 @@ async fn handle_diff_schema(
         .await;
     }
 
-    let fields_map = final_schema
-        .fields()
-        .iter()
-        .enumerate()
-        .map(|(i, f)| (f.name().to_owned(), i))
-        .collect();
+    // check defined_schema_fields
+    let stream_setting = get_settings(org_id, stream_name, stream_type).await;
+    let defined_schema_fields = stream_setting
+        .and_then(|s| s.defined_schema_fields)
+        .unwrap_or_default();
+    let final_schema = SchemaCache::new(final_schema);
+    let final_schema =
+        generate_schema_for_defined_schema_fields(final_schema, &defined_schema_fields);
+
     // update thread cache
-    stream_schema_map.insert(
-        stream_name.to_string(),
-        SchemaCache::new(final_schema, fields_map),
-    );
+    stream_schema_map.insert(stream_name.to_string(), final_schema);
 
     Ok(Some(SchemaEvolution {
         schema_compatible: true,
@@ -332,10 +284,44 @@ async fn handle_diff_schema(
     }))
 }
 
+// if defined_schema_fields is not empty, and schema fields greater than defined_schema_fields + 10,
+// then we will use defined_schema_fields
+fn generate_schema_for_defined_schema_fields(
+    schema: SchemaCache,
+    fields: &[String],
+) -> SchemaCache {
+    if fields.is_empty() || schema.fields_map().len() < fields.len() + 10 {
+        return schema;
+    }
+
+    let mut fields: HashSet<_> = fields.iter().collect();
+    if !fields.contains(&CONFIG.common.column_timestamp) {
+        fields.insert(&CONFIG.common.column_timestamp);
+    }
+    if !fields.contains(&CONFIG.common.all_fields_name) {
+        fields.insert(&CONFIG.common.all_fields_name);
+    }
+    let mut new_fields = Vec::with_capacity(fields.len());
+    for field in fields {
+        if let Some(f) = schema.fields_map().get(field) {
+            new_fields.push(schema.schema().fields()[*f].clone());
+        }
+    }
+    SchemaCache::new(Schema::new_with_metadata(
+        new_fields,
+        schema.schema().metadata().clone(),
+    ))
+}
+
 fn get_schema_changes(schema: &SchemaCache, inferred_schema: &Schema) -> (bool, Vec<Field>) {
     let mut is_schema_changed = false;
     let mut field_datatype_delta: Vec<Field> = vec![];
 
+    let stream_setting = unwrap_stream_settings(schema.schema());
+    let defined_schema_fields = stream_setting
+        .and_then(|s| s.defined_schema_fields)
+        .unwrap_or_default();
+
     for item in inferred_schema.fields.iter() {
         let item_name = item.name();
         let item_data_type = item.data_type();
@@ -345,6 +331,9 @@ fn get_schema_changes(schema: &SchemaCache, inferred_schema: &Schema) -> (bool,
                 is_schema_changed = true;
             }
             Some(idx) => {
+                if !defined_schema_fields.is_empty() && !defined_schema_fields.contains(item_name) {
+                    continue;
+                }
                 let existing_field: Arc<Field> = schema.schema().fields()[*idx].clone();
                 if existing_field.data_type() != item_data_type {
                     if !CONFIG.common.widening_schema_evolution {
@@ -387,16 +376,8 @@ pub async fn stream_schema_exists(
             let schema = infra::schema::get(org_id, stream_name, stream_type)
                 .await
                 .unwrap();
-            let fields_map = schema
-                .fields()
-                .iter()
-                .enumerate()
-                .map(|(i, f)| (f.name().to_owned(), i))
-                .collect();
-            stream_schema_map.insert(
-                stream_name.to_string(),
-                SchemaCache::new(schema.clone(), fields_map),
-            );
+
+            stream_schema_map.insert(stream_name.to_string(), SchemaCache::new(schema.clone()));
             schema
         }
     };
@@ -437,16 +418,8 @@ mod tests {
             Field::new("_timestamp", DataType::Int64, false),
         ]);
         let mut map: HashMap<String, SchemaCache> = HashMap::new();
-        let fields_map = schema
-            .fields()
-            .iter()
-            .enumerate()
-            .map(|(i, f)| (f.name().to_owned(), i))
-            .collect();
-        map.insert(
-            stream_name.to_string(),
-            SchemaCache::new(schema, fields_map),
-        );
+
+        map.insert(stream_name.to_string(), SchemaCache::new(schema));
         let result = check_for_schema(
             org_name,
             stream_name,
@@ -457,7 +430,7 @@ mod tests {
         )
         .await
         .unwrap();
-        assert!(result.0.schema_compatible);
+        assert!(result.schema_compatible);
     }
 
     #[tokio::test]

src/service/search/datafusion/exec.rs

@@ -13,7 +13,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-use std::{str::FromStr, sync::Arc};
+use std::{collections::HashSet, str::FromStr, sync::Arc};
 
 use arrow_schema::Field;
 use config::{
@@ -22,7 +22,10 @@ use config::{
         sql,
         stream::{FileKey, FileMeta, StreamType},
     },
-    utils::{flatten, json, parquet::new_parquet_writer, schema::infer_json_schema_from_values},
+    utils::{
+        flatten, json, parquet::new_parquet_writer, schema::infer_json_schema_from_values,
+        schema_ext::SchemaExt,
+    },
     CONFIG, PARQUET_BATCH_SIZE,
 };
 use datafusion::{
@@ -1148,13 +1151,18 @@ pub async fn merge_parquet_files(
             } else {
                 config::SQL_FULL_TEXT_SEARCH_FIELDS.to_vec()
             };
+            let schema_fields = schema.as_ref().simple_fields();
 
-            // add _timestamp column to columns_to_index
-            if !columns_to_index.contains(&CONFIG.common.column_timestamp.to_string()) {
-                columns_to_index.push(CONFIG.common.column_timestamp.to_string());
-            }
+            let schema_fields: HashSet<&str> = schema_fields.iter().map(|f| f.0.as_str()).collect();
+
+            columns_to_index.retain(|f| schema_fields.contains(f.as_str()));
 
             if !columns_to_index.is_empty() {
+                // add _timestamp column to columns_to_index
+                if !columns_to_index.contains(&CONFIG.common.column_timestamp.to_string()) {
+                    columns_to_index.push(CONFIG.common.column_timestamp.to_string());
+                }
+
                 let selected_column_indices: Vec<usize> = columns_to_index
                     .iter()
                     .filter_map(|name| batch.schema().index_of(name).ok())

src/service/traces/mod.rs

@@ -28,7 +28,7 @@ use config::{
     utils::{flatten, json, schema_ext::SchemaExt},
     CONFIG, DISTINCT_FIELDS,
 };
-use infra::schema::unwrap_partition_time_level;
+use infra::schema::{unwrap_partition_time_level, SchemaCache};
 use opentelemetry::trace::{SpanId, TraceId};
 use opentelemetry_proto::tonic::{
     collector::trace::v1::{
@@ -52,7 +52,7 @@ use crate::{
             distinct_values::DvItem, trace_list_index::TraceListItem, write, MetadataItem,
             MetadataType,
         },
-        schema::{check_for_schema, stream_schema_exists, SchemaCache},
+        schema::{check_for_schema, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };

src/service/traces/otlp_http.rs

@@ -27,7 +27,7 @@ use config::{
     utils::{flatten, json, schema_ext::SchemaExt},
     CONFIG, DISTINCT_FIELDS,
 };
-use infra::schema::unwrap_partition_time_level;
+use infra::schema::{unwrap_partition_time_level, SchemaCache};
 use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
 use prost::Message;
 
@@ -46,7 +46,7 @@ use crate::{
             distinct_values::DvItem, trace_list_index::TraceListItem, write, MetadataItem,
             MetadataType,
         },
-        schema::{check_for_schema, stream_schema_exists, SchemaCache},
+        schema::{check_for_schema, stream_schema_exists},
         usage::report_request_usage_stats,
     },
 };