Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

sql: allow array builtins to operate on tuples #70332

Merged
merged 1 commit into from
Sep 21, 2021
+240 −5
Merged

sql: allow array builtins to operate on tuples #70332

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/generated/sql/aggregates.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@
</span></td></tr>
<tr><td><a name="array_agg"></a><code>array_agg(arg1: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Aggregates the selected values into an array.</p>
</span></td></tr>
<tr><td><a name="array_agg"></a><code>array_agg(arg1: tuple) &rarr; tuple[]</code></td><td><span class="funcdesc"><p>Aggregates the selected values into an array.</p>
</span></td></tr>
<tr><td><a name="array_agg"></a><code>array_agg(arg1: varbit) &rarr; varbit[]</code></td><td><span class="funcdesc"><p>Aggregates the selected values into an array.</p>
</span></td></tr>
<tr><td><a name="avg"></a><code>avg(arg1: <a href="decimal.html">decimal</a>) &rarr; <a href="decimal.html">decimal</a></code></td><td><span class="funcdesc"><p>Calculates the average of the selected values.</p>
Expand Down
14 changes: 14 additions & 0 deletions docs/generated/sql/functions.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@
</span></td></tr>
<tr><td><a name="array_append"></a><code>array_append(array: timetz[], elem: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Appends <code>elem</code> to <code>array</code>, returning the result.</p>
</span></td></tr>
<tr><td><a name="array_append"></a><code>array_append(array: tuple[], elem: tuple) &rarr; tuple[]</code></td><td><span class="funcdesc"><p>Appends <code>elem</code> to <code>array</code>, returning the result.</p>
</span></td></tr>
<tr><td><a name="array_append"></a><code>array_append(array: varbit[], elem: varbit) &rarr; varbit[]</code></td><td><span class="funcdesc"><p>Appends <code>elem</code> to <code>array</code>, returning the result.</p>
</span></td></tr>
<tr><td><a name="array_cat"></a><code>array_cat(left: <a href="bool.html">bool</a>[], right: <a href="bool.html">bool</a>[]) &rarr; <a href="bool.html">bool</a>[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
Expand Down Expand Up @@ -81,6 +83,8 @@
</span></td></tr>
<tr><td><a name="array_cat"></a><code>array_cat(left: timetz[], right: timetz[]) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
</span></td></tr>
<tr><td><a name="array_cat"></a><code>array_cat(left: tuple[], right: tuple[]) &rarr; tuple[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
</span></td></tr>
<tr><td><a name="array_cat"></a><code>array_cat(left: varbit[], right: varbit[]) &rarr; varbit[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
</span></td></tr>
<tr><td><a name="array_length"></a><code>array_length(input: anyelement[], array_dimension: <a href="int.html">int</a>) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Calculates the length of <code>input</code> on the provided <code>array_dimension</code>. However, because CockroachDB doesn’t yet support multi-dimensional arrays, the only supported <code>array_dimension</code> is <strong>1</strong>.</p>
Expand Down Expand Up @@ -125,6 +129,8 @@
</span></td></tr>
<tr><td><a name="array_position"></a><code>array_position(array: timetz[], elem: timetz) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Return the index of the first occurrence of <code>elem</code> in <code>array</code>.</p>
</span></td></tr>
<tr><td><a name="array_position"></a><code>array_position(array: tuple[], elem: tuple) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Return the index of the first occurrence of <code>elem</code> in <code>array</code>.</p>
</span></td></tr>
<tr><td><a name="array_position"></a><code>array_position(array: varbit[], elem: varbit) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Return the index of the first occurrence of <code>elem</code> in <code>array</code>.</p>
</span></td></tr>
<tr><td><a name="array_positions"></a><code>array_positions(array: <a href="bool.html">bool</a>[], elem: <a href="bool.html">bool</a>) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
Expand Down Expand Up @@ -165,6 +171,8 @@
</span></td></tr>
<tr><td><a name="array_positions"></a><code>array_positions(array: timetz[], elem: timetz) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
</span></td></tr>
<tr><td><a name="array_positions"></a><code>array_positions(array: tuple[], elem: tuple) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
</span></td></tr>
<tr><td><a name="array_positions"></a><code>array_positions(array: varbit[], elem: varbit) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
</span></td></tr>
<tr><td><a name="array_prepend"></a><code>array_prepend(elem: <a href="bool.html">bool</a>, array: <a href="bool.html">bool</a>[]) &rarr; <a href="bool.html">bool</a>[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
Expand Down Expand Up @@ -205,6 +213,8 @@
</span></td></tr>
<tr><td><a name="array_prepend"></a><code>array_prepend(elem: timetz, array: timetz[]) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
</span></td></tr>
<tr><td><a name="array_prepend"></a><code>array_prepend(elem: tuple, array: tuple[]) &rarr; tuple[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
</span></td></tr>
<tr><td><a name="array_prepend"></a><code>array_prepend(elem: varbit, array: varbit[]) &rarr; varbit[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
</span></td></tr>
<tr><td><a name="array_remove"></a><code>array_remove(array: <a href="bool.html">bool</a>[], elem: <a href="bool.html">bool</a>) &rarr; <a href="bool.html">bool</a>[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
Expand Down Expand Up @@ -245,6 +255,8 @@
</span></td></tr>
<tr><td><a name="array_remove"></a><code>array_remove(array: timetz[], elem: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
</span></td></tr>
<tr><td><a name="array_remove"></a><code>array_remove(array: tuple[], elem: tuple) &rarr; tuple[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
</span></td></tr>
<tr><td><a name="array_remove"></a><code>array_remove(array: varbit[], elem: varbit) &rarr; varbit[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
</span></td></tr>
<tr><td><a name="array_replace"></a><code>array_replace(array: <a href="bool.html">bool</a>[], toreplace: <a href="bool.html">bool</a>, replacewith: <a href="bool.html">bool</a>) &rarr; <a href="bool.html">bool</a>[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>
Expand Down Expand Up @@ -285,6 +297,8 @@
</span></td></tr>
<tr><td><a name="array_replace"></a><code>array_replace(array: timetz[], toreplace: timetz, replacewith: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>
</span></td></tr>
<tr><td><a name="array_replace"></a><code>array_replace(array: tuple[], toreplace: tuple, replacewith: tuple) &rarr; tuple[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>
</span></td></tr>
<tr><td><a name="array_replace"></a><code>array_replace(array: varbit[], toreplace: varbit, replacewith: varbit) &rarr; varbit[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>
</span></td></tr>
<tr><td><a name="array_to_string"></a><code>array_to_string(input: anyelement[], delim: <a href="string.html">string</a>) &rarr; <a href="string.html">string</a></code></td><td><span class="funcdesc"><p>Join an array into a string with a delimiter.</p>
Expand Down
8 changes: 8 additions & 0 deletions pkg/sql/distsql_physical_planner.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,14 @@ func (v *distSQLExprCheckVisitor) VisitPre(expr tree.Expr) (recurse bool, newExp
v.err = newQueryNotSupportedErrorf("cast to %s is not supported by distsql", t.Type)
return false, expr
}
case *tree.DArray:
// We need to check for arrays of untyped tuples here since constant-folding
// on builtin functions sometimes produces this.
if t.ResolvedType().ArrayContents() == types.AnyTuple {
v.err = newQueryNotSupportedErrorf("array %s cannot be executed with distsql", t)
return false, expr
}

}
return true, expr
}
Expand Down
157 changes: 157 additions & 0 deletions pkg/sql/logictest/testdata/logic_test/array
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1746,6 +1746,8 @@ INSERT INTO t VALUES (
'{0101, 11}',
'{12.34, 45.67}');

subtest array_tuples

# Test for #32715: able to distribute queries with arrays of tuples.

statement ok
Expand Down Expand Up @@ -1781,3 +1783,158 @@ SELECT ARRAY[(k, 'foo'), (1, v)] FROM kv
{"(3,foo)","(1,three)"}
{"(4,foo)","(1,four)"}
{"(5,foo)","(1,)"}

# Test array builtins on record types

query T rowsort
SELECT array_cat(ARRAY[ROW(10,'fish')], ARRAY[(k,v)]) FROM kv
----
{"(10,fish)","(1,one)"}
{"(10,fish)","(2,two)"}
{"(10,fish)","(3,three)"}
{"(10,fish)","(4,four)"}
{"(10,fish)","(5,)"}

query T rowsort
SELECT array_cat(ARRAY[ROW(1,NULL)], ARRAY[(k,v)]) FROM kv
----
{"(1,)","(1,one)"}
{"(1,)","(2,two)"}
{"(1,)","(3,three)"}
{"(1,)","(4,four)"}
{"(1,)","(5,)"}

query T rowsort
SELECT array_cat(ARRAY[NULL::record], ARRAY[(k,v)]) FROM kv
----
{NULL,"(1,one)"}
{NULL,"(2,two)"}
{NULL,"(3,three)"}
{NULL,"(4,four)"}
{NULL,"(5,)"}

query T
SELECT array_cat(ARRAY[ROW(1,2)], ARRAY[NULL::record])
----
{"(1,2)",NULL}

query T
SELECT array_agg(ROW(1, 2))
----
{"(1,2)"}

query T
SELECT array_agg(ROW(k,v) ORDER BY k) FROM kv
----
{"(1,one)","(2,two)","(3,three)","(4,four)","(5,)"}

query T rowsort
SELECT array_append(ARRAY[ROW(10,'fish')], (k,v)) FROM kv
----
{"(10,fish)","(1,one)"}
{"(10,fish)","(2,two)"}
{"(10,fish)","(3,three)"}
{"(10,fish)","(4,four)"}
{"(10,fish)","(5,)"}

query T
SELECT array_append(ARRAY[ROW(1,2)], NULL::record)
----
{"(1,2)",NULL}

query T rowsort
SELECT array_append(NULL::record[], (k,v)) FROM kv
----
{"(1,one)"}
{"(2,two)"}
{"(3,three)"}
{"(4,four)"}
{"(5,)"}

query T
SELECT array_append(NULL::record[], NULL::record)
----
{NULL}

query T rowsort
SELECT array_prepend((k,v), ARRAY[ROW(10,'fish'), ROW(11,'zebra')]) FROM kv
----
{"(1,one)","(10,fish)","(11,zebra)"}
{"(2,two)","(10,fish)","(11,zebra)"}
{"(3,three)","(10,fish)","(11,zebra)"}
{"(4,four)","(10,fish)","(11,zebra)"}
{"(5,)","(10,fish)","(11,zebra)"}

query T
SELECT array_prepend(NULL::record, ARRAY[ROW(10,'fish'), ROW(11,'zebra')])
----
{NULL,"(10,fish)","(11,zebra)"}

query T
SELECT array_prepend(ROW(1,2), NULL::record[])
----
{"(1,2)"}

query T
SELECT array_prepend(NULL::record, NULL::record[])
----
{NULL}

query T
SELECT array_remove(ARRAY[ROW(1,'cat'), ROW(10,'fish'), ROW(11,'zebra')], ROW(10,'fish'))
----
{"(1,cat)","(11,zebra)"}

query T
SELECT array_remove(ARRAY[ROW(1,'cat'), ROW(10,'fish'), ROW(11,'zebra'), NULL], NULL::record)
----
{"(1,cat)","(10,fish)","(11,zebra)"}

query T
SELECT array_remove(NULL::record[], NULL::record)
----
NULL

query T
SELECT array_replace(ARRAY[ROW(1,'cat'), ROW(10,'fish'), ROW(11,'zebra')], ROW(10,'fish'), ROW(2,'dog'))
----
{"(1,cat)","(2,dog)","(11,zebra)"}

query TT
SELECT array_replace(ARRAY[ROW(1,'cat'), NULL, ROW(11,'zebra')], NULL::record, ROW(2,'dog')), array_replace(NULL::record[], ROW(10,'fish'), ROW(2,'dog'))
----
{"(1,cat)","(2,dog)","(11,zebra)"} NULL

# ARRAY_POSITION function

query I
SELECT array_position(ARRAY[ROW(1,'cat'), ROW(10,'fish'), ROW(11,'zebra')], ROW(11,'zebra'))
----
3

query I
SELECT array_position(ARRAY[ROW(1,'cat'), ROW(10,'fish'), ROW(11,'zebra')], ROW(33,'hippo'))
----
NULL

query I
SELECT array_position(NULL::record[], ROW(33,'hippo'))
----
NULL

# ARRAY_POSITIONS function

query T
SELECT array_positions(ARRAY[ROW(1,'cat'), ROW(11,'zebra'), ROW(10,'fish'), ROW(11,'zebra')], ROW(11,'zebra'))
----
{2,4}

query T
SELECT array_positions(ARRAY[ROW(1,'cat'), ROW(11,'zebra'), ROW(10,'fish'), ROW(11,'zebra')], ROW(33,'hippo'))
----
{}

query T
SELECT array_positions(NULL::record[], ROW(33,'hippo'))
----
NULL
34 changes: 32 additions & 2 deletions pkg/sql/opt/exec/execbuilder/testdata/subquery_correlated
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,38 @@ CREATE TABLE o (o_id INT PRIMARY KEY, c_id INT, ship TEXT);

# We can't decorrelate cases which don't use a scalar type in the
# ARRAY(...) operator.
statement error can't execute a correlated ARRAY\(...\) over tuple\{int, string\}
SELECT
query T
EXPLAIN SELECT
c_id,
ARRAY(SELECT (o_id, ship) FROM o WHERE o.c_id = c.c_id ORDER BY o_id)
FROM c ORDER BY c_id
----
distribution: local
vectorized: true
·
• sort
│ order: +c_id
│
└── • render
│
└── • group
│ group by: c_id
│
└── • sort
│ order: +o_id
│
└── • hash join (left outer)
│ equality: (c_id) = (c_id)
│ left cols are key
│
├── • scan
│ missing stats
│ table: c@primary
│ spans: FULL SCAN
│
└── • render
│
└── • scan
missing stats
table: o@primary
spans: FULL SCAN
1 change: 1 addition & 0 deletions pkg/sql/opt/optbuilder/testdata/aggregate
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,7 @@ array_agg(time) -> time[]
array_agg(timetz) -> timetz[]
array_agg(jsonb) -> jsonb[]
array_agg(varbit) -> varbit[]
array_agg(tuple) -> tuple[]
array_agg(bool) -> bool[]

# With an explicit cast, this works as expected.
Expand Down
16 changes: 15 additions & 1 deletion pkg/sql/opt/optbuilder/testdata/scalar
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1261,7 +1261,21 @@ project
build
SELECT ARRAY(SELECT (y, 2) FROM u ORDER BY x) FROM v
----
error (0A000): unimplemented: can't execute a correlated ARRAY(...) over tuple{int[], int}
project
├── columns: array:10
├── scan v
│ └── columns: y:1 v.rowid:2!null v.crdb_internal_mvcc_timestamp:3 v.tableoid:4
└── projections
└── array-flatten col=9 [as=array:10]
└── sort
├── columns: "?column?":9 [hidden: x:5]
├── ordering: +5
└── project
├── columns: "?column?":9 x:5
├── scan u
│ └── columns: x:5 u.rowid:6!null u.crdb_internal_mvcc_timestamp:7 u.tableoid:8
└── projections
└── (y:1, 2) [as="?column?":9]

build
SELECT ARRAY(SELECT y FROM u ORDER BY x) FROM v
Expand Down
6 changes: 5 additions & 1 deletion pkg/sql/sem/builtins/builtins.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7051,12 +7051,16 @@ var similarOverloads = []tree.Overload{
}

func arrayBuiltin(impl func(*types.T) tree.Overload) builtinDefinition {
overloads := make([]tree.Overload, 0, len(types.Scalar))
overloads := make([]tree.Overload, 0, len(types.Scalar)+1)
for _, typ := range types.Scalar {
if ok, _ := types.IsValidArrayElementType(typ); ok {
overloads = append(overloads, impl(typ))
}
}
// Prevent usage in DistSQL because it cannot handle arrays of untyped tuples.
tupleOverload := impl(types.AnyTuple)
tupleOverload.DistsqlBlocklist = true
overloads = append(overloads, tupleOverload)
return builtinDefinition{
props: tree.FunctionProperties{Category: categoryArray},
overloads: overloads,
Expand Down
2 changes: 1 addition & 1 deletion pkg/sql/sem/tree/expr.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1460,7 +1460,7 @@ func (node *FuncExpr) IsWindowFunctionApplication() bool {

// IsDistSQLBlocklist returns whether the function is not supported by DistSQL.
func (node *FuncExpr) IsDistSQLBlocklist() bool {
return node.fnProps != nil && node.fnProps.DistsqlBlocklist
return (node.fn != nil && node.fn.DistsqlBlocklist) || (node.fnProps != nil && node.fnProps.DistsqlBlocklist)
}

// CanHandleNulls returns whether or not the function can handle null
Expand Down
5 changes: 5 additions & 0 deletions pkg/sql/sem/tree/overload.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,11 @@ type Overload struct {
// Oid is the cached oidHasher.BuiltinOid result for this Overload. It's
// populated at init-time.
Oid oid.Oid

// DistsqlBlocklist is set to true when a function cannot be evaluated in
// DistSQL. One example is when the type information for function arguments
// cannot be recovered.
DistsqlBlocklist bool
}

// params implements the overloadImpl interface.
Expand Down