Reorganize UDF data codecs to prepare for better serialization. (#463)

To better separate concerns and eliminate cross-package circular
dependencies, this change moves the main handling of the `tiledb_json`
format to the `_results` package, and restructures the way codecs are
defined and handled.

This will allow us to introduce `tiledb_json` as a format for
serializing UDF results, which will allow us to use the `tiledb_json`
format to serialize Pandas dataframes, with custom encoding/decoding
logic based on Arrow, instead of using native pickles.  This will allow
us to offer client/server compatibility across Pandas versions, without
the hacks we currently resort to.

src/tiledb/cloud/_results/codecs.py

@@ -0,0 +1,216 @@
+import abc
+import base64
+import json
+from typing import Any, Generic, Tuple, Type, TypeVar
+
+import attrs
+import cloudpickle
+import pyarrow
+import urllib3
+from typing_extensions import Self
+
+# This is a circular dependency since we need to be able to decode `tiledb_json`
+# format data.
+from . import tiledb_json
+from . import types
+
+_ARROW_VERSION = pyarrow.MetadataVersion.V5
+_PICKLE_PROTOCOL = 4
+_T = TypeVar("_T")
+
+
+class Codec(Generic[_T], metaclass=abc.ABCMeta):
+ """Translates objects to bytes and vice versa. Purely classmethods."""
+
+ __slots__ = ()
+
+ NAME: str
+ """The name to use for the codec, as used in ``result_format``."""
+
+ MIME: str
+ """The MIME type identifying this codec."""
+
+ @classmethod
+ @abc.abstractmethod
+ def encode(cls, obj: _T) -> bytes:
+ raise NotImplementedError()
+
+ @classmethod
+ @abc.abstractmethod
+ def decode(cls, data: bytes) -> _T:
+ raise NotImplementedError()
+
+ @classmethod
+ def decode_base64(cls, data: str) -> _T:
+ data_bytes = base64.b64decode(data)
+ return cls.decode(data_bytes)
+
+ @classmethod
+ def encode_base64(cls, obj: _T) -> str:
+ data_bytes = cls.encode(obj)
+ return base64.b64encode(data_bytes).decode("utf-8")
+
+
+class ArrowCodec(Codec[pyarrow.Table]):
+ """Encodes Arrow data into its default stream format."""
+
+ NAME = "arrow"
+ MIME = "application/vnd.apache.arrow.stream"
+
+ @classmethod
+ def encode(cls, tbl: pyarrow.Table) -> bytes:
+ sink = pyarrow.BufferOutputStream()
+ writer = pyarrow.RecordBatchStreamWriter(
+ sink,
+ tbl.schema,
+ options=pyarrow.ipc.IpcWriteOptions(
+ metadata_version=_ARROW_VERSION,
+ compression="zstd",
+ ),
+ )
+ writer.write(tbl)
+ return sink.getvalue()
+
+ @classmethod
+ def decode(cls, data: bytes) -> pyarrow.Table:
+ # If a UDF didn't return any rows, there will not have been any batches
+ # of data to write to the output, and thus it will not include any content
+ # at all. (SQL queries will include headers.)
+ if not data:
+ # In this case, we need to return an empty table.
+ return pyarrow.Table.from_pydict({})
+ reader = pyarrow.RecordBatchStreamReader(data)
+ return reader.read_all()
+
+
+class BytesCodec(Codec[bytes]):
+ """Does nothing to bytes."""
+
+ NAME = "bytes"
+ MIME = "application/octet-stream"
+
+ @classmethod
+ def encode(cls, obj: bytes) -> bytes:
+ return obj
+
+ @classmethod
+ def decode(cls, data: bytes) -> bytes:
+ return data
+
+
+class JSONCodec(Codec[object]):
+ """Dumps/loads JSON."""
+
+ NAME = "json"
+ MIME = "application/json"
+
+ @classmethod
+ def encode(cls, obj: object) -> bytes:
+ return json.dumps(obj).encode("utf-8")
+
+ @classmethod
+ def decode(cls, data: bytes) -> object:
+ return json.loads(data)
+
+
+class PickleCodec(Codec[object]):
+ """Pickles objects using CloudPickle."""
+
+ NAME = "python_pickle"
+ MIME = "application/vnd.tiledb.python-pickle"
+
+ @classmethod
+ def encode(cls, obj: object) -> bytes:
+ return cloudpickle.dumps(obj, protocol=_PICKLE_PROTOCOL)
+
+ @classmethod
+ def decode(cls, data: bytes) -> object:
+ return cloudpickle.loads(data)
+
+
+class TileDBJSONCodec(Codec[object]):
+ """Serializes objects with TileDB JSON."""
+
+ NAME = "tiledb_json"
+ MIME = "application/vnd.tiledb.udf-data+json"
+
+ @classmethod
+ def encode(cls, obj: object) -> bytes:
+ return tiledb_json.dumps(obj)
+
+ @classmethod
+ def decode(cls, data: bytes) -> object:
+ return tiledb_json.loads(data)
+
+
+ALL_CODECS: Tuple[Type[Codec[Any]], ...] = (
+ ArrowCodec,
+ BytesCodec,
+ JSONCodec,
+ PickleCodec,
+ TileDBJSONCodec,
+)
+"""Every codec we have."""
+CODECS_BY_FORMAT = {c.NAME: c for c in ALL_CODECS}
+CODECS_BY_FORMAT["native"] = PickleCodec
+CODECS_BY_MIME = {c.MIME: c for c in ALL_CODECS}
+
+
+@attrs.define(frozen=True, slots=False)
+class BinaryBlob:
+ """Container for a binary-encoded value, decoded on-demand.
+
+ This is used to store results obtained from the server, such that it is not
+ necessary to decode them between stages, and they only are decoded upon
+ request.
+ """
+
+ format: str
+ """The TileDB Cloud name of the data's format (see ``CODECS_BY_FORMAT``)."""
+ data: bytes
+ """The binary data itself."""
+
+ @classmethod
+ def from_response(cls, resp: urllib3.HTTPResponse) -> Self:
+ """Reads a urllib3 response into an encoded result."""
+ full_mime = resp.getheader("Content-type") or "application/octet-stream"
+ mime, _, _ = full_mime.partition(";")
+ mime = mime.strip()
+ try:
+ format_name = CODECS_BY_MIME[mime].NAME
+ except KeyError:
+ format_name = "mime:" + mime
+ data = resp.data
+ return cls(format_name, data)
+
+ def decode(self) -> types.NativeValue:
+ """Decodes this result into native Python data."""
+ # This is not lock-protected because we're ok with decoding twice.
+ try:
+ return self.__dict__["_decoded"]
+ except KeyError:
+ pass
+ try:
+ loader = CODECS_BY_FORMAT[self.format].decode
+ except KeyError:
+ raise ValueError(f"Cannot decode {self.format!r} data")
+ self.__dict__["_decoded"] = loader(self.data)
+ return self.__dict__["_decoded"]
+
+ def _tdb_to_json(self) -> types.TileDBJSONValue:
+ return types.TileDBJSONValue(
+ {
+ "__tdbudf__": "immediate",
+ "format": self.format,
+ "base64_data": base64.b64encode(self.data).decode("ascii"),
+ }
+ )
+
+ @classmethod
+ def of(cls, obj: Any) -> Self:
+ """Turns a non–JSON-encodable object into a ``BinaryBlob``."""
+ if isinstance(obj, bytes):
+ return cls("bytes", obj)
+ if isinstance(obj, pyarrow.Table):
+ return cls("arrow", ArrowCodec.encode(obj))
+ return cls("python_pickle", PickleCodec.encode(obj))

src/tiledb/cloud/_results/decoders.py

@@ -2,15 +2,15 @@
 
 import abc
 import dataclasses
-import json
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
-import cloudpickle
 import pyarrow
 
 from tiledb.cloud import tiledb_cloud_error as tce
 from tiledb.cloud.rest_api import models
 
+from . import codecs
+
 if TYPE_CHECKING:
  import pandas
 
@@ -25,24 +25,6 @@ def decode(self, data: bytes) -> Any:
  raise NotImplementedError()
 
 
-def _load_arrow(data: bytes) -> pyarrow.Table:
- # If a UDF didn't return any rows, there will not have been any batches
- # of data to write to the output, and thus it will not include any content
- # at all. (SQL queries will include headers.)
- if not data:
- # In this case, we need to return an empty table.
- return pyarrow.Table.from_pydict({})
- reader = pyarrow.RecordBatchStreamReader(data)
- return reader.read_all()
-
-
-_DECODE_FNS = {
- models.ResultFormat.NATIVE: cloudpickle.loads,
- models.ResultFormat.JSON: json.loads,
- models.ResultFormat.ARROW: _load_arrow,
-}
-
-
 @dataclasses.dataclass(frozen=True)
 class Decoder(AbstractDecoder[_T]):
  """General decoder for the formats we support.
@@ -55,10 +37,10 @@ class Decoder(AbstractDecoder[_T]):
 
  def decode(self, data: bytes) -> _T:
  try:
- decoder = _DECODE_FNS[self.format]
+ codec = codecs.CODECS_BY_FORMAT[self.format]
  except KeyError:
  raise tce.TileDBCloudError(f"{self.format!r} is not a valid result format.")
- return decoder(data)
+ return codec.decode(data)
 
 
 @dataclasses.dataclass(frozen=True)

src/tiledb/cloud/_results/tiledb_json.py

@@ -0,0 +1,125 @@
+import json
+from typing import Any, Dict, Optional
+
+from .._common import visitor
+from . import codecs
+
+SENTINEL_KEY = "__tdbudf__"
+ESCAPE_CODE = "__escape__"
+
+
+class Encoder(visitor.ReplacingVisitor):
+ """Turns arbitrary Python values into TileDB JSON.
+
+ This escapes arbitrary native values so that they can be JSON-serialized.
+ It should only be used with ``NativeValue``s—that is, values that are
+ already JSON-serializable, like ``RegisteredArg``s or ``CallArg``s, should
+ *not* be passed to an ``Escaper``. The base implementation will return
+ fully self-contained JSON-serializable objects, i.e. ``CallArg``s.
+ """
+
+ def maybe_replace(self, arg) -> Optional[visitor.Replacement]:
+ if is_jsonable_shallow(arg):
+ if isinstance(arg, dict):
+ if SENTINEL_KEY in arg:
+ return visitor.Replacement(
+ {
+ SENTINEL_KEY: ESCAPE_CODE,
+ ESCAPE_CODE: {k: self.visit(v) for (k, v) in arg.items()},
+ }
+ )
+ return None
+ return visitor.Replacement(codecs.BinaryBlob.of(arg)._tdb_to_json())
+
+
+class Decoder(visitor.ReplacingVisitor):
+ """A general-purpose replacer to decode sentinel-containing structures.
+
+ This descends through data structures and replaces dictionaries containing
+ ``__tdbudf__`` values with the unescaped values. This base implementation
+ handles only the basics; you can create a derived version to handle specific
+ situations (building arguments, replacing values, etc.).
+
+ The data that is returned from this is generally a ``types.NativeValue``.
+ """
+
+ def maybe_replace(self, arg) -> Optional[visitor.Replacement]:
+ if not isinstance(arg, dict):
+ return None
+ try:
+ sentinel_name = arg[SENTINEL_KEY]
+ except KeyError:
+ return None
+ return self._replace_sentinel(sentinel_name, arg)
+
+ def _replace_sentinel(
+ self,
+ kind: str,
+ value: Dict[str, Any],
+ ) -> Optional[visitor.Replacement]:
+ """The base implementation of a sentinel-replacer.
+
+ It is passed the kind and value of a ``__tdbudf__``–containing object::
+
+ # Given this:
+ the_object = {"__tdbudf__": "node_data", "data": "abc"}
+ # This will be called:
+ self._replace_sentinel("node_data", the_object)
+
+ This implementation handles replacing values that do not require any
+ external information. Derived implementations should handle their own
+ keys and end with a call to
+ ``return super()._replace_sentinel(kind, value)``.
+ """
+ if kind == ESCAPE_CODE:
+ # An escaped value.
+ inner_value = value[ESCAPE_CODE]
+ # We can't just visit `inner_value` here, since `inner_value`
+ # is the thing which has the ESCAPE_CODE key that is being escaped.
+ return visitor.Replacement(
+ {k: self.visit(v) for (k, v) in inner_value.items()}
+ )
+ if kind == "immediate":
+ # "immediate" values are values of the format
+ fmt = value["format"]
+ base64d = value["base64_data"]
+ return visitor.Replacement(
+ codecs.CODECS_BY_FORMAT[fmt].decode_base64(base64d)
+ )
+ raise ValueError(f"Unknown sentinel type {kind!r}")
+
+
+def dumps(obj: object) -> bytes:
+ """Dumps an object to TileDB UDF–encoded JSON."""
+ enc = Encoder()
+ result_json = enc.visit(obj)
+ return json.dumps(result_json).encode("utf-8")
+
+
+def loads(data: bytes) -> object:
+ """Loads TileDB UDF–encoded JSON to an object."""
+ data_json = json.loads(data)
+ dec = Decoder()
+ return dec.visit(data_json)
+
+
+_NATIVE_JSONABLE = (
+ str,
+ int,
+ bool,
+ float,
+ type(None),
+ list,
+ tuple,
+)
+"""Types whose direct values can always be turned into JSON."""
+
+
+def is_jsonable_shallow(obj) -> bool:
+ if isinstance(obj, _NATIVE_JSONABLE):
+ return True
+ if not isinstance(obj, dict):
+ # Apart from the above types, only dicts are JSONable.
+ return False
+ # For a dict to be JSONable, all keys must be strings.
+ return all(isinstance(key, str) for key in obj)