Add write support to json from pd.DataFrame (AmpX-AI#31)

add json write support

Author: ampx-mg

src/fsql/api.py

@@ -147,6 +147,9 @@ def write_object(
         elif format == "csv":
             with fs.open(url_suff, "wb") as fd:
                 data.to_csv(fd)
+        elif format == "json":
+            with fs.open(url_suff, "wb") as fd:
+                data.to_json(fd)
         else:
             raise ValueError(f"unsupported format for dataframe writing: {format}")
     elif isinstance(data, io.StringIO) or isinstance(data, io.BytesIO):

tests/conftest.py

@@ -1,4 +1,5 @@
 """Sets a mock/fake of the S3 filesystem for any `fsql`-based usage."""
+import json
 import os
 
 import fsspec
@@ -14,6 +15,10 @@ def put_s3_file(self, data, url):
         with self.s3fs.open(url, "wb") as fd:
             fd.write(data)
 
+    def read_json_file(self, url):
+        with self.s3fs.open(url, "r") as fd:
+            return json.load(fd)
+
 
 @pytest.fixture
 def helper():

tests/test_write_object.py

@@ -1,4 +1,5 @@
 import io
+import json
 
 import pandas as pd
 import pytest
@@ -87,3 +88,28 @@ def test_write_vanilla_bytes(tmpdir):
     with open(path, "rb") as f:
         extracted = f.read()
     assert extracted == data
+
+
+def test_write_json_s3(helper):
+    """Writes a dataframe as json, tests that read works."""
+    bucket = "test-bouquet"
+    fs = helper.s3fs
+    fs.mkdir(bucket)
+
+    input = pd.DataFrame({"k1": [1, 2], "k2": [3, 4]}, index=["one", "two"])
+    write_object("s3://test-bouquet/my_df.json", input, format="json")
+    output = helper.read_json_file("s3://test-bouquet/my_df.json")
+    expected_output = {"k1": {"one": 1, "two": 2}, "k2": {"one": 3, "two": 4}}
+    assert output == expected_output
+
+
+def test_write_json(tmpdir):
+    """Writes a json file, tests that read works."""
+    input = pd.DataFrame({"k1": [1, 2], "k2": [3, 4]}, index=["one", "two"])
+    path_base = tmpdir.join("my_file.json")
+    url = f"file://{path_base}"
+    write_object(url, input, format="json")
+    with open(path_base, "r") as f:
+        output = json.load(f)
+    expected_output = {"k1": {"one": 1, "two": 2}, "k2": {"one": 3, "two": 4}}
+    assert output == expected_output