chore: update datafusion to 34 (#2222)

Cargo.toml

@@ -76,10 +76,8 @@ clap = { version = "4.1", default-features = false, features = [
 cloudevents-sdk = { version = "0.7.0", features = ["actix"] }
 csv = "1.2.1"
 dashmap = { version = "5.4", features = ["serde"] }
-datafusion = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b648d4e22e82989c65523e62312e1995a1543888", version = "33", features = [
-  "simd",
-] }
-datafusion-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "b648d4e22e82989c65523e62312e1995a1543888", version = "33" }
+datafusion = { version = "34", features = ["simd"] }
+datafusion-expr = { version = "34" }
 arrow = { version = "49", features = ["simd", "ipc_compression"] }
 arrow-schema = { version = "49", features = ["serde"] }
 parquet = { version = "49", features = ["arrow", "async"] }
@@ -141,7 +139,7 @@ simd-json = "0.13"
 sha256 = "1.4.0"
 sled = "0.34"
 snap = "1"
-sqlparser = { version = "0.39", features = ["serde", "visitor"] }
+sqlparser = { version = "0.40", features = ["serde", "visitor"] }
 sqlx = { version = "0.7", features = [
   "runtime-tokio-rustls",
   "postgres",
@@ -168,7 +166,7 @@ utoipa = { version = "4", features = ["actix_extras", "openapi_extensions"] }
 utoipa-swagger-ui = { version = "4", features = ["actix-web"] }
 uuid = { version = "1.2", features = ["v4", "fast-rng", "macro-diagnostics"] }
 vector-enrichment = { package = "enrichment", git = "https://github.com/openobserve/vector", rev = "66667dd291482a440c5eb2032ef3cbfb7377b53b" }
-vrl = { version = "0.8.1", features = ["value", "compiler","test"] }
+vrl = { version = "0.8.1", features = ["value", "compiler", "test"] }
 walkdir = "2"
 xxhash-rust = { version = "0.8", features = ["xxh3"] }
 zstd = "0.13"

src/cli/basic/cli.rs

@@ -13,13 +13,19 @@
 // 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 crate::cli::data::{Context, export, import};
-use crate::cli::data::cli::args;
-use crate::cli::data::cli::Cli;
-use crate::common::{infra, meta, migration};
-use crate::common::infra::config::{CONFIG, USERS};
-use crate::common::utils::file::set_permission;
-use crate::service::{compact, db, file_list, users};
+use crate::{
+    cli::data::{
+        cli::{args, Cli},
+        export, import, Context,
+    },
+    common::{
+        infra,
+        infra::config::{CONFIG, USERS},
+        meta, migration,
+        utils::file::set_permission,
+    },
+    service::{compact, db, file_list, users},
+};
 
 pub async fn cli() -> Result<bool, anyhow::Error> {
     let app = clap::Command::new("openobserve")
@@ -120,7 +126,7 @@ pub async fn cli() -> Result<bool, anyhow::Error> {
                             is_ldap: false,
                         },
                     )
-                        .await?;
+                    .await?;
                 }
                 "user" => {
                     db::user::reset().await?;
@@ -223,4 +229,4 @@ pub async fn cli() -> Result<bool, anyhow::Error> {
 
     println!("command {name} execute succeeded");
     Ok(true)
-}
+}

src/cli/basic/mod.rs

@@ -13,4 +13,4 @@
 // 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/>.
 
-pub mod cli;
+pub mod cli;

src/cli/data/cli.rs

@@ -29,7 +29,7 @@ pub struct Cli {
     pub end_time: i64,
 }
 
-pub fn args() -> impl IntoIterator<Item=impl Into<clap::Arg>> {
+pub fn args() -> impl IntoIterator<Item = impl Into<clap::Arg>> {
     [
         clap::Arg::new("stream").long("stream").alias("c").default_value("stream")
             .help("for stream data, maybe we can export stream, users, metadata, but for now we can only support stream"),
@@ -87,18 +87,16 @@ fn handle_args(app: clap::ArgMatches) -> Cli {
 
     let start = match app.get_one::<String>("start_time") {
         Some(time) => time.parse::<i64>().unwrap(),
-        None => {
-            Local::now()
-                .date_naive()
-                .and_hms_milli_opt(0, 0, 0, 0)
-                .unwrap()
-                .timestamp_micros()
-        }
+        None => Local::now()
+            .date_naive()
+            .and_hms_milli_opt(0, 0, 0, 0)
+            .unwrap()
+            .timestamp_micros(),
     };
 
     let end = match app.get_one::<String>("end_time") {
         Some(time) => time.parse::<i64>().unwrap(),
-        None => { (Local::now() - Duration::hours(1)).timestamp_micros() }
+        None => (Local::now() - Duration::hours(1)).timestamp_micros(),
     };
 
     Cli {
@@ -112,4 +110,4 @@ fn handle_args(app: clap::ArgMatches) -> Cli {
         start_time: start,
         end_time: end,
     }
-}
+}

src/cli/data/export.rs

@@ -13,21 +13,24 @@
 // 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 std::fs::{self};
-use std::fs::File;
-use std::path::Path;
+use std::{
+    collections::HashMap,
+    fs::{self, File},
+    path::Path,
+};
 
 use actix_web::web::Query;
 use ahash::AHashMap;
 
-use crate::cli::data::cli::Cli;
-use crate::cli::data::Context;
-use crate::common::infra::config::CONFIG;
-use crate::common::meta::search::Request;
-use crate::common::meta::StreamType;
-use crate::common::utils::http::get_stream_type_from_request;
-use crate::service::search as SearchService;
+use crate::{
+    cli::data::{cli::Cli, Context},
+    common::{
+        infra::config::CONFIG,
+        meta::{search::Request, StreamType},
+        utils::http::get_stream_type_from_request,
+    },
+    service::search as SearchService,
+};
 
 pub struct Export {}
 
@@ -72,7 +75,11 @@ impl Context for Export {
                     }
                     let path = Path::new(c.data.as_str());
                     fs::create_dir_all(path)?;
-                    let file = File::create(path.join(format!("{}.{}", chrono::Local::now().timestamp_micros(), c.file_type)))?;
+                    let file = File::create(path.join(format!(
+                        "{}.{}",
+                        chrono::Local::now().timestamp_micros(),
+                        c.file_type
+                    )))?;
                     serde_json::to_writer_pretty(file, &res.hits)?;
                     Ok(true)
                 }
@@ -83,4 +90,4 @@ impl Context for Export {
             }
         })
     }
-}
+}

src/cli/data/import.rs

@@ -17,10 +17,11 @@ use std::fs;
 
 use bytes::Bytes;
 
-use crate::cli::data::cli::Cli;
-use crate::cli::data::Context;
-use crate::common::meta::ingestion::IngestionRequest;
-use crate::service::logs;
+use crate::{
+    cli::data::{cli::Cli, Context},
+    common::meta::ingestion::IngestionRequest,
+    service::logs,
+};
 
 pub struct Import {}
 
@@ -38,15 +39,22 @@ fn read_files_in_directory(c: Cli, dir_path: &str) -> Result<bool, anyhow::Error
             let path = entry.path();
             if path.is_file() {
                 let content = fs::read(&path)?;
-                if let Err(e) = logs::ingest::ingest(&c.org, &c.stream_name,
-                                                     IngestionRequest::JSON(&Bytes::from(content)), 0).await {
+                if let Err(e) = logs::ingest::ingest(
+                    &c.org,
+                    &c.stream_name,
+                    IngestionRequest::JSON(&Bytes::from(content)),
+                    0,
+                )
+                .await
+                {
                     eprintln!("insert data fail {:?}: {:?}", path, e);
                     return Ok(false);
                 }
-            } else if path.is_dir() && !read_files_in_directory(c.clone(), &path.to_string_lossy())? {
+            } else if path.is_dir() && !read_files_in_directory(c.clone(), &path.to_string_lossy())?
+            {
                 return Ok(false);
             }
         }
         Ok(true)
     })
-}
+}

src/cli/data/mod.rs

@@ -13,58 +13,69 @@
 // 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 crate::cli::data::cli::Cli;
-
 pub mod cli;
 pub mod export;
 pub mod import;
 
 pub trait Context {
-    fn operator(c: Cli) -> Result<bool, anyhow::Error>;
+    fn operator(c: cli::Cli) -> Result<bool, anyhow::Error>;
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::cli::data::cli::Cli;
-    use crate::cli::data::Context;
-    use crate::cli::data::export::Export;
-    use crate::cli::data::import::Import;
+    use super::*;
 
+    #[test]
     fn test_export_operator() {
         let args = vec![
             "openobserve",
-            "--c", "stream",
-            "--o", "default",
-            "--st", "logs",
-            "--s", "default",
-            "--t", "json",
-            "--f", "day",
-            "--d", "./json",
+            "--c",
+            "stream",
+            "--o",
+            "default",
+            "--st",
+            "logs",
+            "--s",
+            "default",
+            "--t",
+            "json",
+            "--f",
+            "day",
+            "--d",
+            "./json",
         ];
 
-        let cli = Cli::args(args);
+        let cli = cli::Cli::args(args);
 
-        if let Err(err) = Export::operator(cli.clone()) {
+        if let Err(err) = export::Export::operator(cli.clone()) {
             println!("Error: {}", err);
         }
     }
 
+    #[test]
     fn test_import_operator() {
         let args = vec![
             "openobserve",
-            "--c", "stream",
-            "--o", "default",
-            "--st", "logs",
-            "--s", "default",
-            "--t", "json",
-            "--f", "day",
-            "--d", "./json",
+            "--c",
+            "stream",
+            "--o",
+            "default",
+            "--st",
+            "logs",
+            "--s",
+            "default",
+            "--t",
+            "json",
+            "--f",
+            "day",
+            "--d",
+            "./json",
         ];
 
-        let cli = Cli::args(args);
+        let cli = cli::Cli::args(args);
 
-        if let Err(err) = Import::operator(cli.clone()) {
+        if let Err(err) = import::Import::operator(cli.clone()) {
             println!("Error: {}", err);
         }
     }
-}
+}

src/cli/mod.rs

@@ -15,4 +15,3 @@
 
 pub mod basic;
 pub mod data;
-

src/common/meta/http.rs

@@ -113,7 +113,7 @@ impl HttpResponse {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use actix_web::http;
 
     use super::*;

src/common/meta/search.rs

@@ -247,7 +247,7 @@ impl Response {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use super::*;
 
     #[test]

src/common/meta/stream.rs

@@ -392,7 +392,7 @@ pub struct StreamDeleteFields {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use super::*;
 
     #[test]

src/common/utils/hasher.rs

@@ -55,7 +55,7 @@ pub fn get_fields_key_xxh3(fields: &[Field]) -> String {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
 
     use arrow_schema::{DataType, Field};
 

src/handler/grpc/mod.rs

@@ -204,7 +204,7 @@ impl From<Vec<json::Value>> for cluster_rpc::UsageData {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use std::collections::HashMap;
 
     use super::*;

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

@@ -24,9 +24,8 @@ use crate::{
         meta::{self, http::HttpResponse as MetaHttpResponse},
         utils::time::{parse_milliseconds, parse_str_to_timestamp_micros},
     },
-    service::{metrics, promql},
+    service::{metrics, promql, promql::MetricsQueryRequest},
 };
-use crate::service::promql::MetricsQueryRequest;
 
 /// prometheus remote-write endpoint for metrics
 #[utoipa::path(
@@ -622,9 +621,9 @@ fn validate_metadata_params(
             Ok(parser::Expr::VectorSelector(sel)) => {
                 let err = if sel.name.is_none()
                     && sel
-                    .matchers
-                    .find_matcher_value(meta::prom::NAME_LABEL)
-                    .is_none()
+                        .matchers
+                        .find_matcher_value(meta::prom::NAME_LABEL)
+                        .is_none()
                 {
                     Some("match[] argument must start with a metric name, e.g. `match[]=up`")
                 } else if sel.offset.is_some() {
@@ -727,7 +726,11 @@ fn search_timeout(timeout: Option<String>) -> i64 {
     }
 }
 
-async fn search(org_id: &str, timeout: i64, req: &MetricsQueryRequest) -> Result<HttpResponse, Error> {
+async fn search(
+    org_id: &str,
+    timeout: i64,
+    req: &MetricsQueryRequest,
+) -> Result<HttpResponse, Error> {
     match promql::search::search(org_id, req, timeout).await {
         Ok(data) => Ok(HttpResponse::Ok().json(promql::QueryResponse {
             status: promql::Status::Success,
@@ -751,4 +754,4 @@ async fn search(org_id: &str, timeout: i64, req: &MetricsQueryRequest) -> Result
             }))
         }
     }
-}
+}

src/main.rs

@@ -19,36 +19,18 @@ use std::{
     net::SocketAddr,
     str::FromStr,
     sync::{
-        Arc,
         atomic::{AtomicU16, Ordering},
+        Arc,
     },
     time::Duration,
 };
 
-use actix_web::{App, http::KeepAlive, HttpServer, middleware, web};
+use actix_web::{http::KeepAlive, middleware, web, App, HttpServer};
 use actix_web_opentelemetry::RequestTracing;
 use chrono::Local;
 use log::LevelFilter;
-use opentelemetry::{
-    KeyValue,
-    sdk::{propagation::TraceContextPropagator, Resource, trace as sdktrace},
-};
-use opentelemetry_otlp::WithExportConfig;
-use opentelemetry_proto::tonic::collector::{
-    logs::v1::logs_service_server::LogsServiceServer,
-    metrics::v1::metrics_service_server::MetricsServiceServer,
-    trace::v1::trace_service_server::TraceServiceServer,
-};
-#[cfg(feature = "profiling")]
-use pyroscope::PyroscopeAgent;
-#[cfg(feature = "profiling")]
-use pyroscope_pprofrs::{pprof_backend, PprofConfig};
-use tokio::sync::oneshot;
-use tonic::codec::CompressionEncoding;
-use tracing_subscriber::{prelude::*, Registry};
-use uaparser::UserAgentParser;
-
 use openobserve::{
+    cli::basic::cli,
     common::{
         infra::{
             self, cluster,
@@ -80,7 +62,24 @@ use openobserve::{
     job, router,
     service::{db, distinct_values},
 };
-use openobserve::cli::basic::cli;
+use opentelemetry::{
+    sdk::{propagation::TraceContextPropagator, trace as sdktrace, Resource},
+    KeyValue,
+};
+use opentelemetry_otlp::WithExportConfig;
+use opentelemetry_proto::tonic::collector::{
+    logs::v1::logs_service_server::LogsServiceServer,
+    metrics::v1::metrics_service_server::MetricsServiceServer,
+    trace::v1::trace_service_server::TraceServiceServer,
+};
+#[cfg(feature = "profiling")]
+use pyroscope::PyroscopeAgent;
+#[cfg(feature = "profiling")]
+use pyroscope_pprofrs::{pprof_backend, PprofConfig};
+use tokio::sync::oneshot;
+use tonic::codec::CompressionEncoding;
+use tracing_subscriber::{prelude::*, Registry};
+use uaparser::UserAgentParser;
 
 #[cfg(feature = "mimalloc")]
 #[global_allocator]
@@ -91,23 +90,23 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
 static USER_AGENT_REGEX_FILE: &[u8] = include_bytes!(concat!(
-env!("CARGO_MANIFEST_DIR"),
-"/ua_regex/regexes.yaml"
+    env!("CARGO_MANIFEST_DIR"),
+    "/ua_regex/regexes.yaml"
 ));
 
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
     #[cfg(feature = "profiling")]
-        let agent = PyroscopeAgent::builder(
+    let agent = PyroscopeAgent::builder(
         &CONFIG.profiling.pyroscope_server_url,
         &CONFIG.profiling.pyroscope_project_name,
     )
-        .tags([("Host", "Rust")].to_vec())
-        .backend(pprof_backend(PprofConfig::new().sample_rate(100)))
-        .build()
-        .expect("Failed to setup pyroscope agent");
+    .tags([("Host", "Rust")].to_vec())
+    .backend(pprof_backend(PprofConfig::new().sample_rate(100)))
+    .build()
+    .expect("Failed to setup pyroscope agent");
     #[cfg(feature = "profiling")]
-        let agent_running = agent.start().expect("Failed to start pyroscope agent");
+    let agent_running = agent.start().expect("Failed to start pyroscope agent");
 
     if cli::cli().await? {
         return Ok(());
@@ -229,7 +228,7 @@ async fn main() -> Result<(), anyhow::Error> {
     log::info!("server stopped");
 
     #[cfg(feature = "profiling")]
-        let agent_ready = agent_running.stop().unwrap();
+    let agent_ready = agent_running.stop().unwrap();
     #[cfg(feature = "profiling")]
     agent_ready.shutdown();
 
@@ -394,11 +393,11 @@ async fn init_http_server() -> Result<(), anyhow::Error> {
             ))
             .wrap(RequestTracing::new())
     })
-        .keep_alive(KeepAlive::Timeout(Duration::from_secs(
-            CONFIG.limit.keep_alive,
-        )))
-        .client_request_timeout(Duration::from_secs(CONFIG.limit.request_timeout))
-        .bind(haddr)?;
+    .keep_alive(KeepAlive::Timeout(Duration::from_secs(
+        CONFIG.limit.keep_alive,
+    )))
+    .client_request_timeout(Duration::from_secs(CONFIG.limit.request_timeout))
+    .bind(haddr)?;
 
     server
         .workers(CONFIG.limit.http_worker_num)
@@ -438,4 +437,4 @@ fn enable_tracing() -> Result<(), anyhow::Error> {
         .with(tracing_opentelemetry::layer().with_tracer(tracer))
         .init();
     Ok(())
-}
+}

src/service/file_list.rs

@@ -358,7 +358,7 @@ async fn delete_parquet_file_s3(key: &str, file_list_only: bool) -> Result<(), a
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use super::*;
 
     #[actix_web::test]

src/service/functions.rs

@@ -329,7 +329,7 @@ fn remove_stream_fn_from_cache(key: &str, fn_name: &str) {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use super::*;
 
     #[actix_web::test]

src/service/logs/otlp_grpc.rs

@@ -509,7 +509,7 @@ pub async fn handle_grpc_request(
 }
 
 #[cfg(test)]
-pub mod test {
+mod tests {
 
     use opentelemetry_proto::tonic::{
         collector::logs::v1::ExportLogsServiceRequest,

src/service/logs/syslog.rs

@@ -305,7 +305,7 @@ fn message_to_value(message: Message<&str>) -> json::Value {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use std::net::{IpAddr, Ipv4Addr};
 
     use super::*;

src/service/schema.rs

@@ -854,7 +854,7 @@ pub fn filter_schema_null_fields(schema: &mut Schema) {
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use ahash::AHashMap;
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
 

src/service/search/datafusion/exec.rs

@@ -1214,7 +1214,7 @@ fn apply_query_fn(
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use arrow::array::{Int32Array, NullArray, StringArray};
     use arrow_schema::Field;
 

src/service/stream.rs

@@ -435,7 +435,7 @@ async fn _get_stream_stats(
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
 
     use super::*;

src/service/usage/ingestion_service.rs

@@ -95,7 +95,7 @@ pub async fn ingest(
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use rand::{rngs::StdRng, seq::SliceRandom, SeedableRng};
 
     #[test]