Implement count() method with filter clauses for SpannerStructuredTable

cdapio · Sep 12, 2024 · 7a9ba4a · 7a9ba4a
1 parent 0a21d6f
commit 7a9ba4a
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/...torage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java b/...torage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java
@@ -430,7 +430,7 @@ public void updateAll(Range keyRange, Collection<Field<?>> fields) throws Invali
 
   @Override
   public long count(Collection<Range> keyRanges) {
-    try (ResultSet resultSet = transactionContext.executeQuery(getCountStatement(keyRanges))) {
+    try (ResultSet resultSet = transactionContext.executeQuery(getCountStatement(keyRanges, Collections.emptyList()))) {
       if (!resultSet.next()) {
         return 0L;
       }
@@ -441,7 +441,7 @@ public long count(Collection<Range> keyRanges) {
   @Override
   public long count(Collection<Range> keyRanges, Collection<Field<?>> filterIndexes)
       throws InvalidFieldException, IOException {
-    try (ResultSet resultSet = transactionContext.executeQuery(getCountStatement(keyRanges))) {
+    try (ResultSet resultSet = transactionContext.executeQuery(getCountStatement(keyRanges, filterIndexes))) {
       if (!resultSet.next()) {
         return 0L;
       }
@@ -454,18 +454,21 @@ public void close() {
     // No-op
   }
 
-  private Statement getCountStatement(Collection<Range> ranges) {
+  private Statement getCountStatement(Collection<Range> ranges, Collection<Field<?>> filterIndexes) {
     Map<String, Value> parameters = new HashMap<>();
     String whereClause = getRangesWhereClause(ranges, parameters);
+    Statement.Builder builder = Statement.newBuilder("SELECT COUNT(*) FROM " + escapeName(schema.getTableId().getName()));
 
-    if (whereClause == null) {
-      return Statement.of("SELECT COUNT(*) FROM " + escapeName(schema.getTableId().getName()));
+    if (whereClause != null) {
+      builder.append(" WHERE " + whereClause);
+      parameters.forEach((name, value) -> builder.bind(name).to(value));
+    }
+    if (!filterIndexes.isEmpty()){
+      if (whereClause != null) {
+        builder.append(" AND ");
+      }
+      builder.append(getIndexesFilterClause(filterIndexes, parameters));
     }
-
-    Statement.Builder builder = Statement.newBuilder(
-        "SELECT COUNT(*) FROM " + escapeName(schema.getTableId().getName())
-            + " WHERE " + whereClause);
-    parameters.forEach((name, value) -> builder.bind(name).to(value));
     return builder.build();
   }