From f27e45ee4288c041327c7bd98e7f910879556448 Mon Sep 17 00:00:00 2001
From: Damien de Lemeny <contact@ddelemeny.me>
Date: Thu, 14 Mar 2024 16:11:47 -0500
Subject: [PATCH] Reduce default terms agg size, add shard_size
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-Authored-By: François Massot <francois.massot@gmail.com>
---
 pkg/quickwit/client/models.go | 1 +
 pkg/quickwit/data_query.go    | 5 +++++
 src/QueryBuilder/elastic.ts   | 7 +++++--
 src/dataquery.gen.ts          | 1 +
 4 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/pkg/quickwit/client/models.go b/pkg/quickwit/client/models.go
index b7e069e..6cfcbfc 100644
--- a/pkg/quickwit/client/models.go
+++ b/pkg/quickwit/client/models.go
@@ -260,6 +260,7 @@ type FiltersAggregation struct {
 type TermsAggregation struct {
 	Field       string                 `json:"field"`
 	Size        int                    `json:"size"`
+	ShardSize   int                    `json:"shard_size"`
 	Order       map[string]interface{} `json:"order"`
 	MinDocCount *int                   `json:"min_doc_count,omitempty"`
 	Missing     *string                `json:"missing,omitempty"`
diff --git a/pkg/quickwit/data_query.go b/pkg/quickwit/data_query.go
index 44d1c4e..752329f 100644
--- a/pkg/quickwit/data_query.go
+++ b/pkg/quickwit/data_query.go
@@ -229,6 +229,11 @@ func addTermsAgg(aggBuilder es.AggBuilder, bucketAgg *BucketAgg, metrics []*Metr
 		} else {
 			a.Size = stringToIntWithDefaultValue(bucketAgg.Settings.Get("size").MustString(), defaultSize)
 		}
+		if shard_size, err := bucketAgg.Settings.Get("shard_size").Int(); err == nil {
+			a.ShardSize = shard_size
+		} else {
+			a.ShardSize = stringToIntWithDefaultValue(bucketAgg.Settings.Get("shard_size").MustString(), defaultSize)
+		}
 
 		if minDocCount, err := bucketAgg.Settings.Get("min_doc_count").Int(); err == nil {
 			a.MinDocCount = &minDocCount
diff --git a/src/QueryBuilder/elastic.ts b/src/QueryBuilder/elastic.ts
index 96f9340..907f66e 100644
--- a/src/QueryBuilder/elastic.ts
+++ b/src/QueryBuilder/elastic.ts
@@ -11,6 +11,7 @@ type OrderByType =  '_key' | '_term' | '_count'
 function getTermsAgg(
   fieldName: string,
   size: number,
+  shard_size: number,
   orderBy: OrderByType = "_key",
   order: 'asc'|'desc' = 'asc'
   ): BucketAggregation {
@@ -20,6 +21,7 @@ function getTermsAgg(
     field: fieldName,
     settings:{
       size: size.toString(),
+      shard_size: shard_size.toString(),
       order: order,
       orderBy: orderBy,
     }
@@ -31,7 +33,8 @@ export function getDataQuery(queryDef: TermsQuery, refId: string): Elasticsearch
     {id:"count1", type:'count'}
   ];
 
-  // Default behaviour is to order results by { _key: asc }
+  // Default behaviour is to order results by { _key: asc } and get 100 terms,
+  // shard_size is set to 100 to limit terms fetched by shard (by default it's 100 * 10).
   // queryDef.order allows selection of asc/desc
   // queryDef.orderBy allows selection of doc_count ordering (defaults desc)
 
@@ -55,7 +58,7 @@ export function getDataQuery(queryDef: TermsQuery, refId: string): Elasticsearch
 
   const bucketAggs: BucketAggregation[] = [];
   if (queryDef.field) {
-    bucketAggs.push(getTermsAgg(queryDef.field, 500, orderBy, order))
+    bucketAggs.push(getTermsAgg(queryDef.field, 100, 100, orderBy, order))
   }
 
   return {
diff --git a/src/dataquery.gen.ts b/src/dataquery.gen.ts
index efe68c0..bd5dc47 100644
--- a/src/dataquery.gen.ts
+++ b/src/dataquery.gen.ts
@@ -71,6 +71,7 @@ export interface Terms extends BucketAggregationWithField {
   settings?: {
     order?: TermsOrder;
     size?: string;
+    shard_size?: string;
     min_doc_count?: string;
     orderBy?: string;
     missing?: string;