[HOPSWORKS-1762] Add support for Spark-Tfrecord (#94)

* Spark-TFRecord
logicalclocks · Aug 26, 2020 · d0e83ab · d0e83ab
1 parent 141677e
commit d0e83ab
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/python/hsfs/engine/spark.py b/python/hsfs/engine/spark.py
@@ -197,6 +197,9 @@ def write_options(self, data_format, provided_options):
         if data_format.lower() == "tfrecords":
             options = dict(recordType="Example")
             options.update(provided_options)
+        elif data_format.lower() == "tfrecord":
+            options = dict(recordType="Example")
+            options.update(provided_options)
         elif data_format.lower() == "csv":
             options = dict(delimiter=",", header="true")
             options.update(provided_options)
@@ -212,6 +215,9 @@ def read_options(self, data_format, provided_options):
         if data_format.lower() == "tfrecords":
             options = dict(recordType="Example", **provided_options)
             options.update(provided_options)
+        elif data_format.lower() == "tfrecord":
+            options = dict(recordType="Example")
+            options.update(provided_options)
         elif data_format.lower() == "csv":
             options = dict(delimiter=",", header="true", inferSchema="true")
             options.update(provided_options)
@@ -235,7 +241,7 @@ def parse_schema(self, dataframe):
 
     def parse_schema_dict(self, dataframe):
         return {
-            feat["name"]: feature.Feature(
+            feat.name: feature.Feature(
                 feat.name,
                 feat.dataType.simpleString(),
                 feat.metadata.get("description", ""),