Converting _BigtableXXXDataset to DatasetV2 type (the version going to be used in TF 2.0).

PiperOrigin-RevId: 239524076
This commit is contained in:
Henry Tan 2019-03-20 19:33:41 -07:00 committed by TensorFlower Gardener
parent 996aeecdaf
commit c261779aff

View File

@ -221,7 +221,7 @@ class BigtableTable(object):
A `tf.data.Dataset`. containing `tf.string` Tensors corresponding to all
of the row keys matching that prefix.
"""
return dataset_ops.DatasetV1Adapter(_BigtablePrefixKeyDataset(self, prefix))
return _BigtablePrefixKeyDataset(self, prefix)
def sample_keys(self):
"""Retrieves a sampling of row keys from the Bigtable table.
@ -233,7 +233,7 @@ class BigtableTable(object):
Returns:
A `tf.data.Dataset` returning string row keys.
"""
return dataset_ops.DatasetV1Adapter(_BigtableSampleKeysDataset(self))
return _BigtableSampleKeysDataset(self)
def scan_prefix(self, prefix, probability=None, columns=None, **kwargs):
"""Retrieves row (including values) from the Bigtable service.
@ -278,8 +278,7 @@ class BigtableTable(object):
"""
probability = _normalize_probability(probability)
normalized = _normalize_columns(columns, kwargs)
return dataset_ops.DatasetV1Adapter(
_BigtableScanDataset(self, prefix, "", "", normalized, probability))
return _BigtableScanDataset(self, prefix, "", "", normalized, probability)
def scan_range(self, start, end, probability=None, columns=None, **kwargs):
"""Retrieves rows (including values) from the Bigtable service.
@ -324,8 +323,7 @@ class BigtableTable(object):
"""
probability = _normalize_probability(probability)
normalized = _normalize_columns(columns, kwargs)
return dataset_ops.DatasetV1Adapter(
_BigtableScanDataset(self, "", start, end, normalized, probability))
return _BigtableScanDataset(self, "", start, end, normalized, probability)
def parallel_scan_prefix(self,
prefix,
@ -381,8 +379,7 @@ class BigtableTable(object):
"""
probability = _normalize_probability(probability)
normalized = _normalize_columns(columns, kwargs)
ds = dataset_ops.DatasetV1Adapter(
_BigtableSampleKeyPairsDataset(self, prefix, "", ""))
ds = _BigtableSampleKeyPairsDataset(self, prefix, "", "")
return self._make_parallel_scan_dataset(ds, num_parallel_scans, probability,
normalized)
@ -444,8 +441,7 @@ class BigtableTable(object):
"""
probability = _normalize_probability(probability)
normalized = _normalize_columns(columns, kwargs)
ds = dataset_ops.DatasetV1Adapter(
_BigtableSampleKeyPairsDataset(self, "", start, end))
ds = _BigtableSampleKeyPairsDataset(self, "", start, end)
return self._make_parallel_scan_dataset(ds, num_parallel_scans, probability,
normalized)