gpdatacached 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,1635 @@
1
+ """Optional pandas ``Series`` / ``DataFrame`` support for gpdatacached.
2
+
3
+ pandas is an optional dependency. Register it explicitly to make
4
+ :class:`pd.Series` and :class:`pd.DataFrame` cacheable GPDC types with live
5
+ :class:`SeriesObject` / :class:`DataFrameObject` wrappers::
6
+
7
+ from gpdatacached.extensions.pandas_support import register
8
+ register()
9
+
10
+ Scope at a glance
11
+ -----------------
12
+ * **Row indexes**: ``RangeIndex`` (default int index, stored compactly) and any
13
+ labelled index including ``pd.MultiIndex`` (stored one ``collection:list``
14
+ child per level). ``.index_names`` always returns a tuple.
15
+ * **DataFrame columns**: single-level (scalar names) or ``pd.MultiIndex``
16
+ (tuple names); the row and column axes are independent and compose (a
17
+ DataFrame may have both a row and a column MultiIndex).
18
+ * **Value dtypes**: int/float (incl. nullable ``Int64``/``Float64``), bool,
19
+ str/string, object, datetime (naive & tz-aware), category.
20
+ * **Mutation**: append-focused only — row ``extend`` (Series & DataFrame) and,
21
+ on DataFrame, column ``__setitem__`` / ``__delitem__`` (structural
22
+ add/replace/drop). No cell-level or in-place row mutation.
23
+
24
+ Usage guidance — read selectively, then work locally
25
+ -----------------------------------------------------
26
+ gpdatacached is a **cache shared across processes**, not an in-process pandas
27
+ workspace. The intended pattern is: *one writer caches a large structure; many
28
+ reader processes each pull the subset they need for the current task, then
29
+ operate on that local copy.*
30
+
31
+ * **Fetch your subset in one selective pass, then work locally.** Use the
32
+ online accessors below to read just the rows/columns you need; do all
33
+ further processing (filtering, joining, aggregating) on the returned plain
34
+ ``pandas`` object. Do not ping the cache per-row or per-cell.
35
+ * **Online reads are supported only when meaningfully cheaper than full
36
+ restore.** Anything that approaches the cost of ``.value`` (e.g.
37
+ ``.loc[slice]``, ``.loc[list]``, ``.iloc[list]``) is rejected with a
38
+ ``TypeError`` pointing at ``.value``.
39
+
40
+ Per-accessor cost (the decision guide)::
41
+
42
+ .iloc[int] O(1) a single positional value/cell
43
+ .iloc[a:b] O(b-a) a positional slice
44
+ .loc[key] O(N), selective depth-k prefix label scan; decodes
45
+ only matched rows but reads every
46
+ index level — do NOT put in a hot loop
47
+ df[col] / df[tuple] O(N), one column (DataFrame) a single column
48
+ df[level0] / df[list] O(N_cols) + O(matches·N)
49
+ column-axis selection (cheap; few cols)
50
+ .value full restore O(N) — the escape hatch
51
+
52
+ Decision rule:
53
+
54
+ * **One-shot selective read** (rows/columns wanted ≪ total) → use the online
55
+ accessor (``.iloc`` / ``.loc`` / ``__getitem__``).
56
+ * **Repeated or heavy work on the same data** — sorting, many label lookups,
57
+ complex pandas semantics, or any operation the online API rejects → call
58
+ ``.value`` once to materialize the local ``pandas`` object and work on it.
59
+
60
+ Notes
61
+ -----
62
+ * ``.loc`` matches labels with Python ``==``; ``nan`` labels therefore do not
63
+ match (``nan != nan``).
64
+ * Returned partial results retain the full stored index (no level is dropped).
65
+ """
66
+
67
+ from __future__ import annotations
68
+
69
+ import numpy as np
70
+ import pandas as pd
71
+
72
+ from ..container import GPDCContainerObject
73
+ from ..codec import GPDCContainerCodec
74
+ from ..lifecycle import policy_from_meta
75
+ from ..meta import ObjectMeta, _encode_extra_value
76
+ from ..types import (
77
+ GPDCTypeRegistry,
78
+ _decode_scalar,
79
+ decode_element,
80
+ encode_element,
81
+ encode_key,
82
+ encode_session,
83
+ )
84
+
85
+ __all__ = [
86
+ "SeriesObject",
87
+ "SeriesCodec",
88
+ "UnsupportedSeriesDtypeError",
89
+ "DataFrameObject",
90
+ "DataFrameCodec",
91
+ "UnsupportedDataFrameError",
92
+ "UnsupportedElementTypeError",
93
+ "register",
94
+ ]
95
+
96
+
97
+ _ELEMENT_SCAN_FAMILIES = {"object_or_string", "category"}
98
+
99
+
100
+ class UnsupportedElementTypeError(TypeError):
101
+ """A Series/DataFrame element is not encodable by a GPDCScalarCodec subclass.
102
+
103
+ Raised when an element is a container type (list/dict/set/container-registered
104
+ pydantic model) or an unsupported type. Series/DataFrame elements must be
105
+ scalar types encodable by a GPDCScalarCodec subclass.
106
+ """
107
+
108
+ _GUIDANCE = (
109
+ "\nSeries/DataFrame elements must be scalar types encodable by a "
110
+ "GPDCScalarCodec subclass (int/float/bool/str/datetime/date/time/tuple/None, or a "
111
+ "scalar-registered custom type). Container types are not allowed as elements."
112
+ )
113
+
114
+ def __init__(self, message: str) -> None:
115
+ super().__init__(message + self._GUIDANCE)
116
+
117
+
118
+ def _to_python_native(value):
119
+ """Normalize a numpy scalar to its Python native equivalent.
120
+
121
+ GPDCTypeRegistry.lookup_by_python_type is inconsistent on raw numpy scalars
122
+ (it rejects np.int64 / np.bool_ as unsupported). Mirrors the .tolist()
123
+ conversion the encode path already relies on. `isinstance(v, np.generic)`
124
+ is True for every numpy scalar dtype and `.item()` maps each to its native.
125
+ """
126
+ if isinstance(value, np.generic):
127
+ return value.item()
128
+ return value
129
+
130
+
131
+ def _is_na_sentinel(value) -> bool:
132
+ """True for missing-value sentinels that are always allowed as elements.
133
+
134
+ None / pandas NA / pandas NaT are valid NA markers, consistent with
135
+ _python_values/_python_labels which normalize these to None on the encode
136
+ path. Uses identity checks rather than pd.isna: pd.isna on a list/dict
137
+ returns a bool-ambiguous array, and those containers must fall through to
138
+ the codec lookup below so they get rejected.
139
+ """
140
+ return value is None or value is pd.NA or value is pd.NaT
141
+
142
+
143
+ def _assert_scalar_encodable(values, context: str) -> None:
144
+ """Raise UnsupportedElementTypeError if any value is not encodable by a
145
+ GPDCScalarCodec subclass.
146
+
147
+ `context` is a human-readable description for the error message (e.g.
148
+ "Series element", "DataFrame column 'o' element", "index label",
149
+ "category value"). NA sentinels (None / pandas NA / NaT) are always allowed.
150
+ """
151
+ from ..codec import GPDCScalarCodec
152
+ from ..types import GPDCTypeRegistry
153
+
154
+ for raw_value in values:
155
+ if _is_na_sentinel(raw_value):
156
+ continue
157
+ value = _to_python_native(raw_value)
158
+ if _is_na_sentinel(value):
159
+ continue
160
+ try:
161
+ entry = GPDCTypeRegistry.lookup_by_python_type(value)
162
+ except TypeError:
163
+ raise UnsupportedElementTypeError(
164
+ f"{context} of type {type(raw_value).__name__!r} is not a supported scalar type."
165
+ )
166
+ if not issubclass(entry.codec_class, GPDCScalarCodec):
167
+ raise UnsupportedElementTypeError(
168
+ f"{context} of type {type(raw_value).__name__!r} is a container type; only scalar elements are allowed."
169
+ )
170
+
171
+
172
+ class UnsupportedSeriesDtypeError(TypeError):
173
+ """Raised when a pandas dtype cannot be stored element-wise by SeriesObject.
174
+
175
+ Guides the user to register a custom codec for the dtype.
176
+ """
177
+
178
+ def __init__(self, dtype: str, *, is_index: bool) -> None:
179
+ target = "index dtype" if is_index else "Series dtype"
180
+ super().__init__(
181
+ f"{target} {dtype!r} is not supported by SeriesObject.\n"
182
+ "Supported dtypes: int64/Int64, float64/Float64, bool/boolean, "
183
+ "str/string, object, datetime64 (naive & tz-aware), category.\n"
184
+ "To cache this dtype, write a custom codec and register it via "
185
+ "GPDCTypeRegistry.register(...)."
186
+ )
187
+
188
+
189
+ class UnsupportedDataFrameError(TypeError):
190
+ """Raised when a DataFrame cannot be stored element-wise by DataFrameObject.
191
+
192
+ The message disambiguates the cause (per-column dtype, row index, or
193
+ duplicate column names) and guides the user to register a custom codec.
194
+ """
195
+
196
+ _GUIDANCE = (
197
+ "\nSupported dtypes: int64/Int64, float64/Float64, bool/boolean, "
198
+ "str/string, object, datetime64 (naive & tz-aware), category.\n"
199
+ "To cache this, write a custom codec and register it via "
200
+ "GPDCTypeRegistry.register(...)."
201
+ )
202
+
203
+ def __init__(self, message: str) -> None:
204
+ super().__init__(message + self._GUIDANCE)
205
+
206
+
207
+ def _value_dtype_family(dtype) -> str | None:
208
+ """Classify a pandas value dtype into a reconstruction family, or None if unsupported."""
209
+ if isinstance(dtype, pd.CategoricalDtype):
210
+ return "category"
211
+ if isinstance(dtype, pd.SparseDtype):
212
+ return None
213
+ if pd.api.types.is_datetime64_any_dtype(dtype):
214
+ return "datetime"
215
+ if pd.api.types.is_bool_dtype(dtype):
216
+ return "bool"
217
+ if pd.api.types.is_integer_dtype(dtype):
218
+ return "integer"
219
+ if pd.api.types.is_float_dtype(dtype):
220
+ return "float"
221
+ string_dtype = getattr(pd, "StringDtype", None)
222
+ if (string_dtype is not None and isinstance(dtype, string_dtype)) or dtype.kind in ("O", "U", "S"):
223
+ return "object_or_string"
224
+ return None
225
+
226
+
227
+ def _index_kind(index) -> str | None:
228
+ """Classify a pandas index: 'range', 'levels', or None if unsupported.
229
+
230
+ 'levels' covers every labelled index incl. MultiIndex (L>=1).
231
+ """
232
+ if isinstance(index, pd.RangeIndex):
233
+ return "range"
234
+ if isinstance(index, pd.MultiIndex):
235
+ for i in range(index.nlevels):
236
+ if _value_dtype_family(index.get_level_values(i).dtype) is None:
237
+ return None
238
+ return "levels"
239
+ return "levels" if _value_dtype_family(index.dtype) is not None else None
240
+
241
+
242
+ def _validate_dataframe(value: pd.DataFrame) -> tuple[list, str]:
243
+ """Validate a DataFrame for storage.
244
+
245
+ Returns ``(column_names, index_kind)``. Raises ``UnsupportedDataFrameError``
246
+ on duplicate column names, an unsupported per-column dtype, or an
247
+ unsupported row-index dtype. ``index_kind`` is ``_index_kind``'s result
248
+ (``range`` or ``levels``; single-level labelled and MultiIndex both yield
249
+ ``levels``).
250
+ """
251
+ names = list(value.columns)
252
+ seen: set = set()
253
+ for n in names:
254
+ if n in seen:
255
+ raise UnsupportedDataFrameError(f"Duplicate column names are not supported by DataFrameObject: {names!r}.")
256
+ seen.add(n)
257
+ for col in names:
258
+ if _value_dtype_family(value[col].dtype) is None:
259
+ raise UnsupportedDataFrameError(
260
+ f"Column {col!r} has dtype {value[col].dtype} which is not supported by DataFrameObject."
261
+ )
262
+ kind = _index_kind(value.index)
263
+ if kind is None:
264
+ raise UnsupportedDataFrameError(
265
+ f"DataFrame row index dtype {value.index.dtype} is not supported by DataFrameObject."
266
+ )
267
+ return names, kind
268
+
269
+
270
+ def _python_values(series: pd.Series) -> list:
271
+ """Materialize Series values as Python natives with NA -> None.
272
+
273
+ Numpy scalars (e.g. np.int64/np.bool_ inside an object-dtype Series, which
274
+ .tolist() leaves untouched) are normalized to Python natives so the element
275
+ codec lookup succeeds. Container values (list/dict/set) are not numpy
276
+ scalars and pass through unchanged.
277
+ """
278
+ is_na = pd.isna(series)
279
+ return [None if na else _to_python_native(x) for x, na in zip(series.tolist(), list(is_na))]
280
+
281
+
282
+ def _python_labels(index) -> list:
283
+ """Materialize Index labels as Python natives with NA -> None.
284
+
285
+ Numpy scalars are normalized to Python natives (see _python_values).
286
+ """
287
+ is_na = pd.isna(index)
288
+ return [None if na else _to_python_native(x) for x, na in zip(index.tolist(), list(is_na))]
289
+
290
+
291
+ def _level_labels_and_dtype(index, level_i: int) -> tuple[list, str]:
292
+ """Return (python-native labels, dtype-string) for one MultiIndex level."""
293
+ level_vals = index.get_level_values(level_i)
294
+ return _python_labels(level_vals), str(level_vals.dtype)
295
+
296
+
297
+ def _infer_promote_dtype(labels: list) -> str:
298
+ """Honest combined dtype for a range→levels promote: int64 if all appended
299
+ labels cast to int64, else object. (RangeIndex values are always int, so the
300
+ combined level dtype is determined by whether the appended labels are int64.)"""
301
+ try:
302
+ pd.array(labels, dtype="int64")
303
+ return "int64"
304
+ except (TypeError, ValueError):
305
+ return "object"
306
+
307
+
308
+ def _build_levels_index(fields: dict, namespace, prefix: str = "") -> "pd.Index | pd.MultiIndex":
309
+ """Rebuild a MultiIndex (L>=1) from per-level ref fields."""
310
+ names = _decode_scalar(fields[f"{prefix}index_names"])
311
+ dtypes = _decode_scalar(fields[f"{prefix}index_dtypes"])
312
+ L = len(names)
313
+ arrays = []
314
+ for i in range(L):
315
+ labels = decode_element(fields[f"{prefix}index_l{i}"], namespace).value
316
+ try:
317
+ arrays.append(pd.array(labels, dtype=dtypes[i]))
318
+ except (TypeError, ValueError):
319
+ # Create-time validation already enforced dtype<->labels consistency,
320
+ # so reaching here means corruption or a future new dtype; on a cache
321
+ # read path it is safer to infer than to raise.
322
+ arrays.append(pd.array(labels))
323
+ if L == 1:
324
+ return pd.Index(arrays[0], name=names[0])
325
+ return pd.MultiIndex.from_arrays(arrays, names=names)
326
+
327
+
328
+ def _build_columns_index(fields: dict, column_names) -> "pd.Index | pd.MultiIndex | list":
329
+ """Build the DataFrame column index. MultiIndex (with level names) when
330
+ ``~columns_names`` is present, else the plain column-name list."""
331
+ if "~columns_names" in fields:
332
+ col_names = _decode_scalar(fields["~columns_names"])
333
+ return pd.MultiIndex.from_tuples(list(column_names), names=col_names)
334
+ return list(column_names)
335
+
336
+
337
+ def _reconstruct_series(
338
+ values: list,
339
+ index_obj,
340
+ dtype_str: str,
341
+ name,
342
+ namespace,
343
+ cats_ref: str | None = None,
344
+ ordered=None,
345
+ ) -> pd.Series:
346
+ """Build a pd.Series from decoded values + index + dtype metadata.
347
+
348
+ Shared by SeriesCodec.read_raw, DataFrameCodec.read_raw, and DataFrameObject
349
+ column reads. ``cats_ref`` is the encoded ``ref:`` string pointing at a
350
+ categories ListObject (category dtype only).
351
+ """
352
+ family = _value_dtype_family(pd.api.types.pandas_dtype(dtype_str))
353
+ if family == "category":
354
+ cats = decode_element(cats_ref, namespace).value
355
+ cat = pd.Categorical(values, categories=cats, ordered=ordered)
356
+ return pd.Series(cat, index=index_obj, name=name)
357
+ if family == "datetime":
358
+ return pd.Series(values, index=index_obj, name=name).astype(dtype_str)
359
+ return pd.Series(values, index=index_obj, dtype=dtype_str, name=name)
360
+
361
+
362
+ def _reconstruct_column(values, index_obj, fields, col, namespace) -> pd.Series:
363
+ """Reconstruct one DataFrame column from decoded values + raw-hash metadata.
364
+
365
+ Shared by DataFrameCodec.read_raw, DataFrameObject.__getitem__,
366
+ _positional_row_slice, and _DfLocIndexer. Reads the per-column
367
+ ``~dt:``/``~cat:``/``~ord:`` metadata from ``fields`` and delegates to
368
+ _reconstruct_series. ``values`` is the already-decoded list of cell values
369
+ (caller-supplied, since each read path fetches values differently).
370
+ """
371
+ col_key = encode_key(col)
372
+ dtype_str = _decode_scalar(fields[f"~dt:{col_key}"])
373
+ family = _value_dtype_family(pd.api.types.pandas_dtype(dtype_str))
374
+ cats_ref = fields.get(f"~cat:{col_key}") if family == "category" else None
375
+ ordered = _decode_scalar(fields[f"~ord:{col_key}"]) if family == "category" else None
376
+ return _reconstruct_series(values, index_obj, dtype_str, col, namespace, cats_ref, ordered)
377
+
378
+
379
+ def _reconstruct(values: list, index_obj, fields: dict, namespace) -> pd.Series:
380
+ """Series-specific wrapper over _reconstruct_series (reads Series raw-hash fields)."""
381
+ return _reconstruct_series(
382
+ values,
383
+ index_obj,
384
+ _decode_scalar(fields["dtype"]),
385
+ _decode_scalar(fields["name"]),
386
+ namespace,
387
+ cats_ref=fields.get("categories"),
388
+ ordered=_decode_scalar(fields["ordered"]) if "ordered" in fields else None,
389
+ )
390
+
391
+
392
+ def _build_index(fields: dict, namespace, prefix: str = "") -> "pd.Index":
393
+ """Build a RangeIndex from raw-hash range fields.
394
+
395
+ Handles the ``range`` index kind only. ``levels`` indexes (single-level
396
+ labelled or MultiIndex) are rebuilt via ``_build_levels_index`` directly
397
+ at their call sites (``DataFrameCodec.read_raw``, ``SeriesCodec.read_raw``,
398
+ ``DataFrameObject._build_full_row_index``). ``prefix`` is "" for Series
399
+ raw-hash field names (``index_kind``) and "~" for DataFrame raw-hash field
400
+ names (``~index_kind``).
401
+ """
402
+ index_names = _decode_scalar(fields[f"{prefix}index_names"])
403
+ name = index_names[0] if index_names else None
404
+ start = _decode_scalar(fields[f"{prefix}range_start"])
405
+ stop = _decode_scalar(fields[f"{prefix}range_stop"])
406
+ step = _decode_scalar(fields[f"{prefix}range_step"])
407
+ return pd.RangeIndex(start, stop, step, name=name)
408
+
409
+
410
+ def _normalize_extend_input(other) -> tuple[list, list, int]:
411
+ """Return (labels, data_values, nlevels) from extend input.
412
+
413
+ ``labels`` is a list of tuple labels (one per appended row); each tuple's
414
+ length is the appended index's level count. ``nlevels`` is that level count.
415
+ Accepts a pd.Series, a list of pd.Series, or an iterable of (index, value)
416
+ tuples. The level-count check is the caller's responsibility.
417
+ """
418
+ if isinstance(other, pd.Series):
419
+ idx = other.index
420
+ nlevels = getattr(idx, "nlevels", 1)
421
+ if isinstance(idx, pd.MultiIndex):
422
+ labels = [tuple(r) for r in idx]
423
+ else:
424
+ labels = [(r,) for r in idx]
425
+ return labels, _python_values(other), nlevels
426
+ if isinstance(other, (list, tuple)) and len(other) > 0 and isinstance(other[0], pd.Series):
427
+ idx0 = other[0].index
428
+ nlevels = getattr(idx0, "nlevels", 1)
429
+ labels: list = []
430
+ vals: list = []
431
+ for s in other:
432
+ si = s.index
433
+ if isinstance(si, pd.MultiIndex):
434
+ labels.extend(tuple(r) for r in si)
435
+ else:
436
+ labels.extend((r,) for r in si)
437
+ vals.extend(_python_values(s))
438
+ return labels, vals, nlevels
439
+ pairs = list(other)
440
+ bad = [p for p in pairs if not (isinstance(p, tuple) and len(p) == 2)]
441
+ if bad:
442
+ raise TypeError("extend input must be a Series, a list of Series, or (index, value) tuples")
443
+ labels = []
444
+ vals = []
445
+ nlevels = None
446
+ for lbl, v in pairs:
447
+ if not isinstance(lbl, tuple):
448
+ lbl = (lbl,)
449
+ if nlevels is None:
450
+ nlevels = len(lbl)
451
+ elif len(lbl) != nlevels:
452
+ raise ValueError("inconsistent index label depths in extend input")
453
+ labels.append(lbl)
454
+ vals.append(v)
455
+ if nlevels is None:
456
+ nlevels = 1
457
+ return labels, vals, nlevels
458
+
459
+
460
+ def _is_range_indexed_series(other) -> bool:
461
+ """True when `other` is a Series with a RangeIndex (so a range Series can stay a range)."""
462
+ return isinstance(other, pd.Series) and isinstance(other.index, pd.RangeIndex)
463
+
464
+
465
+ def _normalize_df_extend_input(other, columns: list, expect_levels: int | None = None) -> tuple[dict, list[list], int]:
466
+ """Normalize extend input for a DataFrame.
467
+
468
+ Returns ``(per_col, index_labels_per_level, nlevels)`` where ``per_col`` maps
469
+ each column to its python-native values, ``index_labels_per_level`` is a list
470
+ of ``nlevels`` lists (one per index level, each the appended labels for that
471
+ level), and ``nlevels`` is the appended index's level count.
472
+
473
+ For dict / list-of-dict inputs (which carry no index), positional labels
474
+ ``0..n-1`` are synthesized (range-compatible). A genuine MultiIndex-row
475
+ DataFrame (``expect_levels >= 2``) rejects dict/list inputs with a TypeError
476
+ pointing at ``.value`` — those forms synthesise positional labels a
477
+ MultiIndex cannot supply. Raises ValueError on column-set mismatch.
478
+ """
479
+ if isinstance(other, pd.DataFrame):
480
+ other_cols = list(other.columns)
481
+ if set(other_cols) != set(columns):
482
+ raise ValueError(f"extend input columns {other_cols!r} do not match {columns!r}")
483
+ per_col = {c: _python_values(other[c]) for c in columns}
484
+ idx = other.index
485
+ nlevels = getattr(idx, "nlevels", 1)
486
+ if isinstance(idx, pd.MultiIndex):
487
+ index_labels_per_level = [_python_labels(idx.get_level_values(i)) for i in range(nlevels)]
488
+ else:
489
+ index_labels_per_level = [_python_labels(idx)]
490
+ nlevels = 1
491
+ return per_col, index_labels_per_level, nlevels
492
+ multi_rejected = expect_levels is not None and expect_levels > 1
493
+ if isinstance(other, dict):
494
+ if multi_rejected:
495
+ raise TypeError(
496
+ "extend a MultiIndex-row DataFrame with a DataFrame only; "
497
+ "for other inputs call .value for the full pandas object and operate on it locally."
498
+ )
499
+ if set(other.keys()) != set(columns):
500
+ raise ValueError(f"extend input columns {list(other.keys())!r} do not match {columns!r}")
501
+ per_col = {}
502
+ for c in columns:
503
+ vals = other[c]
504
+ per_col[c] = _python_values(vals) if isinstance(vals, pd.Series) else list(vals)
505
+ lengths = {len(v) for v in per_col.values()}
506
+ if len(lengths) > 1:
507
+ raise ValueError(
508
+ f"extend dict-of-lists columns have mismatched lengths: { {c: len(per_col[c]) for c in columns} }"
509
+ )
510
+ n = next(iter(lengths)) if lengths else 0
511
+ return per_col, [list(range(n))], 1
512
+ if isinstance(other, (list, tuple)):
513
+ if multi_rejected:
514
+ raise TypeError(
515
+ "extend a MultiIndex-row DataFrame with a DataFrame only; "
516
+ "for other inputs call .value for the full pandas object and operate on it locally."
517
+ )
518
+ rows = list(other)
519
+ bad = [r for r in rows if not isinstance(r, dict)]
520
+ if bad:
521
+ raise TypeError("extend input must be a DataFrame, dict, or list of dict rows")
522
+ per_col = {c: [] for c in columns}
523
+ for r in rows:
524
+ if set(r.keys()) != set(columns):
525
+ raise ValueError(f"extend row columns {list(r.keys())!r} do not match {columns!r}")
526
+ for c in columns:
527
+ per_col[c].append(r[c])
528
+ return per_col, [list(range(len(rows)))], 1
529
+ raise TypeError("extend input must be a DataFrame, dict, or list of dict rows")
530
+
531
+
532
+ class DataFrameCodec(GPDCContainerCodec):
533
+ """Codec for pandas DataFrame values stored as a Redis hash of column refs."""
534
+
535
+ @classmethod
536
+ def create_raw(cls, namespace, raw_key, value: pd.DataFrame, owner_name, cache_mode, ttl) -> None:
537
+ from ..lifecycle import policy_from_defaults
538
+
539
+ column_names, kind = _validate_dataframe(value)
540
+ policy = policy_from_defaults(cache_mode, ttl)
541
+ index = value.index
542
+
543
+ # Bind all lists to locals BEFORE entering encode_session so they stay
544
+ # alive for the whole session. Passing temporaries inline into
545
+ # encode_element lets CPython reuse a freed id() for the next list,
546
+ # which collides with the encode memo (keyed by id) and aliases
547
+ # distinct content (e.g. an index list silently reuses a column's ref).
548
+ col_values = {col: _python_values(value[col]) for col in column_names}
549
+ col_cats = {
550
+ col: value[col].cat.categories.tolist()
551
+ for col in column_names
552
+ if _value_dtype_family(value[col].dtype) == "category"
553
+ }
554
+
555
+ # Validate elements are scalar-encodable BEFORE encoding (no Redis write yet).
556
+ for col in column_names:
557
+ if _value_dtype_family(value[col].dtype) in _ELEMENT_SCAN_FAMILIES:
558
+ _assert_scalar_encodable(col_values[col], f"DataFrame column {col!r} element")
559
+ if col in col_cats:
560
+ _assert_scalar_encodable(col_cats[col], f"DataFrame column {col!r} category value")
561
+
562
+ fields: dict[str, str] = {}
563
+ with encode_session(namespace):
564
+ for col in column_names:
565
+ col_key = encode_key(col)
566
+ fields[col_key] = encode_element(col_values[col], namespace, owner_name, policy)
567
+ if col in col_cats:
568
+ fields[f"~cat:{col_key}"] = encode_element(col_cats[col], namespace, owner_name, policy)
569
+ if kind == "levels":
570
+ # Bind all level label lists to locals before encoding (id-aliasing guard).
571
+ level_labels = [_level_labels_and_dtype(index, i) for i in range(index.nlevels)]
572
+ for i, (labels, _dt) in enumerate(level_labels):
573
+ _assert_scalar_encodable(labels, f"DataFrame index level {i} label")
574
+ fields[f"~index_l{i}"] = encode_element(labels, namespace, owner_name, policy)
575
+
576
+ for col in column_names:
577
+ col_key = encode_key(col)
578
+ col_series = value[col]
579
+ fields[f"~dt:{col_key}"] = _encode_extra_value(str(col_series.dtype))
580
+ if _value_dtype_family(col_series.dtype) == "category":
581
+ fields[f"~ord:{col_key}"] = _encode_extra_value(bool(col_series.cat.ordered))
582
+
583
+ fields["~columns"] = _encode_extra_value(tuple(column_names))
584
+ # Column level-names stored only for MultiIndex columns; a single-level
585
+ # column index .name is out of scope (see Plan 2 Task 1 / spec §3.2).
586
+ if isinstance(value.columns, pd.MultiIndex):
587
+ fields["~columns_names"] = _encode_extra_value(tuple(value.columns.names))
588
+ fields["~index_kind"] = _encode_extra_value(kind)
589
+ fields["~index_names"] = _encode_extra_value(tuple(index.names))
590
+ if kind == "levels":
591
+ fields["~index_dtypes"] = _encode_extra_value(tuple(dt for _, dt in level_labels))
592
+ else:
593
+ fields["~range_start"] = _encode_extra_value(int(index.start))
594
+ fields["~range_stop"] = _encode_extra_value(int(index.stop))
595
+ fields["~range_step"] = _encode_extra_value(int(index.step))
596
+
597
+ namespace._redis().hset(raw_key, mapping=fields)
598
+
599
+ @classmethod
600
+ def read_raw(cls, namespace, raw_key) -> pd.DataFrame:
601
+ fields = namespace._redis().hgetall(raw_key)
602
+ if not fields:
603
+ return pd.DataFrame()
604
+ column_names = _decode_scalar(fields["~columns"])
605
+ index_kind = _decode_scalar(fields["~index_kind"])
606
+ if index_kind == "range":
607
+ index_obj = _build_index(fields, namespace, prefix="~")
608
+ else: # levels
609
+ index_obj = _build_levels_index(fields, namespace, prefix="~")
610
+ data_dict = {}
611
+ for col in column_names:
612
+ col_key = encode_key(col)
613
+ values = decode_element(fields[col_key], namespace).deepcopy()
614
+ data_dict[col] = _reconstruct_column(values, index_obj, fields, col, namespace)
615
+ columns = _build_columns_index(fields, column_names)
616
+ return pd.DataFrame(data_dict, columns=columns, index=index_obj)
617
+
618
+ @classmethod
619
+ def delete_raw(cls, namespace, raw_key) -> None:
620
+ namespace._redis().delete(raw_key)
621
+
622
+
623
+ class _DfILocIndexer:
624
+ """Positional row indexer for DataFrameObject (read-only)."""
625
+
626
+ def __init__(self, owner: "DataFrameObject") -> None:
627
+ self._owner = owner
628
+
629
+ def __getitem__(self, key):
630
+ owner = self._owner
631
+ if isinstance(key, (int, np.integer)):
632
+ fields = owner._raw_fields()
633
+ columns = owner._column_names()
634
+ values = []
635
+ for col in columns:
636
+ col_list = owner._get_column_list(col)
637
+ if col_list is None:
638
+ raise IndexError("DataFrame has no rows")
639
+ values.append(col_list[int(key)])
640
+ index_obj = owner._build_full_row_index()[int(key)]
641
+ return pd.Series(values, index=_build_columns_index(fields, columns), name=index_obj)
642
+ if isinstance(key, slice):
643
+ return owner._positional_row_slice(key)
644
+ raise TypeError(
645
+ ".iloc list (positional fancy) indexing is not supported on this cache object; "
646
+ "use .iloc with an int or slice, or call .value for the full pandas object."
647
+ )
648
+
649
+
650
+ class _DfLocIndexer:
651
+ """Label-based row indexer for DataFrameObject (read-only).
652
+
653
+ Supports scalar / tuple depth-k prefix selection on a MultiIndex (or
654
+ single-level labelled) row index. Rejects slice and list with a TypeError
655
+ pointing to ``.value`` (the cost-bar from the design spec). Full-depth
656
+ (``k == L``) with exactly one matching row returns that row as a
657
+ ``pd.Series``; partial depth or multiple matches return a ``pd.DataFrame``.
658
+ """
659
+
660
+ def __init__(self, owner: "DataFrameObject") -> None:
661
+ self._owner = owner
662
+
663
+ def __getitem__(self, key):
664
+ owner = self._owner
665
+ if isinstance(key, slice):
666
+ raise TypeError(
667
+ ".loc slices are not supported on this cache object; "
668
+ "call .value for the full pandas object and operate on it locally, "
669
+ "or use .iloc for positional slicing."
670
+ )
671
+ if isinstance(key, list):
672
+ raise TypeError(
673
+ ".loc label-list selection is not supported on this cache object; "
674
+ "call .value for the full pandas object and operate on it locally."
675
+ )
676
+ fields = owner._raw_fields()
677
+ if _decode_scalar(fields["~index_kind"]) == "range":
678
+ raise TypeError(
679
+ ".loc on a RangeIndex is not supported on this cache object; "
680
+ "use .iloc for positional slicing, or call .value for the full pandas object."
681
+ )
682
+ names = _decode_scalar(fields["~index_names"])
683
+ L = len(names)
684
+ key_tuple = key if isinstance(key, tuple) else (key,)
685
+ k = len(key_tuple)
686
+ if k > L:
687
+ raise KeyError(f"key depth {k} exceeds index levels {L}")
688
+ # Fetch all L level values once (ListObject.value is uncached -> each call is a Redis lrange).
689
+ all_level_values = [owner._get_level_list(i).value for i in range(L)]
690
+ # NA-label matching follows Python == semantics; nan labels won't match (pre-existing behaviour).
691
+ positions = [
692
+ pos for pos in range(len(owner)) if all(all_level_values[i][pos] == key_tuple[i] for i in range(k))
693
+ ]
694
+ if not positions:
695
+ raise KeyError(key)
696
+ columns = owner._column_names()
697
+ # Build the result row index from the SAME all_level_values (matched positions, all L levels).
698
+ arrays = [[lvl[p] for p in positions] for lvl in all_level_values]
699
+ if L == 1:
700
+ sliced_index = pd.Index(arrays[0], name=names[0])
701
+ else:
702
+ sliced_index = pd.MultiIndex.from_arrays(arrays, names=names)
703
+ data_dict = {}
704
+ for col in columns:
705
+ col_list = owner._get_column_list(col)
706
+ values = [col_list[p] for p in positions]
707
+ data_dict[col] = _reconstruct_column(values, sliced_index, fields, col, owner._namespace)
708
+ result = pd.DataFrame(data_dict, columns=_build_columns_index(fields, columns))
709
+ if k == L and len(positions) == 1:
710
+ return result.iloc[0] # single full-depth match -> row as pd.Series
711
+ return result
712
+
713
+
714
+ class DataFrameObject(GPDCContainerObject):
715
+ """Redis-backed pandas DataFrame object.
716
+
717
+ raw_key is a Redis HASH: column value-refs under ``encode_key(col)``, the
718
+ per-level index labels under ``~index_l{i}`` (with ``~index_kind`` /
719
+ ``~index_names`` / ``~index_dtypes`` scalar metadata), and other scalar
720
+ metadata under ``~``-prefixed fields. See the design spec for the full
721
+ layout.
722
+ """
723
+
724
+ TYPE = "collection:pandas_dataframe"
725
+
726
+ def __init__(self, name, namespace, *, canonical_name=None) -> None:
727
+ super().__init__(name, namespace, canonical_name=canonical_name)
728
+ self._column_lists_cache: dict = {}
729
+ # Per-level index ListObject cache (L>=1; the extend pipeline and the
730
+ # depth-k read paths both go through _get_level_list). Cleared on
731
+ # mutation so a stale ref is never read.
732
+ self._level_lists_cache: dict = {}
733
+
734
+ @property
735
+ def type(self) -> str:
736
+ return self.TYPE
737
+
738
+ @property
739
+ def value(self) -> pd.DataFrame:
740
+ """Full materialization as a plain :class:`pandas.DataFrame`.
741
+
742
+ Use this for repeated or heavy work on the same data (sorting, many
743
+ lookups, complex pandas semantics, or operations the online API
744
+ rejects). For a one-shot selective read, prefer :attr:`iloc` / :attr:`loc`
745
+ / ``df[col]``. See the module docstring for the full usage guidance.
746
+ """
747
+ return DataFrameCodec.read_raw(self._namespace, self.raw_key)
748
+
749
+ @value.setter
750
+ def value(self, new_value) -> None:
751
+ self._reject_value_replacement()
752
+
753
+ # ── internal helpers ─────────────────────────────────────────────
754
+
755
+ def _raw_fields(self) -> dict:
756
+ return self._redis().hgetall(self.raw_key)
757
+
758
+ def _column_names(self) -> tuple:
759
+ return _decode_scalar(self._redis().hget(self.raw_key, "~columns"))
760
+
761
+ def _columns_names_or_none(self):
762
+ raw = self._redis().hget(self.raw_key, "~columns_names")
763
+ return _decode_scalar(raw) if raw else None
764
+
765
+ def _select_one_series(self, col, fields) -> pd.Series:
766
+ col_list_obj = self._get_column_list(col)
767
+ values = col_list_obj.deepcopy() if col_list_obj is not None else []
768
+ index_obj = self._build_full_row_index()
769
+ return _reconstruct_column(values, index_obj, fields, col, self._namespace)
770
+
771
+ def _select_columns(self, col_list) -> pd.DataFrame:
772
+ fields = self._raw_fields()
773
+ index_obj = self._build_full_row_index()
774
+ data = {}
775
+ for c in col_list:
776
+ col_list_obj = self._get_column_list(c)
777
+ values = col_list_obj.deepcopy() if col_list_obj is not None else []
778
+ data[c] = _reconstruct_column(values, index_obj, fields, c, self._namespace)
779
+ columns = _build_columns_index(fields, col_list)
780
+ return pd.DataFrame(data, columns=columns, index=index_obj)
781
+
782
+ def _index_kind(self) -> str:
783
+ return _decode_scalar(self._redis().hget(self.raw_key, "~index_kind"))
784
+
785
+ def _get_level_list(self, level_i: int):
786
+ """Return the per-level index ListObject wrapper, or None when the level is absent."""
787
+ if level_i not in self._level_lists_cache:
788
+ ref = self._redis().hget(self.raw_key, f"~index_l{level_i}")
789
+ self._level_lists_cache[level_i] = decode_element(ref, self._namespace) if ref else None
790
+ return self._level_lists_cache[level_i]
791
+
792
+ def _build_full_row_index(self):
793
+ """Rebuild the full row pd.Index / pd.MultiIndex from raw-hash fields."""
794
+ fields = self._raw_fields()
795
+ kind = _decode_scalar(fields["~index_kind"])
796
+ if kind == "range":
797
+ return _build_index(fields, self._namespace, prefix="~")
798
+ return _build_levels_index(fields, self._namespace, prefix="~")
799
+
800
+ def _get_column_list(self, col):
801
+ """Return a column's value ListObject wrapper (cached)."""
802
+ if col not in self._column_lists_cache:
803
+ ref = self._redis().hget(self.raw_key, encode_key(col))
804
+ self._column_lists_cache[col] = decode_element(ref, self._namespace) if ref else None
805
+ return self._column_lists_cache[col]
806
+
807
+ def _invalidate_caches(self) -> None:
808
+ self._column_lists_cache = {}
809
+ self._level_lists_cache = {}
810
+
811
+ def _promote_df_range_to_labels(self, fields: dict, owner_name, policy, level_dtype="int64") -> None:
812
+ """Materialize a RangeIndex into an L=1 ``levels`` index in place.
813
+
814
+ Emits ``~index_l0`` / ``~index_kind=levels`` / ``~index_dtypes`` (matching
815
+ the new levels storage layout; ``read_raw`` no longer reads ``labels``).
816
+ ``level_dtype`` lets the extend caller pass an honest combined dtype so
817
+ subsequent extends cast against a truthful dtype instead of a stale
818
+ int64 (mirrors the Series Task-5 fix). Defaults to ``int64`` so today's
819
+ single-level extend caller behaves as before.
820
+
821
+ Note: commits outside the extend pipeline. If the subsequent append
822
+ pipeline fails, the index stays promoted (range→levels) with no rows
823
+ added — semantically valid but representationally changed. Mirrors
824
+ SeriesObject.extend's promote path.
825
+ """
826
+ start = _decode_scalar(fields["~range_start"])
827
+ stop = _decode_scalar(fields["~range_stop"])
828
+ step = _decode_scalar(fields["~range_step"])
829
+ labels = list(pd.RangeIndex(start, stop, step))
830
+ with encode_session(self._namespace):
831
+ ref = encode_element(labels, self._namespace, owner_name, policy)
832
+ r = self._redis()
833
+ mapping = {
834
+ "~index_l0": ref,
835
+ "~index_kind": _encode_extra_value("levels"),
836
+ "~index_dtypes": _encode_extra_value((level_dtype,)),
837
+ }
838
+ r.hdel(self.raw_key, "~range_start", "~range_stop", "~range_step")
839
+ r.hset(self.raw_key, mapping=mapping)
840
+ self._level_lists_cache = {}
841
+
842
+ def extend(self, other) -> None:
843
+ """Append rows in place. Accepts a DataFrame, a dict of lists, or a list of dicts.
844
+
845
+ Mutates the cached object in Redis and returns None. For a genuine
846
+ MultiIndex-row DataFrame (L>=2), only a DataFrame with a matching
847
+ MultiIndex is accepted (dict/list inputs synthesise positional labels a
848
+ MultiIndex cannot supply -> TypeError; use ``.value``).
849
+ """
850
+ from ..lifecycle import policy_from_meta, policy_from_defaults
851
+
852
+ r = self._redis()
853
+ fields = self._raw_fields()
854
+ columns = list(_decode_scalar(fields["~columns"]))
855
+ index_kind = _decode_scalar(fields["~index_kind"])
856
+ names = _decode_scalar(fields["~index_names"])
857
+
858
+ expect_levels = len(names) if index_kind == "levels" else None
859
+ per_col, index_labels_per_level, got_levels = _normalize_df_extend_input(
860
+ other, columns, expect_levels=expect_levels
861
+ )
862
+ n = len(index_labels_per_level[0]) if index_labels_per_level else 0
863
+
864
+ # Range-compatible iff the appended input carries positional labels
865
+ # (a DataFrame with a RangeIndex, or a dict / list-of-dicts which
866
+ # synthesize 0..n-1). Mirrors Series's _is_range_indexed_series.
867
+ if isinstance(other, pd.DataFrame):
868
+ appended_is_range = isinstance(other.index, pd.RangeIndex)
869
+ else:
870
+ appended_is_range = True
871
+
872
+ meta = ObjectMeta.load(r, self._canonical_redis_key())
873
+ policy = policy_from_meta(meta) if meta else policy_from_defaults("permanent", 0)
874
+ owner_name = self._data_owner_name()
875
+
876
+ # Per-column validate/cast BEFORE any Redis write (promote/encode) so a
877
+ # rejected input leaves the cached object untouched. Bind all casted
878
+ # python_values to a dict before the encode loop so each stays referenced
879
+ # for the whole encode_session (mirrors the create_raw id-aliasing fix).
880
+ python_values_by_col: dict[str, list] = {}
881
+ for col in columns:
882
+ col_key = encode_key(col)
883
+ dtype_str = _decode_scalar(fields[f"~dt:{col_key}"])
884
+ family = _value_dtype_family(pd.api.types.pandas_dtype(dtype_str))
885
+ raw_values = per_col[col]
886
+ if family in _ELEMENT_SCAN_FAMILIES:
887
+ _assert_scalar_encodable(raw_values, f"DataFrame column {col!r} element")
888
+ if family == "category":
889
+ existing_cats = decode_element(fields[f"~cat:{col_key}"], self._namespace).value
890
+ for v in raw_values:
891
+ if not pd.isna(v) and v not in existing_cats:
892
+ raise ValueError(
893
+ f"Cannot extend category column {col!r}: value {v!r} is not in "
894
+ f"existing categories {existing_cats!r}."
895
+ )
896
+ # Normalize NA sentinels (pd.NA / np.nan / None) to None so
897
+ # encode_element can store them; the non-category branch gets
898
+ # this for free via _python_values, but category skips that cast.
899
+ python_values_by_col[col] = [None if pd.isna(v) else v for v in raw_values]
900
+ else:
901
+ try:
902
+ casted = pd.Series(raw_values, dtype=dtype_str)
903
+ except (ValueError, TypeError) as exc:
904
+ raise ValueError(f"Cannot extend column {col!r} with dtype {dtype_str!r}: {exc}") from exc
905
+ python_values_by_col[col] = _python_values(casted)
906
+
907
+ # RangeIndex: continue the range, or promote to L=1 levels.
908
+ if index_kind == "range":
909
+ range_step = _decode_scalar(fields["~range_step"])
910
+ if range_step == 1 and appended_is_range:
911
+ # range-compatible: just bump range_stop + append columns.
912
+ pipe = r.pipeline(transaction=True)
913
+ with encode_session(self._namespace):
914
+ encoded_per_col = {
915
+ col: [encode_element(v, self._namespace, owner_name, policy) for v in python_values_by_col[col]]
916
+ for col in columns
917
+ }
918
+ for col in columns:
919
+ enc = encoded_per_col[col]
920
+ if enc:
921
+ pipe.rpush(self._get_column_list(col).raw_key, *enc)
922
+ new_stop = _decode_scalar(fields["~range_stop"]) + n
923
+ pipe.hset(self.raw_key, "~range_stop", _encode_extra_value(new_stop))
924
+ pipe.execute()
925
+ self._invalidate_caches()
926
+ self._touch()
927
+ return
928
+ # Promote to L=1 levels. Validate BEFORE the promote (a Redis write)
929
+ # so a rejected input leaves the object untouched (spec §5).
930
+ if got_levels != 1:
931
+ raise ValueError(f"appended index has {got_levels} levels, stored has 1")
932
+ appended_index_labels = index_labels_per_level[0]
933
+ _assert_scalar_encodable(appended_index_labels, "DataFrame index label")
934
+ # Compute the honest combined level dtype: old RangeIndex values are
935
+ # int, so the level is int64 iff the appended labels are int64-castable,
936
+ # else object. Writing the honest dtype here means subsequent extends
937
+ # cast against a truthful dtype instead of a stale int64 (spec §6:
938
+ # infer, do not reject mixed). Mirrors the Series Task-5 fix.
939
+ promote_dtype = _infer_promote_dtype(appended_index_labels)
940
+ self._promote_df_range_to_labels(fields, owner_name, policy, level_dtype=promote_dtype)
941
+ fields = self._raw_fields()
942
+ index_kind = "levels"
943
+ names = _decode_scalar(fields["~index_names"])
944
+
945
+ # levels path (L >= 1)
946
+ L = len(names)
947
+ if got_levels != L:
948
+ raise ValueError(f"appended index has {got_levels} levels, stored has {L}")
949
+ dtypes = _decode_scalar(fields["~index_dtypes"])
950
+ per_level_new: list[list] = []
951
+ for i in range(L):
952
+ col_labels = index_labels_per_level[i]
953
+ # Validate raw labels BEFORE the dtype cast: a string-dtype cast would
954
+ # silently stringify container labels, so scalar-encodability must be
955
+ # confirmed first (mirrors SeriesObject.extend).
956
+ _assert_scalar_encodable(col_labels, f"DataFrame index level {i} label")
957
+ try:
958
+ casted_lvl = pd.array(col_labels, dtype=dtypes[i])
959
+ except (TypeError, ValueError) as exc:
960
+ raise ValueError(f"Cannot extend index level {i} ({names[i]!r}) to dtype {dtypes[i]!r}: {exc}") from exc
961
+ per_level_new.append([None if pd.isna(v) else _to_python_native(v) for v in casted_lvl])
962
+
963
+ # One MULTI/EXEC pipeline: RPUSH each column AND each level list lockstep.
964
+ pipe = r.pipeline(transaction=True)
965
+ with encode_session(self._namespace):
966
+ encoded_per_col = {
967
+ col: [encode_element(v, self._namespace, owner_name, policy) for v in python_values_by_col[col]]
968
+ for col in columns
969
+ }
970
+ encoded_levels = [
971
+ [encode_element(v, self._namespace, owner_name, policy) for v in per_level_new[i]] for i in range(L)
972
+ ]
973
+ for col in columns:
974
+ enc = encoded_per_col[col]
975
+ if enc:
976
+ pipe.rpush(self._get_column_list(col).raw_key, *enc)
977
+ for i in range(L):
978
+ if encoded_levels[i]:
979
+ pipe.rpush(self._get_level_list(i).raw_key, *encoded_levels[i])
980
+ pipe.execute()
981
+
982
+ self._invalidate_caches()
983
+ self._touch()
984
+
985
+ def deepcopy(self) -> pd.DataFrame:
986
+ return self.value.copy(deep=True)
987
+
988
+ def _current_item_count(self) -> int:
989
+ return len(self)
990
+
991
+ # ── public read API ──────────────────────────────────────────────
992
+
993
+ def __len__(self) -> int:
994
+ if self._index_kind() == "range":
995
+ fields = self._raw_fields()
996
+ return _decode_scalar(fields["~range_stop"]) - _decode_scalar(fields["~range_start"])
997
+ idx = self._get_level_list(0)
998
+ return len(idx) if idx is not None else 0
999
+
1000
+ @property
1001
+ def columns(self) -> tuple:
1002
+ return self._column_names()
1003
+
1004
+ @property
1005
+ def index_names(self) -> tuple:
1006
+ return _decode_scalar(self._redis().hget(self.raw_key, "~index_names"))
1007
+
1008
+ def __contains__(self, col) -> bool:
1009
+ try:
1010
+ encoded = encode_key(col)
1011
+ except TypeError:
1012
+ return False
1013
+ return bool(self._redis().hexists(self.raw_key, encoded))
1014
+
1015
+ def __iter__(self):
1016
+ return iter(self._column_names())
1017
+
1018
+ def __getitem__(self, col):
1019
+ """Select column(s) from the cached DataFrame.
1020
+
1021
+ - ``df[col]`` (single-level, scalar) → that column as a ``pd.Series``.
1022
+ - ``df[(a, b, ...)]`` (multi-level, full-depth tuple) → that column as
1023
+ a ``pd.Series``.
1024
+ - ``df[level0]`` or ``df[(a, b)]`` (multi-level, partial prefix) → a
1025
+ sub-``pd.DataFrame`` of the matching columns.
1026
+ - ``df[list_of_keys]`` → a sub-``pd.DataFrame`` (always, even one key).
1027
+
1028
+ Cost is O(N) per selected column but only over the *selected* columns
1029
+ — the rest of the DataFrame is not read — so this is the right tool
1030
+ for a one-shot selective column read (cheap: column count is small).
1031
+ For repeated or heavy work on the same data, call :attr:`value` and
1032
+ work on the local ``pandas`` object. See the module docstring.
1033
+ """
1034
+ fields = self._raw_fields()
1035
+ columns = self._column_names() # tuple of scalars (single) or tuples (multi)
1036
+ col_names = self._columns_names_or_none()
1037
+ if isinstance(col, list):
1038
+ if col_names is not None:
1039
+ nlevels = len(col_names)
1040
+ matched = []
1041
+ missing = []
1042
+ for item in col:
1043
+ key_tuple = item if isinstance(item, tuple) else (item,)
1044
+ k = len(key_tuple)
1045
+ hits = [c for c in columns if tuple(c[:k]) == key_tuple] if k <= nlevels else []
1046
+ if hits:
1047
+ matched.extend(hits)
1048
+ else:
1049
+ missing.append(item)
1050
+ # List misses raise a bare KeyError (pandas-consistent); the scalar branch
1051
+ # adds a depth message because the failure mode there is less obvious.
1052
+ if missing:
1053
+ raise KeyError(missing)
1054
+ return self._select_columns(matched)
1055
+ missing = [c for c in col if c not in columns]
1056
+ if missing:
1057
+ raise KeyError(missing)
1058
+ return self._select_columns(col)
1059
+ if col_names is not None:
1060
+ nlevels = len(col_names)
1061
+ key_tuple = col if isinstance(col, tuple) else (col,)
1062
+ k = len(key_tuple)
1063
+ if k > nlevels:
1064
+ raise KeyError(f"column key depth {k} exceeds column levels {nlevels}")
1065
+ matched = [c for c in columns if tuple(c[:k]) == key_tuple]
1066
+ if not matched:
1067
+ raise KeyError(col)
1068
+ if k == nlevels:
1069
+ # full-depth: duplicate column tuples are rejected at create → exactly one match
1070
+ return self._select_one_series(matched[0], fields)
1071
+ return self._select_columns(matched)
1072
+ # single-level: exact scalar key
1073
+ if col not in columns:
1074
+ raise KeyError(col)
1075
+ return self._select_one_series(col, fields)
1076
+
1077
+ def __setitem__(self, col, values) -> None:
1078
+ """Add or replace a column. ``len(values)`` must equal the current row count."""
1079
+ from ..lifecycle import policy_from_meta, policy_from_defaults
1080
+
1081
+ r = self._redis()
1082
+ fields = self._raw_fields()
1083
+ is_multi = "~columns_names" in fields
1084
+ if is_multi:
1085
+ nlevels = len(_decode_scalar(fields["~columns_names"]))
1086
+ if not isinstance(col, tuple):
1087
+ raise TypeError(
1088
+ "partial/list column mutation is not supported on a MultiIndex-column "
1089
+ "DataFrame; call .value for the full pandas object and operate on it locally."
1090
+ )
1091
+ if len(col) != nlevels:
1092
+ raise ValueError(f"column key has {len(col)} levels, columns have {nlevels}")
1093
+ elif isinstance(col, (tuple, list)):
1094
+ raise TypeError(
1095
+ "tuple/list column keys are not supported on a single-level-column "
1096
+ "DataFrame; call .value for the full pandas object and operate on it locally."
1097
+ )
1098
+ columns = list(_decode_scalar(fields["~columns"]))
1099
+ owner_name = self._data_owner_name()
1100
+ meta = ObjectMeta.load(r, self._canonical_redis_key())
1101
+ policy = policy_from_meta(meta) if meta else policy_from_defaults("permanent", 0)
1102
+
1103
+ expected = self.__len__()
1104
+ if expected == 0 and len(values) > 0:
1105
+ raise ValueError(
1106
+ "cannot introduce rows by adding a column to an empty DataFrame; create the DataFrame with initial data"
1107
+ )
1108
+ if len(values) != expected:
1109
+ raise ValueError(f"column length {len(values)} does not match row count {expected}")
1110
+
1111
+ inferred = pd.Series(values)
1112
+ family = _value_dtype_family(inferred.dtype)
1113
+ if family is None:
1114
+ raise UnsupportedDataFrameError(
1115
+ f"Column {col!r} has dtype {inferred.dtype!r} which is not supported by DataFrameObject."
1116
+ )
1117
+
1118
+ data_values = _python_values(inferred)
1119
+
1120
+ # Validate elements are scalar-encodable BEFORE releasing any old column
1121
+ # refs (a mutation), so a rejected input leaves the object untouched.
1122
+ if family in _ELEMENT_SCAN_FAMILIES:
1123
+ _assert_scalar_encodable(data_values, f"DataFrame column {col!r} element")
1124
+ if family == "category":
1125
+ _assert_scalar_encodable(
1126
+ inferred.cat.categories.tolist(),
1127
+ f"DataFrame column {col!r} category value",
1128
+ )
1129
+
1130
+ col_key = encode_key(col)
1131
+ replacing = col in columns
1132
+ old_family = None
1133
+ if replacing:
1134
+ old_ref = r.hget(self.raw_key, col_key)
1135
+ self._release_element_ref(old_ref)
1136
+ old_dtype_str = _decode_scalar(fields[f"~dt:{col_key}"])
1137
+ old_family = _value_dtype_family(pd.api.types.pandas_dtype(old_dtype_str))
1138
+ if old_family == "category":
1139
+ self._release_element_ref(r.hget(self.raw_key, f"~cat:{col_key}"))
1140
+
1141
+ with encode_session(self._namespace):
1142
+ mapping = {
1143
+ col_key: encode_element(data_values, self._namespace, owner_name, policy),
1144
+ f"~dt:{col_key}": _encode_extra_value(str(inferred.dtype)),
1145
+ }
1146
+ if family == "category":
1147
+ mapping[f"~cat:{col_key}"] = encode_element(
1148
+ inferred.cat.categories.tolist(), self._namespace, owner_name, policy
1149
+ )
1150
+ mapping[f"~ord:{col_key}"] = _encode_extra_value(bool(inferred.cat.ordered))
1151
+
1152
+ if replacing and old_family == "category" and family != "category":
1153
+ r.hdel(self.raw_key, f"~cat:{col_key}", f"~ord:{col_key}")
1154
+ r.hset(self.raw_key, mapping=mapping)
1155
+
1156
+ if not replacing:
1157
+ columns.append(col)
1158
+ r.hset(self.raw_key, "~columns", _encode_extra_value(tuple(columns)))
1159
+
1160
+ self._invalidate_caches()
1161
+ self._touch()
1162
+
1163
+ def __delitem__(self, col) -> None:
1164
+ """Drop a column. The shared index is retained (even for the last column)."""
1165
+ r = self._redis()
1166
+ fields = self._raw_fields()
1167
+ is_multi = "~columns_names" in fields
1168
+ if is_multi:
1169
+ nlevels = len(_decode_scalar(fields["~columns_names"]))
1170
+ if not isinstance(col, tuple):
1171
+ raise TypeError(
1172
+ "partial/list column deletion is not supported on a MultiIndex-column "
1173
+ "DataFrame; call .value for the full pandas object and operate on it locally."
1174
+ )
1175
+ if len(col) != nlevels:
1176
+ raise ValueError(f"column key has {len(col)} levels, columns have {nlevels}")
1177
+ elif isinstance(col, (tuple, list)):
1178
+ raise TypeError(
1179
+ "tuple/list column keys are not supported on a single-level-column "
1180
+ "DataFrame; call .value for the full pandas object and operate on it locally."
1181
+ )
1182
+ columns = list(_decode_scalar(fields["~columns"]))
1183
+ if col not in columns:
1184
+ raise KeyError(col)
1185
+ col_key = encode_key(col)
1186
+ old_ref = r.hget(self.raw_key, col_key)
1187
+ dtype_str = _decode_scalar(fields[f"~dt:{col_key}"])
1188
+ family = _value_dtype_family(pd.api.types.pandas_dtype(dtype_str))
1189
+ old_cat_ref = r.hget(self.raw_key, f"~cat:{col_key}") if family == "category" else None
1190
+
1191
+ fields_to_del = [col_key, f"~dt:{col_key}"]
1192
+ if family == "category":
1193
+ fields_to_del.extend([f"~cat:{col_key}", f"~ord:{col_key}"])
1194
+ r.hdel(self.raw_key, *fields_to_del)
1195
+ self._release_element_ref(old_ref)
1196
+ if old_cat_ref:
1197
+ self._release_element_ref(old_cat_ref)
1198
+
1199
+ columns.remove(col)
1200
+ r.hset(self.raw_key, "~columns", _encode_extra_value(tuple(columns)))
1201
+
1202
+ self._invalidate_caches()
1203
+ self._touch()
1204
+
1205
+ @property
1206
+ def iloc(self) -> "_DfILocIndexer":
1207
+ return _DfILocIndexer(self)
1208
+
1209
+ @property
1210
+ def loc(self) -> "_DfLocIndexer":
1211
+ """Row label accessor (depth-k prefix on MultiIndex rows).
1212
+
1213
+ O(N) selective scan — fine for a one-shot partial read, not for hot
1214
+ loops; for repeated lookups call :attr:`value`. Slices and label-lists
1215
+ are rejected (call :attr:`value`). Column selection uses
1216
+ ``df[...]`` (cheap; few columns). See the module docstring.
1217
+ """
1218
+ return _DfLocIndexer(self)
1219
+
1220
+ def _positional_row_slice(self, sl: slice) -> pd.DataFrame:
1221
+ columns = self._column_names()
1222
+ fields = self._raw_fields()
1223
+ kind = _decode_scalar(fields["~index_kind"])
1224
+ if kind == "range":
1225
+ index_obj = _build_index(fields, self._namespace, prefix="~")[sl]
1226
+ else:
1227
+ names = _decode_scalar(fields["~index_names"])
1228
+ arrays = [self._get_level_list(i)[sl] for i in range(len(names))]
1229
+ if len(names) == 1:
1230
+ index_obj = pd.Index(arrays[0], name=names[0])
1231
+ else:
1232
+ index_obj = pd.MultiIndex.from_arrays(arrays, names=names)
1233
+ data_dict = {}
1234
+ for col in columns:
1235
+ col_list = self._get_column_list(col)
1236
+ values = col_list[sl] if col_list is not None else []
1237
+ data_dict[col] = _reconstruct_column(values, index_obj, fields, col, self._namespace)
1238
+ return pd.DataFrame(data_dict, columns=_build_columns_index(fields, columns))
1239
+
1240
+
1241
+ def register() -> None:
1242
+ """Register the pandas Series + DataFrame types with GPDCTypeRegistry (idempotent)."""
1243
+ GPDCTypeRegistry.register(
1244
+ python_type=pd.Series,
1245
+ gpdc_type="collection:pandas_series",
1246
+ object_class=SeriesObject,
1247
+ codec_class=SeriesCodec,
1248
+ priority=0,
1249
+ )
1250
+ GPDCTypeRegistry.register(
1251
+ python_type=pd.DataFrame,
1252
+ gpdc_type="collection:pandas_dataframe",
1253
+ object_class=DataFrameObject,
1254
+ codec_class=DataFrameCodec,
1255
+ priority=0,
1256
+ )
1257
+
1258
+
1259
+ class SeriesCodec(GPDCContainerCodec):
1260
+ @classmethod
1261
+ def create_raw(cls, namespace, raw_key, value: pd.Series, owner_name, cache_mode, ttl) -> None:
1262
+ from ..lifecycle import policy_from_defaults
1263
+
1264
+ family = _value_dtype_family(value.dtype)
1265
+ if family is None:
1266
+ raise UnsupportedSeriesDtypeError(str(value.dtype), is_index=False)
1267
+ index = value.index
1268
+ kind = _index_kind(index)
1269
+ if kind is None:
1270
+ raise UnsupportedSeriesDtypeError(str(index.dtype), is_index=True)
1271
+
1272
+ policy = policy_from_defaults(cache_mode, ttl)
1273
+ data_values = _python_values(value)
1274
+ categories = value.cat.categories.tolist() if family == "category" else None
1275
+
1276
+ if family in _ELEMENT_SCAN_FAMILIES:
1277
+ _assert_scalar_encodable(data_values, "Series element")
1278
+ if categories is not None:
1279
+ _assert_scalar_encodable(categories, "Series category value")
1280
+
1281
+ fields: dict[str, str] = {}
1282
+ with encode_session(namespace):
1283
+ fields["data"] = encode_element(data_values, namespace, owner_name, policy)
1284
+ if family == "category":
1285
+ fields["categories"] = encode_element(categories, namespace, owner_name, policy)
1286
+ if kind == "levels":
1287
+ # Bind all level label lists to locals before encoding (id-aliasing guard).
1288
+ level_labels = [_level_labels_and_dtype(index, i) for i in range(index.nlevels)]
1289
+ for i, (labels, _dtype) in enumerate(level_labels):
1290
+ _assert_scalar_encodable(labels, f"Series index level {i} label")
1291
+ fields[f"index_l{i}"] = encode_element(labels, namespace, owner_name, policy)
1292
+
1293
+ fields["dtype"] = _encode_extra_value(str(value.dtype))
1294
+ fields["name"] = _encode_extra_value(value.name)
1295
+ fields["index_kind"] = _encode_extra_value(kind)
1296
+ fields["index_names"] = _encode_extra_value(tuple(index.names))
1297
+ if kind == "levels":
1298
+ fields["index_dtypes"] = _encode_extra_value(tuple(dt for _, dt in level_labels))
1299
+ if family == "category":
1300
+ fields["ordered"] = _encode_extra_value(bool(value.cat.ordered))
1301
+ if kind == "range":
1302
+ fields["range_start"] = _encode_extra_value(int(index.start))
1303
+ fields["range_stop"] = _encode_extra_value(int(index.stop))
1304
+ fields["range_step"] = _encode_extra_value(int(index.step))
1305
+
1306
+ namespace._redis().hset(raw_key, mapping=fields)
1307
+
1308
+ @classmethod
1309
+ def read_raw(cls, namespace, raw_key) -> pd.Series:
1310
+ fields = namespace._redis().hgetall(raw_key)
1311
+ if not fields:
1312
+ return pd.Series([], dtype="object")
1313
+ data = decode_element(fields["data"], namespace)
1314
+ values = data.deepcopy()
1315
+ index_kind = _decode_scalar(fields["index_kind"])
1316
+ if index_kind == "range":
1317
+ index_obj = pd.RangeIndex(
1318
+ _decode_scalar(fields["range_start"]),
1319
+ _decode_scalar(fields["range_stop"]),
1320
+ _decode_scalar(fields["range_step"]),
1321
+ name=_decode_scalar(fields["index_names"])[0],
1322
+ )
1323
+ else: # levels
1324
+ index_obj = _build_levels_index(fields, namespace)
1325
+ return _reconstruct(values, index_obj, fields, namespace)
1326
+
1327
+ @classmethod
1328
+ def delete_raw(cls, namespace, raw_key) -> None:
1329
+ namespace._redis().delete(raw_key)
1330
+
1331
+
1332
+ _UNLOADED = object()
1333
+
1334
+
1335
+ class _ILocIndexer:
1336
+ def __init__(self, owner: "SeriesObject") -> None:
1337
+ self._owner = owner
1338
+
1339
+ def __getitem__(self, key):
1340
+ owner = self._owner
1341
+ data = owner._get_data_list()
1342
+ if data is None:
1343
+ raise IndexError("SeriesObject has no data")
1344
+ if isinstance(key, (int, np.integer)):
1345
+ return data[int(key)]
1346
+ if isinstance(key, slice):
1347
+ values = data[key]
1348
+ return owner._build_series_from(values, owner._positional_index_slice(key))
1349
+ raise TypeError(
1350
+ ".iloc list (positional fancy) indexing is not supported on this cache object; "
1351
+ "use .iloc with an int or slice, or call .value for the full pandas object."
1352
+ )
1353
+
1354
+
1355
+ class _LocIndexer:
1356
+ def __init__(self, owner: "SeriesObject") -> None:
1357
+ self._owner = owner
1358
+
1359
+ def __getitem__(self, key):
1360
+ owner = self._owner
1361
+ if isinstance(key, slice):
1362
+ raise TypeError(
1363
+ ".loc slices are not supported on this cache object; "
1364
+ "call .value for the full pandas object and operate on it locally, "
1365
+ "or use .iloc for positional slicing."
1366
+ )
1367
+ if isinstance(key, list):
1368
+ raise TypeError(
1369
+ ".loc label-list selection is not supported on this cache object; "
1370
+ "call .value for the full pandas object and operate on it locally."
1371
+ )
1372
+ fields = owner._raw_fields()
1373
+ if _decode_scalar(fields["index_kind"]) == "range":
1374
+ raise TypeError(
1375
+ ".loc on a RangeIndex is not supported on this cache object; "
1376
+ "use .iloc for positional slicing, or call .value for the full pandas object."
1377
+ )
1378
+ names = _decode_scalar(fields["index_names"])
1379
+ L = len(names)
1380
+ # normalise key to a depth-k tuple
1381
+ key_tuple = key if isinstance(key, tuple) else (key,)
1382
+ k = len(key_tuple)
1383
+ if k > L:
1384
+ raise KeyError(f"key depth {k} exceeds index levels {L}")
1385
+ # Fetch all L level values once (ListObject.value is uncached -> each call is a Redis lrange).
1386
+ all_level_values = [owner._get_level_list(i).value for i in range(L)]
1387
+ # NA-label matching follows Python == semantics; nan labels won't match (pre-existing behaviour).
1388
+ positions = [
1389
+ pos for pos in range(len(owner)) if all(all_level_values[i][pos] == key_tuple[i] for i in range(k))
1390
+ ]
1391
+ if not positions:
1392
+ raise KeyError(key)
1393
+ # full-depth + exactly one match -> scalar value (pandas semantics)
1394
+ if k == L and len(positions) == 1:
1395
+ return owner._get_data_list()[positions[0]]
1396
+ data = owner._get_data_list()
1397
+ values = [data[p] for p in positions]
1398
+ # Build the result index from the SAME all_level_values (matched positions, all L levels).
1399
+ arrays = [[lvl[p] for p in positions] for lvl in all_level_values]
1400
+ if L == 1:
1401
+ sliced_index = pd.Index(arrays[0], name=names[0])
1402
+ else:
1403
+ sliced_index = pd.MultiIndex.from_arrays(arrays, names=names)
1404
+ return owner._build_series_from(values, sliced_index)
1405
+
1406
+
1407
+ class SeriesObject(GPDCContainerObject):
1408
+ TYPE = "collection:pandas_series"
1409
+
1410
+ def __init__(self, name, namespace, *, canonical_name=None) -> None:
1411
+ super().__init__(name, namespace, canonical_name=canonical_name)
1412
+ self._data_list_cache = _UNLOADED
1413
+ self._level_lists_cache: dict = {}
1414
+
1415
+ @property
1416
+ def type(self) -> str:
1417
+ return self.TYPE
1418
+
1419
+ @property
1420
+ def value(self) -> pd.Series:
1421
+ """Full materialization as a plain :class:`pandas.Series`.
1422
+
1423
+ Use this for repeated or heavy work on the same data (sorting, many
1424
+ label lookups, complex pandas semantics, or operations the online API
1425
+ rejects). For a one-shot selective read, prefer :attr:`iloc` / :attr:`loc`.
1426
+ See the module docstring for the full usage guidance.
1427
+ """
1428
+ return SeriesCodec.read_raw(self._namespace, self.raw_key)
1429
+
1430
+ @value.setter
1431
+ def value(self, new_value) -> None:
1432
+ self._reject_value_replacement()
1433
+
1434
+ def _get_data_list(self):
1435
+ if self._data_list_cache is _UNLOADED:
1436
+ ref = self._redis().hget(self.raw_key, "data")
1437
+ self._data_list_cache = decode_element(ref, self._namespace) if ref else None
1438
+ return self._data_list_cache
1439
+
1440
+ def _get_level_list(self, level_i: int):
1441
+ """Return the per-level index child ListObject, or None when the level is absent."""
1442
+ if level_i not in self._level_lists_cache:
1443
+ ref = self._redis().hget(self.raw_key, f"index_l{level_i}")
1444
+ self._level_lists_cache[level_i] = decode_element(ref, self._namespace) if ref else None
1445
+ return self._level_lists_cache[level_i]
1446
+
1447
+ def _raw_fields(self) -> dict:
1448
+ return self._redis().hgetall(self.raw_key)
1449
+
1450
+ def __len__(self) -> int:
1451
+ data = self._get_data_list()
1452
+ return len(data) if data is not None else 0
1453
+
1454
+ @property
1455
+ def dtype(self) -> str:
1456
+ return _decode_scalar(self._redis().hget(self.raw_key, "dtype"))
1457
+
1458
+ @property
1459
+ def name(self):
1460
+ return _decode_scalar(self._redis().hget(self.raw_key, "name"))
1461
+
1462
+ @property
1463
+ def index_names(self) -> tuple:
1464
+ return _decode_scalar(self._redis().hget(self.raw_key, "index_names"))
1465
+
1466
+ def deepcopy(self) -> pd.Series:
1467
+ return self.value.copy(deep=True)
1468
+
1469
+ def _current_item_count(self) -> int:
1470
+ return len(self)
1471
+
1472
+ @property
1473
+ def iloc(self) -> "_ILocIndexer":
1474
+ return _ILocIndexer(self)
1475
+
1476
+ @property
1477
+ def loc(self) -> "_LocIndexer":
1478
+ """Row label accessor (depth-k prefix on a MultiIndex).
1479
+
1480
+ O(N) selective scan — fine for a one-shot partial read, but do not use
1481
+ in a hot loop; for repeated lookups call :attr:`value` and work on the
1482
+ local ``pandas`` object. Slices and label-lists are rejected (call
1483
+ :attr:`value`). See the module docstring.
1484
+ """
1485
+ return _LocIndexer(self)
1486
+
1487
+ def _positional_index_slice(self, sl: slice):
1488
+ """Return the pd.Index for a positional slice, honoring RangeIndex and MultiIndex."""
1489
+ fields = self._raw_fields()
1490
+ index_kind = _decode_scalar(fields["index_kind"])
1491
+ names = _decode_scalar(fields["index_names"])
1492
+ if index_kind == "range":
1493
+ full = pd.RangeIndex(
1494
+ _decode_scalar(fields["range_start"]),
1495
+ _decode_scalar(fields["range_stop"]),
1496
+ _decode_scalar(fields["range_step"]),
1497
+ name=names[0],
1498
+ )
1499
+ return full[sl]
1500
+ # levels: slice each level list
1501
+ arrays = [self._get_level_list(i)[sl] for i in range(len(names))]
1502
+ if len(names) == 1:
1503
+ return pd.Index(arrays[0], name=names[0])
1504
+ return pd.MultiIndex.from_arrays(arrays, names=names)
1505
+
1506
+ def _build_series_from(self, values: list, index_obj) -> pd.Series:
1507
+ return _reconstruct(values, index_obj, self._raw_fields(), self._namespace)
1508
+
1509
+ def extend(self, other) -> None:
1510
+ """Append rows in place. Accepts a Series, a list of Series, or (index, value) tuples.
1511
+
1512
+ Mutates the cached object in Redis and returns None. (Differs from
1513
+ pandas pd.concat, which is non-mutating and returns a new Series.)
1514
+ """
1515
+ from ..lifecycle import policy_from_defaults
1516
+
1517
+ fields = self._raw_fields()
1518
+ dtype_str = _decode_scalar(fields["dtype"])
1519
+ if isinstance(pd.api.types.pandas_dtype(dtype_str), pd.CategoricalDtype):
1520
+ raise TypeError(
1521
+ "Extending a categorical-dtype Series is not supported; "
1522
+ "create a new Series with the desired categories instead."
1523
+ )
1524
+ index_kind = _decode_scalar(fields["index_kind"])
1525
+ names = _decode_scalar(fields["index_names"])
1526
+
1527
+ labels, data_values, got_levels = _normalize_extend_input(other)
1528
+
1529
+ try:
1530
+ casted = pd.Series(data_values, dtype=dtype_str)
1531
+ except (ValueError, TypeError) as exc:
1532
+ raise ValueError(f"Cannot extend Series of dtype {dtype_str!r} with given values: {exc}") from exc
1533
+ new_data = _python_values(casted)
1534
+
1535
+ # Validate elements are scalar-encodable BEFORE any Redis write (promote/encode)
1536
+ # so a rejected input leaves the cached object untouched.
1537
+ series_family = _value_dtype_family(pd.api.types.pandas_dtype(dtype_str))
1538
+ if series_family in _ELEMENT_SCAN_FAMILIES:
1539
+ _assert_scalar_encodable(new_data, "Series element")
1540
+
1541
+ r = self._redis()
1542
+ meta = ObjectMeta.load(r, self._canonical_redis_key())
1543
+ policy = policy_from_meta(meta) if meta else policy_from_defaults("permanent", 0)
1544
+ owner_name = self._data_owner_name()
1545
+
1546
+ if index_kind == "range":
1547
+ # Range path: promote to L=1 levels if the appended input isn't range-compatible,
1548
+ # else just bump range_stop. Mirror the pre-existing range behaviour.
1549
+ if not _is_range_indexed_series(other):
1550
+ # A range Series promotes to L=1 levels. Validate the appended input
1551
+ # BEFORE the promote (a Redis write) so a rejected input leaves the
1552
+ # object untouched (spec §5).
1553
+ if got_levels != 1:
1554
+ raise ValueError(f"appended index has {got_levels} levels, stored has 1")
1555
+ appended_index_labels = [lbl[0] for lbl in labels]
1556
+ _assert_scalar_encodable(appended_index_labels, "Series index label")
1557
+ # Compute the honest combined level dtype: old RangeIndex values are int,
1558
+ # so the level is int64 iff the appended labels are int64-castable, else
1559
+ # object. Writing the honest dtype here means subsequent extends cast
1560
+ # against a truthful dtype instead of a stale int64 (spec §6: infer, do
1561
+ # not reject mixed).
1562
+ promote_dtype = _infer_promote_dtype(appended_index_labels)
1563
+ self._promote_range_to_labels(fields, owner_name, policy, level_dtype=promote_dtype)
1564
+ fields = self._raw_fields()
1565
+ names = _decode_scalar(fields["index_names"])
1566
+ else:
1567
+ data_list = self._get_data_list()
1568
+ new_stop = _decode_scalar(fields["range_stop"]) + len(new_data)
1569
+ pipe = r.pipeline(transaction=True)
1570
+ with encode_session(self._namespace):
1571
+ encoded_data = [encode_element(v, self._namespace, owner_name, policy) for v in new_data]
1572
+ if encoded_data:
1573
+ pipe.rpush(data_list.raw_key, *encoded_data)
1574
+ pipe.hset(self.raw_key, "range_stop", _encode_extra_value(new_stop))
1575
+ pipe.execute()
1576
+ self._data_list_cache = _UNLOADED
1577
+ self._level_lists_cache = {}
1578
+ self._touch()
1579
+ return
1580
+
1581
+ # levels path (L >= 1)
1582
+ L = len(names)
1583
+ if got_levels != L:
1584
+ raise ValueError(f"appended index has {got_levels} levels, stored has {L}")
1585
+ dtypes = _decode_scalar(fields["index_dtypes"])
1586
+ per_level_new: list[list] = []
1587
+ for i in range(L):
1588
+ col_labels = [lbl[i] for lbl in labels]
1589
+ # Validate raw labels BEFORE the dtype cast: a string-dtype cast would
1590
+ # silently stringify container labels (e.g. [1, 2] -> "[1, 2]"), so the
1591
+ # cast must not run before scalar-encodability is confirmed. The promote
1592
+ # branch already validated L=0 labels; this re-check is a harmless no-op
1593
+ # for that path but required for the direct levels path.
1594
+ _assert_scalar_encodable(col_labels, f"Series index level {i} label")
1595
+ try:
1596
+ casted_lvl = pd.array(col_labels, dtype=dtypes[i])
1597
+ except (TypeError, ValueError) as exc:
1598
+ raise ValueError(f"Cannot extend index level {i} ({names[i]!r}) to dtype {dtypes[i]!r}: {exc}") from exc
1599
+ per_level_new.append([None if pd.isna(v) else _to_python_native(v) for v in casted_lvl])
1600
+
1601
+ data_list = self._get_data_list()
1602
+ pipe = r.pipeline(transaction=True)
1603
+ with encode_session(self._namespace):
1604
+ encoded_data = [encode_element(v, self._namespace, owner_name, policy) for v in new_data]
1605
+ encoded_levels = [
1606
+ [encode_element(v, self._namespace, owner_name, policy) for v in per_level_new[i]] for i in range(L)
1607
+ ]
1608
+ if encoded_data:
1609
+ pipe.rpush(data_list.raw_key, *encoded_data)
1610
+ for i in range(L):
1611
+ if encoded_levels[i]:
1612
+ pipe.rpush(self._get_level_list(i).raw_key, *encoded_levels[i])
1613
+ pipe.execute()
1614
+ self._data_list_cache = _UNLOADED
1615
+ self._level_lists_cache = {}
1616
+ self._touch()
1617
+
1618
+ def _promote_range_to_labels(self, fields: dict, owner_name, policy, level_dtype="int64") -> None:
1619
+ start = _decode_scalar(fields["range_start"])
1620
+ stop = _decode_scalar(fields["range_stop"])
1621
+ step = _decode_scalar(fields["range_step"])
1622
+ labels = list(pd.RangeIndex(start, stop, step))
1623
+ with encode_session(self._namespace):
1624
+ ref = encode_element(labels, self._namespace, owner_name, policy)
1625
+ r = self._redis()
1626
+ mapping = {
1627
+ "index_l0": ref,
1628
+ "index_kind": _encode_extra_value("levels"),
1629
+ "index_dtypes": _encode_extra_value((level_dtype,)),
1630
+ }
1631
+ for k in ("range_start", "range_stop", "range_step"):
1632
+ if k in fields:
1633
+ r.hdel(self.raw_key, k)
1634
+ r.hset(self.raw_key, mapping=mapping)
1635
+ self._level_lists_cache = {}