// Copyright 2023 Molecula Corp. All rights reserved. package planner import ( "context" "fmt" "strings" "time" pilosa "github.com/featurebasedb/featurebase/v3" "github.com/featurebasedb/featurebase/v3/dax" "github.com/featurebasedb/featurebase/v3/pql" "github.com/featurebasedb/featurebase/v3/sql3" "github.com/featurebasedb/featurebase/v3/sql3/parser" "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) // PlanOpPQLDistinctScan plan operator handles a PQL distinct scan // i.e. a scan with only one column that is used in a distinct query type PlanOpPQLDistinctScan struct { planner *ExecutionPlanner tableName string column string filter types.PlanExpression topExpr types.PlanExpression warnings []string } func NewPlanOpPQLDistinctScan(p *ExecutionPlanner, tableName string, column string) (*PlanOpPQLDistinctScan, error) { if strings.EqualFold(string(dax.PrimaryKeyFieldName), column) { return nil, sql3.NewErrInternalf("non _id column required") } return &PlanOpPQLDistinctScan{ planner: p, tableName: tableName, column: column, warnings: make([]string, 0), }, nil } func (p *PlanOpPQLDistinctScan) Plan() map[string]interface{} { result := make(map[string]interface{}) result["_op"] = fmt.Sprintf("%T", p) result["_schema"] = p.Schema().Plan() result["tableName"] = p.tableName if p.topExpr != nil { result["topExpr"] = p.topExpr.Plan() } if p.filter != nil { result["filter"] = p.filter.Plan() } result["column"] = p.column return result } func (p *PlanOpPQLDistinctScan) String() string { return "" } func (p *PlanOpPQLDistinctScan) AddWarning(warning string) { p.warnings = append(p.warnings, warning) } func (p *PlanOpPQLDistinctScan) Warnings() []string { return p.warnings } func (p *PlanOpPQLDistinctScan) Name() string { return p.tableName } func (p *PlanOpPQLDistinctScan) IsFilterable() bool { return true } func (p *PlanOpPQLDistinctScan) UpdateFilters(filterCondition types.PlanExpression) (types.PlanOperator, error) { p.filter = filterCondition return p, nil } func (p *PlanOpPQLDistinctScan) UpdateTimeQuantumFilters(filters ...types.PlanExpression) (types.PlanOperator, error) { return p, nil } func (p *PlanOpPQLDistinctScan) Schema() types.Schema { result := make(types.Schema, 0) tname := dax.TableName(p.tableName) table, err := p.planner.schemaAPI.TableByName(context.Background(), tname) if err != nil { return result } for _, fld := range table.Fields { if strings.EqualFold(string(fld.Name), p.column) { result = append(result, &types.PlannerColumn{ ColumnName: string(fld.Name), RelationName: p.tableName, Type: fieldSQLDataType(pilosa.FieldToFieldInfo(fld)), }) break } } return result } func (p *PlanOpPQLDistinctScan) Children() []types.PlanOperator { return []types.PlanOperator{} } func (p *PlanOpPQLDistinctScan) Iterator(ctx context.Context, row types.Row) (types.RowIterator, error) { return &distinctScanRowIter{ planner: p.planner, tableName: p.tableName, column: p.column, predicate: p.filter, topExpr: p.topExpr, }, nil } func (p *PlanOpPQLDistinctScan) WithChildren(children ...types.PlanOperator) (types.PlanOperator, error) { return nil, nil } type distinctScanRowIter struct { planner *ExecutionPlanner tableName string column string predicate types.PlanExpression topExpr types.PlanExpression result []interface{} rowWidth int columnDataType parser.ExprDataType } var _ types.RowIterator = (*distinctScanRowIter)(nil) func (i *distinctScanRowIter) Next(ctx context.Context) (types.Row, error) { if i.result == nil { err := i.planner.checkAccess(ctx, i.tableName, accessTypeReadData) if err != nil { return nil, err } //go get the schema def and map names to indexes in the resultant row tname := dax.TableName(i.tableName) table, err := i.planner.schemaAPI.TableByName(ctx, tname) if err != nil { if isTableNotFoundError(err) { return nil, sql3.NewErrInternalf("table not found '%s'", i.tableName) } return nil, err } i.rowWidth = 1 for _, fld := range table.Fields { if strings.EqualFold(i.column, string(fld.Name)) { i.columnDataType = fieldSQLDataType(pilosa.FieldToFieldInfo(fld)) break } } var cond *pql.Call cond, err = i.planner.generatePQLCallFromExpr(ctx, i.predicate) if err != nil { return nil, err } if cond == nil { cond = &pql.Call{Name: "All"} } if i.topExpr != nil { _, ok := i.topExpr.(*intLiteralPlanExpression) if !ok { return nil, sql3.NewErrInternalf("unexpected top expression type: %T", i.topExpr) } pqlValue, err := planExprToValue(i.topExpr) if err != nil { return nil, err } cond = &pql.Call{ Name: "Limit", Children: []*pql.Call{cond}, Args: map[string]interface{}{"limit": pqlValue}, Type: pql.PrecallGlobal, } } call := &pql.Call{ Name: "Distinct", Args: map[string]interface{}{"field": i.column}, Children: []*pql.Call{cond}, } queryResponse, err := i.planner.executor.Execute(ctx, table, &pql.Query{Calls: []*pql.Call{call}}, nil, nil) if err != nil { return nil, err } switch res := queryResponse.Results[0].(type) { case *pilosa.Row: result := make([]interface{}, 0) if len(res.Keys) > 0 { for _, n := range res.Keys { result = append(result, n) } } else { for _, n := range res.Columns() { result = append(result, int64(n)) } } i.result = result case pilosa.SignedRow: result := make([]interface{}, 0) negs := res.Neg.Columns() pos := res.Pos.Columns() for _, n := range negs { result = append(result, -(int64(n))) } for _, n := range pos { result = append(result, int64(n)) } i.result = result case pilosa.DistinctTimestamp: result := make([]interface{}, 0) for _, n := range res.Values { if tm, err := time.ParseInLocation(time.RFC3339Nano, n, time.UTC); err == nil { result = append(result, tm) } else { return nil, sql3.NewErrInternalf("unable to convert to time.Time: %v", n) } } i.result = result default: return nil, sql3.NewErrInternalf("unexpected Distinct() result type: %T", res) } } if len(i.result) > 0 { row := make([]interface{}, i.rowWidth) result := i.result[0] switch t := i.columnDataType.(type) { case *parser.DataTypeBool: val, ok := result.(int64) if !ok { return nil, sql3.NewErrInternalf("unexpected type for column value '%T'", result) } row[0] = val == 1 case *parser.DataTypeDecimal: val, ok := result.(int64) if !ok { return nil, sql3.NewErrInternalf("unexpected type for column value '%T'", result) } row[0] = pql.NewDecimal(val, t.Scale) case *parser.DataTypeIDSet: val, ok := result.(int64) if !ok { return nil, sql3.NewErrInternalf("unexpected type for column value '%T'", result) } row[0] = []uint64{uint64(val)} case *parser.DataTypeStringSet: val, ok := result.(string) if !ok { return nil, sql3.NewErrInternalf("unexpected type for column value '%T'", result) } row[0] = []string{val} default: row[0] = result } // Move to next result element. i.result = i.result[1:] return row, nil } return nil, types.ErrNoMoreRows }