fix: restore failing assign, named agg, memory usage, and combine-first tests

Agent-Logs-Url: https://github.com/githubnext/tsessebe/sessions/78d95042-4c45-4c08-be04-02e8afe42b80

Co-authored-by: mrjf <180956+mrjf@users.noreply.github.com>

Author: Copilot

src/core/frame.ts

@@ -302,16 +302,27 @@ export class DataFrame {
    * const df2 = df.assign({ c: [7, 8, 9] });
    * ```
    */
-  assign(newCols: Readonly<Record<string, readonly Scalar[] | Series<Scalar>>>): DataFrame {
-    const colMap = new Map<string, Series<Scalar>>(this._columns);
-    for (const [name, val] of Object.entries(newCols)) {
-      if (val instanceof Series) {
-        colMap.set(name, val);
-      } else {
-        colMap.set(name, new Series({ data: val, index: this.index }));
-      }
+  assign(
+    newCols: Readonly<
+      Record<
+        string,
+        readonly Scalar[] | Series<Scalar> | ((df: DataFrame) => readonly Scalar[] | Series<Scalar>)
+      >
+    >,
+  ): DataFrame {
+    let working: DataFrame = this;
+    for (const [name, spec] of Object.entries(newCols)) {
+      const resolved: readonly Scalar[] | Series<Scalar> =
+        typeof spec === "function" ? spec(working) : spec;
+      const series: Series<Scalar> =
+        resolved instanceof Series
+          ? resolved
+          : new Series({ data: resolved, index: working.index });
+      const colMap = new Map<string, Series<Scalar>>(working._columns);
+      colMap.set(name, series);
+      working = new DataFrame(colMap, working.index);
     }
-    return new DataFrame(colMap, this.index);
+    return working;
   }
 
   /** Drop one or more columns by name.  Returns a new DataFrame. */

src/groupby/groupby.ts

@@ -21,6 +21,7 @@ import { RangeIndex } from "../core/index.ts";
 import { DataFrame } from "../core/index.ts";
 import { Series } from "../core/index.ts";
 import type { Label, Scalar } from "../types.ts";
+import type { NamedAggSpec } from "./named_agg.ts";
 
 // ─── types ────────────────────────────────────────────────────────────────────
 
@@ -304,6 +305,53 @@ export class DataFrameGroupBy {
     return this._runAgg(colSpecs, asIndex);
   }
 
+  /**
+   * Aggregate each group using named aggregation specs.
+   *
+   * Each key in `spec` becomes the output column name; the `NamedAgg` value
+   * specifies which source column to aggregate and how.
+   */
+  aggNamed(spec: NamedAggSpec, asIndex = true): DataFrame {
+    const groupKeys = this._groups.map((g) => g.key);
+    const resultCols: Record<string, Scalar[]> = {};
+
+    if (!asIndex) {
+      if (this._by.length === 1) {
+        const byCol = this._by[0] as string;
+        resultCols[byCol] = groupKeys.slice();
+      } else {
+        for (const by of this._by) {
+          resultCols[by] = [];
+        }
+        for (const g of this._groups) {
+          const parts = (g.key as string).split("__SEP__");
+          this._by.forEach((by, idx) => {
+            const byArr = resultCols[by];
+            if (byArr !== undefined) {
+              byArr.push(parts[idx] ?? null);
+            }
+          });
+        }
+      }
+    }
+
+    for (const [outCol, namedSpec] of Object.entries(spec)) {
+      const fn = resolveAgg(namedSpec.aggfunc);
+      const srcVals = this._df.col(namedSpec.column).values as readonly Scalar[];
+      resultCols[outCol] = this._groups.map((g) =>
+        fn(g.positions.map((p) => srcVals[p] as Scalar)),
+      );
+    }
+
+    const rowIdx: Index<Label> = asIndex
+      ? new Index<Label>(groupKeys)
+      : defaultIndex(groupKeys.length);
+
+    return DataFrame.fromColumns(resultCols as Record<string, readonly Scalar[]>, {
+      index: rowIdx,
+    });
+  }
+
   /** Shorthand for `agg("sum")` — numeric columns only, like pandas. */
   sum(): DataFrame {
     const cols = this._numericValueCols();

src/stats/combine_first.ts

@@ -149,8 +149,8 @@ export function combineFirstDataFrame(self: DataFrame, other: DataFrame): DataFr
   const resultColMap = new Map<string, Series<Scalar>>();
 
   for (const colName of unionCols) {
-    const selfHasCol = self.hasColumn(colName);
-    const otherHasCol = other.hasColumn(colName);
+    const selfHasCol = self.has(colName);
+    const otherHasCol = other.has(colName);
 
     const data: Scalar[] = [];
 

src/stats/memory_usage.ts

@@ -176,7 +176,7 @@ export function dataFrameMemoryUsage(
     values.push(indexMemoryBytes(df.index, deep));
   }
 
-  for (const [colName, col] of df) {
+  for (const [colName, col] of df.items()) {
     names.push(colName);
     let mem: number;
     if (deep) {

tests/core/assign.test.ts

@@ -207,7 +207,7 @@ describe("dataFrameAssign — property tests", () => {
           }
 
           const df2 = dataFrameAssign(df, spec);
-          expect(df2.columns.length).toBe(1 + nNew);
+          expect(df2.columns.size).toBe(1 + nNew);
           expect(df2.shape[0]).toBe(nRows);
         },
       ),
@@ -256,10 +256,10 @@ describe("dataFrameAssign — property tests", () => {
         fc.array(fc.integer({ min: 0, max: 100 }), { minLength: 1, maxLength: 8 }),
         (arr) => {
           const df = DataFrame.fromColumns({ x: arr });
-          const origColCount = df.columns.length;
+          const origColCount = df.columns.size;
           dataFrameAssign(df, { y: arr, z: arr });
 
-          expect(df.columns.length).toBe(origColCount);
+          expect(df.columns.size).toBe(origColCount);
         },
       ),
     );

tests/stats/memory_usage.test.ts

@@ -14,6 +14,16 @@ import {
 } from "../../src/index.ts";
 import type { Scalar } from "../../src/index.ts";
 
+function dfFromSeries(columns: Record<string, Series<Scalar>>): DataFrame {
+  const colMap = new Map<string, Series<Scalar>>();
+  for (const [name, series] of Object.entries(columns)) {
+    colMap.set(name, series);
+  }
+  const firstCol = Object.values(columns)[0];
+  const index = new RangeIndex(firstCol?.size ?? 0);
+  return new DataFrame(colMap, index);
+}
+
 // ─── seriesMemoryUsage ────────────────────────────────────────────────────────
 
 describe("seriesMemoryUsage", () => {
@@ -155,7 +165,7 @@ describe("seriesMemoryUsage", () => {
 
 describe("dataFrameMemoryUsage", () => {
   it("returns Series indexed by column names with Index entry", () => {
-    const df = new DataFrame({
+    const df = dfFromSeries({
       a: new Series<number>({ data: [1, 2, 3], dtype: Dtype.int32 }),
       b: new Series<number>({ data: [4, 5, 6], dtype: Dtype.float64 }),
     });
@@ -167,15 +177,15 @@ describe("dataFrameMemoryUsage", () => {
   });
 
   it("Index row = 24 bytes (RangeIndex)", () => {
-    const df = new DataFrame({
+    const df = dfFromSeries({
       x: new Series<number>({ data: [1, 2], dtype: Dtype.int64 }),
     });
     const mu = dataFrameMemoryUsage(df);
     expect(mu.at("Index")).toBe(24);
   });
 
   it("column bytes = n × itemsize for fixed-width", () => {
-    const df = new DataFrame({
+    const df = dfFromSeries({
       a: new Series<number>({ data: [10, 20, 30], dtype: Dtype.int32 }),
       b: new Series<number>({ data: [1.0, 2.0, 3.0], dtype: Dtype.float64 }),
     });
@@ -185,15 +195,15 @@ describe("dataFrameMemoryUsage", () => {
   });
 
   it("string column = n × 8 bytes (pointers) when shallow", () => {
-    const df = new DataFrame({
+    const df = dfFromSeries({
       s: new Series<Scalar>({ data: ["hello", "world"], dtype: Dtype.string }),
     });
     const mu = dataFrameMemoryUsage(df);
     expect(mu.at("s")).toBe(2 * 8);
   });
 
   it("index=false excludes 'Index' row", () => {
-    const df = new DataFrame({
+    const df = dfFromSeries({
       a: new Series<number>({ data: [1, 2], dtype: Dtype.int32 }),
     });
     const mu = dataFrameMemoryUsage(df, { index: false });
@@ -203,7 +213,7 @@ describe("dataFrameMemoryUsage", () => {
   });
 
   it("deep=true string column reflects actual string sizes", () => {
-    const df = new DataFrame({
+    const df = dfFromSeries({
       s: new Series<Scalar>({ data: ["hi", "hello"], dtype: Dtype.string }),
     });
     const mu = dataFrameMemoryUsage(df, { deep: true, index: false });
@@ -212,12 +222,12 @@ describe("dataFrameMemoryUsage", () => {
   });
 
   it("result Series name is 'memory_usage'", () => {
-    const df = new DataFrame({ x: new Series<number>({ data: [1], dtype: Dtype.int64 }) });
+    const df = dfFromSeries({ x: new Series<number>({ data: [1], dtype: Dtype.int64 }) });
     expect(dataFrameMemoryUsage(df).name).toBe("memory_usage");
   });
 
   it("total() matches sum of all column bytes (no index)", () => {
-    const df = new DataFrame({
+    const df = dfFromSeries({
       a: new Series<number>({ data: [1, 2, 3, 4], dtype: Dtype.int32 }),
       b: new Series<number>({ data: [1.0, 2.0, 3.0, 4.0], dtype: Dtype.float64 }),
     });
@@ -227,7 +237,7 @@ describe("dataFrameMemoryUsage", () => {
   });
 
   it("empty DataFrame returns only Index row", () => {
-    const df = new DataFrame({});
+    const df = dfFromSeries({});
     const mu = dataFrameMemoryUsage(df);
     expect(mu.index.size).toBe(1);
     expect(mu.at("Index")).toBe(24);
@@ -245,7 +255,7 @@ describe("dataFrameMemoryUsage", () => {
               dtype: Dtype.int32,
             });
           });
-          const df = new DataFrame(cols);
+          const df = dfFromSeries(cols);
           const mu = dataFrameMemoryUsage(df);
           return mu.values.every((v) => v >= 0);
         },
@@ -265,7 +275,7 @@ describe("dataFrameMemoryUsage", () => {
               dtype: Dtype.int32,
             });
           });
-          const df = new DataFrame(cols);
+          const df = dfFromSeries(cols);
           const mu = dataFrameMemoryUsage(df, { index: false });
           const expected = colSizes.reduce((s, n) => s + n * 4, 0);
           return mu.sum() === expected;