Core: JSONL/NDJSON support, and invoke single resource on project file

Author: amotl

CHANGES.md

@@ -2,7 +2,9 @@
 
 ## Unreleased
 - CLI: Added CLI interface, per `tikray` program
-- Core: Started using `orjson` package
+- Core: Started using `orjson` and `orjsonl` packages, for performance
+  reasons and JSONL/NDJSON support
+- Core: Optionally invoke single resource on project file
 
 ## 2025/02/05 v0.0.23
 - Renamed package to `loko`, then `tikray`

docs/backlog.md

@@ -1,6 +1,11 @@
 # Backlog
 
 ## Iteration +1
+- Support for JSONL
+- Read and write compressed files. orjson uses `xopen` already, which can do it transparently
+- Allow processing larger-than-memory files
+
+## Iteration +2
 - [ ] Documentation: jqlang stdlib's `to_object` function for substructure management
 - [ ] Documentation: Type casting
   `echo '{"a": 42, "b": {}, "c": []}' | jq -c '.|= (.b |= objects | .c |= objects)'`
@@ -36,8 +41,7 @@ and ._id != "55d71c8ce4b02210dc47b10f"
 #and (.value.urls | type) != "object"
 ```
 
-
-## Iteration +2
+## Iteration +3
 - [ ] CLI interface
 - [ ] Documentation: Add Python example to "Synopsis" section on /index.html
 - [ ] Documentation: Compare with Seatunnel
@@ -55,7 +59,7 @@ Demonstrate more use cases, like...
   - https://github.com/MeltanoLabs/meltano-map-transform/issues/252
 - [ ] Use JSONPath, see https://sdk.meltano.com/en/v0.39.1/code_samples.html#use-a-jsonpath-expression-to-extract-the-next-page-url-from-a-hateoas-response
 
-## Iteration +3
+## Iteration +4
 - [ ] Moksha transformations on Buckets
 - [ ] Fluent API interface
   ```python

docs/cli.md

@@ -20,18 +20,41 @@ tikray \
 
 ### Single collection
 
-Process a single collection / file.
+Process a single resource (collection, file) using a Tikray collection transformation.
 ```shell
 tikray -t transformation-collection.yaml -i eai-warehouse.json
 ```
 
 ### Multiple collections
 
-Process multiple collections / files from a directory.
+Process multiple resources (collections, files) from a directory.
+The Tikray project file enumerates multiple transformation rules per resource.
 ```shell
 tikray -t examples/transformation-project.yaml -i examples/acme -o tmp/acme
 ```
 
+If you are using a Tikray project file, but would like to only invoke a
+single-resource transformation on it, you need to explicitly specify the
+resource address using `--address`/`-a`, so the engine will only select
+this particular collection.
+```shell
+tikray -t examples/transformation-project.yaml -i examples/acme/conversation.json -a acme.conversation
+```
+
+### JSONL support
+
+Tikray supports reading and writing the JSONL/NDJSON format, i.e. newline-
+delimited JSON. Tikray will automatically use this mode if it receives
+input files suffixed with `.jsonl` or `.ndjson`, or if you explicitly
+toggle the mode per `--jsonl` option flag.
+```shell
+tikray -t transformation.yaml -i input.jsonl
+```
+```shell
+tikray -t transformation.yaml -i input.json -a acme.conversation --jsonl -o output.jsonl
+```
+
+
 ## Example
 
 `$ cat eai-warehouse.json`

pyproject.toml

@@ -101,6 +101,7 @@ dependencies = [
   "jq<1.9",
   "jsonpointer<4",
   "orjson<4",
+  "orjsonl<2",
   "python-dateutil<2.10",
   "pyyaml<7",
   "toolz<0.13",

src/tikray/cli.py

@@ -14,23 +14,26 @@
 @click.option("--transformation", "-t", type=Path, required=True, help="Transformation YAML file")
 @click.option("--input", "-i", "input_", type=Path, required=True, help="Input file or directory")
 @click.option("--output", "-o", type=Path, required=False, help="Output file or directory")
+@click.option("--address", "-a", type=str, required=False, help="Select specific collection address")
+@click.option("--jsonl", type=bool, is_flag=True, required=False, help="Select JSONL/NDJSON processing")
 @click.version_option()
 @click.pass_context
-def cli(ctx: click.Context, transformation: Path, input_: Path, output: Path) -> None:
+def cli(ctx: click.Context, transformation: Path, input_: Path, output: Path, address: str, jsonl: bool) -> None:
     setup_logging()
     tdata = yaml.safe_load(transformation.read_text())
     type_ = tdata["meta"]["type"]
 
-    if type_ in ["tikray-project", "zyp-project"]:
+    if type_ in ["tikray-collection", "zyp-collection"] or address is not None:  # noqa: RET506
+        return process_collection(transformation, input_, output, address, jsonl)
+
+    elif type_ in ["tikray-project", "zyp-project"]:
         if not input_.is_dir():
             raise click.ClickException(f"Input is not a directory: {input_}")
         if output is None:
             raise click.ClickException("Processing multiple collections requires an output directory")
         if not output.is_dir():
             raise click.ClickException(f"Output is not a directory: {output}")
-        return process_project(transformation, input_, output)
+        return process_project(transformation, input_, output, jsonl)
 
-    elif type_ in ["tikray-collection", "zyp-collection"]:  # noqa: RET506
-        return process_collection(transformation, input_, output)
     else:
         raise NotImplementedError(f"Unknown transformation type: {type_}")

src/tikray/core.py

@@ -1,18 +1,15 @@
 import logging
-import sys
 import typing as t
 from pathlib import Path
 
-import orjson as json
-
 from tikray.model.collection import CollectionAddress, CollectionTransformation
 from tikray.model.project import TransformationProject
-from tikray.util.data import jd
+from tikray.util.data import load_json, save_json
 
 logger = logging.getLogger(__name__)
 
 
-def process_project(transformation: Path, input_: Path, output: Path):
+def process_project(transformation: Path, input_: Path, output: Path, use_jsonl: bool = False):
     logger.info(f"Using transformation '{transformation}' on multi-collection input '{input_}'")
 
     project = TransformationProject.from_yaml(transformation.read_text())
@@ -24,22 +21,29 @@ def process_project(transformation: Path, input_: Path, output: Path):
         except KeyError as ex:
             logger.warning(f"Could not find transformation definition for collection: {ex}")
             continue
-        data = json.loads(Path(item).read_text())
+        data = load_json(Path(item), use_jsonl=use_jsonl)
         output_path = output / item.name
-        with open(output_path, "wb") as output_stream:
-            output_stream.write(jd(tikray_transformation.apply(data)))
-            logger.info(f"Processed output: {output_path}")
+        save_json(tikray_transformation.apply(data), output_path, use_jsonl=use_jsonl)
+        logger.info(f"Processed output: {output_path}")
 
 
-def process_collection(transformation: Path, input_: Path, output: t.Optional[Path] = None):
+def process_collection(
+    transformation: Path,
+    input_: Path,
+    output: t.Optional[Path] = None,
+    address: t.Optional[str] = None,
+    use_jsonl: bool = False,
+):
     logger.info(f"Using transformation '{transformation}' on single-collection input '{input_}'")
-    data = json.loads(input_.read_text())
     ct = CollectionTransformation.from_yaml(transformation.read_text())
+    if address is not None:
+        pt = TransformationProject.from_yaml(transformation.read_text())
+        ct = pt.get(CollectionAddress(*address.split(".")))
+    logger.info(f"Processing input: {input_}")
+    data = load_json(input_, use_jsonl=use_jsonl)
     result = ct.apply(data)
-    output_stream = sys.stdout.buffer
     if output is not None:
         if output.is_dir():
             output = output / input_.name
-        output_stream = open(output, "wb")
-    output_stream.write(jd(result))
-    output_stream.flush()
+    save_json(result, output, use_jsonl=use_jsonl)
+    logger.info(f"Processed output: {output or 'stdout'}")

src/tikray/util/data.py

@@ -1,6 +1,9 @@
+import sys
 import typing as t
+from pathlib import Path
 
-import orjson as json
+import orjson
+import orjsonl
 
 
 def no_privates_no_nulls_no_empties(key, value) -> bool:
@@ -22,5 +25,26 @@ def no_disabled_false(key, value):
     return not (key.name == "disabled" and value is False)
 
 
-def jd(thing: t.Any):
-    return json.dumps(thing)
+def load_json(path: Path, use_jsonl: bool = False) -> t.Any:
+    if path.suffix in [".jsonl", ".ndjson"] or use_jsonl:
+        return orjsonl.load(path)
+    else:
+        return orjson.loads(path.read_text())
+
+
+def save_json(data: t.Any, path: t.Optional[Path] = None, use_jsonl: bool = False) -> None:
+    # Sanity checks.
+    if use_jsonl and not path:
+        raise NotImplementedError("JSONL not supported on STDOUT yet, please raise an issue")
+
+    # Output JSONL.
+    if path and (path.suffix in [".jsonl", ".ndjson"] or use_jsonl):
+        orjsonl.save(path, data)
+
+    # Output JSON.
+    elif path:
+        with open(path, "wb") as stream:
+            stream.write(orjson.dumps(data))
+
+    else:
+        sys.stdout.buffer.write(orjson.dumps(data))

tests/test_cli.py

@@ -8,10 +8,12 @@
     {"id": 34, "meta": {"name": "bar", "location": "BY"}, "data": {"value": -84.01}},
 ]
 
+acme_conversation_reference = json.loads((Path("tests") / "examples" / "conversation.json").read_text())
+
 
 def test_cli_collection_stdout_success(cli_runner):
     """
-    CLI test: Invoke `tikray` with example data.
+    CLI test: Single resource to STDOUT.
     """
 
     result = cli_runner.invoke(
@@ -26,7 +28,7 @@ def test_cli_collection_stdout_success(cli_runner):
 
 def test_cli_collection_file_output_success(cli_runner, tmp_path):
     """
-    CLI test: Invoke `tikray` with example data.
+    CLI test: Single resource to file.
     """
 
     output_path = tmp_path / "output.json"
@@ -43,7 +45,7 @@ def test_cli_collection_file_output_success(cli_runner, tmp_path):
 
 def test_cli_collection_directory_output_success(cli_runner, tmp_path):
     """
-    CLI test: Invoke `tikray` with example data.
+    CLI test: Single resource to directory.
     """
 
     result = cli_runner.invoke(
@@ -60,7 +62,7 @@ def test_cli_collection_directory_output_success(cli_runner, tmp_path):
 
 def test_cli_project_success(cli_runner, tmp_path):
     """
-    CLI test: Invoke `tikray` with example data.
+    CLI test: Multiple resources (project) to directory.
     """
 
     result = cli_runner.invoke(
@@ -69,14 +71,28 @@ def test_cli_project_success(cli_runner, tmp_path):
         catch_exceptions=False,
     )
     assert result.exit_code == 0
-    reference = json.loads((Path("tests") / "examples" / "conversation.json").read_text())
     output = json.loads(Path(tmp_path / "conversation.json").read_text())
-    assert output == reference
+    assert output == acme_conversation_reference
+
+
+def test_cli_collection_from_project_file_success(cli_runner, tmp_path):
+    """
+    CLI test: Single resource from Tikray project file.
+    """
+
+    result = cli_runner.invoke(
+        cli,
+        args="-t examples/transformation-project.yaml -i examples/acme/conversation.json -a acme.conversation",
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    data = json.loads(result.output)
+    assert data == acme_conversation_reference
 
 
 def test_cli_project_warning_no_transformation(cli_runner, tmp_path, caplog):
     """
-    CLI test: Invoke `tikray` with example data.
+    CLI test: Verify processing multiple resources emits warnings on missing ones.
     """
 
     project = tmp_path / "project"

tests/test_util.py

@@ -1,8 +1,12 @@
+from pathlib import Path
+
 import jmespath
 import jq
 import jsonpointer
+import orjsonl
 import pytest
 import transon
+from tikray.util.data import load_json, save_json
 from tikray.util.expression import compile_expression
 from tikray.util.locator import to_pointer
 
@@ -48,3 +52,17 @@ def test_compile_expression_unknown():
     with pytest.raises(TypeError) as ex:
         compile_expression(type="foobar", expression=None)
     assert ex.match("Compilation failed. Type must be either jmes or jq or transon: foobar")
+
+
+def test_load_jsonl_by_suffix(tmp_path: Path):
+    data = load_json(Path("examples/eai-warehouse.json"))
+    tmp_path = tmp_path / "testdrive.jsonl"
+    orjsonl.save(tmp_path, [data])
+    assert load_json(tmp_path) == [data]
+
+
+def test_load_jsonl_by_flag(tmp_path: Path):
+    data = load_json(Path("examples/eai-warehouse.json"))
+    tmp_path = tmp_path / "testdrive.json"
+    save_json([data], tmp_path, use_jsonl=True)
+    assert load_json(tmp_path, use_jsonl=True) == [data]