fix: prefix partition key with inverted index (#4962)

src/handler/grpc/flight/mod.rs

@@ -29,7 +29,7 @@ use arrow_flight::{
     HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
 };
 use arrow_schema::Schema;
-use config::meta::search::ScanStats;
+use config::{meta::search::ScanStats, metrics};
 use datafusion::{
     common::{DataFusionError, Result},
     execution::SendableRecordBatchStream,
@@ -284,6 +284,16 @@ impl Drop for FlightSenderStream {
             self.trace_id,
             end
         );
+
+        // metrics
+        let time = self.start.elapsed().as_secs_f64();
+        metrics::GRPC_RESPONSE_TIME
+            .with_label_values(&["/search/flight/do_get", "200", "", "", ""])
+            .observe(time);
+        metrics::GRPC_INCOMING_REQUESTS
+            .with_label_values(&["/search/flight/do_get", "200", "", "", ""])
+            .inc();
+
         if let Some(defer) = self.defer.take() {
             drop(defer);
         } else {

src/handler/grpc/request/ingest.rs

@@ -14,7 +14,7 @@
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 use actix_web::http::StatusCode;
-use config::utils::json;
+use config::{metrics, utils::json};
 use proto::cluster_rpc::{
     ingest_server::Ingest, IngestionRequest, IngestionResponse, IngestionType, StreamType,
 };
@@ -31,6 +31,7 @@ impl Ingest for Ingester {
         &self,
         request: Request<IngestionRequest>,
     ) -> Result<Response<IngestionResponse>, Status> {
+        let start = std::time::Instant::now();
         // let metadata = request.metadata().clone();
         let req = request.into_inner();
         let org_id = req.org_id;
@@ -131,6 +132,16 @@ impl Ingest for Ingester {
                 message: err.to_string(),
             },
         };
+
+        // metrics
+        let time = start.elapsed().as_secs_f64();
+        metrics::GRPC_RESPONSE_TIME
+            .with_label_values(&["/ingest/inner", "200", "", "", ""])
+            .observe(time);
+        metrics::GRPC_INCOMING_REQUESTS
+            .with_label_values(&["/ingest/inner", "200", "", "", ""])
+            .inc();
+
         Ok(Response::new(reply))
     }
 }

src/handler/grpc/request/logs.rs

@@ -13,6 +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 config::metrics;
 use opentelemetry_proto::tonic::collector::logs::v1::{
     logs_service_server::LogsService, ExportLogsServiceRequest, ExportLogsServiceResponse,
 };
@@ -27,7 +28,9 @@ impl LogsService for LogsServer {
         &self,
         request: tonic::Request<ExportLogsServiceRequest>,
     ) -> Result<tonic::Response<ExportLogsServiceResponse>, tonic::Status> {
+        let start = std::time::Instant::now();
         let cfg = config::get_config();
+
         let metadata = request.metadata().clone();
         let msg = format!(
             "Please specify organization id with header key '{}' ",
@@ -64,9 +67,20 @@ impl LogsService for LogsServer {
         )
         .await
         {
-            Ok(_) => Ok(Response::new(ExportLogsServiceResponse {
-                partial_success: None,
-            })),
+            Ok(_) => {
+                // metrics
+                let time = start.elapsed().as_secs_f64();
+                metrics::GRPC_RESPONSE_TIME
+                    .with_label_values(&["/otlp/v1/logs", "200", "", "", ""])
+                    .observe(time);
+                metrics::GRPC_INCOMING_REQUESTS
+                    .with_label_values(&["/otlp/v1/logs", "200", "", "", ""])
+                    .inc();
+
+                Ok(Response::new(ExportLogsServiceResponse {
+                    partial_success: None,
+                }))
+            }
             Err(e) => Err(Status::internal(e.to_string())),
         }
     }

src/handler/grpc/request/metrics/ingester.rs

@@ -13,6 +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 config::metrics;
 use opentelemetry_proto::tonic::collector::metrics::v1::{
     metrics_service_server::MetricsService, ExportMetricsServiceRequest,
     ExportMetricsServiceResponse,
@@ -28,7 +29,9 @@ impl MetricsService for MetricsIngester {
         &self,
         request: tonic::Request<ExportMetricsServiceRequest>,
     ) -> Result<tonic::Response<ExportMetricsServiceResponse>, tonic::Status> {
+        let start = std::time::Instant::now();
         let cfg = config::get_config();
+
         let metadata = request.metadata().clone();
         let msg = format!(
             "Please specify organization id with header key '{}' ",
@@ -51,6 +54,14 @@ impl MetricsService for MetricsIngester {
         )
         .await;
         if resp.is_ok() {
+            // metrics
+            let time = start.elapsed().as_secs_f64();
+            metrics::GRPC_RESPONSE_TIME
+                .with_label_values(&["/otlp/v1/metrics", "200", "", "", ""])
+                .observe(time);
+            metrics::GRPC_INCOMING_REQUESTS
+                .with_label_values(&["/otlp/v1/metrics", "200", "", "", ""])
+                .inc();
             return Ok(Response::new(ExportMetricsServiceResponse {
                 partial_success: None,
             }));

src/handler/grpc/request/traces.rs

@@ -13,6 +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 config::metrics;
 use opentelemetry_proto::tonic::collector::trace::v1::{
     trace_service_server::TraceService, ExportTraceServiceRequest, ExportTraceServiceResponse,
 };
@@ -29,7 +30,9 @@ impl TraceService for TraceServer {
         &self,
         request: tonic::Request<ExportTraceServiceRequest>,
     ) -> Result<tonic::Response<ExportTraceServiceResponse>, tonic::Status> {
+        let start = std::time::Instant::now();
         let cfg = config::get_config();
+
         let metadata = request.metadata().clone();
         let msg = format!(
             "Please specify organization id with header key '{}' ",
@@ -59,6 +62,14 @@ impl TraceService for TraceServer {
         )
         .await;
         if resp.is_ok() {
+            // metrics
+            let time = start.elapsed().as_secs_f64();
+            metrics::GRPC_RESPONSE_TIME
+                .with_label_values(&["/otlp/v1/traces", "200", "", "", ""])
+                .observe(time);
+            metrics::GRPC_INCOMING_REQUESTS
+                .with_label_values(&["/otlp/v1/traces", "200", "", "", ""])
+                .inc();
             return Ok(Response::new(ExportTraceServiceResponse {
                 partial_success: None,
             }));

src/infra/src/db/mysql.rs

@@ -293,7 +293,7 @@ impl super::Db for MysqlDb {
         } else {
             DB_QUERY_NUMS.with_label_values(&["select", "meta"]).inc();
             match sqlx::query_as::<_,super::MetaRecord>(
-                r#"SELECT id, module, key1, key2, start_dt, value FROM meta WHERE module = ? AND key1 = ? AND key2 = ? ORDER BY id DESC;"#
+                r#"SELECT id, module, key1, key2, start_dt, value FROM meta WHERE module = ? AND key1 = ? AND key2 = ? ORDER BY start_dt DESC, id DESC;"#
             )
             .bind(&module)
             .bind(&key1)

src/infra/src/db/postgres.rs

@@ -247,7 +247,7 @@ impl super::Db for PostgresDb {
         } else {
             DB_QUERY_NUMS.with_label_values(&["select", "meta"]).inc();
             match sqlx::query_as::<_,super::MetaRecord>(
-                r#"SELECT id, module, key1, key2, start_dt, value FROM meta WHERE module = $1 AND key1 = $2 AND key2 = $3 ORDER BY id DESC;"#
+                r#"SELECT id, module, key1, key2, start_dt, value FROM meta WHERE module = $1 AND key1 = $2 AND key2 = $3 ORDER BY start_dt DESC, id DESC;"#
             )
             .bind(&module)
             .bind(&key1)

src/infra/src/db/sqlite.rs

@@ -349,7 +349,7 @@ impl super::Db for SqliteDb {
             }
         } else {
             match sqlx::query_as::<_,super::MetaRecord>(
-                r#"SELECT id, module, key1, key2, start_dt, value FROM meta WHERE module = $1 AND key1 = $2 AND key2 = $3 ORDER BY id DESC;"#
+                r#"SELECT id, module, key1, key2, start_dt, value FROM meta WHERE module = $1 AND key1 = $2 AND key2 = $3 ORDER BY start_dt DESC, id DESC;"#
             )
             .bind(&module)
             .bind(&key1)

src/job/files/idx.rs

@@ -22,8 +22,10 @@ use arrow::array::{
 use arrow_schema::{DataType, Schema};
 use config::{
     ider,
-    meta::stream::{FileMeta, StreamPartition, StreamType},
-    utils::parquet::new_parquet_writer,
+    meta::stream::{FileMeta, StreamPartition, StreamPartitionType, StreamType},
+    utils::{
+        parquet::new_parquet_writer, record_batch_ext::concat_batches, schema::format_partition_key,
+    },
     FILE_EXT_PARQUET,
 };
 use hashbrown::HashMap;
@@ -46,15 +48,19 @@ fn generate_index_file_name_from_compacted_file(
         file_columns[4], file_columns[5], file_columns[6], file_columns[7]
     );
     let file_name = ider::generate();
-    format!("files/{stream_key}/{file_date}/{prefix}/{file_name}{FILE_EXT_PARQUET}")
+    if prefix.is_empty() {
+        format!("files/{stream_key}/{file_date}/{file_name}{FILE_EXT_PARQUET}")
+    } else {
+        format!("files/{stream_key}/{file_date}/{prefix}/{file_name}{FILE_EXT_PARQUET}")
+    }
 }
 
 pub(crate) async fn write_parquet_index_to_disk(
     batches: Vec<arrow::record_batch::RecordBatch>,
     file_size: u64,
     org_id: &str,
-    stream_name: &str,
     stream_type: StreamType,
+    stream_name: &str,
     file_name: &str,
     caller: &str,
 ) -> Result<Vec<(String, FileMeta)>, anyhow::Error> {
@@ -69,18 +75,26 @@ pub(crate) async fn write_parquet_index_to_disk(
         batches.iter().map(|b| b.num_rows()).sum::<usize>()
     );
 
+    let stream_settings = infra::schema::get_settings(org_id, stream_name, stream_type)
+        .await
+        .unwrap_or_default();
+    let term_prefix_partition = stream_settings.partition_keys.iter().any(|partition| {
+        partition.field == "term"
+            && partition.types == StreamPartitionType::Prefix
+            && !partition.disabled
+    });
+
     // partition the record batches
-    let partitioned_batches = generate_prefixed_batches(schema.clone(), batches)?;
+    let partitioned_batches = if term_prefix_partition {
+        generate_prefixed_batches(schema.clone(), batches)?
+    } else {
+        let mut partitioned_batches = HashMap::new();
+        let batch = concat_batches(schema.clone(), batches)?;
+        partitioned_batches.insert("".to_string(), batch);
+        partitioned_batches
+    };
 
-    for (prefix, batch) in partitioned_batches.iter() {
-        log::debug!(
-            "write_parquet_index_to_disk: prefix: {}, row count: {}",
-            prefix,
-            batch.num_rows()
-        );
-    }
     let mut ret = Vec::new();
-
     for (prefix, batch) in partitioned_batches.into_iter() {
         // write metadata
         let mut file_meta = FileMeta {
@@ -140,7 +154,7 @@ pub(crate) async fn write_parquet_index_to_disk(
     Ok(ret)
 }
 
-/// Generate prefix batches from   record batches
+/// Generate prefix batches from record batches
 fn generate_prefixed_batches(
     schema: Arc<Schema>,
     batches: Vec<RecordBatch>,
@@ -174,7 +188,7 @@ fn generate_prefixed_batches(
                         let prefix = if term.is_empty() {
                             String::new()
                         } else {
-                            partition.get_partition_key(term)
+                            format_partition_key(&partition.get_partition_key(term))
                         };
                         let entry = partition_buf.entry(prefix).or_default();
                         let builder = entry
@@ -198,7 +212,7 @@ fn generate_prefixed_batches(
                         let prefix = if term.is_empty() {
                             String::new()
                         } else {
-                            partition.get_partition_key(term)
+                            format_partition_key(&partition.get_partition_key(term))
                         };
                         let entry = partition_buf.entry(prefix).or_default();
                         let builder = entry
@@ -219,7 +233,7 @@ fn generate_prefixed_batches(
                         let prefix = if term.is_empty() {
                             String::new()
                         } else {
-                            partition.get_partition_key(term)
+                            format_partition_key(&partition.get_partition_key(term))
                         };
                         let entry = partition_buf.entry(prefix).or_default();
                         let builder = entry
@@ -243,7 +257,7 @@ fn generate_prefixed_batches(
                         let prefix = if term.is_empty() {
                             String::new()
                         } else {
-                            partition.get_partition_key(term)
+                            format_partition_key(&partition.get_partition_key(term))
                         };
                         let entry = partition_buf.entry(prefix).or_default();
                         let builder = entry

src/job/files/parquet.rs

@@ -153,19 +153,19 @@ async fn scan_wal_files(
         }
     });
     let mut files_num = 0;
-    let mut last_time = start.elapsed().as_millis();
+    // let mut last_time = start.elapsed().as_millis();
     loop {
         match rx.recv().await {
             None => {
                 break;
             }
             Some(files) => {
-                log::debug!(
-                    "[INGESTER:JOB] scan files get batch: {}, took: {} ms",
-                    files.len(),
-                    start.elapsed().as_millis() - last_time
-                );
-                last_time = start.elapsed().as_millis();
+                // log::debug!(
+                //     "[INGESTER:JOB] scan files get batch: {}, took: {} ms",
+                //     files.len(),
+                //     start.elapsed().as_millis() - last_time
+                // );
+                // last_time = start.elapsed().as_millis();
                 files_num += files.len();
                 match prepare_files(files).await {
                     Err(e) => {
@@ -268,16 +268,21 @@ async fn prepare_files(
             continue;
         }
         let prefix = file_key[..file_key.rfind('/').unwrap()].to_string();
+        // remove thread_id from prefix
+        // eg: files/default/logs/olympics/0/2023/08/21/08/8b8a5451bbe1c44b/
+        let mut columns = prefix.split('/').collect::<Vec<&str>>();
+        columns.remove(4);
+        let prefix = columns.join("/");
         let partition = partition_files_with_size.entry(prefix).or_default();
         partition.push(FileKey::new(&file_key, parquet_meta, false));
         // mark the file as processing
         // log::debug!("Processing files created: {:?}", file_key);
         PROCESSING_FILES.write().await.insert(file_key);
     }
-    log::debug!(
-        "[INGESTER:JOB] move files get partitions: {}",
-        partition_files_with_size.len()
-    );
+    // log::debug!(
+    //     "[INGESTER:JOB] move files get partitions: {}",
+    //     partition_files_with_size.len()
+    // );
 
     Ok(partition_files_with_size)
 }
@@ -291,15 +296,15 @@ async fn move_files(
         return Ok(());
     }
 
-    let columns = prefix.splitn(9, '/').collect::<Vec<&str>>();
-    // eg: files/default/logs/olympics/0/2023/08/21/08/8b8a5451bbe1c44b/
-    // eg: files/default/traces/default/0/2023/09/04/05/default/service_name=ingester/
+    let columns = prefix.split('/').collect::<Vec<&str>>();
+    // removed thread_id from prefix, so there is no thread_id in the path
+    // eg: files/default/logs/olympics/2023/08/21/08/8b8a5451bbe1c44b/
+    // eg: files/default/traces/default/2023/09/04/05/default/service_name=ingester/
     // let _ = columns[0].to_string(); // files/
     let org_id = columns[1].to_string();
     let stream_type = StreamType::from(columns[2]);
     let stream_name = columns[3].to_string();
-    // let _thread_id = columns[4].to_string();
-    let prefix_date = format!("{}-{}-{}", columns[5], columns[6], columns[7]);
+    let prefix_date = format!("{}-{}-{}", columns[4], columns[5], columns[6]);
 
     // log::debug!("[INGESTER:JOB:{thread_id}] check deletion for partition: {}", prefix);
 
@@ -460,10 +465,10 @@ async fn move_files(
         }
     }
 
-    log::debug!(
-        "[INGESTER:JOB:{thread_id}] start merging for partition: {}",
-        prefix
-    );
+    // log::debug!(
+    //     "[INGESTER:JOB:{thread_id}] start merging for partition: {}",
+    //     prefix
+    // );
 
     // start merge files and upload to s3
     loop {
@@ -781,8 +786,8 @@ async fn merge_files(
                             inverted_idx_batch.clone(),
                             new_file_key.clone(),
                             &org_id,
-                            &stream_name,
                             stream_type,
+                            &stream_name,
                             &full_text_search_fields,
                             &index_fields,
                         )
@@ -819,8 +824,8 @@ pub(crate) async fn generate_index_on_ingester(
     inverted_idx_batch: RecordBatch,
     new_file_key: String,
     org_id: &str,
-    stream_name: &str,
     stream_type: StreamType,
+    stream_name: &str,
     full_text_search_fields: &[String],
     index_fields: &[String],
 ) -> Result<(), anyhow::Error> {
@@ -833,6 +838,7 @@ pub(crate) async fn generate_index_on_ingester(
     let record_batches = prepare_index_record_batches(
         inverted_idx_batch,
         org_id,
+        stream_type,
         stream_name,
         &new_file_key,
         full_text_search_fields,
@@ -988,8 +994,8 @@ pub(crate) async fn generate_index_on_compactor(
     inverted_idx_batch: RecordBatch,
     new_file_key: String,
     org_id: &str,
-    stream_name: &str,
     stream_type: StreamType,
+    stream_name: &str,
     full_text_search_fields: &[String],
     index_fields: &[String],
 ) -> Result<Vec<(String, FileMeta)>, anyhow::Error> {
@@ -1002,6 +1008,7 @@ pub(crate) async fn generate_index_on_compactor(
     let mut record_batches = prepare_index_record_batches(
         inverted_idx_batch,
         org_id,
+        stream_type,
         stream_name,
         &new_file_key,
         full_text_search_fields,
@@ -1012,7 +1019,7 @@ pub(crate) async fn generate_index_on_compactor(
     }
     let schema = record_batches.first().unwrap().schema();
 
-    let prefix_to_remove = format!("files/{}/logs/{}/", org_id, stream_name);
+    let prefix_to_remove = format!("files/{}/{}/{}/", org_id, stream_type, stream_name);
     let len_of_columns_to_invalidate = file_list_to_invalidate.len();
 
     let _timestamp: ArrayRef = Arc::new(Int64Array::from(
@@ -1073,8 +1080,8 @@ pub(crate) async fn generate_index_on_compactor(
         record_batches,
         original_file_size,
         org_id,
-        &index_stream_name,
         StreamType::Index,
+        &index_stream_name,
         &new_file_key,
         "index_creator",
     )
@@ -1087,6 +1094,7 @@ pub(crate) async fn generate_index_on_compactor(
 fn prepare_index_record_batches(
     inverted_idx_batch: RecordBatch,
     org_id: &str,
+    stream_type: StreamType,
     stream_name: &str,
     new_file_key: &str,
     full_text_search_fields: &[String],
@@ -1107,7 +1115,7 @@ fn prepare_index_record_batches(
         Field::new("segment_ids", DataType::Binary, true), // bitmap
     ]));
 
-    let prefix_to_remove = format!("files/{}/logs/{}/", org_id, stream_name);
+    let prefix_to_remove = format!("files/{}/{}/{}/", org_id, stream_type, stream_name);
     let file_name_without_prefix = new_file_key.trim_start_matches(&prefix_to_remove);
     let mut indexed_record_batches_to_merge = Vec::new();
 

src/service/compact/merge.rs

@@ -873,8 +873,8 @@ pub async fn merge_files(
                             inverted_idx_batch.clone(),
                             new_file_key.clone(),
                             org_id,
-                            stream_name,
                             stream_type,
+                            stream_name,
                             &full_text_search_fields,
                             &index_fields,
                         )