Skip to content

Commit b823790

Browse files
zachmuzachmu
committed
make set
1 parent 0ea893b commit b823790

File tree

4 files changed

+405
-0
lines changed

4 files changed

+405
-0
lines changed

‎enginetest/queries/queries.go

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5698,6 +5698,110 @@ SELECT * FROM cte WHERE d = 2;`,
56985698
{string("1,1,0,0")},
56995699
},
57005700
},
5701+
{
5702+
Query: `SELECT MAKE_SET(1, "a", "b", "c")`,
5703+
Expected: []sql.Row{
5704+
{string("a")},
5705+
},
5706+
},
5707+
{
5708+
Query: `SELECT MAKE_SET(1 | 4, "hello", "nice", "world")`,
5709+
Expected: []sql.Row{
5710+
{string("hello,world")},
5711+
},
5712+
},
5713+
{
5714+
Query: `SELECT MAKE_SET(0, "a", "b", "c")`,
5715+
Expected: []sql.Row{
5716+
{string("")},
5717+
},
5718+
},
5719+
{
5720+
Query: `SELECT MAKE_SET(3, "a", "b", "c")`,
5721+
Expected: []sql.Row{
5722+
{string("a,b")},
5723+
},
5724+
},
5725+
{
5726+
Query: `SELECT MAKE_SET(5, "a", "b", "c")`,
5727+
Expected: []sql.Row{
5728+
{string("a,c")},
5729+
},
5730+
},
5731+
{
5732+
Query: `SELECT MAKE_SET(7, "a", "b", "c")`,
5733+
Expected: []sql.Row{
5734+
{string("a,b,c")},
5735+
},
5736+
},
5737+
{
5738+
Query: `SELECT MAKE_SET(1024, "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k")`,
5739+
Expected: []sql.Row{
5740+
{string("k")},
5741+
},
5742+
},
5743+
{
5744+
Query: `SELECT MAKE_SET(1025, "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k")`,
5745+
Expected: []sql.Row{
5746+
{string("a,k")},
5747+
},
5748+
},
5749+
{
5750+
Query: `SELECT MAKE_SET(7, "a", NULL, "c")`,
5751+
Expected: []sql.Row{
5752+
{string("a,c")},
5753+
},
5754+
},
5755+
{
5756+
Query: `SELECT MAKE_SET(7, NULL, "b", "c")`,
5757+
Expected: []sql.Row{
5758+
{string("b,c")},
5759+
},
5760+
},
5761+
{
5762+
Query: `SELECT MAKE_SET(NULL, "a", "b", "c")`,
5763+
Expected: []sql.Row{
5764+
{nil},
5765+
},
5766+
},
5767+
{
5768+
Query: `SELECT MAKE_SET("5", "a", "b", "c")`,
5769+
Expected: []sql.Row{
5770+
{string("a,c")},
5771+
},
5772+
},
5773+
{
5774+
Query: `SELECT MAKE_SET(5.7, "a", "b", "c")`,
5775+
Expected: []sql.Row{
5776+
{string("b,c")},
5777+
},
5778+
},
5779+
{
5780+
Query: `SELECT MAKE_SET(-1, "a", "b", "c")`,
5781+
Expected: []sql.Row{
5782+
{string("a,b,c")},
5783+
},
5784+
},
5785+
{
5786+
Query: `SELECT MAKE_SET(16, "a", "b", "c")`,
5787+
Expected: []sql.Row{
5788+
{string("")},
5789+
},
5790+
},
5791+
{
5792+
Query: `SELECT MAKE_SET(3, "", "test", "")`,
5793+
Expected: []sql.Row{
5794+
{string(",test")},
5795+
},
5796+
},
5797+
{
5798+
Query: `SELECT MAKE_SET(i, "first", "second", "third") FROM mytable ORDER BY i`,
5799+
Expected: []sql.Row{
5800+
{string("first")},
5801+
{string("second")},
5802+
{string("first,second")},
5803+
},
5804+
},
57015805
{
57025806
Query: "SELECT version()",
57035807
Expected: []sql.Row{

‎sql/expression/function/make_set.go

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
// Copyright 2020-2024 Dolthub, Inc.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package function
16+
17+
import (
18+
"fmt"
19+
"strings"
20+
21+
"github.com/dolthub/go-mysql-server/sql"
22+
"github.com/dolthub/go-mysql-server/sql/types"
23+
)
24+
25+
// MakeSet implements the SQL function MAKE_SET() which returns a comma-separated set of strings
26+
// where the corresponding bit in bits is set
27+
type MakeSet struct {
28+
bits sql.Expression
29+
values []sql.Expression
30+
}
31+
32+
var _ sql.FunctionExpression = (*MakeSet)(nil)
33+
var _ sql.CollationCoercible = (*MakeSet)(nil)
34+
35+
// NewMakeSet creates a new MakeSet expression
36+
func NewMakeSet(args ...sql.Expression) (sql.Expression, error) {
37+
if len(args) < 2 {
38+
return nil, sql.ErrInvalidArgumentNumber.New("MAKE_SET", "2 or more", len(args))
39+
}
40+
41+
return &MakeSet{
42+
bits: args[0],
43+
values: args[1:],
44+
}, nil
45+
}
46+
47+
// FunctionName implements sql.FunctionExpression
48+
func (m *MakeSet) FunctionName() string {
49+
return "make_set"
50+
}
51+
52+
// Description implements sql.FunctionExpression
53+
func (m *MakeSet) Description() string {
54+
return "returns a set string (a string containing substrings separated by , characters) consisting of the strings that have the corresponding bit in bits set."
55+
}
56+
57+
// Children implements the Expression interface
58+
func (m *MakeSet) Children() []sql.Expression {
59+
children := []sql.Expression{m.bits}
60+
children = append(children, m.values...)
61+
return children
62+
}
63+
64+
// Resolved implements the Expression interface
65+
func (m *MakeSet) Resolved() bool {
66+
for _, child := range m.Children() {
67+
if !child.Resolved() {
68+
return false
69+
}
70+
}
71+
return true
72+
}
73+
74+
// IsNullable implements the Expression interface
75+
func (m *MakeSet) IsNullable() bool {
76+
return m.bits.IsNullable()
77+
}
78+
79+
// Type implements the Expression interface
80+
func (m *MakeSet) Type() sql.Type {
81+
return types.LongText
82+
}
83+
84+
// CollationCoercibility implements the interface sql.CollationCoercible
85+
func (m *MakeSet) CollationCoercibility(ctx *sql.Context) (collation sql.CollationID, coercibility byte) {
86+
// Start with highest coercibility (most coercible)
87+
collation = sql.Collation_Default
88+
coercibility = 5
89+
90+
for _, value := range m.values {
91+
valueCollation, valueCoercibility := sql.GetCoercibility(ctx, value)
92+
collation, coercibility = sql.ResolveCoercibility(collation, coercibility, valueCollation, valueCoercibility)
93+
}
94+
95+
return collation, coercibility
96+
}
97+
98+
// String implements the Expression interface
99+
func (m *MakeSet) String() string {
100+
children := m.Children()
101+
childStrs := make([]string, len(children))
102+
for i, child := range children {
103+
childStrs[i] = child.String()
104+
}
105+
return fmt.Sprintf("make_set(%s)", strings.Join(childStrs, ", "))
106+
}
107+
108+
// WithChildren implements the Expression interface
109+
func (m *MakeSet) WithChildren(children ...sql.Expression) (sql.Expression, error) {
110+
return NewMakeSet(children...)
111+
}
112+
113+
// Eval implements the Expression interface
114+
func (m *MakeSet) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
115+
bitsVal, err := m.bits.Eval(ctx, row)
116+
if err != nil {
117+
return nil, err
118+
}
119+
if bitsVal == nil {
120+
return nil, nil
121+
}
122+
123+
// Convert bits to uint64
124+
bitsInt, _, err := types.Uint64.Convert(ctx, bitsVal)
125+
if err != nil {
126+
return nil, err
127+
}
128+
bits := bitsInt.(uint64)
129+
130+
var result []string
131+
132+
// Check each value argument against the corresponding bit
133+
for i, valueExpr := range m.values {
134+
// Check if bit i is set
135+
if (bits & (1 << uint(i))) != 0 {
136+
val, err := valueExpr.Eval(ctx, row)
137+
if err != nil {
138+
return nil, err
139+
}
140+
// Skip NULL values
141+
if val != nil {
142+
valStr, _, err := types.LongText.Convert(ctx, val)
143+
if err != nil {
144+
return nil, err
145+
}
146+
result = append(result, valStr.(string))
147+
}
148+
}
149+
}
150+
151+
return strings.Join(result, ","), nil
152+
}

0 commit comments

Comments
 (0)