feat: version filters (#6)

This PR adds support for Bigtable version filters.

Author: kboroszko

tensorflow_io/core/BUILD

@@ -184,6 +184,8 @@ cc_library(
         "kernels/bigtable/bigtable_row_range.h",
         "kernels/bigtable/bigtable_row_set.cc",
         "kernels/bigtable/bigtable_row_set.h",
+        "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

@@ -23,6 +23,7 @@ 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;
 
@@ -148,6 +149,7 @@ class Iterator : public DatasetIterator<Dataset> {
         reader_(this->dataset()->CreateTable().ReadRows(
             this->dataset()->row_set(),
             cbt::Filter::Chain(CreateColumnsFilter(columns_),
+                               this->dataset()->filter(),
                                cbt::Filter::Latest(1)))),
         it_(this->reader_.begin()),
         column_to_idx_(CreateColumnToIdxMap(columns_)) {
@@ -277,11 +279,12 @@ class Dataset : public DatasetBase {
  public:
   Dataset(OpKernelContext* ctx,
           const std::shared_ptr<cbt::DataClient>& data_client,
-          cbt::RowSet row_set, std::string table_id,
+          cbt::RowSet row_set, cbt::Filter filter, std::string table_id,
           std::vector<std::string> columns)
       : DatasetBase(DatasetContext(ctx)),
         data_client_(data_client),
         row_set_(std::move(row_set)),
+        filter_(std::move(filter)),
         table_id_(table_id),
         columns_(columns) {
     dtypes_.push_back(DT_STRING);
@@ -319,6 +322,8 @@ class Dataset : public DatasetBase {
     return table;
   }
 
+  const cbt::Filter& filter() const { return filter_; }
+
  protected:
   Status AsGraphDefInternal(SerializationContext* ctx,
                             DatasetGraphDefBuilder* b,
@@ -332,6 +337,7 @@ class Dataset : public DatasetBase {
  private:
   std::shared_ptr<cbt::DataClient> const& data_client_;
   const cbt::RowSet row_set_;
+  cbt::Filter filter_;
   const std::string table_id_;
   const std::vector<std::string> columns_;
   DataTypeVector dtypes_;
@@ -357,8 +363,14 @@ class BigtableDatasetOp : public DatasetOpKernel {
                    GetResourceFromContext(ctx, "row_set", &row_set_resource));
     core::ScopedUnref row_set_resource_unref_(row_set_resource);
 
+    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->data_client(),
-                          row_set_resource->row_set(), table_id_, columns_);
+                          row_set_resource->row_set(),
+                          filter_resource->filter(), table_id_, columns_);
   }
 
  private:

tensorflow_io/core/kernels/bigtable/bigtable_version_filters.cc

@@ -0,0 +1,86 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow_io/core/kernels/bigtable/bigtable_version_filters.h"
+
+namespace cbt = ::google::cloud::bigtable;
+
+namespace tensorflow {
+namespace io {
+
+class BigtableLatestFilterOp
+    : public AbstractBigtableResourceOp<BigtableFilterResource> {
+ public:
+  explicit BigtableLatestFilterOp(OpKernelConstruction* ctx)
+      : AbstractBigtableResourceOp<BigtableFilterResource>(ctx) {
+    VLOG(1) << "BigtableLatestFilterOp ctor ";
+  }
+
+ private:
+  StatusOr<BigtableFilterResource*> CreateResource() override {
+    return new BigtableFilterResource(cbt::Filter::Latest(1));
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("BigtableLatestFilter").Device(DEVICE_CPU),
+                        BigtableLatestFilterOp);
+    
+class BigtableTimestampRangeFilterOp
+    : public AbstractBigtableResourceOp<BigtableFilterResource> {
+ public:
+  explicit BigtableTimestampRangeFilterOp(OpKernelConstruction* ctx)
+      : AbstractBigtableResourceOp<BigtableFilterResource>(ctx) {
+    VLOG(1) << "BigtableTimestampRangeFilterOp ctor ";
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("start_ts_us", &start_ts_us_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("end_ts_us", &end_ts_us_));
+  }
+
+ private:
+  StatusOr<BigtableFilterResource*> CreateResource() override {
+    return new BigtableFilterResource(cbt::Filter::TimestampRangeMicros(start_ts_us_, end_ts_us_));
+  }
+
+ private:
+  int64_t start_ts_us_;
+  int64_t end_ts_us_;
+};
+
+REGISTER_KERNEL_BUILDER(Name("BigtableTimestampRangeFilter").Device(DEVICE_CPU),
+                        BigtableTimestampRangeFilterOp);
+
+class BigtablePrintFilterOp : public OpKernel {
+ public:
+  explicit BigtablePrintFilterOp(OpKernelConstruction* context)
+      : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    BigtableFilterResource* resource;
+    OP_REQUIRES_OK(context,
+                   GetResourceFromContext(context, "filter", &resource));
+    core::ScopedUnref unref(resource);
+
+    // Create an output tensor
+    Tensor* output_tensor = NULL;
+    OP_REQUIRES_OK(context, context->allocate_output(0, {1}, &output_tensor));
+    auto output_v = output_tensor->tensor<tstring, 1>();
+
+    output_v(0) = resource->ToString();
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("BigtablePrintFilter").Device(DEVICE_CPU),
+                        BigtablePrintFilterOp);
+
+} // namespace io
+} // namespace tensorflow

tensorflow_io/core/kernels/bigtable/bigtable_version_filters.h

@@ -0,0 +1,63 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef BIGTABLE_VERSION_FILTERS_H
+#define BIGTABLE_VERSION_FILTERS_H
+
+#include "absl/memory/memory.h"
+#include "google/cloud/bigtable/table.h"
+#include "google/cloud/bigtable/table_admin.h"
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/dataset.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/resource_op_kernel.h"
+#include "tensorflow_io/core/kernels/bigtable/bigtable_resource_kernel.h"
+
+
+namespace tensorflow {
+namespace io {
+
+class BigtableFilterResource : public ResourceBase {
+ public:
+  explicit BigtableFilterResource(google::cloud::bigtable::Filter filter)
+      : filter_(std::move(filter)) {
+    VLOG(1) << "BigtableFilterResource ctor";
+  }
+
+  ~BigtableFilterResource() { VLOG(1) << "BigtableFilterResource dtor"; }
+
+  std::string ToString() const {
+    std::string res;
+    google::protobuf::TextFormat::PrintToString(filter_.as_proto(), &res);
+    return res;
+  }
+
+  const google::cloud::bigtable::Filter& filter() const { return filter_; }
+
+  string DebugString() const override {
+    return "BigtableFilterResource:{" + ToString() + "}";
+  }
+
+ private:
+  const google::cloud::bigtable::Filter filter_;
+};
+
+
+}  // namespace io
+}  // namespace tensorflow
+
+#endif /* BIGTABLE_ROW_SET_H */

tensorflow_io/core/ops/bigtable_ops.cc

@@ -28,6 +28,7 @@ REGISTER_OP("BigtableClient")
 REGISTER_OP("BigtableDataset")
     .Input("client: resource")
     .Input("row_set: resource")
+    .Input("filter: resource")
     .Attr("table_id: string")
     .Attr("columns: list(string) >= 1")
     .Output("handle: variant")
@@ -106,3 +107,24 @@ REGISTER_OP("BigtableSplitRowSetEvenly")
       c->set_output(0, c->Vector(c->UnknownDim()));
       return tensorflow::Status::OK();
     });
+
+REGISTER_OP("BigtableLatestFilter")
+    .Attr("container: string = ''")
+    .Attr("shared_name: string = ''")
+    .Output("filter: resource")
+    .SetIsStateful()
+    .SetShapeFn(shape_inference::ScalarShape);
+
+REGISTER_OP("BigtableTimestampRangeFilter")
+    .Attr("container: string = ''")
+    .Attr("shared_name: string = ''")
+    .Attr("start_ts_us: int")
+    .Attr("end_ts_us: int")
+    .Output("filter: resource")
+    .SetIsStateful()
+    .SetShapeFn(shape_inference::ScalarShape);
+
+REGISTER_OP("BigtablePrintFilter")
+    .Input("filter: resource")
+    .Output("output: string")
+    .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
 from tensorflow.python.data.ops import dataset_ops
@@ -34,14 +35,22 @@ 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], row_set: RowSet):
-        return _BigtableDataset(self._client_resource, self._table_id, columns, row_set)
+    def read_rows(
+        self,
+        columns: List[str],
+        row_set: RowSet,
+        filter: filters.BigtableFilter = filters.latest(),
+    ):
+        return _BigtableDataset(
+            self._client_resource, self._table_id, columns, row_set, filter
+        )
 
     def parallel_read_rows(
         self,
         columns: List[str],
         num_parallel_calls=tf.data.AUTOTUNE,
         row_set: RowSet = from_rows_or_ranges(infinite()),
+        filter: filters.BigtableFilter = filters.latest(),
     ):
 
         print("calling parallel read_rows with row_set:", row_set)
@@ -50,7 +59,7 @@ def parallel_read_rows(
         )
 
         def map_func(idx):
-            return self.read_rows(columns, RowSet(samples[idx]))
+            return self.read_rows(columns, RowSet(samples[idx]), filter)
 
         # We interleave a dataset of sample's indexes instead of a dataset of
         # samples, because Dataset.from_tensor_slices attempts to copy the
@@ -69,14 +78,20 @@ 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], row_set: RowSet,
+        self,
+        client_resource,
+        table_id: str,
+        columns: List[str],
+        row_set: RowSet,
+        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, row_set._impl, table_id, columns
+            client_resource, row_set._impl, filter._impl, table_id, columns
         )
         super().__init__(variant_tensor)