From ed73fcd5ff3b4c9663ed7d5fb16dc53be04d55d8 Mon Sep 17 00:00:00 2001 From: Ashish Kolhe <35160958+oasisk@users.noreply.github.com> Date: Sat, 11 May 2024 22:42:43 +0530 Subject: [PATCH] feat: only write defined schema fields to parquet (#3454) Co-authored-by: Hengfei Yang --- src/config/src/utils/json.rs | 30 ++++ src/config/src/utils/record_batch_ext.rs | 17 +- src/config/src/utils/schema_ext.rs | 28 ++-- src/handler/http/request/search/mod.rs | 2 +- src/handler/http/request/status/mod.rs | 2 +- src/infra/src/schema/mod.rs | 63 ++++++- src/ingester/src/entry.rs | 15 ++ src/ingester/src/lib.rs | 17 +- src/ingester/src/writer.rs | 38 ++++- src/job/files/parquet.rs | 3 +- src/service/compact/file_list.rs | 2 +- src/service/db/schema.rs | 17 +- src/service/enrichment_table/mod.rs | 7 +- src/service/ingestion/mod.rs | 10 +- src/service/logs/bulk.rs | 58 +++---- src/service/logs/ingest.rs | 5 +- src/service/logs/mod.rs | 204 +++++++++++++---------- src/service/logs/multi.rs | 5 +- src/service/logs/otlp_grpc.rs | 5 +- src/service/logs/otlp_http.rs | 3 +- src/service/logs/syslog.rs | 5 +- src/service/metrics/json.rs | 20 +-- src/service/metrics/otlp_grpc.rs | 4 +- src/service/metrics/otlp_http.rs | 4 +- src/service/metrics/prom.rs | 4 +- src/service/mod.rs | 19 +-- src/service/schema.rs | 203 ++++++++++------------ src/service/search/datafusion/exec.rs | 20 ++- src/service/traces/mod.rs | 4 +- src/service/traces/otlp_http.rs | 4 +- 30 files changed, 488 insertions(+), 330 deletions(-) diff --git a/src/config/src/utils/json.rs b/src/config/src/utils/json.rs index 9f36558d7c..9666465544 100644 --- a/src/config/src/utils/json.rs +++ b/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::().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::().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::().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 { diff --git a/src/config/src/utils/record_batch_ext.rs b/src/config/src/utils/record_batch_ext.rs index b5300adc0d..919c25f396 100644 --- a/src/config/src/utils/record_batch_ext.rs +++ b/src/config/src/utils/record_batch_ext.rs @@ -47,14 +47,20 @@ pub fn convert_json_to_record_batch( data: &[Arc], ) -> Result { // 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)> = 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)) } } diff --git a/src/config/src/utils/schema_ext.rs b/src/config/src/utils/schema_ext.rs index 8336e4c287..26f1c04618 100644 --- a/src/config/src/utils/schema_ext.rs +++ b/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::>(); Schema::new(fields) } diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs index 63e1abf80f..12f4c0b2b4 100644 --- a/src/handler/http/request/search/mod.rs +++ b/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() }; diff --git a/src/handler/http/request/status/mod.rs b/src/handler/http/request/status/mod.rs index 1f3c3b4aae..599d7dd6ef 100644 --- a/src/handler/http/request/status/mod.rs +++ b/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) diff --git a/src/infra/src/schema/mod.rs b/src/infra/src/schema/mod.rs index 73efcd6105..9097b16768 100644 --- a/src/infra/src/schema/mod.rs +++ b/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 . +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>> = Lazy::new(Default::default); pub static STREAM_SCHEMAS_COMPRESSED: Lazy>> = Lazy::new(Default::default); -pub static STREAM_SCHEMAS_LATEST: Lazy> = Lazy::new(Default::default); +pub static STREAM_SCHEMAS_LATEST: Lazy> = + Lazy::new(Default::default); pub static STREAM_SCHEMAS_FIELDS: Lazy)>> = Lazy::new(Default::default); pub static STREAM_SETTINGS: Lazy> = 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 { + 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, + 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 { + &self.fields_map + } +} + pub fn is_widening_conversion(from: &DataType, to: &DataType) -> bool { let allowed_type = match from { DataType::Boolean => vec![DataType::Utf8], diff --git a/src/ingester/src/entry.rs b/src/ingester/src/entry.rs index a5757ebe9a..d4f3ace5f7 100644 --- a/src/ingester/src/entry.rs +++ b/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> { 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, diff --git a/src/ingester/src/lib.rs b/src/ingester/src/lib.rs index 7b292fa90f..d6b3d524be 100644 --- a/src/ingester/src/lib.rs +++ b/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(()) } diff --git a/src/ingester/src/writer.rs b/src/ingester/src/writer.rs index 4d6e125f68..f2a3578cce 100644 --- a/src/ingester/src/writer.rs +++ b/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, 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, + 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(()) } diff --git a/src/job/files/parquet.rs b/src/job/files/parquet.rs index eb094633e1..89bd8a4068 100644 --- a/src/job/files/parquet.rs +++ b/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, }, diff --git a/src/service/compact/file_list.rs b/src/service/compact/file_list.rs index f9184e00b3..97d7485b33 100644 --- a/src/service/compact/file_list.rs +++ b/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!( diff --git a/src/service/db/schema.rs b/src/service/db/schema.rs index 4854f9f6c0..d64fd96779 100644 --- a/src/service/db/schema.rs +++ b/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 diff --git a/src/service/enrichment_table/mod.rs b/src/service/enrichment_table/mod.rs index 24e61196f0..f0b21b5178 100644 --- a/src/service/enrichment_table/mod.rs +++ b/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, }, }; diff --git a/src/service/ingestion/mod.rs b/src/service/ingestion/mod.rs index 9d56532fb6..90a7fb2b6d 100644 --- a/src/service/ingestion/mod.rs +++ b/src/service/ingestion/mod.rs @@ -14,7 +14,7 @@ // along with this program. If not, see . 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>, + user_defined_schema_map: &mut HashMap>, ) { 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(), + ); } } } diff --git a/src/service/logs/bulk.rs b/src/service/logs/bulk.rs index 6261552541..e895fd9c24 100644 --- a/src/service/logs/bulk.rs +++ b/src/service/logs/bulk.rs @@ -14,7 +14,7 @@ // along with this program. If not, see . 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> = HashMap::new(); - let mut user_defined_schema_map: HashMap> = HashMap::new(); + let mut user_defined_schema_map: HashMap> = 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 diff --git a/src/service/logs/ingest.rs b/src/service/logs/ingest.rs index 243e9af5b2..94b5d06331 100644 --- a/src/service/logs/ingest.rs +++ b/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, }, }; diff --git a/src/service/logs/mod.rs b/src/service/logs/mod.rs index ff7c79d0b0..cec068811d 100644 --- a/src/service/logs/mod.rs +++ b/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 . -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, 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::() { - 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::() { - 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::() { + 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::() { + 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::() { - 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::() { + 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::() { - 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::() { + 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::>() + .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, @@ -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>, } -pub fn refactor_map(original_map: &mut Map, defined_schema_keys: &[String]) { - let mut non_schema_map = HashMap::new(); +pub fn refactor_map( + original_map: Map, + defined_schema_keys: &HashSet, +) -> Map { + 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)] diff --git a/src/service/logs/multi.rs b/src/service/logs/multi.rs index 3f689229d7..240109b53b 100644 --- a/src/service/logs/multi.rs +++ b/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, }, }; diff --git a/src/service/logs/otlp_grpc.rs b/src/service/logs/otlp_grpc.rs index b8646b68a9..8ac66c1d04 100644 --- a/src/service/logs/otlp_grpc.rs +++ b/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, }, }; diff --git a/src/service/logs/otlp_http.rs b/src/service/logs/otlp_http.rs index 9689727a97..44fa550988 100644 --- a/src/service/logs/otlp_http.rs +++ b/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, }, }; diff --git a/src/service/logs/syslog.rs b/src/service/logs/syslog.rs index fd5b5988d9..80b230d768 100644 --- a/src/service/logs/syslog.rs +++ b/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, }, }; diff --git a/src/service/metrics/json.rs b/src/service/metrics/json.rs index b03112e91e..1a74d104a7 100644 --- a/src/service/metrics/json.rs +++ b/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 diff --git a/src/service/metrics/otlp_grpc.rs b/src/service/metrics/otlp_grpc.rs index a947998d79..6b546566ba 100644 --- a/src/service/metrics/otlp_grpc.rs +++ b/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, }, }; diff --git a/src/service/metrics/otlp_http.rs b/src/service/metrics/otlp_http.rs index 2644ef5034..3f3d9c3c10 100644 --- a/src/service/metrics/otlp_http.rs +++ b/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, }, }; diff --git a/src/service/metrics/prom.rs b/src/service/metrics/prom.rs index caf9422c97..e1fbd09991 100644 --- a/src/service/metrics/prom.rs +++ b/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, }, diff --git a/src/service/mod.rs b/src/service/mod.rs index c948c82560..0a0c1335b3 100644 --- a/src/service/mod.rs +++ b/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, + schema_map: &mut HashMap, ) -> String { let mut stream_name = params.stream_name.to_string(); - let schema = infra::schema::get(¶ms.org_id, &stream_name, params.stream_type) + let schema = infra::schema::get_cache(¶ms.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(¶ms.org_id, &stream_name, params.stream_type) + infra::schema::get_cache(¶ms.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(); diff --git a/src/service/schema.rs b/src/service/schema.rs index 9e5df9518c..405acc8c44 100644 --- a/src/service/schema.rs +++ b/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, - hash_key: String, -} - -impl SchemaCache { - pub fn new(schema: Schema, fields_map: HashMap) -> 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 { - &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, record_val: Vec<&Map>, record_ts: i64, -) -> Result<(SchemaEvolution, Option)> { +) -> Result { 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) { let mut is_schema_changed = false; let mut field_datatype_delta: Vec = 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 = 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 = 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] diff --git a/src/service/search/datafusion/exec.rs b/src/service/search/datafusion/exec.rs index 815b9c2e4b..fdbf9ba6b5 100644 --- a/src/service/search/datafusion/exec.rs +++ b/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 . -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 = columns_to_index .iter() .filter_map(|name| batch.schema().index_of(name).ok()) diff --git a/src/service/traces/mod.rs b/src/service/traces/mod.rs index ddae7a12dd..9906f0b768 100644 --- a/src/service/traces/mod.rs +++ b/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, }, }; diff --git a/src/service/traces/otlp_http.rs b/src/service/traces/otlp_http.rs index 1d47d58f4f..fec8cc61a6 100644 --- a/src/service/traces/otlp_http.rs +++ b/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, }, };