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