refactor: override_inferred_data_type

Author: Devdutt Shenoi

src/event/format/json.rs

@@ -22,17 +22,31 @@
 use anyhow::anyhow;
 use arrow_array::RecordBatch;
 use arrow_json::reader::{infer_json_schema_from_iterator, ReaderBuilder};
-use arrow_schema::{DataType, Field, Fields, Schema};
+use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit};
 use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc};
 use datafusion::arrow::util::bit_util::round_upto_multiple_of_64;
 use itertools::Itertools;
 use serde_json::Value;
 use std::{collections::HashMap, sync::Arc};
 use tracing::error;
 
-use super::EventFormat;
+use super::{update_field_type_in_schema, EventFormat};
 use crate::{metadata::SchemaVersion, storage::StreamType, utils::arrow::get_field};
 
+static TIME_FIELD_NAME_PARTS: [&str; 11] = [
+    "time",
+    "date",
+    "timestamp",
+    "created",
+    "received",
+    "ingested",
+    "collected",
+    "start",
+    "end",
+    "ts",
+    "dt",
+];
+
 pub struct Event {
     pub json: Value,
     pub p_timestamp: DateTime<Utc>,
@@ -87,12 +101,10 @@ impl EventFormat for Event {
                     .map_err(|err| {
                         anyhow!("Could not infer schema for this event due to err {:?}", err)
                     })?;
-                let new_infer_schema = super::update_field_type_in_schema(
+                let new_infer_schema = update_field_type_in_schema(
                     Arc::new(infer_schema),
                     Some(stream_schema),
                     time_partition,
-                    Some(&value_arr),
-                    schema_version,
                 );
                 infer_schema = Schema::new(new_infer_schema.fields().clone());
                 Schema::try_merge(vec![
@@ -222,6 +234,54 @@ fn extract_and_parse_time(
     Ok(parsed_time.naive_utc())
 }
 
+// From Schema v1 onwards, convert json fields with name containig "date"/"time" and having
+// a string value parseable into timestamp as timestamp type and all numbers as float64.
+pub fn override_inferred_data_type(
+    inferred_schema: Schema,
+    log_record: &Value,
+    schema_version: SchemaVersion,
+) -> Schema {
+    let Value::Object(map) = log_record else {
+        return inferred_schema;
+    };
+    let updated_fields = inferred_schema
+        .fields()
+        .iter()
+        .map(|field| {
+            let field_name = field.name().as_str();
+            match (schema_version, map.get(field.name())) {
+                // in V1 for new fields in json named "time"/"date" or such and having inferred
+                // type string, that can be parsed as timestamp, use the timestamp type.
+                // NOTE: support even more datetime string formats
+                (SchemaVersion::V1, Some(Value::String(s)))
+                    if TIME_FIELD_NAME_PARTS
+                        .iter()
+                        .any(|part| field_name.to_lowercase().contains(part))
+                        && field.data_type() == &DataType::Utf8
+                        && (DateTime::parse_from_rfc3339(s).is_ok()
+                            || DateTime::parse_from_rfc2822(s).is_ok()) =>
+                {
+                    // Update the field's data type to Timestamp
+                    Field::new(
+                        field_name,
+                        DataType::Timestamp(TimeUnit::Millisecond, None),
+                        true,
+                    )
+                }
+                // in V1 for new fields in json with inferred type number, cast as float64.
+                (SchemaVersion::V1, Some(Value::Number(_))) if field.data_type().is_numeric() => {
+                    // Update the field's data type to Float64
+                    Field::new(field_name, DataType::Float64, true)
+                }
+                // Return the original field if no update is needed
+                _ => Field::new(field_name, field.data_type().clone(), true),
+            }
+        })
+        .collect_vec();
+
+    Schema::new(updated_fields)
+}
+
 // Returns arrow schema with the fields that are present in the request body
 // This schema is an input to convert the request body to arrow record batch
 fn derive_arrow_schema(

src/event/format/mod.rs

@@ -28,7 +28,6 @@ use arrow_array::RecordBatch;
 use arrow_schema::{DataType, Field, Schema, TimeUnit};
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
 
 use crate::{
     metadata::SchemaVersion,
@@ -40,19 +39,6 @@ use super::{Event, DEFAULT_TIMESTAMP_KEY};
 
 pub mod json;
 
-static TIME_FIELD_NAME_PARTS: [&str; 11] = [
-    "time",
-    "date",
-    "timestamp",
-    "created",
-    "received",
-    "ingested",
-    "collected",
-    "start",
-    "end",
-    "ts",
-    "dt",
-];
 type EventSchema = Vec<Arc<Field>>;
 
 /// Source of the logs, used to perform special processing for certain sources
@@ -167,8 +153,7 @@ pub trait EventFormat: Sized {
         if !Self::is_schema_matching(new_schema.clone(), storage_schema, static_schema_flag) {
             return Err(anyhow!("Schema mismatch"));
         }
-        new_schema =
-            update_field_type_in_schema(new_schema, None, time_partition, None, schema_version);
+        new_schema = update_field_type_in_schema(new_schema, None, time_partition);
 
         let rb = Self::decode(data, new_schema.clone())?;
 
@@ -269,8 +254,6 @@ pub fn update_field_type_in_schema(
     inferred_schema: Arc<Schema>,
     existing_schema: Option<&HashMap<String, Arc<Field>>>,
     time_partition: Option<&String>,
-    log_records: Option<&Vec<Value>>,
-    schema_version: SchemaVersion,
 ) -> Arc<Schema> {
     let mut updated_schema = inferred_schema.clone();
     let existing_field_names = get_existing_field_names(inferred_schema.clone(), existing_schema);
@@ -280,13 +263,6 @@ pub fn update_field_type_in_schema(
         updated_schema = override_existing_timestamp_fields(existing_schema, updated_schema);
     }
 
-    if let Some(log_records) = log_records {
-        for log_record in log_records {
-            updated_schema =
-                override_data_type(updated_schema.clone(), log_record.clone(), schema_version);
-        }
-    }
-
     let Some(time_partition) = time_partition else {
         return updated_schema;
     };
@@ -309,51 +285,3 @@ pub fn update_field_type_in_schema(
         .collect();
     Arc::new(Schema::new(new_schema))
 }
-
-// From Schema v1 onwards, convert json fields with name containig "date"/"time" and having
-// a string value parseable into timestamp as timestamp type and all numbers as float64.
-pub fn override_data_type(
-    inferred_schema: Arc<Schema>,
-    log_record: Value,
-    schema_version: SchemaVersion,
-) -> Arc<Schema> {
-    let Value::Object(map) = log_record else {
-        return inferred_schema;
-    };
-    let updated_schema: Vec<Field> = inferred_schema
-        .fields()
-        .iter()
-        .map(|field| {
-            let field_name = field.name().as_str();
-            match (schema_version, map.get(field.name())) {
-                // in V1 for new fields in json named "time"/"date" or such and having inferred
-                // type string, that can be parsed as timestamp, use the timestamp type.
-                // NOTE: support even more datetime string formats
-                (SchemaVersion::V1, Some(Value::String(s)))
-                    if TIME_FIELD_NAME_PARTS
-                        .iter()
-                        .any(|part| field_name.to_lowercase().contains(part))
-                        && field.data_type() == &DataType::Utf8
-                        && (DateTime::parse_from_rfc3339(s).is_ok()
-                            || DateTime::parse_from_rfc2822(s).is_ok()) =>
-                {
-                    // Update the field's data type to Timestamp
-                    Field::new(
-                        field_name,
-                        DataType::Timestamp(TimeUnit::Millisecond, None),
-                        true,
-                    )
-                }
-                // in V1 for new fields in json with inferred type number, cast as float64.
-                (SchemaVersion::V1, Some(Value::Number(_))) if field.data_type().is_numeric() => {
-                    // Update the field's data type to Float64
-                    Field::new(field_name, DataType::Float64, true)
-                }
-                // Return the original field if no update is needed
-                _ => Field::new(field_name, field.data_type().clone(), true),
-            }
-        })
-        .collect();
-
-    Arc::new(Schema::new(updated_schema))
-}

src/handlers/http/logstream.rs

@@ -16,34 +16,35 @@
  *
  */
 
-use self::error::StreamError;
-use super::cluster::utils::{IngestionStats, QueriedStats, StorageStats};
-use super::query::update_schema_when_distributed;
-use crate::event::format::override_data_type;
-use crate::hottier::{HotTierManager, StreamHotTier, CURRENT_HOT_TIER_VERSION};
-use crate::metadata::SchemaVersion;
-use crate::metrics::{EVENTS_INGESTED_DATE, EVENTS_INGESTED_SIZE_DATE, EVENTS_STORAGE_SIZE_DATE};
-use crate::parseable::{StreamNotFound, PARSEABLE};
-use crate::rbac::role::Action;
-use crate::rbac::Users;
-use crate::stats::{event_labels_date, storage_size_labels_date, Stats};
-use crate::storage::retention::Retention;
-use crate::storage::{StreamInfo, StreamType};
-use crate::utils::actix::extract_session_key_from_req;
-use crate::{stats, validator, LOCK_EXPECT};
+use std::fs::remove_dir_all;
 
 use actix_web::http::StatusCode;
 use actix_web::web::{Json, Path};
 use actix_web::{web, HttpRequest, Responder};
 use arrow_json::reader::infer_json_schema_from_iterator;
 use bytes::Bytes;
 use chrono::Utc;
+use error::StreamError;
 use itertools::Itertools;
 use serde_json::{json, Value};
-use std::fs;
-use std::sync::Arc;
 use tracing::warn;
 
+use crate::event::format::json::override_inferred_data_type;
+use crate::hottier::{HotTierManager, StreamHotTier, CURRENT_HOT_TIER_VERSION};
+use crate::metadata::SchemaVersion;
+use crate::metrics::{EVENTS_INGESTED_DATE, EVENTS_INGESTED_SIZE_DATE, EVENTS_STORAGE_SIZE_DATE};
+use crate::parseable::{StreamNotFound, PARSEABLE};
+use crate::rbac::role::Action;
+use crate::rbac::Users;
+use crate::stats::{self, event_labels_date, storage_size_labels_date, Stats};
+use crate::storage::retention::Retention;
+use crate::storage::{StreamInfo, StreamType};
+use crate::utils::actix::extract_session_key_from_req;
+use crate::{validator, LOCK_EXPECT};
+
+use super::cluster::utils::{IngestionStats, QueriedStats, StorageStats};
+use super::query::update_schema_when_distributed;
+
 pub async fn delete(stream_name: Path<String>) -> Result<impl Responder, StreamError> {
     let stream_name = stream_name.into_inner();
     // Error out if stream doesn't exist in memory, or in the case of query node, in storage as well
@@ -57,7 +58,7 @@ pub async fn delete(stream_name: Path<String>) -> Result<impl Responder, StreamE
     objectstore.delete_stream(&stream_name).await?;
     // Delete from staging
     let stream_dir = PARSEABLE.get_or_create_stream(&stream_name);
-    if fs::remove_dir_all(&stream_dir.data_path).is_err() {
+    if remove_dir_all(&stream_dir.data_path).is_err() {
         warn!(
             "failed to delete local data for stream {}. Clean {} manually",
             stream_name,
@@ -113,10 +114,11 @@ pub async fn detect_schema(Json(json): Json<Value>) -> Result<impl Responder, St
         }
     };
 
-    let mut schema = Arc::new(infer_json_schema_from_iterator(log_records.iter().map(Ok)).unwrap());
-    for log_record in log_records {
-        schema = override_data_type(schema, log_record, SchemaVersion::V1);
+    let mut schema = infer_json_schema_from_iterator(log_records.iter().map(Ok)).unwrap();
+    for log_record in log_records.iter() {
+        schema = override_inferred_data_type(schema, log_record, SchemaVersion::V1);
     }
+
     Ok((web::Json(schema), StatusCode::OK))
 }