Geo type refactor (#761)

Signed-off-by: Rafał Strzaliński <[email protected]>

Author: nablaone

quesma/model/geo_type.go

@@ -0,0 +1,14 @@
+// Copyright Quesma, licensed under the Elastic License 2.0.
+// SPDX-License-Identifier: Elastic-2.0
+package model
+
+const QuesmaGeoLatFunction = "__quesma_geo_lat"
+const QuesmaGeoLonFunction = "__quesma_geo_lon"
+
+func NewGeoLat(propertyName string) Expr {
+	return NewFunction(QuesmaGeoLatFunction, NewColumnRef(propertyName))
+}
+
+func NewGeoLon(propertyName string) Expr {
+	return NewFunction(QuesmaGeoLonFunction, NewColumnRef(propertyName))
+}

quesma/model/metrics_aggregations/top_hits.go

@@ -4,12 +4,8 @@ package metrics_aggregations
 
 import (
 	"context"
-	"encoding/json"
 	"quesma/logger"
 	"quesma/model"
-	"quesma/schema"
-	"strconv"
-	"strings"
 )
 
 type TopHits struct {
@@ -50,49 +46,11 @@ func (query TopHits) TranslateSqlResponseToJson(rows []model.QueryResultRow) mod
 		sourceMap := model.JsonMap{}
 
 		for _, col := range valuesForHits {
-			var withoutQuotes string
-			if unquoted, err := strconv.Unquote(col.ColName); err == nil {
-				withoutQuotes = unquoted
-			} else {
-				withoutQuotes = col.ColName
-			}
-			colName, _ := strings.CutPrefix(withoutQuotes, `windowed_`)
 
-			if col.ColType.Name == schema.QuesmaTypePoint.Name {
-				hits := make(model.JsonMap)
-				// TODO suffixes (::lat, ::lon) hardcoded for now
-				// due to insufficient information in the schema
-				if strings.Contains(col.ColName, "::lon") {
-					hits["lon"] = col.ExtractValue(query.ctx)
-					colName = strings.TrimSuffix(col.ColName, "::lon")
-				}
-				if strings.Contains(col.ColName, "::lat") {
-					hits["lat"] = col.ExtractValue(query.ctx)
-					colName = strings.TrimSuffix(col.ColName, "::lat")
-				}
-				if _, ok := sourceMap[colName]; ok {
-					currentHits := sourceMap[colName].(model.JsonMap)
-					for k, v := range currentHits {
-						hits[k] = v
-					}
-					sourceMap[colName] = hits
-				} else {
-					sourceMap[colName] = hits
-				}
+			value := col.ExtractValue(query.ctx)
+
+			sourceMap[col.ColName] = value
 
-			} else {
-				value := col.ExtractValue(query.ctx)
-				// TODO: this is hack, we should not assume this is location
-				if strings.HasSuffix(col.ColName, "Location") {
-					if valueStr, ok := value.(string); ok {
-						var valueJson model.JsonMap
-						if err := json.Unmarshal([]byte(valueStr), &valueJson); err == nil {
-							value = valueJson
-						}
-					}
-				}
-				sourceMap[col.ColName] = value
-			}
 		}
 
 		elem := model.JsonMap{

quesma/queryparser/pancake_aggregation_parser_buckets.go

@@ -236,21 +236,21 @@ func (cw *ClickhouseQueryTranslator) pancakeTryBucketAggregation(aggregation *pa
 
 		zoomLiteral := model.NewLiteral(precisionZoom)
 
-		// TODO columns names should be created according to the schema
-		var lon = model.AsString(field)
-		lon = strings.Trim(lon, "\"")
-		lon = lon + "::lon"
-		var lat = model.AsString(field)
-		lat = strings.Trim(lat, "\"")
-		lat = lat + "::lat"
-		toFloatFunLon := model.NewFunction("toFloat64", model.NewColumnRef(lon))
+		fieldName, err := strconv.Unquote(model.AsString(field))
+		if err != nil {
+			return false, err
+		}
+		lon := model.NewGeoLon(fieldName)
+		lat := model.NewGeoLat(fieldName)
+
+		toFloatFunLon := model.NewFunction("toFloat64", lon)
 		var infixX model.Expr
 		infixX = model.NewParenExpr(model.NewInfixExpr(toFloatFunLon, "+", model.NewLiteral(180.0)))
 		infixX = model.NewParenExpr(model.NewInfixExpr(infixX, "/", model.NewLiteral(360.0)))
 		infixX = model.NewInfixExpr(infixX, "*",
 			model.NewFunction("POWER", model.NewLiteral(2), zoomLiteral))
 		xTile := model.NewFunction("FLOOR", infixX)
-		toFloatFunLat := model.NewFunction("toFloat64", model.NewColumnRef(lat))
+		toFloatFunLat := model.NewFunction("toFloat64", lat)
 		radians := model.NewFunction("RADIANS", toFloatFunLat)
 		tan := model.NewFunction("TAN", radians)
 		cos := model.NewFunction("COS", radians)

quesma/queryparser/pancake_aggregation_parser_metrics.go

@@ -108,8 +108,8 @@ func generateMetricSelectedColumns(ctx context.Context, metricsAggr metricsAggre
 		if col, ok := firstExpr.(model.ColumnRef); ok {
 			colName := col.ColumnName
 			// TODO we have create columns according to the schema
-			latColumn := model.NewColumnRef(colName + "::lat")
-			lonColumn := model.NewColumnRef(colName + "::lon")
+			latColumn := model.NewGeoLat(colName)
+			lonColumn := model.NewGeoLon(colName)
 			castLat := model.NewFunction("CAST", latColumn, model.NewLiteral(fmt.Sprintf("'%s'", "Float")))
 			castLon := model.NewFunction("CAST", lonColumn, model.NewLiteral(fmt.Sprintf("'%s'", "Float")))
 			result = append(result, model.NewFunction("avgOrNull", castLat))

quesma/queryparser/pancake_sql_query_generation_test.go

@@ -143,6 +143,7 @@ func TestPancakeQueryGeneration(t *testing.T) {
 
 			actualMinusExpected, expectedMinusActual := util.MapDifference(pancakeJson,
 				expectedAggregationsPart, acceptableDifference, true, true)
+
 			if len(actualMinusExpected) != 0 {
 				pp.Println("ACTUAL diff", actualMinusExpected)
 			}

quesma/queryparser/pancake_top_hits.go

@@ -7,7 +7,6 @@ import (
 	"quesma/model"
 	"quesma/model/metrics_aggregations"
 	"strconv"
-	"strings"
 )
 
 func (p *pancakeSqlQueryGenerator) quotedLiteral(name string) model.LiteralExpr {
@@ -83,25 +82,14 @@ func (p *pancakeSqlQueryGenerator) generateTopHitsQuery(aggregation *pancakeMode
 	groupTableName := "group_table"
 	hitTableName := "hit_table"
 
-	hitTableLiteral := func(reference string) model.Expr {
-		return model.NewLiteral(strconv.Quote(hitTableName) + "." + strconv.Quote(reference))
-	}
 	groupTableLiteral := func(reference string) model.Expr {
 		return model.NewLiteral(strconv.Quote(groupTableName) + "." + strconv.Quote(reference))
 	}
 
 	convertColumnRefToHitTable := func(expr model.Expr) model.Expr {
 		switch exprTyped := expr.(type) {
 		case model.ColumnRef:
-			// TODO: hack alert, we treat geo here in unique way
-			if strings.HasSuffix(exprTyped.ColumnName, "Location") {
-				return model.NewFunction("map",
-					model.NewLiteral("'lat'"),
-					hitTableLiteral(exprTyped.ColumnName+"::lat"),
-					model.NewLiteral("'lon'"),
-					hitTableLiteral(exprTyped.ColumnName+"::lon"),
-				)
-			}
+
 			return model.ColumnRef{
 				TableAlias: hitTableName,
 				ColumnName: exprTyped.ColumnName,

quesma/quesma/schema_transformer.go

@@ -152,63 +152,117 @@ func (s *SchemaCheckPass) applyIpTransformations(indexSchema schema.Schema, quer
 	return query, nil
 }
 
-func (s *SchemaCheckPass) applyGeoTransformations(currentSchema schema.Schema, query *model.Query) (*model.Query, error) {
-	fromTable := query.TableName
+func (s *SchemaCheckPass) applyGeoTransformations(schemaInstance schema.Schema, query *model.Query) (*model.Query, error) {
+
+	replace := make(map[string]model.Expr)
+
+	for _, field := range schemaInstance.Fields {
+		if field.Type.Name == schema.QuesmaTypePoint.Name {
+
+			lon := model.NewColumnRef(field.InternalPropertyName.AsString() + "::lon")
+			lat := model.NewColumnRef(field.InternalPropertyName.AsString() + "::lat")
+
+			// This is a workaround. Clickhouse Point is defined as Tuple. We need to know the type of the tuple.
+			// In this step we merge two columns into single map here. Map is in elastic format.
+
+			// In this point we assume that Quesma point type is stored into two separate columns.
+			replace[field.PropertyName.AsString()] = model.NewFunction("map",
+				model.NewLiteral("'lat'"),
+				lat,
+				model.NewLiteral("'lon'"),
+				lon)
+
+			// these a just if we need multifields support
+			replace[field.PropertyName.AsString()+".lat"] = lat
+			replace[field.PropertyName.AsString()+".lon"] = lon
+
+			// if the point is stored as a single column, we need to extract the lat and lon
+			//replace[field.PropertyName.AsString()] = model.NewFunction("give_me_point", model.NewColumnRef(field.InternalPropertyName.AsString()))
+			//replace[field.PropertyName.AsString()+".lat"] = model.NewFunction("give_me_lat", model.NewColumnRef(field.InternalPropertyName.AsString()))
+			//replace[field.PropertyName.AsString()+".lon"] = model.NewFunction("give_me_lon", model.NewColumnRef(field.InternalPropertyName.AsString()))
+
+		}
+	}
 
 	visitor := model.NewBaseVisitor()
-	visitor.OverrideVisitSelectCommand = func(b *model.BaseExprVisitor, e model.SelectCommand) interface{} {
-		if s.schemaRegistry == nil {
-			logger.Error().Msg("Schema registry is not set")
-			return e
-		}
-		schemaInstance, exists := s.schemaRegistry.FindSchema(schema.TableName(fromTable))
-		if !exists {
-			logger.Error().Msgf("Schema fot table %s not found", fromTable)
-			return e
-		}
-		var groupBy []model.Expr
-		for _, expr := range e.GroupBy {
-			groupByExpr := expr.Accept(b).(model.Expr)
-			if col, ok := expr.(model.ColumnRef); ok {
-				// This checks if the column is of type point
-				// and if it is, it appends the lat and lon columns to the group by clause
-				field := schemaInstance.Fields[schema.FieldName(col.ColumnName)]
-				if field.Type.Name == schema.QuesmaTypePoint.Name {
-					// TODO suffixes ::lat, ::lon are hardcoded for now
-					groupBy = append(groupBy, model.NewColumnRef(field.InternalPropertyName.AsString()+"::lat"))
-					groupBy = append(groupBy, model.NewColumnRef(field.InternalPropertyName.AsString()+"::lon"))
-				} else {
-					groupBy = append(groupBy, groupByExpr)
+
+	visitor.OverrideVisitColumnRef = func(b *model.BaseExprVisitor, e model.ColumnRef) interface{} {
+		if expr, ok := replace[e.ColumnName]; ok {
+			return expr
+		}
+		return e
+	}
+
+	visitor.OverrideVisitFunction = func(b *model.BaseExprVisitor, e model.FunctionExpr) interface{} {
+
+		var suffix string
+		switch e.Name {
+		case model.QuesmaGeoLatFunction:
+			suffix = ".lat"
+		case model.QuesmaGeoLonFunction:
+			suffix = ".lon"
+		}
+
+		if suffix != "" && len(e.Args) == 1 {
+			if col, ok := e.Args[0].(model.ColumnRef); ok {
+				if expr, ok := replace[col.ColumnName+suffix]; ok {
+					return expr
 				}
-			} else {
-				groupBy = append(groupBy, groupByExpr)
 			}
 		}
-		var columns []model.Expr
-		for _, expr := range e.Columns {
+
+		return model.NewFunction(e.Name, b.VisitChildren(e.Args)...)
+	}
+
+	visitor.OverrideVisitSelectCommand = func(v *model.BaseExprVisitor, query model.SelectCommand) interface{} {
+		var columns, groupBy []model.Expr
+		var orderBy []model.OrderByExpr
+		from := query.FromClause
+		where := query.WhereClause
+
+		for _, expr := range query.Columns {
+			var alias string
 			if col, ok := expr.(model.ColumnRef); ok {
-				// This checks if the column is of type point
-				// and if it is, it appends the lat and lon columns to the select clause
-				field := schemaInstance.Fields[schema.FieldName(col.ColumnName)]
-				if field.Type.Name == schema.QuesmaTypePoint.Name {
-					// TODO suffixes ::lat, ::lon are hardcoded for now
-					columns = append(columns, model.NewColumnRef(field.InternalPropertyName.AsString()+"::lat"))
-					columns = append(columns, model.NewColumnRef(field.InternalPropertyName.AsString()+"::lon"))
-				} else {
-					columns = append(columns, expr.Accept(b).(model.Expr))
+				if _, ok := replace[col.ColumnName]; ok {
+					alias = col.ColumnName
 				}
-			} else {
-				columns = append(columns, expr.Accept(b).(model.Expr))
 			}
+
+			col := expr.Accept(v).(model.Expr)
+
+			if alias != "" {
+				col = model.NewAliasedExpr(col, alias)
+			}
+
+			columns = append(columns, col)
+		}
+		for _, expr := range query.GroupBy {
+			groupBy = append(groupBy, expr.Accept(v).(model.Expr))
+		}
+		for _, expr := range query.OrderBy {
+			orderBy = append(orderBy, expr.Accept(v).(model.OrderByExpr))
+		}
+		if query.FromClause != nil {
+			from = query.FromClause.Accept(v).(model.Expr)
+		}
+		if query.WhereClause != nil {
+			where = query.WhereClause.Accept(v).(model.Expr)
 		}
 
-		var fromClause model.Expr
-		if e.FromClause != nil {
-			fromClause = e.FromClause.Accept(b).(model.Expr)
+		var namedCTEs []*model.CTE
+		if query.NamedCTEs != nil {
+			for _, cte := range query.NamedCTEs {
+				namedCTEs = append(namedCTEs, cte.Accept(v).(*model.CTE))
+			}
 		}
 
-		return model.NewSelectCommand(columns, groupBy, e.OrderBy,
-			fromClause, e.WhereClause, e.LimitBy, e.Limit, e.SampleLimit, e.IsDistinct, e.NamedCTEs)
+		var limitBy []model.Expr
+		if query.LimitBy != nil {
+			for _, expr := range query.LimitBy {
+				limitBy = append(limitBy, expr.Accept(v).(model.Expr))
+			}
+		}
+		return model.NewSelectCommand(columns, groupBy, orderBy, from, where, limitBy, query.Limit, query.SampleLimit, query.IsDistinct, namedCTEs)
 	}
 
 	expr := query.SelectCommand.Accept(visitor)
@@ -329,9 +383,6 @@ func (s *SchemaCheckPass) applyWildcardExpansion(indexSchema schema.Schema, quer
 
 		cols := make([]string, 0, len(indexSchema.Fields))
 		for _, col := range indexSchema.Fields {
-			if col.Type.Name == schema.QuesmaTypePoint.Name { // Temporary workaround for kibana_flights
-				continue
-			}
 			cols = append(cols, col.InternalPropertyName.AsString())
 		}
 		sort.Strings(cols)
@@ -520,38 +571,6 @@ func (s *SchemaCheckPass) Transform(queries []*model.Query) ([]*model.Query, err
 	return queries, nil
 }
 
-type GeoIpResultTransformer struct {
-	schemaRegistry schema.Registry
-	fromTable      string
-}
-
-func (g *GeoIpResultTransformer) Transform(result [][]model.QueryResultRow) ([][]model.QueryResultRow, error) {
-	if g.schemaRegistry == nil {
-		logger.Error().Msg("Schema registry is not set")
-		return result, nil
-	}
-	schemaInstance, exists := g.schemaRegistry.FindSchema(schema.TableName(g.fromTable))
-	if !exists {
-		logger.Error().Msgf("Schema fot table %s not found", g.fromTable)
-		return result, nil
-	}
-	for i, rows := range result {
-		for j, row := range rows {
-			for k, col := range row.Cols {
-				if strings.Contains(col.ColName, "::lat") {
-					colType := schemaInstance.Fields[schema.FieldName(strings.TrimSuffix(col.ColName, "::lat"))].Type
-					result[i][j].Cols[k].ColType = colType
-				}
-				if strings.Contains(col.ColName, "::lon") {
-					colType := schemaInstance.Fields[schema.FieldName(strings.TrimSuffix(col.ColName, "::lon"))].Type
-					result[i][j].Cols[k].ColType = colType
-				}
-			}
-		}
-	}
-	return result, nil
-}
-
 // ArrayResultTransformer is a transformer that transforms array columns into string representation
 type ArrayResultTransformer struct {
 }

quesma/quesma/schema_transformer_test.go

@@ -258,8 +258,11 @@ func Test_ipRangeTransform(t *testing.T) {
 			TableName: "kibana_sample_data_flights",
 			SelectCommand: model.SelectCommand{
 				FromClause: model.NewTableRef("kibana_sample_data_flights"),
-				Columns: []model.Expr{model.NewColumnRef("DestLocation::lat"),
-					model.NewColumnRef("DestLocation::lon")},
+				Columns: []model.Expr{model.NewAliasedExpr(model.NewFunction("map",
+					model.NewLiteral("'lat'"),
+					model.NewColumnRef("DestLocation::lat"),
+					model.NewLiteral("'lon'"),
+					model.NewColumnRef("DestLocation::lon")), "DestLocation")},
 			}},
 	}
 	queries := [][]*model.Query{

quesma/quesma/search.go

@@ -764,7 +764,6 @@ func (q *QueryRunner) postProcessResults(table *clickhouse.Table, results [][]mo
 		transformer model.ResultTransformer
 	}{
 		{"replaceColumNamesWithFieldNames", &replaceColumNamesWithFieldNames{}},
-		{"geoIpResultTransformer", &GeoIpResultTransformer{schemaRegistry: q.schemaRegistry, fromTable: table.Name}},
 		{"arrayResultTransformer", &ArrayResultTransformer{}},
 	}