feat: add support for Search statistics (#2787)

googleapis · Jul 17, 2023 · 344f695 · 344f695
1 parent 07760b7
commit 344f695
Show file tree

Hide file tree

Showing 5 changed files with 285 additions and 0 deletions.
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/IndexUnusedReason.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/IndexUnusedReason.java
@@ -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();
+  }
+}
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
@@ -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;
 
     /**
@@ -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() {}
 
@@ -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());
           }
@@ -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();
@@ -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;
     }
 
@@ -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.
@@ -743,6 +763,7 @@ ToStringHelper toStringHelper() {
           .add("queryPlan", queryPlan)
           .add("timeline", timeline)
           .add("schema", schema)
+          .add("searchStats", searchStats)
           .add("queryParameters", queryParameters);
     }
 
@@ -765,6 +786,7 @@ public final int hashCode() {
           totalBytesProcessed,
           queryPlan,
           schema,
+          searchStats,
           queryParameters);
     }
 
@@ -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);
       }

diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/SearchStats.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/SearchStats.java
@@ -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();
+  }
+}
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java
@@ -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)
@@ -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()
@@ -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()
@@ -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());
   }

diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -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 {