Skip to content

Commit

Permalink
feat: add support for user defined TVFs (#1278)
Browse files Browse the repository at this point in the history
Allow BigQuery client library users to define table-valued functions (TVFs) that return table data.
  • Loading branch information
stephaniewang526 committed May 4, 2021
1 parent a58dd7c commit 89958e9
Show file tree
Hide file tree
Showing 7 changed files with 218 additions and 1 deletion.
2 changes: 1 addition & 1 deletion google-cloud-bigquery/clirr-ignored-differences.xml
Expand Up @@ -5,6 +5,6 @@
<difference>
<differenceType>7013</differenceType>
<className>com/google/cloud/bigquery/RoutineInfo$Builder</className>
<method>com.google.cloud.bigquery.RoutineInfo$Builder setDeterminismLevel(java.lang.String)</method>
<method>com.google.cloud.bigquery.RoutineInfo$Builder setReturnTableType(com.google.cloud.bigquery.StandardSQLTableType)</method>
</difference>
</differences>
Expand Up @@ -111,6 +111,12 @@ public Builder setReturnType(StandardSQLDataType returnType) {
return this;
}

@Override
public Builder setReturnTableType(StandardSQLTableType returnTableType) {
infoBuilder.setReturnTableType(returnTableType);
return this;
}

@Override
public Builder setImportedLibraries(List<String> libraries) {
infoBuilder.setImportedLibraries(libraries);
Expand Down
Expand Up @@ -67,6 +67,7 @@ public Routine apply(RoutineInfo routineInfo) {
private final String language;
private final List<RoutineArgument> argumentList;
private final StandardSQLDataType returnType;
private final StandardSQLTableType returnTableType;
private final List<String> importedLibrariesList;
private final String body;

Expand Down Expand Up @@ -113,6 +114,9 @@ public abstract static class Builder {
*/
public abstract Builder setReturnType(StandardSQLDataType returnType);

/** Optional. Set only if Routine is a "TABLE_VALUED_FUNCTION". */
public abstract Builder setReturnTableType(StandardSQLTableType returnTableType);

/**
* Optional. If language = "JAVASCRIPT", this field stores the path of the imported JAVASCRIPT
* libraries as a list of gs:// URLs.
Expand Down Expand Up @@ -159,6 +163,7 @@ static class BuilderImpl extends Builder {
private String language;
private List<RoutineArgument> argumentList;
private StandardSQLDataType returnType;
private StandardSQLTableType returnTableType;
private List<String> importedLibrariesList;
private String body;

Expand All @@ -175,6 +180,7 @@ static class BuilderImpl extends Builder {
this.language = routineInfo.language;
this.argumentList = routineInfo.argumentList;
this.returnType = routineInfo.returnType;
this.returnTableType = routineInfo.returnTableType;
this.importedLibrariesList = routineInfo.importedLibrariesList;
this.body = routineInfo.body;
}
Expand All @@ -195,6 +201,9 @@ static class BuilderImpl extends Builder {
if (routinePb.getReturnType() != null) {
this.returnType = StandardSQLDataType.fromPb(routinePb.getReturnType());
}
if (routinePb.getReturnTableType() != null) {
this.returnTableType = StandardSQLTableType.fromPb(routinePb.getReturnTableType());
}
if (routinePb.getImportedLibraries() == null) {
this.importedLibrariesList = Collections.emptyList();
} else {
Expand Down Expand Up @@ -263,6 +272,12 @@ public Builder setReturnType(StandardSQLDataType returnType) {
return this;
}

@Override
public Builder setReturnTableType(StandardSQLTableType returnTableType) {
this.returnTableType = returnTableType;
return this;
}

@Override
public Builder setImportedLibraries(List<String> importedLibrariesList) {
this.importedLibrariesList = importedLibrariesList;
Expand Down Expand Up @@ -292,6 +307,7 @@ public RoutineInfo build() {
this.language = builder.language;
this.argumentList = builder.argumentList;
this.returnType = builder.returnType;
this.returnTableType = builder.returnTableType;
this.importedLibrariesList = builder.importedLibrariesList;
this.body = builder.body;
}
Expand Down Expand Up @@ -350,6 +366,11 @@ public StandardSQLDataType getReturnType() {
return returnType;
}

/** If specified, returns the table type returned from the routine. */
public StandardSQLTableType getReturnTableType() {
return returnTableType;
}

/**
* Returns the list of imported libraries for the routine. Only relevant for routines implemented
* using the JAVASCRIPT language.
Expand Down Expand Up @@ -381,6 +402,7 @@ public String toString() {
.add("language", language)
.add("arguments", argumentList)
.add("returnType", returnType)
.add("returnTableType", returnTableType)
.add("importedLibrariesList", importedLibrariesList)
.add("body", body)
.toString();
Expand All @@ -399,6 +421,7 @@ public int hashCode() {
language,
argumentList,
returnType,
returnTableType,
importedLibrariesList,
body);
}
Expand Down Expand Up @@ -448,6 +471,9 @@ Routine toPb() {
if (getReturnType() != null) {
routinePb.setReturnType(getReturnType().toPb());
}
if (getReturnTableType() != null) {
routinePb.setReturnTableType(getReturnTableType().toPb());
}
return routinePb;
}

Expand Down
@@ -0,0 +1,70 @@
/*
* Copyright 2021 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.google.cloud.bigquery;

import com.google.api.services.bigquery.model.StandardSqlTableType;
import com.google.auto.value.AutoValue;
import com.google.common.collect.Lists;
import java.io.Serializable;
import java.util.List;

/** Represents Standard SQL table type information. */
@AutoValue
public abstract class StandardSQLTableType implements Serializable {

@AutoValue.Builder
public abstract static class Builder {

/** Sets the columns in this table type. */
public abstract Builder setColumns(List<StandardSQLField> columns);

/** Creates a {@code StandardSQLTableType} object. */
public abstract StandardSQLTableType build();
}

/** Returns the columns in this table type. */
public abstract List<StandardSQLField> getColumns();

public abstract Builder toBuilder();

/** Returns a builder for a {@code StandardSQLTableType} object. */
public static Builder newBuilder() {
return new AutoValue_StandardSQLTableType.Builder();
}

/** Returns a builder for a {@code StandardSQLTableType} object with the specified columns. */
public static StandardSQLTableType.Builder newBuilder(List<StandardSQLField> columns) {
return newBuilder().setColumns(columns);
}

static StandardSQLTableType fromPb(
com.google.api.services.bigquery.model.StandardSqlTableType tableTypePb) {
StandardSQLTableType.Builder builder = newBuilder();
if (tableTypePb.getColumns() != null) {
builder.setColumns(
Lists.transform(tableTypePb.getColumns(), StandardSQLField.FROM_PB_FUNCTION));
}
return builder.build();
}

StandardSqlTableType toPb() {
StandardSqlTableType tableType = new StandardSqlTableType();
if (getColumns() != null) {
tableType.setColumns(Lists.transform(getColumns(), StandardSQLField.TO_PB_FUNCTION));
}
return tableType;
}
}
Expand Up @@ -38,9 +38,11 @@
public class RoutineTest {

private static final RoutineId ROUTINE_ID = RoutineId.of("dataset", "routine");
private static final RoutineId ROUTINE_ID_TVF = RoutineId.of("dataset", "tvf_routine");
private static final String DETERMINISM_LEVEL = "DETERMINISTIC";
private static final String ETAG = "etag";
private static final String ROUTINE_TYPE = "SCALAR_FUNCTION";
private static final String ROUTINE_TYPE_TVF = "TABLE_VALUED_FUNCTION";
private static final Long CREATION_TIME = 10L;
private static final Long LAST_MODIFIED_TIME = 20L;
private static final String LANGUAGE = "SQL";
Expand All @@ -56,6 +58,18 @@ public class RoutineTest {
private static final StandardSQLDataType RETURN_TYPE =
StandardSQLDataType.newBuilder("FLOAT64").build();

private static final StandardSQLField COLUMN_1 =
StandardSQLField.newBuilder("COLUMN_1", StandardSQLDataType.newBuilder("STRING").build())
.build();
private static final StandardSQLField COLUMN_2 =
StandardSQLField.newBuilder("COLUMN_2", StandardSQLDataType.newBuilder("FLOAT64").build())
.build();

private static final List<StandardSQLField> COLUMN_LIST = ImmutableList.of(COLUMN_1, COLUMN_2);

private static final StandardSQLTableType RETURN_TABLE_TYPE =
StandardSQLTableType.newBuilder(COLUMN_LIST).build();

private static final List<String> IMPORTED_LIBRARIES =
ImmutableList.of("gs://foo", "gs://bar", "gs://baz");

Expand All @@ -75,11 +89,19 @@ public class RoutineTest {
.setBody(BODY)
.build();

private static final RoutineInfo ROUTINE_INFO_TVF =
RoutineInfo.newBuilder(ROUTINE_ID_TVF)
.setBody(BODY)
.setRoutineType(ROUTINE_TYPE_TVF)
.setReturnTableType(RETURN_TABLE_TYPE)
.build();

@Rule public MockitoRule rule;

private BigQuery bigquery;
private BigQueryOptions mockOptions;
private Routine expectedRoutine;
private Routine expectedRoutineTvf;
private Routine routine;

@Before
Expand All @@ -88,6 +110,7 @@ public void setUp() {
mockOptions = mock(BigQueryOptions.class);
when(bigquery.getOptions()).thenReturn(mockOptions);
expectedRoutine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
expectedRoutineTvf = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO_TVF));
routine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
}

Expand All @@ -114,6 +137,7 @@ public void testBuilder() {
@Test
public void testToBuilder() {
compareRoutineInfo(expectedRoutine, expectedRoutine.toBuilder().build());
compareRoutineInfo(expectedRoutineTvf, expectedRoutineTvf.toBuilder().build());
}

@Test
Expand Down Expand Up @@ -200,6 +224,7 @@ public void compareRoutineInfo(RoutineInfo expected, RoutineInfo value) {
assertEquals(expected.getLanguage(), value.getLanguage());
assertEquals(expected.getArguments(), value.getArguments());
assertEquals(expected.getReturnType(), value.getReturnType());
assertEquals(expected.getReturnTableType(), value.getReturnTableType());
assertEquals(expected.getImportedLibraries(), value.getImportedLibraries());
assertEquals(expected.getBody(), value.getBody());
assertEquals(expected.hashCode(), value.hashCode());
Expand Down
@@ -0,0 +1,60 @@
/*
* Copyright 2021 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.cloud.bigquery;

import static org.junit.Assert.*;

import com.google.common.collect.ImmutableList;
import java.util.List;
import org.junit.Test;

public class StandardSQLTableTypeTest {

private static final StandardSQLField COLUMN_1 =
StandardSQLField.newBuilder("COLUMN_1", StandardSQLDataType.newBuilder("STRING").build())
.build();
private static final StandardSQLField COLUMN_2 =
StandardSQLField.newBuilder("COLUMN_2", StandardSQLDataType.newBuilder("FLOAT64").build())
.build();

private static final List<StandardSQLField> COLUMN_LIST = ImmutableList.of(COLUMN_1, COLUMN_2);
private static final StandardSQLTableType TABLE_TYPE =
StandardSQLTableType.newBuilder(COLUMN_LIST).build();

@Test
public void testToBuilder() {
compareStandardSQLTableType(TABLE_TYPE, TABLE_TYPE.toBuilder().build());
}

@Test
public void testBuilder() {
assertEquals(COLUMN_1, TABLE_TYPE.getColumns().get(0));
assertEquals(COLUMN_2, TABLE_TYPE.getColumns().get(1));
}

@Test
public void testToAndFromPb() {
compareStandardSQLTableType(TABLE_TYPE, StandardSQLTableType.fromPb(TABLE_TYPE.toPb()));
}

private void compareStandardSQLTableType(
StandardSQLTableType expected, StandardSQLTableType value) {
assertEquals(expected, value);
assertEquals(expected.getColumns(), value.getColumns());
assertEquals(expected.hashCode(), value.hashCode());
}
}
Expand Up @@ -87,6 +87,8 @@
import com.google.cloud.bigquery.RoutineInfo;
import com.google.cloud.bigquery.Schema;
import com.google.cloud.bigquery.StandardSQLDataType;
import com.google.cloud.bigquery.StandardSQLField;
import com.google.cloud.bigquery.StandardSQLTableType;
import com.google.cloud.bigquery.StandardTableDefinition;
import com.google.cloud.bigquery.Table;
import com.google.cloud.bigquery.TableDataWriteChannel;
Expand Down Expand Up @@ -1676,6 +1678,34 @@ public void testRoutineAPICreationJavascriptUDF() {
assertEquals(routine.getReturnType(), StandardSQLDataType.newBuilder("STRING").build());
}

@Test
public void testRoutineAPICreationTVF() {
String routineName = RemoteBigQueryHelper.generateRoutineName();
RoutineId routineId = RoutineId.of(ROUTINE_DATASET, routineName);
List<StandardSQLField> columns =
ImmutableList.of(
StandardSQLField.newBuilder("x", StandardSQLDataType.newBuilder("INT64").build())
.build());
StandardSQLTableType returnTableType = StandardSQLTableType.newBuilder(columns).build();
RoutineInfo routineInfo =
RoutineInfo.newBuilder(routineId)
.setRoutineType("TABLE_VALUED_FUNCTION")
.setLanguage("SQL")
.setArguments(
ImmutableList.of(
RoutineArgument.newBuilder()
.setName("filter")
.setDataType(StandardSQLDataType.newBuilder("INT64").build())
.build()))
.setReturnTableType(returnTableType)
.setBody("SELECT x FROM UNNEST([1,2,3]) x WHERE x = filter")
.build();
Routine routine = bigquery.create(routineInfo);
assertNotNull(routine);
assertEquals(routine.getRoutineType(), "TABLE_VALUED_FUNCTION");
assertEquals(routine.getReturnTableType(), returnTableType);
}

@Test
public void testAuthorizeRoutine() {
String routineName = RemoteBigQueryHelper.generateRoutineName();
Expand Down

0 comments on commit 89958e9

Please sign in to comment.