From 31d6fe9c6971a0eb3fe1d96fa3ede869e7c58090 Mon Sep 17 00:00:00 2001
From: Tetsuya Shioda <shioda.tetsuya@lab.ntt.co.jp>
Date: Wed, 9 Jan 2019 10:35:55 +0900
Subject: [PATCH] add clear_row method to recommender service

---
 jubakit/recommender.py           | 12 ++++++++++++
 jubakit/test/test_recommender.py | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/jubakit/recommender.py b/jubakit/recommender.py
index 26842a7..86fb4e4 100644
--- a/jubakit/recommender.py
+++ b/jubakit/recommender.py
@@ -55,6 +55,17 @@ def _client_class(cls):
   def _embedded_class(cls):
     return jubatus.embedded.Recommender
 
+  def clear_row(self, dataset):
+    """
+    Removes the given rows from the recommendation table.
+    """
+    cli = self._client()
+    for (idx, (row_id, d)) in dataset:
+      if row_id is None:
+        raise RuntimeError('dataset must have `id`.')
+      result = cli.clear_row(row_id)
+      yield (idx, row_id, result)
+
   def update_row(self, dataset):
     """
     Update data points to the recommender model using the given dataset.
@@ -169,6 +180,7 @@ def decode_row(self, dataset):
       result = cli.decode_row(row_id)
       yield (idx, row_id, result)
 
+
 class Config(GenericConfig):
   """
   Configuration to run Recommender service.
diff --git a/jubakit/test/test_recommender.py b/jubakit/test/test_recommender.py
index 342a960..dca0240 100644
--- a/jubakit/test/test_recommender.py
+++ b/jubakit/test/test_recommender.py
@@ -51,6 +51,26 @@ def test_simple(self):
   def test_embedded(self):
     recommender = Recommender.run(Config(), embedded=True)
 
+  def test_clear_row(self):
+    recommender = Recommender.run(Config())
+    loader = StubLoader()
+
+    # dataset must have id when execute `clear_row`.
+    schema = Schema({'v': Schema.NUMBER})
+    dataset = Dataset(loader, schema)
+    def func():
+      for _ in recommender.clear_row(dataset): pass
+    self.assertRaises(RuntimeError, lambda: func())
+
+    schema = Schema({'v': Schema.ID})
+    dataset = Dataset(loader, schema)
+
+    # expect to get False when table is empty.
+    for (idx, row_id, result) in recommender.clear_row(dataset):
+      self.assertEqual(result, True)
+
+    recommender.stop()
+
   def test_update_row(self):
     recommender = Recommender.run(Config())
     loader = StubLoader()