fix(datastore): query limit without deleted rows (#2822)

star-whale · Oct 8, 2023 · 4be25a5 · 4be25a5
1 parent 3a48759
commit 4be25a5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/server/controller/src/main/java/ai/starwhale/mlops/datastore/DataStore.java b/server/controller/src/main/java/ai/starwhale/mlops/datastore/DataStore.java
@@ -260,7 +260,11 @@ public RecordList query(DataStoreQueryRequest req) {
                 --skipCount;
             }
             while (iterator.hasNext() && limitCount > 0) {
-                results.add(iterator.next());
+                var r = iterator.next();
+                if (r.isDeleted()) {
+                    continue;
+                }
+                results.add(r);
                 --limitCount;
             }
             String lastKey;
@@ -290,7 +294,6 @@ public RecordList query(DataStoreQueryRequest req) {
                 columnSchemaMap = null;
             }
             var records = results.stream()
-                    .filter(r -> !r.isDeleted())
                     .map(r -> RecordEncoder.encodeRecord(r.getValues(), req.isRawResult(), req.isEncodeWithType()))
                     .collect(Collectors.toList());
             return new RecordList(columnSchemaMap,

diff --git a/server/controller/src/test/java/ai/starwhale/mlops/datastore/DataStoreTest.java b/server/controller/src/test/java/ai/starwhale/mlops/datastore/DataStoreTest.java
@@ -245,6 +245,17 @@ public void testQuery() {
                         .columnValueHints(List.of("1", "2", "3", "4", "5"))
                         .build())));
 
+        // query with limit without deleted rows
+        this.dataStore.update("t1", desc, List.of(Map.of("k", "0", "-", "1")));
+        recordList = this.dataStore.query(DataStoreQueryRequest.builder().tableName("t1").limit(2).build());
+        assertThat("test",
+                   recordList.getRecords(),
+                   is(List.of(Map.of("k", "1", "a", "00000004"),
+                              Map.of("k", "2", "a", "00000003"))));
+
+        // revert deleted rows
+        this.dataStore.update("t1", desc, List.of(Map.of("k", "0", "a", "5")));
+
         recordList = this.dataStore.query(DataStoreQueryRequest.builder()
                 .tableName("t1")
                 .filter(TableQueryFilter.builder()