New blosc2.struct() for declaring columns (useful for Arrow/Parquet struct<...> type)

Author: FrancescAlted

doc/reference/ctable.rst

@@ -549,14 +549,35 @@ Text & binary
     bytes
     vlstring
     vlbytes
+    struct
     list
 
 .. autoclass:: string
 .. autoclass:: bytes
 .. autofunction:: vlstring
 .. autofunction:: vlbytes
+.. autofunction:: struct
 .. autofunction:: list
 
+Struct columns
+--------------
+
+Struct columns are declared with :func:`blosc2.struct` and store one dictionary
+(or ``None`` when nullable) per row in batched variable-length storage.  They are
+also used when importing top-level Arrow/Parquet ``struct<...>`` columns::
+
+    from dataclasses import dataclass
+    import blosc2 as b2
+
+    @dataclass
+    class Product:
+        properties: dict = b2.field(
+            b2.struct({"code": b2.int32(), "label": b2.vlstring()}, nullable=True)
+        )
+
+    table.append([{"code": 1, "label": "fresh"}])
+    table.append([None])
+
 List columns
 ------------
 

src/blosc2/ctable.py

@@ -535,7 +535,7 @@ def is_list(self) -> bool:
     def is_varlen_scalar(self) -> bool:
         """True if this column holds variable-length scalar strings or bytes."""
         col = self._table._schema.columns_by_name.get(self._col_name)
-        return col is not None and isinstance(col.spec, (VLStringSpec, VLBytesSpec))
+        return col is not None and isinstance(col.spec, (VLStringSpec, VLBytesSpec, StructSpec))
 
     @property
     def _valid_rows(self):
@@ -1581,7 +1581,7 @@ def _is_list_column(col: CompiledColumn) -> bool:
 
     @staticmethod
     def _is_varlen_scalar_column(col: CompiledColumn) -> bool:
-        return isinstance(col.spec, (VLStringSpec, VLBytesSpec))
+        return isinstance(col.spec, (VLStringSpec, VLBytesSpec, StructSpec))
 
     @staticmethod
     def _is_list_spec(spec: SchemaSpec) -> bool:
@@ -1644,7 +1644,7 @@ def _resolve_nullable_specs(
         for col in schema.columns:
             spec = col.spec
             if (
-                isinstance(spec, (ListSpec, VLStringSpec, VLBytesSpec))
+                isinstance(spec, (ListSpec, VLStringSpec, VLBytesSpec, StructSpec))
                 or getattr(spec, "null_value", None) is not None
             ):
                 continue
@@ -2913,14 +2913,21 @@ def _arrow_type_to_spec(  # noqa: C901
             for field in pa_type:
                 child_col = None
                 if arrow_col is not None:
-                    child_col = arrow_col.field(field.name)
+                    combined = (
+                        arrow_col.combine_chunks() if hasattr(arrow_col, "combine_chunks") else arrow_col
+                    )
+                    child_col = combined.field(field.name)
                 child_string_max_length = string_max_length
                 if field.type in (pa.string(), pa.large_string(), pa.utf8(), pa.large_utf8()):
                     child_string_max_length = max(string_max_length or 1, 1_000_000)
                 fields[field.name] = CTable._arrow_type_to_spec(
-                    pa, field.type, child_col, string_max_length=child_string_max_length
+                    pa,
+                    field.type,
+                    child_col,
+                    string_max_length=child_string_max_length,
+                    nullable=field.nullable,
                 )
-            return b2s.struct(fields)
+            return b2s.struct(fields, nullable=nullable)
 
         if pa_type in (pa.string(), pa.large_string(), pa.utf8(), pa.large_utf8()):
             if string_max_length is None:
@@ -2938,7 +2945,7 @@ def _arrow_type_to_spec(  # noqa: C901
 
         raise TypeError(
             f"No blosc2 spec for Arrow type {pa_type!r}. Supported: int8/16/32/64, "
-            "uint8/16/32/64, float32/64, bool, string, binary, and list."
+            "uint8/16/32/64, float32/64, bool, string, binary, list, and struct."
         )
 
     @staticmethod
@@ -3197,8 +3204,11 @@ def from_arrow(
         When *string_max_length* is ``None`` (the default), scalar Arrow
         ``string`` / ``large_string`` columns are imported as
         :func:`~blosc2.vlstring` columns and ``binary`` / ``large_binary``
-        columns are imported as :func:`~blosc2.vlbytes` columns.  Null values
-        are represented as native ``None`` with no sentinel needed.
+        columns are imported as :func:`~blosc2.vlbytes` columns.  Arrow
+        ``struct`` columns are imported as :func:`~blosc2.struct` columns backed
+        by batched variable-length storage.  Null values for these variable-
+        length scalar columns are represented as native ``None`` with no
+        sentinel needed.
 
         When *string_max_length* is set to a positive integer, scalar string
         and binary columns are imported as fixed-width
@@ -3208,9 +3218,10 @@ def from_arrow(
         :func:`~blosc2.vlstring` / :func:`~blosc2.vlbytes` columns.
 
         ``blosc2_batch_size`` controls how many rows are buffered before
-        BatchArray-backed imported columns (list columns and varlen scalar
-        columns) are flushed to their backend.  Set it to ``None`` to keep
-        those columns pending until the final flush.
+        BatchArray-backed imported columns (list columns and variable-length
+        scalar columns such as ``vlstring``, ``vlbytes``, and ``struct``) are
+        flushed to their backend.  Set it to ``None`` to keep those columns
+        pending until the final flush.
         """
         pa = cls._require_pyarrow("from_arrow()")
         if blosc2_batch_size is not None and blosc2_batch_size <= 0:
@@ -3314,7 +3325,9 @@ def from_parquet(
 
         This method delegates the actual table construction to
         :meth:`CTable.from_arrow`, so Arrow schema handling, nullable-column support,
-        and Blosc2 write tuning follow the same rules as that method.
+        and Blosc2 write tuning follow the same rules as that method.  Top-level
+        Arrow ``struct<...>`` columns are imported as :func:`~blosc2.struct`
+        columns backed by batched variable-length storage.
 
         Parameters
         ----------
@@ -3900,7 +3913,7 @@ def _fetch_col_at_positions(self, name: str, positions: np.ndarray):
             )
         col = self._cols[name]
         spec = self._schema.columns_by_name[name].spec
-        if self._is_list_spec(spec) or isinstance(spec, (VLStringSpec, VLBytesSpec)):
+        if self._is_list_spec(spec) or isinstance(spec, (VLStringSpec, VLBytesSpec, StructSpec)):
             return col[positions]
         return col[positions]
 
@@ -4473,7 +4486,7 @@ def _normalise_sort_keys(
             dtype = self._col_dtype(name)
             if dtype is None:
                 cc = self._schema.columns_by_name.get(name)
-                if cc is not None and isinstance(cc.spec, (VLStringSpec, VLBytesSpec)):
+                if cc is not None and isinstance(cc.spec, (VLStringSpec, VLBytesSpec, StructSpec)):
                     raise TypeError(
                         f"Column {name!r} is a varlen scalar column and does not support sort ordering."
                     )
@@ -5364,10 +5377,10 @@ def create_index(  # noqa: C901
         col_arr = self._cols[col_name]
         if isinstance(self._schema.columns_by_name[col_name].spec, ListSpec):
             raise ValueError(f"Cannot create an index on list column {col_name!r} in V1.")
-        if isinstance(self._schema.columns_by_name[col_name].spec, (VLStringSpec, VLBytesSpec)):
+        if isinstance(self._schema.columns_by_name[col_name].spec, (VLStringSpec, VLBytesSpec, StructSpec)):
             raise NotImplementedError(
-                f"Cannot create an index on varlen scalar column {col_name!r}: "
-                "indexing for vlstring/vlbytes columns is not supported yet."
+                f"Cannot create an index on variable-length scalar column {col_name!r}: "
+                "indexing for vlstring/vlbytes/struct columns is not supported yet."
             )
         is_persistent = self._storage.index_anchor_path(col_name) is not None
 
@@ -5756,6 +5769,8 @@ def _dtype_info_label(dtype: np.dtype | None, spec: SchemaSpec | None = None) ->
             return "vlstring"
         if isinstance(spec, VLBytesSpec):
             return "vlbytes"
+        if isinstance(spec, StructSpec):
+            return spec.display_label()
         if isinstance(spec, ListSpec):
             return spec.display_label()
         if dtype is None:
@@ -6010,7 +6025,7 @@ def _guard_varlen_scalar_expression(self, expr: str) -> None:
                 rf"(?<!\w){re.escape(col.name)}(?!\w)", expr
             ):
                 raise NotImplementedError(
-                    f"Column {col.name!r} is a vlstring/vlbytes column; "
+                    f"Column {col.name!r} is a variable-length scalar column (vlstring/vlbytes/struct); "
                     "lazy expressions are not supported yet."
                 )
 

src/blosc2/scalar_array.py

@@ -9,8 +9,8 @@
 """Internal variable-length scalar column adapter over BatchArray.
 
 This module is *not* part of the public API.  It provides row-wise scalar
-semantics (one str or bytes value per row) backed by batched msgpack storage
-via :class:`blosc2.BatchArray`.
+semantics (one str, bytes, or struct dict value per row) backed by batched
+msgpack storage via :class:`blosc2.BatchArray`.
 
 Physical layout: each chunk in the backing BatchArray stores a list of
 scalar values, e.g. ``["foo", None, "bar", "baz"]``.  Nulls are represented
@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+import os
 from bisect import bisect_right
 from collections import defaultdict
 from typing import TYPE_CHECKING, Any
@@ -35,9 +36,15 @@
 
 def _role_metadata_for_spec(spec) -> dict[str, Any]:
     """Return the fixed metadata role tag for a CTable varlen scalar backend."""
+    if spec.python_type is str:
+        py_type = "str"
+    elif spec.python_type is bytes:
+        py_type = "bytes"
+    else:
+        py_type = "struct"
     return {
         "version": 1,
-        "py_type": "str" if spec.python_type is str else "bytes",
+        "py_type": py_type,
         "nullable": bool(getattr(spec, "nullable", False)),
         "batch_rows": getattr(spec, "batch_rows", 2048),
     }
@@ -61,7 +68,12 @@ def _validate_role_metadata(backend: BatchArray, spec) -> None:
         # still identifies the logical role, so keep reopen tolerant.
         return
     role = meta[_CTABLE_VARLEN_SCALAR_META_KEY]
-    expected_py_type = "str" if spec.python_type is str else "bytes"
+    if spec.python_type is str:
+        expected_py_type = "str"
+    elif spec.python_type is bytes:
+        expected_py_type = "bytes"
+    else:
+        expected_py_type = "struct"
     if role.get("py_type") != expected_py_type:
         raise ValueError(
             f"Varlen scalar backend type mismatch: expected {expected_py_type!r}, "
@@ -81,6 +93,7 @@ def _make_backend(spec) -> BatchArray:
 
 def _make_persistent_backend(spec, urlpath: str, mode: str, *, cparams=None, dparams=None) -> BatchArray:
     """Create or open a persistent BatchArray for a varlen scalar spec."""
+    os.makedirs(os.path.dirname(urlpath), exist_ok=True)
     kwargs: dict[str, Any] = {}
     if cparams is not None:
         kwargs["cparams"] = cparams
@@ -115,20 +128,23 @@ class _ScalarVarLenArray:
     Parameters
     ----------
     spec:
-        A :class:`~blosc2.schema.VLStringSpec` or
-        :class:`~blosc2.schema.VLBytesSpec` describing this column.
+        A :class:`~blosc2.schema.VLStringSpec`,
+        :class:`~blosc2.schema.VLBytesSpec`, or
+        :class:`~blosc2.schema.StructSpec` describing this column.
     backend:
         Pre-constructed :class:`~blosc2.BatchArray`.  If ``None``, a fresh
         in-memory backend is created from *spec*.
     """
 
     def __init__(self, spec, backend: BatchArray | None = None) -> None:
-        from blosc2.schema import VLBytesSpec, VLStringSpec
+        from blosc2.schema import StructSpec, VLBytesSpec, VLStringSpec
 
-        if not isinstance(spec, (VLStringSpec, VLBytesSpec)):
-            raise TypeError(f"_ScalarVarLenArray requires a VLStringSpec or VLBytesSpec, got {type(spec)!r}")
+        if not isinstance(spec, (VLStringSpec, VLBytesSpec, StructSpec)):
+            raise TypeError(
+                f"_ScalarVarLenArray requires a VLStringSpec, VLBytesSpec, or StructSpec, got {type(spec)!r}"
+            )
         self._spec = spec
-        self._py_type: type = spec.python_type  # str or bytes
+        self._py_type: type = spec.python_type  # str, bytes, or dict
         self._nullable: bool = getattr(spec, "nullable", False)
         self._batch_rows: int = int(getattr(spec, "batch_rows", 2048) or 2048)
 
@@ -190,10 +206,13 @@ def _coerce(self, value: Any) -> Any:
             if isinstance(value, str):
                 return value
             raise TypeError(f"Expected str for vlstring column, got {type(value).__name__!r}.")
-        # bytes
-        if isinstance(value, (bytes, bytearray, memoryview)):
-            return bytes(value)
-        raise TypeError(f"Expected bytes for vlbytes column, got {type(value).__name__!r}.")
+        if self._py_type is bytes:
+            if isinstance(value, (bytes, bytearray, memoryview)):
+                return bytes(value)
+            raise TypeError(f"Expected bytes for vlbytes column, got {type(value).__name__!r}.")
+        from blosc2.list_array import _coerce_struct_item
+
+        return _coerce_struct_item(self._spec, value)
 
     def _flush_full_batches(self) -> None:
         """Flush as many full batches as possible from _pending."""

src/blosc2/schema.py

@@ -363,12 +363,17 @@ def to_metadata_dict(self) -> dict[str, Any]:
 
 
 class StructSpec(SchemaSpec):
-    """Logical schema descriptor for dict-like structured values."""
+    """Logical schema descriptor for dict-like structured values.
+
+    Top-level CTable struct columns are stored as row-wise dictionaries in a
+    batched variable-length backend.  Struct specs can also be used as
+    :func:`list` item specs for Arrow ``list<struct<...>>`` columns.
+    """
 
     python_type = dict
     dtype = None
 
-    def __init__(self, fields: dict[str, SchemaSpec]):
+    def __init__(self, fields: dict[str, SchemaSpec], *, nullable: bool = False):
         if not isinstance(fields, dict) or not fields:
             raise TypeError("StructSpec fields must be a non-empty dict")
         for name, spec in fields.items():
@@ -377,6 +382,7 @@ def __init__(self, fields: dict[str, SchemaSpec]):
             if not isinstance(spec, SchemaSpec):
                 raise TypeError("StructSpec field values must be SchemaSpec instances")
         self.fields = dict(fields)
+        self.nullable = nullable
 
     def to_pydantic_kwargs(self) -> dict[str, Any]:
         return {}
@@ -385,6 +391,7 @@ def to_metadata_dict(self) -> dict[str, Any]:
         return {
             "kind": "struct",
             "fields": [{"name": name, **spec.to_metadata_dict()} for name, spec in self.fields.items()],
+            "nullable": self.nullable,
         }
 
     def display_label(self) -> str:
@@ -399,7 +406,7 @@ def from_metadata_dict(cls, data: dict[str, Any]) -> StructSpec:
             field = dict(field)
             name = field.pop("name")
             fields[name] = spec_from_metadata_dict(field)
-        return cls(fields)
+        return cls(fields, nullable=data.get("nullable", False))
 
 
 class ListSpec(SchemaSpec):
@@ -646,9 +653,13 @@ def vlbytes(
     )
 
 
-def struct(fields: dict[str, SchemaSpec]) -> StructSpec:
-    """Build a structured schema descriptor for nested CTable values."""
-    return StructSpec(fields)
+def struct(fields: dict[str, SchemaSpec], *, nullable: bool = False) -> StructSpec:
+    """Build a structured schema descriptor for dict-like CTable values.
+
+    Top-level struct columns store one dictionary (or ``None`` when nullable)
+    per row.  Struct specs may also be nested as list item specs.
+    """
+    return StructSpec(fields, nullable=nullable)
 
 
 def list(

src/blosc2/schema_compiler.py

@@ -98,7 +98,7 @@ def compute_display_width(spec: SchemaSpec) -> int:
     """Return a reasonable terminal display width for *spec*'s column."""
     if isinstance(spec, (VLStringSpec, VLBytesSpec)):
         return 40
-    if isinstance(spec, ListSpec):
+    if isinstance(spec, (ListSpec, StructSpec)):
         return max(40, len(spec.display_label()) + 4)
     dtype = spec.dtype
     if dtype is None:
@@ -389,7 +389,7 @@ def spec_from_metadata_dict(data: dict[str, Any]) -> SchemaSpec:
         item_spec = spec_from_metadata_dict(data.pop("item"))
         return ListSpec(item_spec, **data)
     if kind == "struct":
-        return StructSpec.from_metadata_dict({"fields": data.pop("fields")})
+        return StructSpec.from_metadata_dict({"fields": data.pop("fields"), **data})
     spec_cls = _KIND_TO_SPEC.get(kind)
     if spec_cls is None:
         raise ValueError(f"Unknown column kind {kind!r}")