fix: apply log query limit after expressions (#8496)

Signed-off-by: discord9 <discord9@163.com>

src/query/src/log_query/planner.rs

@@ -34,8 +34,6 @@ use crate::log_query::error::{
     UnknownTableSnafu,
 };
 
-const DEFAULT_LIMIT: usize = 1000;
-
 pub struct LogQueryPlanner {
     table_provider: DfTableSourceProvider,
     session_state: SessionState,
@@ -99,19 +97,18 @@ impl LogQueryPlanner {
                 .context(DataFusionPlanningSnafu)?;
         }
 
-        // Apply limit
-        plan_builder = plan_builder
-            .limit(
-                query.limit.skip.unwrap_or(0),
-                Some(query.limit.fetch.unwrap_or(DEFAULT_LIMIT)),
-            )
-            .context(DataFusionPlanningSnafu)?;
-
         // Apply log expressions
         for expr in &query.exprs {
             plan_builder = self.process_log_expr(plan_builder, expr)?;
         }
 
+        // Apply pagination to the final result after all log expressions.
+        if query.limit.skip.is_some() || query.limit.fetch.is_some() {
+            plan_builder = plan_builder
+                .limit(query.limit.skip.unwrap_or(0), query.limit.fetch)
+                .context(DataFusionPlanningSnafu)?;
+        }
+
         // Build the final plan
         let plan = plan_builder.build().context(DataFusionPlanningSnafu)?;
 
@@ -849,7 +846,7 @@ mod tests {
         };
 
         let plan = planner.query_to_plan(log_query).await.unwrap();
-        let expected = "Limit: skip=10, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
+        let expected = "Limit: skip=10, fetch=None [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
 \n  Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
 \n    TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]";
 
@@ -884,9 +881,8 @@ mod tests {
         };
 
         let plan = planner.query_to_plan(log_query).await.unwrap();
-        let expected = "Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
-\n  Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
-\n    TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]";
+        let expected = "Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
+\n  TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]";
 
         assert_eq!(plan.display_indent_schema().to_string(), expected);
     }
@@ -931,8 +927,8 @@ mod tests {
         };
 
         let plan = planner.query_to_plan(log_query).await.unwrap();
-        let expected = "Aggregate: groupBy=[[greptime.public.test_table.host]], aggr=[[count(greptime.public.test_table.message) AS count_result]] [host:Utf8;N, count_result:Int64]\
-\n  Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
+        let expected = "Limit: skip=0, fetch=100 [host:Utf8;N, count_result:Int64]\
+\n  Aggregate: groupBy=[[greptime.public.test_table.host]], aggr=[[count(greptime.public.test_table.message) AS count_result]] [host:Utf8;N, count_result:Int64]\
 \n    Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
 \n      TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]";
 
@@ -971,8 +967,8 @@ mod tests {
         };
 
         let plan = planner.query_to_plan(log_query).await.unwrap();
-        let expected = "Projection: date_trunc(Utf8(\"day\"), greptime.public.test_table.timestamp) AS time_bucket [time_bucket:Timestamp(ms)]\
-        \n  Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
+        let expected = "Limit: skip=0, fetch=100 [time_bucket:Timestamp(ms)]\
+        \n  Projection: date_trunc(Utf8(\"day\"), greptime.public.test_table.timestamp) AS time_bucket [time_bucket:Timestamp(ms)]\
         \n    Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
         \n      TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]";
 
@@ -1054,9 +1050,9 @@ mod tests {
         };
 
         let plan = planner.query_to_plan(log_query).await.unwrap();
-        let expected = "Aggregate: groupBy=[[2__date_histogram__time_bucket]], aggr=[[count(2__date_histogram__time_bucket) AS count_result]] [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\
-\n  Projection: date_bin(Utf8(\"30 seconds\"), greptime.public.test_table.timestamp) AS 2__date_histogram__time_bucket [2__date_histogram__time_bucket:Timestamp(ns);N]\
-\n    Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
+        let expected = "Limit: skip=0, fetch=None [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\
+\n  Aggregate: groupBy=[[2__date_histogram__time_bucket]], aggr=[[count(2__date_histogram__time_bucket) AS count_result]] [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\
+\n    Projection: date_bin(Utf8(\"30 seconds\"), greptime.public.test_table.timestamp) AS 2__date_histogram__time_bucket [2__date_histogram__time_bucket:Timestamp(ns);N]\
 \n      Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\
 \n        TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]";
 

src/servers/src/http/logs.rs

@@ -19,24 +19,31 @@ use axum::extract::State;
 use axum::response::{IntoResponse, Response};
 use axum::{Extension, Json};
 use common_telemetry::tracing;
-use log_query::LogQuery;
+use log_query::{Limit, LogQuery};
 use session::context::{Channel, QueryContext};
 
 use crate::http::result::greptime_result_v1::GreptimedbV1Response;
 use crate::query_handler::LogQueryHandlerRef;
 
+const DEFAULT_FETCH: usize = 1000;
+
+fn apply_default_fetch(limit: &mut Limit) {
+    limit.fetch.get_or_insert(DEFAULT_FETCH);
+}
+
 #[axum_macros::debug_handler]
 #[tracing::instrument(skip_all, fields(protocol = "http", request_type = "logs"))]
 pub async fn logs(
     State(handler): State<LogQueryHandlerRef>,
     Extension(mut query_ctx): Extension<QueryContext>,
-    Json(params): Json<LogQuery>,
+    Json(mut params): Json<LogQuery>,
 ) -> Response {
     let exec_start = Instant::now();
     let db = query_ctx.get_db_string();
 
     query_ctx.set_channel(Channel::Log);
     let query_ctx = Arc::new(query_ctx);
+    apply_default_fetch(&mut params.limit);
 
     let _timer = crate::metrics::METRIC_HTTP_LOGS_ELAPSED
         .with_label_values(&[db.as_str()])
@@ -48,3 +55,34 @@ pub async fn logs(
     resp.with_execution_time(exec_start.elapsed().as_millis() as u64)
         .into_response()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_apply_default_fetch() {
+        let mut limit = Limit {
+            skip: None,
+            fetch: None,
+        };
+        apply_default_fetch(&mut limit);
+        assert_eq!(limit.fetch, Some(DEFAULT_FETCH));
+
+        let mut limit = Limit {
+            skip: Some(10),
+            fetch: None,
+        };
+        apply_default_fetch(&mut limit);
+        assert_eq!(limit.skip, Some(10));
+        assert_eq!(limit.fetch, Some(DEFAULT_FETCH));
+
+        let mut limit = Limit {
+            skip: Some(10),
+            fetch: Some(42),
+        };
+        apply_default_fetch(&mut limit);
+        assert_eq!(limit.skip, Some(10));
+        assert_eq!(limit.fetch, Some(42));
+    }
+}

tests-integration/tests/http.rs

@@ -44,7 +44,7 @@ use common_memory_manager::OnExhaustedPolicy;
 use common_options::plugin_options::StandaloneFlag;
 use flate2::Compression;
 use flate2::write::GzEncoder;
-use log_query::{Context, Limit, LogQuery, TimeFilter};
+use log_query::{AggFunc, Context, Limit, LogExpr, LogQuery, TimeFilter};
 use loki_proto::logproto::{EntryAdapter, LabelPairAdapter, PushRequest, StreamAdapter};
 use loki_proto::prost_types::Timestamp;
 use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
@@ -7880,6 +7880,84 @@ pub async fn test_log_query(store_type: StorageType) {
         ]
     );
 
+    let res = client
+        .get("/v1/sql?sql=create table logs_limit (`ts` timestamp time index, `message` string);")
+        .send()
+        .await;
+    assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await);
+    let res = client
+        .get("/v1/sql?sql=insert into logs_limit select number, 'hello' from numbers limit 1001;")
+        .send()
+        .await;
+    assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await);
+
+    let aggregate_query = LogQuery {
+        table: TableName::new("greptime", "public", "logs_limit"),
+        time_filter: TimeFilter {
+            start: Some("1970-01-01T00:00:00Z".to_string()),
+            end: Some("1970-01-01T00:00:02Z".to_string()),
+            span: None,
+        },
+        limit: Limit {
+            skip: None,
+            fetch: None,
+        },
+        columns: vec![],
+        filters: Default::default(),
+        context: Context::None,
+        exprs: vec![LogExpr::AggrFunc {
+            expr: vec![AggFunc::new(
+                "count".to_string(),
+                vec![LogExpr::NamedIdent("message".to_string())],
+                Some("count_result".to_string()),
+            )],
+            by: vec![],
+        }],
+    };
+    let res = client
+        .post("/v1/logs")
+        .header("Content-Type", "application/json")
+        .body(serde_json::to_string(&aggregate_query).unwrap())
+        .send()
+        .await;
+    assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await);
+    assert_eq!(get_rows_from_output(&res.text().await), "[[1001]]");
+
+    let mut output_query = LogQuery {
+        table: TableName::new("greptime", "public", "logs_limit"),
+        time_filter: TimeFilter {
+            start: Some("1970-01-01T00:00:00Z".to_string()),
+            end: Some("1970-01-01T00:00:02Z".to_string()),
+            span: None,
+        },
+        limit: Limit {
+            skip: None,
+            fetch: None,
+        },
+        columns: vec!["ts".to_string(), "message".to_string()],
+        filters: Default::default(),
+        context: Context::None,
+        exprs: vec![],
+    };
+    let res = client
+        .post("/v1/logs")
+        .header("Content-Type", "application/json")
+        .body(serde_json::to_string(&output_query).unwrap())
+        .send()
+        .await;
+    assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await);
+    assert_eq!(get_output_row_count(&res.text().await), 1000);
+
+    output_query.limit.fetch = Some(17);
+    let res = client
+        .post("/v1/logs")
+        .header("Content-Type", "application/json")
+        .body(serde_json::to_string(&output_query).unwrap())
+        .send()
+        .await;
+    assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await);
+    assert_eq!(get_output_row_count(&res.text().await), 17);
+
     guard.remove_all().await;
 }
 
@@ -9662,6 +9740,18 @@ fn get_rows_from_output(output: &str) -> String {
         .to_string()
 }
 
+fn get_output_row_count(output: &str) -> usize {
+    let resp: Value = serde_json::from_str(output).unwrap();
+    resp.get("output")
+        .and_then(Value::as_array)
+        .and_then(|v| v.first())
+        .and_then(|v| v.get("records"))
+        .and_then(|v| v.get("rows"))
+        .and_then(Value::as_array)
+        .unwrap()
+        .len()
+}
+
 fn compress_vec_with_gzip(data: Vec<u8>) -> Vec<u8> {
     let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
     encoder.write_all(&data).unwrap();