diff --git a/cdap-storage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java b/cdap-storage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java index 0acd660cdf6..eaad8e28245 100644 --- a/cdap-storage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java +++ b/cdap-storage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java @@ -430,7 +430,7 @@ public void updateAll(Range keyRange, Collection> fields) throws Invali @Override public long count(Collection 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 keyRanges) { @Override public long count(Collection keyRanges, Collection> 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 ranges) { + private Statement getCountStatement(Collection ranges, Collection> filterIndexes) { Map 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(); }