Skip to content

Commit

Permalink
feat: add support for Search statistics (#2787)
Browse files Browse the repository at this point in the history
  • Loading branch information
farhan0102 authored Jul 17, 2023
1 parent 07760b7 commit 344f695
Show file tree
Hide file tree
Showing 5 changed files with 285 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
/*
* Copyright 2023 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.TableReference;
import com.google.auto.value.AutoValue;
import java.io.Serializable;
import javax.annotation.Nullable;

/** Represents Reason of why the index was not used in a SQL search. */
@AutoValue
public abstract class IndexUnusedReason implements Serializable {

@AutoValue.Builder
public abstract static class Builder {

/**
* Specifies the name of the unused search index, if available.
*
* @param indexName indexName or {@code null} for none
*/
public abstract Builder setIndexName(String indexName);

/**
* Specifies the high-level reason for the scenario when no search index was used.
*
* @param code code or {@code null} for none
*/
public abstract Builder setCode(String code);

/**
* Free form human-readable reason for the scenario when no search index was used.
*
* @param message message or {@code null} for none
*/
public abstract Builder setMessage(String message);

/**
* Specifies the base table involved in the reason that no search index was used.
*
* @param tableReference tableReference or {@code null} for none
*/
public abstract Builder setBaseTable(TableReference tableReference);

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

public abstract Builder toBuilder();

public static Builder newBuilder() {
return new AutoValue_IndexUnusedReason.Builder();
}

/**
* Returns the name of the unused search index, if available.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract String getIndexName();

/**
* Returns the high-level reason for the scenario when no search index was used.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract String getCode();

/**
* Returns free form human-readable reason for the scenario when no search index was used.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract String getMessage();

/**
* Returns the base table involved in the reason that no search index was used.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract TableReference getBaseTable();

com.google.api.services.bigquery.model.IndexUnusedReason toPb() {
com.google.api.services.bigquery.model.IndexUnusedReason indexUnusedReason =
new com.google.api.services.bigquery.model.IndexUnusedReason();
if (getIndexName() != null) {
indexUnusedReason.setIndexName(indexUnusedReason.getIndexName());
}
if (getCode() != null) {
indexUnusedReason.setCode(indexUnusedReason.getCode());
}
if (getMessage() != null) {
indexUnusedReason.setMessage(indexUnusedReason.getMessage());
}
if (getBaseTable() != null) {
indexUnusedReason.setBaseTable(indexUnusedReason.getBaseTable());
}
return indexUnusedReason;
}

static IndexUnusedReason fromPb(
com.google.api.services.bigquery.model.IndexUnusedReason indexUnusedReason) {
Builder builder = newBuilder();
if (indexUnusedReason.getIndexName() != null) {
builder.setIndexName(indexUnusedReason.getIndexName());
}
if (indexUnusedReason.getCode() != null) {
builder.setCode(indexUnusedReason.getCode());
}
if (indexUnusedReason.getMessage() != null) {
builder.setMessage(indexUnusedReason.getMessage());
}
if (indexUnusedReason.getBaseTable() != null) {
builder.setBaseTable(indexUnusedReason.getBaseTable());
}
return builder.build();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -340,6 +340,7 @@ public static class QueryStatistics extends JobStatistics {
private final List<QueryStage> queryPlan;
private final List<TimelineSample> timeline;
private final Schema schema;
private final SearchStats searchStats;
private final List<QueryParameter> queryParameters;

/**
Expand Down Expand Up @@ -424,6 +425,7 @@ static final class Builder extends JobStatistics.Builder<QueryStatistics, Builde
private List<TimelineSample> timeline;
private Schema schema;
private List<QueryParameter> queryParameters;
private SearchStats searchStats;

private Builder() {}

Expand Down Expand Up @@ -471,6 +473,9 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
if (statisticsPb.getQuery().getSchema() != null) {
this.schema = Schema.fromPb(statisticsPb.getQuery().getSchema());
}
if (statisticsPb.getQuery().getSearchStatistics() != null) {
this.searchStats = SearchStats.fromPb(statisticsPb.getQuery().getSearchStatistics());
}
if (statisticsPb.getQuery().getDmlStats() != null) {
this.dmlStats = DmlStats.fromPb(statisticsPb.getQuery().getDmlStats());
}
Expand Down Expand Up @@ -572,6 +577,11 @@ Builder setSchema(Schema schema) {
return self();
}

Builder setSearchStats(SearchStats searchStats) {
this.searchStats = searchStats;
return self();
}

Builder setQueryParameters(List<QueryParameter> queryParameters) {
this.queryParameters = queryParameters;
return self();
Expand Down Expand Up @@ -603,6 +613,7 @@ private QueryStatistics(Builder builder) {
this.queryPlan = builder.queryPlan;
this.timeline = builder.timeline;
this.schema = builder.schema;
this.searchStats = builder.searchStats;
this.queryParameters = builder.queryParameters;
}

Expand Down Expand Up @@ -724,6 +735,15 @@ public Schema getSchema() {
return schema;
}

/**
* Statistics for a search query. Populated as part of JobStatistics2. Provides information
* about how indexes are used in search queries. If an index is not used, you can retrieve
* debugging information about the reason why.
*/
public SearchStats getSearchStats() {
return searchStats;
}

/**
* Standard SQL only: Returns a list of undeclared query parameters detected during a dry run
* validation.
Expand All @@ -743,6 +763,7 @@ ToStringHelper toStringHelper() {
.add("queryPlan", queryPlan)
.add("timeline", timeline)
.add("schema", schema)
.add("searchStats", searchStats)
.add("queryParameters", queryParameters);
}

Expand All @@ -765,6 +786,7 @@ public final int hashCode() {
totalBytesProcessed,
queryPlan,
schema,
searchStats,
queryParameters);
}

Expand Down Expand Up @@ -807,6 +829,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
if (schema != null) {
queryStatisticsPb.setSchema(schema.toPb());
}
if (searchStats != null) {
queryStatisticsPb.setSearchStatistics(searchStats.toPb());
}
if (queryParameters != null) {
queryStatisticsPb.setUndeclaredQueryParameters(queryParameters);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
/*
* Copyright 2023 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.SearchStatistics;
import com.google.auto.value.AutoValue;
import java.io.Serializable;
import java.util.List;
import java.util.stream.Collectors;
import javax.annotation.Nullable;

/** Represents Search statistics information of a search query. */
@AutoValue
public abstract class SearchStats implements Serializable {

@AutoValue.Builder
public abstract static class Builder {

/**
* Specifies index usage mode for the query.
*
* @param indexUsageMode, has three modes UNUSED, PARTIALLY_USED, and FULLY_USED
*/
public abstract Builder setIndexUsageMode(String indexUsageMode);

/**
* When index_usage_mode is UNUSED or PARTIALLY_USED, this field explains why index was not used
* in all or part of the search query. If index_usage_mode is FULLY_USED, this field is not
* populated.
*
* @param indexUnusedReasons
*/
public abstract Builder setIndexUnusedReasons(List<IndexUnusedReason> indexUnusedReasons);

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

public abstract Builder toBuilder();

public static Builder newBuilder() {
return new AutoValue_SearchStats.Builder();
}

@Nullable
public abstract String getIndexUsageMode();

@Nullable
public abstract List<IndexUnusedReason> getIndexUnusedReasons();

SearchStatistics toPb() {
SearchStatistics searchStatistics = new SearchStatistics();
if (getIndexUsageMode() != null) {
searchStatistics.setIndexUsageMode(getIndexUsageMode());
}
if (getIndexUnusedReasons() != null) {
searchStatistics.setIndexUnusedReason(
getIndexUnusedReasons().stream()
.map(IndexUnusedReason::toPb)
.collect(Collectors.toList()));
}
return searchStatistics;
}

static SearchStats fromPb(SearchStatistics searchStatistics) {
Builder builder = newBuilder();
if (searchStatistics.getIndexUsageMode() != null) {
builder.setIndexUsageMode(searchStatistics.getIndexUsageMode());
}
if (searchStatistics.getIndexUnusedReason() != null) {
builder.setIndexUnusedReasons(
searchStatistics.getIndexUnusedReason().stream()
.map(IndexUnusedReason::fromPb)
.collect(Collectors.toList()));
}
return builder.build();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,9 @@ public class JobStatisticsTest {
ImmutableList.of(TIMELINE_SAMPLE1, TIMELINE_SAMPLE2);
private static final List<QueryStage> QUERY_PLAN = ImmutableList.of(QUERY_STAGE);
private static final Schema SCHEMA = Schema.of(Field.of("column", LegacySQLTypeName.DATETIME));
private static final String UNUSED_INDEX_USAGE_MODE = "UNUSED";
private static final SearchStats SEARCH_STATS =
SearchStats.newBuilder().setIndexUsageMode(UNUSED_INDEX_USAGE_MODE).build();
private static final QueryStatistics QUERY_STATISTICS =
QueryStatistics.newBuilder()
.setCreationTimestamp(CREATION_TIME)
Expand All @@ -182,6 +185,7 @@ public class JobStatisticsTest {
.setQueryPlan(QUERY_PLAN)
.setTimeline(TIMELINE)
.setSchema(SCHEMA)
.setSearchStats(SEARCH_STATS)
.build();
private static final QueryStatistics QUERY_STATISTICS_INCOMPLETE =
QueryStatistics.newBuilder()
Expand All @@ -190,6 +194,7 @@ public class JobStatisticsTest {
.setStartTime(START_TIME)
.setBillingTier(BILLING_TIER)
.setCacheHit(CACHE_HIT)
.setSearchStats(SEARCH_STATS)
.build();
private static final ScriptStackFrame STATEMENT_STACK_FRAME =
ScriptStackFrame.newBuilder()
Expand Down Expand Up @@ -407,6 +412,8 @@ private void compareQueryStatistics(QueryStatistics expected, QueryStatistics va
assertEquals(expected.getQueryPlan(), value.getQueryPlan());
assertEquals(expected.getReferencedTables(), value.getReferencedTables());
assertEquals(expected.getSchema(), value.getSchema());
assertEquals(
expected.getSearchStats().getIndexUsageMode(), value.getSearchStats().getIndexUsageMode());
assertEquals(expected.getStatementType(), value.getStatementType());
assertEquals(expected.getTimeline(), value.getTimeline());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5088,6 +5088,31 @@ public void testQueryJobWithLabels() throws InterruptedException, TimeoutExcepti
}
}

@Test
public void testQueryJobWithSearchReturnsSearchStatistics() throws InterruptedException {
String tableName = "test_query_job_table";
String query =
"SELECT * FROM "
+ TABLE_ID.getTable()
+ " WHERE search(StringField, \"stringValue\")";
TableId destinationTable = TableId.of(DATASET, tableName);
try {
QueryJobConfiguration configuration =
QueryJobConfiguration.newBuilder(query)
.setDefaultDataset(DatasetId.of(DATASET))
.setDestinationTable(destinationTable)
.build();
Job remoteJob = bigquery.create(JobInfo.of(configuration));
remoteJob = remoteJob.waitFor();
assertNull(remoteJob.getStatus().getError());
JobStatistics.QueryStatistics stats = remoteJob.getStatistics();
assertNotNull(stats.getSearchStats());
assertEquals(stats.getSearchStats().getIndexUsageMode(), "UNUSED");
} finally {
bigquery.delete(destinationTable);
}
}

/* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQueryJobWithRangePartitioning() throws InterruptedException {
Expand Down

0 comments on commit 344f695

Please sign in to comment.