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

Issue:
[DATA] Improve schema parsing in ParquetDataset
This patch improves schema parsing, which supports fixed-length list and case-insensitive field names. Closes #53 #51 


Discussion / hints:
## Test Results
  72 files  +3    72 suites  +3   4m 56s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "duration of all tests") +4s
106 tests +2  106 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests") +2    0 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") ±0  0 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") ±0 
318 runs  +6  227 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests") +6  91 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") ±0  0 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") ±0 

Results for commit 386e163d. ± Comparison against base commit b888ed5b.

:recycle: This comment has been updated with latest results.

Fix patch:
diff --git a/hybridbackend/cpp/tensorflow/arrow/arrow.cc b/hybridbackend/cpp/tensorflow/arrow/arrow.cc
index 95fcb07f..5abc4192 100644
--- a/hybridbackend/cpp/tensorflow/arrow/arrow.cc
+++ b/hybridbackend/cpp/tensorflow/arrow/arrow.cc
@@ -71,15 +71,32 @@ class ArrowPrimitiveTensorBuffer : public TensorBuffer {
 #endif
 
 ::arrow::Status MakeTensorFromArrowBuffer(
-    DataType dtype, const std::shared_ptr<::arrow::Buffer>& arrow_buffer,
-    Tensor* tensor) {
-  const TensorShape shape = {arrow_buffer->size() / DataTypeSize(dtype)};
+    const DataType dtype, const PartialTensorShape& shape,
+    const std::shared_ptr<::arrow::Buffer>& arrow_buffer, Tensor* tensor) {
+  const int64 total_num_elems = arrow_buffer->size() / DataTypeSize(dtype);
+  int64 shape_num_elems = shape.num_elements();
+  if (TF_PREDICT_FALSE(shape_num_elems < 0)) {
+    return ::arrow::Status::Invalid("Field shape is not fully defined");
+  }
+  int64 dim0 = total_num_elems;
+  if (TF_PREDICT_FALSE(shape_num_elems > 0)) {
+    dim0 = total_num_elems / shape_num_elems;
+    if (TF_PREDICT_FALSE(dim0 * shape_num_elems != total_num_elems)) {
+      return ::arrow::Status::Invalid("Field shape mismatch with actual data");
+    }
+  }
+  TensorShape actual_shape;
+  if (!TF_PREDICT_TRUE(
+          PartialTensorShape({dim0}).Concatenate(shape).AsTensorShape(
+              &actual_shape))) {
+    return ::arrow::Status::Invalid("Field shape is not fully defined");
+  }
 
 #if HYBRIDBACKEND_ARROW_ZEROCOPY
   // NOTE: Alignment is 64 in Arrow 4.x, same to EIGEN_MAX_ALIGN_BYTES. See:
   // https://github.com/apache/arrow/blob/apache-arrow-4.0.1/cpp/src/arrow/memory_pool.cc#L97
   if (TF_PREDICT_FALSE(!CHECK_EIGEN_ALIGN(arrow_buffer->data()))) {
-    *tensor = Tensor(dtype, shape);
+    *tensor = Tensor(dtype, actual_shape);
     std::memcpy(const_cast<char*>(tensor->tensor_data().data()),
                 arrow_buffer->data(), arrow_buffer->size());
     return ::arrow::Status::OK();
@@ -88,10 +105,10 @@ ::arrow::Status MakeTensorFromArrowBuffer(
   ArrowPrimitiveTensorBuffer* tensor_buffer =
       new ArrowPrimitiveTensorBuffer(arrow_buffer);
   core::ScopedUnref unref(tensor_buffer);
-  *tensor = Tensor(dtype, shape, tensor_buffer);
+  *tensor = Tensor(dtype, actual_shape, tensor_buffer);
   return ::arrow::Status::OK();
 #else
-  *tensor = Tensor(dtype, shape);
+  *tensor = Tensor(dtype, actual_shape);
   std::memcpy(const_cast<char*>(tensor->tensor_data().data()),
               arrow_buffer->data(), arrow_buffer->size());
   return ::arrow::Status::OK();
@@ -99,17 +116,35 @@ ::arrow::Status MakeTensorFromArrowBuffer(
 }
 
 ::arrow::Status MakeStringTensorFromArrowArray(
-    const ::arrow::StringArray& array, Tensor* tensor) {
+    const PartialTensorShape& shape, const ::arrow::StringArray& array,
+    Tensor* tensor) {
   if (array.null_count() != 0) {
     return arrow::Status::Invalid("Null elements not supported");
   }
 
-  const auto num_strings = array.length();
+  const auto total_num_elems = array.length();
+  int64 shape_num_elems = shape.num_elements();
+  if (TF_PREDICT_FALSE(shape_num_elems < 0)) {
+    return ::arrow::Status::Invalid("Field shape is not fully defined");
+  }
+  int64 dim0 = total_num_elems;
+  if (TF_PREDICT_FALSE(shape_num_elems > 0)) {
+    dim0 = total_num_elems / shape_num_elems;
+    if (TF_PREDICT_FALSE(dim0 * shape_num_elems != total_num_elems)) {
+      return ::arrow::Status::Invalid("Field shape mismatch with actual data");
+    }
+  }
+  TensorShape actual_shape;
+  if (!TF_PREDICT_TRUE(
+          PartialTensorShape({dim0}).Concatenate(shape).AsTensorShape(
+              &actual_shape))) {
+    return ::arrow::Status::Invalid("Field shape is not fully defined");
+  }
 
-  *tensor = Tensor(DT_STRING, TensorShape({num_strings}));
+  *tensor = Tensor(DT_STRING, actual_shape);
   auto tensor_vec = tensor->vec<std::string>();
 
-  for (auto i = 0; i < num_strings; ++i) {
+  for (auto i = 0; i < total_num_elems; ++i) {
     int string_size;
     auto string_data = array.GetValue(i, &string_size);
     tensor_vec(i).assign(reinterpret_cast<const char*>(string_data),
@@ -119,39 +154,40 @@ ::arrow::Status MakeStringTensorFromArrowArray(
 }
 
 // Primitive Arrow arrays have validity and value buffers.
-#define RAGGED_TENSOR_BUILDER_PRIMITIVE_VISIT(ARRAY_CLASS)                    \
-  ::arrow::Status Visit(const ARRAY_CLASS& array) override {                  \
-    if (TF_PREDICT_FALSE(ragged_rank_ != 0)) {                                \
-      return ::arrow::Status::Invalid("Inconsistent ragged rank");            \
-    }                                                                         \
-    Tensor tensor;                                                            \
-    auto st =                                                                 \
-        MakeTensorFromArrowBuffer(dtype_, array.data()->buffers[1], &tensor); \
-    if (!st.ok()) {                                                           \
-      return st;                                                              \
-    }                                                                         \
-    ragged_tensor_.push_front(std::move(tensor));                             \
-    return ::arrow::Status::OK();                                             \
-  }
-
-#define RAGGED_TENSOR_BUILDER_STRING_VISIT(ARRAY_CLASS)            \
-  ::arrow::Status Visit(const ARRAY_CLASS& array) override {       \
-    if (TF_PREDICT_FALSE(ragged_rank_ != 0)) {                     \
-      return ::arrow::Status::Invalid("Inconsistent ragged rank"); \
-    }                                                              \
-    Tensor tensor;                                                 \
-    auto st = MakeStringTensorFromArrowArray(array, &tensor);      \
-    if (!st.ok()) {                                                \
-      return st;                                                   \
-    }                                                              \
-    ragged_tensor_.push_front(std::move(tensor));                  \
-    return ::arrow::Status::OK();                                  \
+#define RAGGED_TENSOR_BUILDER_PRIMITIVE_VISIT(ARRAY_CLASS)                  \
+  ::arrow::Status Visit(const ARRAY_CLASS& array) override {                \
+    if (TF_PREDICT_FALSE(ragged_rank_ != 0)) {                              \
+      return ::arrow::Status::Invalid("Inconsistent ragged rank");          \
+    }                                                                       \
+    Tensor tensor;                                                          \
+    auto st = MakeTensorFromArrowBuffer(dtype_, shape_,                     \
+                                        array.data()->buffers[1], &tensor); \
+    if (!st.ok()) {                                                         \
+      return st;                                                            \
+    }                                                                       \
+    ragged_tensor_.push_front(std::move(tensor));                           \
+    return ::arrow::Status::OK();                                           \
+  }
+
+#define RAGGED_TENSOR_BUILDER_STRING_VISIT(ARRAY_CLASS)               \
+  ::arrow::Status Visit(const ARRAY_CLASS& array) override {          \
+    if (TF_PREDICT_FALSE(ragged_rank_ != 0)) {                        \
+      return ::arrow::Status::Invalid("Inconsistent ragged rank");    \
+    }                                                                 \
+    Tensor tensor;                                                    \
+    auto st = MakeStringTensorFromArrowArray(shape_, array, &tensor); \
+    if (!st.ok()) {                                                   \
+      return st;                                                      \
+    }                                                                 \
+    ragged_tensor_.push_front(std::move(tensor));                     \
+    return ::arrow::Status::OK();                                     \
   }
 
 class RaggedTensorBuilder : public ::arrow::ArrayVisitor {
  public:
-  RaggedTensorBuilder(DataType dtype, int32 ragged_rank)
-      : dtype_(dtype), ragged_rank_(ragged_rank) {}
+  RaggedTensorBuilder(const DataType dtype, const int32 ragged_rank,
+                      const PartialTensorShape& shape)
+      : dtype_(dtype), ragged_rank_(ragged_rank), shape_(shape) {}
 
   ::arrow::Status Build(const std::shared_ptr<::arrow::Array>& array,
                         std::vector<Tensor>* output_tensors) {
@@ -167,8 +203,8 @@ class RaggedTensorBuilder : public ::arrow::ArrayVisitor {
   ::arrow::Status Visit(const ::arrow::ListArray& array) override {
     --ragged_rank_;
     Tensor tensor;
-    auto st =
-        MakeTensorFromArrowBuffer(DT_INT32, array.value_offsets(), &tensor);
+    auto st = MakeTensorFromArrowBuffer(DT_INT32, PartialTensorShape({}),
+                                        array.value_offsets(), &tensor);
     if (!st.ok()) {
       return st;
     }
@@ -191,6 +227,7 @@ class RaggedTensorBuilder : public ::arrow::ArrayVisitor {
  private:
   const DataType dtype_;
   int32 ragged_rank_;
+  const PartialTensorShape shape_;
   std::deque<Tensor> ragged_tensor_;
 };
 
@@ -230,7 +267,8 @@ Status MakeDataTypeAndRaggedRankFromArrowDataType(
 }
 
 Status MakeTensorsFromArrowArray(
-    DataType dtype, int32 ragged_rank,
+    const DataType dtype, const int32 ragged_rank,
+    const PartialTensorShape& shape,
     const std::shared_ptr<::arrow::Array>& arrow_array,
     std::vector<Tensor>* output_tensors) {
   if (TF_PREDICT_FALSE(arrow_array->null_count() != 0)) {
@@ -241,7 +279,7 @@ Status MakeTensorsFromArrowArray(
     return errors::Internal("Arrow array has zero non-offset not supported");
   }
 
-  RaggedTensorBuilder builder(dtype, ragged_rank);
+  RaggedTensorBuilder builder(dtype, ragged_rank, shape);
   TF_RETURN_IF_ARROW_ERROR(builder.Build(arrow_array, output_tensors));
   return Status::OK();
 }
diff --git a/hybridbackend/cpp/tensorflow/arrow/arrow.h b/hybridbackend/cpp/tensorflow/arrow/arrow.h
index 30dba44a..639a52c9 100644
--- a/hybridbackend/cpp/tensorflow/arrow/arrow.h
+++ b/hybridbackend/cpp/tensorflow/arrow/arrow.h
@@ -93,7 +93,8 @@ Status MakeDataTypeAndRaggedRankFromArrowDataType(
     int32* ragged_rank);
 
 Status MakeTensorsFromArrowArray(
-    DataType type, int32 ragged_rank,
+    const DataType type, const int32 ragged_rank,
+    const PartialTensorShape& shape,
     const std::shared_ptr<::arrow::Array>& arrow_array,
     std::vector<Tensor>* output_tensors);
 
diff --git a/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.cc b/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.cc
index b0f1fb0b..bf99c56b 100644
--- a/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.cc
+++ b/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.cc
@@ -33,6 +33,7 @@ class ParquetBatchReader::Impl {
        const std::vector<string>& field_names,
        const DataTypeVector& field_dtypes,
        const std::vector<int32>& field_ragged_ranks,
+       const std::vector<PartialTensorShape>& field_shapes,
        const int64 partition_count, const int64 partition_index,
        const bool drop_remainder)
       : filename_(filename),
@@ -40,6 +41,7 @@ class ParquetBatchReader::Impl {
         field_names_(field_names),
         field_dtypes_(field_dtypes),
         field_ragged_ranks_(field_ragged_ranks),
+        field_shapes_(field_shapes),
         partition_count_(partition_count),
         partition_index_(partition_index),
         drop_remainder_(drop_remainder) {}
@@ -126,7 +128,8 @@ class ParquetBatchReader::Impl {
     auto arrays = batch->columns();
     for (size_t i = 0; i < arrays.size(); ++i) {
       TF_RETURN_IF_ERROR(MakeTensorsFromArrowArray(
-          field_dtypes_[i], field_ragged_ranks_[i], arrays[i], output_tensors));
+          field_dtypes_[i], field_ragged_ranks_[i], field_shapes_[i], arrays[i],
+          output_tensors));
     }
 
     return Status::OK();
@@ -141,6 +144,7 @@ class ParquetBatchReader::Impl {
   std::vector<string> field_names_;
   DataTypeVector field_dtypes_;
   std::vector<int32> field_ragged_ranks_;
+  std::vector<PartialTensorShape> field_shapes_;
   int64 partition_count_;
   int64 partition_index_;
   bool drop_remainder_;
@@ -156,11 +160,13 @@ class ParquetBatchReader::Impl {
 ParquetBatchReader::ParquetBatchReader(
     const string& filename, const int64 batch_size,
     const std::vector<string>& field_names, const DataTypeVector& field_dtypes,
-    const std::vector<int32>& field_ragged_ranks, const int64 partition_count,
-    const int64 partition_index, const bool drop_remainder)
+    const std::vector<int32>& field_ragged_ranks,
+    const std::vector<PartialTensorShape>& field_shapes,
+    const int64 partition_count, const int64 partition_index,
+    const bool drop_remainder)
     : pimpl_(new ParquetBatchReader::Impl(
           filename, batch_size, field_names, field_dtypes, field_ragged_ranks,
-          partition_count, partition_index, drop_remainder)) {}
+          field_shapes, partition_count, partition_index, drop_remainder)) {}
 
 Status ParquetBatchReader::Open() { return pimpl_->Open(); }
 
diff --git a/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.h b/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.h
index bc6d083c..6ea0e817 100644
--- a/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.h
+++ b/hybridbackend/cpp/tensorflow/arrow/parquet_batch_reader.h
@@ -31,6 +31,7 @@ class ParquetBatchReader {
                      const std::vector<string>& field_names,
                      const DataTypeVector& field_dtypes,
                      const std::vector<int32>& field_ragged_ranks,
+                     const std::vector<PartialTensorShape>& field_shapes,
                      const int64 partition_count, const int64 partition_index,
                      const bool drop_remainder);
 
diff --git a/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc b/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc
index f7690b8b..53167511 100644
--- a/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc
+++ b/hybridbackend/cpp/tensorflow/ops/parquet_dataset_ops.cc
@@ -36,6 +36,7 @@ REGISTER_OP("ParquetTabularDataset")
     .Attr("field_names: list(string) >= 1")
     .Attr("field_dtypes: list(type) >= 1")
     .Attr("field_ragged_ranks: list(int) >= 1")
+    .Attr("field_shapes: list(shape) >= 1")
     .Attr("partition_count: int = 1")
     .Attr("partition_index: int = 0")
     .Attr("drop_remainder: bool = false")
@@ -56,6 +57,7 @@ batch_size: Maxium number of samples in an output batch.
 field_names: List of field names to read.
 field_dtypes: List of data types for each field.
 field_ragged_ranks: List of ragged rank for each field.
+field_shapes: List of shapes for each field.
 partition_count: Count of row group partitions.
 partition_index: Index of row group partitions.
 drop_remainder: If True, only keep batches with exactly `batch_size` samples.
@@ -72,6 +74,7 @@ class ParquetTabularDatasetOp : public DatasetOpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("field_dtypes", &field_dtypes_));
     OP_REQUIRES_OK(ctx,
                    ctx->GetAttr("field_ragged_ranks", &field_ragged_ranks_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("field_shapes", &field_shapes_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("partition_count", &partition_count_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("partition_index", &partition_index_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("drop_remainder", &drop_remainder_));
@@ -85,6 +88,7 @@ class ParquetTabularDatasetOp : public DatasetOpKernel {
   std::vector<string> field_names_;
   DataTypeVector field_dtypes_;
   std::vector<int32> field_ragged_ranks_;
+  std::vector<PartialTensorShape> field_shapes_;
   int64 partition_count_;
   int64 partition_index_;
   bool drop_remainder_;
@@ -96,6 +100,7 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
           const std::vector<string>& field_names,
           const DataTypeVector& field_dtypes,
           const std::vector<int32>& field_ragged_ranks,
+          const std::vector<PartialTensorShape>& field_shapes,
           const int64 partition_count, const int64 partition_index,
           const bool drop_remainder)
       : DatasetBase(DatasetContext(ctx)),
@@ -104,25 +109,26 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
         field_names_(std::move(field_names)),
         field_dtypes_(std::move(field_dtypes)),
         field_ragged_ranks_(std::move(field_ragged_ranks)),
+        field_shapes_(std::move(field_shapes)),
         partition_count_(partition_count),
         partition_index_(partition_index),
         drop_remainder_(drop_remainder) {
+    const int64 actual_batch_size(drop_remainder ? batch_size : -1);
     int64 num_outputs = field_names.size();
     for (int64 i = 0; i < field_names.size(); ++i) {
       output_dtypes_.push_back(std::move(field_dtypes[i]));
+      output_shapes_.push_back(PartialTensorShape({actual_batch_size})
+                                   .Concatenate(field_shapes_[i]));
       for (int64 j = 0; j < field_ragged_ranks_[i]; ++j) {
         output_dtypes_.push_back(DT_INT32);
+        output_shapes_.push_back(PartialTensorShape({actual_batch_size}));
       }
       num_outputs += field_ragged_ranks_[i];
     }
-    int64 actual_batch_size(drop_remainder ? batch_size : -1);
-    for (size_t i = 0; i < num_outputs; ++i) {
-      output_shapes_.push_back(PartialTensorShape({actual_batch_size}));
-    }
 
     reader_ = absl::make_unique<ParquetBatchReader>(
         filename_, batch_size_, field_names_, field_dtypes_,
-        field_ragged_ranks_, partition_count_, partition_index_,
+        field_ragged_ranks_, field_shapes_, partition_count_, partition_index_,
         drop_remainder_);
   }
 
@@ -160,6 +166,8 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
     b->BuildAttrValue(field_dtypes_, &field_dtypes);
     AttrValue field_ragged_ranks;
     b->BuildAttrValue(field_ragged_ranks_, &field_ragged_ranks);
+    AttrValue field_shapes;
+    b->BuildAttrValue(field_ragged_ranks_, &field_shapes);
     AttrValue partition_count;
     b->BuildAttrValue(partition_count_, &partition_count);
     AttrValue partition_index;
@@ -171,6 +179,7 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
                       {{"field_names", field_names},
                        {"field_dtypes", field_dtypes},
                        {"field_ragged_ranks", field_ragged_ranks},
+                       {"field_shapes", field_shapes},
                        {"partition_count", partition_count},
                        {"partition_index", partition_index},
                        {"drop_remainder", drop_remainder}},
@@ -185,6 +194,7 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
   const std::vector<string> field_names_;
   const DataTypeVector field_dtypes_;
   const std::vector<int32> field_ragged_ranks_;
+  const std::vector<PartialTensorShape> field_shapes_;
   const int64 partition_count_;
   const int64 partition_index_;
   const bool drop_remainder_;
@@ -203,9 +213,10 @@ void ParquetTabularDatasetOp::MakeDataset(OpKernelContext* ctx,
   OP_REQUIRES(ctx, batch_size > 0,
               errors::InvalidArgument("batch_size must be greater than zero."));
 
-  Dataset* ds = new Dataset(
-      ctx, filename, batch_size, field_names_, field_dtypes_,
-      field_ragged_ranks_, partition_count_, partition_index_, drop_remainder_);
+  Dataset* ds =
+      new Dataset(ctx, filename, batch_size, field_names_, field_dtypes_,
+                  field_ragged_ranks_, field_shapes_, partition_count_,
+                  partition_index_, drop_remainder_);
   OP_REQUIRES_OK(ctx, ds->Open());
   *output = ds;
 }
diff --git a/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc b/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc
index 94e7d800..0a1cab87 100644
--- a/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc
+++ b/hybridbackend/cpp/tensorflow/ops/rebatch_dataset_ops.cc
@@ -234,7 +234,13 @@ void RecalculateSplit(Tensor* split, int32 value) {
     sdata[i] += value;
   }
   if (offset < ssize) {
-    auto rslice = split->Slice(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);
   }
 }
@@ -313,8 +319,14 @@ class RebatchTabularDatasetOp::Dataset::Iterator
     for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
       const int64 rank = field_ranks_[fid];
       if (rank == 0) {
-        output_tensors->push_back(
-            input_tensors[cur].Slice(row_start, row_limit));
+        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;
       }
@@ -322,7 +334,14 @@ class RebatchTabularDatasetOp::Dataset::Iterator
       int64 start = row_start;
       int64 limit = row_limit;
       for (size_t i = 1; i < rank + 1; ++i) {
-        auto s = input_tensors[cur + i].Slice(start, limit + 1);
+        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;
@@ -330,7 +349,13 @@ class RebatchTabularDatasetOp::Dataset::Iterator
         limit = sdata[slimit];
         output_tensors->push_back(std::move(s));
       }
-      output_tensors->at(cur) = input_tensors[cur].Slice(start, limit);
+      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.
@@ -394,15 +419,28 @@ class RebatchTabularDatasetOp::Dataset::Iterator
     for (size_t fid = 0; fid < field_ranks_.size(); ++fid) {
       const int64 rank = field_ranks_[fid];
       if (rank == 0) {
-        tensor_queues_[cur].push_back(
-            input_tensors[cur].Slice(row_start, row_limit));
+        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) {
-        auto sliced = input_tensors[cur + i].Slice(start, limit + 1);
+        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;
@@ -410,7 +448,13 @@ class RebatchTabularDatasetOp::Dataset::Iterator
         limit = sdata[slimit];
         tensor_queues_[cur + i].push_back(std::move(sliced));
       }
-      tensor_queues_[cur].push_back(input_tensors[cur].Slice(start, limit));
+      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);
@@ -505,8 +549,14 @@ class RebatchTabularDatasetOp::Dataset::Iterator
               const_cast<char*>(t.tensor_data().data()));
           memcpy(sdata + tstart, tdata + 1, t.TotalBytes() - sizeof(int32));
           const int32 tsize = t.NumElements() - 1;
-          auto tpart =
-              output_tensors->at(cur + i).Slice(tstart, tstart + tsize);
+          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;
         }
diff --git a/hybridbackend/tensorflow/data/dataframe.py b/hybridbackend/tensorflow/data/dataframe.py
index f14b32e8..a40775cf 100644
--- a/hybridbackend/tensorflow/data/dataframe.py
+++ b/hybridbackend/tensorflow/data/dataframe.py
@@ -67,6 +67,8 @@ def __init__(self, name, dtype=None, ragged_rank=None, shape=None):
           raise ValueError(
             f'Field {name} is a nested list ({ragged_rank}) '
             f'with shape {shape}')
+      else:
+        shape = tensor_shape.TensorShape({})
       self._shape = shape
 
     @property
@@ -101,10 +103,7 @@ def __repr__(self):
             dtypestr = f'list^{self._ragged_rank}<{dtypestr}>'
           elif self._ragged_rank > 0:
             dtypestr = f'list<{dtypestr}>'
-      if self._shape is None:
-        shapestr = 'unknown'
-      else:
-        shapestr = str(self._shape)
+      shapestr = str(self._shape)
       return f'{self._name} (dtype={dtypestr}, shape={shapestr})'
 
     def map(self, func, rank=None):
@@ -133,17 +132,13 @@ def output_types(self):
 
     @property
     def output_shapes(self):
-      if self._shape is None:
-        return self.map(lambda _: tensor_shape.vector(None))
       return self.map(
         lambda i: tensor_shape.vector(None).concatenate(self._shape) if i == 0
         else tensor_shape.vector(None))
 
     @property
     def output_specs(self):
-      shape = tensor_shape.vector(None)
-      if self._shape is not None:
-        shape = shape.concatenate(self._shape)
+      shape = tensor_shape.vector(None).concatenate(self._shape)
       specs = [tensor_spec.TensorSpec(shape, dtype=self._dtype)]
       specs += [
         tensor_spec.TensorSpec([None], dtype=dtypes.int32)
diff --git a/hybridbackend/tensorflow/data/parquet_dataset_v1.py b/hybridbackend/tensorflow/data/parquet_dataset_v1.py
index c03e9f31..0ae03516 100644
--- a/hybridbackend/tensorflow/data/parquet_dataset_v1.py
+++ b/hybridbackend/tensorflow/data/parquet_dataset_v1.py
@@ -64,6 +64,7 @@ def __init__(
     self._field_dtypes = nest.flatten({f.name: f.dtype for f in self._fields})
     self._field_ragged_ranks = nest.flatten(
       {f.name: f.ragged_rank for f in self._fields})
+    self._field_shapes = nest.flatten({f.name: f.shape for f in self._fields})
     self._partition_count = partition_count
     self._partition_index = partition_index
     self._drop_remainder = drop_remainder
@@ -76,6 +77,7 @@ def _as_variant_tensor(self):
       field_names=self._field_names,
       field_dtypes=self._field_dtypes,
       field_ragged_ranks=self._field_ragged_ranks,
+      field_shapes=self._field_shapes,
       partition_count=self._partition_count,
       partition_index=self._partition_index,
       drop_remainder=self._drop_remainder)
diff --git a/hybridbackend/tensorflow/data/parquet_dataset_v2.py b/hybridbackend/tensorflow/data/parquet_dataset_v2.py
index 94859889..04857820 100644
--- a/hybridbackend/tensorflow/data/parquet_dataset_v2.py
+++ b/hybridbackend/tensorflow/data/parquet_dataset_v2.py
@@ -123,6 +123,7 @@ def __init__(
     self._field_dtypes = nest.flatten({f.name: f.dtype for f in self._fields})
     self._field_ragged_ranks = nest.flatten(
       {f.name: f.ragged_rank for f in self._fields})
+    self._field_shapes = nest.flatten({f.name: f.shape for f in self._fields})
     self._partition_count = partition_count
     self._partition_index = partition_index
     self._drop_remainder = drop_remainder
@@ -133,6 +134,7 @@ def __init__(
       field_names=self._field_names,
       field_dtypes=self._field_dtypes,
       field_ragged_ranks=self._field_ragged_ranks,
+      field_shapes=self._field_shapes,
       partition_count=self._partition_count,
       partition_index=self._partition_index,
       drop_remainder=self._drop_remainder)
diff --git a/hybridbackend/tensorflow/data/parquet_validate.py b/hybridbackend/tensorflow/data/parquet_validate.py
new file mode 100755
index 00000000..7267aab8
--- /dev/null
+++ b/hybridbackend/tensorflow/data/parquet_validate.py
@@ -0,0 +1,96 @@
+#!/usr/bin/env python
+
+# 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'''Validates Parquet files.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import hybridbackend.tensorflow as hb
+import os
+import sys
+import tensorflow as tf
+
+
+def _main(args):
+  r'''Entry function.
+  '''
+  if not args.input_files:
+    return
+  num_files = len(args.input_files)
+  sys.stderr.write(f'[Parquet] Validating schema of {num_files} files...\n')
+  fields = None
+  if args.fields:
+    fields = args.fields.split(',')
+  prev_fields = None
+  for f in args.input_files:
+    actual_fields = hb.data.ParquetDataset.read_schema(
+      f, fields=fields, lower=args.lower)
+    if prev_fields:
+      if len(actual_fields) != len(prev_fields):
+        raise ValueError(
+          f'{f} has {len(actual_fields)} fields while '
+          f'{args.input_files[0]} has {len(prev_fields)} fields')
+      prev_f = args.input_files[0]
+      for idx, field in enumerate(actual_fields):
+        prev_field = prev_fields[idx]
+        if field.name != prev_field.name:
+          raise ValueError(
+            f'name of field {idx} of {f} and {prev_f} are different: '
+            f'{field.name} vs {prev_field.name}')
+        if field.dtype != prev_field.dtype:
+          raise ValueError(
+            f'dtype of field {idx} of {f} and {prev_f} are different: '
+            f'{field.dtype} vs {prev_field.dtype}')
+        if field.ragged_rank != prev_field.ragged_rank:
+          raise ValueError(
+            f'ragged_rank of field {idx} of {f} and {prev_f} are different: '
+            f'{field.ragged_rank} vs {prev_field.ragged_rank}')
+        if field.shape != prev_field.shape:
+          raise ValueError(
+            f'shape of field {idx} of {f} and {prev_f} are different: '
+            f'{field.shape} vs {prev_field.shape}')
+    else:
+      prev_fields = actual_fields
+  sys.stderr.write('[Parquet] Schema validated.\n')
+
+  with tf.Graph().as_default() as graph:
+    peek_ops = []
+    for f in args.input_files:
+      ds = hb.data.ParquetDataset(f, batch_size=args.peek, fields=prev_fields)
+      batch = hb.data.make_one_shot_iterator(ds).get_next()
+      peek_ops.append(batch)
+
+    with tf.Session(graph=graph) as sess:
+      for fidx, f in enumerate(args.input_files):
+        sys.stderr.write(
+          f'[Parquet] Validating {f} ({fidx + 1} of {num_files})...\n')
+        sess.run(peek_ops[fidx])
+        sys.stderr.write(f'[Parquet] {f} validated.\n')
+
+
+if __name__ == '__main__':
+  os.environ['CUDA_VISIBLE_DEVICES'] = ''
+  parser = argparse.ArgumentParser()
+  parser.add_argument('input_files', nargs='+')
+  parser.add_argument('--fields', type=str, default='')
+  parser.add_argument('--lower', default=False, action='store_true')
+  parser.add_argument('--peek', type=int, default=32)
+  _main(parser.parse_args())


Test patch:
diff --git a/hybridbackend/tensorflow/data/parquet_dataset_reshape_test.py b/hybridbackend/tensorflow/data/parquet_dataset_reshape_test.py
index 4510c4a7..02c63b4e 100644
--- a/hybridbackend/tensorflow/data/parquet_dataset_reshape_test.py
+++ b/hybridbackend/tensorflow/data/parquet_dataset_reshape_test.py
@@ -44,9 +44,9 @@ def setUp(self):  # pylint: disable=invalid-name
         [
           np.random.randint(
             0, 100,
-            size=(np.random.randint(1, 5),),
+            size=(4,) if icol == 0 else (np.random.randint(1, 5),),
             dtype=np.int64)
-          for _ in xrange(num_cols)]
+          for icol in xrange(num_cols)]
         for _ in xrange(100)], dtype=object),
       columns=[f'col{c}' for c in xrange(num_cols)])
     self._df.to_parquet(self._filename)
@@ -62,8 +62,8 @@ def test_reshape(self):
         [self._filename],
         batch_size=batch_size,
         fields=[
-          hb.data.DataFrame.Field('col2', shape=[100]),
-          hb.data.DataFrame.Field('col0', shape=[100])])
+          hb.data.DataFrame.Field('col2'),
+          hb.data.DataFrame.Field('col0', shape=[4])])
       ds = ds.prefetch(4)
       batch = hb.data.make_one_shot_iterator(ds).get_next()
       batch = hb.data.DataFrame.to_sparse(batch)
@@ -81,6 +81,7 @@ def test_reshape(self):
           expected_values.extend(item)
           expected_splits.append(expected_splits[-1] + len(item))
         expected = np.array(expected_values)
+        expected = np.reshape(expected, (batch_size, 4))
         actual = result['col0']
         np.testing.assert_allclose(actual, expected)
 
diff --git a/hybridbackend/tensorflow/data/rebatch_dataset_test.py b/hybridbackend/tensorflow/data/rebatch_dataset_test.py
index ec597bc6..61f74006 100644
--- a/hybridbackend/tensorflow/data/rebatch_dataset_test.py
+++ b/hybridbackend/tensorflow/data/rebatch_dataset_test.py
@@ -46,8 +46,7 @@ def build_row(br, er, bseq, eseq):
       b = np.random.randint(
         br, er, size=(np.random.randint(bseq, eseq),), dtype=np.int64)
       c = np.random.randint(br, er, dtype=np.int64)
-      d = np.random.randint(
-        br, er, size=(np.random.randint(bseq, eseq),), dtype=np.int64)
+      d = np.random.randint(br, er, size=(4,), dtype=np.int64)
       e = np.array(list(map(
         str,
         np.random.randint(
@@ -230,6 +229,34 @@ def test_thread_pool(self):
         np.testing.assert_equal(
           actual.nested_row_splits, expected.nested_row_splits)
 
+  def test_ragged_with_shape(self):
+    micro_batch_size = 20
+    batch_size = 32
+    with tf.Graph().as_default() as graph:
+      ds = hb.data.ParquetDataset(
+        self._filename,
+        micro_batch_size,
+        fields=[hb.data.DataFrame.Field('D', shape=[4])])
+      ds = ds.apply(hb.data.rebatch(batch_size))
+      ds = ds.prefetch(4)
+      batch = hb.data.make_one_shot_iterator(ds).get_next()
+
+    d = self._df['D']
+    with tf.Session(graph=graph) as sess:
+      for i in xrange(3):
+        result = sess.run(batch)
+        start_row = i * batch_size
+        end_row = (i + 1) * batch_size
+        expected_items = d[start_row:end_row].to_numpy().tolist()
+        expected_values = []
+        expected_splits = [0]
+        for item in expected_items:
+          expected_values.extend(item)
+          expected_splits.append(expected_splits[-1] + len(item))
+        expected = np.reshape(np.array(expected_values), (-1, 4))
+        actual = result['D']
+        np.testing.assert_allclose(actual.values, expected)
+
 
 if __name__ == '__main__':
   hbtest.main(f'{__file__}.xml')