Repository: DeepRec-AI/HybridBackend (language: python)

Issue:
ParquetDataset benchmark add tfrecord data
# Summary
the benchmark  for ParquetDataset  is only compare ParquetDataset vs TextLineDataset. add TFRecordDataset, please.
# Installation environment
- GPU model and memory:
- OS Platform:
- Docker version:
- GCC/CUDA/cuDNN version:
- Python/conda version:
- TensorFlow/PyTorch version:

# Willing to contribute

Yes




Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 2956d108..685d24bb 100644
--- a/README.md
+++ b/README.md
@@ -11,11 +11,8 @@ recommender systems on heterogeneous cluster.
 ## Features
 
 - Memory-efficient loading of categorical data
-
 - GPU-efficient orchestration of embedding layers
-
 - Communication-efficient training and evaluation at scale
-
 - Easy to use with existing AI workflows
 
 ## Usage
@@ -26,8 +23,8 @@ A minimal example:
 import tensorflow as tf
 import hybridbackend.tensorflow as hb
 
-ds = hb.data.ParquetDataset(filenames, batch_size=batch_size)
-ds = ds.apply(hb.data.parse())
+ds = hb.data.Dataset.from_parquet(filenames)
+ds = ds.batch(batch_size)
 # ...
 
 with tf.device('/gpu:0'):
@@ -44,16 +41,16 @@ more information.
 
 `pip install {PACKAGE}`
 
-`{PACKAGE}` | Dependency | Python  | CUDA | GLIBC | Data Opt. | Embedding Opt. | Parallelism Opt.
------------ | ---------- | ------- | ---- | ----- | --------- | -------------- | -----------------
-[hybridbackend-deeprec2208-cu114](https://pypi.org/project/hybridbackend-deeprec2208-cu114/) | [DeepRec 22.08](https://github.com/alibaba/DeepRec/tree/deeprec2208) `1` | 3.6 | 11.4 | >=2.27 | &check; | &check; | &check;
-[hybridbackend-tf115-cu118](https://pypi.org/project/hybridbackend-tf115-cu118/) | [TensorFlow 1.15](https://github.com/NVIDIA/tensorflow) `2` | 3.8 | 11.8 | >=2.31 | &check; | &check; | &check;
-[hybridbackend-tf115-cu100](https://pypi.org/project/hybridbackend-tf115-cu100/) | [TensorFlow 1.15](https://github.com/tensorflow/tensorflow/tree/r1.15) | 3.6 | 10.0 | >=2.27 | &check; | &check; | &cross;
-[hybridbackend-tf115-cpu](https://pypi.org/project/hybridbackend-tf115-cpu/) | [TensorFlow 1.15](https://github.com/tensorflow/tensorflow/tree/r1.15) | 3.6 | - | >=2.24 | &check; | &cross; | &cross;
+| `{PACKAGE}`                                                                             | Dependency                                                              | Python | CUDA | GLIBC  | Data Opt. | Embedding Opt. | Parallelism Opt. |
+| ----------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- | ------ | ---- | ------ | --------- | -------------- | ---------------- |
+| [hybridbackend-tf115-cu118](https://pypi.org/project/hybridbackend-tf115-cu118/)             | [TensorFlow 1.15](https://github.com/NVIDIA/tensorflow) `1`              | 3.8    | 11.8 | >=2.31 | &check;   | &check;        | &check;          |
+| [hybridbackend-tf115-cu100](https://pypi.org/project/hybridbackend-tf115-cu100/)             | [TensorFlow 1.15](https://github.com/tensorflow/tensorflow/tree/r1.15)     | 3.6    | 10.0 | >=2.27 | &check;   | &check;        | &cross;          |
+| [hybridbackend-tf115-cpu](https://pypi.org/project/hybridbackend-tf115-cpu/)                 | [TensorFlow 1.15](https://github.com/tensorflow/tensorflow/tree/r1.15)     | 3.6    | -    | >=2.24 | &check;   | &cross;        | &cross;          |
+| [hybridbackend-deeprec2208-cu114](https://pypi.org/project/hybridbackend-deeprec2208-cu114/) | [DeepRec 22.08](https://github.com/alibaba/DeepRec/tree/deeprec2208) `2` | 3.6    | 11.4 | >=2.27 | &check;   | &check;        | &check;          |
 
-> `1`: Suggested docker image: `dsw-registry.cn-shanghai.cr.aliyuncs.com/pai/tensorflow-training:1.15PAI-gpu-py36-cu114-ubuntu18.04`
+> `1`: Suggested docker image: `nvcr.io/nvidia/tensorflow:22.12-tf1-py3`
 
-> `2`: Suggested docker image: `nvcr.io/nvidia/tensorflow:22.12-tf1-py3`
+> `2`: Suggested docker image: `dsw-registry.cn-shanghai.cr.aliyuncs.com/pai/tensorflow-training:1.15PAI-gpu-py36-cu114-ubuntu18.04`
 
 ### Method 2: Build from source
 
@@ -66,13 +63,11 @@ HybridBackend is licensed under the [Apache 2.0 License](LICENSE).
 ## Community
 
 - Please see [Contributing Guide](https://github.com/alibaba/HybridBackend/blob/main/CONTRIBUTING.md)
-before your first contribution.
-
+  before your first contribution.
 - Please [register as an adopter](https://github.com/alibaba/HybridBackend/blob/main/ADOPTERS.md)
-if your organization is interested in adoption. We will discuss
-[RoadMap](https://github.com/alibaba/HybridBackend/blob/main/ROADMAP.md) with
-registered adopters in advance.
-
+  if your organization is interested in adoption. We will discuss
+  [RoadMap](https://github.com/alibaba/HybridBackend/blob/main/ROADMAP.md) with
+  registered adopters in advance.
 - Please cite [HybridBackend](https://ieeexplore.ieee.org/document/9835450) in your publications if it helps:
 
   ```text
@@ -90,4 +85,4 @@ registered adopters in advance.
 If you would like to share your experiences with others, you are welcome to
 contact us in DingTalk:
 
-[![dingtalk](https://github.com/alibaba/HybridBackend/raw/main/docs/images/dingtalk.png)](https://qr.dingtalk.com/action/joingroup?code=v1,k1,VouhbeuTwXYEgaLzSOE8o6VF2kTHVJ8lw5h93WbZW8o=&_dt_no_comment=1&origin=11)
+[![dingtalk](https://github.com/alibaba/HybridBackend/raw/main/docs/images/dingtalk.png)](https://qr.dingtalk.com/action/joingroup?code=v1,k1,VouhbeuTwXYEgaLzSOE8o6VF2kTHVJ8lw5h93WbZW8o=&_dt_no_comment=1&origin=11)
\ No newline at end of file
diff --git a/hybridbackend/__init__.py b/hybridbackend/__init__.py
index c191563f..9cd253dc 100644
--- a/hybridbackend/__init__.py
+++ b/hybridbackend/__init__.py
@@ -20,6 +20,6 @@
 from __future__ import division
 from __future__ import print_function
 
-__version__ = '0.7.0a2'
+__version__ = '0.8.0'
 __author__ = 'Alibaba Group Holding Limited'
 __copyright__ = '2021 Alibaba Group Holding Limited'
diff --git a/hybridbackend/run.py b/hybridbackend/run.py
index 7d2e936c..9d2b3f0a 100755
--- a/hybridbackend/run.py
+++ b/hybridbackend/run.py
@@ -69,6 +69,8 @@ def run(command):
     command: Function or command to run
   '''
   visible_devices = _query_visible_devices()
+  local_world_size_str = str(len(visible_devices))
+
   port = int(os.getenv('HB_RUN_BASE_PORT', '20001'))
   device_to_ports = []
   for d in visible_devices:
@@ -126,6 +128,8 @@ def run(command):
     new_tf_config['task']['type'] = task_type
     new_tf_config['task']['index'] = task_id
     os.environ['TF_CONFIG'] = json.dumps(new_tf_config)
+    os.environ['TF_TASK_TYPE'] = str(task_type)
+    os.environ['TF_TASK_INDEX'] = str(task_id)
     os.environ['CUDA_VISIBLE_DEVICES'] = ''
     os.environ['HB_OP_OPTIMIZATION_DISABLED'] = '1'
     if callable(command):
@@ -165,7 +169,10 @@ def run(command):
       gpu_tf_config['task']['index'] = gpu_index
     gpu_env = os.environ.copy()
     gpu_env['TF_CONFIG'] = json.dumps(gpu_tf_config)
+    gpu_env['TF_TASK_TYPE'] = gpu_tf_config['task']['type']
+    gpu_env['TF_TASK_INDEX'] = str(gpu_tf_config['task']['index'])
     gpu_env['CUDA_VISIBLE_DEVICES'] = device
+    gpu_env['LOCAL_WORLD_SIZE'] = local_world_size_str
     if interop_threads_gpu:
       gpu_env['TF_NUM_INTEROP_THREADS'] = str(interop_threads_gpu)
     if intraop_threads_gpu:
diff --git a/hybridbackend/tensorflow/benchmarks/data_benchmark_parquet.py b/hybridbackend/tensorflow/benchmarks/data_benchmark_parquet.py
index 9b4f9fab..9c914fb9 100644
--- a/hybridbackend/tensorflow/benchmarks/data_benchmark_parquet.py
+++ b/hybridbackend/tensorflow/benchmarks/data_benchmark_parquet.py
@@ -51,13 +51,13 @@ def benchmark(params):
   with tf.Graph().as_default():
     step = tf.train.get_or_create_global_step()
     if params.baseline:
-      ds = hb.data.TabularDataset.from_parquet(params.filenames)
+      ds = hb.data.Dataset.from_parquet(params.filenames)
       ds = ds.map(lambda data: data)  # Prevent fusion
       if params.shuffle:
         ds = ds.shuffle(params.batch_size * 10)
       ds = ds.batch(params.batch_size, drop_remainder=True)
     else:
-      ds = hb.data.TabularDataset.from_parquet(params.filenames)
+      ds = hb.data.Dataset.from_parquet(params.filenames)
       if params.shuffle:
         ds = ds.shuffle_batch(
           params.batch_size, drop_remainder=True,
diff --git a/hybridbackend/tensorflow/common/slice_sum.cu.cc b/hybridbackend/tensorflow/common/slice_sum.cu.cc
new file mode 100644
index 00000000..29726690
--- /dev/null
+++ b/hybridbackend/tensorflow/common/slice_sum.cu.cc
@@ -0,0 +1,103 @@
+/* Copyright 2021 Alibaba Group Holding Limited. 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.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <limits>
+
+#include <tensorflow/core/framework/register_types.h>
+#include <tensorflow/core/framework/tensor.h>
+#include <tensorflow/core/public/version.h>
+
+#include "hybridbackend/common/atomic.cu.h"
+#include "hybridbackend/tensorflow/common/device_functions.h"
+#include "hybridbackend/tensorflow/common/slice_sum.h"
+
+namespace tensorflow {
+
+using GPUDevice = Eigen::GpuDevice;
+
+namespace hybridbackend {
+
+namespace functor {
+
+template <typename T, int32 N = 256>
+__global__ void SliceSumKernel(const int32 num_rows, const int32 num_cols,
+                               const int32 col, const T* input, T* output_total,
+                               T* output) {
+  for (int32 idx : CudaGridRangeX(num_rows)) {
+    const T v = input[idx * num_cols + col];
+    output[idx] = v;
+    atomicAdd(output_total, v);
+  }
+}
+
+template <typename T>
+struct SliceSum<GPUDevice, T> {
+  void operator()(const int32 num_rows, const int32 num_cols, const int32 col,
+                  const T* input, T* output_total, T* output,
+                  const Eigen::GpuDevice& d) {
+    CudaLaunch(SliceSumKernel<T>, num_rows, 0, d, nullptr, num_rows, num_cols,
+               col, input, output_total, output);
+  }
+};
+
+template struct SliceSum<GPUDevice, int32>;
+template struct SliceSum<GPUDevice, int64>;
+template struct SliceSum<GPUDevice, uint32>;
+template struct SliceSum<GPUDevice, uint64>;
+
+template <typename T, int32 N = 256>
+__global__ void GroupSliceSumKernel(const int32 num_rows, const int32 num_cols,
+                                    const int32 col, const int32 num_inputs,
+                                    const T* inputs, T* output_totals,
+                                    T** outputs) {
+  for (int32 idx : CudaGridRangeX(num_inputs * num_rows)) {
+    const int32 s = idx / num_rows;
+    const int32 sidx = idx % num_rows;
+    const T v = inputs[idx * num_cols + col];
+    outputs[s][sidx] = v;
+    atomicAdd(output_totals + s, v);
+  }
+}
+
+template <typename T>
+struct SliceSumN<GPUDevice, T> {
+  void operator()(const int32 num_rows, const int32 num_cols, const int32 col,
+                  const int32 num_inputs, const T* inputs, T* output_totals,
+                  T** outputs, const Eigen::GpuDevice& d) {
+    CudaLaunch(GroupSliceSumKernel<T>, num_inputs * num_rows, 0, d, nullptr,
+               num_rows, num_cols, col, num_inputs, inputs, output_totals,
+               outputs);
+  }
+};
+
+template struct SliceSumN<GPUDevice, int32>;
+template struct SliceSumN<GPUDevice, int64>;
+template struct SliceSumN<GPUDevice, uint32>;
+template struct SliceSumN<GPUDevice, uint64>;
+
+}  // namespace functor
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/data/__init__.py b/hybridbackend/tensorflow/data/__init__.py
index bff117b3..77672e20 100644
--- a/hybridbackend/tensorflow/data/__init__.py
+++ b/hybridbackend/tensorflow/data/__init__.py
@@ -29,8 +29,7 @@
 from hybridbackend.tensorflow.data.prefetch.ops import Iterator
 from hybridbackend.tensorflow.data.rebatch.dataset import RebatchDataset
 from hybridbackend.tensorflow.data.rebatch.dataset import rebatch
-from hybridbackend.tensorflow.data.rectify.dataset import rectify
-from hybridbackend.tensorflow.data.tabular.dataset import TabularDataset
+from hybridbackend.tensorflow.data.tabular.dataset import Dataset
 
 # HybridBackend operators must be loaded before TensorFlow operators to
 # make AWS SDK implementation correct.
diff --git a/hybridbackend/tensorflow/data/rebatch/buffer.h b/hybridbackend/tensorflow/data/rebatch/buffer.h
new file mode 100644
index 00000000..4efbb725
--- /dev/null
+++ b/hybridbackend/tensorflow/data/rebatch/buffer.h
@@ -0,0 +1,79 @@
+/* Copyright 2021 Alibaba Group Holding Limited. 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 HYBRIDBACKEND_TENSORFLOW_DATA_REBATCH_BUFFER_H_
+#define HYBRIDBACKEND_TENSORFLOW_DATA_REBATCH_BUFFER_H_
+
+#include <deque>
+#include <vector>
+
+#include <tensorflow/core/framework/tensor.h>
+#include <tensorflow/core/lib/random/philox_random.h>
+#include <tensorflow/core/lib/random/random.h>
+#include <tensorflow/core/lib/random/random_distributions.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+struct RebatchBufferItem {
+ public:
+  RebatchBufferItem(int64 batch_size, const std::vector<Tensor>& components)
+      : batch_size(batch_size), components(components) {}
+  int64 batch_size;
+  std::vector<Tensor> components;
+};
+
+class RebatchBuffer {
+ public:
+  RebatchBuffer(const DataTypeVector& output_dtypes,
+                const std::vector<PartialTensorShape>& output_shapes,
+                const std::vector<int32>& field_ranks);
+
+  int64 size() const { return size_; }
+
+  Status Put(const std::vector<Tensor>& input_tensors, const int64 num_rows);
+
+  Status PutSlice(const std::vector<Tensor>& input_tensors,
+                  const int64 row_start, const int64 row_limit);
+
+  Status Shuffle(random::SingleSampleAdapter<random::PhiloxRandom>& generator,
+                 const int64 num_rows);
+
+  Status Take(Allocator* alloc, std::vector<Tensor>* output_tensors,
+              const int64 num_rows);
+
+ private:
+  Status TakeDense(Allocator* alloc, std::vector<Tensor>* output_tensors,
+                   std::vector<Tensor>* residual_tensors, const int64 num_rows,
+                   const int64 remained_rows, const int64 rank,
+                   const int64 col);
+
+  Status TakeSparse(Allocator* alloc, std::vector<Tensor>* output_tensors,
+                    std::vector<Tensor>* residual_tensors, const int64 num_rows,
+                    const int64 remained_rows, const int64 rank,
+                    const int64 col);
+
+  const DataTypeVector& output_dtypes_;
+  const std::vector<PartialTensorShape>& output_shapes_;
+  const std::vector<int32> field_ranks_;
+
+  int64 size_;
+  std::deque<RebatchBufferItem> items_;
+};
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW_DATA_REBATCH_BUFFER_H_
diff --git a/hybridbackend/tensorflow/data/rebatch/dataset.py b/hybridbackend/tensorflow/data/rebatch/dataset.py
index 86f91b14..725f6cff 100644
--- a/hybridbackend/tensorflow/data/rebatch/dataset.py
+++ b/hybridbackend/tensorflow/data/rebatch/dataset.py
@@ -23,6 +23,8 @@
 import inspect
 
 # pylint: disable=ungrouped-imports
+from hybridbackend.tensorflow.data.dataframe import input_fields
+
 try:
   from tensorflow.python.data.ops.dataset_ops import DatasetV2 as _dataset  # pylint: disable=unused-import
 
@@ -43,28 +45,19 @@
 
 def rebatch(
     batch_size,
-    min_batch_size=None,
-    fields=None,
     drop_remainder=False,
-    num_parallel_scans=1):
+    fields=None):
   r'''Create a `RebatchDataset`.
 
   Args:
     batch_size: Maxium number of samples in an output batch.
-    min_batch_size: (Optional.) Minimum number of samples in a non-final
-      batch. Same to `batch_size` by default.
-    fields: (Optional.) List of DataFrame fields. Fetched from `input_dataset`
-      by default.
     drop_remainder: (Optional.) If True, smaller final batch is dropped.
       `False` by default.
-    num_parallel_scans: (Optional.) Number of concurrent scans against fields
-        of input dataset.
+    fields: (Optional.) List of DataFrame fields. Fetched from `input_dataset`
+      by default.
   '''
   def _apply_fn(dataset):
     return RebatchDataset(
-      dataset, batch_size,
-      min_batch_size=min_batch_size,
-      fields=fields,
-      drop_remainder=drop_remainder,
-      num_parallel_scans=num_parallel_scans)
+      dataset, input_fields(dataset, fields), batch_size,
+      drop_remainder=drop_remainder)
   return _apply_fn
diff --git a/hybridbackend/tensorflow/data/rebatch/dataset_v1.py b/hybridbackend/tensorflow/data/rebatch/dataset_v1.py
index 897b6f70..8e6a53f5 100644
--- a/hybridbackend/tensorflow/data/rebatch/dataset_v1.py
+++ b/hybridbackend/tensorflow/data/rebatch/dataset_v1.py
@@ -23,6 +23,7 @@
 from __future__ import print_function
 
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import random_seed
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
@@ -30,49 +31,47 @@
 from tensorflow.python.util import nest
 
 from hybridbackend.tensorflow.common import oplib as _ops
-from hybridbackend.tensorflow.data.dataframe import input_fields
 
 
 class RebatchDatasetV1(dataset_ops.Dataset):
   r'''A dataset that adjusts batches.
   '''
   def __init__(
-      self, input_dataset,
-      batch_size,
-      min_batch_size=None,
-      fields=None,
+      self, input_dataset, fields, batch_size,
       drop_remainder=False,
-      num_parallel_scans=1):
+      shuffle_buffer_size=None,
+      shuffle_seed=None,
+      reshuffle_each_iteration=None):
     r'''Create a `RebatchDatasetV1`.
 
     Args:
       input_dataset: A dataset outputs batches.
-      batch_size: Maxium number of samples in an output batch.
-      min_batch_size: (Optional.) Minimum number of samples in a non-final
-        batch. Same to `batch_size` by default.
-      fields: (Optional.) List of DataFrame fields. Fetched from `input_dataset`
+      fields: List of DataFrame fields. Fetched from `input_dataset`
         by default.
+      batch_size: Maxium number of samples in an output batch.
       drop_remainder: (Optional.) If True, smaller final batch is dropped.
         `False` by default.
-      num_parallel_scans: (Optional.) Number of concurrent scans against fields
-        of input dataset.
+      shuffle_buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the
+        number of elements from this dataset from which the new
+        dataset will sample.
+      shuffle_seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing
+        the random seed that will be used to create the distribution. See
+        @{tf.set_random_seed} for behavior.
+      reshuffle_each_iteration: (Optional.) A boolean, which if true indicates
+        that the dataset should be pseudorandomly reshuffled each time it is
+        iterated over. (Defaults to `True`.)
     '''
     self._input_dataset = input_dataset
+    self._fields = fields
     self._batch_size = ops.convert_to_tensor(
       batch_size,
       dtype=dtypes.int64,
       name='batch_size')
-    if min_batch_size is None:
-      min_batch_size = batch_size
-    self._min_batch_size = ops.convert_to_tensor(
-      min_batch_size,
-      dtype=dtypes.int64,
-      name='min_batch_size')
-    self._fields = input_fields(input_dataset, fields)
     self._drop_remainder = drop_remainder
-    if num_parallel_scans == dataset_ops.AUTOTUNE:
-      num_parallel_scans = len(self._fields)
-    self._num_parallel_scans = num_parallel_scans
+    self._shuffle_buffer_size = ops.convert_to_tensor(
+      shuffle_buffer_size or 0, dtype=dtypes.int64, name='buffer_size')
+    self._seed, self._seed2 = random_seed.get_seed(shuffle_seed)
+    self._reshuffle_each_iteration = reshuffle_each_iteration or True
     super().__init__()
 
   @property
@@ -83,22 +82,22 @@ def fields(self):
   def drop_remainder(self):
     return self._drop_remainder
 
-  @property
-  def num_parallel_scans(self):
-    return self._num_parallel_scans
-
   def _as_variant_tensor(self):
-    return _ops.hb_rebatch_tabular_dataset(
+    return _ops.hb_rebatch_tabular_dataset_v2(
       self._input_dataset._as_variant_tensor(),  # pylint: disable=protected-access
       self._batch_size,
-      self._min_batch_size,
+      self._shuffle_buffer_size,
+      self._seed,
+      self._seed2,
       field_ids=nest.flatten({
         f.name: f.map(lambda _, j=idx: j)
         for idx, f in enumerate(self._fields)}),
       field_ragged_indices=nest.flatten(
         {f.name: f.ragged_indices for f in self._fields}),
       drop_remainder=self._drop_remainder,
-      num_parallel_scans=self._num_parallel_scans)
+      reshuffle_each_iteration=self._reshuffle_each_iteration,
+      output_types=self.output_types,
+      output_shapes=self.output_shapes)
 
   def _inputs(self):
     return [self._input_dataset]
diff --git a/hybridbackend/tensorflow/data/rebatch/dataset_v2.py b/hybridbackend/tensorflow/data/rebatch/dataset_v2.py
index 50f7b747..3eaf6937 100644
--- a/hybridbackend/tensorflow/data/rebatch/dataset_v2.py
+++ b/hybridbackend/tensorflow/data/rebatch/dataset_v2.py
@@ -23,65 +23,67 @@
 from __future__ import print_function
 
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import random_seed
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.util import nest
 
 from hybridbackend.tensorflow.common import oplib as _ops
-from hybridbackend.tensorflow.data.dataframe import input_fields
 
 
 class RebatchDatasetV2(dataset_ops.DatasetV2):
   r'''A dataset that adjusts batches.
   '''
   def __init__(
-      self, input_dataset,
-      batch_size,
-      min_batch_size=None,
-      fields=None,
+      self, input_dataset, fields, batch_size,
       drop_remainder=False,
-      num_parallel_scans=1):
+      shuffle_buffer_size=None,
+      shuffle_seed=None,
+      reshuffle_each_iteration=None):
     r'''Create a `RebatchDatasetV2`.
 
     Args:
       input_dataset: A dataset outputs batches.
-      batch_size: Maxium number of samples in an output batch.
-      min_batch_size: (Optional.) Minimum number of samples in a non-final
-        batch. Same to `batch_size` by default.
-      fields: (Optional.) List of DataFrame fields. Fetched from `input_dataset`
+      fields: List of DataFrame fields. Fetched from `input_dataset`
         by default.
+      batch_size: Maxium number of samples in an output batch.
       drop_remainder: (Optional.) If True, smaller final batch is dropped.
         `False` by default.
-      num_parallel_scans: (Optional.) Number of concurrent scans against fields
-        of input dataset.
+      shuffle_buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the
+        number of elements from this dataset from which the new
+        dataset will sample.
+      shuffle_seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing
+        the random seed that will be used to create the distribution. See
+        @{tf.set_random_seed} for behavior.
+      reshuffle_each_iteration: (Optional.) A boolean, which if true indicates
+        that the dataset should be pseudorandomly reshuffled each time it is
+        iterated over. (Defaults to `True`.)
     '''
     self._input_dataset = input_dataset
+    self._fields = fields
     self._batch_size = ops.convert_to_tensor(
       batch_size,
       dtype=dtypes.int64,
       name='batch_size')
-    if min_batch_size is None:
-      min_batch_size = batch_size
-    self._min_batch_size = ops.convert_to_tensor(
-      min_batch_size,
-      dtype=dtypes.int64,
-      name='min_batch_size')
-    self._fields = input_fields(input_dataset, fields)
     self._drop_remainder = drop_remainder
-    if num_parallel_scans == dataset_ops.AUTOTUNE:
-      num_parallel_scans = len(self._fields)
-    self._num_parallel_scans = num_parallel_scans
-    self._impl = _ops.hb_rebatch_tabular_dataset(
+    self._shuffle_buffer_size = ops.convert_to_tensor(
+      shuffle_buffer_size or 0, dtype=dtypes.int64, name='buffer_size')
+    self._seed, self._seed2 = random_seed.get_seed(shuffle_seed)
+    self._reshuffle_each_iteration = reshuffle_each_iteration or True
+    self._impl = _ops.hb_rebatch_tabular_dataset_v2(
       self._input_dataset._variant_tensor,  # pylint: disable=protected-access
       self._batch_size,
-      self._min_batch_size,
+      self._shuffle_buffer_size,
+      self._seed,
+      self._seed2,
       field_ids=nest.flatten({
         f.name: f.map(lambda _, j=idx: j)
         for idx, f in enumerate(self._fields)}),
       field_ragged_indices=nest.flatten(
         {f.name: f.ragged_indices for f in self._fields}),
       drop_remainder=self._drop_remainder,
-      num_parallel_scans=self._num_parallel_scans)
+      reshuffle_each_iteration=self._reshuffle_each_iteration,
+      **self._flat_structure)
     super().__init__(self._impl)
 
   @property
@@ -92,10 +94,6 @@ def fields(self):
   def drop_remainder(self):
     return self._drop_remainder
 
-  @property
-  def num_parallel_scans(self):
-    return self._num_parallel_scans
-
   def _inputs(self):
     return [self._input_dataset]
 
diff --git a/hybridbackend/tensorflow/data/rebatch/rebatch_buffer.cc b/hybridbackend/tensorflow/data/rebatch/rebatch_buffer.cc
new file mode 100644
index 00000000..7f4698ce
--- /dev/null
+++ b/hybridbackend/tensorflow/data/rebatch/rebatch_buffer.cc
@@ -0,0 +1,426 @@
+/* Copyright 2021 Alibaba Group Holding Limited. 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/core/framework/register_types.h>
+#include <tensorflow/core/public/version.h>
+
+#include <algorithm>
+#include <deque>
+#include <vector>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/common/eigen.h"
+#include "hybridbackend/tensorflow/data/rebatch/buffer.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+namespace {
+
+template <typename T>
+Status CopyValues(T* dest, const T* src, int64 num_values,
+                  bool /* can_move */) {
+  static_assert(is_simple_type<T>::value, "Memcpy requires a simple type.");
+  memcpy(dest, src, num_values * sizeof(T));
+  return Status::OK();
+}
+
+template <>
+Status CopyValues<tstring>(tstring* dest, const tstring* src, int64 num_values,
+                           bool can_move) {
+  if (can_move) {
+    for (int64 i = 0; i < num_values; ++i) {
+      *dest++ = std::move(*src++);
+    }
+  } else {
+    std::copy_n(src, num_values, dest);
+  }
+  return Status::OK();
+}
+
+template <>
+Status CopyValues<Variant>(Variant* dest, const Variant* src, int64 num_values,
+                           bool can_move) {
+  if (can_move) {
+    for (int64 i = 0; i < num_values; ++i) {
+      *dest++ = std::move(*src++);
+    }
+  } else {
+    std::copy_n(src, num_values, dest);
+  }
+  return Status::OK();
+}
+
+template <>
+Status CopyValues<ResourceHandle>(ResourceHandle* dest,
+                                  const ResourceHandle* src, int64 num_values,
+                                  bool /* can_move */) {
+  std::copy_n(src, num_values, dest);
+  return Status::OK();
+}
+
+template <>
+Status CopyValues<Eigen::half>(Eigen::half* dest, const Eigen::half* src,
+                               int64 num_values, bool /* can_move */) {
+  std::copy_n(src, num_values, dest);
+  return Status::OK();
+}
+
+Status CopyToSlicesFromTensor(Tensor* dest, const int64 dest_index,
+                              const Tensor& src, const bool can_move) {
+  // can_move can be set by `src.RefCountIsOne();` if this function is a friend
+  // of `Tensor` class
+  if (src.dim_size(0) < 1) {
+    return Status::OK();
+  }
+  DCHECK_NE(dest->dim_size(0), 0);
+  DCHECK_NE(src.dim_size(0), 0);
+  DCHECK_GE(dest_index, 0);
+  const int64 num_values_per_record = dest->NumElements() / dest->dim_size(0);
+  if (TF_PREDICT_FALSE(num_values_per_record !=
+                       (src.NumElements() / src.dim_size(0)))) {
+    return errors::Internal(
+        "Cannot perform copy: shape of tensors does not match. "
+        " Shapes are: [src]: ",
+        src.shape().DebugString(), ", [dest]: ", dest->shape().DebugString());
+  }
+  if (TF_PREDICT_FALSE(dest_index + src.dim_size(0) > dest->dim_size(0))) {
+    return errors::Internal(
+        "Cannot perform copy: shape of dest tensor is too small. "
+        " Shapes are: [src]: ",
+        src.shape().DebugString(), ", [dest]: ", dest->shape().DebugString(),
+        " Destination index is ", dest_index);
+  }
+
+#define HANDLE_TYPE(T)                                                    \
+  case DataTypeToEnum<T>::value: {                                        \
+    const T* src_ptr = src.unaligned_flat<T>().data();                    \
+    T* dest_ptr = dest->unaligned_flat<T>().data() +                      \
+                  (num_values_per_record * dest_index);                   \
+    return CopyValues<T>(dest_ptr, src_ptr, src.NumElements(), can_move); \
+  }
+
+  switch (src.dtype()) {
+    TF_CALL_ALL_TYPES(HANDLE_TYPE);
+    TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
+    TF_CALL_uint32(HANDLE_TYPE);
+    TF_CALL_uint64(HANDLE_TYPE);
+#undef HANDLE_TYPE
+    default:
+      return errors::Unimplemented(
+          "CopyToSlicesFromTensor Unhandled data type: ", src.dtype());
+  }
+}
+
+}  // namespace
+
+RebatchBuffer::RebatchBuffer(
+    const DataTypeVector& output_dtypes,
+    const std::vector<PartialTensorShape>& output_shapes,
+    const std::vector<int32>& field_ranks)
+    : output_dtypes_(output_dtypes),
+      output_shapes_(output_shapes),
+      field_ranks_(field_ranks),
+      size_(0) {}
+
+Status RebatchBuffer::Put(const std::vector<Tensor>& input_tensors,
+                          const int64 num_rows) {
+  if (TF_PREDICT_FALSE(num_rows == 0)) {
+    return Status::OK();
+  }
+
+  items_.emplace_back(num_rows, input_tensors);
+  size_ += num_rows;
+  return Status::OK();
+}
+
+Status RebatchBuffer::PutSlice(const std::vector<Tensor>& input_tensors,
+                               const int64 row_start, const int64 row_limit) {
+  if (TF_PREDICT_FALSE(row_start == row_limit)) {
+    return Status::OK();
+  }
+
+  std::vector<Tensor> sliced_input_tensors(output_dtypes_.size());
+  int64 col = 0;
+  for (int64 rank : field_ranks_) {
+    int64 start = row_start;
+    int64 limit = row_limit;
+    if (rank != 0) {
+      for (size_t split_idx = 1; split_idx < rank + 1; ++split_idx) {
+        auto split_slice =
+            input_tensors[col + split_idx].Slice(start, (limit + 1));
+        const int64 slice_limit = (limit - start);
+        start = split_slice.unaligned_flat<int32>()(0);
+        limit = split_slice.unaligned_flat<int32>()(slice_limit);
+        sliced_input_tensors[col + split_idx] = std::move(split_slice);
+      }
+    }
+    const auto input_shape = input_tensors[col].shape();
+    int64 input_base_elems = 1;
+    if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
+      input_base_elems = input_shape.num_elements() / input_shape.dim_size(0);
+    }
+    sliced_input_tensors[col] = input_tensors[col].Slice(
+        start / input_base_elems, limit / input_base_elems);
+    col += (rank + 1);
+  }
+
+  const int64 num_rows = row_limit - row_start;
+  items_.emplace_back(num_rows, std::move(sliced_input_tensors));
+  size_ += num_rows;
+  return Status::OK();
+}
+
+Status RebatchBuffer::Shuffle(
+    random::SingleSampleAdapter<random::PhiloxRandom>& generator,
+    const int64 num_rows) {
+  if (size_ < num_rows) {
+    return errors::InvalidArgument("Not enough rows (", size_,
+                                   ") in buffer to shuffle a batch (", num_rows,
+                                   ")");
+  }
+
+  for (int64 row = 0; row < num_rows; ++row) {
+    int64 picked = row + generator() % (size_ - row);
+    std::swap(items_[picked], items_[row]);
+  }
+  return Status::OK();
+}
+
+Status RebatchBuffer::Take(Allocator* alloc,
+                           std::vector<Tensor>* output_tensors,
+                           const int64 num_rows) {
+  if (TF_PREDICT_FALSE(size_ < num_rows)) {
+    return errors::InvalidArgument("Not enough rows (", size_,
+                                   ") in buffer to populate a batch (",
+                                   num_rows, ")");
+  }
+
+  // Compute remained rows
+  int64 remained_rows = 0;
+  int64 num_dirty_rows = 0;
+  for (int64 row = 0; num_dirty_rows < items_.size(); ++num_dirty_rows) {
+    if (row + items_[num_dirty_rows].batch_size > num_rows) {
+      remained_rows = (num_rows - row);
+      break;
+    }
+    row += items_[num_dirty_rows].batch_size;
+  }
+
+  const size_t num_components = output_dtypes_.size();
+  output_tensors->clear();
+  output_tensors->resize(num_components);
+  std::vector<Tensor> residual_tensors;
+  residual_tensors.resize(num_components);
+  int64 col = 0;
+  for (int64 rank : field_ranks_) {
+    if (rank == 0) {
+      TF_RETURN_IF_ERROR(TakeDense(alloc, output_tensors, &residual_tensors,
+                                   num_rows, remained_rows, rank, col));
+    } else {
+      TF_RETURN_IF_ERROR(TakeSparse(alloc, output_tensors, &residual_tensors,
+                                    num_rows, remained_rows, rank, col));
+    }
+    col += (rank + 1);
+  }
+
+  for (int64 idx = 0; idx < num_dirty_rows; ++idx) {
+    items_.pop_front();
+  }
+  if (remained_rows > 0) {
+    int64 residual_rows = items_.front().batch_size - remained_rows;
+    items_.pop_front();
+    if (residual_rows > 0) {
+      items_.emplace_front(residual_rows, std::move(residual_tensors));
+    }
+  }
+
+  size_ -= num_rows;
+  return Status::OK();
+}
+
+Status RebatchBuffer::TakeDense(Allocator* alloc,
+                                std::vector<Tensor>* output_tensors,
+                                std::vector<Tensor>* residual_tensors,
+                                const int64 num_rows, const int64 remained_rows,
+                                const int64 rank, const int64 col) {
+  // Create output
+  PartialTensorShape output_pshape(output_shapes_[col]);
+  output_pshape.set_dim(0, num_rows);
+  TensorShape output_shape;
+  output_pshape.AsTensorShape(&output_shape);
+  (*output_tensors)[col] = Tensor(alloc, output_dtypes_[col], output_shape);
+  if (!output_tensors->at(col).IsInitialized()) {
+    return errors::ResourceExhausted(
+        "Failed to allocate memory for output component ", col);
+  }
+
+  // Populate output
+  for (int64 idx = 0, row = 0; idx < items_.size(); ++idx) {
+    auto& item = items_[idx].components[col];
+    if (row + items_[idx].batch_size > num_rows) {
+      if (remained_rows == 0) {
+        break;
+      }
+      auto sliced_input = item.Slice(0, remained_rows);
+      TF_RETURN_IF_ERROR(CopyToSlicesFromTensor(&(output_tensors->at(col)), row,
+                                                std::move(sliced_input), true));
+      (*residual_tensors)[col] =
+          item.Slice(remained_rows, items_[idx].batch_size);
+      break;
+    }
+    TF_RETURN_IF_ERROR(
+        CopyToSlicesFromTensor(&(output_tensors->at(col)), row, item, true));
+    row += items_[idx].batch_size;
+  }
+
+  return Status::OK();
+}
+
+Status RebatchBuffer::TakeSparse(Allocator* alloc,
+                                 std::vector<Tensor>* output_tensors,
+                                 std::vector<Tensor>* residual_tensors,
+                                 const int64 num_rows,
+                                 const int64 remained_rows, const int64 rank,
+                                 const int64 col) {
+  // Create and populate output splits
+  int64 remained_dim0_size = 1 + remained_rows;
+  for (size_t split_idx = 1; split_idx < rank + 1; ++split_idx) {
+    int64 next_remained_dim0_size = 0;
+    int64 dim0_size = 0;
+    for (int64 idx = 0, row = 0; idx < items_.size(); ++idx) {
+      if (row + items_[idx].batch_size > num_rows) {
+        next_remained_dim0_size =
+            items_[idx].components[col + split_idx].unaligned_flat<int32>()(
+                remained_dim0_size - 1);
+        dim0_size += (remained_dim0_size - 1);
+        break;
+      }
+      dim0_size += (items_[idx].components[col + split_idx].dim_size(0) - 1);
+      row += items_[idx].batch_size;
+    }
+
+    PartialTensorShape split_pshape(output_shapes_[col + split_idx]);
+    split_pshape.set_dim(0, 1 + dim0_size);
+    TensorShape split_shape;
+    split_pshape.AsTensorShape(&split_shape);
+    (*output_tensors)[col + split_idx] = Tensor(alloc, DT_INT32, split_shape);
+
+    if (!output_tensors->at(col + split_idx).IsInitialized()) {
+      return errors::ResourceExhausted(
+          "Failed to allocate memory for output component ", col);
+    }
+
+    (*output_tensors)[col + split_idx].unaligned_flat<int32>()(0) = 0;
+    int64 dim0_index = 0;
+    for (int64 idx = 0, row = 0; idx < items_.size(); ++idx) {
+      auto& split = items_[idx].components[col + split_idx];
+      int32 output_last = output_tensors->at(col + split_idx)
+                              .unaligned_flat<int32>()(dim0_index);
+      int32 input_first = split.unaligned_flat<int32>()(0);
+      if (row + items_[idx].batch_size > num_rows) {
+        if (remained_rows == 0) {
+          break;
+        }
+        auto sliced_input_split = split.Slice(1, remained_dim0_size);
+        TF_RETURN_IF_ERROR(CopyToSlicesFromTensor(
+            &(output_tensors->at(col + split_idx)), 1 + dim0_index,
+            std::move(sliced_input_split), true));
+        auto sliced_output_split =
+            output_tensors->at(col + split_idx)
+                .Slice(1 + dim0_index, remained_dim0_size + dim0_index);
+        sliced_output_split.unaligned_flat<int32>() +=
+            sliced_output_split.unaligned_flat<int32>().constant(output_last -
+                                                                 input_first);
+        auto residual_input_split =
+            split.Slice(remained_dim0_size - 1, split.dim_size(0));
+        (*residual_tensors)[col + split_idx] = std::move(residual_input_split);
+        break;
+      }
+      auto sliced_input_split = split.Slice(1, split.dim_size(0));
+      TF_RETURN_IF_ERROR(CopyToSlicesFromTensor(
+          &(output_tensors->at(col + split_idx)), 1 + dim0_index,
+          std::move(sliced_input_split), true));
+      auto sliced_output_split =
+          output_tensors->at(col + split_idx)
+              .Slice(1 + dim0_index, split.dim_size(0) + dim0_index);
+      sliced_output_split.unaligned_flat<int32>() +=
+          sliced_output_split.unaligned_flat<int32>().constant(output_last -
+                                                               input_first);
+      dim0_index += sliced_input_split.dim_size(0);
+      row += items_[idx].batch_size;
+    }
+
+    remained_dim0_size = next_remained_dim0_size;
+  }
+
+  // Create and populate ouput values
+  int64 values_dim0_size = 0;
+  for (int64 idx = 0, row = 0; idx < items_.size(); ++idx) {
+    if (row + items_[idx].batch_size > num_rows) {
+      break;
+    }
+    values_dim0_size += items_[idx].components[col].dim_size(0);
+    row += items_[idx].batch_size;
+  }
+  PartialTensorShape base_pshape(output_shapes_[col]);
+  base_pshape.set_dim(0, 1);
+  if (remained_rows > 0) {
+    values_dim0_size += remained_dim0_size / base_pshape.num_elements();
+  }
+  PartialTensorShape output_pshape(output_shapes_[col]);
+  output_pshape.set_dim(0, values_dim0_size);
+  TensorShape output_shape;
+  output_pshape.AsTensorShape(&output_shape);
+  (*output_tensors)[col] = Tensor(alloc, output_dtypes_[col], output_shape);
+  if (!output_tensors->at(col).IsInitialized()) {
+    return errors::ResourceExhausted(
+        "Failed to allocate memory for output component ", col);
+  }
+
+  int64 dim0_index = 0;
+  for (int64 idx = 0, row = 0; idx < items_.size(); ++idx) {
+    auto& values = items_[idx].components[col];
+    if (row + items_[idx].batch_size > num_rows) {
+      if (remained_rows == 0) {
+        break;
+      }
+      const auto values_shape = values.shape();
+      int64 values_base_elems = 1;
+      if (TF_PREDICT_FALSE(values_shape.dims() > 1)) {
+        values_base_elems =
+            values_shape.num_elements() / values_shape.dim_size(0);
+      }
+      auto sliced_values =
+          values.Slice(0, remained_dim0_size / values_base_elems);
+      TF_RETURN_IF_ERROR(
+          CopyToSlicesFromTensor(&(output_tensors->at(col)), dim0_index,
+                                 std::move(sliced_values), true));
+      (*residual_tensors)[col] = values.Slice(
+          remained_dim0_size / values_base_elems, values.dim_size(0));
+      break;
+    }
+
+    TF_RETURN_IF_ERROR(CopyToSlicesFromTensor(&(output_tensors->at(col)),
+                                              dim0_index, values, true));
+    dim0_index += values.dim_size(0);
+    row += items_[idx].batch_size;
+  }
+
+  return Status::OK();
+}
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
diff --git a/hybridbackend/tensorflow/data/rebatch/rebatch_dataset.cc b/hybridbackend/tensorflow/data/rebatch/rebatch_dataset.cc
deleted file mode 100644
index 2e774506..00000000
--- a/hybridbackend/tensorflow/data/rebatch/rebatch_dataset.cc
+++ /dev/null
@@ -1,913 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. 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 <absl/strings/str_cat.h>
-
-#include <tensorflow/core/framework/common_shape_fns.h>
-#include <tensorflow/core/framework/op.h>
-#include <tensorflow/core/framework/op_def_builder.h>
-#include <tensorflow/core/framework/partial_tensor_shape.h>
-#include <tensorflow/core/framework/shape_inference.h>
-#include <tensorflow/core/framework/tensor.h>
-#include <tensorflow/core/lib/core/blocking_counter.h>
-#include <tensorflow/core/lib/core/threadpool.h>
-#include <tensorflow/core/platform/file_system.h>
-#include <tensorflow/core/public/version.h>
-
-#include <deque>
-#include <map>
-#include <unordered_set>
-#include <vector>
-
-#include "hybridbackend/tensorflow/common/dataset.h"
-#include "hybridbackend/tensorflow/common/eigen.h"
-
-namespace tensorflow {
-namespace hybridbackend {
-
-constexpr char kRebatchTabularDatasetWorkerPool[] =
-    "rebatch_tabular_dataset_worker_pool";
-
-REGISTER_OP("HbRebatchTabularDataset")
-    .Output("handle: variant")
-    .Input("input_dataset: variant")
-    .Input("batch_size: int64")
-    .Input("min_batch_size: int64")
-    .Attr("field_ids: list(int) >= 1")
-    .Attr("field_ragged_indices: list(int) >= 1")
-    .Attr("drop_remainder: bool")
-    .Attr("num_parallel_scans: int = 1")
-    .SetIsStateful()  // NOTE: Source dataset ops must be marked stateful to
-                      // inhibit constant folding.
-    .SetShapeFn([](shape_inference::InferenceContext* c) {
-      shape_inference::ShapeHandle unused;
-      // batch_size should be a scalar.
-      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
-      // min_batch_size should be a scalar.
-      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
-
-      return shape_inference::ScalarShape(c);
-    })
-    .Doc(R"doc(
-A dataset that resizes batches from another tabular dataset.
-
-handle: The handle to reference the dataset.
-input_dataset: Input batch dataset.
-batch_size: Maxium number of samples in an output batch.
-min_batch_size: Minimum number of samples in an non-final batch.
-field_ids: A list of tensor indices to indicate the type of a tensor is
-  values (0), batch splits (1) or other splits (>1).
-field_ragged_indices: A list of indices to indicate the type of a tensor is
-  values (0), batch splits (1) or other splits (>1).
-drop_remainder: If True, only keep batches with exactly `batch_size` samples.
-num_parallel_scans: Number of concurrent scans against fields of input dataset.
-)doc");
-
-class RebatchTabularDatasetOp : public UnaryDatasetOpKernel {
- public:
-  explicit RebatchTabularDatasetOp(OpKernelConstruction* ctx)
-      : UnaryDatasetOpKernel(ctx),
-        drop_remainder_(false),
-        num_parallel_scans_(0) {
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("field_ids", &field_ids_));
-    OP_REQUIRES_OK(
-        ctx, ctx->GetAttr("field_ragged_indices", &field_ragged_indices_));
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("drop_remainder", &drop_remainder_));
-    OP_REQUIRES_OK(ctx,
-                   ctx->GetAttr("num_parallel_scans", &num_parallel_scans_));
-  }
-
-  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
-                   DatasetBase** output) override;
-
- private:
-  class Dataset;
-
-  std::vector<int> field_ids_;
-  std::vector<int> field_ragged_indices_;
-  bool drop_remainder_;
-  int num_parallel_scans_;
-};
-
-class RebatchTabularDatasetOp::Dataset : public DatasetBase {
- public:
-  Dataset(OpKernelContext* ctx, const DatasetBase* input_dataset,
-          const int64 batch_size, const int64 min_batch_size,
-          const std::vector<int>& field_ids,
-          const std::vector<int>& field_ragged_indices,
-          const bool drop_remainder, const int num_parallel_scans)
-      : DatasetBase(DatasetContext(ctx)),
-        input_dataset_(input_dataset),
-        batch_size_(batch_size),
-        min_batch_size_(min_batch_size),
-        field_ids_(field_ids),
-        field_ragged_indices_(field_ragged_indices),
-        drop_remainder_(drop_remainder),
-        num_parallel_scans_(num_parallel_scans) {
-    input_dataset_->Ref();
-
-    const auto& input_shapes = input_dataset_->output_shapes();
-    output_shapes_.reserve(input_shapes.size());
-    for (const auto& input_shape : input_shapes) {
-      auto partial_shape = input_shape;
-      partial_shape.RemoveDim(0);
-      if (drop_remainder_) {
-        output_shapes_.emplace_back(
-            PartialTensorShape({batch_size_}).Concatenate(partial_shape));
-      } else {
-        output_shapes_.emplace_back(
-            PartialTensorShape({-1}).Concatenate(partial_shape));
-      }
-    }
-  }
-
-  ~Dataset() override { input_dataset_->Unref(); }
-
-  std::unique_ptr<IteratorBase> MakeIteratorInternal(
-      const string& prefix) const override;
-
-  const DataTypeVector& output_dtypes() const override {
-    return input_dataset_->output_dtypes();
-  }
-
-  const std::vector<PartialTensorShape>& output_shapes() const override {
-    return output_shapes_;
-  }
-
-  string DebugString() const override {
-    return "RebatchTabularDatasetOp::Dataset";
-  }
-
- protected:
-  Status AsGraphDefInternal(SerializationContext* ctx,
-                            DatasetGraphDefBuilder* b,
-                            Node** output) const override {
-    Node* input_dataset = nullptr;
-    TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_dataset_, &input_dataset));
-    Node* batch_size = nullptr;
-    TF_RETURN_IF_ERROR(b->AddScalar(batch_size_, &batch_size));
-    Node* min_batch_size = nullptr;
-    TF_RETURN_IF_ERROR(b->AddScalar(min_batch_size_, &min_batch_size));
-    AttrValue field_ids;
-    b->BuildAttrValue(field_ids_, &field_ids);
-    AttrValue field_ragged_indices;
-    b->BuildAttrValue(field_ragged_indices_, &field_ragged_indices);
-    AttrValue drop_remainder;
-    b->BuildAttrValue(drop_remainder_, &drop_remainder);
-    AttrValue num_parallel_scans;
-    b->BuildAttrValue(num_parallel_scans_, &num_parallel_scans);
-
-    TF_RETURN_IF_ERROR(b->AddDataset(
-        this, {{0, input_dataset}, {1, batch_size}, {2, min_batch_size}}, {},
-        {{"field_ids", field_ids},
-         {"field_ragged_indices", field_ragged_indices},
-         {"drop_remainder", drop_remainder},
-         {"num_parallel_scans", num_parallel_scans}},
-        output));
-    return Status::OK();
-  }
-
- private:
-  class Iterator;
-  const DatasetBase* const input_dataset_;
-  const int64 batch_size_;
-  const int64 min_batch_size_;
-  const std::vector<int> field_ids_;
-  const std::vector<int> field_ragged_indices_;
-  const bool drop_remainder_;
-  const int num_parallel_scans_;
-
-  std::vector<PartialTensorShape> output_shapes_;
-};
-
-void RebatchTabularDatasetOp::MakeDataset(OpKernelContext* ctx,
-                                          DatasetBase* input,
-                                          DatasetBase** output) {
-  int64 batch_size = 0;
-  OP_REQUIRES_OK(ctx, PARSE_SCALAR(ctx, "batch_size", &batch_size));
-  OP_REQUIRES(ctx, batch_size > 0,
-              errors::InvalidArgument("batch_size must be greater than zero."));
-
-  int64 min_batch_size = 0;
-  OP_REQUIRES_OK(ctx, PARSE_SCALAR(ctx, "min_batch_size", &min_batch_size));
-  OP_REQUIRES(
-      ctx, min_batch_size > 0,
-      errors::InvalidArgument("min_batch_size must be greater than zero."));
-
-  OP_REQUIRES(ctx, batch_size >= min_batch_size,
-              errors::InvalidArgument(
-                  "batch_size must be greater than min_batch_size."));
-
-  OP_REQUIRES(
-      ctx, (!drop_remainder_) || (batch_size == min_batch_size),
-      errors::InvalidArgument("drop_remainder must be false if batch_size is "
-                              "different than min_batch_size."));
-
-  *output =
-      new Dataset(ctx, input, batch_size, min_batch_size, field_ids_,
-                  field_ragged_indices_, drop_remainder_, num_parallel_scans_);
-}
-
-namespace {
-void RecalculateSplit(Tensor* split, int32 value) {
-  int32* sdata =
-      reinterpret_cast<int32*>(const_cast<char*>(split->tensor_data().data()));
-  if (TF_PREDICT_FALSE(value == 0)) {
-    return;
-  }
-  if (CHECK_EIGEN_ALIGN(sdata)) {
-    split->flat<int32>() += split->flat<int32>().constant(value);
-    return;
-  }
-  intptr_t sdata_ptr = reinterpret_cast<intptr_t>(sdata);
-  size_t offset = EIGEN_MAX_ALIGN_BYTES - sdata_ptr % EIGEN_MAX_ALIGN_BYTES;
-  offset /= sizeof(int32);
-  const int32 ssize = split->NumElements();
-  if (offset > ssize) {
-    offset = ssize;
-  }
-  for (size_t i = 0; i < offset; ++i) {
-    sdata[i] += value;
-  }
-  if (offset < ssize) {
-    const auto input_shape = split->shape();
-    int64 input_base_elems = 1;
-    if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
-      input_base_elems = input_shape.num_elements() / input_shape.dim_size(0);
-    }
-    auto rslice =
-        split->Slice(offset / input_base_elems, ssize / input_base_elems);
-    rslice.flat<int32>() += rslice.flat<int32>().constant(value);
-  }
-}
-}  // namespace
-
-class RebatchTabularDatasetOp::Dataset::Iterator
-    : public DatasetIterator<RebatchTabularDatasetOp::Dataset> {
- public:
-  explicit Iterator(const Params& params)
-      : DatasetIterator<RebatchTabularDatasetOp::Dataset>(params),
-        queue_batch_size_(0) {
-    tensor_queues_.resize(dataset()->field_ragged_indices_.size());
-  }
-
-  Status Initialize(IteratorContext* ctx) override {
-    mutex_lock l(mu_);
-    field_ranks_.clear();
-    int64 prev_field_id = -1;
-    int64 prev_ragged_index = -1;
-    for (size_t i = 0; i < dataset()->field_ids_.size(); ++i) {
-      const int64 field_id = dataset()->field_ids_[i];
-      const int64 ragged_index = dataset()->field_ragged_indices_[i];
-      auto output_dtype = dataset()->output_dtypes()[i];
-      if (TF_PREDICT_FALSE(ragged_index > 1 && output_dtype != DT_INT32)) {
-        return errors::InvalidArgument("Output tensor ", i,
-                                       " must be DT_INT32 not ",
-                                       DataTypeString(output_dtype));
-      }
-      if (ragged_index <= prev_ragged_index) {
-        if (TF_PREDICT_FALSE(field_id == prev_field_id)) {
-          return errors::Internal("Invalid `fields` serialization");
-        }
-        field_ranks_.push_back(prev_ragged_index);
-      }
-      prev_field_id = field_id;
-      prev_ragged_index = ragged_index;
-    }
-    field_ranks_.push_back(prev_ragged_index);
-    if (dataset()->num_parallel_scans_ > 1) {
-#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) < 1015L
-      thread_pool_.reset(new thread::ThreadPool(
-          Env::Default(), ThreadOptions(),
-          kRebatchTabularDatasetWorkerPool /* name */,
-          dataset()->num_parallel_scans_ /* num_threads */,
-          false /* low_latency_hint */));
-#else
-      thread_pool_ = ctx->CreateThreadPool(
-          kRebatchTabularDatasetWorkerPool /* name */,
-          dataset()->num_parallel_scans_ /* num_threads */);
-#endif
-    }
-    return dataset()->input_dataset_->MakeIterator(ctx, prefix(), &input_impl_);
-  }
-
-  int64 GetBatchSize(const std::vector<Tensor>& input_tensors) {
-    if (TF_PREDICT_FALSE(input_tensors.size() == 0)) {
-      return 0;
-    }
-    if (field_ranks_[0] > 0) {
-      return input_tensors[1].dim_size(0) - 1;
-    }
-    return input_tensors[0].dim_size(0);
-  }
-
-  Status Redirect(std::vector<Tensor>* output_tensors,
-                  const std::vector<Tensor>& input_tensors) {
-    *output_tensors = std::move(input_tensors);
-    return Status::OK();
-  }
-
-  Status SliceAndRedirect(std::vector<Tensor>* output_tensors,
-                          const std::vector<Tensor>& input_tensors,
-                          const int64 row_start, const int64 row_limit) {
-    // Slice input tensors.
-    int64 cur = 0;
-    for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-      const int64 rank = field_ranks_[fid];
-      if (rank == 0) {
-        const auto input_shape = input_tensors[cur].shape();
-        int64 input_base_elems = 1;
-        if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
-          input_base_elems =
-              input_shape.num_elements() / input_shape.dim_size(0);
-        }
-        output_tensors->push_back(input_tensors[cur].Slice(
-            row_start / input_base_elems, row_limit / input_base_elems));
-        cur += (rank + 1);
-        continue;
-      }
-      output_tensors->emplace_back();
-      int64 start = row_start;
-      int64 limit = row_limit;
-      for (size_t i = 1; i < rank + 1; ++i) {
-        const auto input_shape = input_tensors[cur + i].shape();
-        int64 input_base_elems = 1;
-        if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
-          input_base_elems =
-              input_shape.num_elements() / input_shape.dim_size(0);
-        }
-        auto s = input_tensors[cur + i].Slice(start / input_base_elems,
-                                              (limit + 1) / input_base_elems);
-        const int32* sdata =
-            reinterpret_cast<int32*>(const_cast<char*>(s.tensor_data().data()));
-        const int64 slimit = limit - start;
-        start = sdata[0];
-        limit = sdata[slimit];
-        output_tensors->push_back(std::move(s));
-      }
-      const auto input_shape = input_tensors[cur].shape();
-      int64 input_base_elems = 1;
-      if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
-        input_base_elems = input_shape.num_elements() / input_shape.dim_size(0);
-      }
-      output_tensors->at(cur) = input_tensors[cur].Slice(
-          start / input_base_elems, limit / input_base_elems);
-      cur += (rank + 1);
-    }
-    // Recalculate splits.
-    auto recalc_splits = [output_tensors](int64 rank, int64 cur) {
-      for (size_t i = 1; i < rank + 1; ++i) {
-        auto s = output_tensors->at(cur + i);
-        const int32* sdata =
-            reinterpret_cast<int32*>(const_cast<char*>(s.tensor_data().data()));
-        RecalculateSplit(&s, -sdata[0]);
-      }
-    };
-    if (thread_pool_) {
-      BlockingCounter counter(field_ranks_.size());
-      cur = 0;
-      for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-        const int64 rank = field_ranks_[fid];
-        if (rank == 0) {
-          cur += (rank + 1);
-          continue;
-        }
-        thread_pool_->Schedule([&recalc_splits, &counter, rank, cur]() {
-          recalc_splits(rank, cur);
-          counter.DecrementCount();
-        });
-        cur += (rank + 1);
-      }
-      counter.Wait();
-    } else {
-      cur = 0;
-      for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-        const int64 rank = field_ranks_[fid];
-        if (rank == 0) {
-          cur += (rank + 1);
-          continue;
-        }
-        recalc_splits(rank, cur);
-        cur += (rank + 1);
-      }
-    }
-    return Status::OK();
-  }
-
-  Status Enqueue(const std::vector<Tensor>& input_tensors,
-                 const int64 batch_size) {
-    if (TF_PREDICT_FALSE(batch_size == 0)) {
-      return Status::OK();
-    }
-    for (size_t input_idx = 0; input_idx < input_tensors.size(); ++input_idx) {
-      tensor_queues_[input_idx].push_back(input_tensors[input_idx]);
-    }
-    queue_batch_size_ += batch_size;
-    return Status::OK();
-  }
-
-  Status SliceAndEnqueue(const std::vector<Tensor>& input_tensors,
-                         const int64 row_start, const int64 row_limit) {
-    if (TF_PREDICT_FALSE(row_limit == row_start)) {
-      return Status::OK();
-    }
-    int64 cur = 0;
-    for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-      const int64 rank = field_ranks_[fid];
-      if (rank == 0) {
-        const auto input_shape = input_tensors[cur].shape();
-        int64 input_base_elems = 1;
-        if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
-          input_base_elems =
-              input_shape.num_elements() / input_shape.dim_size(0);
-        }
-        tensor_queues_[cur].push_back(input_tensors[cur].Slice(
-            row_start / input_base_elems, row_limit / input_base_elems));
-        cur += (rank + 1);
-        continue;
-      }
-      int64 start = row_start;
-      int64 limit = row_limit;
-      for (size_t i = 1; i < rank + 1; ++i) {
-        const auto input_shape = input_tensors[cur + i].shape();
-        int64 input_base_elems = 1;
-        if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
-          input_base_elems =
-              input_shape.num_elements() / input_shape.dim_size(0);
-        }
-        auto sliced = input_tensors[cur + i].Slice(
-            start / input_base_elems, (limit + 1) / input_base_elems);
-        int32* sdata = reinterpret_cast<int32*>(
-            const_cast<char*>(sliced.tensor_data().data()));
-        const int64 slimit = limit - start;
-        start = sdata[0];
-        limit = sdata[slimit];
-        tensor_queues_[cur + i].push_back(std::move(sliced));
-      }
-      const auto input_shape = input_tensors[cur].shape();
-      int64 input_base_elems = 1;
-      if (TF_PREDICT_FALSE(input_shape.dims() > 1)) {
-        input_base_elems = input_shape.num_elements() / input_shape.dim_size(0);
-      }
-      tensor_queues_[cur].push_back(input_tensors[cur].Slice(
-          start / input_base_elems, limit / input_base_elems));
-      cur += (rank + 1);
-    }
-    queue_batch_size_ += (row_limit - row_start);
-    return Status::OK();
-  }
-
-  Status Dequeue(IteratorContext* ctx, std::vector<Tensor>* output_tensors,
-                 const int64 batch_size) {
-    AllocatorAttributes host_alloc_attrs;
-    host_alloc_attrs.set_on_host(true);
-    Allocator* alloc = ctx->allocator(host_alloc_attrs);
-    int64 cur = 0;
-    for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-      // Create value tensor.
-      auto output_dtype = dataset()->output_dtypes()[cur];
-      PartialTensorShape output_pshape(dataset()->output_shapes_[cur]);
-      int64 dim0_size = 0;
-      for (auto& t : tensor_queues_[cur]) {
-        dim0_size += t.dim_size(0);
-        if (dim0_size > batch_size) {
-          dim0_size = batch_size;
-          break;
-        }
-      }
-      output_pshape.set_dim(0, dim0_size);
-      TensorShape output_shape;
-      output_pshape.AsTensorShape(&output_shape);
-      if (TF_PREDICT_FALSE(output_dtype == DT_STRING)) {
-        if (TF_PREDICT_FALSE(!TensorShapeUtils::IsVector(output_shape))) {
-          return errors::InvalidArgument(
-              "Tensors in string field must be vector");
-        }
-      }
-      output_tensors->emplace_back(alloc, output_dtype, output_shape);
-
-      const int64 rank = field_ranks_[fid];
-      if (rank == 0) {
-        cur += (rank + 1);
-        continue;
-      }
-
-      // Create split tensors.
-      for (size_t i = 1; i < rank + 1; ++i) {
-        PartialTensorShape split_pshape(dataset()->output_shapes_[cur + i]);
-        int64 dim0_size = 0;
-        for (auto& t : tensor_queues_[cur + i]) {
-          dim0_size += (t.dim_size(0) - 1);
-          if (dim0_size > batch_size) {
-            dim0_size = batch_size;
-            break;
-          }
-        }
-        split_pshape.set_dim(0, 1 + dim0_size);
-
-        TensorShape split_shape;
-        split_pshape.AsTensorShape(&split_shape);
-        output_tensors->emplace_back(alloc, DT_INT32, split_shape);
-      }
-      cur += (rank + 1);
-    }
-
-    auto flush_queue = [this, output_tensors, batch_size](int64 rank,
-                                                          int64 cur) {
-      // Copy to value tensor.
-      if (TF_PREDICT_FALSE(dataset()->output_dtypes()[cur] == DT_STRING)) {
-        int output_cur = 0;
-        int64 dim0_size = 0;
-        int64 num_pops = 0;
-        for (auto& t : tensor_queues_[cur]) {
-          dim0_size += t.dim_size(0);
-          int64 tensor_size = t.NumElements();
-          if (dim0_size > batch_size) {
-            PartialTensorShape residual_pshape(dataset()->output_shapes_[cur]);
-            residual_pshape.set_dim(0, batch_size + t.dim_size(0) - dim0_size);
-            tensor_size = residual_pshape.num_elements();
-          }
-          for (int i = 0; i < tensor_size; ++i) {
-            output_tensors->at(cur).vec<string>()(output_cur + i) =
-                t.unaligned_flat<string>()(i);
-          }
-          output_cur += tensor_size;
-          ++num_pops;
-          if (dim0_size > batch_size) {
-            break;
-          }
-        }
-        for (int64 i = 0; i < num_pops; ++i) {
-          tensor_queues_[cur].pop_front();
-        }
-      } else {
-        char* output_ptr =
-            const_cast<char*>(output_tensors->at(cur).tensor_data().data());
-        int64 dim0_size = 0;
-        int64 num_pops = 0;
-        for (auto& t : tensor_queues_[cur]) {
-          dim0_size += t.dim_size(0);
-          int64 tensor_bytes = t.TotalBytes();
-          if (dim0_size > batch_size) {
-            PartialTensorShape residual_pshape(dataset()->output_shapes_[cur]);
-            residual_pshape.set_dim(0, batch_size + t.dim_size(0) - dim0_size);
-            tensor_bytes =
-                residual_pshape.num_elements() * DataTypeSize(t.dtype());
-          }
-          assert(output_tensors->at(cur).TotalBytes() >= tensor_bytes);
-          memcpy(output_ptr, t.tensor_data().data(), tensor_bytes);
-          output_ptr += tensor_bytes;
-          ++num_pops;
-          if (dim0_size > batch_size) {
-            break;
-          }
-        }
-        for (int64 i = 0; i < num_pops; ++i) {
-          tensor_queues_[cur].pop_front();
-        }
-      }
-
-      if (rank == 0) {
-        return;
-      }
-
-      // Recalculate splits.
-      for (size_t i = 1; i < rank + 1; ++i) {
-        int32* sdata = reinterpret_cast<int32*>(const_cast<char*>(
-            output_tensors->at(cur + i).tensor_data().data()));
-        sdata[0] = 0;
-        int32 tstart = 1;
-        int64 dim0_size = 0;
-        int64 num_pops = 0;
-        for (auto& t : tensor_queues_[cur + i]) {
-          dim0_size += (t.dim_size(0) - 1);
-          int32* tdata = reinterpret_cast<int32*>(
-              const_cast<char*>(t.tensor_data().data()));
-          int32 tsize = t.NumElements() - 1;
-          if (dim0_size > batch_size) {
-            PartialTensorShape residual_pshape(t.shape());
-            residual_pshape.set_dim(0, batch_size + t.dim_size(0) - dim0_size);
-            tsize = residual_pshape.num_elements() - 1;
-          }
-          memcpy(sdata + tstart, tdata + 1, tsize * DataTypeSize(DT_INT32));
-          const auto output_shape = output_tensors->at(cur + i).shape();
-          int64 input_base_elems = 1;
-          if (TF_PREDICT_FALSE(output_shape.dims() > 1)) {
-            input_base_elems =
-                output_shape.num_elements() / output_shape.dim_size(0);
-          }
-          auto tpart = output_tensors->at(cur + i).Slice(
-              tstart / input_base_elems, (tstart + tsize) / input_base_elems);
-          RecalculateSplit(&tpart, sdata[tstart - 1] - tdata[0]);
-          tstart += tsize;
-          ++num_pops;
-          if (dim0_size > batch_size) {
-            break;
-          }
-        }
-        for (int64 n = 0; n < num_pops; ++n) {
-          tensor_queues_[cur + i].pop_front();
-        }
-      }
-    };
-
-    if (thread_pool_) {
-      BlockingCounter counter(field_ranks_.size());
-      cur = 0;
-      for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-        const int64 rank = field_ranks_[fid];
-        thread_pool_->Schedule([&flush_queue, &counter, rank, cur]() {
-          flush_queue(rank, cur);
-          counter.DecrementCount();
-        });
-        cur += (rank + 1);
-      }
-      counter.Wait();
-    } else {
-      cur = 0;
-      for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-        const int64 rank = field_ranks_[fid];
-        flush_queue(rank, cur);
-        cur += (rank + 1);
-      }
-    }
-
-    if (queue_batch_size_ > batch_size) {
-      queue_batch_size_ -= batch_size;
-    } else {
-      queue_batch_size_ = 0;
-    }
-    return Status::OK();
-  }
-
-  Status DequeueAll(IteratorContext* ctx, std::vector<Tensor>* output_tensors) {
-    AllocatorAttributes host_alloc_attrs;
-    host_alloc_attrs.set_on_host(true);
-    Allocator* alloc = ctx->allocator(host_alloc_attrs);
-    int64 cur = 0;
-    for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-      // Create value tensor.
-      auto output_dtype = dataset()->output_dtypes()[cur];
-      PartialTensorShape output_pshape(dataset()->output_shapes_[cur]);
-      int64 dim0_size = 0;
-      for (auto& t : tensor_queues_[cur]) {
-        dim0_size += t.dim_size(0);
-      }
-      output_pshape.set_dim(0, dim0_size);
-      TensorShape output_shape;
-      output_pshape.AsTensorShape(&output_shape);
-      if (TF_PREDICT_FALSE(output_dtype == DT_STRING)) {
-        if (TF_PREDICT_FALSE(!TensorShapeUtils::IsVector(output_shape))) {
-          return errors::InvalidArgument(
-              "Tensors in string field must be vector");
-        }
-      }
-      output_tensors->emplace_back(alloc, output_dtype, output_shape);
-
-      const int64 rank = field_ranks_[fid];
-      if (rank == 0) {
-        cur += (rank + 1);
-        continue;
-      }
-
-      // Create split tensors.
-      for (size_t i = 1; i < rank + 1; ++i) {
-        PartialTensorShape split_pshape(dataset()->output_shapes_[cur + i]);
-        int64 dim0_size = 1;
-        for (auto& t : tensor_queues_[cur + i]) {
-          dim0_size += (t.dim_size(0) - 1);
-        }
-        split_pshape.set_dim(0, dim0_size);
-        TensorShape split_shape;
-        split_pshape.AsTensorShape(&split_shape);
-        output_tensors->emplace_back(alloc, DT_INT32, split_shape);
-      }
-      cur += (rank + 1);
-    }
-
-    auto flush_queue = [this, output_tensors](int64 rank, int64 cur) {
-      // Copy to value tensor.
-      if (TF_PREDICT_FALSE(dataset()->output_dtypes()[cur] == DT_STRING)) {
-        int output_cur = 0;
-        for (auto& t : tensor_queues_[cur]) {
-          const int tensor_size = t.NumElements();
-          for (int i = 0; i < tensor_size; ++i) {
-            output_tensors->at(cur).vec<string>()(output_cur + i) =
-                t.unaligned_flat<string>()(i);
-          }
-          output_cur += tensor_size;
-        }
-      } else {
-        char* output_ptr =
-            const_cast<char*>(output_tensors->at(cur).tensor_data().data());
-        for (auto& t : tensor_queues_[cur]) {
-          const auto tensor_bytes = t.TotalBytes();
-          assert(output_tensors->at(cur).TotalBytes() >= tensor_bytes);
-          memcpy(output_ptr, t.tensor_data().data(), tensor_bytes);
-          output_ptr += tensor_bytes;
-        }
-      }
-      tensor_queues_[cur].clear();
-
-      if (rank == 0) {
-        return;
-      }
-
-      // Recalculate splits.
-      for (size_t i = 1; i < rank + 1; ++i) {
-        int32* sdata = reinterpret_cast<int32*>(const_cast<char*>(
-            output_tensors->at(cur + i).tensor_data().data()));
-        sdata[0] = 0;
-        int32 tstart = 1;
-        for (auto& t : tensor_queues_[cur + i]) {
-          int32* tdata = reinterpret_cast<int32*>(
-              const_cast<char*>(t.tensor_data().data()));
-          memcpy(sdata + tstart, tdata + 1, t.TotalBytes() - sizeof(int32));
-          const int32 tsize = t.NumElements() - 1;
-          const auto output_shape = output_tensors->at(cur + i).shape();
-          int64 input_base_elems = 1;
-          if (TF_PREDICT_FALSE(output_shape.dims() > 1)) {
-            input_base_elems =
-                output_shape.num_elements() / output_shape.dim_size(0);
-          }
-          auto tpart = output_tensors->at(cur + i).Slice(
-              tstart / input_base_elems, (tstart + tsize) / input_base_elems);
-          RecalculateSplit(&tpart, sdata[tstart - 1] - tdata[0]);
-          tstart += tsize;
-        }
-        tensor_queues_[cur + i].clear();
-      }
-    };
-
-    if (thread_pool_) {
-      BlockingCounter counter(field_ranks_.size());
-      cur = 0;
-      for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-        const int64 rank = field_ranks_[fid];
-        thread_pool_->Schedule([&flush_queue, &counter, rank, cur]() {
-          flush_queue(rank, cur);
-          counter.DecrementCount();
-        });
-        cur += (rank + 1);
-      }
-      counter.Wait();
-    } else {
-      cur = 0;
-      for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
-        const int64 rank = field_ranks_[fid];
-        flush_queue(rank, cur);
-        cur += (rank + 1);
-      }
-    }
-
-    queue_batch_size_ = 0;
-    return Status::OK();
-  }
-
-  Status GetNextInternal(IteratorContext* ctx,
-                         std::vector<Tensor>* output_tensors,
-                         bool* end_of_sequence) override {
-    mutex_lock l(mu_);
-    if (queue_batch_size_ > dataset()->batch_size_) {
-      // Queue is not empty and very large.
-      TF_RETURN_IF_ERROR(Dequeue(ctx, output_tensors, dataset()->batch_size_));
-      return Status::OK();
-    }
-    *end_of_sequence = false;
-    while (!*end_of_sequence) {
-      std::vector<Tensor> input_tensors;
-      int64 input_batch_size = 0;
-      bool input_not_found = true;
-      if (input_impl_) {
-        // Current input is ready.
-        bool input_end_of_sequence = false;
-        TF_RETURN_IF_ERROR(
-            input_impl_->GetNext(ctx, &input_tensors, &input_end_of_sequence));
-        if (input_end_of_sequence) {
-          input_impl_.reset();
-        }
-        input_batch_size = GetBatchSize(input_tensors);
-        input_not_found = false;
-      }
-      int64 output_batch_size = queue_batch_size_ + input_batch_size;
-      if (output_batch_size < dataset()->min_batch_size_) {
-        // Merge batches.
-        if (!input_impl_ && input_not_found) {
-          // All inputs are read.
-          if (TF_PREDICT_FALSE(dataset()->drop_remainder_)) {
-            *end_of_sequence = true;
-            return Status::OK();
-          }
-          if (TF_PREDICT_FALSE(output_batch_size == 0)) {
-            *end_of_sequence = true;
-            return Status::OK();
-          }
-          TF_RETURN_IF_ERROR(DequeueAll(ctx, output_tensors));
-          return Status::OK();
-        } else {
-          // Not all inputs are not read.
-          TF_RETURN_IF_ERROR(Enqueue(input_tensors, input_batch_size));
-          continue;
-        }
-      } else if (output_batch_size <= dataset()->batch_size_) {
-        // Redirect or merge batches.
-        if (queue_batch_size_ == 0) {
-          // Queue is empty.
-          TF_RETURN_IF_ERROR(Redirect(output_tensors, input_tensors));
-          return Status::OK();
-        } else if (!input_impl_ && input_not_found) {
-          // Queue is not empty, all inputs are read.
-          if (TF_PREDICT_FALSE(dataset()->drop_remainder_ &&
-                               output_batch_size < dataset()->batch_size_)) {
-            *end_of_sequence = true;
-            return Status::OK();
-          }
-          TF_RETURN_IF_ERROR(DequeueAll(ctx, output_tensors));
-          return Status::OK();
-        } else {
-          // Queue is not empty, not all inputs are consumed.
-          TF_RETURN_IF_ERROR(Enqueue(input_tensors, input_batch_size));
-          TF_RETURN_IF_ERROR(DequeueAll(ctx, output_tensors));
-          return Status::OK();
-        }
-      } else {
-        // Split batches.
-        if (queue_batch_size_ == 0) {
-          // Queue is empty.
-          TF_RETURN_IF_ERROR(SliceAndRedirect(output_tensors, input_tensors, 0,
-                                              dataset()->batch_size_));
-          TF_RETURN_IF_ERROR(SliceAndEnqueue(
-              input_tensors, dataset()->batch_size_, input_batch_size));
-          return Status::OK();
-        } else if (queue_batch_size_ < dataset()->batch_size_) {
-          // Queue is not empty and small.
-          const auto remained_batch_size =
-              dataset()->batch_size_ - queue_batch_size_;
-          TF_RETURN_IF_ERROR(
-              SliceAndEnqueue(input_tensors, 0, remained_batch_size));
-          TF_RETURN_IF_ERROR(DequeueAll(ctx, output_tensors));
-          TF_RETURN_IF_ERROR(SliceAndEnqueue(input_tensors, remained_batch_size,
-                                             input_batch_size));
-          return Status::OK();
-        } else if (queue_batch_size_ == dataset()->batch_size_) {
-          // Queue is ready to dequeue.
-          TF_RETURN_IF_ERROR(DequeueAll(ctx, output_tensors));
-          TF_RETURN_IF_ERROR(Enqueue(input_tensors, input_batch_size));
-          return Status::OK();
-        } else {
-          // Queue is not empty and large.
-          TF_RETURN_IF_ERROR(
-              Dequeue(ctx, output_tensors, dataset()->batch_size_));
-          TF_RETURN_IF_ERROR(Enqueue(input_tensors, input_batch_size));
-          return Status::OK();
-        }
-      }
-    }
-    return Status::OK();
-  }
-
- protected:
-  Status SaveInternal(IteratorStateWriter* writer) override {
-    return errors::Unimplemented("SaveInternal is currently not supported");
-  }
-
-  Status RestoreInternal(IteratorContext* ctx,
-                         IteratorStateReader* reader) override {
-    return errors::Unimplemented("RestoreInternal is currently not supported");
-  }
-
- private:
-  mutex mu_;
-  std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
-  std::vector<int64> field_ranks_;
-  std::vector<std::deque<Tensor>> tensor_queues_;
-  int64 queue_batch_size_;
-  std::unique_ptr<thread::ThreadPool> thread_pool_;
-};
-
-std::unique_ptr<IteratorBase>
-RebatchTabularDatasetOp::Dataset::MakeIteratorInternal(
-    const string& prefix) const {
-  return std::unique_ptr<IteratorBase>(
-      new Iterator({this, absl::StrCat(prefix, "::RebatchTabular")}));
-}
-
-REGISTER_KERNEL_BUILDER(Name("HbRebatchTabularDataset").Device(DEVICE_CPU),
-                        RebatchTabularDatasetOp);
-
-WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS("HbRebatchTabularDataset");
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
diff --git a/hybridbackend/tensorflow/data/rectify/rectify_dataset.cc b/hybridbackend/tensorflow/data/rebatch/rebatch_dataset_v2.cc
similarity index 52%
rename from hybridbackend/tensorflow/data/rectify/rectify_dataset.cc
rename to hybridbackend/tensorflow/data/rebatch/rebatch_dataset_v2.cc
index 3596c6b9..75732dda 100644
--- a/hybridbackend/tensorflow/data/rectify/rectify_dataset.cc
+++ b/hybridbackend/tensorflow/data/rebatch/rebatch_dataset_v2.cc
@@ -21,23 +21,28 @@ limitations under the License.
 #include <tensorflow/core/framework/partial_tensor_shape.h>
 #include <tensorflow/core/framework/shape_inference.h>
 #include <tensorflow/core/framework/tensor.h>
+#include <tensorflow/core/lib/core/blocking_counter.h>
+#include <tensorflow/core/lib/core/threadpool.h>
+#include <tensorflow/core/platform/file_system.h>
 #include <tensorflow/core/public/version.h>
 
 #include <deque>
+#include <map>
+#include <unordered_set>
 #include <vector>
 
 #include "hybridbackend/common/env.h"
 #include "hybridbackend/tensorflow/common/dataset.h"
-#include "hybridbackend/tensorflow/data/rectify/queue.h"
+#include "hybridbackend/tensorflow/common/eigen.h"
+#include "hybridbackend/tensorflow/data/rebatch/buffer.h"
 
 namespace tensorflow {
 namespace hybridbackend {
 
-constexpr char kRectifyDatasetWorkerPool[] = "rectify_dataset_worker_pool";
-constexpr char kRectifyDatasetFastPathEnv[] =
-    "HB_DATA_RECTIFY_FAST_PATH_DISABLED";
+constexpr char kRebatchDatasetFastPathEnv[] =
+    "HB_DATA_REBATCH_FASTPATH_DISABLED";
 
-REGISTER_OP("HbRectifyDataset")
+REGISTER_OP("HbRebatchTabularDatasetV2")
     .Output("handle: variant")
     .Input("input_dataset: variant")
     .Input("batch_size: int64")
@@ -46,7 +51,8 @@ REGISTER_OP("HbRectifyDataset")
     .Input("shuffle_seed2: int64")
     .Attr("drop_remainder: bool")
     .Attr("reshuffle_each_iteration: bool = true")
-    .Attr("output_kinds: list(int) >= 1")
+    .Attr("field_ids: list(int) >= 1")
+    .Attr("field_ragged_indices: list(int) >= 1")
     .Attr("output_types: list(type) >= 1")
     .Attr("output_shapes: list(shape) >= 1")
     .SetShapeFn([](shape_inference::InferenceContext* c) {
@@ -62,22 +68,35 @@ REGISTER_OP("HbRectifyDataset")
       return shape_inference::ScalarShape(c);
     })
     .Doc(R"doc(
-A dataset that rectifies samples.
+A dataset that resizes batches from another tabular dataset.
 
 handle: The handle to reference the dataset.
 input_dataset: Input batch dataset.
 batch_size: Maxium number of samples in an output batch.
+shuffle_buffer_size: Buffer size to shuffle.
+shuffle_seed: Seed for shuffling.
+shuffle_seed2: Seed for shuffling.
 drop_remainder: If True, only keep batches with exactly `batch_size` samples.
+reshuffle_each_iteration: If true, the dataset should be pseudorandomly
+  reshuffled each time it is iterated over.
+field_ids: A list of tensor indices to indicate the type of a tensor is
+  values (0), batch splits (1) or other splits (>1).
+field_ragged_indices: A list of indices to indicate the type of a tensor is
+  values (0), batch splits (1) or other splits (>1).
 )doc");
 
-class RectifyDatasetOp : public UnaryDatasetOpKernel {
+class RebatchTabularDatasetV2Op : public UnaryDatasetOpKernel {
  public:
-  explicit RectifyDatasetOp(OpKernelConstruction* ctx)
-      : UnaryDatasetOpKernel(ctx), drop_remainder_(false) {
+  explicit RebatchTabularDatasetV2Op(OpKernelConstruction* ctx)
+      : UnaryDatasetOpKernel(ctx),
+        drop_remainder_(false),
+        reshuffle_each_iteration_(false) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("drop_remainder", &drop_remainder_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("reshuffle_each_iteration",
                                      &reshuffle_each_iteration_));
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_kinds", &output_kinds_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("field_ids", &field_ids_));
+    OP_REQUIRES_OK(
+        ctx, ctx->GetAttr("field_ragged_indices", &field_ragged_indices_));
   }
 
   void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
@@ -88,16 +107,18 @@ class RectifyDatasetOp : public UnaryDatasetOpKernel {
 
   bool drop_remainder_;
   bool reshuffle_each_iteration_;
-  std::vector<int> output_kinds_;
+  std::vector<int> field_ids_;
+  std::vector<int> field_ragged_indices_;
 };
 
-class RectifyDatasetOp::Dataset : public DatasetBase {
+class RebatchTabularDatasetV2Op::Dataset : public DatasetBase {
  public:
   Dataset(OpKernelContext* ctx, const DatasetBase* input_dataset,
           const int64 batch_size, const bool drop_remainder,
           const int64 shuffle_buffer_size, const int64 shuffle_seed,
           const int64 shuffle_seed2, const bool reshuffle_each_iteration,
-          const std::vector<int>& output_kinds)
+          const std::vector<int>& field_ids,
+          const std::vector<int>& field_ragged_indices)
       : DatasetBase(DatasetContext(ctx)),
         input_dataset_(input_dataset),
         batch_size_(batch_size),
@@ -106,27 +127,46 @@ class RectifyDatasetOp::Dataset : public DatasetBase {
         shuffle_seed_(shuffle_seed),
         shuffle_seed2_(shuffle_seed2),
         reshuffle_each_iteration_(reshuffle_each_iteration),
-        output_kinds_(output_kinds) {
+        field_ids_(field_ids),
+        field_ragged_indices_(field_ragged_indices) {
     input_dataset_->Ref();
 
     const auto& input_shapes = input_dataset_->output_shapes();
     output_shapes_.reserve(input_shapes.size());
-    for (size_t i = 0; i < input_shapes.size(); ++i) {
-      auto kind = output_kinds[i];
-      if (kind == kTensorOrSparseTensorValues) {
-        auto partial_shape = input_shapes[i];
-        partial_shape.RemoveDim(0);
-        if (drop_remainder_) {
-          output_shapes_.emplace_back(
-              PartialTensorShape({batch_size_}).Concatenate(partial_shape));
-        } else {
-          output_shapes_.emplace_back(
-              PartialTensorShape({-1}).Concatenate(partial_shape));
-        }
+    for (const auto& input_shape : input_shapes) {
+      auto partial_shape = input_shape;
+      partial_shape.RemoveDim(0);
+      if (drop_remainder_) {
+        output_shapes_.emplace_back(
+            PartialTensorShape({batch_size_}).Concatenate(partial_shape));
       } else {
-        output_shapes_.emplace_back(input_shapes[i]);
+        output_shapes_.emplace_back(
+            PartialTensorShape({-1}).Concatenate(partial_shape));
+      }
+    }
+
+    int64 prev_field_id = -1;
+    int64 prev_ragged_index = -1;
+    for (size_t i = 0; i < field_ids.size(); ++i) {
+      const int64 field_id = field_ids[i];
+      const int64 ragged_index = field_ragged_indices[i];
+      auto output_dtype = output_dtypes()[i];
+      if (TF_PREDICT_FALSE(ragged_index > 1 && output_dtype != DT_INT32)) {
+        LOG(ERROR) << "Output tensor " << i << " must be DT_INT32 not "
+                   << DataTypeString(output_dtype);
+        return;
+      }
+      if (ragged_index <= prev_ragged_index) {
+        if (TF_PREDICT_FALSE(field_id == prev_field_id)) {
+          LOG(ERROR) << "Invalid `fields` serialization";
+          return;
+        }
+        field_ranks_.push_back(prev_ragged_index);
       }
+      prev_field_id = field_id;
+      prev_ragged_index = ragged_index;
     }
+    field_ranks_.push_back(prev_ragged_index);
   }
 
   ~Dataset() override { input_dataset_->Unref(); }
@@ -142,20 +182,8 @@ class RectifyDatasetOp::Dataset : public DatasetBase {
     return output_shapes_;
   }
 
-  const std::vector<int>& output_kinds() const { return output_kinds_; }
-
-  string DebugString() const override { return "RectifyDatasetOp::Dataset"; }
-
-  int64 Cardinality() const override {
-    int64 n = input_dataset_->Cardinality();
-    if (n == data::kInfiniteCardinality || n == data::kUnknownCardinality) {
-      return n;
-    }
-    return data::kUnknownCardinality;
-  }
-
-  Status CheckExternalState() const override {
-    return input_dataset_->CheckExternalState();
+  string DebugString() const override {
+    return "RebatchTabularDatasetV2Op::Dataset";
   }
 
  protected:
@@ -177,8 +205,10 @@ class RectifyDatasetOp::Dataset : public DatasetBase {
     b->BuildAttrValue(drop_remainder_, &drop_remainder);
     AttrValue reshuffle_each_iteration;
     b->BuildAttrValue(reshuffle_each_iteration_, &reshuffle_each_iteration);
-    AttrValue output_kinds;
-    b->BuildAttrValue(output_kinds_, &output_kinds);
+    AttrValue field_ids;
+    b->BuildAttrValue(field_ids_, &field_ids);
+    AttrValue field_ragged_indices;
+    b->BuildAttrValue(field_ragged_indices_, &field_ragged_indices);
 
     TF_RETURN_IF_ERROR(
         b->AddDataset(this,
@@ -190,7 +220,8 @@ class RectifyDatasetOp::Dataset : public DatasetBase {
                       {},
                       {{"drop_remainder", drop_remainder},
                        {"reshuffle_each_iteration", reshuffle_each_iteration},
-                       {"output_kinds", output_kinds}},
+                       {"field_ids", field_ids},
+                       {"field_ragged_indices", field_ragged_indices}},
                       output));
     return Status::OK();
   }
@@ -204,17 +235,21 @@ class RectifyDatasetOp::Dataset : public DatasetBase {
   const int64 shuffle_seed_;
   const int64 shuffle_seed2_;
   const bool reshuffle_each_iteration_;
+  const std::vector<int> field_ids_;
+  const std::vector<int> field_ragged_indices_;
 
-  const std::vector<int> output_kinds_;
   std::vector<PartialTensorShape> output_shapes_;
+  std::vector<int32> field_ranks_;
 };
 
-void RectifyDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
-                                   DatasetBase** output) {
+void RebatchTabularDatasetV2Op::MakeDataset(OpKernelContext* ctx,
+                                            DatasetBase* input,
+                                            DatasetBase** output) {
   int64 batch_size = 0;
   OP_REQUIRES_OK(ctx, PARSE_SCALAR(ctx, "batch_size", &batch_size));
   OP_REQUIRES(ctx, batch_size > 0,
               errors::InvalidArgument("batch_size must be greater than zero."));
+
   int64 shuffle_buffer_size = 0;
   OP_REQUIRES_OK(
       ctx, PARSE_SCALAR(ctx, "shuffle_buffer_size", &shuffle_buffer_size));
@@ -223,30 +258,28 @@ void RectifyDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
   int64 shuffle_seed2 = 0;
   OP_REQUIRES_OK(ctx, PARSE_SCALAR(ctx, "shuffle_seed2", &shuffle_seed2));
 
-  *output = new Dataset(ctx, input, batch_size, drop_remainder_,
-                        shuffle_buffer_size, shuffle_seed, shuffle_seed2,
-                        reshuffle_each_iteration_, output_kinds_);
+  *output =
+      new Dataset(ctx, input, batch_size, drop_remainder_, shuffle_buffer_size,
+                  shuffle_seed, shuffle_seed2, reshuffle_each_iteration_,
+                  field_ids_, field_ragged_indices_);
 }
 
-class RectifyDatasetOp::Dataset::Iterator
-    : public DatasetIterator<RectifyDatasetOp::Dataset> {
+class RebatchTabularDatasetV2Op::Dataset::Iterator
+    : public DatasetIterator<RebatchTabularDatasetV2Op::Dataset> {
  public:
   explicit Iterator(const Params& params)
-      : DatasetIterator<RectifyDatasetOp::Dataset>(params),
-        queue_(dataset()->shuffle_buffer_size_, dataset()->shuffle_seed_,
-               dataset()->shuffle_seed2_, dataset()->reshuffle_each_iteration_,
-               dataset()->output_dtypes(), dataset()->output_shapes(),
-               dataset()->output_kinds()),
-        first_values_idx_(0) {
-    for (first_values_idx_ = 0;
-         first_values_idx_ < dataset()->output_kinds().size();
-         ++first_values_idx_) {
-      if (dataset()->output_kinds()[first_values_idx_] ==
-          kTensorOrSparseTensorValues) {
-        break;
-      }
+      : DatasetIterator<RebatchTabularDatasetV2Op::Dataset>(params),
+        buffer_(dataset()->output_dtypes(), dataset()->output_shapes(),
+                dataset()->field_ranks_),
+        shuffle_seed_(dataset()->shuffle_seed_),
+        shuffle_seed2_(dataset()->shuffle_seed2_),
+        parent_generator_(dataset()->shuffle_seed_, dataset()->shuffle_seed2_),
+        generator_(&parent_generator_) {
+    if (dataset()->reshuffle_each_iteration_) {
+      // TODO: support restore
+      shuffle_seed_ = generator_();
+      shuffle_seed2_ = generator_();
     }
-    DCHECK(first_values_idx_ < dataset()->output_kinds().size());
   }
 
   Status Initialize(IteratorContext* ctx) override {
@@ -259,32 +292,23 @@ class RectifyDatasetOp::Dataset::Iterator
     mutex_lock l(mu_);
     Allocator* alloc = ctx->allocator({});
 
-    int64 batch_size = dataset()->batch_size_;
-    int64 shuffle_buffer_size =
-        std::max(batch_size, dataset()->shuffle_buffer_size_);
-    if (queue_.size() >= shuffle_buffer_size) {
-      return queue_.Pop(batch_size, output_tensors, alloc);
-    }
-
-    if (!input_impl_) {
-      if (queue_.size() > batch_size) {
-        *end_of_sequence = false;
-        TF_RETURN_IF_ERROR(queue_.Pop(batch_size, output_tensors, alloc));
-        return Status::OK();
+    if (buffer_.size() >= dataset()->batch_size_ &&
+        buffer_.size() >= dataset()->shuffle_buffer_size_) {
+      if (dataset()->shuffle_buffer_size_ > 0) {
+        TF_RETURN_IF_ERROR(buffer_.Shuffle(generator_, dataset()->batch_size_));
       }
-      if (!dataset()->drop_remainder_ && queue_.size() > 0) {
-        *end_of_sequence = false;
-        TF_RETURN_IF_ERROR(queue_.Pop(queue_.size(), output_tensors, alloc));
-        return Status::OK();
-      }
-      *end_of_sequence = true;
-      return Status::OK();
+      return buffer_.Take(alloc, output_tensors, dataset()->batch_size_);
     }
 
     *end_of_sequence = false;
-    while (!*end_of_sequence) {
-      if (queue_.size() >= shuffle_buffer_size) {
-        return queue_.Pop(batch_size, output_tensors, alloc);
+    while (input_impl_ && !*end_of_sequence) {
+      if (buffer_.size() >= dataset()->batch_size_ &&
+          buffer_.size() >= dataset()->shuffle_buffer_size_) {
+        if (dataset()->shuffle_buffer_size_ > 0) {
+          TF_RETURN_IF_ERROR(
+              buffer_.Shuffle(generator_, dataset()->batch_size_));
+        }
+        return buffer_.Take(alloc, output_tensors, dataset()->batch_size_);
       }
 
       std::vector<Tensor> input_tensors;
@@ -298,32 +322,56 @@ class RectifyDatasetOp::Dataset::Iterator
                 "component.");
           }
         }
-        const int64 input_batch_size =
-            input_tensors[first_values_idx_].dim_size(0);
+
+        // Compute input batch size
+        int64 input_batch_size = 0;
+        if (TF_PREDICT_TRUE(input_tensors.size() > 0)) {
+          if (dataset()->field_ranks_[0] > 0) {
+            input_batch_size = input_tensors[1].dim_size(0) - 1;
+          } else {
+            input_batch_size = input_tensors[0].dim_size(0);
+          }
+        }
+
         // Fast path for same batch size
-        static const bool kRectifyFastPathDisabled =
-            ::hybridbackend::EnvVarGetBool(kRectifyDatasetFastPathEnv, false);
-        if (TF_PREDICT_TRUE(!kRectifyFastPathDisabled)) {
-          if (dataset()->shuffle_buffer_size_ < 1 && queue_.size() == 0 &&
-              batch_size == input_batch_size) {
+        static const bool kRebatchFaspathDisabled =
+            ::hybridbackend::EnvVarGetBool(kRebatchDatasetFastPathEnv, false);
+        if (TF_PREDICT_TRUE(!kRebatchFaspathDisabled)) {
+          if (dataset()->shuffle_buffer_size_ < 1 && buffer_.size() == 0 &&
+              dataset()->batch_size_ == input_batch_size) {
             *output_tensors = std::move(input_tensors);
             return Status::OK();
           }
         }
-        queue_.Push(input_batch_size, std::move(input_tensors));
+
+        if (dataset()->shuffle_buffer_size_ > 0) {
+          // Insert input rows to buffer
+          for (int64 row = 0; row < input_batch_size; ++row) {
+            TF_RETURN_IF_ERROR(buffer_.PutSlice(input_tensors, row, row + 1));
+          }
+        } else {
+          // Insert input batch to buffer
+          TF_RETURN_IF_ERROR(
+              buffer_.Put(std::move(input_tensors), input_batch_size));
+        }
       }
     }
-    if (queue_.size() > batch_size) {
+    input_impl_.reset();
+    if (buffer_.size() > dataset()->batch_size_) {
       *end_of_sequence = false;
-      input_impl_.reset();
-      return queue_.Pop(batch_size, output_tensors, alloc);
+      if (dataset()->shuffle_buffer_size_ > 0) {
+        TF_RETURN_IF_ERROR(buffer_.Shuffle(generator_, dataset()->batch_size_));
+      }
+      return buffer_.Take(alloc, output_tensors, dataset()->batch_size_);
     }
-    if (!dataset()->drop_remainder_ && queue_.size() > 0) {
+    if (buffer_.size() > 0 && !dataset()->drop_remainder_) {
       *end_of_sequence = false;
-      input_impl_.reset();
-      return queue_.Pop(queue_.size(), output_tensors, alloc);
+      if (dataset()->shuffle_buffer_size_ > 0) {
+        TF_RETURN_IF_ERROR(buffer_.Shuffle(generator_, buffer_.size()));
+      }
+      return buffer_.Take(alloc, output_tensors, buffer_.size());
     }
-    input_impl_.reset();
+    *end_of_sequence = true;
     return Status::OK();
   }
 
@@ -340,20 +388,25 @@ class RectifyDatasetOp::Dataset::Iterator
  private:
   mutex mu_;
   std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
-  RectifyQueue queue_ GUARDED_BY(mu_);
-  int64 first_values_idx_ GUARDED_BY(mu_);
+  RebatchBuffer buffer_ GUARDED_BY(mu_);
+
+  int64 shuffle_seed_;
+  int64 shuffle_seed2_;
+  random::PhiloxRandom parent_generator_;
+  random::SingleSampleAdapter<random::PhiloxRandom> generator_;
 };
 
-std::unique_ptr<IteratorBase> RectifyDatasetOp::Dataset::MakeIteratorInternal(
+std::unique_ptr<IteratorBase>
+RebatchTabularDatasetV2Op::Dataset::MakeIteratorInternal(
     const string& prefix) const {
   return std::unique_ptr<IteratorBase>(
-      new Iterator({this, absl::StrCat(prefix, "::Rectify")}));
+      new Iterator({this, absl::StrCat(prefix, "::RebatchTabularV2")}));
 }
 
-REGISTER_KERNEL_BUILDER(Name("HbRectifyDataset").Device(DEVICE_CPU),
-                        RectifyDatasetOp);
+REGISTER_KERNEL_BUILDER(Name("HbRebatchTabularDatasetV2").Device(DEVICE_CPU),
+                        RebatchTabularDatasetV2Op);
 
-WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS("HbRectifyDataset");
+WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS("HbRebatchTabularDatasetV2");
 
 }  // namespace hybridbackend
 }  // namespace tensorflow
diff --git a/hybridbackend/tensorflow/data/rectify/__init__.py b/hybridbackend/tensorflow/data/rectify/__init__.py
deleted file mode 100644
index 7489ae72..00000000
--- a/hybridbackend/tensorflow/data/rectify/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. 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.
-# =============================================================================
-
-r'''Rectifying related classes and functions.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
diff --git a/hybridbackend/tensorflow/data/rectify/dataset.py b/hybridbackend/tensorflow/data/rectify/dataset.py
deleted file mode 100644
index 4fa8960f..00000000
--- a/hybridbackend/tensorflow/data/rectify/dataset.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. 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.
-# =============================================================================
-
-r'''Dataset that resizes batches of tensors.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import inspect
-
-# pylint: disable=ungrouped-imports
-try:
-  from tensorflow.python.data.ops.dataset_ops import DatasetV2 as _dataset  # pylint: disable=unused-import
-
-  from hybridbackend.tensorflow.data.rectify.dataset_v2 import \
-    RectifyDatasetV2 as RectifyDataset
-  if inspect.isabstract(RectifyDataset):
-    raise ImportError
-  RectifyDataset.__module__ = __name__
-  RectifyDataset.__name__ = 'RectifyDataset'
-except ImportError:
-  from hybridbackend.tensorflow.data.rectify.dataset_v1 import \
-    RectifyDatasetV1 as RectifyDataset
-  RectifyDataset.__module__ = __name__
-  RectifyDataset.__name__ = 'RectifyDataset'
-  assert not inspect.isabstract(RectifyDataset)
-# pylint: enable=ungrouped-imports
-
-
-def rectify(batch_size, shuffle_buffer_size=None, drop_remainder=False):
-  r'''Create a `RectifyDataset`.
-
-  Args:
-    batch_size: Maxium number of samples in an output batch.
-    shuffle_buffer_size: (Optional.) A `tf.int64` scalar `tf.Tensor`,
-      representing the number of elements from this dataset from which the new
-      dataset will sample.
-    drop_remainder: (Optional.) If True, smaller final batch is dropped.
-      `False` by default.
-  '''
-  def _apply_fn(dataset):
-    return RectifyDataset(
-      dataset, batch_size,
-      shuffle_buffer_size=shuffle_buffer_size,
-      drop_remainder=drop_remainder)
-  return _apply_fn
diff --git a/hybridbackend/tensorflow/data/rectify/dataset_v1.py b/hybridbackend/tensorflow/data/rectify/dataset_v1.py
deleted file mode 100644
index 0c31f1a6..00000000
--- a/hybridbackend/tensorflow/data/rectify/dataset_v1.py
+++ /dev/null
@@ -1,217 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. 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.
-# =============================================================================
-
-r'''Dataset that resizes batches of tensors.
-
-This class is compatible with TensorFlow 1.12.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import random_seed
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.framework import tensor_util
-from tensorflow.python.util import nest
-
-from hybridbackend.tensorflow.common import oplib as _ops
-from hybridbackend.tensorflow.framework.ops import TensorKinds
-
-
-class _RectifyDatasetV1(dataset_ops.Dataset):
-  r'''A dataset that adjusts batches.
-  '''
-  def __init__(
-      self, input_dataset, output_kinds, batch_size,
-      drop_remainder=False,
-      shuffle_buffer_size=None,
-      shuffle_seed=None,
-      reshuffle_each_iteration=None):
-    r'''Create a `RectifyDatasetV2`.
-
-    Args:
-      input_dataset: A dataset outputs batches.
-      batch_size: Maxium number of samples in an output batch.
-      drop_remainder: (Optional.) If True, smaller final batch is dropped.
-        `False` by default.
-      shuffle_buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the
-        number of elements from this dataset from which the new
-        dataset will sample.
-      shuffle_seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing
-        the random seed that will be used to create the distribution. See
-        @{tf.set_random_seed} for behavior.
-      reshuffle_each_iteration: (Optional.) A boolean, which if true indicates
-        that the dataset should be pseudorandomly reshuffled each time it is
-        iterated over. (Defaults to `True`.)
-    '''
-    self._input_dataset = input_dataset
-    self._batch_size = batch_size
-    self._drop_remainder = drop_remainder
-    self._shuffle_buffer_size = ops.convert_to_tensor(
-      shuffle_buffer_size or 0, dtype=dtypes.int64, name='buffer_size')
-    self._seed, self._seed2 = random_seed.get_seed(shuffle_seed)
-    self._reshuffle_each_iteration = reshuffle_each_iteration or True
-    self._output_kinds = output_kinds
-
-    dim0_shape = tensor_shape.TensorShape([
-      tensor_util.constant_value(self._batch_size)
-      if self._drop_remainder
-      else None])
-    self._output_shapes = []
-    for idx, kind in enumerate(self._output_kinds):
-      input_shape = self._input_dataset.output_shapes[idx]
-      if kind == TensorKinds.INDICES:
-        self._output_shapes.append(input_shape)
-      elif kind == TensorKinds.VALUES:
-        self._output_shapes.append(dim0_shape.concatenate(input_shape[1:]))
-      elif kind == TensorKinds.DENSE_SHAPE:
-        self._output_shapes.append(input_shape)
-      else:
-        raise ValueError('Unknown tensor kind: ', kind)
-    self._output_shapes = tuple(self._output_shapes)
-
-    super().__init__()
-
-  def _inputs(self):
-    return [self._input_dataset]
-
-  def _as_variant_tensor(self):
-    return _ops.hb_rectify_dataset(
-      self._input_dataset._as_variant_tensor(),  # pylint: disable=protected-access
-      self._batch_size,
-      self._shuffle_buffer_size,
-      self._seed,
-      self._seed2,
-      drop_remainder=self._drop_remainder,
-      reshuffle_each_iteration=self._reshuffle_each_iteration,
-      output_kinds=self._output_kinds,
-      output_types=self.output_types,
-      output_shapes=self.output_shapes)
-
-  @property
-  def output_classes(self):
-    return self._input_dataset.output_classes
-
-  @property
-  def output_types(self):
-    return self._input_dataset.output_types
-
-  @property
-  def output_shapes(self):
-    r'''Output shapes.
-    '''
-    return self._output_shapes
-
-
-class RectifyDatasetV1(dataset_ops.Dataset):
-  r'''A dataset that adjusts batches.
-  '''
-  VERSION = 2001
-
-  def __init__(
-      self, input_dataset, batch_size,
-      drop_remainder=False,
-      shuffle_buffer_size=None,
-      shuffle_seed=None,
-      reshuffle_each_iteration=None):
-    r'''Create a `RectifyDatasetV1`.
-
-    Args:
-      input_dataset: A dataset outputs batches.
-      batch_size: Maxium number of samples in an output batch.
-      drop_remainder: (Optional.) If True, smaller final batch is dropped.
-        `False` by default.
-      shuffle_buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the
-        number of elements from this dataset from which the new
-        dataset will sample.
-      shuffle_seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing
-        the random seed that will be used to create the distribution. See
-        @{tf.set_random_seed} for behavior.
-      reshuffle_each_iteration: (Optional.) A boolean, which if true indicates
-        that the dataset should be pseudorandomly reshuffled each time it is
-        iterated over. (Defaults to `True`.)
-    '''
-    flat_shapes = input_dataset._flat_shapes  # pylint: disable=protected-access
-    if any(s.ndims == 0 for s in flat_shapes):
-      raise ValueError('Cannot rectify an input with scalar components.')
-    known_batch_dim = tensor_shape.Dimension(None)
-    for s in flat_shapes:
-      try:
-        known_batch_dim = known_batch_dim.merge_with(s[0])
-      except ValueError as ex:
-        raise ValueError(
-          'Cannot rectify an input whose components have different '
-          'batch sizes.') from ex
-
-    self._input_dataset = input_dataset
-    self._batch_size = ops.convert_to_tensor(
-      batch_size,
-      dtype=dtypes.int64,
-      name='batch_size')
-    self._drop_remainder = drop_remainder
-
-    flattened_classes = nest.flatten(input_dataset.output_classes)
-    flattened_kinds = []
-    for cls in flattened_classes:
-      if cls == ops.Tensor:
-        flattened_kinds.append(TensorKinds.VALUES)
-      elif cls == sparse_tensor.SparseTensor:
-        flattened_kinds.append(
-          sparse_tensor.SparseTensorValue(
-            TensorKinds.INDICES, TensorKinds.VALUES, TensorKinds.DENSE_SHAPE))
-      else:
-        raise ValueError(
-          'hb.data.rectify cannot support input datasets with outputs other '
-          'than tensors or sparse tensors')
-    normalized_kinds = nest.flatten(flattened_kinds)
-    normalized_dataset = input_dataset.map(TensorKinds.normalize)
-    rectify_dataset = _RectifyDatasetV1(  # pylint: disable=abstract-class-instantiated
-      normalized_dataset, normalized_kinds,
-      batch_size, drop_remainder=drop_remainder,
-      shuffle_buffer_size=shuffle_buffer_size,
-      shuffle_seed=shuffle_seed,
-      reshuffle_each_iteration=reshuffle_each_iteration)
-    self._impl = rectify_dataset.map(
-      lambda *args: TensorKinds.denormalize(
-        input_dataset.element_spec, flattened_kinds, args))
-
-    super().__init__()
-
-  @property
-  def drop_remainder(self):
-    return self._drop_remainder
-
-  def _as_variant_tensor(self):
-    return self._impl._as_variant_tensor()  # pylint: disable=protected-access
-
-  def _inputs(self):
-    return [self._input_dataset]
-
-  @property
-  def output_classes(self):
-    return self._impl.output_classes
-
-  @property
-  def output_types(self):
-    return self._impl.output_types
-
-  @property
-  def output_shapes(self):
-    return self._impl.output_shapes
diff --git a/hybridbackend/tensorflow/data/rectify/dataset_v2.py b/hybridbackend/tensorflow/data/rectify/dataset_v2.py
deleted file mode 100644
index dc881dcd..00000000
--- a/hybridbackend/tensorflow/data/rectify/dataset_v2.py
+++ /dev/null
@@ -1,197 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. 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.
-# =============================================================================
-
-r'''Dataset that resizes batches of tensors.
-
-This class is compatible with TensorFlow 1.15.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import random_seed
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.framework import tensor_spec
-from tensorflow.python.framework import tensor_util
-from tensorflow.python.util import nest
-
-from hybridbackend.tensorflow.common import oplib as _ops
-from hybridbackend.tensorflow.framework.ops import TensorKinds
-
-
-class _RectifyDatasetV2(dataset_ops.DatasetV2):
-  r'''A dataset that adjusts batches.
-  '''
-  def __init__(
-      self, input_dataset, output_kinds, batch_size,
-      drop_remainder=False,
-      shuffle_buffer_size=None,
-      shuffle_seed=None,
-      reshuffle_each_iteration=None):
-    r'''Create a `RectifyDatasetV2`.
-
-    Args:
-      input_dataset: A dataset outputs batches.
-      batch_size: Maxium number of samples in an output batch.
-      drop_remainder: (Optional.) If True, smaller final batch is dropped.
-        `False` by default.
-      shuffle_buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the
-        number of elements from this dataset from which the new
-        dataset will sample.
-      shuffle_seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing
-        the random seed that will be used to create the distribution. See
-        @{tf.set_random_seed} for behavior.
-      reshuffle_each_iteration: (Optional.) A boolean, which if true indicates
-        that the dataset should be pseudorandomly reshuffled each time it is
-        iterated over. (Defaults to `True`.)
-    '''
-    self._input_dataset = input_dataset
-    self._batch_size = batch_size
-    self._drop_remainder = drop_remainder
-    self._shuffle_buffer_size = ops.convert_to_tensor(
-      shuffle_buffer_size or 0, dtype=dtypes.int64, name='buffer_size')
-    self._seed, self._seed2 = random_seed.get_seed(shuffle_seed)
-    self._reshuffle_each_iteration = reshuffle_each_iteration or True
-    self._output_kinds = output_kinds
-
-    dim0_shape = tensor_shape.TensorShape([
-      tensor_util.constant_value(self._batch_size)
-      if self._drop_remainder
-      else None])
-    self._element_spec = []
-    for idx, kind in enumerate(self._output_kinds):
-      spec = self._input_dataset.element_spec[idx]
-      if kind == TensorKinds.INDICES:
-        self._element_spec.append(spec)
-      elif kind == TensorKinds.VALUES:
-        self._element_spec.append(
-          tensor_spec.TensorSpec(
-            dim0_shape.concatenate(spec.shape[1:]), spec.dtype, spec.name))
-      elif kind == TensorKinds.DENSE_SHAPE:
-        self._element_spec.append(spec)
-      else:
-        raise ValueError('Unknown tensor kind: ', kind)
-    self._element_spec = tuple(self._element_spec)
-
-    super().__init__(
-      _ops.hb_rectify_dataset(
-        input_dataset._variant_tensor,  # pylint: disable=protected-access
-        batch_size,
-        self._shuffle_buffer_size,
-        self._seed,
-        self._seed2,
-        drop_remainder=drop_remainder,
-        reshuffle_each_iteration=self._reshuffle_each_iteration,
-        output_kinds=output_kinds,
-        **self._flat_structure))
-
-  def _inputs(self):
-    return [self._input_dataset]
-
-  @property
-  def element_spec(self):
-    r'''Element Specification.
-    '''
-    return self._element_spec
-
-
-class RectifyDatasetV2(dataset_ops.DatasetV2):  # pylint: disable=abstract-method
-  r'''A dataset that adjusts batches.
-  '''
-  VERSION = 2002
-
-  def __init__(
-      self, input_dataset, batch_size,
-      drop_remainder=False,
-      shuffle_buffer_size=None,
-      shuffle_seed=None,
-      reshuffle_each_iteration=None):
-    r'''Create a `RectifyDatasetV2`.
-
-    Args:
-      input_dataset: A dataset outputs batches.
-      batch_size: Maxium number of samples in an output batch.
-      drop_remainder: (Optional.) If True, smaller final batch is dropped.
-        `False` by default.
-      shuffle_buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the
-        number of elements from this dataset from which the new
-        dataset will sample.
-      shuffle_seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing
-        the random seed that will be used to create the distribution. See
-        @{tf.set_random_seed} for behavior.
-      reshuffle_each_iteration: (Optional.) A boolean, which if true indicates
-        that the dataset should be pseudorandomly reshuffled each time it is
-        iterated over. (Defaults to `True`.)
-    '''
-    flat_shapes = input_dataset._flat_shapes  # pylint: disable=protected-access
-    if any(s.ndims == 0 for s in flat_shapes):
-      raise ValueError('Cannot rectify an input with scalar components.')
-    known_batch_dim = tensor_shape.Dimension(None)
-    for s in flat_shapes:
-      try:
-        known_batch_dim = known_batch_dim.merge_with(s[0])
-      except ValueError as ex:
-        raise ValueError(
-          'Cannot rectify an input whose components have different '
-          'batch sizes.') from ex
-
-    self._input_dataset = input_dataset
-    self._batch_size = ops.convert_to_tensor(
-      batch_size,
-      dtype=dtypes.int64,
-      name='batch_size')
-    self._drop_remainder = drop_remainder
-
-    flattened_specs = nest.flatten(input_dataset.element_spec)
-    flattened_kinds = []
-    for spec in flattened_specs:
-      if isinstance(spec, tensor_spec.TensorSpec):
-        flattened_kinds.append(TensorKinds.VALUES)
-      elif isinstance(spec, sparse_tensor.SparseTensorSpec):
-        flattened_kinds.append(
-          sparse_tensor.SparseTensorValue(
-            TensorKinds.INDICES, TensorKinds.VALUES, TensorKinds.DENSE_SHAPE))
-      else:
-        raise ValueError(
-          'hb.data.rectify cannot support input datasets with outputs other '
-          'than tensors or sparse tensors')
-    normalized_kinds = nest.flatten(flattened_kinds)
-    normalized_dataset = input_dataset.map(TensorKinds.normalize)
-    rectify_dataset = _RectifyDatasetV2(  # pylint: disable=abstract-class-instantiated
-      normalized_dataset, normalized_kinds,
-      batch_size, drop_remainder=drop_remainder,
-      shuffle_buffer_size=shuffle_buffer_size,
-      shuffle_seed=shuffle_seed,
-      reshuffle_each_iteration=reshuffle_each_iteration)
-    self._impl = rectify_dataset.map(
-      lambda *args: TensorKinds.denormalize(
-        input_dataset.element_spec, flattened_kinds, args))
-    super().__init__(self._impl._variant_tensor)  # pylint: disable=protected-access
-
-  @property
-  def drop_remainder(self):
-    return self._drop_remainder
-
-  def _inputs(self):
-    return [self._input_dataset]
-
-  @property
-  def element_spec(self):
-    return self._impl.element_spec  # pylint: disable=protected-access
diff --git a/hybridbackend/tensorflow/data/rectify/queue.h b/hybridbackend/tensorflow/data/rectify/queue.h
deleted file mode 100644
index 6db804ff..00000000
--- a/hybridbackend/tensorflow/data/rectify/queue.h
+++ /dev/null
@@ -1,76 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. 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 HYBRIDBACKEND_TENSORFLOW_DATA_RECTIFY_QUEUE_H_
-#define HYBRIDBACKEND_TENSORFLOW_DATA_RECTIFY_QUEUE_H_
-
-#include <deque>
-#include <vector>
-
-#include <tensorflow/core/framework/tensor.h>
-#include <tensorflow/core/lib/random/philox_random.h>
-#include <tensorflow/core/lib/random/random.h>
-#include <tensorflow/core/lib/random/random_distributions.h>
-
-namespace tensorflow {
-namespace hybridbackend {
-
-enum TensorKinds {
-  kSparseTensorIndices = 0,
-  kTensorOrSparseTensorValues = 1,
-  kSparseTensorDenseShape = 2
-};
-
-struct RectifyQueueItem {
- public:
-  RectifyQueueItem(int64 batch_size, const std::vector<Tensor>& components)
-      : batch_size(batch_size), components(components) {}
-  int64 batch_size;
-  std::vector<Tensor> components;
-};
-
-class RectifyQueue {
- public:
-  RectifyQueue(const int64 shuffle_buffer_size, const int64 shuffle_seed,
-               const int64 shuffle_seed2, const bool reshuffle_each_iteration,
-               const DataTypeVector& output_dtypes,
-               const std::vector<PartialTensorShape>& output_shapes,
-               const std::vector<int>& output_kinds);
-
-  int64 size() const { return size_; }
-
-  Status Push(const int64 input_batch_size,
-              const std::vector<Tensor>& input_tensors);
-
-  Status Pop(const int64 output_batch_size, std::vector<Tensor>* output_tensors,
-             Allocator* alloc);
-
- private:
-  const int64 shuffle_buffer_size_;
-  const DataTypeVector& output_dtypes_;
-  const std::vector<PartialTensorShape>& output_shapes_;
-  const std::vector<int>& output_kinds_;
-  int64 shuffle_seed_;
-  int64 shuffle_seed2_;
-  int64 size_;
-  std::deque<RectifyQueueItem> queue_;
-  random::PhiloxRandom parent_generator_;
-  random::SingleSampleAdapter<random::PhiloxRandom> generator_;
-};
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // HYBRIDBACKEND_TENSORFLOW_DATA_RECTIFY_QUEUE_H_
diff --git a/hybridbackend/tensorflow/data/rectify/rectify_queue.cc b/hybridbackend/tensorflow/data/rectify/rectify_queue.cc
deleted file mode 100644
index 961118d9..00000000
--- a/hybridbackend/tensorflow/data/rectify/rectify_queue.cc
+++ /dev/null
@@ -1,286 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. 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/core/framework/register_types.h>
-#include <tensorflow/core/public/version.h>
-
-#include <algorithm>
-#include <deque>
-#include <vector>
-
-#include "hybridbackend/common/env.h"
-#include "hybridbackend/tensorflow/data/rectify/queue.h"
-
-namespace tensorflow {
-namespace hybridbackend {
-
-namespace {
-template <typename T>
-Status CopyValues(T* dest, const T* src, int64 num_values,
-                  bool /* can_move */) {
-  static_assert(is_simple_type<T>::value, "Memcpy requires a simple type.");
-  memcpy(dest, src, num_values * sizeof(T));
-  return Status::OK();
-}
-
-template <>
-Status CopyValues<tstring>(tstring* dest, const tstring* src, int64 num_values,
-                           bool can_move) {
-  if (can_move) {
-    for (int64 i = 0; i < num_values; ++i) {
-      *dest++ = std::move(*src++);
-    }
-  } else {
-    std::copy_n(src, num_values, dest);
-  }
-  return Status::OK();
-}
-
-template <>
-Status CopyValues<Variant>(Variant* dest, const Variant* src, int64 num_values,
-                           bool can_move) {
-  if (can_move) {
-    for (int64 i = 0; i < num_values; ++i) {
-      *dest++ = std::move(*src++);
-    }
-  } else {
-    std::copy_n(src, num_values, dest);
-  }
-  return Status::OK();
-}
-
-template <>
-Status CopyValues<ResourceHandle>(ResourceHandle* dest,
-                                  const ResourceHandle* src, int64 num_values,
-                                  bool /* can_move */) {
-  std::copy_n(src, num_values, dest);
-  return Status::OK();
-}
-
-template <>
-Status CopyValues<Eigen::half>(Eigen::half* dest, const Eigen::half* src,
-                               int64 num_values, bool /* can_move */) {
-  std::copy_n(src, num_values, dest);
-  return Status::OK();
-}
-
-Status CopyToSlicesFromTensor(Tensor* dest, const int64 dest_index,
-                              const Tensor& src, const bool can_move) {
-  // can_move can be set by `src.RefCountIsOne();` if this function is a friend
-  // of `Tensor` class
-  if (src.dim_size(0) < 1) {
-    return Status::OK();
-  }
-  DCHECK_NE(dest->dim_size(0), 0);
-  DCHECK_NE(src.dim_size(0), 0);
-  DCHECK_GE(dest_index, 0);
-  const int64 num_values_per_record = dest->NumElements() / dest->dim_size(0);
-  if (TF_PREDICT_FALSE(num_values_per_record !=
-                       (src.NumElements() / src.dim_size(0)))) {
-    return errors::Internal(
-        "Cannot perform copy: shape of tensors does not match. "
-        " Shapes are: [src]: ",
-        src.shape().DebugString(), ", [dest]: ", dest->shape().DebugString());
-  }
-  if (TF_PREDICT_FALSE(dest_index + src.dim_size(0) > dest->dim_size(0))) {
-    return errors::Internal(
-        "Cannot perform copy: shape of dest tensor is too small. "
-        " Shapes are: [src]: ",
-        src.shape().DebugString(), ", [dest]: ", dest->shape().DebugString(),
-        " Destination index is ", dest_index);
-  }
-
-#define HANDLE_TYPE(T)                                                    \
-  case DataTypeToEnum<T>::value: {                                        \
-    const T* src_ptr = src.unaligned_flat<T>().data();                    \
-    T* dest_ptr = dest->unaligned_flat<T>().data() +                      \
-                  (num_values_per_record * dest_index);                   \
-    return CopyValues<T>(dest_ptr, src_ptr, src.NumElements(), can_move); \
-  }
-
-  switch (src.dtype()) {
-    TF_CALL_ALL_TYPES(HANDLE_TYPE);
-    TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
-#undef HANDLE_TYPE
-    default:
-      return errors::Unimplemented(
-          "CopyToSlicesFromTensor Unhandled data type: ", src.dtype());
-  }
-}
-}  // namespace
-
-RectifyQueue::RectifyQueue(const int64 shuffle_buffer_size,
-                           const int64 shuffle_seed, const int64 shuffle_seed2,
-                           const bool reshuffle_each_iteration,
-                           const DataTypeVector& output_dtypes,
-                           const std::vector<PartialTensorShape>& output_shapes,
-                           const std::vector<int>& output_kinds)
-    : shuffle_buffer_size_(shuffle_buffer_size),
-      output_dtypes_(output_dtypes),
-      output_shapes_(output_shapes),
-      output_kinds_(output_kinds),
-      shuffle_seed_(shuffle_seed),
-      shuffle_seed2_(shuffle_seed2),
-      size_(0),
-      parent_generator_(shuffle_seed, shuffle_seed2),
-      generator_(&parent_generator_) {
-  if (reshuffle_each_iteration) {
-    shuffle_seed_ = generator_();
-    shuffle_seed2_ = generator_();
-  }
-}
-
-Status RectifyQueue::Push(const int64 input_batch_size,
-                          const std::vector<Tensor>& input_tensors) {
-  if (shuffle_buffer_size_ > 0) {
-    for (int64 row = 0; row < input_batch_size; ++row) {
-      std::vector<Tensor> partial_input_tensors;
-      partial_input_tensors.reserve(output_dtypes_.size());
-      for (int64 col = 0; col < output_dtypes_.size(); ++col) {
-        if (output_kinds_[col] == kSparseTensorDenseShape) {
-          auto input_shape = input_tensors[col];
-          input_shape.flat<int64>()(0) = 1;
-          partial_input_tensors.push_back(std::move(input_shape));
-        } else {
-          partial_input_tensors.push_back(
-              std::move(input_tensors[col].Slice(row, row + 1)));
-        }
-      }
-      queue_.emplace_back(1, std::move(partial_input_tensors));
-    }
-
-    size_ += input_batch_size;
-    return Status::OK();
-  }
-
-  queue_.emplace_back(input_batch_size, input_tensors);
-  size_ += input_batch_size;
-  return Status::OK();
-}
-
-Status RectifyQueue::Pop(const int64 output_batch_size,
-                         std::vector<Tensor>* output_tensors,
-                         Allocator* alloc) {
-  const size_t num_components = output_dtypes_.size();
-  output_tensors->clear();
-  output_tensors->reserve(num_components);
-  for (size_t i = 0; i < num_components; ++i) {
-    PartialTensorShape output_pshape(output_shapes_[i]);
-    auto kind = output_kinds_[i];
-    if (kind != kSparseTensorDenseShape) {
-      output_pshape.set_dim(0, output_batch_size);
-    }
-    TensorShape output_shape;
-    output_pshape.AsTensorShape(&output_shape);
-    output_tensors->emplace_back(alloc, output_dtypes_[i], output_shape);
-    if (!output_tensors->back().IsInitialized()) {
-      return errors::ResourceExhausted(
-          "Failed to allocate memory for the batch of component ", i);
-    }
-  }
-
-  if (queue_.size() < 1) {
-    return Status::OK();
-  }
-
-  if (shuffle_buffer_size_ > 0) {
-    for (int64 row = 0; row < output_batch_size; ++row) {
-      int64 picked = row + generator_() % (queue_.size() - row);
-      std::swap(queue_[picked], queue_[row]);
-    }
-  }
-
-  int64 output_index = 0;
-  while (output_index < output_batch_size - queue_.front().batch_size) {
-    auto& front_item = queue_.front();
-    if (front_item.batch_size < 1) {
-      queue_.pop_front();
-      continue;
-    }
-    for (size_t i = 0; i < num_components; ++i) {
-      auto kind = output_kinds_[i];
-      if (kind == kTensorOrSparseTensorValues) {
-        TF_RETURN_IF_ERROR(
-            CopyToSlicesFromTensor(&(output_tensors->at(i)), output_index,
-                                   std::move(front_item.components[i]), true));
-      } else if (kind == kSparseTensorIndices) {
-        auto input = std::move(front_item.components[i]);
-        int64* input_ptr = input.unaligned_flat<int64>().data();
-        const int64 num_rows = input.shape().dim_size(0);
-        const int64 row_elements = input.shape().num_elements() / num_rows;
-        for (int64 row = 0; row < num_rows; ++row) {
-          input_ptr[row * row_elements] += (output_index - input_ptr[0]);
-        }
-        TF_RETURN_IF_ERROR(CopyToSlicesFromTensor(&(output_tensors->at(i)),
-                                                  output_index, input, true));
-      } else if (output_index == 0) {
-        output_tensors->at(i) = std::move(front_item.components[i]);
-      }
-    }
-    output_index += front_item.batch_size;
-    queue_.pop_front();
-  }
-
-  if (output_index < output_batch_size) {
-    auto& remained_item = queue_.front();
-    std::vector<Tensor> remained_components;
-    remained_components.reserve(num_components);
-    for (size_t i = 0; i < num_components; ++i) {
-      auto kind = output_kinds_[i];
-      if (kind == kTensorOrSparseTensorValues) {
-        auto sliced_input = remained_item.components[i].Slice(
-            0, output_batch_size - output_index);
-        TF_RETURN_IF_ERROR(CopyToSlicesFromTensor(
-            &(output_tensors->at(i)), output_index, sliced_input, true));
-        remained_components.push_back(remained_item.components[i].Slice(
-            output_batch_size - output_index, remained_item.batch_size));
-      } else if (kind == kSparseTensorIndices) {
-        auto sliced_input = remained_item.components[i].Slice(
-            0, output_batch_size - output_index);
-        int64* input_ptr = sliced_input.unaligned_flat<int64>().data();
-        const int64 num_rows = sliced_input.shape().dim_size(0);
-        const int64 row_elements =
-            sliced_input.shape().num_elements() / num_rows;
-        for (int64 row = 0; row < num_rows; ++row) {
-          input_ptr[row * row_elements] += (output_index - input_ptr[0]);
-        }
-        TF_RETURN_IF_ERROR(CopyToSlicesFromTensor(
-            &(output_tensors->at(i)), output_index, sliced_input, true));
-        remained_components.push_back(remained_item.components[i].Slice(
-            output_batch_size - output_index, remained_item.batch_size));
-      }
-    }
-    int64 remained_batch_size =
-        remained_item.batch_size - output_batch_size + output_index;
-    output_index += remained_item.batch_size;
-    queue_.pop_front();
-    queue_.emplace_front(remained_batch_size, remained_components);
-  }
-
-  for (size_t i = 0; i < num_components; ++i) {
-    auto kind = output_kinds_[i];
-    if (kind == kSparseTensorDenseShape) {
-      output_tensors->at(i).flat<int64>()(0) = output_batch_size;
-    }
-  }
-
-  size_ -= output_batch_size;
-  return Status::OK();
-}
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
diff --git a/hybridbackend/tensorflow/data/tabular/core.py b/hybridbackend/tensorflow/data/tabular/core.py
index 1041fba8..e2e4620e 100644
--- a/hybridbackend/tensorflow/data/tabular/core.py
+++ b/hybridbackend/tensorflow/data/tabular/core.py
@@ -29,7 +29,7 @@
 from tensorflow.python.framework import ops
 
 from hybridbackend.tensorflow.data.dataframe import parse
-from hybridbackend.tensorflow.data.rectify.dataset import RectifyDataset
+from hybridbackend.tensorflow.data.rebatch.dataset import RebatchDataset
 
 
 class TabularDatasetCreator(object):  # pylint: disable=useless-object-inheritance
@@ -136,22 +136,19 @@ def num_parallel_reads(self):
   def num_parallel_parser_calls(self):
     return self._num_parallel_reads
 
-  def read(self):
-    r'''Read records in the dataset.
+  def _create_dataset(self, batch_size):
+    r'''Create dataset for specific batch size in parallel.
     '''
     def _creator(filename):
       filename = ops.convert_to_tensor(filename, dtypes.string, name='filename')
-      ds = self._fn(filename, 256)
-      ds = ds.apply(parse(pad=self._sparse_to_dense))
-      return ds.unbatch()
+      return self._fn(filename, batch_size)
 
     if self._num_parallel_reads == 1:
       return self._filenames.flat_map(_creator)
     if (self._num_parallel_reads is None
         or self._num_parallel_reads == dataset_ops.AUTOTUNE):
       return self._filenames.interleave(
-        _creator,
-        num_parallel_calls=self._num_parallel_reads)
+        _creator, num_parallel_calls=self._num_parallel_reads)
     return readers.ParallelInterleaveDataset(
       self._filenames, _creator,
       cycle_length=self._num_parallel_reads,
@@ -160,27 +157,12 @@ def _creator(filename):
       buffer_output_elements=None,
       prefetch_input_elements=1)
 
-  def _create_batched_dataset(self, batch_size):
-    r'''Create batched dataset for specific batch size.
+  def read(self):
+    r'''Read records in the dataset.
     '''
-    def _creator(filename):
-      filename = ops.convert_to_tensor(filename, dtypes.string, name='filename')
-      ds = self._fn(filename, batch_size)
-      return ds.apply(parse(pad=self._sparse_to_dense))
-
-    if self._num_parallel_reads == 1:
-      return self._filenames.flat_map(_creator)
-    if (self._num_parallel_reads is None
-        or self._num_parallel_reads == dataset_ops.AUTOTUNE):
-      return self._filenames.interleave(
-        _creator, num_parallel_calls=self._num_parallel_reads)
-    return readers.ParallelInterleaveDataset(
-      self._filenames, _creator,
-      cycle_length=self._num_parallel_reads,
-      block_length=1,
-      sloppy=True,
-      buffer_output_elements=None,
-      prefetch_input_elements=1)
+    ds = self._create_dataset(256)  # Prefetch 256 rows to speed up
+    ds = ds.apply(parse(pad=self._sparse_to_dense))
+    return ds.unbatch()
 
   def batch(self, batch_size, drop_remainder=False):
     r'''Combines consecutive elements of this dataset into batches.
@@ -199,13 +181,17 @@ def batch(self, batch_size, drop_remainder=False):
         whether the last batch should be dropped in the case its has fewer than
         `batch_size` elements; the default behavior is not to drop the smaller
         batch.
+      allow_smaller: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
+        whether a batch could be smaller.
 
     Returns:
       Dataset: A `Dataset`.
     '''
-    return RectifyDataset(
-      self._create_batched_dataset(batch_size),
-      batch_size, drop_remainder=drop_remainder)
+    ds = self._create_dataset(batch_size)
+    ds = RebatchDataset(
+      ds, self._fields, batch_size,
+      drop_remainder=drop_remainder)
+    return ds.apply(parse(pad=self._sparse_to_dense))
 
   def shuffle_batch(
       self, batch_size,
@@ -238,10 +224,11 @@ def shuffle_batch(
     '''
     if buffer_size is None:
       buffer_size = batch_size
-    return RectifyDataset(
-      self._create_batched_dataset(batch_size),
-      batch_size,
+    ds = self._create_dataset(batch_size)
+    ds = RebatchDataset(
+      ds, self._fields, batch_size,
       drop_remainder=drop_remainder,
       shuffle_buffer_size=buffer_size,
       shuffle_seed=seed,
       reshuffle_each_iteration=reshuffle_each_iteration)
+    return ds.apply(parse(pad=self._sparse_to_dense))
diff --git a/hybridbackend/tensorflow/data/tabular/dataset.py b/hybridbackend/tensorflow/data/tabular/dataset.py
index 06917ba3..4003d563 100644
--- a/hybridbackend/tensorflow/data/tabular/dataset.py
+++ b/hybridbackend/tensorflow/data/tabular/dataset.py
@@ -23,12 +23,12 @@
 # pylint: disable=ungrouped-imports
 try:
   from hybridbackend.tensorflow.data.tabular.dataset_v2 import \
-    TabularDatasetV2 as TabularDataset
-  TabularDataset.__module__ = __name__
-  TabularDataset.__name__ = 'TabularDataset'
+    TabularDatasetV2 as Dataset
+  Dataset.__module__ = __name__
+  Dataset.__name__ = 'TabularDataset'
 except ImportError:
   from hybridbackend.tensorflow.data.tabular.dataset_v1 import \
-    TabularDatasetV1 as TabularDataset
-  TabularDataset.__module__ = __name__
-  TabularDataset.__name__ = 'TabularDataset'
+    TabularDatasetV1 as Dataset
+  Dataset.__module__ = __name__
+  Dataset.__name__ = 'TabularDataset'
 # pylint: enable=ungrouped-imports
diff --git a/hybridbackend/tensorflow/data/tabular/dataset_v1.py b/hybridbackend/tensorflow/data/tabular/dataset_v1.py
index 5c4ab52e..eaba3e6e 100644
--- a/hybridbackend/tensorflow/data/tabular/dataset_v1.py
+++ b/hybridbackend/tensorflow/data/tabular/dataset_v1.py
@@ -27,6 +27,7 @@
 from tensorflow.python.data.ops import dataset_ops
 
 from hybridbackend.tensorflow.data.parquet.dataset_v1 import _ParquetDatasetV1
+from hybridbackend.tensorflow.data.parquet.schema import parquet_fields
 from hybridbackend.tensorflow.data.parquet.schema import \
   parquet_filenames_and_fields
 from hybridbackend.tensorflow.data.tabular.core import TabularDatasetCreator
@@ -41,6 +42,20 @@ class TabularDatasetV1(dataset_ops.Dataset):
   def from_parquet(cls, filenames, **kwargs):
     return ParquetRecordDatasetV1(filenames, **kwargs)
 
+  @classmethod
+  def schema_from_parquet(cls, filename, fields=None, lower=False):
+    r'''Read schema from a parquet file.
+
+    Args:
+      filename: Path of the parquet file.
+      fields: Existing field definitions or field names.
+      lower: Convert field name to lower case if not found.
+
+    Returns:
+      Field definition list.
+    '''
+    return parquet_fields(filename, fields, lower=lower)
+
   def __init__(
       self, filenames,
       fields=None,
diff --git a/hybridbackend/tensorflow/data/tabular/dataset_v2.py b/hybridbackend/tensorflow/data/tabular/dataset_v2.py
index 6509d4dc..3b13e630 100644
--- a/hybridbackend/tensorflow/data/tabular/dataset_v2.py
+++ b/hybridbackend/tensorflow/data/tabular/dataset_v2.py
@@ -27,6 +27,7 @@
 from tensorflow.python.data.ops import dataset_ops
 
 from hybridbackend.tensorflow.data.parquet.dataset_v2 import _ParquetDatasetV2
+from hybridbackend.tensorflow.data.parquet.schema import parquet_fields
 from hybridbackend.tensorflow.data.parquet.schema import \
   parquet_filenames_and_fields
 from hybridbackend.tensorflow.data.tabular.core import TabularDatasetCreator
@@ -41,6 +42,20 @@ class TabularDatasetV2(dataset_ops.DatasetV2):  # pylint: disable=abstract-metho
   def from_parquet(cls, filenames, **kwargs):
     return ParquetRecordDatasetV2(filenames, **kwargs)
 
+  @classmethod
+  def schema_from_parquet(cls, filename, fields=None, lower=False):
+    r'''Read schema from a parquet file.
+
+    Args:
+      filename: Path of the parquet file.
+      fields: Existing field definitions or field names.
+      lower: Convert field name to lower case if not found.
+
+    Returns:
+      Field definition list.
+    '''
+    return parquet_fields(filename, fields, lower=lower)
+
   def __init__(
       self, filenames,
       fields=None,


Test patch:
diff --git a/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_nested_test.py b/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_nested_test.py
index b99998db..7e491bb6 100644
--- a/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_nested_test.py
+++ b/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_nested_test.py
@@ -54,6 +54,7 @@ def test_read(self):
       ds = hb.data.ParquetDataset(
         [self._filename],
         batch_size=2)
+      ds = ds.apply(hb.data.rebatch(2))
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -64,10 +65,8 @@ def test_read(self):
 
   def test_apply_to_sparse(self):
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
-        [self._filename],
-        batch_size=2)
-      ds = ds.apply(hb.data.parse())
+      ds = hb.data.Dataset.from_parquet([self._filename])
+      ds = ds.batch(2)
       batch = tf.data.make_one_shot_iterator(ds).get_next()['A']
       baseline = tf.ragged.constant(self._data.to_pylist()).to_sparse()
 
@@ -79,10 +78,8 @@ def test_apply_to_sparse(self):
 
   def test_apply_to_tensor(self):
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
-        [self._filename],
-        batch_size=2)
-      ds = ds.apply(hb.data.parse(pad=True))
+      ds = hb.data.Dataset.from_parquet([self._filename], sparse_to_dense=True)
+      ds = ds.batch(2)
       batch = tf.data.make_one_shot_iterator(ds).get_next()['A']
       baseline = tf.ragged.constant(self._data.to_pylist()).to_tensor()
 
diff --git a/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_test.py b/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_test.py
index d60f39ba..4ba88e30 100644
--- a/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_test.py
+++ b/hybridbackend/tensorflow/data/tests/parquet_dataset_ragged_test.py
@@ -63,6 +63,7 @@ def test_read(self):
         [self._filename],
         batch_size=batch_size,
         fields=['col2', 'col0'])
+      ds = ds.apply(hb.data.rebatch(batch_size))
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -89,13 +90,12 @@ def test_read(self):
   def test_to_sparse(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
+      ds = hb.data.Dataset.from_parquet(
         [self._filename],
-        batch_size=batch_size,
         fields=['col2', 'col0'])
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
-      batch = hb.data.DataFrame.parse(batch)
 
     c = self._df['col0']
     with tf.Session(graph=graph) as sess:
@@ -121,11 +121,10 @@ def test_to_sparse(self):
   def test_map_to_sparse(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
+      ds = hb.data.Dataset.from_parquet(
         [self._filename],
-        batch_size=batch_size,
         fields=['col2', 'col0'])
-      ds = ds.map(hb.data.DataFrame.parse)
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -153,11 +152,10 @@ def test_map_to_sparse(self):
   def test_apply_to_sparse(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
+      ds = hb.data.Dataset.from_parquet(
         [self._filename],
-        batch_size=batch_size,
         fields=['col2', 'col0'])
-      ds = ds.apply(hb.data.parse())
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -185,11 +183,10 @@ def test_apply_to_sparse(self):
   def test_feedable_iterator(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
+      ds = hb.data.Dataset.from_parquet(
         [self._filename],
-        batch_size=batch_size,
         fields=['col2', 'col0'])
-      ds = ds.apply(hb.data.parse())
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       it = tf.data.make_one_shot_iterator(ds)
       handle_tensor = it.string_handle()
@@ -230,6 +227,7 @@ def test_read_and_map(self):
         [self._filename],
         batch_size=batch_size,
         fields=['col2', 'col0'])
+      ds = ds.apply(hb.data.rebatch(batch_size))
 
       def _parse(values):
         return values['col0'], values['col2']
diff --git a/hybridbackend/tensorflow/data/tests/parquet_dataset_reshape_test.py b/hybridbackend/tensorflow/data/tests/parquet_dataset_reshape_test.py
index eaa2185b..b84bd64d 100644
--- a/hybridbackend/tensorflow/data/tests/parquet_dataset_reshape_test.py
+++ b/hybridbackend/tensorflow/data/tests/parquet_dataset_reshape_test.py
@@ -59,15 +59,14 @@ def tearDown(self):  # pylint: disable=invalid-name
   def test_reshape(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
+      ds = hb.data.Dataset.from_parquet(
         [self._filename],
-        batch_size=batch_size,
         fields=[
           hb.data.DataFrame.Field('col2'),
           hb.data.DataFrame.Field('col0', shape=[4])])
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
-      batch = hb.data.DataFrame.parse(batch)
 
     c = self._df['col0']
     with tf.Session(graph=graph) as sess:
diff --git a/hybridbackend/tensorflow/data/tests/parquet_dataset_string_test.py b/hybridbackend/tensorflow/data/tests/parquet_dataset_string_test.py
index fdd85cfc..7c62e75c 100644
--- a/hybridbackend/tensorflow/data/tests/parquet_dataset_string_test.py
+++ b/hybridbackend/tensorflow/data/tests/parquet_dataset_string_test.py
@@ -128,8 +128,7 @@ def test_to_sparse(self):
 
   def test_unbatch_and_to_sparse(self):
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(self._filename)
-      ds = ds.map(hb.data.DataFrame.unbatch_and_to_sparse)
+      ds = hb.data.Dataset.from_parquet(self._filename)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
diff --git a/hybridbackend/tensorflow/data/tests/parquet_dataset_test.py b/hybridbackend/tensorflow/data/tests/parquet_dataset_test.py
index 0612ca5e..402933c1 100644
--- a/hybridbackend/tensorflow/data/tests/parquet_dataset_test.py
+++ b/hybridbackend/tensorflow/data/tests/parquet_dataset_test.py
@@ -56,11 +56,11 @@ def tearDown(self):  # pylint: disable=invalid-name
   def test_read(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
+      ds = hb.data.Dataset.from_parquet(
         self._filename,
-        batch_size=batch_size,
         fields=[hb.data.DataFrame.Field('A', tf.int64),
                 hb.data.DataFrame.Field('C', tf.int64)])
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -77,7 +77,8 @@ def test_read(self):
   def test_schema_auto_detection_read(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset([self._filename], batch_size=batch_size)
+      ds = hb.data.Dataset.from_parquet([self._filename])
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -92,10 +93,8 @@ def test_schema_auto_detection_read(self):
   def test_dtype_auto_detection_read(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(
-        [self._filename],
-        batch_size=batch_size,
-        fields=['B', 'C'])
+      ds = hb.data.Dataset.from_parquet([self._filename], fields=['B', 'C'])
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -110,13 +109,11 @@ def test_dtype_auto_detection_read(self):
   def test_dtype_auto_detection_read_lower(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      actual_fields = hb.data.ParquetDataset.read_schema(
+      actual_fields = hb.data.Dataset.schema_from_parquet(
         self._filename, ['B', 'D'], lower=True)
       fld = actual_fields[1].name
-      ds = hb.data.ParquetDataset(
-        [self._filename],
-        batch_size=batch_size,
-        fields=actual_fields)
+      ds = hb.data.Dataset.from_parquet([self._filename], fields=actual_fields)
+      ds = ds.batch(batch_size)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
diff --git a/hybridbackend/tensorflow/data/tests/rebatch_dataset_test.py b/hybridbackend/tensorflow/data/tests/rebatch_dataset_test.py
index 2d28aff3..0314cfd7 100644
--- a/hybridbackend/tensorflow/data/tests/rebatch_dataset_test.py
+++ b/hybridbackend/tensorflow/data/tests/rebatch_dataset_test.py
@@ -101,30 +101,6 @@ def test_expilict_batch(self):
         np.testing.assert_equal(result['A'], a[start_row:end_row].to_numpy())
         np.testing.assert_equal(result['C'], c[start_row:end_row].to_numpy())
 
-  def test_min_batch(self):
-    micro_batch_size = 20
-    batch_size = 32
-    min_batch_size = 30
-    with tf.Graph().as_default() as graph:
-      ds = hb.data.ParquetDataset(self._filename, micro_batch_size)
-      ds = ds.apply(hb.data.rebatch(batch_size, min_batch_size))
-      ds = ds.prefetch(4)
-      batch = tf.data.make_one_shot_iterator(ds).get_next()
-
-    a = self._df['A']
-    c = self._df['C']
-    with tf.Session(graph=graph) as sess:
-      end_row = 0
-      for _ in xrange(3):
-        result = sess.run(batch)
-        aresult = result['A']
-        cresult = result['C']
-        bs = len(aresult)
-        start_row = end_row
-        end_row = start_row + bs
-        np.testing.assert_equal(aresult, a[start_row:end_row].to_numpy())
-        np.testing.assert_equal(cresult, c[start_row:end_row].to_numpy())
-
   def test_ragged(self):
     micro_batch_size = 20
     batch_size = 32
@@ -186,14 +162,10 @@ def test_to_sparse(self):
           len(expected.nested_row_splits[0]))
 
   def test_shuffle_micro_batch(self):
-    micro_batch_size = 20
     batch_size = 32
     with tf.Graph().as_default() as graph:
-      srcds = hb.data.ParquetDataset(self._filename, micro_batch_size)
-      ds = srcds.shuffle(4)
-      ds = ds.apply(
-        hb.data.rebatch(batch_size, fields=srcds.fields, num_parallel_scans=2))
-      ds = ds.apply(hb.data.parse())
+      ds = hb.data.Dataset.from_parquet(self._filename)
+      ds = ds.shuffle_batch(batch_size, buffer_size=80)
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()
 
@@ -206,7 +178,7 @@ def test_thread_pool(self):
     batch_size = 32
     with tf.Graph().as_default() as graph:
       ds = hb.data.ParquetDataset(self._filename, micro_batch_size)
-      ds = ds.apply(hb.data.rebatch(batch_size, num_parallel_scans=3))
+      ds = ds.apply(hb.data.rebatch(batch_size))
       ds = ds.prefetch(4)
       batch = tf.data.make_one_shot_iterator(ds).get_next()