We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Test_s3 failing
```
      with s3_context(s3_base=s3_base, bucket=engine, files=files) as s3fs:
            # Create nvt.Dataset from mock s3 paths
            url = f"s3://{engine}" if engine == "parquet" else f"s3://{engine}/*"
            dataset = nvt.Dataset(url, engine=engine, storage_options=s3so)
    
            # Check that the iteration API works
            columns = mycols_pq if engine == "parquet" else mycols_csv
            gdf = nvt.dispatch._concat(list(dataset.to_iter()))[columns]
            assert_eq(gdf.reset_index(drop=True), df.reset_index(drop=True))
    
            cat_names = ["name-cat", "name-string"] if engine == "parquet" else ["name-string"]
            cont_names = ["x", "y", "id"]
            label_name = ["label"]
    
            conts = cont_names >> ops.FillMissing() >> ops.Clip(min_value=0) >> ops.LogOp()
            cats = cat_names >> ops.Categorify(cat_cache="host")
    
            processor = nvt.Workflow(conts + cats + label_name)
            processor.fit(dataset)
    
            # make sure we can write out the dataset back to S3
            # (https://github.com/NVIDIA-Merlin/NVTabular/issues/1214)
>           processor.transform(dataset).to_parquet(f"s3://{engine}/output")

/nvtabular/tests/unit/test_s3.py:111: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/nvtabular/nvtabular/io/dataset.py:906: in to_parquet
    self.schema.write(output_path)
/nvtabular/nvtabular/graph/schema.py:154: in write
    return PbTxt_SchemaWriter.write(self, schema_path)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

cls = <class 'nvtabular.graph.schema_io.schema_writer_pbtxt.PbTxt_SchemaWriter'>
schema = [{'name': 'x', 'tags': [<Tags.CONTINUOUS: 'continuous'>], 'properties': {}, 'dtype': <class 'float'>, '_is_list': Fals...int'>, '_is_list': False}, {'name': 'label', 'tags': [], 'properties': {}, 'dtype': dtype('int64'), '_is_list': False}]
schema_path = PosixPath('s3:/csv/output')

    @classmethod
    def write(cls, schema, schema_path):
        schema_path = Path(schema_path)
        if not schema_path.is_dir():
>           raise ValueError(f"The path provided is not a valid directory: {schema_path}")
E           ValueError: The path provided is not a valid directory: s3:/csv/output

/nvtabular/nvtabular/graph/schema_io/schema_writer_pbtxt.py:45: ValueError
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `nvtabular/graph/schema_io/schema_writer_pbtxt.py`
Content:
```
1 #
2 # Copyright (c) 2021, NVIDIA CORPORATION.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 #
16 import os
17 from pathlib import Path
18 
19 import numpy
20 
21 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
22 from google.protobuf import json_format, text_format  # noqa
23 from google.protobuf.any_pb2 import Any  # noqa
24 from google.protobuf.struct_pb2 import Struct  # noqa
25 from tensorflow_metadata.proto.v0 import schema_pb2  # noqa
26 
27 import nvtabular as nvt  # noqa
28 from nvtabular.graph.schema_io.schema_writer_base import SchemaWriter  # noqa
29 from nvtabular.graph.tags import Tags  # noqa
30 
31 
32 class PbTxt_SchemaWriter(SchemaWriter):
33     @classmethod
34     def _read(cls, schema_path):
35         with open(schema_path, "r") as f:
36             schema = schema_pb2.Schema()
37             text_format.Parse(f.read(), schema)
38 
39         return schema
40 
41     @classmethod
42     def write(cls, schema, schema_path):
43         schema_path = Path(schema_path)
44         if not schema_path.is_dir():
45             raise ValueError(f"The path provided is not a valid directory: {schema_path}")
46 
47         # traverse list of column schema
48         schema_file = schema_pb2.Schema()
49         features = []
50         for col_name, col_schema in schema.column_schemas.items():
51             features.append(create_protobuf_feature(col_schema))
52         schema_file.feature.extend(features)
53 
54         with open(schema_path / "schema.pbtxt", "w") as f:
55             f.write(text_format.MessageToString(schema_file))
56         return schema
57 
58     @classmethod
59     def load(cls, schema_path):
60         columns = []
61         if isinstance(schema_path, (str, Path)):
62             if isinstance(schema_path, str):
63                 schema_path = Path(schema_path)
64             if schema_path.is_dir():
65                 schema_path = schema_path / "schema.pbtxt"
66             schema = cls._read(schema_path)
67 
68         for feat in schema.feature:
69             _is_list = False
70             dtype = None
71             properties = {}
72             tags = list(feat.annotation.tag) or []
73             # only one item should ever be in extra_metadata
74             if len(feat.annotation.extra_metadata) > 1:
75                 raise ValueError(
76                     f"{feat.name}: extra_metadata should have 1 item, has \
77                     {len(feat.annotation.extra_metadata)}"
78                 )
79             if feat.annotation.extra_metadata:
80                 properties = json_format.MessageToDict(feat.annotation.extra_metadata[0])["value"]
81             # what domain
82             # load the domain values
83             shape_name = feat.WhichOneof("shape_type")
84             if shape_name:
85                 _is_list = True
86             field_name = feat.WhichOneof("domain_info")
87             if field_name:
88                 domain_values = getattr(feat, field_name)
89                 # if zero no values were passed
90                 if domain_values.max > 0:
91                     properties["domain"] = {"min": domain_values.min, "max": domain_values.max}
92                 if feat.type:
93                     if feat.type == 2:
94                         dtype = numpy.int
95                     elif feat.type == 3:
96                         dtype = numpy.float
97             columns.append(
98                 nvt.ColumnSchema(
99                     feat.name, tags=tags, properties=properties, dtype=dtype, _is_list=_is_list
100                 )
101             )
102 
103         return nvt.Schema(columns)
104 
105 
106 def register_extra_metadata(column_schema, feature):
107     filtered_properties = {k: v for k, v in column_schema.properties.items() if k != "domain"}
108     msg_struct = Struct()
109     # must pack message into "Any" type
110     any_pack = Any()
111     any_pack.Pack(json_format.ParseDict(filtered_properties, msg_struct))
112     # extra_metadata only takes type "Any" messages
113     feature.annotation.extra_metadata.add().CopyFrom(any_pack)
114     return feature
115 
116 
117 def register_list(column_schema, feature):
118     if str(column_schema._is_list):
119         min_length, max_length = None, None
120         if "value_count" in column_schema.properties:
121             min_length = column_schema.properties["value_count"]["min"]
122             max_length = column_schema.properties["value_count"]["max"]
123         if min_length and max_length and min_length == max_length:
124             shape = schema_pb2.FixedShape()
125             dim = shape.dim.add()
126             dim.size = min_length
127             feature.shape.CopyFrom(shape)
128         elif min_length and max_length and min_length < max_length:
129             feature.value_count.CopyFrom(schema_pb2.ValueCount(min=min_length, max=max_length))
130         else:
131             # if no min max available set dummy value, to signal this is list
132             feature.value_count.CopyFrom(schema_pb2.ValueCount(min=0, max=0))
133     return feature
134 
135 
136 def set_protobuf_float(column_schema, feature):
137     domain = column_schema.properties.get("domain", {})
138     feature.float_domain.CopyFrom(
139         schema_pb2.FloatDomain(
140             name=column_schema.name,
141             min=domain.get("min", None),
142             max=domain.get("max", None),
143         )
144     )
145     feature.type = schema_pb2.FeatureType.FLOAT
146     return feature
147 
148 
149 def set_protobuf_int(column_schema, feature):
150     domain = column_schema.properties.get("domain", {})
151     feature.int_domain.CopyFrom(
152         schema_pb2.IntDomain(
153             name=column_schema.name,
154             min=domain.get("min", None),
155             max=domain.get("max", None),
156             is_categorical=(
157                 Tags.CATEGORICAL in column_schema.tags
158                 or Tags.CATEGORICAL.value in column_schema.tags
159             ),
160         )
161     )
162     feature.type = schema_pb2.FeatureType.INT
163     return feature
164 
165 
166 def register_dtype(column_schema, feature):
167     #  column_schema is a dict, changes are held
168     #  TODO: this double check can be refactored
169     if column_schema.dtype:
170         if column_schema._is_list:
171             feature = proto_dict["list"](column_schema, feature)
172         if hasattr(column_schema.dtype, "kind"):
173             string_name = numpy.core._dtype._kind_name(column_schema.dtype)
174         elif hasattr(column_schema.dtype, "item"):
175             string_name = type(column_schema.dtype(1).item()).__name__
176         elif isinstance(column_schema.dtype, str):
177             string_name = column_schema.dtype
178         elif hasattr(column_schema.dtype, "__name__"):
179             string_name = column_schema.dtype.__name__
180         else:
181             raise TypeError(f"unsupported dtype for column schema: {column_schema.dtype}")
182 
183         if string_name in proto_dict:
184             feature = proto_dict[string_name](column_schema, feature)
185     return feature
186 
187 
188 proto_dict = {
189     "list": register_list,
190     "float": set_protobuf_float,
191     "int": set_protobuf_int,
192     "uint": set_protobuf_int,
193 }
194 
195 
196 def create_protobuf_feature(column_schema):
197     feature = schema_pb2.Feature()
198     feature.name = column_schema.name
199     feature = register_dtype(column_schema, feature)
200     annotation = feature.annotation
201     annotation.tag.extend(
202         [tag.value if hasattr(tag, "value") else tag for tag in column_schema.tags]
203     )
204     # can be instantiated with no values
205     # if  so, unnecessary to dump
206     # import pdb; pdb.set_trace()
207     if len(column_schema.properties) > 0:
208         feature = register_extra_metadata(column_schema, feature)
209     return feature
210
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/nvtabular/graph/schema_io/schema_writer_pbtxt.py b/nvtabular/graph/schema_io/schema_writer_pbtxt.py
--- a/nvtabular/graph/schema_io/schema_writer_pbtxt.py
+++ b/nvtabular/graph/schema_io/schema_writer_pbtxt.py
@@ -16,6 +16,7 @@
 import os
 from pathlib import Path
 
+import fsspec
 import numpy
 
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
@@ -40,9 +41,7 @@
 
     @classmethod
     def write(cls, schema, schema_path):
-        schema_path = Path(schema_path)
-        if not schema_path.is_dir():
-            raise ValueError(f"The path provided is not a valid directory: {schema_path}")
+        fs = fsspec.get_fs_token_paths(schema_path)[0]
 
         # traverse list of column schema
         schema_file = schema_pb2.Schema()
@@ -51,9 +50,16 @@
             features.append(create_protobuf_feature(col_schema))
         schema_file.feature.extend(features)
 
-        with open(schema_path / "schema.pbtxt", "w") as f:
-            f.write(text_format.MessageToString(schema_file))
-        return schema
+        try:
+            with fs.open(fs.sep.join([str(schema_path), "schema.pbtxt"]), "w") as f:
+                f.write(text_format.MessageToString(schema_file))
+            return schema
+        except Exception as e:
+            if not fs.isdir(schema_path):
+                raise ValueError(
+                    f"The path provided is not a valid directory: {schema_path}"
+                ) from e
+            raise
 
     @classmethod
     def load(cls, schema_path):

{"golden_diff": "diff --git a/nvtabular/graph/schema_io/schema_writer_pbtxt.py b/nvtabular/graph/schema_io/schema_writer_pbtxt.py\n--- a/nvtabular/graph/schema_io/schema_writer_pbtxt.py\n+++ b/nvtabular/graph/schema_io/schema_writer_pbtxt.py\n@@ -16,6 +16,7 @@\n import os\n from pathlib import Path\n \n+import fsspec\n import numpy\n \n os.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"\n@@ -40,9 +41,7 @@\n \n     @classmethod\n     def write(cls, schema, schema_path):\n-        schema_path = Path(schema_path)\n-        if not schema_path.is_dir():\n-            raise ValueError(f\"The path provided is not a valid directory: {schema_path}\")\n+        fs = fsspec.get_fs_token_paths(schema_path)[0]\n \n         # traverse list of column schema\n         schema_file = schema_pb2.Schema()\n@@ -51,9 +50,16 @@\n             features.append(create_protobuf_feature(col_schema))\n         schema_file.feature.extend(features)\n \n-        with open(schema_path / \"schema.pbtxt\", \"w\") as f:\n-            f.write(text_format.MessageToString(schema_file))\n-        return schema\n+        try:\n+            with fs.open(fs.sep.join([str(schema_path), \"schema.pbtxt\"]), \"w\") as f:\n+                f.write(text_format.MessageToString(schema_file))\n+            return schema\n+        except Exception as e:\n+            if not fs.isdir(schema_path):\n+                raise ValueError(\n+                    f\"The path provided is not a valid directory: {schema_path}\"\n+                ) from e\n+            raise\n \n     @classmethod\n     def load(cls, schema_path):\n", "issue": "[BUG] Test_s3 failing\n```\r\n      with s3_context(s3_base=s3_base, bucket=engine, files=files) as s3fs:\r\n            # Create nvt.Dataset from mock s3 paths\r\n            url = f\"s3://{engine}\" if engine == \"parquet\" else f\"s3://{engine}/*\"\r\n            dataset = nvt.Dataset(url, engine=engine, storage_options=s3so)\r\n    \r\n            # Check that the iteration API works\r\n            columns = mycols_pq if engine == \"parquet\" else mycols_csv\r\n            gdf = nvt.dispatch._concat(list(dataset.to_iter()))[columns]\r\n            assert_eq(gdf.reset_index(drop=True), df.reset_index(drop=True))\r\n    \r\n            cat_names = [\"name-cat\", \"name-string\"] if engine == \"parquet\" else [\"name-string\"]\r\n            cont_names = [\"x\", \"y\", \"id\"]\r\n            label_name = [\"label\"]\r\n    \r\n            conts = cont_names >> ops.FillMissing() >> ops.Clip(min_value=0) >> ops.LogOp()\r\n            cats = cat_names >> ops.Categorify(cat_cache=\"host\")\r\n    \r\n            processor = nvt.Workflow(conts + cats + label_name)\r\n            processor.fit(dataset)\r\n    \r\n            # make sure we can write out the dataset back to S3\r\n            # (https://github.com/NVIDIA-Merlin/NVTabular/issues/1214)\r\n>           processor.transform(dataset).to_parquet(f\"s3://{engine}/output\")\r\n\r\n/nvtabular/tests/unit/test_s3.py:111: \r\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \r\n/nvtabular/nvtabular/io/dataset.py:906: in to_parquet\r\n    self.schema.write(output_path)\r\n/nvtabular/nvtabular/graph/schema.py:154: in write\r\n    return PbTxt_SchemaWriter.write(self, schema_path)\r\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \r\n\r\ncls = <class 'nvtabular.graph.schema_io.schema_writer_pbtxt.PbTxt_SchemaWriter'>\r\nschema = [{'name': 'x', 'tags': [<Tags.CONTINUOUS: 'continuous'>], 'properties': {}, 'dtype': <class 'float'>, '_is_list': Fals...int'>, '_is_list': False}, {'name': 'label', 'tags': [], 'properties': {}, 'dtype': dtype('int64'), '_is_list': False}]\r\nschema_path = PosixPath('s3:/csv/output')\r\n\r\n    @classmethod\r\n    def write(cls, schema, schema_path):\r\n        schema_path = Path(schema_path)\r\n        if not schema_path.is_dir():\r\n>           raise ValueError(f\"The path provided is not a valid directory: {schema_path}\")\r\nE           ValueError: The path provided is not a valid directory: s3:/csv/output\r\n\r\n/nvtabular/nvtabular/graph/schema_io/schema_writer_pbtxt.py:45: ValueError\r\n```\n", "before_files": [{"content": "#\n# Copyright (c) 2021, NVIDIA CORPORATION.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\nimport os\nfrom pathlib import Path\n\nimport numpy\n\nos.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"\nfrom google.protobuf import json_format, text_format  # noqa\nfrom google.protobuf.any_pb2 import Any  # noqa\nfrom google.protobuf.struct_pb2 import Struct  # noqa\nfrom tensorflow_metadata.proto.v0 import schema_pb2  # noqa\n\nimport nvtabular as nvt  # noqa\nfrom nvtabular.graph.schema_io.schema_writer_base import SchemaWriter  # noqa\nfrom nvtabular.graph.tags import Tags  # noqa\n\n\nclass PbTxt_SchemaWriter(SchemaWriter):\n    @classmethod\n    def _read(cls, schema_path):\n        with open(schema_path, \"r\") as f:\n            schema = schema_pb2.Schema()\n            text_format.Parse(f.read(), schema)\n\n        return schema\n\n    @classmethod\n    def write(cls, schema, schema_path):\n        schema_path = Path(schema_path)\n        if not schema_path.is_dir():\n            raise ValueError(f\"The path provided is not a valid directory: {schema_path}\")\n\n        # traverse list of column schema\n        schema_file = schema_pb2.Schema()\n        features = []\n        for col_name, col_schema in schema.column_schemas.items():\n            features.append(create_protobuf_feature(col_schema))\n        schema_file.feature.extend(features)\n\n        with open(schema_path / \"schema.pbtxt\", \"w\") as f:\n            f.write(text_format.MessageToString(schema_file))\n        return schema\n\n    @classmethod\n    def load(cls, schema_path):\n        columns = []\n        if isinstance(schema_path, (str, Path)):\n            if isinstance(schema_path, str):\n                schema_path = Path(schema_path)\n            if schema_path.is_dir():\n                schema_path = schema_path / \"schema.pbtxt\"\n            schema = cls._read(schema_path)\n\n        for feat in schema.feature:\n            _is_list = False\n            dtype = None\n            properties = {}\n            tags = list(feat.annotation.tag) or []\n            # only one item should ever be in extra_metadata\n            if len(feat.annotation.extra_metadata) > 1:\n                raise ValueError(\n                    f\"{feat.name}: extra_metadata should have 1 item, has \\\n                    {len(feat.annotation.extra_metadata)}\"\n                )\n            if feat.annotation.extra_metadata:\n                properties = json_format.MessageToDict(feat.annotation.extra_metadata[0])[\"value\"]\n            # what domain\n            # load the domain values\n            shape_name = feat.WhichOneof(\"shape_type\")\n            if shape_name:\n                _is_list = True\n            field_name = feat.WhichOneof(\"domain_info\")\n            if field_name:\n                domain_values = getattr(feat, field_name)\n                # if zero no values were passed\n                if domain_values.max > 0:\n                    properties[\"domain\"] = {\"min\": domain_values.min, \"max\": domain_values.max}\n                if feat.type:\n                    if feat.type == 2:\n                        dtype = numpy.int\n                    elif feat.type == 3:\n                        dtype = numpy.float\n            columns.append(\n                nvt.ColumnSchema(\n                    feat.name, tags=tags, properties=properties, dtype=dtype, _is_list=_is_list\n                )\n            )\n\n        return nvt.Schema(columns)\n\n\ndef register_extra_metadata(column_schema, feature):\n    filtered_properties = {k: v for k, v in column_schema.properties.items() if k != \"domain\"}\n    msg_struct = Struct()\n    # must pack message into \"Any\" type\n    any_pack = Any()\n    any_pack.Pack(json_format.ParseDict(filtered_properties, msg_struct))\n    # extra_metadata only takes type \"Any\" messages\n    feature.annotation.extra_metadata.add().CopyFrom(any_pack)\n    return feature\n\n\ndef register_list(column_schema, feature):\n    if str(column_schema._is_list):\n        min_length, max_length = None, None\n        if \"value_count\" in column_schema.properties:\n            min_length = column_schema.properties[\"value_count\"][\"min\"]\n            max_length = column_schema.properties[\"value_count\"][\"max\"]\n        if min_length and max_length and min_length == max_length:\n            shape = schema_pb2.FixedShape()\n            dim = shape.dim.add()\n            dim.size = min_length\n            feature.shape.CopyFrom(shape)\n        elif min_length and max_length and min_length < max_length:\n            feature.value_count.CopyFrom(schema_pb2.ValueCount(min=min_length, max=max_length))\n        else:\n            # if no min max available set dummy value, to signal this is list\n            feature.value_count.CopyFrom(schema_pb2.ValueCount(min=0, max=0))\n    return feature\n\n\ndef set_protobuf_float(column_schema, feature):\n    domain = column_schema.properties.get(\"domain\", {})\n    feature.float_domain.CopyFrom(\n        schema_pb2.FloatDomain(\n            name=column_schema.name,\n            min=domain.get(\"min\", None),\n            max=domain.get(\"max\", None),\n        )\n    )\n    feature.type = schema_pb2.FeatureType.FLOAT\n    return feature\n\n\ndef set_protobuf_int(column_schema, feature):\n    domain = column_schema.properties.get(\"domain\", {})\n    feature.int_domain.CopyFrom(\n        schema_pb2.IntDomain(\n            name=column_schema.name,\n            min=domain.get(\"min\", None),\n            max=domain.get(\"max\", None),\n            is_categorical=(\n                Tags.CATEGORICAL in column_schema.tags\n                or Tags.CATEGORICAL.value in column_schema.tags\n            ),\n        )\n    )\n    feature.type = schema_pb2.FeatureType.INT\n    return feature\n\n\ndef register_dtype(column_schema, feature):\n    #  column_schema is a dict, changes are held\n    #  TODO: this double check can be refactored\n    if column_schema.dtype:\n        if column_schema._is_list:\n            feature = proto_dict[\"list\"](column_schema, feature)\n        if hasattr(column_schema.dtype, \"kind\"):\n            string_name = numpy.core._dtype._kind_name(column_schema.dtype)\n        elif hasattr(column_schema.dtype, \"item\"):\n            string_name = type(column_schema.dtype(1).item()).__name__\n        elif isinstance(column_schema.dtype, str):\n            string_name = column_schema.dtype\n        elif hasattr(column_schema.dtype, \"__name__\"):\n            string_name = column_schema.dtype.__name__\n        else:\n            raise TypeError(f\"unsupported dtype for column schema: {column_schema.dtype}\")\n\n        if string_name in proto_dict:\n            feature = proto_dict[string_name](column_schema, feature)\n    return feature\n\n\nproto_dict = {\n    \"list\": register_list,\n    \"float\": set_protobuf_float,\n    \"int\": set_protobuf_int,\n    \"uint\": set_protobuf_int,\n}\n\n\ndef create_protobuf_feature(column_schema):\n    feature = schema_pb2.Feature()\n    feature.name = column_schema.name\n    feature = register_dtype(column_schema, feature)\n    annotation = feature.annotation\n    annotation.tag.extend(\n        [tag.value if hasattr(tag, \"value\") else tag for tag in column_schema.tags]\n    )\n    # can be instantiated with no values\n    # if  so, unnecessary to dump\n    # import pdb; pdb.set_trace()\n    if len(column_schema.properties) > 0:\n        feature = register_extra_metadata(column_schema, feature)\n    return feature\n", "path": "nvtabular/graph/schema_io/schema_writer_pbtxt.py"}], "after_files": [{"content": "#\n# Copyright (c) 2021, NVIDIA CORPORATION.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\nimport os\nfrom pathlib import Path\n\nimport fsspec\nimport numpy\n\nos.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"\nfrom google.protobuf import json_format, text_format  # noqa\nfrom google.protobuf.any_pb2 import Any  # noqa\nfrom google.protobuf.struct_pb2 import Struct  # noqa\nfrom tensorflow_metadata.proto.v0 import schema_pb2  # noqa\n\nimport nvtabular as nvt  # noqa\nfrom nvtabular.graph.schema_io.schema_writer_base import SchemaWriter  # noqa\nfrom nvtabular.graph.tags import Tags  # noqa\n\n\nclass PbTxt_SchemaWriter(SchemaWriter):\n    @classmethod\n    def _read(cls, schema_path):\n        with open(schema_path, \"r\") as f:\n            schema = schema_pb2.Schema()\n            text_format.Parse(f.read(), schema)\n\n        return schema\n\n    @classmethod\n    def write(cls, schema, schema_path):\n        fs = fsspec.get_fs_token_paths(schema_path)[0]\n\n        # traverse list of column schema\n        schema_file = schema_pb2.Schema()\n        features = []\n        for col_name, col_schema in schema.column_schemas.items():\n            features.append(create_protobuf_feature(col_schema))\n        schema_file.feature.extend(features)\n\n        try:\n            with fs.open(fs.sep.join([str(schema_path), \"schema.pbtxt\"]), \"w\") as f:\n                f.write(text_format.MessageToString(schema_file))\n            return schema\n        except Exception as e:\n            if not fs.isdir(schema_path):\n                raise ValueError(\n                    f\"The path provided is not a valid directory: {schema_path}\"\n                ) from e\n            raise\n\n    @classmethod\n    def load(cls, schema_path):\n        columns = []\n        if isinstance(schema_path, (str, Path)):\n            if isinstance(schema_path, str):\n                schema_path = Path(schema_path)\n            if schema_path.is_dir():\n                schema_path = schema_path / \"schema.pbtxt\"\n            schema = cls._read(schema_path)\n\n        for feat in schema.feature:\n            _is_list = False\n            dtype = None\n            properties = {}\n            tags = list(feat.annotation.tag) or []\n            # only one item should ever be in extra_metadata\n            if len(feat.annotation.extra_metadata) > 1:\n                raise ValueError(\n                    f\"{feat.name}: extra_metadata should have 1 item, has \\\n                    {len(feat.annotation.extra_metadata)}\"\n                )\n            if feat.annotation.extra_metadata:\n                properties = json_format.MessageToDict(feat.annotation.extra_metadata[0])[\"value\"]\n            # what domain\n            # load the domain values\n            shape_name = feat.WhichOneof(\"shape_type\")\n            if shape_name:\n                _is_list = True\n            field_name = feat.WhichOneof(\"domain_info\")\n            if field_name:\n                domain_values = getattr(feat, field_name)\n                # if zero no values were passed\n                if domain_values.max > 0:\n                    properties[\"domain\"] = {\"min\": domain_values.min, \"max\": domain_values.max}\n                if feat.type:\n                    if feat.type == 2:\n                        dtype = numpy.int\n                    elif feat.type == 3:\n                        dtype = numpy.float\n            columns.append(\n                nvt.ColumnSchema(\n                    feat.name, tags=tags, properties=properties, dtype=dtype, _is_list=_is_list\n                )\n            )\n\n        return nvt.Schema(columns)\n\n\ndef register_extra_metadata(column_schema, feature):\n    filtered_properties = {k: v for k, v in column_schema.properties.items() if k != \"domain\"}\n    msg_struct = Struct()\n    # must pack message into \"Any\" type\n    any_pack = Any()\n    any_pack.Pack(json_format.ParseDict(filtered_properties, msg_struct))\n    # extra_metadata only takes type \"Any\" messages\n    feature.annotation.extra_metadata.add().CopyFrom(any_pack)\n    return feature\n\n\ndef register_list(column_schema, feature):\n    if str(column_schema._is_list):\n        min_length, max_length = None, None\n        if \"value_count\" in column_schema.properties:\n            min_length = column_schema.properties[\"value_count\"][\"min\"]\n            max_length = column_schema.properties[\"value_count\"][\"max\"]\n        if min_length and max_length and min_length == max_length:\n            shape = schema_pb2.FixedShape()\n            dim = shape.dim.add()\n            dim.size = min_length\n            feature.shape.CopyFrom(shape)\n        elif min_length and max_length and min_length < max_length:\n            feature.value_count.CopyFrom(schema_pb2.ValueCount(min=min_length, max=max_length))\n        else:\n            # if no min max available set dummy value, to signal this is list\n            feature.value_count.CopyFrom(schema_pb2.ValueCount(min=0, max=0))\n    return feature\n\n\ndef set_protobuf_float(column_schema, feature):\n    domain = column_schema.properties.get(\"domain\", {})\n    feature.float_domain.CopyFrom(\n        schema_pb2.FloatDomain(\n            name=column_schema.name,\n            min=domain.get(\"min\", None),\n            max=domain.get(\"max\", None),\n        )\n    )\n    feature.type = schema_pb2.FeatureType.FLOAT\n    return feature\n\n\ndef set_protobuf_int(column_schema, feature):\n    domain = column_schema.properties.get(\"domain\", {})\n    feature.int_domain.CopyFrom(\n        schema_pb2.IntDomain(\n            name=column_schema.name,\n            min=domain.get(\"min\", None),\n            max=domain.get(\"max\", None),\n            is_categorical=(\n                Tags.CATEGORICAL in column_schema.tags\n                or Tags.CATEGORICAL.value in column_schema.tags\n            ),\n        )\n    )\n    feature.type = schema_pb2.FeatureType.INT\n    return feature\n\n\ndef register_dtype(column_schema, feature):\n    #  column_schema is a dict, changes are held\n    #  TODO: this double check can be refactored\n    if column_schema.dtype:\n        if column_schema._is_list:\n            feature = proto_dict[\"list\"](column_schema, feature)\n        if hasattr(column_schema.dtype, \"kind\"):\n            string_name = numpy.core._dtype._kind_name(column_schema.dtype)\n        elif hasattr(column_schema.dtype, \"item\"):\n            string_name = type(column_schema.dtype(1).item()).__name__\n        elif isinstance(column_schema.dtype, str):\n            string_name = column_schema.dtype\n        elif hasattr(column_schema.dtype, \"__name__\"):\n            string_name = column_schema.dtype.__name__\n        else:\n            raise TypeError(f\"unsupported dtype for column schema: {column_schema.dtype}\")\n\n        if string_name in proto_dict:\n            feature = proto_dict[string_name](column_schema, feature)\n    return feature\n\n\nproto_dict = {\n    \"list\": register_list,\n    \"float\": set_protobuf_float,\n    \"int\": set_protobuf_int,\n    \"uint\": set_protobuf_int,\n}\n\n\ndef create_protobuf_feature(column_schema):\n    feature = schema_pb2.Feature()\n    feature.name = column_schema.name\n    feature = register_dtype(column_schema, feature)\n    annotation = feature.annotation\n    annotation.tag.extend(\n        [tag.value if hasattr(tag, \"value\") else tag for tag in column_schema.tags]\n    )\n    # can be instantiated with no values\n    # if  so, unnecessary to dump\n    # import pdb; pdb.set_trace()\n    if len(column_schema.properties) > 0:\n        feature = register_extra_metadata(column_schema, feature)\n    return feature\n", "path": "nvtabular/graph/schema_io/schema_writer_pbtxt.py"}]}