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

Issue:
rebatch api produce an Check failed: limit <= dim0_size error
# Current behavior

After rebatch(), data iterator get_next() produce an error:

```python
F tensorflow/core/framework/tensor.cc:833] Check failed: limit <= dim0_size (8194 vs. 8193)
```

# Expected behavior

no error

# System information
- OS Platform and Distribution: Ubuntu 18.04.5 LTS
- TensorFlow version: 1.15.0
- Python version: 3.6
- CUDA/cuDNN version: 10.1
- RAM: 94G
- GPU model and memory: Tesla T4, 16G

# Code to reproduce

**Step 1:** Generate a parquet file by running following code
```python
import numpy as np
import pandas as pd
import random

data_list = []
for i in range(1, 10000):
    int_feature = random.randint(1, 100)
    # float_feature = random.random()
    array_feature = [random.randint(1, 10) for x in range(0, 4)]
    data_list.append([int_feature, array_feature])

df = pd.DataFrame(data_list, columns=["int_feature", "array_feature"])
df.to_parquet("parquet_sample_file.parquet")
```

**Step 2:** Load generated parquet in step 1 by HybridBackend

```python
import tensorflow as tf
import hybridbackend.tensorflow as hb


filenames_ds = tf.data.Dataset.from_tensor_slices(['file1.snappy.parquet', 'file2.snappy.parquet', ... 'fileN.snappy.parquet'])


hb_fields = []
hb_fields.append(hb.data.DataFrame.Field("feature1", tf.int64, ragged_rank=0))
hb_fields.append(hb.data.DataFrame.Field("feature2", tf.float32, ragged_rank=1))
hb_fields.append(hb.data.DataFrame.Field("feature3", tf.int64, ragged_rank=1))

ds = filenames_ds.apply(hb.data.read_parquet(8192, hb_fields, num_parallel_reads=tf.data.experimental.AUTOTUNE))
iterator = ds.apply(hb.data.rebatch(8192, fields=hb_fields))

it = iterator.make_one_shot_iterator()
item = it.get_next()

batch_size_dict = {}
with tf.Session() as sess:
    print("======  start ======")
    total_batch_size = 0
    while True:
        try:
            batch = sess.run(item)
            batch_size = len(batch['mod_series'])
            batch_size_dict[batch_size] = batch_size_dict.get(batch_size, 0) + 1
        except tf.errors.OutOfRangeError:
            break

```

Running above code in a pyhon3 shell, an error shall be thrown:

```python
F tensorflow/core/framework/tensor.cc:833] Check failed: limit <= dim0_size (8194 vs. 8193)
```

# Willing to contribute

Yes




Discussion / hints:
Thanks for reporting, can you provide a sample file for reproducing this issue?
> Thanks for reporting, can you provide a sample file for reproducing this issue?

(1) Generate a parquet file by running following code

```python
import numpy as np
import pandas as pd
import random

data_list = []
for i in range(1, 10000):
    int_feature = random.randint(1, 100)
    # float_feature = random.random()
    array_feature = [random.randint(1, 10) for x in range(0, 4)]
    data_list.append([int_feature, array_feature])

df = pd.DataFrame(data_list, columns=["int_feature", "array_feature"])
df.to_parquet("parquet_sample_file.parquet")
```

(2) Load generated parquet file by HybridBackend will reproduce this issue

```python
filenames_ds = tf.data.Dataset.from_tensor_slices(["parquet_sample_file.parquet"])

hb_fields = []
hb_fields.append(hb.data.DataFrame.Field("int_feature", tf.int64, ragged_rank=0))
# hb_fields.append(hb.data.DataFrame.Field("float_feature", tf.float32, ragged_rank=0))
hb_fields.append(hb.data.DataFrame.Field("array_feature", tf.int64, ragged_rank=1))

iterator = filenames_ds.apply(hb.data.read_parquet(100, hb_fields, num_parallel_reads=tf.data.experimental.AUTOTUNE))
iterator = iterator.apply(hb.data.rebatch(100, fields=hb_fields)).repeat(30)

iterator = iterator.make_one_shot_iterator()
item = iterator.get_next()
with tf.Session() as sess:
    print("======  start ======")
    total_batch_size = 0
    while True:
        try:
            a = sess.run(item)
        except tf.errors.OutOfRangeError:
            break
```

Fix patch:
diff --git a/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc b/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc
index c901f3e7..f7690b8b 100644
--- a/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc
+++ b/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc
@@ -127,7 +127,7 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
   }
 
   Status Open() {
-    VLOG(0) << "Starting to read " << filename_ << " ...";
+    VLOG(1) << "Starting to read " << filename_ << " ...";
     return reader_->Open();
   }
 
diff --git a/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc b/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc
index d79f6cc9..94e7d800 100644
--- a/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc
+++ b/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc
@@ -294,7 +294,7 @@ class RebatchTabularDatasetOp::Dataset::Iterator
       return 0;
     }
     if (field_ranks_[0] > 0) {
-      return input_tensors[1].dim_size(0);
+      return input_tensors[1].dim_size(0) - 1;
     }
     return input_tensors[0].dim_size(0);
   }


Test patch:
diff --git a/tests/tensorflow/data/parquet_dataset_rebatch_test.py b/tests/tensorflow/data/parquet_dataset_rebatch_test.py
index 073183f8..a84bf426 100644
--- a/tests/tensorflow/data/parquet_dataset_rebatch_test.py
+++ b/tests/tensorflow/data/parquet_dataset_rebatch_test.py
@@ -22,7 +22,10 @@
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
 import os
+import random
 from six.moves import xrange  # pylint: disable=redefined-builtin
 import tempfile
 import tensorflow as tf
@@ -234,6 +237,39 @@ def test_thread_pool(self):
           actual.nested_row_splits, expected.nested_row_splits)
 
 
+class ParquetDatasetSequenceRebatchTest(unittest.TestCase):
+  def setUp(self):  # pylint: disable=invalid-name
+    self._workspace = tempfile.mkdtemp()
+    self._filename = os.path.join(self._workspace, 'seqtest.parquet')
+    self._nrows = 1000
+    self._ncols = 10
+    self._data = {
+      'clicks': [
+        [random.randint(0, 100) for col in range(self._ncols)]
+        for row in range(self._nrows)]}
+    pq.write_table(pa.Table.from_pydict(self._data), self._filename)
+
+  def tearDown(self):  # pylint: disable=invalid-name
+    os.remove(self._filename)
+
+  def test_ragged(self):
+    batch_size = 8
+    with tf.Graph().as_default() as graph:
+      ds = ParquetDataset(self._filename, batch_size=batch_size)
+      ds = ds.apply(rebatch(batch_size))
+      batch = make_one_shot_iterator(ds).get_next()
+
+    clicks = self._data['clicks']
+    with tf.Session(graph=graph) as sess:
+      for i in xrange(3):
+        actual = sess.run(batch['clicks'])
+        start_row = i * batch_size
+        end_row = (i + 1) * batch_size
+        expected = clicks[start_row:end_row]
+        expected_values = [v for sublist in expected for v in sublist]
+        np.testing.assert_equal(actual.values, expected_values)
+
+
 if __name__ == '__main__':
   register(['cpu', 'data'])
   os.environ['CUDA_VISIBLE_DEVICES'] = ''