internal/cmd/query: unify JSON output for aggregated and non-aggregated results

Remove the special-cased aggregation branch from the -f=json code path.
All query results now use tableRowAtIndex() to produce NDJSON objects,
regardless of whether the result has aggregation metadata.

Previously, aggregated results without groups output only the first
field's bare scalar value (table.Columns[0][0]), silently discarding
other fields. Aggregated results with groups errored entirely.

A server-side optimizer change (axiomhq/axiom#16816) now preserves
aggregation metadata on fields that were previously stripped, causing
multi-field queries (e.g. summarize ... | project a, b, c) to hit the
broken scalar path.

The fix deletes the aggregation branch and lets all results flow through
the existing, correct NDJSON loop. Table output format is unaffected.

Author: lukasmalkmus

internal/cmd/query/query.go

@@ -290,19 +290,9 @@ func run(ctx context.Context, opts *options) error {
 
 	table := res.Tables[0]
 
-	// Deal with JSON output format. It only works for non-aggregated results OR
-	// an aggregated result which produces a single value (it has no groups).
+	// Deal with JSON output format. Outputs NDJSON (one JSON object per row)
+	// for all results, regardless of whether they are aggregated.
 	if opts.Format == iofmt.JSON.String() {
-		if tableHasAggregation(table) {
-			if len(table.Groups) > 1 || (len(table.Columns) > 1 && len(table.Columns[0]) > 1) {
-				return errors.New("JSON output format is not supported for aggregated results with groups")
-			}
-			if opts.IO.IsStdoutTTY() {
-				fmt.Fprint(opts.IO.Out(), headerText)
-			}
-			return iofmt.FormatToJSON(opts.IO.Out(), table.Columns[0][0], opts.IO.ColorEnabled())
-		}
-
 		if opts.IO.IsStdoutTTY() {
 			fmt.Fprint(opts.IO.Out(), headerText)
 		}

internal/cmd/query/query_test.go

@@ -3,6 +3,8 @@ package query
 import (
 	"testing"
 	"time"
+
+	"github.com/axiomhq/axiom-go/axiom/query"
 )
 
 func TestParseDuration(t *testing.T) {
@@ -54,6 +56,83 @@ func TestParseDuration(t *testing.T) {
 	}
 }
 
+func TestTableRowAtIndex(t *testing.T) {
+	table := query.Table{
+		Fields: []query.Field{
+			{Name: "organizations"},
+			{Name: "datasets"},
+			{Name: "spans"},
+		},
+		Columns: []query.Column{
+			{"102 (+2)"},
+			{"132 (+3)"},
+			{"49.3M (+2.5M)"},
+		},
+	}
+
+	row := tableRowAtIndex(table, 0)
+
+	if got := row["organizations"]; got != "102 (+2)" {
+		t.Errorf("organizations = %v, want %q", got, "102 (+2)")
+	}
+	if got := row["datasets"]; got != "132 (+3)" {
+		t.Errorf("datasets = %v, want %q", got, "132 (+3)")
+	}
+	if got := row["spans"]; got != "49.3M (+2.5M)" {
+		t.Errorf("spans = %v, want %q", got, "49.3M (+2.5M)")
+	}
+}
+
+func TestTableRowAtIndexMultipleRows(t *testing.T) {
+	table := query.Table{
+		Fields: []query.Field{
+			{Name: "status"},
+			{Name: "count_", Aggregation: &query.Aggregation{Op: query.OpCount}},
+		},
+		Columns: []query.Column{
+			{"ok", "err"},
+			{float64(42), float64(3)},
+		},
+	}
+
+	row0 := tableRowAtIndex(table, 0)
+	if got := row0["status"]; got != "ok" {
+		t.Errorf("row0 status = %v, want %q", got, "ok")
+	}
+	if got := row0["count_"]; got != float64(42) {
+		t.Errorf("row0 count_ = %v, want %v", got, 42)
+	}
+
+	row1 := tableRowAtIndex(table, 1)
+	if got := row1["status"]; got != "err" {
+		t.Errorf("row1 status = %v, want %q", got, "err")
+	}
+	if got := row1["count_"]; got != float64(3) {
+		t.Errorf("row1 count_ = %v, want %v", got, 3)
+	}
+}
+
+func TestTableHasAggregation(t *testing.T) {
+	noAgg := query.Table{
+		Fields: []query.Field{
+			{Name: "org"},
+			{Name: "spans"},
+		},
+	}
+	if tableHasAggregation(noAgg) {
+		t.Error("expected no aggregation for plain fields")
+	}
+
+	withAgg := query.Table{
+		Fields: []query.Field{
+			{Name: "count_", Aggregation: &query.Aggregation{Op: query.OpCount}},
+		},
+	}
+	if !tableHasAggregation(withAgg) {
+		t.Error("expected aggregation when field has Aggregation set")
+	}
+}
+
 func TestParseDurationErrors(t *testing.T) {
 	tests := []string{
 		"",