diff --git a/CHANGELOG.md b/CHANGELOG.md index f6019d8..a4b1828 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,22 @@ All notable changes to this project will be documented in this file. --- +## [Unreleased] + +### Added + +### Fixed + +- **`SafeOperation` never actually recovered panics** — `recover()` was called inside a helper function invoked by the deferred closure, where Go defines it as a no-op, so a panic inside the wrapped function escaped to the caller — the exact failure mode the "no panics" design promise forbids. The recovery helper is now installed directly as the deferred function (`defer recoverFromPanic(op, &err)`) and is covered by tests that drive string, error, arbitrary-value, and genuine runtime panics through `SafeOperation`. + +### Added + +- **Column derivation & transforms** — `Apply(newColumn, fn)` derives a column from whole rows, `Map(src, dst, fn)` transforms one column element-wise, `FillNA(column, value)` replaces missing values, and `DropNA(columns...)` drops rows with missing values (all columns when none given). "Missing" means the column type's zero value, the library's existing CSV/JSONL convention — documented explicitly, since a legitimate 0 is indistinguishable from missing in an int column. Derived-column types are inferred from the returned values (all ints → `Int64Type`, int/float mix → `Float64Type`, uniform string/bool/time kept, other mixes fall back to `StringType` with canonical formatting; `nil` fills the zero value). An existing target column is replaced in place, receivers are never mutated, and a panic inside the user function is captured as an error on the returned frame instead of escaping. + +- **Join operations** — `Join(other, on, how)` and `JoinOn(other, leftKeys, rightKeys, how)` with all four join types (`InnerJoin`, `LeftJoin`, `RightJoin`, `OuterJoin`). Hash-join implementation with typed fast paths for single int64/float64/string/bool/time keys and a length-prefixed tuple encoding for composite keys (the same collision-safe scheme GroupBy uses). Key columns appear once in the output under the left frame's name; colliding non-key columns are suffixed `_left`/`_right`. An int64 key may join a float64 key (matched in float64 space, output promoted to `Float64Type`); any other key-type mismatch is an error carried via `Error()`. A key equal to its column type's zero value is treated as missing and never matches, mirroring the CSV/JSONL missing-cell convention; unmatched sides zero-fill the same way. Output order is deterministic: left-frame row order (multi-matches expand in right-frame order), then unmatched right rows in right-frame order. A 100k×10k single-key join runs in ~3 ms with ~60 allocations. + +--- + ## [1.0.8] — 2026-07-16 ### Documentation diff --git a/README.md b/README.md index 063a8f8..f0aad30 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,34 @@ ranked := df.SortBy( ) ``` +### Joining DataFrames + +```go +// Combine two frames on a shared key column +result := orders.Join(customers, []string{"customer_id"}, otters.InnerJoin) + +// All four join types: InnerJoin, LeftJoin, RightJoin, OuterJoin +all := orders.Join(customers, []string{"customer_id"}, otters.LeftJoin) + +// Keys named differently in each frame +res := employees.JoinOn(departments, []string{"dept_code"}, []string{"code"}, otters.InnerJoin) + +// Composite keys join on the tuple of all key columns +byRegionYear := sales.Join(targets, []string{"region", "year"}, otters.InnerJoin) +``` + +Join semantics worth knowing: + +- Key columns appear **once** in the output (under the left frame's name). Non-key + columns present in both frames are suffixed `_left` / `_right`. +- An `int64` key can join a `float64` key — matching happens in float64 space and + the output key column is promoted to `float64`. Any other type mismatch errors. +- A key equal to its type's zero value (`""`, `0`, `0.0`, `false`, zero time) is + treated as **missing** and never matches — consistent with how CSV/JSONL fill + missing cells. Unmatched sides are zero-filled the same way. +- Output order is deterministic: left-frame row order, then (for right/outer) + unmatched right rows in right-frame order. + ### Aggregations and Statistics ```go @@ -157,9 +185,22 @@ fmt.Println(summary) ### Data Transformation ```go -// Create new columns -df_with_bonus := df.Copy() -// Add 10% bonus calculation (implementation coming soon) +// Derive a column from the whole row +withBonus := df.Apply("bonus", func(row map[string]any) any { + return row["salary"].(float64) * 0.10 +}) + +// Element-wise transform of a single column (newColumn == srcColumn replaces it) +upper := df.Map("name", "name_upper", func(v any) any { + return strings.ToUpper(v.(string)) +}) + +// Missing-value handling. Otters has no explicit null: the column type's +// zero value ("", 0, 0.0, false, zero time) is the missing-value sentinel, +// so a legitimate 0 in an int column is indistinguishable from missing. +filled := df.FillNA("department", "Unassigned") // replace zero values +complete := df.DropNA("salary", "department") // drop rows with zero values +allComplete := df.DropNA() // check every column // Rename columns clean_df := df.RenameColumn("hired_date", "start_date") @@ -168,6 +209,13 @@ clean_df := df.RenameColumn("hired_date", "start_date") essential := df.Drop("internal_id", "notes") ``` +The derived column's type is inferred from what the function returns: all +integers → `int64`, integer/float mix → `float64`, uniform string/bool/time +keep their type, any other mix falls back to `string`. Returning `nil` fills +that row with the zero value. A panic inside your function is caught and +carried as an error on the returned DataFrame — it will not crash your +program. + ## 🏗️ API Reference ### DataFrame Creation @@ -297,6 +345,7 @@ Coming from Pandas? Here's how Otters compares: | `df[df.age > 25]` | `df.Filter("age", ">", 25)` | Explicit syntax | | `df[['name', 'age']]` | `df.Select("name", "age")` | Method-based selection | | `df.sort_values()` | `df.Sort("column", true)` | Simple sort syntax | +| `df.merge(other)` | `df.Join(other, keys, how)` | Explicit join type | | `df.describe()` | `df.Describe()` | Similar functionality | ## 🚧 Roadmap @@ -311,11 +360,12 @@ Coming from Pandas? Here's how Otters compares: - [x] Simple query strings (`Query("age > 25")`) and `Where` - [x] Statistics (describe, median, variance, quantiles, correlation, value counts) - [x] Lazy views for chained operations (`df.Lazy()...Collect()`) +- [x] Join operations (inner, left, right, outer) — `Join` / `JoinOn` +- [x] Column derivation & transforms — `Apply`, `Map`, `FillNA`, `DropNA` - [x] Fluent API with error handling ### 🔄 Coming Soon -- [ ] Join operations (inner, left, right, outer) - [ ] More file formats (JSON arrays, Parquet) - [ ] Data visualization helpers - [ ] Streaming operations for large files diff --git a/err.go b/err.go index a08829e..a76de03 100644 --- a/err.go +++ b/err.go @@ -223,28 +223,34 @@ func (df *DataFrame) Error() error { return df.err } -// recoverFromPanic recovers from panics and converts them to OtterErrors -func recoverFromPanic(op string) error { +// panicToError converts a recovered panic value into an OtterError. +func panicToError(op string, r any) error { + switch v := r.(type) { + case error: + return wrapError(op, v) + case string: + return newOpError(op, v) + default: + return newOpError(op, fmt.Sprintf("panic: %v", r)) + } +} + +// recoverFromPanic converts an in-flight panic into an error assigned to +// *errp. It must be installed directly as the deferred function: +// +// defer recoverFromPanic("Op", &err) +// +// Calling it from inside another deferred closure does not work: recover +// only intercepts a panic when called directly by a deferred function. +func recoverFromPanic(op string, errp *error) { if r := recover(); r != nil { - switch v := r.(type) { - case error: - return wrapError(op, v) - case string: - return newOpError(op, v) - default: - return newOpError(op, fmt.Sprintf("panic: %v", r)) - } + *errp = panicToError(op, r) } - return nil } // SafeOperation wraps a function to handle panics and convert them to errors func SafeOperation(op string, fn func() error) (err error) { - defer func() { - if panicErr := recoverFromPanic(op); panicErr != nil { - err = panicErr - } - }() + defer recoverFromPanic(op, &err) return fn() } diff --git a/err_test.go b/err_test.go index db8291d..f50e115 100644 --- a/err_test.go +++ b/err_test.go @@ -129,10 +129,26 @@ func TestDataFrame_ErrorMethods(t *testing.T) { } func TestRecoverFromPanic(t *testing.T) { - // Test that recoverFromPanic works when there's no panic - err := recoverFromPanic("TestOp") - if err != nil { - t.Error("recoverFromPanic should return nil when no panic") + // No-panic path: the error must stay untouched. + run := func() (err error) { + defer recoverFromPanic("TestOp", &err) + return nil + } + if err := run(); err != nil { + t.Error("recoverFromPanic should leave err nil when no panic occurred") + } + + // Panic path: err is populated with an OtterError. + boom := func() (err error) { + defer recoverFromPanic("TestOp", &err) + panic("kaboom") + } + err := boom() + if err == nil { + t.Fatal("recoverFromPanic should convert the panic into an error") + } + if !strings.Contains(err.Error(), "kaboom") { + t.Errorf("recovered error should carry the panic message, got: %v", err) } } @@ -225,3 +241,68 @@ func TestSentinelErrors(t *testing.T) { }) }() } + +// TestSafeOperationRecoversPanics is the core "no panics escape" guarantee: +// a panic inside the wrapped function must come back as an error, not crash. +func TestSafeOperationRecoversPanics(t *testing.T) { + // String panic. + err := SafeOperation("StringPanic", func() error { panic("exploded") }) + if err == nil { + t.Fatal("SafeOperation must convert a string panic into an error") + } + if !strings.Contains(err.Error(), "exploded") { + t.Errorf("recovered error should mention the panic message, got: %v", err) + } + var oe *OtterError + if !errors.As(err, &oe) || oe.Op != "StringPanic" { + t.Errorf("recovered error should be an *OtterError carrying the op, got: %#v", err) + } + + // Error panic: the original error must remain reachable via errors.Is. + cause := errors.New("root cause") + err = SafeOperation("ErrorPanic", func() error { panic(cause) }) + if err == nil { + t.Fatal("SafeOperation must convert an error panic into an error") + } + if !errors.Is(err, cause) { + t.Errorf("recovered error should wrap the panicked error, got: %v", err) + } + + // Arbitrary value panic. + err = SafeOperation("IntPanic", func() error { panic(42) }) + if err == nil { + t.Fatal("SafeOperation must convert a non-error panic into an error") + } + if !strings.Contains(err.Error(), "42") { + t.Errorf("recovered error should mention the panic value, got: %v", err) + } +} + +// TestSafeOperationRecoversRuntimePanic drives a genuine runtime panic +// (index out of range) through the safety net. +func TestSafeOperationRecoversRuntimePanic(t *testing.T) { + err := SafeOperation("Runtime", func() error { + s := make([]int, 0) + i := len(s) - 1 + _ = s[i] // index out of range at runtime + return nil + }) + if err == nil { + t.Fatal("SafeOperation must recover a runtime panic") + } + if !strings.Contains(err.Error(), "range") { + t.Errorf("recovered error should carry the runtime message, got: %v", err) + } +} + +// TestSentinelMatchingSurvivesRecoveredPath guards the Round-2 sentinel fix: +// errors produced after a recovery still match via errors.Is. +func TestSentinelMatchingSurvivesRecoveredPath(t *testing.T) { + df, _ := NewDataFrameFromMap(map[string]any{"a": []int64{1}}) + err := SafeOperation("Wrapped", func() error { + return df.validateColumnExists("ghost") + }) + if !errors.Is(err, ErrColumnNotFound) { + t.Errorf("sentinel matching lost through SafeOperation, got: %v", err) + } +} diff --git a/example_join_test.go b/example_join_test.go new file mode 100644 index 0000000..adb1ad2 --- /dev/null +++ b/example_join_test.go @@ -0,0 +1,55 @@ +package otters + +import ( + "fmt" +) + +// Example_join demonstrates combining two DataFrames on a key column. +func Example_join() { + orderID, _ := NewSeries("order_id", []int64{1, 2, 3}) + orderCust, _ := NewSeries("customer_id", []int64{10, 20, 30}) + amount, _ := NewSeries("amount", []float64{250, 125.5, 300}) + orders, _ := NewDataFrameFromSeries(orderID, orderCust, amount) + + custID, _ := NewSeries("customer_id", []int64{10, 20, 40}) + name, _ := NewSeries("name", []string{"Alice", "Bob", "Carol"}) + customers, _ := NewDataFrameFromSeries(custID, name) + + // Inner join: only orders with a matching customer survive. + matched := orders.Join(customers, []string{"customer_id"}, InnerJoin) + if matched.Error() != nil { + fmt.Println("join failed:", matched.Error()) + return + } + fmt.Print(matched) + + // Left join keeps every order; unmatched rows zero-fill customer columns. + all := orders.Join(customers, []string{"customer_id"}, LeftJoin) + rows, cols := all.Shape() + fmt.Printf("left join keeps all orders: %d rows, %d columns\n", rows, cols) + + // Output: + // order_id customer_id amount name + // 1 10 250 Alice + // 2 20 125.5 Bob + // left join keeps all orders: 3 rows, 4 columns +} + +// ExampleDataFrame_JoinOn joins frames whose key columns are named differently. +func ExampleDataFrame_JoinOn() { + empID, _ := NewSeries("emp", []string{"ann", "bob"}) + deptRef, _ := NewSeries("dept_code", []string{"ENG", "OPS"}) + employees, _ := NewDataFrameFromSeries(empID, deptRef) + + code, _ := NewSeries("code", []string{"ENG", "OPS"}) + deptName, _ := NewSeries("dept_name", []string{"Engineering", "Operations"}) + departments, _ := NewDataFrameFromSeries(code, deptName) + + res := employees.JoinOn(departments, []string{"dept_code"}, []string{"code"}, InnerJoin) + fmt.Print(res) + + // Output: + // emp dept_code dept_name + // ann ENG Engineering + // bob OPS Operations +} diff --git a/example_transform_test.go b/example_transform_test.go new file mode 100644 index 0000000..b5df5bc --- /dev/null +++ b/example_transform_test.go @@ -0,0 +1,29 @@ +package otters + +import ( + "fmt" +) + +// Example_transform demonstrates deriving columns and handling missing values. +func Example_transform() { + name, _ := NewSeries("name", []string{"ann", "", "cy"}) + salary, _ := NewSeries("salary", []float64{1000, 2000, 0}) + df, _ := NewDataFrameFromSeries(name, salary) + + res := df. + FillNA("name", "unknown"). + DropNA("salary"). + Apply("bonus", func(row map[string]any) any { + return row["salary"].(float64) * 0.1 + }) + if res.Error() != nil { + fmt.Println("transform failed:", res.Error()) + return + } + fmt.Print(res) + + // Output: + // name salary bonus + // ann 1000 100 + // unknown 2000 200 +} diff --git a/join.go b/join.go new file mode 100644 index 0000000..77dafeb --- /dev/null +++ b/join.go @@ -0,0 +1,592 @@ +package otters + +import ( + "fmt" + "strconv" + "time" +) + +// JoinType specifies the kind of relational join to perform. +type JoinType int + +const ( + // InnerJoin keeps only rows whose key exists in both frames. + InnerJoin JoinType = iota + // LeftJoin keeps every left row, zero-filling right columns when unmatched. + LeftJoin + // RightJoin keeps every right row, zero-filling left columns when unmatched. + RightJoin + // OuterJoin keeps every row from both frames. + OuterJoin +) + +// String returns the string representation of a JoinType. +func (jt JoinType) String() string { + switch jt { + case InnerJoin: + return "inner" + case LeftJoin: + return "left" + case RightJoin: + return "right" + case OuterJoin: + return "outer" + default: + return "unknown" + } +} + +// Join combines two DataFrames on one or more shared key columns. +// +// Semantics: +// - Key columns must have the same type in both frames, except that an +// int64/float64 pairing is allowed: matching then happens in float64 +// space and the output key column is promoted to Float64Type. +// - A key equal to its column type's zero value ("", 0, 0.0, false, zero +// time) is treated as missing and never matches — such rows are dropped +// by an inner join and zero-filled on the outer side otherwise. For +// composite keys, one missing component makes the whole key missing. +// - Key columns appear once in the output. Non-key columns present in both +// frames are suffixed "_left" / "_right". +// - Row order is deterministic: left-frame row order first (a left row +// matching several right rows expands in right-frame order), then, for +// right/outer joins, unmatched right rows in right-frame order. +// - Unmatched sides are filled with the column type's zero value. +// +// Join never mutates its inputs; errors are carried on the returned frame. +func (df *DataFrame) Join(other *DataFrame, on []string, how JoinType) *DataFrame { + return df.JoinOn(other, on, on, how) +} + +// JoinOn joins two DataFrames whose key columns are named differently. +// leftKeys and rightKeys are matched positionally and must have equal length. +// The output contains each key column once, under its left-frame name; for +// right/outer joins, unmatched right rows contribute their key values into +// that column. See Join for the full semantics. +func (df *DataFrame) JoinOn(other *DataFrame, leftKeys, rightKeys []string, how JoinType) *DataFrame { + const op = "Join" + + if df.err != nil { + return df + } + if other == nil { + return df.setError(newOpError(op, "other DataFrame is nil")) + } + if other.err != nil { + return df.setError(wrapError(op, other.err)) + } + if how < InnerJoin || how > OuterJoin { + return df.setError(newOpError(op, fmt.Sprintf("invalid join type %d", int(how)))) + } + if len(leftKeys) == 0 { + return df.setError(newOpError(op, "at least one key column must be specified")) + } + if len(leftKeys) != len(rightKeys) { + return df.setError(newOpError(op, + fmt.Sprintf("leftKeys and rightKeys must have the same length (%d vs %d)", len(leftKeys), len(rightKeys)))) + } + if col, ok := firstDuplicate(leftKeys); ok { + return df.setError(newColumnError(op, col, "key column specified more than once")) + } + if col, ok := firstDuplicate(rightKeys); ok { + return df.setError(newColumnError(op, col, "key column specified more than once")) + } + if err := df.validateColumnsExist(leftKeys); err != nil { + return df.setError(err) + } + if err := other.validateColumnsExist(rightKeys); err != nil { + return df.setError(err) + } + + spaces, err := joinKeySpaces(df, other, leftKeys, rightKeys) + if err != nil { + return df.setError(err) + } + + leftOut, rightOut, err := joinOutputNames(df, other, leftKeys, rightKeys) + if err != nil { + return df.setError(err) + } + + leftIdx, rightIdx := joinRowPairs(df, other, leftKeys, rightKeys, spaces, how) + + return buildJoinResult(df, other, leftKeys, rightKeys, spaces, leftOut, rightOut, leftIdx, rightIdx) +} + +// joinSpace is the unified comparison space for one key column pair. +type joinSpace int + +const ( + spaceInt joinSpace = iota + spaceFloat + spaceString + spaceBool + spaceTime +) + +// joinKeySpaces validates key type compatibility and returns the comparison +// space for each key pair. Identical types map to their own space; an +// int64/float64 mix is compared (and emitted) in float64 space; any other +// combination is an error. +func joinKeySpaces(left, right *DataFrame, leftKeys, rightKeys []string) ([]joinSpace, error) { + spaces := make([]joinSpace, len(leftKeys)) + for k := range leftKeys { + lt := left.columns[leftKeys[k]].Type + rt := right.columns[rightKeys[k]].Type + + switch { + case lt == rt: + switch lt { + case Int64Type: + spaces[k] = spaceInt + case Float64Type: + spaces[k] = spaceFloat + case StringType: + spaces[k] = spaceString + case BoolType: + spaces[k] = spaceBool + case TimeType: + spaces[k] = spaceTime + default: + return nil, newColumnError("Join", leftKeys[k], "unsupported key column type") + } + case (lt == Int64Type && rt == Float64Type) || (lt == Float64Type && rt == Int64Type): + spaces[k] = spaceFloat + default: + return nil, newColumnError("Join", leftKeys[k], + fmt.Sprintf("key type mismatch: left %q is %s, right %q is %s", + leftKeys[k], lt, rightKeys[k], rt)) + } + } + return spaces, nil +} + +// joinOutputNames computes the output column names. leftOut is indexed like +// left.order; rightOut like right.order, with "" marking right key columns +// (which are dropped — the key appears once under its left name). Key columns +// always keep their name; colliding non-key columns get "_left" / "_right". +func joinOutputNames(left, right *DataFrame, leftKeys, rightKeys []string) (leftOut, rightOut []string, err error) { + leftKeySet := make(map[string]bool, len(leftKeys)) + for _, k := range leftKeys { + leftKeySet[k] = true + } + rightKeySet := make(map[string]bool, len(rightKeys)) + for _, k := range rightKeys { + rightKeySet[k] = true + } + + leftNameSet := make(map[string]bool, len(left.order)) + for _, c := range left.order { + leftNameSet[c] = true + } + rightNonKeySet := make(map[string]bool, len(right.order)) + for _, c := range right.order { + if !rightKeySet[c] { + rightNonKeySet[c] = true + } + } + + leftOut = make([]string, len(left.order)) + for i, c := range left.order { + switch { + case leftKeySet[c]: + leftOut[i] = c + case rightNonKeySet[c]: + leftOut[i] = c + "_left" + default: + leftOut[i] = c + } + } + + rightOut = make([]string, len(right.order)) + for i, c := range right.order { + switch { + case rightKeySet[c]: + rightOut[i] = "" // dropped: key emitted once under its left name + case leftNameSet[c]: + rightOut[i] = c + "_right" + default: + rightOut[i] = c + } + } + + seen := make(map[string]bool, len(leftOut)+len(rightOut)) + for _, name := range leftOut { + if seen[name] { + return nil, nil, newColumnError("Join", name, "output column name collision") + } + seen[name] = true + } + for _, name := range rightOut { + if name == "" { + continue + } + if seen[name] { + return nil, nil, newColumnError("Join", name, "output column name collision") + } + seen[name] = true + } + return leftOut, rightOut, nil +} + +// joinRowPairs computes the output row mapping as parallel slices of left and +// right source row indices; -1 marks the zero-filled side of an unmatched row. +func joinRowPairs(left, right *DataFrame, leftKeys, rightKeys []string, spaces []joinSpace, how JoinType) ([]int, []int) { + if len(leftKeys) == 1 { + ls := left.columns[leftKeys[0]] + rs := right.columns[rightKeys[0]] + switch spaces[0] { + case spaceInt: + lk := ls.Data.([]int64) + rk := rs.Data.([]int64) + return hashJoinPairs(lk, nullMaskInt64(lk), rk, nullMaskInt64(rk), how) + case spaceFloat: + lk := asFloat64Keys(ls) + rk := asFloat64Keys(rs) + return hashJoinPairs(lk, nullMaskFloat64(lk), rk, nullMaskFloat64(rk), how) + case spaceString: + lk := ls.Data.([]string) + rk := rs.Data.([]string) + return hashJoinPairs(lk, nullMaskString(lk), rk, nullMaskString(rk), how) + case spaceBool: + lk := ls.Data.([]bool) + rk := rs.Data.([]bool) + return hashJoinPairs(lk, nullMaskBool(lk), rk, nullMaskBool(rk), how) + case spaceTime: + lt := ls.Data.([]time.Time) + rt := rs.Data.([]time.Time) + return hashJoinPairs(timeKeys(lt), nullMaskTime(lt), timeKeys(rt), nullMaskTime(rt), how) + } + } + + // Composite keys: encode each row's key tuple into a string. + lk, lnull := encodeCompositeKeys(left, leftKeys, spaces) + rk, rnull := encodeCompositeKeys(right, rightKeys, spaces) + return hashJoinPairs(lk, lnull, rk, rnull, how) +} + +// hashJoinPairs is the join core: it indexes the right side's non-null keys, +// probes with each left row in order, and appends unmatched right rows for +// right/outer joins. Null keys never match. +func hashJoinPairs[K comparable](lk []K, lnull []bool, rk []K, rnull []bool, how JoinType) ([]int, []int) { + // Chained hash index: head maps a key to its first right row, next links + // rows sharing that key. Building in reverse makes each chain run in + // right-frame order, and the flat next slice avoids allocating a small + // row slice per distinct key. + head := make(map[K]int, len(rk)) + next := make([]int, len(rk)) + for j := len(rk) - 1; j >= 0; j-- { + if rnull[j] { + continue + } + if h, ok := head[rk[j]]; ok { + next[j] = h + } else { + next[j] = -1 + } + head[rk[j]] = j + } + + var matched []bool + if how == RightJoin || how == OuterJoin { + matched = make([]bool, len(rk)) + } + + leftIdx := make([]int, 0, len(lk)) + rightIdx := make([]int, 0, len(lk)) + + for i, key := range lk { + j := -1 + if !lnull[i] { + if h, ok := head[key]; ok { + j = h + } + } + if j < 0 { + if how == LeftJoin || how == OuterJoin { + leftIdx = append(leftIdx, i) + rightIdx = append(rightIdx, -1) + } + continue + } + for ; j >= 0; j = next[j] { + leftIdx = append(leftIdx, i) + rightIdx = append(rightIdx, j) + if matched != nil { + matched[j] = true + } + } + } + + if matched != nil { + for j := range rk { + if !matched[j] { + leftIdx = append(leftIdx, -1) + rightIdx = append(rightIdx, j) + } + } + } + + return leftIdx, rightIdx +} + +// Null-key masks: the zero value of the key column type marks a missing key. + +func nullMaskInt64(data []int64) []bool { + mask := make([]bool, len(data)) + for i, v := range data { + mask[i] = v == 0 + } + return mask +} + +func nullMaskFloat64(data []float64) []bool { + mask := make([]bool, len(data)) + for i, v := range data { + mask[i] = v == 0 + } + return mask +} + +func nullMaskString(data []string) []bool { + mask := make([]bool, len(data)) + for i, v := range data { + mask[i] = v == "" + } + return mask +} + +func nullMaskBool(data []bool) []bool { + mask := make([]bool, len(data)) + for i, v := range data { + mask[i] = !v + } + return mask +} + +func nullMaskTime(data []time.Time) []bool { + mask := make([]bool, len(data)) + for i, v := range data { + mask[i] = v.IsZero() + } + return mask +} + +// timeKeys maps times to comparable int64 keys. UnixNano compares instants, +// so the same moment in different locations hashes identically (matching +// time.Time.Equal semantics used elsewhere in the library). +func timeKeys(data []time.Time) []int64 { + keys := make([]int64, len(data)) + for i, v := range data { + keys[i] = v.UnixNano() + } + return keys +} + +// asFloat64Keys returns the series data as a float64 slice, converting from +// int64 when the key pair mixes numeric types. Float64 series are returned +// as-is (read-only use). +func asFloat64Keys(s *Series) []float64 { + if s.Type == Float64Type { + return s.Data.([]float64) + } + data := s.Data.([]int64) + out := make([]float64, len(data)) + for i, v := range data { + out[i] = float64(v) + } + return out +} + +// encodeCompositeKeys builds one string key per row from multiple key +// columns. Components are rendered in the pair's comparison space, separated +// by a zero byte; string components are length-prefixed so tuple boundaries +// stay unambiguous (("ab","c") never equals ("a","bc")). A row is null if any +// component is null. +func encodeCompositeKeys(df *DataFrame, keys []string, spaces []joinSpace) ([]string, []bool) { + n := df.length + out := make([]string, n) + null := make([]bool, n) + + series := make([]*Series, len(keys)) + for k, name := range keys { + series[k] = df.columns[name] + } + + buf := make([]byte, 0, 64) + for i := 0; i < n; i++ { + buf = buf[:0] + rowNull := false + for k, s := range series { + if k > 0 { + buf = append(buf, 0) + } + var componentNull bool + buf, componentNull = appendKeyComponent(buf, s, i, spaces[k]) + rowNull = rowNull || componentNull + } + null[i] = rowNull + if !rowNull { + out[i] = string(buf) + } + } + return out, null +} + +// appendKeyComponent renders one key cell into buf in the given comparison +// space and reports whether the cell is the null sentinel. +func appendKeyComponent(buf []byte, s *Series, row int, space joinSpace) ([]byte, bool) { + switch space { + case spaceInt: + v := s.Data.([]int64)[row] + return strconv.AppendInt(buf, v, 10), v == 0 + case spaceFloat: + var v float64 + if s.Type == Int64Type { + v = float64(s.Data.([]int64)[row]) + } else { + v = s.Data.([]float64)[row] + } + return strconv.AppendFloat(buf, v, 'g', -1, 64), v == 0 + case spaceString: + v := s.Data.([]string)[row] + buf = strconv.AppendInt(buf, int64(len(v)), 10) + buf = append(buf, ':') + return append(buf, v...), v == "" + case spaceBool: + v := s.Data.([]bool)[row] + if v { + return append(buf, 't'), false + } + return append(buf, 'f'), true + case spaceTime: + v := s.Data.([]time.Time)[row] + return strconv.AppendInt(buf, v.UnixNano(), 10), v.IsZero() + default: + return buf, true + } +} + +// buildJoinResult materializes the output frame from the row-pair mapping. +func buildJoinResult(left, right *DataFrame, leftKeys, rightKeys []string, spaces []joinSpace, + leftOut, rightOut []string, leftIdx, rightIdx []int) *DataFrame { + + leftKeyPos := make(map[string]int, len(leftKeys)) + for k, name := range leftKeys { + leftKeyPos[name] = k + } + + newDf := NewDataFrame() + newDf.length = len(leftIdx) + + for i, col := range left.order { + var data any + if k, isKey := leftKeyPos[col]; isKey { + data = buildKeyColumnData(left.columns[col], right.columns[rightKeys[k]], spaces[k], leftIdx, rightIdx) + } else { + data = gatherSeriesWithFill(left.columns[col], leftIdx) + } + series, err := newSeriesOwned(leftOut[i], data) + if err != nil { + return left.setError(wrapColumnError("Join", col, err)) + } + newDf.addSeriesUnsafe(series) + } + + for i, col := range right.order { + if rightOut[i] == "" { + continue // right key column: values already emitted under the left key name + } + series, err := newSeriesOwned(rightOut[i], gatherSeriesWithFill(right.columns[col], rightIdx)) + if err != nil { + return left.setError(wrapColumnError("Join", col, err)) + } + newDf.addSeriesUnsafe(series) + } + + return newDf +} + +// buildKeyColumnData assembles a key column: values come from the left row +// when present, otherwise from the right row (right/outer unmatched rows). +// A mixed int64/float64 key pair is emitted as float64. +func buildKeyColumnData(ls, rs *Series, space joinSpace, leftIdx, rightIdx []int) any { + if space == spaceFloat && (ls.Type != Float64Type || rs.Type != Float64Type) { + return gatherKeyColumn(asFloat64Keys(ls), asFloat64Keys(rs), leftIdx, rightIdx, 0) + } + + switch ls.Type { + case StringType: + return gatherKeyColumn(ls.Data.([]string), rs.Data.([]string), leftIdx, rightIdx, "") + case Int64Type: + return gatherKeyColumn(ls.Data.([]int64), rs.Data.([]int64), leftIdx, rightIdx, 0) + case Float64Type: + return gatherKeyColumn(ls.Data.([]float64), rs.Data.([]float64), leftIdx, rightIdx, 0) + case BoolType: + return gatherKeyColumn(ls.Data.([]bool), rs.Data.([]bool), leftIdx, rightIdx, false) + case TimeType: + return gatherKeyColumn(ls.Data.([]time.Time), rs.Data.([]time.Time), leftIdx, rightIdx, time.Time{}) + default: + return nil + } +} + +// gatherKeyColumn picks each output row's key from the left source when its +// index is present, falling back to the right source, then to the zero value. +func gatherKeyColumn[T any](l, r []T, leftIdx, rightIdx []int, zero T) []T { + out := make([]T, len(leftIdx)) + for i := range leftIdx { + switch { + case leftIdx[i] >= 0: + out[i] = l[leftIdx[i]] + case rightIdx[i] >= 0: + out[i] = r[rightIdx[i]] + default: + out[i] = zero + } + } + return out +} + +// gatherSeriesWithFill extracts rows at idx from a series; -1 produces the +// column type's zero value. +func gatherSeriesWithFill(s *Series, idx []int) any { + switch s.Type { + case StringType: + return gatherWithFill(s.Data.([]string), idx, "") + case Int64Type: + return gatherWithFill(s.Data.([]int64), idx, 0) + case Float64Type: + return gatherWithFill(s.Data.([]float64), idx, 0) + case BoolType: + return gatherWithFill(s.Data.([]bool), idx, false) + case TimeType: + return gatherWithFill(s.Data.([]time.Time), idx, time.Time{}) + default: + return nil + } +} + +func gatherWithFill[T any](data []T, idx []int, zero T) []T { + out := make([]T, len(idx)) + for i, j := range idx { + if j >= 0 { + out[i] = data[j] + } else { + out[i] = zero + } + } + return out +} + +// firstDuplicate returns the first name appearing more than once. +func firstDuplicate(names []string) (string, bool) { + seen := make(map[string]bool, len(names)) + for _, n := range names { + if seen[n] { + return n, true + } + seen[n] = true + } + return "", false +} diff --git a/join_bench_test.go b/join_bench_test.go new file mode 100644 index 0000000..2b8aef9 --- /dev/null +++ b/join_bench_test.go @@ -0,0 +1,135 @@ +package otters + +import ( + "strconv" + "testing" +) + +func benchFrames(b *testing.B, nLeft, nRight int) (*DataFrame, *DataFrame) { + b.Helper() + lk := make([]int64, nLeft) + lv := make([]float64, nLeft) + ls := make([]string, nLeft) + for i := range lk { + lk[i] = int64(i%nRight) + 1 + lv[i] = float64(i) + ls[i] = "left-" + strconv.Itoa(i%1000) + } + rk := make([]int64, nRight) + rn := make([]string, nRight) + for j := range rk { + rk[j] = int64(j) + 1 + rn[j] = "name-" + strconv.Itoa(j) + } + + mustSeries := func(name string, data any) *Series { + s, err := NewSeries(name, data) + if err != nil { + b.Fatal(err) + } + return s + } + left, err := NewDataFrameFromSeries( + mustSeries("id", lk), mustSeries("value", lv), mustSeries("tag", ls)) + if err != nil { + b.Fatal(err) + } + right, err := NewDataFrameFromSeries( + mustSeries("id", rk), mustSeries("name", rn)) + if err != nil { + b.Fatal(err) + } + return left, right +} + +func BenchmarkJoinInnerInt64_100k_x_10k(b *testing.B) { + left, right := benchFrames(b, 100_000, 10_000) + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + res := left.Join(right, []string{"id"}, InnerJoin) + if res.Error() != nil { + b.Fatal(res.Error()) + } + } +} + +func BenchmarkJoinLeftInt64_100k_x_10k(b *testing.B) { + left, right := benchFrames(b, 100_000, 10_000) + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + res := left.Join(right, []string{"id"}, LeftJoin) + if res.Error() != nil { + b.Fatal(res.Error()) + } + } +} + +func BenchmarkJoinInnerString_100k_x_10k(b *testing.B) { + nLeft, nRight := 100_000, 10_000 + lk := make([]string, nLeft) + lv := make([]int64, nLeft) + for i := range lk { + lk[i] = "key-" + strconv.Itoa(i%nRight) + lv[i] = int64(i) + } + rk := make([]string, nRight) + rv := make([]int64, nRight) + for j := range rk { + rk[j] = "key-" + strconv.Itoa(j) + rv[j] = int64(j) + } + mustSeries := func(name string, data any) *Series { + s, err := NewSeries(name, data) + if err != nil { + b.Fatal(err) + } + return s + } + left, _ := NewDataFrameFromSeries(mustSeries("k", lk), mustSeries("v", lv)) + right, _ := NewDataFrameFromSeries(mustSeries("k", rk), mustSeries("w", rv)) + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + res := left.Join(right, []string{"k"}, InnerJoin) + if res.Error() != nil { + b.Fatal(res.Error()) + } + } +} + +func BenchmarkJoinInnerComposite_100k_x_10k(b *testing.B) { + nLeft, nRight := 100_000, 10_000 + lk1 := make([]string, nLeft) + lk2 := make([]int64, nLeft) + for i := range lk1 { + lk1[i] = "grp-" + strconv.Itoa(i%100) + lk2[i] = int64(i % 100) + } + // Right keys are unique; the first 100 line up with the left key domain, + // so every left row matches exactly one right row (100k output rows). + rk1 := make([]string, nRight) + rk2 := make([]int64, nRight) + for j := range rk1 { + rk1[j] = "grp-" + strconv.Itoa(j%100) + rk2[j] = int64(j) + } + mustSeries := func(name string, data any) *Series { + s, err := NewSeries(name, data) + if err != nil { + b.Fatal(err) + } + return s + } + left, _ := NewDataFrameFromSeries(mustSeries("g", lk1), mustSeries("n", lk2)) + right, _ := NewDataFrameFromSeries(mustSeries("g", rk1), mustSeries("n", rk2)) + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + res := left.Join(right, []string{"g", "n"}, InnerJoin) + if res.Error() != nil { + b.Fatal(res.Error()) + } + } +} diff --git a/join_test.go b/join_test.go new file mode 100644 index 0000000..cb20926 --- /dev/null +++ b/join_test.go @@ -0,0 +1,717 @@ +package otters + +import ( + "errors" + "reflect" + "testing" + "time" +) + +// --- test helpers --- + +// dfFromSeries builds a DataFrame from (name, data) pairs preserving column order. +func dfFromSeries(t *testing.T, pairs ...any) *DataFrame { + t.Helper() + if len(pairs)%2 != 0 { + t.Fatal("dfFromSeries: pairs must be name, data alternating") + } + var series []*Series + for i := 0; i < len(pairs); i += 2 { + name, ok := pairs[i].(string) + if !ok { + t.Fatalf("dfFromSeries: pair %d name is %T, want string", i, pairs[i]) + } + s, err := NewSeries(name, pairs[i+1]) + if err != nil { + t.Fatalf("dfFromSeries: NewSeries(%s): %v", name, err) + } + series = append(series, s) + } + df, err := NewDataFrameFromSeries(series...) + if err != nil { + t.Fatalf("dfFromSeries: %v", err) + } + return df +} + +// checkColumn asserts a column's full contents match want (a typed slice). +func checkColumn(t *testing.T, df *DataFrame, col string, want any) { + t.Helper() + if df.Error() != nil { + t.Fatalf("checkColumn(%s): DataFrame has error: %v", col, df.Error()) + } + s, err := df.GetSeries(col) + if err != nil { + t.Fatalf("checkColumn(%s): %v", col, err) + } + if !reflect.DeepEqual(s.Data, want) { + t.Errorf("column %s = %#v, want %#v", col, s.Data, want) + } +} + +// checkColumns asserts the exact column names and order. +func checkColumns(t *testing.T, df *DataFrame, want ...string) { + t.Helper() + if df.Error() != nil { + t.Fatalf("checkColumns: DataFrame has error: %v", df.Error()) + } + got := df.Columns() + if !reflect.DeepEqual(got, want) { + t.Errorf("columns = %v, want %v", got, want) + } +} + +// ordersCustomers returns the canonical left/right test frames. +// left "orders": order_id, customer_id (30 unmatched, 0 null), amount +// right "customers": customer_id (40 unmatched, 0 null), name +func ordersCustomers(t *testing.T) (*DataFrame, *DataFrame) { + t.Helper() + left := dfFromSeries(t, + "order_id", []int64{1, 2, 3, 4, 5}, + "customer_id", []int64{10, 20, 30, 10, 0}, + "amount", []float64{100, 200, 300, 400, 500}, + ) + right := dfFromSeries(t, + "customer_id", []int64{10, 20, 40, 0}, + "name", []string{"Alice", "Bob", "Carol", "NullCo"}, + ) + return left, right +} + +// --- JoinType --- + +func TestJoinTypeString(t *testing.T) { + cases := map[JoinType]string{ + InnerJoin: "inner", + LeftJoin: "left", + RightJoin: "right", + OuterJoin: "outer", + JoinType(999): "unknown", + } + for jt, want := range cases { + if got := jt.String(); got != want { + t.Errorf("JoinType(%d).String() = %q, want %q", int(jt), got, want) + } + } +} + +// --- basic join semantics --- + +func TestInnerJoinBasic(t *testing.T) { + left, right := ordersCustomers(t) + res := left.Join(right, []string{"customer_id"}, InnerJoin) + if res.Error() != nil { + t.Fatalf("inner join error: %v", res.Error()) + } + + // Left rows in order: 1(10→Alice), 2(20→Bob), 3(30 no match), 4(10→Alice), 5(0 null key skipped) + checkColumns(t, res, "order_id", "customer_id", "amount", "name") + checkColumn(t, res, "order_id", []int64{1, 2, 4}) + checkColumn(t, res, "customer_id", []int64{10, 20, 10}) + checkColumn(t, res, "amount", []float64{100, 200, 400}) + checkColumn(t, res, "name", []string{"Alice", "Bob", "Alice"}) +} + +func TestLeftJoinBasic(t *testing.T) { + left, right := ordersCustomers(t) + res := left.Join(right, []string{"customer_id"}, LeftJoin) + if res.Error() != nil { + t.Fatalf("left join error: %v", res.Error()) + } + + // All 5 left rows preserved in order; unmatched (30) and null-key (0) rows zero-filled on right. + checkColumn(t, res, "order_id", []int64{1, 2, 3, 4, 5}) + checkColumn(t, res, "customer_id", []int64{10, 20, 30, 10, 0}) + checkColumn(t, res, "amount", []float64{100, 200, 300, 400, 500}) + checkColumn(t, res, "name", []string{"Alice", "Bob", "", "Alice", ""}) +} + +func TestRightJoinBasic(t *testing.T) { + left, right := ordersCustomers(t) + res := left.Join(right, []string{"customer_id"}, RightJoin) + if res.Error() != nil { + t.Fatalf("right join error: %v", res.Error()) + } + + // Matched pairs in left-row order, then unmatched right rows (40, null 0) in right order. + checkColumn(t, res, "order_id", []int64{1, 2, 4, 0, 0}) + checkColumn(t, res, "customer_id", []int64{10, 20, 10, 40, 0}) + checkColumn(t, res, "amount", []float64{100, 200, 400, 0, 0}) + checkColumn(t, res, "name", []string{"Alice", "Bob", "Alice", "Carol", "NullCo"}) +} + +func TestOuterJoinBasic(t *testing.T) { + left, right := ordersCustomers(t) + res := left.Join(right, []string{"customer_id"}, OuterJoin) + if res.Error() != nil { + t.Fatalf("outer join error: %v", res.Error()) + } + + // All left rows in order (incl. unmatched 30 and null 0), then unmatched right rows. + // Null keys never match: left row with key 0 and right row with key 0 stay separate. + checkColumn(t, res, "order_id", []int64{1, 2, 3, 4, 5, 0, 0}) + checkColumn(t, res, "customer_id", []int64{10, 20, 30, 10, 0, 40, 0}) + checkColumn(t, res, "amount", []float64{100, 200, 300, 400, 500, 0, 0}) + checkColumn(t, res, "name", []string{"Alice", "Bob", "", "Alice", "", "Carol", "NullCo"}) +} + +// --- duplicate keys / row multiplication --- + +func TestJoinDuplicateRightKeys(t *testing.T) { + left := dfFromSeries(t, + "id", []int64{1, 2}, + "l", []string{"L1", "L2"}, + ) + right := dfFromSeries(t, + "id", []int64{1, 1, 3}, + "r", []string{"a", "b", "c"}, + ) + res := left.Join(right, []string{"id"}, InnerJoin) + // Left row 1 matches two right rows, emitted in right-frame order. + checkColumn(t, res, "id", []int64{1, 1}) + checkColumn(t, res, "l", []string{"L1", "L1"}) + checkColumn(t, res, "r", []string{"a", "b"}) +} + +func TestJoinDuplicateBothSides(t *testing.T) { + left := dfFromSeries(t, "id", []int64{7, 7}, "l", []string{"x", "y"}) + right := dfFromSeries(t, "id", []int64{7, 7}, "r", []string{"p", "q"}) + res := left.Join(right, []string{"id"}, InnerJoin) + // Cartesian per key: 4 rows, left-major order. + checkColumn(t, res, "l", []string{"x", "x", "y", "y"}) + checkColumn(t, res, "r", []string{"p", "q", "p", "q"}) +} + +// --- composite keys --- + +func TestJoinCompositeKeys(t *testing.T) { + left := dfFromSeries(t, + "region", []string{"East", "East", "West", "West"}, + "year", []int64{2024, 2025, 2024, 2025}, + "sales", []float64{10, 20, 30, 40}, + ) + right := dfFromSeries(t, + "region", []string{"East", "West", "East"}, + "year", []int64{2024, 2025, 2030}, + "target", []float64{11, 44, 99}, + ) + res := left.Join(right, []string{"region", "year"}, InnerJoin) + checkColumns(t, res, "region", "year", "sales", "target") + checkColumn(t, res, "region", []string{"East", "West"}) + checkColumn(t, res, "year", []int64{2024, 2025}) + checkColumn(t, res, "sales", []float64{10, 40}) + checkColumn(t, res, "target", []float64{11, 44}) +} + +func TestJoinCompositeKeysNoCrossTupleConfusion(t *testing.T) { + // Keys ("ab","c") and ("a","bc") must not collide in the tuple encoding. + left := dfFromSeries(t, "k1", []string{"ab"}, "k2", []string{"c"}, "v", []int64{1}) + right := dfFromSeries(t, "k1", []string{"a"}, "k2", []string{"bc"}, "w", []int64{2}) + res := left.Join(right, []string{"k1", "k2"}, InnerJoin) + if res.Error() != nil { + t.Fatalf("join error: %v", res.Error()) + } + if res.Len() != 0 { + t.Errorf("expected 0 rows, got %d — tuple keys collided", res.Len()) + } +} + +func TestJoinCompositeNullComponent(t *testing.T) { + // A key tuple with any null (zero-value) component never matches. + left := dfFromSeries(t, "a", []string{"x", ""}, "b", []int64{1, 2}, "v", []int64{10, 20}) + right := dfFromSeries(t, "a", []string{"x", ""}, "b", []int64{1, 2}, "w", []int64{30, 40}) + res := left.Join(right, []string{"a", "b"}, InnerJoin) + checkColumn(t, res, "v", []int64{10}) + checkColumn(t, res, "w", []int64{30}) +} + +// --- JoinOn with differently-named keys --- + +func TestJoinOnDifferentKeyNames(t *testing.T) { + left := dfFromSeries(t, + "cust", []int64{10, 20, 30}, + "amount", []float64{1, 2, 3}, + ) + right := dfFromSeries(t, + "id", []int64{10, 30}, + "name", []string{"Alice", "Carl"}, + ) + res := left.JoinOn(right, []string{"cust"}, []string{"id"}, InnerJoin) + // Key column appears once, under the left name. + checkColumns(t, res, "cust", "amount", "name") + checkColumn(t, res, "cust", []int64{10, 30}) + checkColumn(t, res, "name", []string{"Alice", "Carl"}) + if res.HasColumn("id") { + t.Error("right key column should not appear in output") + } +} + +func TestJoinOnRightOuterKeyValuesComeFromRight(t *testing.T) { + left := dfFromSeries(t, "cust", []int64{10}, "amount", []float64{1}) + right := dfFromSeries(t, "id", []int64{10, 55}, "name", []string{"Alice", "Eve"}) + res := left.JoinOn(right, []string{"cust"}, []string{"id"}, RightJoin) + // Unmatched right row contributes its key value into the left-named key column. + checkColumn(t, res, "cust", []int64{10, 55}) + checkColumn(t, res, "amount", []float64{1, 0}) + checkColumn(t, res, "name", []string{"Alice", "Eve"}) +} + +// --- column name collisions --- + +func TestJoinNonKeyCollisionSuffixes(t *testing.T) { + left := dfFromSeries(t, "id", []int64{1, 2}, "amount", []float64{10, 20}) + right := dfFromSeries(t, "id", []int64{1, 2}, "amount", []float64{100, 200}) + res := left.Join(right, []string{"id"}, InnerJoin) + checkColumns(t, res, "id", "amount_left", "amount_right") + checkColumn(t, res, "id", []int64{1, 2}) + checkColumn(t, res, "amount_left", []float64{10, 20}) + checkColumn(t, res, "amount_right", []float64{100, 200}) +} + +func TestJoinRightNonKeyCollidesWithLeftKeyName(t *testing.T) { + // Right has a non-key column named like the left key: key keeps its name, + // only the right column is suffixed. + left := dfFromSeries(t, "id", []int64{1}, "v", []int64{9}) + right := dfFromSeries(t, "rid", []int64{1}, "id", []int64{777}) + res := left.JoinOn(right, []string{"id"}, []string{"rid"}, InnerJoin) + checkColumns(t, res, "id", "v", "id_right") + checkColumn(t, res, "id", []int64{1}) + checkColumn(t, res, "id_right", []int64{777}) +} + +func TestJoinSuffixCollisionErrors(t *testing.T) { + // Suffixing "amount" would collide with a pre-existing "amount_left". + left := dfFromSeries(t, "id", []int64{1}, "amount", []float64{1}, "amount_left", []float64{2}) + right := dfFromSeries(t, "id", []int64{1}, "amount", []float64{3}) + res := left.Join(right, []string{"id"}, InnerJoin) + if res.Error() == nil { + t.Fatal("expected error on suffix collision, got none") + } +} + +func TestJoinSelfJoin(t *testing.T) { + df := dfFromSeries(t, "id", []int64{1, 2}, "v", []int64{10, 20}) + res := df.Join(df, []string{"id"}, InnerJoin) + checkColumns(t, res, "id", "v_left", "v_right") + checkColumn(t, res, "v_left", []int64{10, 20}) + checkColumn(t, res, "v_right", []int64{10, 20}) +} + +// --- key type handling --- + +func TestJoinMixedNumericKeysCoerceToFloat64(t *testing.T) { + left := dfFromSeries(t, "k", []int64{1, 2, 3}, "l", []string{"a", "b", "c"}) + right := dfFromSeries(t, "k", []float64{1.0, 2.5, 3.0}, "r", []string{"x", "y", "z"}) + res := left.Join(right, []string{"k"}, InnerJoin) + if res.Error() != nil { + t.Fatalf("mixed numeric join error: %v", res.Error()) + } + // int64 1 and 3 match float64 1.0 and 3.0; 2 vs 2.5 does not. + // Output key column is promoted to Float64Type. + ct, err := res.GetColumnType("k") + if err != nil { + t.Fatal(err) + } + if ct != Float64Type { + t.Errorf("mixed numeric key column type = %v, want Float64Type", ct) + } + checkColumn(t, res, "k", []float64{1, 3}) + checkColumn(t, res, "l", []string{"a", "c"}) + checkColumn(t, res, "r", []string{"x", "z"}) +} + +func TestJoinMixedNumericOuterFillFromRight(t *testing.T) { + left := dfFromSeries(t, "k", []int64{1}, "l", []string{"a"}) + right := dfFromSeries(t, "k", []float64{2.5}, "r", []string{"y"}) + res := left.Join(right, []string{"k"}, OuterJoin) + checkColumn(t, res, "k", []float64{1, 2.5}) + checkColumn(t, res, "l", []string{"a", ""}) + checkColumn(t, res, "r", []string{"", "y"}) +} + +func TestJoinIncompatibleKeyTypesError(t *testing.T) { + left := dfFromSeries(t, "k", []int64{1}, "l", []string{"a"}) + right := dfFromSeries(t, "k", []string{"1"}, "r", []string{"x"}) + res := left.Join(right, []string{"k"}, InnerJoin) + if res.Error() == nil { + t.Fatal("expected key type mismatch error, got none") + } +} + +func TestJoinStringKeys(t *testing.T) { + left := dfFromSeries(t, "name", []string{"a", "b", ""}, "v", []int64{1, 2, 3}) + right := dfFromSeries(t, "name", []string{"b", "", "a"}, "w", []int64{20, 30, 10}) + res := left.Join(right, []string{"name"}, InnerJoin) + // "" is the string null sentinel: never matches. + checkColumn(t, res, "name", []string{"a", "b"}) + checkColumn(t, res, "v", []int64{1, 2}) + checkColumn(t, res, "w", []int64{10, 20}) +} + +func TestJoinTimeKeys(t *testing.T) { + t1 := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + t2 := time.Date(2026, 6, 15, 12, 0, 0, 0, time.UTC) + var zero time.Time + left := dfFromSeries(t, "ts", []time.Time{t1, t2, zero}, "v", []int64{1, 2, 3}) + right := dfFromSeries(t, "ts", []time.Time{t2, zero}, "w", []int64{20, 30}) + res := left.Join(right, []string{"ts"}, InnerJoin) + // Zero time is the null sentinel: never matches. + checkColumn(t, res, "v", []int64{2}) + checkColumn(t, res, "w", []int64{20}) +} + +func TestJoinTimeKeysDifferentLocationsSameInstant(t *testing.T) { + utc := time.Date(2026, 3, 1, 12, 0, 0, 0, time.UTC) + est := utc.In(time.FixedZone("EST", -5*3600)) + left := dfFromSeries(t, "ts", []time.Time{utc}, "v", []int64{1}) + right := dfFromSeries(t, "ts", []time.Time{est}, "w", []int64{2}) + res := left.Join(right, []string{"ts"}, InnerJoin) + // Same instant in different zones must match (time.Time.Equal semantics). + if res.Len() != 1 { + t.Errorf("same-instant times in different zones should join, got %d rows", res.Len()) + } +} + +func TestJoinFloat64Keys(t *testing.T) { + left := dfFromSeries(t, "k", []float64{1.5, 2.5, 0}, "v", []int64{1, 2, 3}) + right := dfFromSeries(t, "k", []float64{2.5, 0}, "w", []int64{20, 30}) + res := left.Join(right, []string{"k"}, InnerJoin) + // 0.0 is the float null sentinel: never matches. + checkColumn(t, res, "v", []int64{2}) + checkColumn(t, res, "w", []int64{20}) +} + +func TestJoinBoolKeys(t *testing.T) { + // Consistent with the zero-value-is-null rule, false keys never match. + left := dfFromSeries(t, "k", []bool{true, false}, "v", []int64{1, 2}) + right := dfFromSeries(t, "k", []bool{true, false}, "w", []int64{10, 20}) + res := left.Join(right, []string{"k"}, InnerJoin) + checkColumn(t, res, "v", []int64{1}) + checkColumn(t, res, "w", []int64{10}) +} + +// --- empty frames --- + +func TestJoinEmptyRight(t *testing.T) { + left := dfFromSeries(t, "id", []int64{1, 2}, "v", []int64{10, 20}) + right := dfFromSeries(t, "id", []int64{}, "w", []string{}) + + inner := left.Join(right, []string{"id"}, InnerJoin) + if inner.Error() != nil { + t.Fatalf("inner join with empty right: %v", inner.Error()) + } + if inner.Len() != 0 { + t.Errorf("inner join with empty right should have 0 rows, got %d", inner.Len()) + } + checkColumns(t, inner, "id", "v", "w") + + lj := left.Join(right, []string{"id"}, LeftJoin) + checkColumn(t, lj, "id", []int64{1, 2}) + checkColumn(t, lj, "v", []int64{10, 20}) + checkColumn(t, lj, "w", []string{"", ""}) +} + +func TestJoinEmptyLeft(t *testing.T) { + left := dfFromSeries(t, "id", []int64{}, "v", []int64{}) + right := dfFromSeries(t, "id", []int64{1}, "w", []string{"x"}) + + rj := left.Join(right, []string{"id"}, RightJoin) + checkColumn(t, rj, "id", []int64{1}) + checkColumn(t, rj, "v", []int64{0}) + checkColumn(t, rj, "w", []string{"x"}) + + inner := left.Join(right, []string{"id"}, InnerJoin) + if inner.Len() != 0 { + t.Errorf("inner join with empty left should have 0 rows, got %d", inner.Len()) + } +} + +func TestJoinBothEmpty(t *testing.T) { + left := dfFromSeries(t, "id", []int64{}, "v", []int64{}) + right := dfFromSeries(t, "id", []int64{}, "w", []int64{}) + res := left.Join(right, []string{"id"}, OuterJoin) + if res.Error() != nil { + t.Fatalf("outer join of empty frames: %v", res.Error()) + } + if res.Len() != 0 { + t.Errorf("expected 0 rows, got %d", res.Len()) + } + checkColumns(t, res, "id", "v", "w") +} + +// --- error handling --- + +func TestJoinValidationErrors(t *testing.T) { + left := dfFromSeries(t, "id", []int64{1}, "v", []int64{2}) + right := dfFromSeries(t, "id", []int64{1}, "w", []int64{3}) + + cases := []struct { + name string + run func() *DataFrame + }{ + {"nil other", func() *DataFrame { return left.Join(nil, []string{"id"}, InnerJoin) }}, + {"no keys", func() *DataFrame { return left.Join(right, []string{}, InnerJoin) }}, + {"nil keys", func() *DataFrame { return left.Join(right, nil, InnerJoin) }}, + {"missing left column", func() *DataFrame { return left.Join(right, []string{"nope"}, InnerJoin) }}, + {"missing right column", func() *DataFrame { return left.JoinOn(right, []string{"id"}, []string{"nope"}, InnerJoin) }}, + {"key count mismatch", func() *DataFrame { return left.JoinOn(right, []string{"id"}, []string{"id", "w"}, InnerJoin) }}, + {"duplicate left key", func() *DataFrame { return left.JoinOn(right, []string{"id", "id"}, []string{"id", "w"}, InnerJoin) }}, + {"duplicate right key", func() *DataFrame { + l := dfFromSeries(t, "a", []int64{1}, "b", []int64{2}) + return l.JoinOn(right, []string{"a", "b"}, []string{"id", "id"}, InnerJoin) + }}, + {"invalid join type", func() *DataFrame { return left.Join(right, []string{"id"}, JoinType(42)) }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + res := tc.run() + if res.Error() == nil { + t.Fatalf("%s: expected error, got none", tc.name) + } + }) + } +} + +func TestJoinMissingColumnErrorIsColumnNotFound(t *testing.T) { + left := dfFromSeries(t, "id", []int64{1}) + right := dfFromSeries(t, "id", []int64{1}) + res := left.Join(right, []string{"ghost"}, InnerJoin) + if !errors.Is(res.Error(), ErrColumnNotFound) { + t.Errorf("missing key column error should match ErrColumnNotFound, got %v", res.Error()) + } +} + +func TestJoinErrorPropagation(t *testing.T) { + good := dfFromSeries(t, "id", []int64{1}) + bad := good.Filter("ghost", "==", 1) // carries an error + + res := bad.Join(good, []string{"id"}, InnerJoin) + if res.Error() == nil { + t.Fatal("join on errored left frame should carry the error") + } + + res2 := good.Join(bad, []string{"id"}, InnerJoin) + if res2.Error() == nil { + t.Fatal("join with errored right frame should carry the error") + } +} + +// --- immutability & determinism --- + +func TestJoinDoesNotMutateInputs(t *testing.T) { + left, right := ordersCustomers(t) + leftBefore := left.String() + rightBefore := right.String() + + _ = left.Join(right, []string{"customer_id"}, OuterJoin) + + if left.String() != leftBefore { + t.Error("Join mutated the left frame") + } + if right.String() != rightBefore { + t.Error("Join mutated the right frame") + } + if left.Error() != nil || right.Error() != nil { + t.Error("Join set an error on an input frame") + } +} + +func TestJoinResultIsIndependentCopy(t *testing.T) { + left := dfFromSeries(t, "id", []int64{1}, "v", []int64{10}) + right := dfFromSeries(t, "id", []int64{1}, "w", []int64{20}) + res := left.Join(right, []string{"id"}, InnerJoin) + + if err := res.Set(0, "v", int64(999)); err != nil { + t.Fatal(err) + } + orig, _ := left.Get(0, "v") + if orig != int64(10) { + t.Error("mutating join result changed the input frame") + } +} + +func TestJoinDeterministic(t *testing.T) { + left, right := ordersCustomers(t) + first := left.Join(right, []string{"customer_id"}, OuterJoin).String() + for i := 0; i < 10; i++ { + again := left.Join(right, []string{"customer_id"}, OuterJoin).String() + if again != first { + t.Fatalf("join output not deterministic on run %d:\n%s\nvs\n%s", i, first, again) + } + } +} + +// --- chaining --- + +func TestJoinChainsWithOtherOps(t *testing.T) { + left, right := ordersCustomers(t) + res := left. + Join(right, []string{"customer_id"}, InnerJoin). + Filter("amount", ">", 150.0). + Sort("order_id", true) + if res.Error() != nil { + t.Fatalf("chained join error: %v", res.Error()) + } + checkColumn(t, res, "order_id", []int64{2, 4}) + checkColumn(t, res, "name", []string{"Bob", "Alice"}) +} + +func TestJoinLargeConsistency(t *testing.T) { + // A larger randomized-shape join cross-checked against a brute-force nested loop. + n, m := 500, 300 + lk := make([]int64, n) + lv := make([]int64, n) + for i := range lk { + lk[i] = int64((i*7)%97) + 1 // avoid 0: null sentinel + lv[i] = int64(i) + } + rk := make([]int64, m) + rv := make([]int64, m) + for j := range rk { + rk[j] = int64((j*13)%89) + 1 + rv[j] = int64(j) + } + left := dfFromSeries(t, "k", lk, "lrow", lv) + right := dfFromSeries(t, "k", rk, "rrow", rv) + + res := left.Join(right, []string{"k"}, InnerJoin) + if res.Error() != nil { + t.Fatalf("join error: %v", res.Error()) + } + + // Brute force expected pairs in left-major, right-minor order. + var wantL, wantR []int64 + for i := 0; i < n; i++ { + for j := 0; j < m; j++ { + if lk[i] == rk[j] { + wantL = append(wantL, lv[i]) + wantR = append(wantR, rv[j]) + } + } + } + checkColumn(t, res, "lrow", wantL) + checkColumn(t, res, "rrow", wantR) +} + +// --- full type-space coverage --- + +func TestJoinCompositeKeysAllSpaces(t *testing.T) { + ts := time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC) + // Composite key spanning every comparison space, including a mixed + // int64/float64 component (left int, right float). + left := dfFromSeries(t, + "s", []string{"x", "y"}, + "i", []int64{1, 2}, + "f", []int64{7, 8}, // int on left, float on right → float space + "b", []bool{true, true}, + "ts", []time.Time{ts, ts}, + "v", []int64{100, 200}, + ) + right := dfFromSeries(t, + "s", []string{"x", "y"}, + "i", []int64{1, 99}, + "f", []float64{7, 8}, + "b", []bool{true, true}, + "ts", []time.Time{ts, ts}, + "w", []int64{111, 222}, + ) + res := left.Join(right, []string{"s", "i", "f", "b", "ts"}, InnerJoin) + if res.Error() != nil { + t.Fatalf("composite all-spaces join error: %v", res.Error()) + } + // Only row 0 matches (right row 1 has i=99). + checkColumn(t, res, "v", []int64{100}) + checkColumn(t, res, "w", []int64{111}) + // Mixed numeric key component is promoted to float64 in the output. + ct, _ := res.GetColumnType("f") + if ct != Float64Type { + t.Errorf("mixed numeric composite key output type = %v, want Float64Type", ct) + } +} + +func TestJoinOuterCompositeKeys(t *testing.T) { + left := dfFromSeries(t, + "a", []string{"x", "y"}, + "n", []int64{1, 2}, + "v", []float64{1.5, 2.5}, + ) + right := dfFromSeries(t, + "a", []string{"x", "z"}, + "n", []int64{1, 3}, + "w", []float64{10, 30}, + ) + res := left.Join(right, []string{"a", "n"}, OuterJoin) + // Left rows in order, then unmatched right row; key values for the + // right-only row come from the right frame. + checkColumn(t, res, "a", []string{"x", "y", "z"}) + checkColumn(t, res, "n", []int64{1, 2, 3}) + checkColumn(t, res, "v", []float64{1.5, 2.5, 0}) + checkColumn(t, res, "w", []float64{10, 0, 30}) +} + +func TestJoinOuterFillsAllColumnTypes(t *testing.T) { + ts := time.Date(2026, 2, 2, 0, 0, 0, 0, time.UTC) + left := dfFromSeries(t, + "id", []int64{1, 2}, + "ls", []string{"a", "b"}, + "lf", []float64{1.1, 2.2}, + "lb", []bool{true, true}, + "lt", []time.Time{ts, ts}, + ) + right := dfFromSeries(t, + "id", []int64{2, 3}, + "rs", []string{"B", "C"}, + "rf", []float64{20, 30}, + "rb", []bool{true, true}, + "rt", []time.Time{ts, ts}, + ) + res := left.Join(right, []string{"id"}, OuterJoin) + if res.Error() != nil { + t.Fatalf("outer join error: %v", res.Error()) + } + var zt time.Time + checkColumn(t, res, "id", []int64{1, 2, 3}) + checkColumn(t, res, "ls", []string{"a", "b", ""}) + checkColumn(t, res, "lf", []float64{1.1, 2.2, 0}) + checkColumn(t, res, "lb", []bool{true, true, false}) + checkColumn(t, res, "lt", []time.Time{ts, ts, zt}) + checkColumn(t, res, "rs", []string{"", "B", "C"}) + checkColumn(t, res, "rf", []float64{0, 20, 30}) + checkColumn(t, res, "rb", []bool{false, true, true}) + checkColumn(t, res, "rt", []time.Time{zt, ts, ts}) +} + +func TestJoinOuterKeyTypesStringBoolTime(t *testing.T) { + ts1 := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + ts2 := time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC) + + // String key, right-only row feeds the key column. + l1 := dfFromSeries(t, "k", []string{"a"}, "v", []int64{1}) + r1 := dfFromSeries(t, "k", []string{"b"}, "w", []int64{2}) + res1 := l1.Join(r1, []string{"k"}, OuterJoin) + checkColumn(t, res1, "k", []string{"a", "b"}) + + // Time key. + l2 := dfFromSeries(t, "k", []time.Time{ts1}, "v", []int64{1}) + r2 := dfFromSeries(t, "k", []time.Time{ts2}, "w", []int64{2}) + res2 := l2.Join(r2, []string{"k"}, OuterJoin) + checkColumn(t, res2, "k", []time.Time{ts1, ts2}) + + // Bool key: true matches true; both sides' false rows are null keys. + l3 := dfFromSeries(t, "k", []bool{true, false}, "v", []int64{1, 2}) + r3 := dfFromSeries(t, "k", []bool{true, false}, "w", []int64{10, 20}) + res3 := l3.Join(r3, []string{"k"}, OuterJoin) + checkColumn(t, res3, "k", []bool{true, false, false}) + checkColumn(t, res3, "v", []int64{1, 2, 0}) + checkColumn(t, res3, "w", []int64{10, 0, 20}) + + // Float key on both sides stays Float64Type. + l4 := dfFromSeries(t, "k", []float64{1.5}, "v", []int64{1}) + r4 := dfFromSeries(t, "k", []float64{2.5}, "w", []int64{2}) + res4 := l4.Join(r4, []string{"k"}, OuterJoin) + checkColumn(t, res4, "k", []float64{1.5, 2.5}) + ct, _ := res4.GetColumnType("k") + if ct != Float64Type { + t.Errorf("float key type = %v, want Float64Type", ct) + } +} diff --git a/json_test.go b/json_test.go new file mode 100644 index 0000000..e502747 --- /dev/null +++ b/json_test.go @@ -0,0 +1,307 @@ +package otters + +import ( + "math" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// --- reading basics --- + +func TestReadJSONFromStringBasic(t *testing.T) { + data := `[ + {"name": "Alice", "age": 30, "score": 95.5, "active": true}, + {"name": "Bob", "age": 25, "score": 87.2, "active": false} + ]` + df, err := ReadJSONFromString(data) + if err != nil { + t.Fatalf("ReadJSONFromString: %v", err) + } + checkColumns(t, df, "name", "age", "score", "active") + checkColumn(t, df, "name", []string{"Alice", "Bob"}) + checkColumn(t, df, "age", []int64{30, 25}) + checkColumn(t, df, "score", []float64{95.5, 87.2}) + checkColumn(t, df, "active", []bool{true, false}) +} + +func TestReadJSONRespectsNativeTypes(t *testing.T) { + // JSON string "123" must stay a string; integer numbers stay Int64Type. + data := `[{"n": 123, "s": "123"}]` + df, err := ReadJSONFromString(data) + if err != nil { + t.Fatal(err) + } + if ct, _ := df.GetColumnType("n"); ct != Int64Type { + t.Errorf("n type = %v, want Int64Type", ct) + } + if ct, _ := df.GetColumnType("s"); ct != StringType { + t.Errorf("s type = %v, want StringType", ct) + } +} + +func TestReadJSONTimeColumn(t *testing.T) { + data := `[{"ts": "2026-01-15T10:30:00Z"}, {"ts": "2026-02-20T08:00:00Z"}]` + df, err := ReadJSONFromString(data) + if err != nil { + t.Fatal(err) + } + if ct, _ := df.GetColumnType("ts"); ct != TimeType { + t.Errorf("ts type = %v, want TimeType", ct) + } + v, _ := df.Get(0, "ts") + want := time.Date(2026, 1, 15, 10, 30, 0, 0, time.UTC) + if !v.(time.Time).Equal(want) { + t.Errorf("ts[0] = %v, want %v", v, want) + } +} + +func TestReadJSONUnionSchemaFirstSeenOrder(t *testing.T) { + data := `[{"a": 1}, {"b": "x", "a": 2}, {"c": true}]` + df, err := ReadJSONFromString(data) + if err != nil { + t.Fatal(err) + } + checkColumns(t, df, "a", "b", "c") + // Missing keys fill with zero values. + checkColumn(t, df, "a", []int64{1, 2, 0}) + checkColumn(t, df, "b", []string{"", "x", ""}) + checkColumn(t, df, "c", []bool{false, false, true}) +} + +func TestReadJSONNullsFillZero(t *testing.T) { + data := `[{"n": 5, "s": "x"}, {"n": null, "s": null}]` + df, err := ReadJSONFromString(data) + if err != nil { + t.Fatal(err) + } + checkColumn(t, df, "n", []int64{5, 0}) + checkColumn(t, df, "s", []string{"x", ""}) +} + +func TestReadJSONTypeConflictPromotesToString(t *testing.T) { + data := `[{"v": 1}, {"v": "two"}, {"v": true}]` + df, err := ReadJSONFromString(data) + if err != nil { + t.Fatal(err) + } + if ct, _ := df.GetColumnType("v"); ct != StringType { + t.Errorf("v type = %v, want StringType", ct) + } + checkColumn(t, df, "v", []string{"1", "two", "true"}) +} + +func TestReadJSONNestedValuesStringify(t *testing.T) { + data := `[{"obj": {"x": 1}, "arr": [1, 2]}]` + df, err := ReadJSONFromString(data) + if err != nil { + t.Fatal(err) + } + checkColumn(t, df, "obj", []string{`{"x":1}`}) + checkColumn(t, df, "arr", []string{"[1,2]"}) +} + +func TestReadJSONEmptyArray(t *testing.T) { + df, err := ReadJSONFromString(`[]`) + if err != nil { + t.Fatalf("empty array should not error: %v", err) + } + if !df.IsEmpty() { + t.Error("empty array should produce an empty DataFrame") + } +} + +func TestReadJSONWhitespaceTolerance(t *testing.T) { + df, err := ReadJSONFromString("\n\t [ {\"a\": 1} ,\n {\"a\": 2} ] \n\t") + if err != nil { + t.Fatalf("whitespace around/inside the array should be fine: %v", err) + } + checkColumn(t, df, "a", []int64{1, 2}) +} + +// --- reading errors --- + +func TestReadJSONErrors(t *testing.T) { + cases := []struct { + name string + data string + }{ + {"top-level object", `{"a": 1}`}, + {"top-level scalar", `42`}, + {"top-level string", `"hello"`}, + {"empty input", ``}, + {"whitespace only", " \n\t "}, + {"non-object element", `[1, 2]`}, + {"string element", `["a"]`}, + {"malformed", `[{"a": }]`}, + {"unterminated array", `[{"a": 1}`}, + {"trailing garbage", `[{"a": 1}] extra`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if _, err := ReadJSONFromString(tc.data); err == nil { + t.Errorf("expected error for %s, got none", tc.name) + } + }) + } +} + +func TestReadJSONErrorReportsElementIndex(t *testing.T) { + _, err := ReadJSONFromString(`[{"a": 1}, {"a": 2}, [1]]`) + if err == nil { + t.Fatal("expected error") + } + oe, ok := err.(*OtterError) + if !ok { + t.Fatalf("expected *OtterError, got %T", err) + } + if oe.Row != 3 { + t.Errorf("error row = %d, want 3 (1-based element index)", oe.Row) + } +} + +func TestReadJSONMissingFile(t *testing.T) { + if _, err := ReadJSON(filepath.Join(t.TempDir(), "nope.json")); err == nil { + t.Error("missing file should error") + } +} + +// --- options --- + +func TestReadJSONOptions(t *testing.T) { + data := `[{"n": 1}, {"n": 2}, {"n": 3}, {"n": 4}]` + + skip, err := ReadJSONFromStringWithOptions(data, JSONOptions{SkipRows: 2}) + if err != nil { + t.Fatal(err) + } + checkColumn(t, skip, "n", []int64{3, 4}) + + max, err := ReadJSONFromStringWithOptions(data, JSONOptions{MaxRows: 2}) + if err != nil { + t.Fatal(err) + } + checkColumn(t, max, "n", []int64{1, 2}) + + both, err := ReadJSONFromStringWithOptions(data, JSONOptions{SkipRows: 1, MaxRows: 2}) + if err != nil { + t.Fatal(err) + } + checkColumn(t, both, "n", []int64{2, 3}) + + past, err := ReadJSONFromStringWithOptions(data, JSONOptions{SkipRows: 10}) + if err != nil { + t.Fatal(err) + } + if past.Len() != 0 { + t.Errorf("skipping past the end should give 0 rows, got %d", past.Len()) + } +} + +// --- file round trips --- + +func TestJSONFileRoundTrip(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "data.json") + + ts := time.Date(2026, 6, 1, 12, 0, 0, 0, time.UTC) + df := dfFromSeries(t, + "id", []int64{1, 2}, + "name", []string{"a", "b"}, + "score", []float64{1.5, 2.5}, + "ok", []bool{true, false}, + "ts", []time.Time{ts, ts}, + ) + if err := df.WriteJSON(path); err != nil { + t.Fatalf("WriteJSON: %v", err) + } + + back, err := ReadJSON(path) + if err != nil { + t.Fatalf("ReadJSON: %v", err) + } + checkColumns(t, back, "id", "name", "score", "ok", "ts") + checkColumn(t, back, "id", []int64{1, 2}) + checkColumn(t, back, "name", []string{"a", "b"}) + checkColumn(t, back, "score", []float64{1.5, 2.5}) + checkColumn(t, back, "ok", []bool{true, false}) + if ct, _ := back.GetColumnType("ts"); ct != TimeType { + t.Errorf("round-tripped ts type = %v, want TimeType", ct) + } +} + +func TestWriteJSONSpecialValues(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "special.json") + + var zeroTime time.Time + ts := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + df := dfFromSeries(t, + "f", []float64{math.NaN(), math.Inf(1), 1.5}, + "t", []time.Time{zeroTime, ts, ts}, + ) + if err := df.WriteJSON(path); err != nil { + t.Fatalf("WriteJSON: %v", err) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + content := string(raw) + // NaN/Inf and zero time serialize as null. + if !strings.Contains(content, `"f":null`) { + t.Errorf("NaN should serialize as null, got: %s", content) + } + if !strings.Contains(content, `"t":null`) { + t.Errorf("zero time should serialize as null, got: %s", content) + } + // Output must be valid JSON that reads back. + if _, err := ReadJSON(path); err != nil { + t.Errorf("written JSON should read back cleanly: %v", err) + } +} + +func TestWriteJSONEmptyFrame(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "empty.json") + df := dfFromSeries(t, "a", []int64{}) + if err := df.WriteJSON(path); err != nil { + t.Fatalf("WriteJSON empty: %v", err) + } + back, err := ReadJSON(path) + if err != nil { + t.Fatalf("ReadJSON of empty output: %v", err) + } + if back.Len() != 0 { + t.Errorf("expected 0 rows, got %d", back.Len()) + } +} + +func TestWriteJSONErroredFrame(t *testing.T) { + df := dfFromSeries(t, "a", []int64{1}).Filter("ghost", "==", 1) + if err := df.WriteJSON(filepath.Join(t.TempDir(), "x.json")); err == nil { + t.Error("WriteJSON on errored frame should return the error") + } +} + +// --- JSONL/JSON shared behavior parity --- + +func TestJSONAndJSONLProduceIdenticalFrames(t *testing.T) { + jsonData := `[{"a": 1, "b": "x"}, {"a": 2, "c": 3.5}]` + jsonlData := "{\"a\": 1, \"b\": \"x\"}\n{\"a\": 2, \"c\": 3.5}\n" + + fromJSON, err := ReadJSONFromString(jsonData) + if err != nil { + t.Fatal(err) + } + fromJSONL, err := ReadJSONLFromString(jsonlData) + if err != nil { + t.Fatal(err) + } + if fromJSON.String() != fromJSONL.String() { + t.Errorf("JSON and JSONL should produce identical frames:\n%s\nvs\n%s", + fromJSON.String(), fromJSONL.String()) + } +} diff --git a/jsonl.go b/jsonl.go index 378886b..3e59212 100644 --- a/jsonl.go +++ b/jsonl.go @@ -105,7 +105,7 @@ func readJSONL(r io.Reader, options JSONLOptions, operation string) (*DataFrame, return nil, wrapError(operation, err) } - return buildDataFrameFromJSONLRows(order, rows, operation) + return buildDataFrameFromObjects(order, rows, operation) } // decodeJSONLine decodes one JSONL line into a value map plus the object's @@ -114,12 +114,28 @@ func decodeJSONLine(line string) (map[string]any, []string, error) { dec := json.NewDecoder(strings.NewReader(line)) dec.UseNumber() + obj, keys, err := decodeJSONObject(dec) + if err != nil { + return nil, nil, err + } + if _, err := dec.Token(); err != io.EOF { + return nil, nil, fmt.Errorf("unexpected data after JSON object") + } + + return obj, keys, nil +} + +// decodeJSONObject decodes the next value from dec, which must be a JSON +// object, returning the value map plus the keys in order of appearance. +// Shared by the JSONL (one object per line) and JSON (array of objects) +// readers. +func decodeJSONObject(dec *json.Decoder) (map[string]any, []string, error) { tok, err := dec.Token() if err != nil { return nil, nil, err } if delim, ok := tok.(json.Delim); !ok || delim != '{' { - return nil, nil, fmt.Errorf("line is not a JSON object") + return nil, nil, fmt.Errorf("value is not a JSON object") } obj := make(map[string]any) @@ -144,15 +160,13 @@ func decodeJSONLine(line string) (map[string]any, []string, error) { if _, err := dec.Token(); err != nil { // consume closing '}' return nil, nil, err } - if _, err := dec.Token(); err != io.EOF { - return nil, nil, fmt.Errorf("unexpected data after JSON object") - } return obj, keys, nil } -// buildDataFrameFromJSONLRows constructs a DataFrame from decoded JSONL rows -func buildDataFrameFromJSONLRows(order []string, rows []map[string]any, operation string) (*DataFrame, error) { +// buildDataFrameFromObjects constructs a DataFrame from decoded JSON objects +// (one per row). Shared by the JSONL and JSON-array readers. +func buildDataFrameFromObjects(order []string, rows []map[string]any, operation string) (*DataFrame, error) { if len(order) == 0 { return NewDataFrame(), nil } @@ -164,8 +178,8 @@ func buildDataFrameFromJSONLRows(order []string, rows []map[string]any, operatio values[i] = row[name] // missing key yields nil, same as JSON null } - colType := inferJSONLColumnType(values) - s, err := buildJSONLSeries(name, values, colType) + colType := inferJSONColumnType(values) + s, err := buildJSONSeries(name, values, colType) if err != nil { return nil, wrapColumnError(operation, name, err) } @@ -175,11 +189,11 @@ func buildDataFrameFromJSONLRows(order []string, rows []map[string]any, operatio return NewDataFrameFromSeries(series...) } -// inferJSONLColumnType picks a column type from decoded JSON values. +// inferJSONColumnType picks a column type from decoded JSON values. // Unlike CSV inference, JSON values carry their own types, so strings are // never reinterpreted as numbers or bools; the only string promotion is to // TimeType when every non-empty value parses as a time. -func inferJSONLColumnType(values []any) ColumnType { +func inferJSONColumnType(values []any) ColumnType { sawNumber := false sawFloat := false sawBool := false @@ -239,9 +253,9 @@ func inferJSONLColumnType(values []any) ColumnType { return StringType } -// buildJSONLSeries converts decoded JSON values into a typed Series. +// buildJSONSeries converts decoded JSON values into a typed Series. // nil (JSON null or missing key) becomes the column type's zero value. -func buildJSONLSeries(name string, values []any, colType ColumnType) (*Series, error) { +func buildJSONSeries(name string, values []any, colType ColumnType) (*Series, error) { switch colType { case Int64Type: data := make([]int64, len(values)) @@ -251,7 +265,7 @@ func buildJSONLSeries(name string, values []any, colType ColumnType) (*Series, e } n, err := v.(json.Number).Int64() if err != nil { - return nil, wrapError("buildJSONLSeries", err) + return nil, wrapError("buildJSONSeries", err) } data[i] = n } @@ -265,7 +279,7 @@ func buildJSONLSeries(name string, values []any, colType ColumnType) (*Series, e } f, err := v.(json.Number).Float64() if err != nil { - return nil, wrapError("buildJSONLSeries", err) + return nil, wrapError("buildJSONSeries", err) } data[i] = f } @@ -293,7 +307,7 @@ func buildJSONLSeries(name string, values []any, colType ColumnType) (*Series, e } t, err := parseTimeValue(trimmed) if err != nil { - return nil, wrapError("buildJSONLSeries", err) + return nil, wrapError("buildJSONSeries", err) } data[i] = t } @@ -304,7 +318,7 @@ func buildJSONLSeries(name string, values []any, colType ColumnType) (*Series, e for i, v := range values { s, err := formatJSONValueAsString(v) if err != nil { - return nil, wrapError("buildJSONLSeries", err) + return nil, wrapError("buildJSONSeries", err) } data[i] = s } @@ -369,7 +383,7 @@ func (df *DataFrame) WriteJSONL(filename string) error { if err != nil { return wrapColumnError("WriteJSONL", colName, err) } - formatted, err := formatValueForJSONL(value) + formatted, err := formatValueForJSON(value) if err != nil { return wrapColumnError("WriteJSONL", colName, err) } @@ -390,8 +404,8 @@ func (df *DataFrame) WriteJSONL(filename string) error { return nil } -// formatValueForJSONL formats a single cell as a JSON value -func formatValueForJSONL(value any) (string, error) { +// formatValueForJSON formats a single cell as a JSON value +func formatValueForJSON(value any) (string, error) { switch v := value.(type) { case string: raw, err := json.Marshal(v) diff --git a/transform.go b/transform.go new file mode 100644 index 0000000..bce9919 --- /dev/null +++ b/transform.go @@ -0,0 +1,432 @@ +package otters + +import ( + "fmt" + "math" + "strconv" + "time" +) + +// Apply derives a column by evaluating fn once per row. fn receives the row's +// values keyed by column name and returns the new cell value. +// +// The derived column's type is inferred from the returned values: all-integer +// results become Int64Type, a mix of integers and floats becomes Float64Type, +// and uniform string/bool/time results keep their type. Any other mix falls +// back to StringType with canonical formatting. A nil return fills with the +// inferred type's zero value; a column of only nils is StringType. +// +// If newColumn already exists it is replaced in place (same position, +// possibly a new type); otherwise the column is appended. The receiver is +// never mutated, and a panic inside fn is captured as an error on the +// returned frame. +func (df *DataFrame) Apply(newColumn string, fn func(row map[string]any) any) (result *DataFrame) { + const op = "Apply" + if df.err != nil { + return df + } + if newColumn == "" { + return df.setError(newOpError(op, "new column name must not be empty")) + } + if fn == nil { + return df.setError(newOpError(op, "fn must not be nil")) + } + + // User code runs inside fn: convert a panic into a carried error instead + // of letting it escape ("no panics" design promise). + defer func() { + if r := recover(); r != nil { + result = df.setError(panicToError(op, r)) + } + }() + + seriesList := make([]*Series, len(df.order)) + for i, name := range df.order { + seriesList[i] = df.columns[name] + } + + values := make([]any, df.length) + for i := 0; i < df.length; i++ { + row := make(map[string]any, len(seriesList)) + for _, s := range seriesList { + v, err := s.Get(i) + if err != nil { + return df.setError(wrapColumnError(op, s.Name, err)) + } + row[s.Name] = v + } + values[i] = fn(row) + } + + return df.withDerivedColumn(op, newColumn, values) +} + +// Map derives a column by transforming srcColumn element-wise with fn. +// newColumn may equal srcColumn to replace it. Type inference, replacement, +// immutability, and panic capture behave exactly like Apply. +func (df *DataFrame) Map(srcColumn, newColumn string, fn func(v any) any) (result *DataFrame) { + const op = "Map" + if df.err != nil { + return df + } + if err := df.validateColumnExists(srcColumn); err != nil { + return df.setError(err) + } + if newColumn == "" { + return df.setError(newOpError(op, "new column name must not be empty")) + } + if fn == nil { + return df.setError(newOpError(op, "fn must not be nil")) + } + + defer func() { + if r := recover(); r != nil { + result = df.setError(panicToError(op, r)) + } + }() + + src := df.columns[srcColumn] + values := make([]any, df.length) + for i := 0; i < df.length; i++ { + v, err := src.Get(i) + if err != nil { + return df.setError(wrapColumnError(op, srcColumn, err)) + } + values[i] = fn(v) + } + + return df.withDerivedColumn(op, newColumn, values) +} + +// FillNA returns a DataFrame with the column's missing values (the column +// type's zero value: "", 0, 0.0, false, zero time) replaced by value. The +// fill value must fit the column type; ints coerce to int64/float64 columns, +// and an integral float may fill an int64 column. Note that a legitimate +// zero (e.g. an actual 0 in an int column) is indistinguishable from missing +// under this convention. +func (df *DataFrame) FillNA(column string, value any) *DataFrame { + const op = "FillNA" + if df.err != nil { + return df + } + if err := df.validateColumnExists(column); err != nil { + return df.setError(err) + } + + newDf := df.Copy() + s := newDf.columns[column] + + switch s.Type { + case Int64Type: + fill, ok := fillValueInt64(value) + if !ok { + return df.setError(newColumnError(op, column, fmt.Sprintf("cannot use %T value %v as an int64 fill", value, value))) + } + data := s.Data.([]int64) + for i, v := range data { + if v == 0 { + data[i] = fill + } + } + case Float64Type: + fill, ok := toFloat64(value) + if !ok { + return df.setError(newColumnError(op, column, fmt.Sprintf("cannot use %T value %v as a float64 fill", value, value))) + } + data := s.Data.([]float64) + for i, v := range data { + if v == 0 { + data[i] = fill + } + } + case StringType: + fill, ok := value.(string) + if !ok { + return df.setError(newColumnError(op, column, fmt.Sprintf("cannot use %T value %v as a string fill", value, value))) + } + data := s.Data.([]string) + for i, v := range data { + if v == "" { + data[i] = fill + } + } + case BoolType: + fill, ok := value.(bool) + if !ok { + return df.setError(newColumnError(op, column, fmt.Sprintf("cannot use %T value %v as a bool fill", value, value))) + } + data := s.Data.([]bool) + for i, v := range data { + if !v { + data[i] = fill + } + } + case TimeType: + fill, ok := value.(time.Time) + if !ok { + return df.setError(newColumnError(op, column, fmt.Sprintf("cannot use %T value %v as a time fill", value, value))) + } + data := s.Data.([]time.Time) + for i, v := range data { + if v.IsZero() { + data[i] = fill + } + } + default: + return df.setError(newColumnError(op, column, "unsupported column type")) + } + + return newDf +} + +// DropNA returns a DataFrame without the rows holding a missing value (the +// column type's zero value) in any of the given columns. With no columns +// given, every column is checked. Row order is preserved. +func (df *DataFrame) DropNA(columns ...string) *DataFrame { + const op = "DropNA" + if df.err != nil { + return df + } + cols := columns + if len(cols) == 0 { + cols = df.order + } + if err := df.validateColumnsExist(cols); err != nil { + return df.setError(err) + } + if df.length == 0 { + return df.Copy() + } + + masks := make([][]bool, len(cols)) + for j, c := range cols { + masks[j] = seriesNullMask(df.columns[c]) + } + + keep := make([]int, 0, df.length) + for i := 0; i < df.length; i++ { + hasNull := false + for _, m := range masks { + if m[i] { + hasNull = true + break + } + } + if !hasNull { + keep = append(keep, i) + } + } + + return df.selectRows(keep, op) +} + +// seriesNullMask marks each row whose value is the column type's zero value. +func seriesNullMask(s *Series) []bool { + switch s.Type { + case Int64Type: + return nullMaskInt64(s.Data.([]int64)) + case Float64Type: + return nullMaskFloat64(s.Data.([]float64)) + case StringType: + return nullMaskString(s.Data.([]string)) + case BoolType: + return nullMaskBool(s.Data.([]bool)) + case TimeType: + return nullMaskTime(s.Data.([]time.Time)) + default: + return make([]bool, s.Length) + } +} + +// withDerivedColumn builds a typed series from boxed values and returns a +// copy of df with it added, or replacing an existing column of the same name. +func (df *DataFrame) withDerivedColumn(op, name string, values []any) *DataFrame { + data, err := anyValuesToTypedSlice(op, values) + if err != nil { + return df.setError(err) + } + series, err := newSeriesOwned(name, data) + if err != nil { + return df.setError(wrapColumnError(op, name, err)) + } + + newDf := df.Copy() + if _, exists := newDf.columns[name]; exists { + newDf.columns[name] = series // replaced in place: order position kept + } else { + newDf.columns[name] = series + newDf.order = append(newDf.order, name) + } + return newDf +} + +// anyValuesToTypedSlice infers the best column type for boxed values and +// materializes the matching typed slice. See Apply for the inference rules. +func anyValuesToTypedSlice(op string, values []any) (any, error) { + var sawInt, sawFloat, sawString, sawBool, sawTime bool + + for i, v := range values { + switch v.(type) { + case nil: + case int, int8, int16, int32, int64, uint, uint8, uint16, uint32: + sawInt = true + case float32, float64: + sawFloat = true + case string: + sawString = true + case bool: + sawBool = true + case time.Time: + sawTime = true + default: + return nil, newRowError(op, i, fmt.Sprintf("unsupported derived value type %T", v)) + } + } + + kinds := 0 + for _, saw := range []bool{sawInt || sawFloat, sawString, sawBool, sawTime} { + if saw { + kinds++ + } + } + + switch { + case kinds > 1: + return derivedStrings(values), nil + case sawFloat: + return derivedFloats(values), nil + case sawInt: + return derivedInts(values), nil + case sawBool: + return derivedBools(values), nil + case sawTime: + return derivedTimes(values), nil + default: + // Only strings, or only nils: a column of no values carries no type + // information, matching InferType's behavior. + return derivedStrings(values), nil + } +} + +func derivedInts(values []any) []int64 { + out := make([]int64, len(values)) + for i, v := range values { + if v == nil { + continue + } + out[i], _ = anyToInt64(v) + } + return out +} + +func derivedFloats(values []any) []float64 { + out := make([]float64, len(values)) + for i, v := range values { + if v == nil { + continue + } + out[i], _ = anyToFloat64(v) + } + return out +} + +func derivedBools(values []any) []bool { + out := make([]bool, len(values)) + for i, v := range values { + if b, ok := v.(bool); ok { + out[i] = b + } + } + return out +} + +func derivedTimes(values []any) []time.Time { + out := make([]time.Time, len(values)) + for i, v := range values { + if t, ok := v.(time.Time); ok { + out[i] = t + } + } + return out +} + +// derivedStrings renders mixed-kind values with the library's canonical +// formatting (the same rendering WriteCSV and String use). +func derivedStrings(values []any) []string { + out := make([]string, len(values)) + for i, v := range values { + switch x := v.(type) { + case nil: + out[i] = "" + case string: + out[i] = x + case bool: + if x { + out[i] = "true" + } else { + out[i] = "false" + } + case float32: + out[i] = strconv.FormatFloat(float64(x), 'g', -1, 64) + case float64: + out[i] = strconv.FormatFloat(x, 'g', -1, 64) + case time.Time: + out[i] = x.String() + default: + n, _ := anyToInt64(x) + out[i] = strconv.FormatInt(n, 10) + } + } + return out +} + +// anyToInt64 widens any supported integer kind to int64. +func anyToInt64(v any) (int64, bool) { + switch x := v.(type) { + case int: + return int64(x), true + case int8: + return int64(x), true + case int16: + return int64(x), true + case int32: + return int64(x), true + case int64: + return x, true + case uint: + return int64(x), true + case uint8: + return int64(x), true + case uint16: + return int64(x), true + case uint32: + return int64(x), true + } + return 0, false +} + +// anyToFloat64 widens any supported numeric kind to float64. +func anyToFloat64(v any) (float64, bool) { + if f, ok := v.(float64); ok { + return f, true + } + if f, ok := v.(float32); ok { + return float64(f), true + } + if n, ok := anyToInt64(v); ok { + return float64(n), true + } + return 0, false +} + +// fillValueInt64 coerces a FillNA value for an int64 column: ints pass +// through, an integral float converts, a fractional float is rejected. +func fillValueInt64(value any) (int64, bool) { + if f, isFloat := value.(float64); isFloat { + if f != math.Trunc(f) { + return 0, false + } + return int64(f), true + } + return anyToInt64(value) +} diff --git a/transform_test.go b/transform_test.go new file mode 100644 index 0000000..176cd02 --- /dev/null +++ b/transform_test.go @@ -0,0 +1,450 @@ +package otters + +import ( + "strings" + "testing" + "time" +) + +// --- Apply --- + +func TestApplyDerivedColumn(t *testing.T) { + df := dfFromSeries(t, + "salary", []float64{1000, 2000}, + "bonus_rate", []float64{0.1, 0.2}, + ) + res := df.Apply("bonus", func(row map[string]any) any { + return row["salary"].(float64) * row["bonus_rate"].(float64) + }) + if res.Error() != nil { + t.Fatalf("Apply error: %v", res.Error()) + } + checkColumns(t, res, "salary", "bonus_rate", "bonus") + checkColumn(t, res, "bonus", []float64{100, 400}) +} + +func TestApplyTypeInference(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 2, 3}) + ts := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + cases := []struct { + name string + fn func(map[string]any) any + wantType ColumnType + want any + }{ + {"all int64", func(r map[string]any) any { return r["n"].(int64) * 2 }, Int64Type, []int64{2, 4, 6}}, + {"plain int", func(r map[string]any) any { return int(r["n"].(int64)) }, Int64Type, []int64{1, 2, 3}}, + {"all float", func(r map[string]any) any { return float64(r["n"].(int64)) / 2 }, Float64Type, []float64{0.5, 1, 1.5}}, + {"int and float mix", func(r map[string]any) any { + if r["n"].(int64) == 2 { + return 2.5 + } + return r["n"].(int64) + }, Float64Type, []float64{1, 2.5, 3}}, + {"all string", func(r map[string]any) any { return "x" }, StringType, []string{"x", "x", "x"}}, + {"all bool", func(r map[string]any) any { return r["n"].(int64) > 1 }, BoolType, []bool{false, true, true}}, + {"all time", func(r map[string]any) any { return ts }, TimeType, []time.Time{ts, ts, ts}}, + {"mixed kinds fall back to string", func(r map[string]any) any { + switch r["n"].(int64) { + case 1: + return int64(7) + case 2: + return "seven" + default: + return true + } + }, StringType, []string{"7", "seven", "true"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + res := df.Apply("out", tc.fn) + if res.Error() != nil { + t.Fatalf("Apply error: %v", res.Error()) + } + ct, _ := res.GetColumnType("out") + if ct != tc.wantType { + t.Fatalf("inferred type = %v, want %v", ct, tc.wantType) + } + checkColumn(t, res, "out", tc.want) + }) + } +} + +func TestApplyNilValuesFillZero(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 2, 3}) + res := df.Apply("out", func(r map[string]any) any { + if r["n"].(int64) == 2 { + return nil + } + return r["n"].(int64) * 10 + }) + checkColumn(t, res, "out", []int64{10, 0, 30}) +} + +func TestApplyAllNilIsStringColumn(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 2}) + res := df.Apply("out", func(map[string]any) any { return nil }) + ct, _ := res.GetColumnType("out") + if ct != StringType { + t.Errorf("all-nil derived column type = %v, want StringType", ct) + } + checkColumn(t, res, "out", []string{"", ""}) +} + +func TestApplyReplacesExistingColumnInPlace(t *testing.T) { + df := dfFromSeries(t, "a", []int64{1, 2}, "b", []int64{3, 4}, "c", []int64{5, 6}) + res := df.Apply("b", func(r map[string]any) any { return "replaced" }) + // Position preserved, type changed. + checkColumns(t, res, "a", "b", "c") + checkColumn(t, res, "b", []string{"replaced", "replaced"}) +} + +func TestApplyRowMapContainsAllColumns(t *testing.T) { + df := dfFromSeries(t, "a", []int64{1}, "b", []string{"x"}, "c", []bool{true}) + res := df.Apply("out", func(r map[string]any) any { + if len(r) != 3 { + t.Errorf("row map has %d entries, want 3", len(r)) + } + if r["a"] != int64(1) || r["b"] != "x" || r["c"] != true { + t.Errorf("row map contents wrong: %#v", r) + } + return int64(0) + }) + if res.Error() != nil { + t.Fatal(res.Error()) + } +} + +func TestApplyErrors(t *testing.T) { + df := dfFromSeries(t, "a", []int64{1}) + + if res := df.Apply("", func(map[string]any) any { return 1 }); res.Error() == nil { + t.Error("empty column name should error") + } + if res := df.Apply("out", nil); res.Error() == nil { + t.Error("nil fn should error") + } + // Unsupported return type. + if res := df.Apply("out", func(map[string]any) any { return struct{}{} }); res.Error() == nil { + t.Error("unsupported return type should error") + } + // Error frame propagates. + bad := df.Filter("ghost", "==", 1) + if res := bad.Apply("out", func(map[string]any) any { return 1 }); res.Error() == nil { + t.Error("Apply on errored frame should carry the error") + } +} + +func TestApplyRecoversPanicInUserFn(t *testing.T) { + df := dfFromSeries(t, "a", []int64{1, 0}) + res := df.Apply("inv", func(r map[string]any) any { + return 10 / r["a"].(int64) // divide by zero on row 2 + }) + if res.Error() == nil { + t.Fatal("panic in user fn must surface as an error, not crash") + } + if !strings.Contains(res.Error().Error(), "divide") { + t.Errorf("error should mention the panic, got: %v", res.Error()) + } +} + +func TestApplyDoesNotMutateReceiver(t *testing.T) { + df := dfFromSeries(t, "a", []int64{1, 2}) + before := df.String() + _ = df.Apply("b", func(map[string]any) any { return 1 }) + _ = df.Apply("a", func(map[string]any) any { return "overwrite" }) + if df.String() != before { + t.Error("Apply mutated the receiver") + } + if df.Width() != 1 { + t.Error("Apply added a column to the receiver") + } +} + +func TestApplyEmptyFrame(t *testing.T) { + df := dfFromSeries(t, "a", []int64{}) + res := df.Apply("out", func(map[string]any) any { return int64(1) }) + if res.Error() != nil { + t.Fatalf("Apply on empty frame: %v", res.Error()) + } + if res.Len() != 0 || !res.HasColumn("out") { + t.Error("Apply on empty frame should add an empty column") + } +} + +// --- Map --- + +func TestMapBasic(t *testing.T) { + df := dfFromSeries(t, "name", []string{"alice", "bob"}, "age", []int64{30, 40}) + res := df.Map("name", "upper", func(v any) any { + return strings.ToUpper(v.(string)) + }) + if res.Error() != nil { + t.Fatalf("Map error: %v", res.Error()) + } + checkColumns(t, res, "name", "age", "upper") + checkColumn(t, res, "upper", []string{"ALICE", "BOB"}) +} + +func TestMapReplaceSourceColumn(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 2}, "s", []string{"a", "b"}) + res := df.Map("n", "n", func(v any) any { return v.(int64) * 100 }) + checkColumns(t, res, "n", "s") + checkColumn(t, res, "n", []int64{100, 200}) +} + +func TestMapTypeChange(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 2}) + res := df.Map("n", "label", func(v any) any { + return "n=" + string(rune('0'+v.(int64))) + }) + ct, _ := res.GetColumnType("label") + if ct != StringType { + t.Errorf("mapped column type = %v, want StringType", ct) + } + checkColumn(t, res, "label", []string{"n=1", "n=2"}) +} + +func TestMapErrors(t *testing.T) { + df := dfFromSeries(t, "a", []int64{1}) + if res := df.Map("ghost", "out", func(v any) any { return v }); res.Error() == nil { + t.Error("missing source column should error") + } + if res := df.Map("a", "", func(v any) any { return v }); res.Error() == nil { + t.Error("empty target column name should error") + } + if res := df.Map("a", "out", nil); res.Error() == nil { + t.Error("nil fn should error") + } +} + +func TestMapRecoversPanic(t *testing.T) { + df := dfFromSeries(t, "a", []string{"x"}) + res := df.Map("a", "out", func(v any) any { + return v.(int64) // wrong type assertion panics + }) + if res.Error() == nil { + t.Fatal("panic in Map fn must surface as an error") + } +} + +// --- FillNA --- + +func TestFillNAInt64(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 0, 3, 0}) + res := df.FillNA("n", int64(-1)) + checkColumn(t, res, "n", []int64{1, -1, 3, -1}) + + // Plain int coerces. + res2 := df.FillNA("n", 9) + checkColumn(t, res2, "n", []int64{1, 9, 3, 9}) +} + +func TestFillNAFloat64(t *testing.T) { + df := dfFromSeries(t, "f", []float64{1.5, 0, 2.5}) + res := df.FillNA("f", 99.9) + checkColumn(t, res, "f", []float64{1.5, 99.9, 2.5}) + + // Int value coerces to float column. + res2 := df.FillNA("f", 7) + checkColumn(t, res2, "f", []float64{1.5, 7, 2.5}) +} + +func TestFillNAString(t *testing.T) { + df := dfFromSeries(t, "s", []string{"a", "", "c"}) + res := df.FillNA("s", "missing") + checkColumn(t, res, "s", []string{"a", "missing", "c"}) +} + +func TestFillNABool(t *testing.T) { + df := dfFromSeries(t, "b", []bool{true, false}) + res := df.FillNA("b", true) + checkColumn(t, res, "b", []bool{true, true}) +} + +func TestFillNATime(t *testing.T) { + ts := time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC) + fill := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) + var zero time.Time + df := dfFromSeries(t, "t", []time.Time{ts, zero}) + res := df.FillNA("t", fill) + checkColumn(t, res, "t", []time.Time{ts, fill}) +} + +func TestFillNAErrors(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 0}, "s", []string{"a", ""}) + + if res := df.FillNA("ghost", 1); res.Error() == nil { + t.Error("missing column should error") + } + if res := df.FillNA("n", "not a number"); res.Error() == nil { + t.Error("string fill on int column should error") + } + if res := df.FillNA("n", 2.5); res.Error() == nil { + t.Error("fractional fill on int column should error") + } + if res := df.FillNA("s", 42); res.Error() == nil { + t.Error("int fill on string column should error") + } + if res := df.FillNA("n", nil); res.Error() == nil { + t.Error("nil fill value should error") + } +} + +func TestFillNADoesNotMutateReceiver(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 0}) + _ = df.FillNA("n", int64(5)) + checkColumn(t, df, "n", []int64{1, 0}) +} + +// --- DropNA --- + +func TestDropNASingleColumn(t *testing.T) { + df := dfFromSeries(t, + "n", []int64{1, 0, 3}, + "s", []string{"a", "b", ""}, + ) + res := df.DropNA("n") + checkColumn(t, res, "n", []int64{1, 3}) + checkColumn(t, res, "s", []string{"a", ""}) +} + +func TestDropNAMultipleColumns(t *testing.T) { + df := dfFromSeries(t, + "n", []int64{1, 0, 3, 4}, + "s", []string{"a", "b", "", "d"}, + ) + res := df.DropNA("n", "s") + // Row 1 dropped (n=0), row 2 dropped (s=""). + checkColumn(t, res, "n", []int64{1, 4}) + checkColumn(t, res, "s", []string{"a", "d"}) +} + +func TestDropNAAllColumnsByDefault(t *testing.T) { + ts := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + var zero time.Time + df := dfFromSeries(t, + "n", []int64{1, 0, 3, 4}, + "f", []float64{1, 2, 0, 4}, + "t", []time.Time{ts, ts, ts, zero}, + ) + res := df.DropNA() + // Only row 0 has no zero values. + checkColumn(t, res, "n", []int64{1}) + checkColumn(t, res, "f", []float64{1}) +} + +func TestDropNAPreservesOrder(t *testing.T) { + df := dfFromSeries(t, "n", []int64{5, 0, 3, 0, 1}) + res := df.DropNA("n") + checkColumn(t, res, "n", []int64{5, 3, 1}) +} + +func TestDropNANoMatches(t *testing.T) { + df := dfFromSeries(t, "n", []int64{0, 0}) + res := df.DropNA("n") + if res.Error() != nil { + t.Fatalf("DropNA error: %v", res.Error()) + } + if res.Len() != 0 { + t.Errorf("expected 0 rows, got %d", res.Len()) + } + if !res.HasColumn("n") { + t.Error("columns should survive even when all rows drop") + } +} + +func TestDropNAErrors(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1}) + if res := df.DropNA("ghost"); res.Error() == nil { + t.Error("missing column should error") + } + bad := df.Filter("ghost", "==", 1) + if res := bad.DropNA(); res.Error() == nil { + t.Error("DropNA on errored frame should carry the error") + } +} + +func TestDropNADoesNotMutateReceiver(t *testing.T) { + df := dfFromSeries(t, "n", []int64{1, 0}) + _ = df.DropNA("n") + if df.Len() != 2 { + t.Error("DropNA mutated the receiver") + } +} + +// --- chaining transforms --- + +func TestTransformChain(t *testing.T) { + df := dfFromSeries(t, + "name", []string{"alice", "", "carol"}, + "salary", []float64{1000, 2000, 0}, + ) + res := df. + FillNA("name", "unknown"). + DropNA("salary"). + Apply("bonus", func(r map[string]any) any { + return r["salary"].(float64) * 0.1 + }) + if res.Error() != nil { + t.Fatalf("chain error: %v", res.Error()) + } + checkColumn(t, res, "name", []string{"alice", "unknown"}) + checkColumn(t, res, "bonus", []float64{100, 200}) +} + +func TestApplyAllNumericWidths(t *testing.T) { + df := dfFromSeries(t, "n", []int64{0, 1, 2, 3, 4, 5, 6, 7, 8}) + res := df.Apply("out", func(r map[string]any) any { + switch r["n"].(int64) { + case 0: + return int8(10) + case 1: + return int16(11) + case 2: + return int32(12) + case 3: + return uint(13) + case 4: + return uint8(14) + case 5: + return uint16(15) + case 6: + return uint32(16) + case 7: + return int(17) + default: + return int64(18) + } + }) + if res.Error() != nil { + t.Fatalf("Apply error: %v", res.Error()) + } + checkColumn(t, res, "out", []int64{10, 11, 12, 13, 14, 15, 16, 17, 18}) + + // float32 promotes to Float64Type alongside ints. + res2 := df.Apply("f", func(r map[string]any) any { + if r["n"].(int64) == 0 { + return float32(1.5) + } + return r["n"].(int64) + }) + ct, _ := res2.GetColumnType("f") + if ct != Float64Type { + t.Errorf("float32+int mix type = %v, want Float64Type", ct) + } + + // Mixed numeric + string exercises canonical float/int formatting. + res3 := df.Head(3).Apply("m", func(r map[string]any) any { + switch r["n"].(int64) { + case 0: + return 2.5 + case 1: + return "txt" + default: + return uint16(9) + } + }) + checkColumn(t, res3, "m", []string{"2.5", "txt", "9"}) +}