Support N-D array columns when constructing a Dataset from Julia data

[CarloLucibello]

Support N-D array columns when constructing a Dataset from Julia data

PR #54 added Dataset(::AbstractDict) / Dataset(::NamedTuple) but rejects array-valued
(multi-dimensional) columns loudly rather than silently transposing them. This issue
scopes out how to lift that restriction, and pairs naturally with the batch-stacking gap
(item 7 in the internal review). Filing to collect design decisions before implementing.

Python / datasets conventions

Feature types for arrays:

• Value(dtype) — scalars
• List(feature) / Sequence(feature) — variable-length (or fixed-length) nested lists
• Array2D/Array3D/Array4D/Array5D(shape, dtype) — fixed-shape tensors; shape is
  Python row-major, first axis may be None for a variable dim
• Image / Audio / Video — modality features

Verified empirically against the pinned datasets 4.x:

1. Inference — from_dict with no explicit features infers List(List(Value)) for a
   column of 2-D arrays, not Array2D. Rectangular data still works; Array2D only adds
   shape/dtype enforcement.
2. Orientation — Arrow/numpy are row-major, Julia is column-major. Our jl2numpy
   reverses axes on write ((2,3)→(3,2)), numpy2jl (DLPack) reverses them back; the
   package already guarantees numpy2jl(jl2numpy(x)) == x.
3. Read format matters:
   • python/default format → array cells come back as nested Python lists (no numpy).
     This is what our "julia" format uses today, so an array column reads as
     Vector{Vector{…}}, never a real Matrix.
   • numpy format → array cells come back as numpy ndarrays, and a range index
     stacks rows into an (N, dims…) tensor.
   • ragged columns under numpy format → object array (can't stack) → must fall back to
     per-row.
4. Per-column formatting works — set_format("numpy", columns=[…], output_all_columns=True)
   gives numpy for the array columns and leaves string/scalar columns as Python, so mixed
   schemas are fine.

The core problem

Our current "julia" format = python default format + py2jl, so array columns come back
as nested lists (Vector{Vector}), never a real N-D array — and a naive jl2numpy write
path does not line up with that nested read. Supporting N-D requires making write and
read symmetric.

The recipe that works (verified):

• Write — each element via jl2numpy (axis-reversed numpy); features inferred
  List(List) or explicit ArrayND with the reversed shape.
• Read — route array columns through numpy format, then numpy2jl. The two axis
  reversals cancel:

Dataset(d)[i]["x"] == d["x"][i]          # single row → real Julia Matrix ✓
Dataset(d)[:]["x"]  → (dims…, N) tensor  # batch stacks, observation axis LAST (MLUtils convention) ✓

This one mechanism solves N-D tensor columns and batch-stacking together.

Design choices to settle

1. Read semantics under the "julia" format

• R1 — leave as nested lists (Vector{Vector}); no true N-D.
• R2 — auto-route numeric array columns through per-column numpy format by default. Real
  N-D + free stacking, but changes today's Vector{Vector} batch semantics.
• R3 — opt-in, e.g. with_format(ds, "julia"; stack=true) or a distinct format; default
  behavior unchanged, incremental.

2. Write path (constructor)

• W1 — reject N-D (status quo).
• W2 — jl2numpy each element, infer List(List). Simplest; rectangular round-trips;
  ragged allowed but won't stack.
• W3 — jl2numpy + explicit ArrayND features (reversed shape). Enforces shape/dtype,
  reliable stacking; needs exposing Array2D…5D (ties into the Features/ClassLabel
  view work).
• W4 — a features= kwarg for user-supplied schema. Most flexible; larger surface.

3. Input shapes to accept for an array column

• Vector{Matrix} (element = one observation), and/or
• a single stacked (dims…, N) array (last axis = observations) — mirrors exactly what
  the read path returns, giving clean symmetry
  (Dataset((; x=batch))[:]["x"] == batch).

4. Ragged / non-numeric columns

• Ragged numeric columns can't numpy-stack → fall back to vector-of-arrays on read.
• Only numeric dtypes go through DLPack; strings/mixed stay on the python-default path.

Recommendation

Implement the symmetric write + read recipe (covers N-D construction and batch stacking
in one pass, sharing the axis-reversal logic), accept both input shapes, and gate the
stacking read behind an opt-in first (R3) so current "julia" semantics don't change —
promoting to default (R2) can follow once it's proven. Start with inferred List(List)
features (W2) and add explicit ArrayND (W3) alongside the planned Features view.

Context: follow-up to #54.

[CarloLucibello]

Python datasets usage examples

Reproducible snippets illustrating the conventions above (verified against the pinned
datasets 4.x; outputs shown as comments).

1. Inferred schema — a column of N-D arrays becomes nested List, not Array2D:

import numpy as np
from datasets import Dataset

ds = Dataset.from_dict({"x": [np.arange(6).reshape(2, 3),
                              np.arange(6, 12).reshape(2, 3)]})
ds.features        # {'x': List(List(Value('int64')))}
ds[0]["x"]         # default format → nested Python lists: [[0, 1, 2], [3, 4, 5]]

2. numpy format — per-row ndarray, and a range index stacks rows:

dn = ds.with_format("numpy")
dn[0]["x"].shape       # (2, 3)      → single row is a 2-D array
dn[0:2]["x"].shape     # (2, 2, 3)   → batch stacked with the observation axis FIRST

3. Explicit fixed-shape schema — Array2D/3D/4D/5D (shape is row-major):

from datasets import Features, Array2D

feats = Features({"x": Array2D(shape=(2, 3), dtype="int64")})
ds2 = Dataset.from_dict({"x": [np.arange(6).reshape(2, 3)]}, features=feats)
ds2.features       # {'x': Array2D(shape=(2, 3), dtype='int64')}

4. Ragged columns — fine per-row, but cannot stack (object array on a batch):

dr = Dataset.from_dict({"x": [[1, 2], [3, 4, 5]]}).with_format("numpy")
dr[0]["x"]             # array([1, 2])          → per-row OK
dr[0:2]["x"].dtype     # dtype('O')             → object array, no numeric stack

5. Per-column formats — array column as numpy, keep the rest as Python:

dm = Dataset.from_dict({"x": [np.arange(4).reshape(2, 2)], "name": ["a"]})
dm.set_format("numpy", columns=["x"], output_all_columns=True)
row = dm[0]
type(row["x"]).__name__   # 'ndarray'   → array column decoded as numpy
row["name"]               # 'a'          → string column stays Python

How this maps onto the Julia wrapper

• Snippet 2 is the read mechanism we'd reuse: numpy format stacks (obs axis first),
  then numpy2jl (DLPack) reverses axes → a Julia (dims…, N) tensor with the obs axis
  last (MLUtils convention). On the write side jl2numpy applies the same reversal, so
  the two cancel and Dataset(d)[i] == d["x"][i].
• Snippet 5 is what makes the julia format able to route only numeric array columns
  through numpy while leaving strings/scalars on the python-default path (options R2/R3).
• Snippets 1 vs 3 are the W2-vs-W3 write choice: infer List(List) (no new exports) or
  emit explicit Array2D…5D (needs exposing the Array* feature types).
• Snippet 4 is why ragged numeric columns must fall back to vector-of-arrays on read.