add filters to python api

Author: kboroszko

tensorflow_io/core/BUILD

@@ -184,6 +184,8 @@ cc_library(
         "kernels/bigtable/bigtable_row_set.h",
         "kernels/bigtable/bigtable_row_range.cc",
         "kernels/bigtable/bigtable_row_set.cc",
+        "kernels/bigtable/bigtable_version_filters.cc",
+        "kernels/bigtable/bigtable_version_filters.h",
         "ops/bigtable_ops.cc",
     ],
     copts = tf_io_copts(),

tensorflow_io/core/kernels/bigtable/bigtable_dataset_kernel.cc

@@ -22,14 +22,16 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/resource_op_kernel.h"
 #include "tensorflow_io/core/kernels/bigtable/bigtable_row_set.h"
+#include "tensorflow_io/core/kernels/bigtable/bigtable_version_filters.h"
 
 namespace cbt = ::google::cloud::bigtable;
 
 namespace tensorflow {
 namespace data {
 namespace {
 
-tensorflow::error::Code GoogleCloudErrorCodeToTfErrorCode(::google::cloud::StatusCode code) {
+tensorflow::error::Code GoogleCloudErrorCodeToTfErrorCode(
+    ::google::cloud::StatusCode code) {
   switch (code) {
     case ::google::cloud::StatusCode::kOk:
       return ::tensorflow::error::OK;
@@ -72,9 +74,9 @@ Status GoogleCloudStatusToTfStatus(const ::google::cloud::Status& status) {
   if (status.ok()) {
     return Status::OK();
   }
-  return Status(GoogleCloudErrorCodeToTfErrorCode(status.code()),
-                strings::StrCat("Error reading from Cloud Bigtable: ",
-                                status.message()));
+  return Status(
+      GoogleCloudErrorCodeToTfErrorCode(status.code()),
+      strings::StrCat("Error reading from Cloud Bigtable: ", status.message()));
 }
 
 class BigtableClientResource : public ResourceBase {
@@ -143,6 +145,7 @@ class Iterator : public DatasetIterator<Dataset> {
             this->dataset()->client_resource().CreateTable(table_id).ReadRows(
                 cbt::RowRange::InfiniteRange(),
                 cbt::Filter::Chain(CreateColumnsFilter(columns_),
+                                   this->dataset()->filter_resource().filter(),
                                    cbt::Filter::Latest(1)))),
         it_(this->reader_.begin()),
         column_to_idx_(CreateColumnToIdxMap(columns_)) {}
@@ -270,10 +273,13 @@ class Iterator : public DatasetIterator<Dataset> {
 class Dataset : public DatasetBase {
  public:
   Dataset(OpKernelContext* ctx, BigtableClientResource* client_resource,
-          std::string table_id, std::vector<std::string> columns)
+          io::BigtableFilterResource* filter_resource, std::string table_id,
+          std::vector<std::string> columns)
       : DatasetBase(DatasetContext(ctx)),
         client_resource_(*client_resource),
         client_resource_unref_(client_resource),
+        filter_resource_(*filter_resource),
+        filter_resource_unref_(filter_resource),
         table_id_(table_id),
         columns_(columns) {
     dtypes_.push_back(DT_STRING);
@@ -300,6 +306,9 @@ class Dataset : public DatasetBase {
   }
 
   BigtableClientResource& client_resource() const { return client_resource_; }
+  io::BigtableFilterResource& filter_resource() const {
+    return filter_resource_;
+  }
 
  protected:
   Status AsGraphDefInternal(SerializationContext* ctx,
@@ -314,6 +323,8 @@ class Dataset : public DatasetBase {
  private:
   BigtableClientResource& client_resource_;
   const core::ScopedUnref client_resource_unref_;
+  io::BigtableFilterResource& filter_resource_;
+  const core::ScopedUnref filter_resource_unref_;
   const std::string table_id_;
   const std::vector<std::string> columns_;
   DataTypeVector dtypes_;
@@ -330,10 +341,15 @@ class BigtableDatasetOp : public DatasetOpKernel {
   void MakeDataset(OpKernelContext* ctx, DatasetBase** output) override {
     VLOG(1) << "Make Dataset";
     BigtableClientResource* client_resource;
-    OP_REQUIRES_OK(
-        ctx, LookupResource(ctx, HandleFromInput(ctx, 0), &client_resource));
+    OP_REQUIRES_OK(ctx,
+                   GetResourceFromContext(ctx, "client", &client_resource));
     core::ScopedUnref client_resource_unref_(client_resource);
-    *output = new Dataset(ctx, client_resource, table_id_, columns_);
+    io::BigtableFilterResource* filter_resource;
+    OP_REQUIRES_OK(ctx,
+                   GetResourceFromContext(ctx, "filter", &filter_resource));
+    core::ScopedUnref filter_resource_unref_(filter_resource);
+    *output =
+        new Dataset(ctx, client_resource, filter_resource, table_id_, columns_);
   }
 
  private:

tensorflow_io/core/kernels/bigtable/bigtable_version_filters.cc

@@ -20,18 +20,16 @@ namespace tensorflow {
 namespace io {
 
 class BigtableLatestFilterOp
-    : public OpKernelCreatingResource<BigtableFilterResource> {
+    : public AbstractBigtableResourceOp<BigtableFilterResource> {
  public:
   explicit BigtableLatestFilterOp(OpKernelConstruction* ctx)
-      : OpKernelCreatingResource<BigtableFilterResource>(ctx) {
+      : AbstractBigtableResourceOp<BigtableFilterResource>(ctx) {
     VLOG(1) << "BigtableLatestFilterOp ctor ";
   }
 
  private:
-  Status CreateResource(BigtableFilterResource** resource)
-      TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
-    *resource = new BigtableFilterResource(cbt::Filter::Latest(1));
-    return Status::OK();
+  StatusOr<BigtableFilterResource*> CreateResource() override {
+    return new BigtableFilterResource(cbt::Filter::Latest(1));
   }
 
  private:
@@ -42,20 +40,18 @@ REGISTER_KERNEL_BUILDER(Name("BigtableLatestFilter").Device(DEVICE_CPU),
                         BigtableLatestFilterOp);
 
 class BigtableTimestampRangeFilterOp
-    : public OpKernelCreatingResource<BigtableFilterResource> {
+    : public AbstractBigtableResourceOp<BigtableFilterResource> {
  public:
   explicit BigtableTimestampRangeFilterOp(OpKernelConstruction* ctx)
-      : OpKernelCreatingResource<BigtableFilterResource>(ctx) {
+      : AbstractBigtableResourceOp<BigtableFilterResource>(ctx) {
     VLOG(1) << "BigtableTimestampRangeFilterOp ctor ";
     OP_REQUIRES_OK(ctx, ctx->GetAttr("start", &start_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("end", &end_));
   }
 
  private:
-  Status CreateResource(BigtableFilterResource** resource)
-      TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
-    *resource = new BigtableFilterResource(cbt::Filter::TimestampRangeMicros(start_, end_));
-    return Status::OK();
+  StatusOr<BigtableFilterResource*> CreateResource() override {
+    return new BigtableFilterResource(cbt::Filter::TimestampRangeMicros(start_, end_));
   }
 
  private:

tensorflow_io/core/ops/bigtable_ops.cc

@@ -27,6 +27,7 @@ REGISTER_OP("BigtableClient")
 
 REGISTER_OP("BigtableDataset")
     .Input("client: resource")
+    .Input("filter: resource")
     .Attr("table_id: string")
     .Attr("columns: list(string) >= 1")
     .Output("handle: variant")
@@ -113,7 +114,7 @@ REGISTER_OP("BigtableSampleRowSets")
       return tensorflow::Status::OK();
     });
 
-REGISTER_OP("BigtableEmptyRowRange")
+REGISTER_OP("BigtableLatestFilter")
     .Attr("container: string = ''")
     .Attr("shared_name: string = ''")
     .Output("filter: resource")
@@ -124,7 +125,7 @@ REGISTER_OP("BigtableTimestampRangeFilter")
     .Attr("container: string = ''")
     .Attr("shared_name: string = ''")
     .Attr("start: int")
-    .Attr("start: int")
+    .Attr("end: int")
     .Output("filter: resource")
     .SetIsStateful()
     .SetShapeFn(shape_inference::ScalarShape);

tensorflow_io/python/ops/bigtable/bigtable_dataset_ops.py

@@ -2,6 +2,7 @@
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.framework import tensor_spec
 from tensorflow_io.python.ops import core_ops
+import tensorflow_io.python.ops.bigtable.bigtable_version_filters as filters
 from tensorflow.python.framework import dtypes
 import tensorflow as tf
 
@@ -28,22 +29,23 @@ def __init__(self, client_resource, table_id: str):
         self._table_id = table_id
         self._client_resource = client_resource
 
-    def read_rows(self, columns: List[str]):
-        return _BigtableDataset(self._client_resource, self._table_id, columns)
+    def read_rows(self, columns: List[str], filter: filters.BigtableFilter = filters.latest()):
+        return _BigtableDataset(self._client_resource, self._table_id, columns, filter)
 
 
 class _BigtableDataset(dataset_ops.DatasetSource):
     """_BigtableDataset represents a dataset that retrieves keys and values."""
 
-    def __init__(self, client_resource, table_id: str, columns: List[str]):
+    def __init__(self, client_resource, table_id: str, columns: List[str], filter):
         self._table_id = table_id
         self._columns = columns
+        self._filter = filter
         self._element_spec = tf.TensorSpec(
             shape=[len(columns)], dtype=dtypes.string
         )
 
         variant_tensor = core_ops.bigtable_dataset(
-            client_resource, table_id, columns
+            client_resource, filter._impl, table_id, columns
         )
         super().__init__(variant_tensor)
 

tensorflow_io/python/ops/bigtable/bigtable_version_filters.py

@@ -34,7 +34,7 @@ def latest():
     Returns:
       pbt_C.Filter: Filter passing only most recent version of a value.
     """
-    return core_ops.bigtable_latest_filter()
+    return BigtableFilter(core_ops.bigtable_latest_filter())
 
 
 def timestamp_range(
@@ -60,7 +60,9 @@ def timestamp_range(
     else:
         end_timestamp = int(end * 1e6)
 
-    return core_ops.bigtable_timestamp_range_filter(start_timestamp, end_timestamp)
+    return BigtableFilter(
+        core_ops.bigtable_timestamp_range_filter(start_timestamp, end_timestamp)
+    )
 
 
 def timestamp_range_micros(
@@ -75,6 +77,8 @@ def timestamp_range_micros(
     Returns:
       pbt_C.Filter: Filter passing only values' versions from the specified range.
     """
-    return core_ops.bigtable_timestamp_range_filter(
-        int(start_timestamp), int(end_timestamp)
+    return BigtableFilter(
+        core_ops.bigtable_timestamp_range_filter(
+            int(start_timestamp), int(end_timestamp)
+        )
     )