Skip to content

Commit 89958e9

Browse files
stephaniewang526stephaniewang526
authored
feat: add support for user defined TVFs (#1278)
Allow BigQuery client library users to define table-valued functions (TVFs) that return table data.
1 parent a58dd7c commit 89958e9

File tree

7 files changed

+218
-1
lines changed

7 files changed

+218
-1
lines changed

‎google-cloud-bigquery/clirr-ignored-differences.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,6 @@
55
<difference>
66
<differenceType>7013</differenceType>
77
<className>com/google/cloud/bigquery/RoutineInfo$Builder</className>
8-
<method>com.google.cloud.bigquery.RoutineInfo$Builder setDeterminismLevel(java.lang.String)</method>
8+
<method>com.google.cloud.bigquery.RoutineInfo$Builder setReturnTableType(com.google.cloud.bigquery.StandardSQLTableType)</method>
99
</difference>
1010
</differences>

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Routine.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,12 @@ public Builder setReturnType(StandardSQLDataType returnType) {
111111
return this;
112112
}
113113

114+
@Override
115+
public Builder setReturnTableType(StandardSQLTableType returnTableType) {
116+
infoBuilder.setReturnTableType(returnTableType);
117+
return this;
118+
}
119+
114120
@Override
115121
public Builder setImportedLibraries(List<String> libraries) {
116122
infoBuilder.setImportedLibraries(libraries);

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ public Routine apply(RoutineInfo routineInfo) {
6767
private final String language;
6868
private final List<RoutineArgument> argumentList;
6969
private final StandardSQLDataType returnType;
70+
private final StandardSQLTableType returnTableType;
7071
private final List<String> importedLibrariesList;
7172
private final String body;
7273

@@ -113,6 +114,9 @@ public abstract static class Builder {
113114
*/
114115
public abstract Builder setReturnType(StandardSQLDataType returnType);
115116

117+
/** Optional. Set only if Routine is a "TABLE_VALUED_FUNCTION". */
118+
public abstract Builder setReturnTableType(StandardSQLTableType returnTableType);
119+
116120
/**
117121
* Optional. If language = "JAVASCRIPT", this field stores the path of the imported JAVASCRIPT
118122
* libraries as a list of gs:// URLs.
@@ -159,6 +163,7 @@ static class BuilderImpl extends Builder {
159163
private String language;
160164
private List<RoutineArgument> argumentList;
161165
private StandardSQLDataType returnType;
166+
private StandardSQLTableType returnTableType;
162167
private List<String> importedLibrariesList;
163168
private String body;
164169

@@ -175,6 +180,7 @@ static class BuilderImpl extends Builder {
175180
this.language = routineInfo.language;
176181
this.argumentList = routineInfo.argumentList;
177182
this.returnType = routineInfo.returnType;
183+
this.returnTableType = routineInfo.returnTableType;
178184
this.importedLibrariesList = routineInfo.importedLibrariesList;
179185
this.body = routineInfo.body;
180186
}
@@ -195,6 +201,9 @@ static class BuilderImpl extends Builder {
195201
if (routinePb.getReturnType() != null) {
196202
this.returnType = StandardSQLDataType.fromPb(routinePb.getReturnType());
197203
}
204+
if (routinePb.getReturnTableType() != null) {
205+
this.returnTableType = StandardSQLTableType.fromPb(routinePb.getReturnTableType());
206+
}
198207
if (routinePb.getImportedLibraries() == null) {
199208
this.importedLibrariesList = Collections.emptyList();
200209
} else {
@@ -263,6 +272,12 @@ public Builder setReturnType(StandardSQLDataType returnType) {
263272
return this;
264273
}
265274

275+
@Override
276+
public Builder setReturnTableType(StandardSQLTableType returnTableType) {
277+
this.returnTableType = returnTableType;
278+
return this;
279+
}
280+
266281
@Override
267282
public Builder setImportedLibraries(List<String> importedLibrariesList) {
268283
this.importedLibrariesList = importedLibrariesList;
@@ -292,6 +307,7 @@ public RoutineInfo build() {
292307
this.language = builder.language;
293308
this.argumentList = builder.argumentList;
294309
this.returnType = builder.returnType;
310+
this.returnTableType = builder.returnTableType;
295311
this.importedLibrariesList = builder.importedLibrariesList;
296312
this.body = builder.body;
297313
}
@@ -350,6 +366,11 @@ public StandardSQLDataType getReturnType() {
350366
return returnType;
351367
}
352368

369+
/** If specified, returns the table type returned from the routine. */
370+
public StandardSQLTableType getReturnTableType() {
371+
return returnTableType;
372+
}
373+
353374
/**
354375
* Returns the list of imported libraries for the routine. Only relevant for routines implemented
355376
* using the JAVASCRIPT language.
@@ -381,6 +402,7 @@ public String toString() {
381402
.add("language", language)
382403
.add("arguments", argumentList)
383404
.add("returnType", returnType)
405+
.add("returnTableType", returnTableType)
384406
.add("importedLibrariesList", importedLibrariesList)
385407
.add("body", body)
386408
.toString();
@@ -399,6 +421,7 @@ public int hashCode() {
399421
language,
400422
argumentList,
401423
returnType,
424+
returnTableType,
402425
importedLibrariesList,
403426
body);
404427
}
@@ -448,6 +471,9 @@ Routine toPb() {
448471
if (getReturnType() != null) {
449472
routinePb.setReturnType(getReturnType().toPb());
450473
}
474+
if (getReturnTableType() != null) {
475+
routinePb.setReturnTableType(getReturnTableType().toPb());
476+
}
451477
return routinePb;
452478
}
453479

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/StandardSQLTableType.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigquery;
17+
18+
import com.google.api.services.bigquery.model.StandardSqlTableType;
19+
import com.google.auto.value.AutoValue;
20+
import com.google.common.collect.Lists;
21+
import java.io.Serializable;
22+
import java.util.List;
23+
24+
/** Represents Standard SQL table type information. */
25+
@AutoValue
26+
public abstract class StandardSQLTableType implements Serializable {
27+
28+
@AutoValue.Builder
29+
public abstract static class Builder {
30+
31+
/** Sets the columns in this table type. */
32+
public abstract Builder setColumns(List<StandardSQLField> columns);
33+
34+
/** Creates a {@code StandardSQLTableType} object. */
35+
public abstract StandardSQLTableType build();
36+
}
37+
38+
/** Returns the columns in this table type. */
39+
public abstract List<StandardSQLField> getColumns();
40+
41+
public abstract Builder toBuilder();
42+
43+
/** Returns a builder for a {@code StandardSQLTableType} object. */
44+
public static Builder newBuilder() {
45+
return new AutoValue_StandardSQLTableType.Builder();
46+
}
47+
48+
/** Returns a builder for a {@code StandardSQLTableType} object with the specified columns. */
49+
public static StandardSQLTableType.Builder newBuilder(List<StandardSQLField> columns) {
50+
return newBuilder().setColumns(columns);
51+
}
52+
53+
static StandardSQLTableType fromPb(
54+
com.google.api.services.bigquery.model.StandardSqlTableType tableTypePb) {
55+
StandardSQLTableType.Builder builder = newBuilder();
56+
if (tableTypePb.getColumns() != null) {
57+
builder.setColumns(
58+
Lists.transform(tableTypePb.getColumns(), StandardSQLField.FROM_PB_FUNCTION));
59+
}
60+
return builder.build();
61+
}
62+
63+
StandardSqlTableType toPb() {
64+
StandardSqlTableType tableType = new StandardSqlTableType();
65+
if (getColumns() != null) {
66+
tableType.setColumns(Lists.transform(getColumns(), StandardSQLField.TO_PB_FUNCTION));
67+
}
68+
return tableType;
69+
}
70+
}

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,11 @@
3838
public class RoutineTest {
3939

4040
private static final RoutineId ROUTINE_ID = RoutineId.of("dataset", "routine");
41+
private static final RoutineId ROUTINE_ID_TVF = RoutineId.of("dataset", "tvf_routine");
4142
private static final String DETERMINISM_LEVEL = "DETERMINISTIC";
4243
private static final String ETAG = "etag";
4344
private static final String ROUTINE_TYPE = "SCALAR_FUNCTION";
45+
private static final String ROUTINE_TYPE_TVF = "TABLE_VALUED_FUNCTION";
4446
private static final Long CREATION_TIME = 10L;
4547
private static final Long LAST_MODIFIED_TIME = 20L;
4648
private static final String LANGUAGE = "SQL";
@@ -56,6 +58,18 @@ public class RoutineTest {
5658
private static final StandardSQLDataType RETURN_TYPE =
5759
StandardSQLDataType.newBuilder("FLOAT64").build();
5860

61+
private static final StandardSQLField COLUMN_1 =
62+
StandardSQLField.newBuilder("COLUMN_1", StandardSQLDataType.newBuilder("STRING").build())
63+
.build();
64+
private static final StandardSQLField COLUMN_2 =
65+
StandardSQLField.newBuilder("COLUMN_2", StandardSQLDataType.newBuilder("FLOAT64").build())
66+
.build();
67+
68+
private static final List<StandardSQLField> COLUMN_LIST = ImmutableList.of(COLUMN_1, COLUMN_2);
69+
70+
private static final StandardSQLTableType RETURN_TABLE_TYPE =
71+
StandardSQLTableType.newBuilder(COLUMN_LIST).build();
72+
5973
private static final List<String> IMPORTED_LIBRARIES =
6074
ImmutableList.of("gs://foo", "gs://bar", "gs://baz");
6175

@@ -75,11 +89,19 @@ public class RoutineTest {
7589
.setBody(BODY)
7690
.build();
7791

92+
private static final RoutineInfo ROUTINE_INFO_TVF =
93+
RoutineInfo.newBuilder(ROUTINE_ID_TVF)
94+
.setBody(BODY)
95+
.setRoutineType(ROUTINE_TYPE_TVF)
96+
.setReturnTableType(RETURN_TABLE_TYPE)
97+
.build();
98+
7899
@Rule public MockitoRule rule;
79100

80101
private BigQuery bigquery;
81102
private BigQueryOptions mockOptions;
82103
private Routine expectedRoutine;
104+
private Routine expectedRoutineTvf;
83105
private Routine routine;
84106

85107
@Before
@@ -88,6 +110,7 @@ public void setUp() {
88110
mockOptions = mock(BigQueryOptions.class);
89111
when(bigquery.getOptions()).thenReturn(mockOptions);
90112
expectedRoutine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
113+
expectedRoutineTvf = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO_TVF));
91114
routine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
92115
}
93116

@@ -114,6 +137,7 @@ public void testBuilder() {
114137
@Test
115138
public void testToBuilder() {
116139
compareRoutineInfo(expectedRoutine, expectedRoutine.toBuilder().build());
140+
compareRoutineInfo(expectedRoutineTvf, expectedRoutineTvf.toBuilder().build());
117141
}
118142

119143
@Test
@@ -200,6 +224,7 @@ public void compareRoutineInfo(RoutineInfo expected, RoutineInfo value) {
200224
assertEquals(expected.getLanguage(), value.getLanguage());
201225
assertEquals(expected.getArguments(), value.getArguments());
202226
assertEquals(expected.getReturnType(), value.getReturnType());
227+
assertEquals(expected.getReturnTableType(), value.getReturnTableType());
203228
assertEquals(expected.getImportedLibraries(), value.getImportedLibraries());
204229
assertEquals(expected.getBody(), value.getBody());
205230
assertEquals(expected.hashCode(), value.hashCode());

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/StandardSQLTableTypeTest.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import static org.junit.Assert.*;
20+
21+
import com.google.common.collect.ImmutableList;
22+
import java.util.List;
23+
import org.junit.Test;
24+
25+
public class StandardSQLTableTypeTest {
26+
27+
private static final StandardSQLField COLUMN_1 =
28+
StandardSQLField.newBuilder("COLUMN_1", StandardSQLDataType.newBuilder("STRING").build())
29+
.build();
30+
private static final StandardSQLField COLUMN_2 =
31+
StandardSQLField.newBuilder("COLUMN_2", StandardSQLDataType.newBuilder("FLOAT64").build())
32+
.build();
33+
34+
private static final List<StandardSQLField> COLUMN_LIST = ImmutableList.of(COLUMN_1, COLUMN_2);
35+
private static final StandardSQLTableType TABLE_TYPE =
36+
StandardSQLTableType.newBuilder(COLUMN_LIST).build();
37+
38+
@Test
39+
public void testToBuilder() {
40+
compareStandardSQLTableType(TABLE_TYPE, TABLE_TYPE.toBuilder().build());
41+
}
42+
43+
@Test
44+
public void testBuilder() {
45+
assertEquals(COLUMN_1, TABLE_TYPE.getColumns().get(0));
46+
assertEquals(COLUMN_2, TABLE_TYPE.getColumns().get(1));
47+
}
48+
49+
@Test
50+
public void testToAndFromPb() {
51+
compareStandardSQLTableType(TABLE_TYPE, StandardSQLTableType.fromPb(TABLE_TYPE.toPb()));
52+
}
53+
54+
private void compareStandardSQLTableType(
55+
StandardSQLTableType expected, StandardSQLTableType value) {
56+
assertEquals(expected, value);
57+
assertEquals(expected.getColumns(), value.getColumns());
58+
assertEquals(expected.hashCode(), value.hashCode());
59+
}
60+
}

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,8 @@
8787
import com.google.cloud.bigquery.RoutineInfo;
8888
import com.google.cloud.bigquery.Schema;
8989
import com.google.cloud.bigquery.StandardSQLDataType;
90+
import com.google.cloud.bigquery.StandardSQLField;
91+
import com.google.cloud.bigquery.StandardSQLTableType;
9092
import com.google.cloud.bigquery.StandardTableDefinition;
9193
import com.google.cloud.bigquery.Table;
9294
import com.google.cloud.bigquery.TableDataWriteChannel;
@@ -1676,6 +1678,34 @@ public void testRoutineAPICreationJavascriptUDF() {
16761678
assertEquals(routine.getReturnType(), StandardSQLDataType.newBuilder("STRING").build());
16771679
}
16781680

1681+
@Test
1682+
public void testRoutineAPICreationTVF() {
1683+
String routineName = RemoteBigQueryHelper.generateRoutineName();
1684+
RoutineId routineId = RoutineId.of(ROUTINE_DATASET, routineName);
1685+
List<StandardSQLField> columns =
1686+
ImmutableList.of(
1687+
StandardSQLField.newBuilder("x", StandardSQLDataType.newBuilder("INT64").build())
1688+
.build());
1689+
StandardSQLTableType returnTableType = StandardSQLTableType.newBuilder(columns).build();
1690+
RoutineInfo routineInfo =
1691+
RoutineInfo.newBuilder(routineId)
1692+
.setRoutineType("TABLE_VALUED_FUNCTION")
1693+
.setLanguage("SQL")
1694+
.setArguments(
1695+
ImmutableList.of(
1696+
RoutineArgument.newBuilder()
1697+
.setName("filter")
1698+
.setDataType(StandardSQLDataType.newBuilder("INT64").build())
1699+
.build()))
1700+
.setReturnTableType(returnTableType)
1701+
.setBody("SELECT x FROM UNNEST([1,2,3]) x WHERE x = filter")
1702+
.build();
1703+
Routine routine = bigquery.create(routineInfo);
1704+
assertNotNull(routine);
1705+
assertEquals(routine.getRoutineType(), "TABLE_VALUED_FUNCTION");
1706+
assertEquals(routine.getReturnTableType(), returnTableType);
1707+
}
1708+
16791709
@Test
16801710
public void testAuthorizeRoutine() {
16811711
String routineName = RemoteBigQueryHelper.generateRoutineName();

0 commit comments

Comments
 (0)