return right away if keys is empty (#3658)

Summary: Pull Request resolved: #3658 X-link: facebookresearch/FBGEMM#734 when reading embedding rows from rocks db, skip reading from a db shard if this db shard doesn't contain any requested rows. Reviewed By: pradeepfn Differential Revision: D68913784 fbshipit-source-id: b5949f64ef4b5f60b9ba0df7f3f326a4aee33bd0
pytorch · Feb 5, 2025 · 9c9adb9 · 9c9adb9
1 parent fbcb73e
commit 9c9adb9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_table_batched_embeddings.h b/fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_table_batched_embeddings.h
@@ -846,6 +846,12 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
                               }
                               key_indices.push_back(i);
                             }
+
+                            // bail if nothing to do
+                            if (key_indices.empty()) {
+                              return;
+                            }
+
                             std::sort(
                                 key_indices.begin(),
                                 key_indices.end(),

diff --git a/fbgemm_gpu/test/tbe/ssd/kv_tensor_wrapper_test.py b/fbgemm_gpu/test/tbe/ssd/kv_tensor_wrapper_test.py
@@ -89,6 +89,9 @@ def test_read_tensor_using_wrapper_from_db(
             )
             self.assertEqual(tensor_wrapper.shape, [E, D])
 
+            # read one row as an extreme case
+            narrowed = tensor_wrapper.narrow(0, 0, 1)
+
             # table has a total of E rows
             # load 1000 rows at a time
             step = 1000