From 2827324993609ecaebc8c35efc1b7554f5240692 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Tue, 22 Dec 2015 12:34:26 -0800
Subject: [PATCH] Implementing Bigtable row filters for sampling and labeling.

---
 gcloud/bigtable/row.py      | 61 +++++++++++++++++++++++++++++++
 gcloud/bigtable/test_row.py | 72 +++++++++++++++++++++++++++++++++++++
 2 files changed, 133 insertions(+)

diff --git a/gcloud/bigtable/row.py b/gcloud/bigtable/row.py
index d2cbbe5d144b..124bb185583c 100644
--- a/gcloud/bigtable/row.py
+++ b/gcloud/bigtable/row.py
@@ -173,6 +173,31 @@ def to_pb(self):
         return data_pb2.RowFilter(row_key_regex_filter=self.regex)
 
 
+class RowSampleFilter(RowFilter):
+    """Matches all cells from a row with probability p.
+
+    :type sample: float
+    :param sample: The probability of matching a cell (must be in the
+                   interval ``[0, 1]``).
+    """
+
+    def __init__(self, sample):
+        self.sample = sample
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return other.sample == self.sample
+
+    def to_pb(self):
+        """Converts the row filter to a protobuf.
+
+        :rtype: :class:`.data_pb2.RowFilter`
+        :returns: The converted current object.
+        """
+        return data_pb2.RowFilter(row_sample_filter=self.sample)
+
+
 class FamilyNameRegexFilter(_RegexFilter):
     """Row filter for a family name regular expression.
 
@@ -522,3 +547,39 @@ def to_pb(self):
         :returns: The converted current object.
         """
         return data_pb2.RowFilter(strip_value_transformer=self.flag)
+
+
+class ApplyLabelFilter(RowFilter):
+    """Filter to apply labels to cells.
+
+    Intended to be used as an intermediate filter on a pre-existing filtered
+    result set. This was if two sets are combined, the label can tell where
+    the cell(s) originated.This allows the client to determine which results
+    were produced from which part of the filter.
+
+    .. note::
+
+        Due to a technical limitation, it is not currently possible to apply
+        multiple labels to a cell.
+
+    :type label: str
+    :param label: Label to apply to cells in the output row. Values must be
+                  at most 15 characters long, and match the pattern
+                  ``[a-z0-9\\-]+``.
+    """
+
+    def __init__(self, label):
+        self.label = label
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return other.label == self.label
+
+    def to_pb(self):
+        """Converts the row filter to a protobuf.
+
+        :rtype: :class:`.data_pb2.RowFilter`
+        :returns: The converted current object.
+        """
+        return data_pb2.RowFilter(apply_label_transformer=self.label)
diff --git a/gcloud/bigtable/test_row.py b/gcloud/bigtable/test_row.py
index 9396d2afdaa5..fdd78a693b7f 100644
--- a/gcloud/bigtable/test_row.py
+++ b/gcloud/bigtable/test_row.py
@@ -192,6 +192,42 @@ def test_to_pb(self):
         self.assertEqual(pb_val, expected_pb)
 
 
+class TestRowSampleFilter(unittest2.TestCase):
+
+    def _getTargetClass(self):
+        from gcloud.bigtable.row import RowSampleFilter
+        return RowSampleFilter
+
+    def _makeOne(self, *args, **kwargs):
+        return self._getTargetClass()(*args, **kwargs)
+
+    def test_constructor(self):
+        sample = object()
+        row_filter = self._makeOne(sample)
+        self.assertTrue(row_filter.sample is sample)
+
+    def test___eq__type_differ(self):
+        sample = object()
+        row_filter1 = self._makeOne(sample)
+        row_filter2 = object()
+        self.assertNotEqual(row_filter1, row_filter2)
+
+    def test___eq__same_value(self):
+        sample = object()
+        row_filter1 = self._makeOne(sample)
+        row_filter2 = self._makeOne(sample)
+        self.assertEqual(row_filter1, row_filter2)
+
+    def test_to_pb(self):
+        from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
+
+        sample = 0.25
+        row_filter = self._makeOne(sample)
+        pb_val = row_filter.to_pb()
+        expected_pb = data_pb2.RowFilter(row_sample_filter=sample)
+        self.assertEqual(pb_val, expected_pb)
+
+
 class TestFamilyNameRegexFilter(unittest2.TestCase):
 
     def _getTargetClass(self):
@@ -591,3 +627,39 @@ def test_to_pb(self):
         pb_val = row_filter.to_pb()
         expected_pb = data_pb2.RowFilter(strip_value_transformer=flag)
         self.assertEqual(pb_val, expected_pb)
+
+
+class TestApplyLabelFilter(unittest2.TestCase):
+
+    def _getTargetClass(self):
+        from gcloud.bigtable.row import ApplyLabelFilter
+        return ApplyLabelFilter
+
+    def _makeOne(self, *args, **kwargs):
+        return self._getTargetClass()(*args, **kwargs)
+
+    def test_constructor(self):
+        label = object()
+        row_filter = self._makeOne(label)
+        self.assertTrue(row_filter.label is label)
+
+    def test___eq__type_differ(self):
+        label = object()
+        row_filter1 = self._makeOne(label)
+        row_filter2 = object()
+        self.assertNotEqual(row_filter1, row_filter2)
+
+    def test___eq__same_value(self):
+        label = object()
+        row_filter1 = self._makeOne(label)
+        row_filter2 = self._makeOne(label)
+        self.assertEqual(row_filter1, row_filter2)
+
+    def test_to_pb(self):
+        from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
+
+        label = u'label'
+        row_filter = self._makeOne(label)
+        pb_val = row_filter.to_pb()
+        expected_pb = data_pb2.RowFilter(apply_label_transformer=label)
+        self.assertEqual(pb_val, expected_pb)