Issue #1290: Stopped using metadata for optimized count path (#1293)

* Issue #1290: Stopped using metadata for optimized count path

* add changes, fix test

* handle requirePartitionFilter

* increase timeout for presubmit

CHANGES.md

@@ -1,6 +1,7 @@
 # Release Notes
 
 ## Next
+* Issue #1290: Stopped using metadata for optimized count path
 
 ## 0.41.0 - 2024-09-05
 

bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryClient.java

@@ -15,8 +15,10 @@
  */
 package com.google.cloud.bigquery.connector.common;
 
+import static com.google.cloud.bigquery.connector.common.BigQueryUtil.getPartitionFields;
 import static com.google.cloud.bigquery.connector.common.BigQueryUtil.getQueryForRangePartitionedTable;
 import static com.google.cloud.bigquery.connector.common.BigQueryUtil.getQueryForTimePartitionedTable;
+import static com.google.cloud.bigquery.connector.common.BigQueryUtil.isBigQueryNativeTable;
 
 import com.google.cloud.BaseServiceException;
 import com.google.cloud.RetryOption;
@@ -604,9 +606,23 @@ public long calculateTableSize(TableId tableId, Optional<String> filter) {
 
   public long calculateTableSize(TableInfo tableInfo, Optional<String> filter) {
     TableDefinition.Type type = tableInfo.getDefinition().getType();
-    if (type == TableDefinition.Type.TABLE && !filter.isPresent()) {
-      return tableInfo.getNumRows().longValue();
-    } else if (type == TableDefinition.Type.EXTERNAL && !filter.isPresent()) {
+    if ((type == TableDefinition.Type.EXTERNAL || type == TableDefinition.Type.TABLE)
+        && !filter.isPresent()) {
+      if (isBigQueryNativeTable(tableInfo)
+          && tableInfo.getRequirePartitionFilter() != null
+          && tableInfo.getRequirePartitionFilter()) {
+        List<String> partitioningFields = getPartitionFields(tableInfo);
+        if (partitioningFields.isEmpty()) {
+          throw new IllegalStateException(
+              "Could not find partitioning columns for table requiring partition filter: "
+                  + tableInfo.getTableId());
+        }
+        String table = fullTableName(tableInfo.getTableId());
+        return getNumberOfRows(
+            String.format(
+                "SELECT COUNT(*) from `%s` WHERE %s IS NOT NULL",
+                table, partitioningFields.get(0)));
+      }
       String table = fullTableName(tableInfo.getTableId());
       return getNumberOfRows(String.format("SELECT COUNT(*) from `%s`", table));
     } else if (type == TableDefinition.Type.VIEW

cloudbuild/cloudbuild.yaml

@@ -117,7 +117,7 @@ steps:
 
 
 # Tests take around 1 hr 15 mins in general.
-timeout: 7200s
+timeout: 9000s
 
 options:
   machineType: 'N1_HIGHCPU_32'

spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/integration/WriteIntegrationTestBase.java

@@ -2664,6 +2664,23 @@ public void testTableDescriptionRemainsUnchanged() {
         .isEqualTo(bq.getTable(testDataset.toString(), testTable).getDescription());
   }
 
+  @Test
+  public void testCountAfterWrite() {
+    IntegrationTestUtils.runQuery(
+        String.format("CREATE TABLE `%s.%s` (name STRING, age INT64)", testDataset, testTable));
+    Dataset<Row> read1Df = spark.read().format("bigquery").load(fullTableName());
+    assertThat(read1Df.count()).isEqualTo(0L);
+
+    Dataset<Row> dfToWrite =
+        spark.createDataFrame(
+            Arrays.asList(RowFactory.create("foo", 10), RowFactory.create("bar", 20)),
+            new StructType().add("name", DataTypes.StringType).add("age", DataTypes.IntegerType));
+    writeToBigQueryAvroFormat(dfToWrite, SaveMode.Append, "false");
+
+    Dataset<Row> read2Df = spark.read().format("bigquery").load(fullTableName());
+    assertThat(read2Df.count()).isEqualTo(2L);
+  }
+
   private TableResult insertAndGetTimestampNTZToBigQuery(LocalDateTime time, String format)
       throws InterruptedException {
     Preconditions.checkArgument(timeStampNTZType.isPresent(), "timestampNTZType not present");