extract schema and batches via schema_arrow, and cleanup redundant logics

Author: tonyalaribe

Cargo.lock

@@ -9,6 +9,7 @@ datafusion = "46.0.0"
 arrow = "54.2.0"
 uuid = { version = "1.13", features = ["v4", "serde"] }
 serde = { version = "1", features = ["derive"] }
+serde_arrow = { version = "0.13.1", features = ["arrow-54"] }
 serde_json = "1.0.138"
 async-trait = "0.1.86"
 env_logger = "0.11.6"
@@ -30,7 +31,6 @@ anyhow = "1.0.95"
 tokio-util = "0.7.13"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 tracing = "0.1.41"
-prometheus = "0.13.4"
 dotenv = "0.15.0"
 task = "0.0.1"
 sqlparser = "0.55.0"

Cargo.toml

@@ -1,17 +1,12 @@
 use std::{collections::HashMap, sync::Arc};
 
 use anyhow::Result;
-use datafusion::{
-    arrow::{
-        array::{ArrayRef, StringArray, TimestampNanosecondArray},
-        datatypes::{DataType, Field, Schema, TimeUnit},
-        record_batch::RecordBatch,
-    },
-    execution::context::SessionContext,
-};
+use datafusion::execution::context::SessionContext;
 use deltalake::{storage::StorageOptions, DeltaOps, DeltaTable, DeltaTableBuilder};
+use serde_arrow::schema::{SchemaLike, TracingOptions};
 use tokio::sync::RwLock;
-use datafusion::prelude::DataFrame;
+
+use crate::persistent_queue::OtelLogsAndSpans;
 
 type ProjectConfig = (String, StorageOptions, Arc<RwLock<DeltaTable>>);
 
@@ -28,10 +23,7 @@ impl Database {
         let mut project_configs = HashMap::new();
 
         let default_options = StorageOptions::default();
-        let table = DeltaTableBuilder::from_uri(storage_uri)
-            .with_allow_http(true)
-            .with_storage_options(default_options.0.clone())
-            .build()?;
+        let table = DeltaTableBuilder::from_uri(storage_uri).with_allow_http(true).with_storage_options(default_options.0.clone()).build()?;
         ctx.register_table("otel_logs_and_spans", Arc::new(table.clone()))?;
         project_configs.insert("default".to_string(), (storage_uri.to_string(), default_options, Arc::new(RwLock::new(table))));
 
@@ -46,75 +38,30 @@ impl Database {
         &self.ctx
     }
 
-    pub async fn query(&self, project_id: &str, sql: &str) -> Result<DataFrame> {
-        let configs = self.project_configs.read().await;
-        if !configs.contains_key(project_id) {
-            return Err(anyhow::anyhow!("Project ID '{}' not found", project_id));
-        }
-        let adjusted_sql = sql.replace("otel_logs_and_spans", &format!("otel_logs_and_spans_{}", project_id));
-        self.ctx
-            .sql(&adjusted_sql)
-            .await
-            .map_err(|e| anyhow::anyhow!("SQL query failed for project '{}': {:?}", project_id, e))
-    }
-
-    fn event_schema() -> Schema {
-        Schema::new(vec![
-            Field::new("traceId", DataType::Utf8, false),
-            Field::new("spanId", DataType::Utf8, false),
-            Field::new("startTimeUnixNano", DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("UTC"))), false),
-            Field::new("endTimeUnixNano", DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("UTC"))), true),
-            Field::new("name", DataType::Utf8, false),
-        ])
-    }
-
-    pub async fn insert_record(
-        &self,
-        project_id: &str,
-        record: &crate::persistent_queue::IngestRecord,
-    ) -> Result<()> {
+    pub async fn insert_records(&self, project_id: &str, records: &Vec<crate::persistent_queue::OtelLogsAndSpans>) -> Result<()> {
         let (_conn_str, _options, table_ref) = {
             let configs = self.project_configs.read().await;
             configs.get(project_id).ok_or_else(|| anyhow::anyhow!("Project ID '{}' not found", project_id))?.clone()
         };
 
-        let arrays: Vec<ArrayRef> = vec![
-            Arc::new(StringArray::from(vec![Some(&record.context___trace_id[..])])),
-            Arc::new(StringArray::from(vec![Some(&record.context___span_id[..])])),
-            Arc::new(TimestampNanosecondArray::from(vec![record.timestamp])),
-            Arc::new(TimestampNanosecondArray::from(vec![record.start_time])),
-            Arc::new(TimestampNanosecondArray::from(vec![record.end_time.unwrap_or(0)])),
-            Arc::new(StringArray::from(vec![Some(&record.name[..])])),
-        ];
-
-        let schema = Self::event_schema();
-        let batch = RecordBatch::try_new(Arc::new(schema), arrays)?;
+        let fields = Vec::<arrow_schema::FieldRef>::from_type::<OtelLogsAndSpans>(TracingOptions::default())?;
+        let batch = serde_arrow::to_record_batch(&fields, &records)?;
 
         let mut table = table_ref.write().await;
         let ops = DeltaOps(table.clone());
         *table = ops.write(vec![batch]).await?;
         Ok(())
     }
 
-    pub async fn register_project(
-        &self,
-        project_id: &str,
-        bucket: &str,
-        access_key: &str,
-        secret_key: &str,
-        endpoint: &str,
-    ) -> Result<()> {
+    pub async fn register_project(&self, project_id: &str, bucket: &str, access_key: &str, secret_key: &str, endpoint: &str) -> Result<()> {
         let conn_str = format!("s3://{}/otel_logs_and_spans_{}", bucket, project_id);
         let mut storage_options = StorageOptions::default();
         storage_options.0.insert("AWS_ACCESS_KEY_ID".to_string(), access_key.to_string());
         storage_options.0.insert("AWS_SECRET_ACCESS_KEY".to_string(), secret_key.to_string());
         storage_options.0.insert("AWS_ENDPOINT".to_string(), endpoint.to_string());
         storage_options.0.insert("AWS_ALLOW_HTTP".to_string(), "true".to_string());
 
-        let table = DeltaTableBuilder::from_uri(&conn_str)
-            .with_storage_options(storage_options.0.clone())
-            .with_allow_http(true)
-            .build()?;
+        let table = DeltaTableBuilder::from_uri(&conn_str).with_storage_options(storage_options.0.clone()).with_allow_http(true).build()?;
 
         self.ctx.register_table(&format!("otel_logs_and_spans_{}", project_id), Arc::new(table.clone()))?;