fix the parquet dataset creation with list_files and map. (#1521)

The change moves the `filename` IsDirectory check to the cc layer from python layer. This allows the users to use tf.data.Dataset.list_files(...) and map() to load multiple parquet files into a dataset.

Author: kvignesh1420

tensorflow_io/core/kernels/parquet_kernels.cc

@@ -31,6 +31,12 @@ class ParquetReadableResource : public ResourceBase {
 
   Status Init(const string& input) {
     mutex_lock l(mu_);
+    Status status = env_->IsDirectory(input);
+    if (status.ok()) {
+      return errors::InvalidArgument(
+          "passing a directory path to 'filename' is not supported. ",
+          "Use 'tf.data.Dataset.list_files()' with a map() operation instead.");
+    }
 
     file_.reset(new SizedRandomAccessFile(env_, input, nullptr, 0));
     TF_RETURN_IF_ERROR(file_->GetFileSize(&file_size_));

tensorflow_io/python/ops/parquet_dataset_ops.py

@@ -27,11 +27,6 @@ def __init__(self, filename, columns=None, internal=True):
         """ParquetIODataset."""
         assert internal
         with tf.name_scope("ParquetIODataset"):
-            if tf.io.gfile.isdir(filename):
-                raise ValueError(
-                    "passing a directory path to 'filename' is not supported. "
-                    "Use 'tf.data.Dataset.list_files()' with a map() operation instead."
-                )
             components, shapes, dtypes = core_ops.io_parquet_readable_info(
                 filename, shared=filename, container="ParquetIODataset"
             )

tests/test_parquet.py

@@ -205,5 +205,44 @@ def test_parquet_data():
         i += 1
 
 
+def test_parquet_dataset_from_file_pattern():
+    """Test the parquet dataset creation process using a file pattern"""
+    df = pd.DataFrame({"pred_0": 0.1 * np.arange(100), "pred_1": -0.1 * np.arange(100)})
+    df.to_parquet("df_x.parquet")
+    df.to_parquet("df_y.parquet")
+    columns = {
+        "pred_0": tf.TensorSpec(tf.TensorShape([]), tf.double),
+        "pred_1": tf.TensorSpec(tf.TensorShape([]), tf.double),
+    }
+
+    def map_fn(file_location):
+        return tfio.IODataset.from_parquet(file_location, columns=columns)
+
+    ds = tf.data.Dataset.list_files("*.parquet").map(map_fn)
+
+    for d in ds:  # loop over the files
+        for item in d:  # loop over items
+            print(item)
+
+
+def test_parquet_dataset_from_dir_failure():
+    """Test the dataset creation failure when a directory is passed
+    instead of a filename."""
+    dir_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_parquet")
+    columns = {
+        "pred_0": tf.TensorSpec(tf.TensorShape([]), tf.double),
+        "pred_1": tf.TensorSpec(tf.TensorShape([]), tf.double),
+    }
+    try:
+        _ = tfio.IODataset.from_parquet(dir_path, columns=columns)
+    except Exception as e:
+        assert (
+            str(e)
+            == "passing a directory path to 'filename' is not supported. "
+            + "Use 'tf.data.Dataset.list_files()' with a map() operation instead. "
+            + "[Op:IO>ParquetReadableInfo]"
+        )
+
+
 if __name__ == "__main__":
     test.main()