From f3a1a006a9eb905cfad77a1ef00b68dbd81f0871 Mon Sep 17 00:00:00 2001 From: jay <38236622+JayJamieson@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:27:36 +1200 Subject: [PATCH 1/4] port over arrow code gen from roost --- arrgen/README.md | 188 +++++++ arrgen/arrgen/README.md | 209 ++++++++ arrgen/arrgen/cmd/arrgen/main.go | 102 ++++ arrgen/arrgen/doc.go | 51 ++ arrgen/arrgen/drift_test.go | 66 +++ arrgen/arrgen/example/example_test.go | 105 ++++ arrgen/arrgen/example/metric.go | 36 ++ arrgen/arrgen/example/metric_arrow.go | 126 +++++ arrgen/arrgen/generate.go | 335 ++++++++++++ arrgen/arrgen/generate_test.go | 222 ++++++++ arrgen/arrgen/go.mod | 35 ++ arrgen/arrgen/go.sum | 44 ++ arrgen/arrgen/internal/gentypes/alloc_test.go | 158 ++++++ arrgen/arrgen/internal/gentypes/bench_test.go | 219 ++++++++ .../internal/gentypes/equivalence_test.go | 449 ++++++++++++++++ .../gentypes/instrumented_race_test.go | 21 + .../internal/gentypes/instrumented_test.go | 24 + arrgen/arrgen/internal/gentypes/row_arrow.go | 475 +++++++++++++++++ arrgen/arrgen/internal/gentypes/types.go | 98 ++++ arrgen/arrgen/license_header.txt | 15 + arrgen/arrgen/mapping.go | 320 ++++++++++++ arrgen/arrgen/tag.go | 140 +++++ arrgen/arrgen/template.go | 138 +++++ arrgen/arrgen/testdata/basic.golden | 226 ++++++++ arrgen/arrgen/testdata/basic/basic.go | 37 ++ arrgen/arrgen/testdata/errors/errors.go | 126 +++++ arrgen/arrgen/testdata/pregen/pregen.go | 37 ++ arrgen/cmd/arrgen/main.go | 104 ++++ arrgen/doc.go | 56 ++ arrgen/drift_test.go | 72 +++ arrgen/example/example_test.go | 107 ++++ arrgen/example/metric.go | 42 ++ arrgen/example/metric_arrow.go | 134 +++++ arrgen/generate.go | 350 +++++++++++++ arrgen/generate_test.go | 228 ++++++++ arrgen/go.mod | 35 ++ arrgen/go.sum | 44 ++ arrgen/internal/gentypes/alloc_test.go | 167 ++++++ arrgen/internal/gentypes/bench_test.go | 221 ++++++++ arrgen/internal/gentypes/equivalence_test.go | 470 +++++++++++++++++ .../gentypes/instrumented_race_test.go | 21 + arrgen/internal/gentypes/instrumented_test.go | 24 + arrgen/internal/gentypes/row_arrow.go | 491 ++++++++++++++++++ arrgen/internal/gentypes/types.go | 103 ++++ arrgen/license_header.txt | 15 + arrgen/mapping.go | 330 ++++++++++++ arrgen/tag.go | 146 ++++++ arrgen/template.go | 147 ++++++ arrgen/testdata/basic.golden | 242 +++++++++ arrgen/testdata/basic/basic.go | 37 ++ arrgen/testdata/errors/errors.go | 128 +++++ arrgen/testdata/pregen/pregen.go | 37 ++ 52 files changed, 7753 insertions(+) create mode 100644 arrgen/README.md create mode 100644 arrgen/arrgen/README.md create mode 100644 arrgen/arrgen/cmd/arrgen/main.go create mode 100644 arrgen/arrgen/doc.go create mode 100644 arrgen/arrgen/drift_test.go create mode 100644 arrgen/arrgen/example/example_test.go create mode 100644 arrgen/arrgen/example/metric.go create mode 100644 arrgen/arrgen/example/metric_arrow.go create mode 100644 arrgen/arrgen/generate.go create mode 100644 arrgen/arrgen/generate_test.go create mode 100644 arrgen/arrgen/go.mod create mode 100644 arrgen/arrgen/go.sum create mode 100644 arrgen/arrgen/internal/gentypes/alloc_test.go create mode 100644 arrgen/arrgen/internal/gentypes/bench_test.go create mode 100644 arrgen/arrgen/internal/gentypes/equivalence_test.go create mode 100644 arrgen/arrgen/internal/gentypes/instrumented_race_test.go create mode 100644 arrgen/arrgen/internal/gentypes/instrumented_test.go create mode 100644 arrgen/arrgen/internal/gentypes/row_arrow.go create mode 100644 arrgen/arrgen/internal/gentypes/types.go create mode 100644 arrgen/arrgen/license_header.txt create mode 100644 arrgen/arrgen/mapping.go create mode 100644 arrgen/arrgen/tag.go create mode 100644 arrgen/arrgen/template.go create mode 100644 arrgen/arrgen/testdata/basic.golden create mode 100644 arrgen/arrgen/testdata/basic/basic.go create mode 100644 arrgen/arrgen/testdata/errors/errors.go create mode 100644 arrgen/arrgen/testdata/pregen/pregen.go create mode 100644 arrgen/cmd/arrgen/main.go create mode 100644 arrgen/doc.go create mode 100644 arrgen/drift_test.go create mode 100644 arrgen/example/example_test.go create mode 100644 arrgen/example/metric.go create mode 100644 arrgen/example/metric_arrow.go create mode 100644 arrgen/generate.go create mode 100644 arrgen/generate_test.go create mode 100644 arrgen/go.mod create mode 100644 arrgen/go.sum create mode 100644 arrgen/internal/gentypes/alloc_test.go create mode 100644 arrgen/internal/gentypes/bench_test.go create mode 100644 arrgen/internal/gentypes/equivalence_test.go create mode 100644 arrgen/internal/gentypes/instrumented_race_test.go create mode 100644 arrgen/internal/gentypes/instrumented_test.go create mode 100644 arrgen/internal/gentypes/row_arrow.go create mode 100644 arrgen/internal/gentypes/types.go create mode 100644 arrgen/license_header.txt create mode 100644 arrgen/mapping.go create mode 100644 arrgen/tag.go create mode 100644 arrgen/template.go create mode 100644 arrgen/testdata/basic.golden create mode 100644 arrgen/testdata/basic/basic.go create mode 100644 arrgen/testdata/errors/errors.go create mode 100644 arrgen/testdata/pregen/pregen.go diff --git a/arrgen/README.md b/arrgen/README.md new file mode 100644 index 000000000..6c231006a --- /dev/null +++ b/arrgen/README.md @@ -0,0 +1,188 @@ + + +# arrgen + +Zero-reflection Arrow encoders, generated from your Go structs. + +`arrgen` is the code-generation counterpart to +[`arrow/array/arreflect`](../arrow/array/arreflect). `arreflect` reads a +struct's `arrow:"..."` tags with reflection on every value; `arrgen` reads the +same tags once, when you run it, and writes typed Go source that appends struct +fields straight into typed Arrow builders. + +It lives in its own Go module, so it is opt-in twice over: nothing in +`github.com/apache/arrow-go/v18` imports it, and the code it emits imports only +`arrow`, `arrow/array` and `arrow/memory` - never `arrgen` itself. Generating is +a build-time step; your binary does not grow a dependency for it. + +## Quick start + +Tag a struct as you would for `arreflect`, and add a `go:generate` line: + +```go +package telemetry + +import "time" + +type Metric struct { + Time time.Time `arrow:"ts"` + Host string `arrow:"host"` + CPU float64 `arrow:"cpu"` + Value *float64 `arrow:"value"` // nullable + Secret string `arrow:"-"` // not a column +} + +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +``` + +```sh +go generate ./... # writes metric_arrow.go next to the type +``` + +Check the result in, as you would `stringer` or `easyjson` output, and +regenerate when the struct changes. If you would rather not add `arrgen` to your +module's dependencies at all, pin it in the directive instead: + +```go +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen@v0.1.0 -type Metric +``` + +### What you get + +```go +// One batch from a slice - the drop-in for arreflect.RecordFromSlice. +rec, err := telemetry.MetricRecordBatch(mem, metrics) + +// Or stream rows in and cut batches where you want them. +a := telemetry.NewMetricAppender(mem) +defer a.Release() +a.Reserve(batchSize) +for row := range rows { + a.Append(&row) + if a.Len() == batchSize { + rec := a.NewRecordBatch() + ... // hand it off + rec.Release() + a.Reserve(batchSize) + } +} + +telemetry.MetricSchema() // the schema, built once at init +a.Err() // first append error; only dictionary columns can fail +``` + +`Append` reads `v` synchronously and never retains it, so a caller draining a +stream can reuse one row variable for every row. + +## Flags + +| Flag | Meaning | +| --- | --- | +| `-type` | Struct type name. Repeatable, or comma-separated. Required. | +| `-output` | Output file, relative to `-dir`. Defaults to `_arrow.go`. | +| `-dir` | Package directory. Defaults to `.`, which is where `go:generate` runs. | +| `-header` | File whose contents are copied above the generated-code marker, for a license header. | + +## Supported fields + +`bool`, `int8/16/32/64`, `int`, `uint8/16/32/64`, `uint`, `float32/64`, +`string`, `[]byte`, `time.Time`, `time.Duration`, `decimal.Decimal32`, +`decimal.Decimal64`, `decimal128.Num`, `decimal256.Num`, and a pointer to any of +those for a nullable column. + +Tag options are `arreflect`'s: a leading column name, `-` to skip the field, the +temporal overrides `date32`, `date64`, `time32`, `time64` and `timestamp`, +`dict`, `view`, `large`, and `decimal(precision,scale)`. + +Anything else such as nested structs, slices other than `[]byte`, arrays, maps, +embedded fields, defined scalar types such as `type ID int64` is a +generate-time error naming the field. Those structs still encode fine with `arreflect` at runtime. + +The generator is equally strict about options that would do nothing: `date32` on +an `int64` field is an error here even though `arreflect` ignores it. + +## Equivalence with arreflect + +The point of matching `arreflect`'s tag dialect is that switching between the +two paths is a call-site change and nothing else. `internal/gentypes` holds a +fixture with one field for every supported column shape and asserts, on every +run, that the generated encoder and `arreflect` produce the same schema and the +same column data for the same input. + +## Performance + +`go test ./internal/gentypes/ -bench . -benchmem`, Go 1.25, linux/arm64. Batch +benchmarks encode 1024 rows per operation. + +| Benchmark | arreflect | arrgen | | +| --- | --- | --- | --- | +| `MetricBatch` (4 columns) | 56.9 µs, 98 allocs | 20.3 µs, 52 allocs | **2.8x faster** | +| `FixedBatch` (5 fixed-width columns) | 59.0 µs, 97 allocs | 17.4 µs, 46 allocs | **3.4x faster** | +| `RowBatch` (46 columns) | 843 µs, 3430 allocs | 413 µs, 3126 allocs | **2.0x faster** | +| `StreamAppend` (per row) | 58.3 ns | 18.6 ns, **0 allocs** | **3.1x faster** | +| `Schema` | 5.6 µs, 132 allocs | 1.5 ns, **0 allocs** | | + +Three things move the numbers: + +- **Field access.** The generated code loads `v.CPU` and calls + `Float64Builder.Append`. The reflection path resolves a field index, produces + a `reflect.Value`, and dispatches on the builder's dynamic type, per value. +- **Builder lookup.** Typed builders are resolved once, in the constructor, so + `Append` does no type assertions at all - not even the + `b.Field(i).(*array.Float64Builder)` a hand-written encoder usually repeats. +- **Schema construction.** The schema is a package-level variable rather than a + struct walk per call. + +Allocation counts converge on wide fixtures because both paths pay for the same +Arrow buffers, and because Arrow's dictionary builders allocate per appended +value regardless of who calls them. The place the difference is unambiguous is +the streaming path: with room reserved, `Append` allocates nothing, and +`TestStreamingAppendIsAllocationFree` asserts exactly that rather than leaving +it to a benchmark to imply. + +The streaming path also removes a constraint rather than a cost: `arreflect` +cannot encode anything until the caller has materialized the whole `[]T`, while +an appender can be fed a row at a time and cut into batches wherever you like. + +## Working on arrgen + +The repository root carries a `go.work` so this module builds against the +arrow-go tree next to it rather than the released version its `go.mod` names: + +```sh +go test ./arrgen/... # workspace: local arrow-go +cd arrgen && GOWORK=off go test ./... # standalone: released arrow-go +``` + +Both are expected to pass. `TestCheckedInFilesAreUpToDate` regenerates every +committed output in this module and fails if it differs, so a struct change +without a `go generate` is caught at the point of the edit. + +The golden file in `testdata` is the review surface for generator changes: +`go test ./arrgen/ -update` rewrites it, and the resulting diff is exactly what +users would see in their own regenerated code. + +### Releasing + +As a nested module, `arrgen` is versioned and tagged independently of the root +module: its tags are `arrgen/vX.Y.Z`, not `vX.Y.Z`. It cannot share the root's +`v18` line, because a module path without a `/vN` suffix is limited to v0 and +v1. Nothing in the release scripts tags it yet - that is a deliberate omission +for maintainers to decide on, not an oversight. diff --git a/arrgen/arrgen/README.md b/arrgen/arrgen/README.md new file mode 100644 index 000000000..41b88a024 --- /dev/null +++ b/arrgen/arrgen/README.md @@ -0,0 +1,209 @@ + + +# arrgen + +Zero-reflection Arrow encoders, generated from your Go structs. + +`arrgen` is the code-generation counterpart to +[`arrow/array/arreflect`](../arrow/array/arreflect). `arreflect` reads a +struct's `arrow:"..."` tags with reflection on every value; `arrgen` reads the +same tags once, when you run it, and writes typed Go source that appends struct +fields straight into typed Arrow builders. + +It lives in its own Go module: nothing in `github.com/apache/arrow-go/v18` +imports it, and the code it emits imports only `arrow`, `arrow/array` and +`arrow/memory` — never `arrgen` itself. Generating is a build-time step; your +binary does not grow a dependency for it. + +## Quick start + +Tag a struct as you would for `arreflect`, and add a `go:generate` line: + +```go +package telemetry + +import "time" + +type Metric struct { + Time time.Time `arrow:"ts"` + Host string `arrow:"host"` + CPU float64 `arrow:"cpu"` + Value *float64 `arrow:"value"` // nullable + Secret string `arrow:"-"` // not a column +} + +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +``` + +```sh +go generate ./... # writes metric_arrow.go next to the type +``` + +Check the result in, as you would `stringer` or `easyjson` output, and +regenerate when the struct changes. If you would rather not add `arrgen` to your +module's dependencies at all, pin it in the directive instead: + +```go +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen@v0.1.0 -type Metric +``` + +### What you get + +```go +// One batch from a slice - the drop-in for arreflect.RecordFromSlice. +rec, err := telemetry.MetricRecordBatch(mem, metrics) + +// Or stream rows in and cut batches where you want them. +a := telemetry.NewMetricAppender(mem) +defer a.Release() +a.Reserve(batchSize) +for row := range rows { + a.Append(&row) + if a.Len() == batchSize { + rec := a.NewRecordBatch() + ... // hand it off + rec.Release() + a.Reserve(batchSize) + } +} + +telemetry.MetricSchema() // the schema, built once at init +a.Err() // first append error; only dictionary columns can fail +``` + +`Append` reads `v` synchronously and never retains it, so a caller draining a +stream can reuse one row variable for every row. + +## Flags + +| Flag | Meaning | +| --- | --- | +| `-type` | Struct type name. Repeatable, or comma-separated. Required. | +| `-output` | Output file, relative to `-dir`. Defaults to `_arrow.go`. | +| `-dir` | Package directory. Defaults to `.`, which is where `go:generate` runs. | +| `-header` | File whose contents are copied above the generated-code marker, for a license header. | + +## Supported fields + +`bool`, `int8/16/32/64`, `int`, `uint8/16/32/64`, `uint`, `float32/64`, +`string`, `[]byte`, `time.Time`, `time.Duration`, `decimal.Decimal32`, +`decimal.Decimal64`, `decimal128.Num`, `decimal256.Num`, and a pointer to any of +those for a nullable column. + +Tag options are `arreflect`'s: a leading column name, `-` to skip the field, the +temporal overrides `date32`, `date64`, `time32`, `time64` and `timestamp`, +`dict`, `view`, `large`, and `decimal(precision,scale)`. + +Anything else — nested structs, slices other than `[]byte`, arrays, maps, +embedded fields, defined scalar types such as `type ID int64` — is a +**generate-time error naming the field**, never a silently dropped column. Those +structs still encode fine with `arreflect` at runtime. + +The generator is equally strict about options that would do nothing: `date32` on +an `int64` field is an error here even though `arreflect` ignores it. + +## Equivalence with arreflect + +The point of matching `arreflect`'s tag dialect is that switching between the +two paths is a call-site change and nothing else. `internal/gentypes` holds a +fixture with one field for every supported column shape and asserts, on every +run, that the generated encoder and `arreflect` produce the same schema and the +same column data for the same input. + +### Two columns where they differ, and why + +`arreflect`'s struct-field path dispatches on `reflect.Kind` before it checks +the types it special-cases, so a Go struct that Arrow models as a scalar reaches +its nested-struct branch, which finds only unexported fields and returns an +empty `struct<>`. The value is then dropped. This affects: + +| Field | `arreflect` infers | `arrgen` emits | +| --- | --- | --- | +| `time.Time` (untagged, or `,timestamp`) | `struct<>` | `timestamp[ns, tz=UTC]` | +| `decimal128.Num` / `decimal256.Num` (untagged) | `struct<>` | `decimal128` / `decimal256` | + +A tag that names an explicit Arrow type — `,date32`, `,decimal(20,3)` — rescues +the field, because `arreflect` overrides the inferred type from the tag +afterwards. `arreflect`'s own tests only pass these types as the top-level +element of a slice, where a different branch handles them correctly. + +`arrgen` emits what `arreflect.InferType[time.Time]` and its top-level +`FromSlice` path both agree on. `TestArreflectStructFieldGaps` pins the current +upstream behavior, so if `arrow/array/arreflect` is fixed, that test fails and +tells us to delete this section. + +## Performance + +`go test ./internal/gentypes/ -bench . -benchmem`, Go 1.25, linux/arm64. Batch +benchmarks encode 1024 rows per operation. + +| Benchmark | arreflect | arrgen | | +| --- | --- | --- | --- | +| `MetricBatch` (4 columns) | 56.9 µs, 98 allocs | 20.3 µs, 52 allocs | **2.8x faster** | +| `FixedBatch` (5 fixed-width columns) | 59.0 µs, 97 allocs | 17.4 µs, 46 allocs | **3.4x faster** | +| `RowBatch` (46 columns) | 843 µs, 3430 allocs | 413 µs, 3126 allocs | **2.0x faster** | +| `StreamAppend` (per row) | 58.3 ns | 18.6 ns, **0 allocs** | **3.1x faster** | +| `Schema` | 5.6 µs, 132 allocs | 1.5 ns, **0 allocs** | | + +Three things move the numbers: + +- **Field access.** The generated code loads `v.CPU` and calls + `Float64Builder.Append`. The reflection path resolves a field index, produces + a `reflect.Value`, and dispatches on the builder's dynamic type, per value. +- **Builder lookup.** Typed builders are resolved once, in the constructor, so + `Append` does no type assertions at all — not even the + `b.Field(i).(*array.Float64Builder)` a hand-written encoder usually repeats. +- **Schema construction.** The schema is a package-level variable rather than a + struct walk per call. + +Allocation counts converge on wide fixtures because both paths pay for the same +Arrow buffers, and because Arrow's dictionary builders allocate per appended +value regardless of who calls them. The place the difference is unambiguous is +the streaming path: with room reserved, `Append` allocates nothing, which +`TestStreamingAppendIsAllocationFree` asserts. + +The streaming path also removes a constraint: `arreflect` cannot encode anything +until the caller has materialized the whole `[]T`, while an appender can be fed +a row at a time and cut into batches wherever you like. + +## Working on arrgen + +The repository root carries a `go.work` so this module builds against the +arrow-go tree next to it rather than the released version its `go.mod` names: + +```sh +go test ./arrgen/... # workspace: local arrow-go +cd arrgen && GOWORK=off go test ./... # standalone: released arrow-go +``` + +Both are expected to pass. `TestCheckedInFilesAreUpToDate` regenerates every +committed output in this module and fails if it differs, so a struct change +without a `go generate` is caught at the point of the edit. + +The golden file in `testdata` is the review surface for generator changes: +`go test ./arrgen/ -update` rewrites it, and the resulting diff is exactly what +users would see in their own regenerated code. + +### Releasing + +As a nested module, `arrgen` is versioned and tagged independently of the root +module: its tags are `arrgen/vX.Y.Z`, not `vX.Y.Z`. It cannot share the root's +`v18` line, because a module path without a `/vN` suffix is limited to v0 and +v1. Nothing in the release scripts tags it yet. diff --git a/arrgen/arrgen/cmd/arrgen/main.go b/arrgen/arrgen/cmd/arrgen/main.go new file mode 100644 index 000000000..31d497b4f --- /dev/null +++ b/arrgen/arrgen/cmd/arrgen/main.go @@ -0,0 +1,102 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Command arrgen generates zero-reflection Arrow appenders for Go struct types. +// +// It reads the same arrow:"..." struct tags that +// github.com/apache/arrow-go/v18/arrow/array/arreflect interprets at runtime, +// but reads them once, when you run it, and writes typed Go source instead. +// +// Usage: +// +// arrgen -type Metric [-type Other] [-output metric_arrow.go] [-header LICENSE.txt] +// +// Typically it is invoked through go:generate, next to the type: +// +// //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +// +// Check the generated file in, as you would stringer or easyjson output. +package main + +import ( + "flag" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/apache/arrow-go/arrgen" +) + +func main() { + if err := run(os.Args[1:]); err != nil { + fmt.Fprintf(os.Stderr, "arrgen: %v\n", err) + os.Exit(1) + } +} + +// typeList collects -type values: both "-type A,B" and "-type A -type B" work. +type typeList []string + +func (t *typeList) String() string { return strings.Join(*t, ",") } + +func (t *typeList) Set(v string) error { + for _, p := range strings.Split(v, ",") { + if p = strings.TrimSpace(p); p != "" { + *t = append(*t, p) + } + } + return nil +} + +func run(argv []string) error { + fs := flag.NewFlagSet("arrgen", flag.ContinueOnError) + var types typeList + fs.Var(&types, "type", "struct type name to generate for; repeatable, or comma-separated (required)") + dir := fs.String("dir", ".", "directory of the package holding the types") + output := fs.String("output", "", "output file, relative to -dir (default _arrow.go)") + header := fs.String("header", "", "file whose contents are copied above the generated-code marker, e.g. a license header") + if err := fs.Parse(argv); err != nil { + return err + } + if len(types) == 0 { + return fmt.Errorf("-type is required") + } + + cfg := arrgen.Config{Dir: *dir, Types: types} + if *header != "" { + b, err := os.ReadFile(*header) + if err != nil { + return fmt.Errorf("reading -header: %w", err) + } + cfg.Header = string(b) + } + + src, err := arrgen.Generate(cfg) + if err != nil { + return err + } + + out := *output + if out == "" { + out = strings.ToLower(types[0]) + "_arrow.go" + } + path := filepath.Join(*dir, out) + if err := os.WriteFile(path, src, 0o644); err != nil { + return fmt.Errorf("writing %s: %w", path, err) + } + return nil +} diff --git a/arrgen/arrgen/doc.go b/arrgen/arrgen/doc.go new file mode 100644 index 000000000..682a70f79 --- /dev/null +++ b/arrgen/arrgen/doc.go @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package arrgen generates zero-reflection Arrow encoders for Go struct types. +// +// It is the code-generation counterpart to +// [github.com/apache/arrow-go/v18/arrow/array/arreflect]: arreflect interprets +// a struct's `arrow:"..."` tags with reflection on every value, while arrgen +// reads the same tags once, at generate time, and emits typed Go source that +// appends struct fields into typed Arrow builders. For every struct arrgen +// accepts, the generated schema equals arreflect.InferSchema and the generated +// record batch equals the one arreflect.RecordFromSlice builds. +// +// # Usage +// +// Put a go:generate directive next to the struct: +// +// //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +// +// then run `go generate ./...`, which writes metric_arrow.go next to the type. +// The emitted code imports only arrow, arrow/array and arrow/memory, never +// arrgen itself. +// +// # Supported field types +// +// bool, int8/16/32/64, int, uint8/16/32/64, uint, float32/64, string, []byte, +// time.Time, time.Duration, decimal.Decimal32, decimal.Decimal64, +// decimal128.Num and decimal256.Num, plus a pointer to any of those for a +// nullable column. +// +// Tag options mirror arreflect: a leading name, "-" to skip a field, the +// temporal overrides date32, date64, time32, time64 and timestamp, dict, view, +// large, and decimal(precision,scale). +// +// Anything else - nested structs, slices other than []byte, arrays, maps, +// embedded fields, named scalar types - is a generate-time error naming the +// field, never a silently dropped column. +package arrgen diff --git a/arrgen/arrgen/drift_test.go b/arrgen/arrgen/drift_test.go new file mode 100644 index 000000000..6c891684f --- /dev/null +++ b/arrgen/arrgen/drift_test.go @@ -0,0 +1,66 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen_test + +import ( + "bytes" + "os" + "path/filepath" + "testing" + + "github.com/apache/arrow-go/arrgen" +) + +// TestCheckedInFilesAreUpToDate regenerates every committed arrgen output in +// this module and fails if the result differs from what is on disk. +func TestCheckedInFilesAreUpToDate(t *testing.T) { + header, err := os.ReadFile("license_header.txt") + if err != nil { + t.Fatalf("reading license header: %v", err) + } + + // Mirrors the go:generate directives next to each type. + targets := []struct { + dir string + types []string + file string + }{ + {"example", []string{"Metric"}, "metric_arrow.go"}, + {"internal/gentypes", []string{"Row", "Fixed"}, "row_arrow.go"}, + } + + for _, target := range targets { + t.Run(target.file, func(t *testing.T) { + got, err := arrgen.Generate(arrgen.Config{ + Dir: target.dir, + Types: target.types, + Header: string(header), + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + path := filepath.Join(target.dir, target.file) + want, err := os.ReadFile(path) + if err != nil { + t.Fatalf("reading %s: %v", path, err) + } + if !bytes.Equal(got, want) { + t.Errorf("%s is out of date; run go generate ./...", path) + } + }) + } +} diff --git a/arrgen/arrgen/example/example_test.go b/arrgen/arrgen/example/example_test.go new file mode 100644 index 000000000..5c7cbcb8c --- /dev/null +++ b/arrgen/arrgen/example/example_test.go @@ -0,0 +1,105 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package example_test + +import ( + "fmt" + "time" + + "github.com/apache/arrow-go/arrgen/example" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// Example encodes a slice of rows in one call, replacing +// arreflect.RecordFromSlice[Metric](metrics, mem). +func Example() { + cpuLoad := 0.75 + metrics := []example.Metric{ + {Time: time.Unix(0, 0).UTC(), Host: "web-1", CPU: 0.5, Value: &cpuLoad}, + {Time: time.Unix(1, 0).UTC(), Host: "web-2", CPU: 1.5}, // Value is nil: a null + } + + rec, err := example.MetricRecordBatch(memory.DefaultAllocator, metrics) + if err != nil { + fmt.Println("encode:", err) + return + } + defer rec.Release() + + fmt.Println("rows:", rec.NumRows()) + for i, col := range rec.Columns() { + fmt.Printf("%s: %v\n", rec.Schema().Field(i).Name, col) + } + // Output: + // rows: 2 + // ts: [0 1000000000] + // host: ["web-1" "web-2"] + // cpu: [0.5 1.5] + // value: [0.75 (null)] +} + +// ExampleMetricAppender streams rows in one at a time and cuts a batch at a +// size boundary. The appender never retains a row, so the loop reuses a single +// variable and allocates nothing. +func ExampleMetricAppender() { + const batchSize = 2 + + a := example.NewMetricAppender(memory.DefaultAllocator) + defer a.Release() + a.Reserve(batchSize) + + var row example.Metric + for i := 0; i < 5; i++ { + row = example.Metric{ + Time: time.Unix(int64(i), 0).UTC(), + Host: fmt.Sprintf("web-%d", i), + CPU: float64(i), + } + a.Append(&row) + + if a.Len() == batchSize { + rec := a.NewRecordBatch() + fmt.Println("batch of", rec.NumRows()) + rec.Release() + a.Reserve(batchSize) + } + } + if a.Len() > 0 { + rec := a.NewRecordBatch() + fmt.Println("final batch of", rec.NumRows()) + rec.Release() + } + if err := a.Err(); err != nil { + fmt.Println("append:", err) + } + // Output: + // batch of 2 + // batch of 2 + // final batch of 1 +} + +// ExampleMetricSchema shows the schema the generator resolved from the tags. +func ExampleMetricSchema() { + for _, f := range example.MetricSchema().Fields() { + fmt.Printf("%s %s nullable=%t\n", f.Name, f.Type, f.Nullable) + } + // Output: + // ts timestamp[ns, tz=UTC] nullable=false + // host utf8 nullable=false + // cpu float64 nullable=false + // value float64 nullable=true +} diff --git a/arrgen/arrgen/example/metric.go b/arrgen/arrgen/example/metric.go new file mode 100644 index 000000000..60bf341f8 --- /dev/null +++ b/arrgen/arrgen/example/metric.go @@ -0,0 +1,36 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package example is a worked, checked-in demonstration of arrgen. +// +// Metric carries the arrow struct tags; metric_arrow.go next to it is what +// arrgen emitted from them, committed as stringer or easyjson output would be. +// Regenerate it with `go generate ./...`. +package example + +import "time" + +// Metric is a small row: a timestamp, two scalar columns, a nullable pointer +// column, and a field kept out of Arrow entirely. +type Metric struct { + Time time.Time `arrow:"ts"` + Host string `arrow:"host"` + CPU float64 `arrow:"cpu"` + Value *float64 `arrow:"value"` // nullable: a nil pointer appends null + Secret string `arrow:"-"` // never leaves the process +} + +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -header ../license_header.txt diff --git a/arrgen/arrgen/example/metric_arrow.go b/arrgen/arrgen/example/metric_arrow.go new file mode 100644 index 000000000..1909fd3a5 --- /dev/null +++ b/arrgen/arrgen/example/metric_arrow.go @@ -0,0 +1,126 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by arrgen. DO NOT EDIT. + +package example + +import ( + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// metricArrowSchema is the Arrow schema of Metric, resolved from its arrow +// struct tags. +var metricArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "host", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: "cpu", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "value", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, +}, nil) + +// MetricSchema returns the Arrow schema encoded by MetricAppender. +func MetricSchema() *arrow.Schema { return metricArrowSchema } + +// MetricAppender converts Metric values into Arrow record batches +// without reflection. +// +// It is not safe for concurrent use. Release it when done. +type MetricAppender struct { + rb *array.RecordBuilder + b0 *array.TimestampBuilder + b1 *array.StringBuilder + b2 *array.Float64Builder + b3 *array.Float64Builder + err error +} + +// NewMetricAppender returns an appender that builds Metric batches with mem, +// or memory.DefaultAllocator when mem is nil. The typed field builders are +// resolved once here, so Append does no type assertions. +func NewMetricAppender(mem memory.Allocator) *MetricAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, metricArrowSchema) + return &MetricAppender{ + rb: rb, + b0: rb.Field(0).(*array.TimestampBuilder), + b1: rb.Field(1).(*array.StringBuilder), + b2: rb.Field(2).(*array.Float64Builder), + b3: rb.Field(3).(*array.Float64Builder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *MetricAppender) Schema() *arrow.Schema { return metricArrowSchema } + +// RecordBuilder returns the underlying builder. Appending through it directly +// is allowed as long as every column stays the same length. +func (a *MetricAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends keeps the append path free of allocations. +func (a *MetricAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller can reuse a single Metric variable across calls. +func (a *MetricAppender) Append(v *Metric) { + a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) + a.b1.Append(v.Host) + a.b2.Append(v.CPU) + if v.Value == nil { + a.b3.AppendNull() + } else { + a.b3.Append(*v.Value) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *MetricAppender) AppendSlice(vs []Metric) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch returns the appended rows as a record batch and resets the +// appender. The caller must Release the batch. +func (a *MetricAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail. The error is sticky: NewRecordBatch does not clear it. +func (a *MetricAppender) Err() error { return a.err } + +// Release releases the appender's builders. +func (a *MetricAppender) Release() { a.rb.Release() } + +// MetricRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Metric](vs, mem). The +// caller must Release the batch. +func MetricRecordBatch(mem memory.Allocator, vs []Metric) (arrow.RecordBatch, error) { + a := NewMetricAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} diff --git a/arrgen/arrgen/generate.go b/arrgen/arrgen/generate.go new file mode 100644 index 000000000..91e63de6d --- /dev/null +++ b/arrgen/arrgen/generate.go @@ -0,0 +1,335 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import ( + "bytes" + "fmt" + "go/format" + "go/token" + "go/types" + "reflect" + "strings" + "unicode" + + "golang.org/x/tools/go/packages" +) + +// Config describes one generator run. +type Config struct { + // Dir is the directory of the package holding the struct types. go:generate + // runs in the package directory, so the command leaves this as ".". + Dir string + // Types are the struct type names to generate for, in output order. + Types []string + // Header, when non-empty, is emitted verbatim above the generated-code + // marker. Use it for a license header; it is expected to be Go comments. + Header string +} + +// Generate resolves cfg.Types in the package at cfg.Dir and returns the +// formatted source of their Arrow appenders. It does not write to disk. +// +// Output depends on nothing but the package's source, so regenerating an +// unchanged package reproduces the file byte for byte. +func Generate(cfg Config) ([]byte, error) { + if len(cfg.Types) == 0 { + return nil, fmt.Errorf("arrgen: no types requested") + } + dir := cfg.Dir + if dir == "" { + dir = "." + } + pkg, loadErrs, err := loadPackage(dir) + if err != nil { + return nil, err + } + + file := genFile{Header: strings.TrimRight(cfg.Header, "\n"), Package: pkg.Name} + declared := make(map[string]string) // generated identifier -> type that claimed it + for _, name := range cfg.Types { + gt, err := genForType(pkg, name) + if err != nil { + return nil, withLoadContext(fmt.Errorf("arrgen: %s: %w", name, err), loadErrs) + } + if err := claimIdents(declared, name, gt); err != nil { + return nil, fmt.Errorf("arrgen: %s: %w", name, err) + } + file.NeedTime = file.NeedTime || gt.needTime + file.Types = append(file.Types, gt) + } + + var buf bytes.Buffer + if err := fileTemplate.Execute(&buf, file); err != nil { + return nil, fmt.Errorf("arrgen: rendering template: %w", err) + } + src, err := format.Source(buf.Bytes()) + if err != nil { + return nil, fmt.Errorf("arrgen: formatting generated source: %w\n--- raw ---\n%s", err, buf.String()) + } + return src, nil +} + +// loadPackage type-checks the package in dir and returns it along with any +// errors the type checker reported. Those errors are returned rather than +// raised: the package a generator runs in usually does not compile yet, since +// the call sites are written before the generated file exists. +func loadPackage(dir string) (*packages.Package, []string, error) { + cfg := &packages.Config{ + Mode: packages.NeedName | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax, + Dir: dir, + } + pkgs, err := packages.Load(cfg, ".") + if err != nil { + return nil, nil, fmt.Errorf("arrgen: loading package in %s: %w", dir, err) + } + if len(pkgs) != 1 { + return nil, nil, fmt.Errorf("arrgen: expected exactly one package in %s, found %d", dir, len(pkgs)) + } + var errs []string + for _, e := range pkgs[0].Errors { + errs = append(errs, e.Error()) + } + if pkgs[0].Types == nil { + return nil, nil, fmt.Errorf("arrgen: package in %s could not be type-checked: %s", dir, strings.Join(errs, "; ")) + } + return pkgs[0], errs, nil +} + +// withLoadContext appends the package's type errors to err as a likely cause. +func withLoadContext(err error, loadErrs []string) error { + if len(loadErrs) == 0 { + return err + } + const maxShown = 5 + shown := loadErrs + suffix := "" + if len(shown) > maxShown { + shown, suffix = shown[:maxShown], fmt.Sprintf("\n ... and %d more", len(loadErrs)-maxShown) + } + return fmt.Errorf("%w\nthe package does not type-check, which may be the cause:\n %s%s", + err, strings.Join(shown, "\n "), suffix) +} + +func genForType(pkg *packages.Package, name string) (genType, error) { + st, err := findStruct(pkg, name) + if err != nil { + return genType{}, err + } + cols, err := collectColumns(st) + if err != nil { + return genType{}, err + } + if len(cols) == 0 { + return genType{}, fmt.Errorf("struct has no Arrow columns") + } + + exported := token.IsExported(name) + gt := genType{ + GoName: name, + SchemaVar: lowerFirst(name) + "ArrowSchema", + SchemaFunc: caseAs(exported, name+"Schema"), + AppenderType: caseAs(exported, name+"Appender"), + CtorName: caseAs(exported, "New"+upperFirst(name)+"Appender"), + BatchFunc: caseAs(exported, name+"RecordBatch"), + } + for i, c := range cols { + gt.needTime = gt.needTime || c.spec.needsTime + gt.AnyFallible = gt.AnyFallible || c.spec.fallible + gt.Fields = append(gt.Fields, genField{ + Index: i, + Name: c.name, + GoField: c.goField, + Nullable: c.nullable, + ArrowType: c.spec.arrowType, + BuilderType: c.spec.builderType, + BuilderVar: fmt.Sprintf("b%d", i), + AppendStmt: renderAppend(i, c), + }) + } + return gt, nil +} + +// claimIdents records the package-level names a type's generated code will +// declare, rejecting a second type that would declare the same one. Two types +// differing only in the case of their first letter collide this way. +func claimIdents(declared map[string]string, name string, gt genType) error { + for _, ident := range []string{gt.SchemaVar, gt.SchemaFunc, gt.AppenderType, gt.CtorName, gt.BatchFunc} { + if prev, taken := declared[ident]; taken { + return fmt.Errorf("generated name %s collides with the one generated for %s; generate the two types into separate files", ident, prev) + } + declared[ident] = name + } + return nil +} + +func findStruct(pkg *packages.Package, name string) (*types.Struct, error) { + obj := pkg.Types.Scope().Lookup(name) + if obj == nil { + return nil, fmt.Errorf("type not found in package %s", pkg.Name) + } + named, ok := obj.Type().(*types.Named) + if !ok { + return nil, fmt.Errorf("not a named type") + } + st, ok := named.Underlying().(*types.Struct) + if !ok { + return nil, fmt.Errorf("not a struct type") + } + return st, nil +} + +// column is one resolved struct field. +type column struct { + name string + goField string + nullable bool + spec colSpec +} + +// collectColumns walks the struct's fields in declaration order, which is the +// order arreflect settles on for a flat struct. Embedded fields are rejected +// rather than promoted, since arreflect resolves promoted names by breadth and +// tag. +func collectColumns(st *types.Struct) ([]column, error) { + var cols []column + seen := make(map[string]string, st.NumFields()) + + for i := 0; i < st.NumFields(); i++ { + f := st.Field(i) + if f.Anonymous() { + return nil, fmt.Errorf("field %s: embedded fields are not supported; give the field a name or encode the struct with arreflect", f.Name()) + } + if !f.Exported() { + continue + } + + tag, hasTag := reflect.StructTag(st.Tag(i)).Lookup("arrow") + var opts tagOpts + if hasTag { + var err error + if opts, err = parseTag(tag); err != nil { + return nil, fmt.Errorf("field %s: %w", f.Name(), err) + } + } + if opts.Skip { + continue + } + if err := opts.validate(); err != nil { + return nil, fmt.Errorf("field %s: %w", f.Name(), err) + } + + name := opts.Name + if name == "" { + name = f.Name() + } + if prev, dup := seen[name]; dup { + return nil, fmt.Errorf("fields %s and %s both map to column %q", prev, f.Name(), name) + } + seen[name] = f.Name() + + spec, nullable, err := resolveColumn(f.Type(), opts) + if err != nil { + return nil, fmt.Errorf("field %s: %w", f.Name(), err) + } + cols = append(cols, column{name: name, goField: f.Name(), nullable: nullable, spec: spec}) + } + return cols, nil +} + +// renderAppend emits the append statement for one column, wrapping it in the +// null guards arreflect applies: a nil pointer is null, and so is a nil []byte +// even when the column itself is not nullable. +func renderAppend(idx int, c column) string { + bld := fmt.Sprintf("a.b%d", idx) + val, recv := "v."+c.goField, "v."+c.goField + if c.nullable { + val, recv = "*v."+c.goField, "(*v."+c.goField+")" + } + + stmt := c.spec.appendStmt(bld, recv, val) + if c.spec.fallible { + stmt = "a.setErr(" + stmt + ")" + } + + var guard string + switch { + case c.nullable && c.spec.nilable: + guard = fmt.Sprintf("v.%s == nil || *v.%s == nil", c.goField, c.goField) + case c.nullable, c.spec.nilable: + guard = fmt.Sprintf("v.%s == nil", c.goField) + default: + return stmt + } + + // A statement that is already a block (a time-of-day column scopes a local) + // sheds its braces so the guard's else does not nest two sets. + inner := strings.TrimSuffix(strings.TrimPrefix(stmt, "{\n"), "\n}") + return fmt.Sprintf("if %s {\n%s.AppendNull()\n} else {\n%s\n}", guard, bld, inner) +} + +type genFile struct { + Header string + Package string + NeedTime bool + Types []genType +} + +type genType struct { + GoName string + SchemaVar string + SchemaFunc string + AppenderType string + CtorName string + BatchFunc string + Fields []genField + AnyFallible bool + + needTime bool +} + +type genField struct { + Index int + Name string + GoField string + Nullable bool + ArrowType string + BuilderType string + BuilderVar string + AppendStmt string +} + +func lowerFirst(s string) string { return mapFirst(s, unicode.ToLower) } +func upperFirst(s string) string { return mapFirst(s, unicode.ToUpper) } + +func mapFirst(s string, f func(rune) rune) string { + if s == "" { + return s + } + r := []rune(s) + r[0] = f(r[0]) + return string(r) +} + +// caseAs gives a generated identifier the same visibility as the struct it was +// derived from. +func caseAs(exported bool, s string) string { + if exported { + return upperFirst(s) + } + return lowerFirst(s) +} diff --git a/arrgen/arrgen/generate_test.go b/arrgen/arrgen/generate_test.go new file mode 100644 index 000000000..a864013a1 --- /dev/null +++ b/arrgen/arrgen/generate_test.go @@ -0,0 +1,222 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen_test + +import ( + "bytes" + "flag" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/apache/arrow-go/arrgen" +) + +var update = flag.Bool("update", false, "rewrite the golden files instead of comparing against them") + +const goldenPath = "testdata/basic.golden" + +// TestGenerateGolden compares the whole emitted file against a checked-in copy, +// so the diff of a generator change shows what users would see in their own +// regenerated code. +func TestGenerateGolden(t *testing.T) { + header, err := os.ReadFile("license_header.txt") + if err != nil { + t.Fatalf("reading license header: %v", err) + } + got, err := arrgen.Generate(arrgen.Config{ + Dir: "testdata/basic", + Types: []string{"Metric", "reading"}, + Header: string(header), + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + + if *update { + if err := os.WriteFile(goldenPath, got, 0o644); err != nil { + t.Fatalf("writing golden: %v", err) + } + return + } + + want, err := os.ReadFile(goldenPath) + if err != nil { + t.Fatalf("reading golden: %v (run: go test ./... -update)", err) + } + if !bytes.Equal(got, want) { + t.Errorf("generated output differs from %s; rerun with -update to accept\n--- got ---\n%s", goldenPath, got) + } +} + +// TestGenerateIsDeterministic checks that the same package in produces the same +// bytes out, which TestCheckedInFilesAreUpToDate depends on. +func TestGenerateIsDeterministic(t *testing.T) { + cfg := arrgen.Config{Dir: "testdata/basic", Types: []string{"Metric", "reading"}} + first, err := arrgen.Generate(cfg) + if err != nil { + t.Fatalf("Generate: %v", err) + } + second, err := arrgen.Generate(cfg) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if !bytes.Equal(first, second) { + t.Error("two runs over the same package produced different output") + } +} + +// TestGenerateHeaderOptional checks that the generated-code marker leads the +// file when no header is configured. +func TestGenerateHeaderOptional(t *testing.T) { + got, err := arrgen.Generate(arrgen.Config{Dir: "testdata/basic", Types: []string{"Metric"}}) + if err != nil { + t.Fatalf("Generate: %v", err) + } + const marker = "// Code generated by arrgen. DO NOT EDIT." + if !bytes.HasPrefix(got, []byte(marker)) { + t.Errorf("output does not start with %q:\n%s", marker, first(got, 200)) + } +} + +// TestGenerateRejects covers the structs arrgen refuses rather than encode +// wrongly or drop a column from. +func TestGenerateRejects(t *testing.T) { + tests := []struct { + typ string + want string + }{ + {"Embedded", "embedded fields are not supported"}, + {"Nested", "nested struct column"}, + {"SliceField", "list column for []int64 is not supported"}, + {"ArrayField", "is not supported"}, + {"MapField", "is not supported"}, + {"NamedScalar", "named type"}, + {"UnknownOption", `unknown option "nope"`}, + {"BadDecimalTag", "is not an integer"}, + {"ShortDecimalTag", "expected decimal(precision,scale)"}, + {"DuplicateNames", `both map to column "same"`}, + {"TemporalOnInt", "only valid on a time.Time field"}, + {"DecimalOnInt", "only valid on a decimal field"}, + {"DictOnBool", "dict is not supported on bool"}, + {"ViewOnInt", "view has no effect"}, + {"LargeOnInt", "large has no effect"}, + {"RunEndEncoded", "ree is not supported"}, + {"DictAndView", "at most one of dict, view, ree"}, + {"DictAndLarge", "cannot be combined with large"}, + {"NoColumns", "no Arrow columns"}, + {"DoublePointer", "multi-level pointer"}, + {"TimeSlice", "list column"}, + {"NotAStruct", "not a struct type"}, + {"Absent", "type not found"}, + } + + for _, tt := range tests { + t.Run(tt.typ, func(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{tt.typ}}) + if err == nil { + t.Fatalf("Generate(%s) succeeded, want an error mentioning %q", tt.typ, tt.want) + } + if !strings.Contains(err.Error(), tt.want) { + t.Errorf("Generate(%s) error = %q, want it to mention %q", tt.typ, err, tt.want) + } + if !strings.Contains(err.Error(), tt.typ) { + t.Errorf("Generate(%s) error = %q, want it to name the type", tt.typ, err) + } + }) + } +} + +// TestGenerateNamesTheField checks that a rejection points at the field, not +// just the struct. +func TestGenerateNamesTheField(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Nested"}}) + if err == nil { + t.Fatal("Generate succeeded, want an error") + } + if !strings.Contains(err.Error(), "field Field") { + t.Errorf("error = %q, want it to name the offending field", err) + } +} + +// TestGenerateRejectsCollidingTypes covers two types in one file whose +// generated names would clash. +func TestGenerateRejectsCollidingTypes(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Collide", "collide"}}) + if err == nil { + t.Fatal("Generate succeeded, want a collision error") + } + if !strings.Contains(err.Error(), "collides with") { + t.Errorf("error = %q, want it to report the collision", err) + } +} + +func TestGenerateRejectsRepeatedType(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Collide", "Collide"}}) + if err == nil { + t.Fatal("Generate succeeded, want a collision error") + } +} + +// TestGenerateBeforePackageCompiles covers the chicken-and-egg case: the call +// sites are written first, so the package does not type-check until the +// generator has run. +func TestGenerateBeforePackageCompiles(t *testing.T) { + got, err := arrgen.Generate(arrgen.Config{Dir: "testdata/pregen", Types: []string{"Sample"}}) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if !bytes.Contains(got, []byte("func SampleRecordBatch(")) { + t.Error("generated output does not declare SampleRecordBatch") + } +} + +// TestGenerateReportsLoadErrors checks that a type which cannot be resolved +// reports the package's own compile errors as the likely cause. +func TestGenerateReportsLoadErrors(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/pregen", Types: []string{"Absent"}}) + if err == nil { + t.Fatal("Generate succeeded, want an error") + } + if !strings.Contains(err.Error(), "does not type-check") { + t.Errorf("error = %q, want it to surface the package's compile errors", err) + } +} + +func TestGenerateNoTypes(t *testing.T) { + if _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/basic"}); err == nil { + t.Error("Generate with no types succeeded, want an error") + } +} + +func TestGenerateMissingPackage(t *testing.T) { + dir := filepath.Join(t.TempDir(), "empty") + if err := os.Mkdir(dir, 0o755); err != nil { + t.Fatal(err) + } + if _, err := arrgen.Generate(arrgen.Config{Dir: dir, Types: []string{"Metric"}}); err == nil { + t.Error("Generate in a directory with no Go files succeeded, want an error") + } +} + +func first(b []byte, n int) []byte { + if len(b) < n { + return b + } + return b[:n] +} diff --git a/arrgen/arrgen/go.mod b/arrgen/arrgen/go.mod new file mode 100644 index 000000000..ad12f8825 --- /dev/null +++ b/arrgen/arrgen/go.mod @@ -0,0 +1,35 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +module github.com/apache/arrow-go/arrgen + +go 1.25.0 + +require ( + github.com/apache/arrow-go/v18 v18.7.0 + golang.org/x/tools v0.45.0 +) + +require ( + github.com/goccy/go-json v0.10.6 // indirect + github.com/google/flatbuffers v25.12.19+incompatible // indirect + github.com/klauspost/cpuid/v2 v2.4.0 // indirect + github.com/zeebo/xxh3 v1.1.0 // indirect + golang.org/x/exp v0.0.0-20260112195511-716be5621a96 // indirect + golang.org/x/mod v0.36.0 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/sys v0.47.0 // indirect +) diff --git a/arrgen/arrgen/go.sum b/arrgen/arrgen/go.sum new file mode 100644 index 000000000..963f34aaf --- /dev/null +++ b/arrgen/arrgen/go.sum @@ -0,0 +1,44 @@ +github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM= +github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/apache/arrow-go/v18 v18.7.0 h1:Vw/i+cJyebUofT7JlqFpe65LrmwxULn166jjwStM4HY= +github.com/apache/arrow-go/v18 v18.7.0/go.mod h1:PM6IigLJkdMwIpeHXnymo+xZ52f42a9EYiLtRel4p/A= +github.com/apache/thrift v0.24.0 h1:zy31L1a49QTNB2bG1BBfMXol3yJrTH975G3pPubQVLQ= +github.com/apache/thrift v0.24.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU= +github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= +github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/klauspost/compress v1.19.0 h1:sXLILfc9jV2QYWkzFOPWStmcUVH2RHEB1JCdY2oVvCQ= +github.com/klauspost/compress v1.19.0/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw= +github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU= +github.com/pierrec/lz4/v4 v4.1.27 h1:+PhzhWDrjRj89TH2sw43nE3+4+W8lSxIuQadEHZyjUk= +github.com/pierrec/lz4/v4 v4.1.27/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= +github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96 h1:Z/6YuSHTLOHfNFdb8zVZomZr7cqNgTJvA8+Qz75D8gU= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96/go.mod h1:nzimsREAkjBCIEFtHiYkrJyT+2uy9YZJB7H1k68CXZU= +golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= +golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= +golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/arrgen/arrgen/internal/gentypes/alloc_test.go b/arrgen/arrgen/internal/gentypes/alloc_test.go new file mode 100644 index 000000000..010172a09 --- /dev/null +++ b/arrgen/arrgen/internal/gentypes/alloc_test.go @@ -0,0 +1,158 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gentypes_test + +import ( + "testing" + + "github.com/apache/arrow-go/arrgen/internal/gentypes" + "github.com/apache/arrow-go/v18/arrow/array/arreflect" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// TestStreamingAppendIsAllocationFree asserts that with room reserved up front, +// appending a row of fixed-width columns allocates nothing. +func TestStreamingAppendIsAllocationFree(t *testing.T) { + skipIfInstrumented(t) + const runs = 2000 + + a := gentypes.NewFixedAppender(memory.DefaultAllocator) + defer a.Release() + a.Reserve(runs + 8) // AllocsPerRun calls the function runs+1 times + + rows := makeFixed(4) + i := 0 + // A single reused variable, as a caller draining a stream would do it. + var row gentypes.Fixed + got := testing.AllocsPerRun(runs, func() { + row = rows[i%len(rows)] + i++ + a.Append(&row) + }) + if got != 0 { + t.Errorf("Append allocated %.2f times per row, want 0", got) + } + a.NewRecordBatch().Release() +} + +// TestSchemaIsAllocationFree checks that the generated schema is a package-level +// variable, where arreflect rebuilds one on every call. +func TestSchemaIsAllocationFree(t *testing.T) { + skipIfInstrumented(t) + if got := testing.AllocsPerRun(1000, func() { _ = gentypes.RowSchema() }); got != 0 { + t.Errorf("RowSchema allocated %.2f times per call, want 0", got) + } + + reflected := testing.AllocsPerRun(100, func() { + if _, err := arreflect.InferSchema[gentypes.Row](); err != nil { + t.Fatal(err) + } + }) + if reflected == 0 { + t.Skip("arreflect.InferSchema no longer allocates; the comparison is moot") + } + t.Logf("schema per call: arrgen 0 allocs, arreflect %.0f allocs", reflected) +} + +// TestBatchEncodingAllocatesLessThanReflection compares whole-batch encoding. +// The threshold is loose: both paths pay for the same Arrow buffers, and the +// point is to catch a regression that erases the advantage, not to pin a number. +func TestBatchEncodingAllocatesLessThanReflection(t *testing.T) { + skipIfInstrumented(t) + const rows = 1024 + fixtures := makeFixed(rows) + mem := memory.DefaultAllocator + + reflected := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := arreflect.RecordFromSlice(fixtures, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + generated := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := gentypes.FixedRecordBatch(mem, fixtures) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + + t.Logf("%d-row batch: arrgen %d allocs / %d ns, arreflect %d allocs / %d ns", + rows, generated.AllocsPerOp(), generated.NsPerOp(), reflected.AllocsPerOp(), reflected.NsPerOp()) + + if generated.AllocsPerOp() >= reflected.AllocsPerOp() { + t.Errorf("generated encoder allocated %d times per batch, want fewer than arreflect's %d", + generated.AllocsPerOp(), reflected.AllocsPerOp()) + } + if got, limit := generated.AllocsPerOp(), reflected.AllocsPerOp()*3/4; got > limit { + t.Errorf("generated encoder allocated %d times per batch, want at most %d (three quarters of arreflect's %d)", + got, limit, reflected.AllocsPerOp()) + } +} + +// TestBatchEncodingIsFasterThanReflection guards the speed claim. Wall-clock +// assertions are noisy on shared CI, so the bar sits well below the ~3.5x +// actually measured on the fixed-width fixture. +func TestBatchEncodingIsFasterThanReflection(t *testing.T) { + skipIfInstrumented(t) + if testing.Short() { + t.Skip("timing comparison skipped in short mode") + } + const rows = 1024 + fixtures := makeFixed(rows) + mem := memory.DefaultAllocator + + reflected := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := arreflect.RecordFromSlice(fixtures, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + generated := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := gentypes.FixedRecordBatch(mem, fixtures) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + if generated.NsPerOp() == 0 || reflected.NsPerOp() == 0 { + t.Skip("benchmark produced no timing data") + } + + speedup := float64(reflected.NsPerOp()) / float64(generated.NsPerOp()) + t.Logf("%d-row batch: arrgen %d ns, arreflect %d ns (%.2fx)", rows, generated.NsPerOp(), reflected.NsPerOp(), speedup) + if speedup < 1.5 { + t.Errorf("generated encoder is only %.2fx faster than arreflect, want at least 1.5x", speedup) + } +} + +func skipIfInstrumented(t *testing.T) { + t.Helper() + if instrumented { + t.Skip("allocation counts and timings are not meaningful under -race or -asan") + } +} diff --git a/arrgen/arrgen/internal/gentypes/bench_test.go b/arrgen/arrgen/internal/gentypes/bench_test.go new file mode 100644 index 000000000..acf2634c3 --- /dev/null +++ b/arrgen/arrgen/internal/gentypes/bench_test.go @@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gentypes_test + +import ( + "fmt" + "testing" + "time" + + "github.com/apache/arrow-go/arrgen/example" + "github.com/apache/arrow-go/arrgen/internal/gentypes" + "github.com/apache/arrow-go/v18/arrow/array/arreflect" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// benchRows is the batch size the encoder benchmarks build: large enough that +// per-batch setup does not dominate, small enough to stay in cache. +const benchRows = 1024 + +func makeMetrics(n int) []example.Metric { + metrics := make([]example.Metric, n) + v := 42.5 + for i := range metrics { + metrics[i] = example.Metric{ + Time: base.Add(time.Duration(i) * time.Second), + Host: fmt.Sprintf("host-%d", i%16), + CPU: float64(i) / 100, + } + if i%4 != 0 { + metrics[i].Value = &v + } + } + return metrics +} + +func makeFixed(n int) []gentypes.Fixed { + rows := make([]gentypes.Fixed, n) + v := 1.5 + for i := range rows { + rows[i] = gentypes.Fixed{ + Timestamp: base.Add(time.Duration(i) * time.Second), + ID: int64(i), + Value: float64(i) * 1.5, + OK: i%2 == 0, + } + if i%3 != 0 { + rows[i].Optional = &v + } + } + return rows +} + +// BenchmarkMetricBatch encodes the same slice of rows into a record batch by +// the reflection path and by the generated one. +func BenchmarkMetricBatch(b *testing.B) { + metrics := makeMetrics(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := arreflect.RecordFromSlice(metrics, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) + + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := example.MetricRecordBatch(mem, metrics) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) +} + +// BenchmarkRowBatch runs the same comparison over the wide fixture: 46 columns +// per row instead of four. +func BenchmarkRowBatch(b *testing.B) { + rows := makeRows(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) + + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := gentypes.RowRecordBatch(mem, rows) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) +} + +// BenchmarkFixedBatch isolates the encoder from Arrow's variable-width buffers: +// every column is fixed width, so the allocations are the encoding strategy's +// own rather than a growing data buffer. +func BenchmarkFixedBatch(b *testing.B) { + rows := makeFixed(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) + + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := gentypes.FixedRecordBatch(mem, rows) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) +} + +// BenchmarkStreamAppend measures the per-row cost of the streaming entry point. +// The reflection path has no single-row API, so its column is the amortized +// cost of encoding a batch of the same rows. +func BenchmarkStreamAppend(b *testing.B) { + rows := makeFixed(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect/batched", func(b *testing.B) { + b.ReportAllocs() + n := 0 + for b.Loop() { + rec, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + n += benchRows + } + b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(n), "ns/row") + }) + + b.Run("arrgen/streamed", func(b *testing.B) { + a := gentypes.NewFixedAppender(mem) + defer a.Release() + b.ReportAllocs() + var row gentypes.Fixed + i := 0 + for b.Loop() { + row = rows[i%benchRows] + i++ + if i%benchRows == 0 { + a.NewRecordBatch().Release() + a.Reserve(benchRows) + } + a.Append(&row) + } + a.NewRecordBatch().Release() + b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(b.N), "ns/row") + }) +} + +// BenchmarkSchema compares arreflect's struct walk against the generated +// package-level schema variable. +func BenchmarkSchema(b *testing.B) { + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + if _, err := arreflect.InferSchema[gentypes.Row](); err != nil { + b.Fatal(err) + } + } + }) + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + _ = gentypes.RowSchema() + } + }) +} diff --git a/arrgen/arrgen/internal/gentypes/equivalence_test.go b/arrgen/arrgen/internal/gentypes/equivalence_test.go new file mode 100644 index 000000000..e99edb56f --- /dev/null +++ b/arrgen/arrgen/internal/gentypes/equivalence_test.go @@ -0,0 +1,449 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gentypes_test + +import ( + "fmt" + "math" + "testing" + "time" + + "github.com/apache/arrow-go/arrgen/example" + "github.com/apache/arrow-go/arrgen/internal/gentypes" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/array/arreflect" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/decimal256" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// arreflectStructFieldGaps names the Row columns arreflect cannot encode as a +// struct field, and the Go type of each. +// +// Its inferArrowType dispatches on reflect.Kind before checking the types it +// special-cases, so a Go struct that Arrow models as a scalar reaches +// inferStructType, which finds only unexported fields and returns an empty +// struct<>. A tag naming an explicit Arrow type rescues the field, so only the +// untagged spellings are listed. These columns are compared by value in +// TestGapColumnValues; see also README.md. +var arreflectStructFieldGaps = map[string]string{ + "ts": "time.Time", + "pts": "*time.Time", + "dec256": "decimal256.Num", + "pdec128": "*decimal128.Num", +} + +// compareColumns asserts that got and want hold the same columns, skipping the +// ones arreflect cannot encode as struct fields. +// +// dictsByValue relaxes dictionary columns to a comparison of the values they +// decode to: an appender reused across batches carries its memo table forward, +// so its later batches encode the same values against a longer dictionary. +func compareColumns(t *testing.T, label string, got, want arrow.RecordBatch, dictsByValue bool) { + t.Helper() + if got.NumRows() != want.NumRows() { + t.Fatalf("%s: NumRows = %d, want %d", label, got.NumRows(), want.NumRows()) + } + if got.NumCols() != want.NumCols() { + t.Fatalf("%s: NumCols = %d, want %d", label, got.NumCols(), want.NumCols()) + } + for i := 0; i < int(got.NumCols()); i++ { + name := got.Schema().Field(i).Name + if _, gap := arreflectStructFieldGaps[name]; gap { + continue + } + gc, wc := got.Column(i), want.Column(i) + if dictsByValue && gc.DataType().ID() == arrow.DICTIONARY { + compareByValueStr(t, label, name, gc, wc) + continue + } + if !array.Equal(gc, wc) { + t.Errorf("%s: column %q differs\n got: %s\nwant: %s", label, name, gc, wc) + } + } +} + +// compareByValueStr compares two arrays by the value each row decodes to, +// ignoring how that value is physically encoded. +func compareByValueStr(t *testing.T, label, name string, got, want arrow.Array) { + t.Helper() + if got.Len() != want.Len() { + t.Errorf("%s: column %q: len = %d, want %d", label, name, got.Len(), want.Len()) + return + } + for i := 0; i < got.Len(); i++ { + if got.IsNull(i) != want.IsNull(i) { + t.Errorf("%s: column %q row %d: null = %t, want %t", label, name, i, got.IsNull(i), want.IsNull(i)) + return + } + if got.IsNull(i) { + continue + } + if got.ValueStr(i) != want.ValueStr(i) { + t.Errorf("%s: column %q row %d = %s, want %s", label, name, i, got.ValueStr(i), want.ValueStr(i)) + return + } + } +} + +// base is a fixed instant with a non-zero time of day and sub-second part, so +// the date32/date64/time32/time64 columns each exercise a different truncation. +var base = time.Date(2024, time.March, 17, 13, 45, 12, 123456789, time.UTC) + +// makeRows builds n deterministic rows. Every third row nulls its pointer +// fields and empties its variable-width ones, so nulls are covered too. +func makeRows(n int) []gentypes.Row { + rows := make([]gentypes.Row, n) + for i := range rows { + null := i%3 == 0 + r := gentypes.Row{ + Bool: i%2 == 0, + Int8: int8(i), + Int16: int16(i * 3), + Int32: int32(i * 7), + Int64: int64(i) * 1e6, + Int: i, + Uint8: uint8(i), + Uint16: uint16(i * 5), + Uint32: uint32(i * 11), + Uint64: uint64(i) * 1e9, + Uint: uint(i), + Float32: float32(i) + 0.25, + Float64: float64(i) * math.Pi, + Str: fmt.Sprintf("host-%d", i%7), + Bin: []byte{byte(i), byte(i >> 8)}, + Timestamp: base.Add(time.Duration(i) * time.Second), + Date32: base.AddDate(0, 0, i), + Date64: base.AddDate(0, 0, i), + Time32: base.Add(time.Duration(i) * time.Minute), + Time64: base.Add(time.Duration(i) * time.Minute), + Duration: time.Duration(i) * time.Millisecond, + Dec32: decimal.Decimal32(i * 100), + Dec64: decimal.Decimal64(i * 10000), + Dec128: decimal128.FromU64(uint64(i) * 1000), + Dec256: decimal256.FromU64(uint64(i) * 2000), + LargeStr: fmt.Sprintf("large-%d", i), + ViewStr: fmt.Sprintf("view-%d", i), + LargeBin: []byte(fmt.Sprintf("lb-%d", i)), + ViewBin: []byte(fmt.Sprintf("vb-%d", i)), + DictStr: fmt.Sprintf("region-%d", i%3), + DictBin: []byte(fmt.Sprintf("k%d", i%4)), + DictInt: int32(i % 5), + DictF64: float64(i % 6), + Untagged: int64(i), + Secret: "never encoded", + } + if null { + // Leave every pointer nil, and null the []byte columns, which map to + // null without a pointer. + r.Bin = nil + r.LargeBin = nil + r.ViewBin = nil + r.DictBin = nil + rows[i] = r + continue + } + b := i%2 == 1 + i64 := int64(i) * 3 + f64 := float64(i) / 3 + s := fmt.Sprintf("p-%d", i) + bin := []byte(fmt.Sprintf("pb-%d", i)) + ts := base.Add(time.Duration(i) * time.Hour) + dur := time.Duration(i) * time.Microsecond + dec := decimal.Decimal32(i) + ds := fmt.Sprintf("pd-%d", i%2) + r.PBool, r.PInt64, r.PF64, r.PStr, r.PBin = &b, &i64, &f64, &s, &bin + dec128 := decimal128.FromU64(uint64(i)) + r.PTS, r.PDate32, r.PTime64, r.PDur, r.PDec32, r.PDictS = &ts, &ts, &ts, &dur, &dec, &ds + r.PDec128 = &dec128 + rows[i] = r + } + return rows +} + +// TestRowSchemaMatchesArreflect checks that the generated schema is what +// arreflect infers, column for column, including names, types and nullability. +func TestRowSchemaMatchesArreflect(t *testing.T) { + want, err := arreflect.InferSchema[gentypes.Row]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + got := gentypes.RowSchema() + if got.NumFields() != want.NumFields() { + t.Fatalf("field count = %d, want %d\n got: %s\nwant: %s", got.NumFields(), want.NumFields(), got, want) + } + for i := 0; i < got.NumFields(); i++ { + gf, wf := got.Field(i), want.Field(i) + if gf.Name != wf.Name { + t.Errorf("field %d: name = %q, want %q", i, gf.Name, wf.Name) + continue + } + if _, gap := arreflectStructFieldGaps[gf.Name]; gap { + continue + } + if !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { + t.Errorf("field %q: got %s nullable=%t, want %s nullable=%t", gf.Name, gf.Type, gf.Nullable, wf.Type, wf.Nullable) + } + } +} + +func TestMetricSchemaMatchesArreflect(t *testing.T) { + want, err := arreflect.InferSchema[example.Metric]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + got := example.MetricSchema() + // "ts" is a bare time.Time: see arreflectStructFieldGaps. + for i := 0; i < got.NumFields(); i++ { + gf, wf := got.Field(i), want.Field(i) + if gf.Name == "ts" { + continue + } + if gf.Name != wf.Name || !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { + t.Errorf("field %d: got %s, want %s", i, gf, wf) + } + } +} + +// TestRowRecordMatchesArreflect checks that for the same input the generated +// encoder and the reflection encoder produce equal columns. +// +// It compares column by column rather than with array.RecordEqual so the +// columns arreflect cannot encode as struct fields can be skipped by name; see +// arreflectStructFieldGaps. +func TestRowRecordMatchesArreflect(t *testing.T) { + for _, n := range []int{0, 1, 2, 3, 64} { + t.Run(fmt.Sprintf("rows=%d", n), func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + rows := makeRows(n) + + want, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + t.Fatalf("arreflect.RecordFromSlice: %v", err) + } + defer want.Release() + + got, err := gentypes.RowRecordBatch(mem, rows) + if err != nil { + t.Fatalf("RowRecordBatch: %v", err) + } + defer got.Release() + + compareColumns(t, fmt.Sprintf("rows=%d", n), got, want, false) + }) + } +} + +// TestArreflectStructFieldGaps pins the upstream behavior that +// arreflectStructFieldGaps works around: if arreflect learns to encode these +// struct fields, this fails and the skip list should shrink to match. +func TestArreflectStructFieldGaps(t *testing.T) { + schema, err := arreflect.InferSchema[gentypes.Row]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + for name, goType := range arreflectStructFieldGaps { + idx := schema.FieldIndices(name) + if len(idx) != 1 { + t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) + } + if dt := schema.Field(idx[0]).Type; dt.ID() != arrow.STRUCT { + t.Errorf("arreflect now infers %s column %q as %s rather than an empty struct; "+ + "remove it from arreflectStructFieldGaps so the strict comparison covers it", goType, name, dt) + } + } +} + +// TestGapColumnValues covers the columns TestRowRecordMatchesArreflect skips. +// There is no reflection-built batch to compare against, so the expected Arrow +// values are spelled out directly. +func TestGapColumnValues(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + ts := base + dec128 := decimal128.FromU64(12345) + dec256 := decimal256.FromU64(67890) + rows := []gentypes.Row{ + {Timestamp: ts, Dec128: dec128, Dec256: dec256, PTS: &ts, PDec128: &dec128}, + {Timestamp: ts.Add(time.Second), Dec128: dec128, Dec256: dec256}, // pointers nil + } + + rec, err := gentypes.RowRecordBatch(mem, rows) + if err != nil { + t.Fatalf("RowRecordBatch: %v", err) + } + defer rec.Release() + + col := func(name string) arrow.Array { + idx := rec.Schema().FieldIndices(name) + if len(idx) != 1 { + t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) + } + return rec.Column(idx[0]) + } + + tsCol := col("ts").(*array.Timestamp) + if got, want := tsCol.Value(0), arrow.Timestamp(ts.UnixNano()); got != want { + t.Errorf("ts[0] = %d, want %d", got, want) + } + if got, want := tsCol.Value(1), arrow.Timestamp(ts.Add(time.Second).UnixNano()); got != want { + t.Errorf("ts[1] = %d, want %d", got, want) + } + if dt := tsCol.DataType().(*arrow.TimestampType); dt.Unit != arrow.Nanosecond || dt.TimeZone != "UTC" { + t.Errorf("ts type = %s, want timestamp[ns, tz=UTC]", dt) + } + + ptsCol := col("pts").(*array.Timestamp) + if ptsCol.IsNull(0) { + t.Error("pts[0] is null, want a value") + } + if !ptsCol.IsNull(1) { + t.Error("pts[1] is not null, want null") + } + + if got := col("dec128").(*array.Decimal128).Value(0); got != dec128 { + t.Errorf("dec128[0] = %v, want %v", got, dec128) + } + if got := col("dec256").(*array.Decimal256).Value(0); got != dec256 { + t.Errorf("dec256[0] = %v, want %v", got, dec256) + } + if !col("pdec128").(*array.Decimal128).IsNull(1) { + t.Error("pdec128[1] is not null, want null") + } +} + +func TestMetricRecordMatchesArreflect(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + v := 3.5 + metrics := []example.Metric{ + {Time: base, Host: "a", CPU: 0.5, Value: &v, Secret: "hidden"}, + {Time: base.Add(time.Second), Host: "b", CPU: 1.5, Value: nil}, + } + + want, err := arreflect.RecordFromSlice(metrics, mem) + if err != nil { + t.Fatalf("arreflect.RecordFromSlice: %v", err) + } + defer want.Release() + + got, err := example.MetricRecordBatch(mem, metrics) + if err != nil { + t.Fatalf("MetricRecordBatch: %v", err) + } + defer got.Release() + + // "ts" is a bare time.Time; see arreflectStructFieldGaps. + for i := 0; i < int(got.NumCols()); i++ { + name := got.Schema().Field(i).Name + if name == "ts" { + continue + } + if !array.Equal(got.Column(i), want.Column(i)) { + t.Errorf("column %q differs\n got: %s\nwant: %s", name, got.Column(i), want.Column(i)) + } + } + if got, want := got.Column(0).(*array.Timestamp).Value(0), arrow.Timestamp(base.UnixNano()); got != want { + t.Errorf("ts[0] = %d, want %d", got, want) + } + if secrets := got.Schema().FieldIndices("Secret"); len(secrets) != 0 { + t.Errorf(`a field tagged arrow:"-" was encoded as column %v`, secrets) + } +} + +// TestAppendMatchesAppendSlice checks that feeding rows one at a time lands in +// the same place as handing over a slice. +func TestAppendMatchesAppendSlice(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + rows := makeRows(16) + + bulk := gentypes.NewRowAppender(mem) + defer bulk.Release() + bulk.AppendSlice(rows) + bulkRec := bulk.NewRecordBatch() + defer bulkRec.Release() + + streamed := gentypes.NewRowAppender(mem) + defer streamed.Release() + streamed.Reserve(len(rows)) + var row gentypes.Row + for i := range rows { + row = rows[i] // a caller reusing one row variable must be safe + streamed.Append(&row) + } + if got, want := streamed.Len(), len(rows); got != want { + t.Errorf("Len() = %d, want %d", got, want) + } + streamedRec := streamed.NewRecordBatch() + defer streamedRec.Release() + + if !array.RecordEqual(streamedRec, bulkRec) { + t.Errorf("streamed batch differs from bulk batch\n got: %s\nwant: %s", streamedRec, bulkRec) + } + if err := streamed.Err(); err != nil { + t.Errorf("Err() = %v, want nil", err) + } +} + +// TestAppenderReusableAcrossBatches checks that NewRecordBatch hands over the +// rows so far and leaves the appender ready for the next batch. +func TestAppenderReusableAcrossBatches(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + rows := makeRows(9) + a := gentypes.NewRowAppender(mem) + defer a.Release() + + for i := 0; i < len(rows); i += 3 { + a.AppendSlice(rows[i : i+3]) + rec := a.NewRecordBatch() + if got := rec.NumRows(); got != 3 { + t.Errorf("batch %d: NumRows() = %d, want 3", i/3, got) + } + want, err := arreflect.RecordFromSlice(rows[i:i+3], mem) + if err != nil { + t.Fatalf("arreflect.RecordFromSlice: %v", err) + } + compareColumns(t, fmt.Sprintf("batch %d", i/3), rec, want, true) + want.Release() + rec.Release() + if got := a.Len(); got != 0 { + t.Errorf("batch %d: Len() after NewRecordBatch = %d, want 0", i/3, got) + } + } +} + +// TestNilAllocatorUsesDefault covers the constructor's nil-allocator contract. +func TestNilAllocatorUsesDefault(t *testing.T) { + a := gentypes.NewRowAppender(nil) + defer a.Release() + a.AppendSlice(makeRows(2)) + rec := a.NewRecordBatch() + defer rec.Release() + if got := rec.NumRows(); got != 2 { + t.Errorf("NumRows() = %d, want 2", got) + } +} diff --git a/arrgen/arrgen/internal/gentypes/instrumented_race_test.go b/arrgen/arrgen/internal/gentypes/instrumented_race_test.go new file mode 100644 index 000000000..dbdd4fe22 --- /dev/null +++ b/arrgen/arrgen/internal/gentypes/instrumented_race_test.go @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build race || asan + +package gentypes_test + +const instrumented = true diff --git a/arrgen/arrgen/internal/gentypes/instrumented_test.go b/arrgen/arrgen/internal/gentypes/instrumented_test.go new file mode 100644 index 000000000..f8d5afc8d --- /dev/null +++ b/arrgen/arrgen/internal/gentypes/instrumented_test.go @@ -0,0 +1,24 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build !race && !asan + +package gentypes_test + +// instrumented reports whether the binary was built with a sanitizer that +// perturbs allocation counts and timings; the allocation and speed assertions +// skip when it is set. +const instrumented = false diff --git a/arrgen/arrgen/internal/gentypes/row_arrow.go b/arrgen/arrgen/internal/gentypes/row_arrow.go new file mode 100644 index 000000000..8a3a5332a --- /dev/null +++ b/arrgen/arrgen/internal/gentypes/row_arrow.go @@ -0,0 +1,475 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by arrgen. DO NOT EDIT. + +package gentypes + +import ( + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// rowArrowSchema is the Arrow schema of Row, resolved from its arrow +// struct tags. +var rowArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "b", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, + {Name: "i8", Type: arrow.PrimitiveTypes.Int8, Nullable: false}, + {Name: "i16", Type: arrow.PrimitiveTypes.Int16, Nullable: false}, + {Name: "i32", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "i64", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "i", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "u8", Type: arrow.PrimitiveTypes.Uint8, Nullable: false}, + {Name: "u16", Type: arrow.PrimitiveTypes.Uint16, Nullable: false}, + {Name: "u32", Type: arrow.PrimitiveTypes.Uint32, Nullable: false}, + {Name: "u64", Type: arrow.PrimitiveTypes.Uint64, Nullable: false}, + {Name: "u", Type: arrow.PrimitiveTypes.Uint64, Nullable: false}, + {Name: "f32", Type: arrow.PrimitiveTypes.Float32, Nullable: false}, + {Name: "f64", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "s", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: "bin", Type: arrow.BinaryTypes.Binary, Nullable: false}, + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "d32", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, + {Name: "d64", Type: arrow.FixedWidthTypes.Date64, Nullable: false}, + {Name: "t32", Type: &arrow.Time32Type{Unit: arrow.Millisecond}, Nullable: false}, + {Name: "t64", Type: &arrow.Time64Type{Unit: arrow.Nanosecond}, Nullable: false}, + {Name: "dur", Type: &arrow.DurationType{Unit: arrow.Nanosecond}, Nullable: false}, + {Name: "dec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: false}, + {Name: "dec64", Type: &arrow.Decimal64Type{Precision: 18, Scale: 4}, Nullable: false}, + {Name: "dec128", Type: &arrow.Decimal128Type{Precision: 20, Scale: 3}, Nullable: false}, + {Name: "dec256", Type: &arrow.Decimal256Type{Precision: 76, Scale: 0}, Nullable: false}, + {Name: "ls", Type: arrow.BinaryTypes.LargeString, Nullable: false}, + {Name: "vs", Type: arrow.BinaryTypes.StringView, Nullable: false}, + {Name: "lb", Type: arrow.BinaryTypes.LargeBinary, Nullable: false}, + {Name: "vb", Type: arrow.BinaryTypes.BinaryView, Nullable: false}, + {Name: "ds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: false}, + {Name: "db", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.Binary}, Nullable: false}, + {Name: "di", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Int32}, Nullable: false}, + {Name: "df", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Float64}, Nullable: false}, + {Name: "Untagged", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "pb", Type: arrow.FixedWidthTypes.Boolean, Nullable: true}, + {Name: "pi64", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "pf64", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + {Name: "ps", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "pbin", Type: arrow.BinaryTypes.Binary, Nullable: true}, + {Name: "pts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: true}, + {Name: "pd32", Type: arrow.FixedWidthTypes.Date32, Nullable: true}, + {Name: "pt64", Type: &arrow.Time64Type{Unit: arrow.Nanosecond}, Nullable: true}, + {Name: "pdur", Type: &arrow.DurationType{Unit: arrow.Nanosecond}, Nullable: true}, + {Name: "pdec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: true}, + {Name: "pdec128", Type: &arrow.Decimal128Type{Precision: 38, Scale: 0}, Nullable: true}, + {Name: "pds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: true}, +}, nil) + +// RowSchema returns the Arrow schema encoded by RowAppender. +func RowSchema() *arrow.Schema { return rowArrowSchema } + +// RowAppender converts Row values into Arrow record batches +// without reflection. +// +// It is not safe for concurrent use. Release it when done. +type RowAppender struct { + rb *array.RecordBuilder + b0 *array.BooleanBuilder + b1 *array.Int8Builder + b2 *array.Int16Builder + b3 *array.Int32Builder + b4 *array.Int64Builder + b5 *array.Int64Builder + b6 *array.Uint8Builder + b7 *array.Uint16Builder + b8 *array.Uint32Builder + b9 *array.Uint64Builder + b10 *array.Uint64Builder + b11 *array.Float32Builder + b12 *array.Float64Builder + b13 *array.StringBuilder + b14 *array.BinaryBuilder + b15 *array.TimestampBuilder + b16 *array.Date32Builder + b17 *array.Date64Builder + b18 *array.Time32Builder + b19 *array.Time64Builder + b20 *array.DurationBuilder + b21 *array.Decimal32Builder + b22 *array.Decimal64Builder + b23 *array.Decimal128Builder + b24 *array.Decimal256Builder + b25 *array.LargeStringBuilder + b26 *array.StringViewBuilder + b27 *array.BinaryBuilder + b28 *array.BinaryViewBuilder + b29 *array.BinaryDictionaryBuilder + b30 *array.BinaryDictionaryBuilder + b31 *array.Int32DictionaryBuilder + b32 *array.Float64DictionaryBuilder + b33 *array.Int64Builder + b34 *array.BooleanBuilder + b35 *array.Int64Builder + b36 *array.Float64Builder + b37 *array.StringBuilder + b38 *array.BinaryBuilder + b39 *array.TimestampBuilder + b40 *array.Date32Builder + b41 *array.Time64Builder + b42 *array.DurationBuilder + b43 *array.Decimal32Builder + b44 *array.Decimal128Builder + b45 *array.BinaryDictionaryBuilder + err error +} + +// NewRowAppender returns an appender that builds Row batches with mem, +// or memory.DefaultAllocator when mem is nil. The typed field builders are +// resolved once here, so Append does no type assertions. +func NewRowAppender(mem memory.Allocator) *RowAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, rowArrowSchema) + return &RowAppender{ + rb: rb, + b0: rb.Field(0).(*array.BooleanBuilder), + b1: rb.Field(1).(*array.Int8Builder), + b2: rb.Field(2).(*array.Int16Builder), + b3: rb.Field(3).(*array.Int32Builder), + b4: rb.Field(4).(*array.Int64Builder), + b5: rb.Field(5).(*array.Int64Builder), + b6: rb.Field(6).(*array.Uint8Builder), + b7: rb.Field(7).(*array.Uint16Builder), + b8: rb.Field(8).(*array.Uint32Builder), + b9: rb.Field(9).(*array.Uint64Builder), + b10: rb.Field(10).(*array.Uint64Builder), + b11: rb.Field(11).(*array.Float32Builder), + b12: rb.Field(12).(*array.Float64Builder), + b13: rb.Field(13).(*array.StringBuilder), + b14: rb.Field(14).(*array.BinaryBuilder), + b15: rb.Field(15).(*array.TimestampBuilder), + b16: rb.Field(16).(*array.Date32Builder), + b17: rb.Field(17).(*array.Date64Builder), + b18: rb.Field(18).(*array.Time32Builder), + b19: rb.Field(19).(*array.Time64Builder), + b20: rb.Field(20).(*array.DurationBuilder), + b21: rb.Field(21).(*array.Decimal32Builder), + b22: rb.Field(22).(*array.Decimal64Builder), + b23: rb.Field(23).(*array.Decimal128Builder), + b24: rb.Field(24).(*array.Decimal256Builder), + b25: rb.Field(25).(*array.LargeStringBuilder), + b26: rb.Field(26).(*array.StringViewBuilder), + b27: rb.Field(27).(*array.BinaryBuilder), + b28: rb.Field(28).(*array.BinaryViewBuilder), + b29: rb.Field(29).(*array.BinaryDictionaryBuilder), + b30: rb.Field(30).(*array.BinaryDictionaryBuilder), + b31: rb.Field(31).(*array.Int32DictionaryBuilder), + b32: rb.Field(32).(*array.Float64DictionaryBuilder), + b33: rb.Field(33).(*array.Int64Builder), + b34: rb.Field(34).(*array.BooleanBuilder), + b35: rb.Field(35).(*array.Int64Builder), + b36: rb.Field(36).(*array.Float64Builder), + b37: rb.Field(37).(*array.StringBuilder), + b38: rb.Field(38).(*array.BinaryBuilder), + b39: rb.Field(39).(*array.TimestampBuilder), + b40: rb.Field(40).(*array.Date32Builder), + b41: rb.Field(41).(*array.Time64Builder), + b42: rb.Field(42).(*array.DurationBuilder), + b43: rb.Field(43).(*array.Decimal32Builder), + b44: rb.Field(44).(*array.Decimal128Builder), + b45: rb.Field(45).(*array.BinaryDictionaryBuilder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *RowAppender) Schema() *arrow.Schema { return rowArrowSchema } + +// RecordBuilder returns the underlying builder. Appending through it directly +// is allowed as long as every column stays the same length. +func (a *RowAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends keeps the append path free of allocations. +func (a *RowAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *RowAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller can reuse a single Row variable across calls. +func (a *RowAppender) Append(v *Row) { + a.b0.Append(v.Bool) + a.b1.Append(v.Int8) + a.b2.Append(v.Int16) + a.b3.Append(v.Int32) + a.b4.Append(v.Int64) + a.b5.Append(int64(v.Int)) + a.b6.Append(v.Uint8) + a.b7.Append(v.Uint16) + a.b8.Append(v.Uint32) + a.b9.Append(v.Uint64) + a.b10.Append(uint64(v.Uint)) + a.b11.Append(v.Float32) + a.b12.Append(v.Float64) + a.b13.Append(v.Str) + if v.Bin == nil { + a.b14.AppendNull() + } else { + a.b14.Append(v.Bin) + } + a.b15.Append(arrow.Timestamp(v.Timestamp.UnixNano())) + a.b16.Append(arrow.Date32FromTime(v.Date32)) + a.b17.Append(arrow.Date64FromTime(v.Date64)) + { + tod := v.Time32.UTC() + a.b18.Append(arrow.Time32(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds() / 1e6)) + } + { + tod := v.Time64.UTC() + a.b19.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) + } + a.b20.Append(arrow.Duration(v.Duration.Nanoseconds())) + a.b21.Append(v.Dec32) + a.b22.Append(v.Dec64) + a.b23.Append(v.Dec128) + a.b24.Append(v.Dec256) + a.b25.Append(v.LargeStr) + a.b26.Append(v.ViewStr) + if v.LargeBin == nil { + a.b27.AppendNull() + } else { + a.b27.Append(v.LargeBin) + } + if v.ViewBin == nil { + a.b28.AppendNull() + } else { + a.b28.Append(v.ViewBin) + } + a.setErr(a.b29.AppendString(v.DictStr)) + if v.DictBin == nil { + a.b30.AppendNull() + } else { + a.setErr(a.b30.Append(v.DictBin)) + } + a.setErr(a.b31.Append(v.DictInt)) + a.setErr(a.b32.Append(v.DictF64)) + a.b33.Append(v.Untagged) + if v.PBool == nil { + a.b34.AppendNull() + } else { + a.b34.Append(*v.PBool) + } + if v.PInt64 == nil { + a.b35.AppendNull() + } else { + a.b35.Append(*v.PInt64) + } + if v.PF64 == nil { + a.b36.AppendNull() + } else { + a.b36.Append(*v.PF64) + } + if v.PStr == nil { + a.b37.AppendNull() + } else { + a.b37.Append(*v.PStr) + } + if v.PBin == nil || *v.PBin == nil { + a.b38.AppendNull() + } else { + a.b38.Append(*v.PBin) + } + if v.PTS == nil { + a.b39.AppendNull() + } else { + a.b39.Append(arrow.Timestamp((*v.PTS).UnixNano())) + } + if v.PDate32 == nil { + a.b40.AppendNull() + } else { + a.b40.Append(arrow.Date32FromTime(*v.PDate32)) + } + if v.PTime64 == nil { + a.b41.AppendNull() + } else { + tod := (*v.PTime64).UTC() + a.b41.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) + } + if v.PDur == nil { + a.b42.AppendNull() + } else { + a.b42.Append(arrow.Duration((*v.PDur).Nanoseconds())) + } + if v.PDec32 == nil { + a.b43.AppendNull() + } else { + a.b43.Append(*v.PDec32) + } + if v.PDec128 == nil { + a.b44.AppendNull() + } else { + a.b44.Append(*v.PDec128) + } + if v.PDictS == nil { + a.b45.AppendNull() + } else { + a.setErr(a.b45.AppendString(*v.PDictS)) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *RowAppender) AppendSlice(vs []Row) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch returns the appended rows as a record batch and resets the +// appender. The caller must Release the batch. +func (a *RowAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail. The error is sticky: NewRecordBatch does not clear it. +func (a *RowAppender) Err() error { return a.err } + +// setErr keeps the first error. +func (a *RowAppender) setErr(err error) { + if err != nil && a.err == nil { + a.err = err + } +} + +// Release releases the appender's builders. +func (a *RowAppender) Release() { a.rb.Release() } + +// RowRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Row](vs, mem). The +// caller must Release the batch. +func RowRecordBatch(mem memory.Allocator, vs []Row) (arrow.RecordBatch, error) { + a := NewRowAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} + +// fixedArrowSchema is the Arrow schema of Fixed, resolved from its arrow +// struct tags. +var fixedArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "val", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "ok", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, + {Name: "opt", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, +}, nil) + +// FixedSchema returns the Arrow schema encoded by FixedAppender. +func FixedSchema() *arrow.Schema { return fixedArrowSchema } + +// FixedAppender converts Fixed values into Arrow record batches +// without reflection. +// +// It is not safe for concurrent use. Release it when done. +type FixedAppender struct { + rb *array.RecordBuilder + b0 *array.TimestampBuilder + b1 *array.Int64Builder + b2 *array.Float64Builder + b3 *array.BooleanBuilder + b4 *array.Float64Builder + err error +} + +// NewFixedAppender returns an appender that builds Fixed batches with mem, +// or memory.DefaultAllocator when mem is nil. The typed field builders are +// resolved once here, so Append does no type assertions. +func NewFixedAppender(mem memory.Allocator) *FixedAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, fixedArrowSchema) + return &FixedAppender{ + rb: rb, + b0: rb.Field(0).(*array.TimestampBuilder), + b1: rb.Field(1).(*array.Int64Builder), + b2: rb.Field(2).(*array.Float64Builder), + b3: rb.Field(3).(*array.BooleanBuilder), + b4: rb.Field(4).(*array.Float64Builder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *FixedAppender) Schema() *arrow.Schema { return fixedArrowSchema } + +// RecordBuilder returns the underlying builder. Appending through it directly +// is allowed as long as every column stays the same length. +func (a *FixedAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends keeps the append path free of allocations. +func (a *FixedAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *FixedAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller can reuse a single Fixed variable across calls. +func (a *FixedAppender) Append(v *Fixed) { + a.b0.Append(arrow.Timestamp(v.Timestamp.UnixNano())) + a.b1.Append(v.ID) + a.b2.Append(v.Value) + a.b3.Append(v.OK) + if v.Optional == nil { + a.b4.AppendNull() + } else { + a.b4.Append(*v.Optional) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *FixedAppender) AppendSlice(vs []Fixed) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch returns the appended rows as a record batch and resets the +// appender. The caller must Release the batch. +func (a *FixedAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail. The error is sticky: NewRecordBatch does not clear it. +func (a *FixedAppender) Err() error { return a.err } + +// Release releases the appender's builders. +func (a *FixedAppender) Release() { a.rb.Release() } + +// FixedRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Fixed](vs, mem). The +// caller must Release the batch. +func FixedRecordBatch(mem memory.Allocator, vs []Fixed) (arrow.RecordBatch, error) { + a := NewFixedAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} diff --git a/arrgen/arrgen/internal/gentypes/types.go b/arrgen/arrgen/internal/gentypes/types.go new file mode 100644 index 000000000..318dade5e --- /dev/null +++ b/arrgen/arrgen/internal/gentypes/types.go @@ -0,0 +1,98 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package gentypes holds the fixtures the arrgen test suite generates from. +// Row carries one field for every column shape arrgen supports, so the +// equivalence tests cover the whole surface. +package gentypes + +import ( + "time" + + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/decimal256" +) + +// Row covers every supported field type, tag option and nullable variant. +type Row struct { + Bool bool `arrow:"b"` + Int8 int8 `arrow:"i8"` + Int16 int16 `arrow:"i16"` + Int32 int32 `arrow:"i32"` + Int64 int64 `arrow:"i64"` + Int int `arrow:"i"` + Uint8 uint8 `arrow:"u8"` + Uint16 uint16 `arrow:"u16"` + Uint32 uint32 `arrow:"u32"` + Uint64 uint64 `arrow:"u64"` + Uint uint `arrow:"u"` + Float32 float32 `arrow:"f32"` + Float64 float64 `arrow:"f64"` + Str string `arrow:"s"` + Bin []byte `arrow:"bin"` + + Timestamp time.Time `arrow:"ts"` + Date32 time.Time `arrow:"d32,date32"` + Date64 time.Time `arrow:"d64,date64"` + Time32 time.Time `arrow:"t32,time32"` + Time64 time.Time `arrow:"t64,time64"` + Duration time.Duration `arrow:"dur"` + + Dec32 decimal.Decimal32 `arrow:"dec32"` + Dec64 decimal.Decimal64 `arrow:"dec64,decimal(18,4)"` + Dec128 decimal128.Num `arrow:"dec128,decimal(20,3)"` + Dec256 decimal256.Num `arrow:"dec256"` + + LargeStr string `arrow:"ls,large"` + ViewStr string `arrow:"vs,view"` + LargeBin []byte `arrow:"lb,large"` + ViewBin []byte `arrow:"vb,view"` + + DictStr string `arrow:"ds,dict"` + DictBin []byte `arrow:"db,dict"` + DictInt int32 `arrow:"di,dict"` + DictF64 float64 `arrow:"df,dict"` + Untagged int64 // no tag: the column takes the Go field name + + PBool *bool `arrow:"pb"` + PInt64 *int64 `arrow:"pi64"` + PF64 *float64 `arrow:"pf64"` + PStr *string `arrow:"ps"` + PBin *[]byte `arrow:"pbin"` + PTS *time.Time `arrow:"pts"` + PDate32 *time.Time `arrow:"pd32,date32"` + PTime64 *time.Time `arrow:"pt64,time64"` + PDur *time.Duration `arrow:"pdur"` + PDec32 *decimal.Decimal32 `arrow:"pdec32"` + PDec128 *decimal128.Num `arrow:"pdec128"` + PDictS *string `arrow:"pds,dict"` + + Secret string `arrow:"-"` // excluded from Arrow entirely +} + +// Fixed holds only fixed-width columns, so appending a row never has to grow a +// variable-length data buffer. It is what the zero-allocation assertion in +// alloc_test.go measures. +type Fixed struct { + Timestamp time.Time `arrow:"ts"` + ID int64 `arrow:"id"` + Value float64 `arrow:"val"` + OK bool `arrow:"ok"` + Optional *float64 `arrow:"opt"` +} + +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Row,Fixed -header ../../license_header.txt -output row_arrow.go diff --git a/arrgen/arrgen/license_header.txt b/arrgen/arrgen/license_header.txt new file mode 100644 index 000000000..f1fc31831 --- /dev/null +++ b/arrgen/arrgen/license_header.txt @@ -0,0 +1,15 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. diff --git a/arrgen/arrgen/mapping.go b/arrgen/arrgen/mapping.go new file mode 100644 index 000000000..f00e6583b --- /dev/null +++ b/arrgen/arrgen/mapping.go @@ -0,0 +1,320 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import ( + "fmt" + "go/types" +) + +// colSpec is everything the template needs to emit one Arrow column: the type +// expression for the schema, the concrete builder the appender caches, and a +// renderer for the append statement. The three are resolved together so they +// cannot drift apart. +type colSpec struct { + arrowType string // Go expression for the arrow.DataType + builderType string // Go type of the cached typed builder + fallible bool // the builder's Append returns an error (dictionaries) + needsTime bool // the append statement references the time package + nilable bool // a nil Go value maps to null even in a non-nullable column + + // appendStmt renders the append of a known-non-nil value. bld is the cached + // builder expression ("a.b3"), val is the value ("v.F" or "*v.F"), and recv + // is val parenthesized so a method call binds to the value, not the pointer. + appendStmt func(bld, recv, val string) string +} + +// optSupport records which tag options a column kind can honor. An option that +// would be ignored is rejected instead, unlike arreflect. +type optSupport struct { + temporal bool + decimal bool + largeView bool + dict bool +} + +func checkOpts(o tagOpts, goType string, s optSupport) error { + switch { + case o.Temporal != "" && !s.temporal: + return fmt.Errorf("%q is only valid on a time.Time field, not %s", o.Temporal, goType) + case o.HasDecimalOpts && !s.decimal: + return fmt.Errorf("decimal(precision,scale) is only valid on a decimal field, not %s", goType) + case o.Large && !s.largeView: + return fmt.Errorf("large has no effect on %s; it is only valid on string and []byte fields", goType) + case o.View && !s.largeView: + return fmt.Errorf("view has no effect on %s; it is only valid on string and []byte fields", goType) + case o.Dict && !s.dict: + return fmt.Errorf("dict is not supported on %s; it is valid on string, []byte, integer and float fields", goType) + case o.Dict && o.Large: + return fmt.Errorf("dict cannot be combined with large: Dictionary is not implemented by arrow-go") + } + return nil +} + +// resolveColumn maps a Go field type and its parsed tag to an Arrow column, +// reporting nullable=true for a pointer field, as arreflect does. +func resolveColumn(t types.Type, opts tagOpts) (spec colSpec, nullable bool, err error) { + t = types.Unalias(t) + if p, ok := t.(*types.Pointer); ok { + nullable = true + t = types.Unalias(p.Elem()) + if _, ok := t.(*types.Pointer); ok { + return colSpec{}, false, fmt.Errorf("multi-level pointer type %s is not supported", t) + } + } + spec, err = baseSpec(t, opts) + return spec, nullable, err +} + +func baseSpec(t types.Type, opts tagOpts) (colSpec, error) { + if named, ok := t.(*types.Named); ok { + return namedSpec(named, opts) + } + if sl, ok := t.(*types.Slice); ok { + if b, ok := types.Unalias(sl.Elem()).(*types.Basic); ok && b.Kind() == types.Uint8 { + return byteSliceSpec(opts) + } + return colSpec{}, fmt.Errorf("list column for %s is not supported; encode it with arreflect", t) + } + if b, ok := t.(*types.Basic); ok { + return basicSpec(b, opts) + } + return colSpec{}, fmt.Errorf("type %s is not supported", t) +} + +// namedSpec handles the named types arreflect recognizes by identity. Every +// other named type, including a defined scalar such as `type ID int64`, is +// unsupported there too. +func namedSpec(named *types.Named, opts tagOpts) (colSpec, error) { + obj := named.Obj() + pkg := "" + if obj.Pkg() != nil { + pkg = obj.Pkg().Path() + } + switch { + case pkg == "time" && obj.Name() == "Time": + return timeSpec(opts) + case pkg == "time" && obj.Name() == "Duration": + if err := checkOpts(opts, "time.Duration", optSupport{}); err != nil { + return colSpec{}, err + } + return colSpec{ + arrowType: "&arrow.DurationType{Unit: arrow.Nanosecond}", + builderType: "*array.DurationBuilder", + appendStmt: func(bld, recv, _ string) string { + return fmt.Sprintf("%s.Append(arrow.Duration(%s.Nanoseconds()))", bld, recv) + }, + }, nil + case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal32": + return decimalSpec(opts, "Decimal32", 9) + case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal64": + return decimalSpec(opts, "Decimal64", 18) + case isDecimalPkg(pkg, "decimal128") && obj.Name() == "Num": + return decimalSpec(opts, "Decimal128", 38) + case isDecimalPkg(pkg, "decimal256") && obj.Name() == "Num": + return decimalSpec(opts, "Decimal256", 76) + } + if _, ok := named.Underlying().(*types.Struct); ok { + return colSpec{}, fmt.Errorf("nested struct column for %s is not supported; encode it with arreflect", named) + } + return colSpec{}, fmt.Errorf("named type %s is not supported; arreflect matches scalar fields by exact type, so a defined type such as this one has no Arrow mapping in either path", named) +} + +func isDecimalPkg(path, name string) bool { + const prefix = "github.com/apache/arrow-go/v18/arrow/" + return path == prefix+name +} + +func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32) (colSpec, error) { + goType := "decimal." + arrowName + if err := checkOpts(opts, goType, optSupport{decimal: true}); err != nil { + return colSpec{}, err + } + precision, scale := defaultPrecision, int32(0) + if opts.HasDecimalOpts { + precision, scale = opts.DecimalPrecision, opts.DecimalScale + } + return colSpec{ + arrowType: fmt.Sprintf("&arrow.%sType{Precision: %d, Scale: %d}", arrowName, precision, scale), + builderType: "*array." + arrowName + "Builder", + appendStmt: simpleAppend(), + }, nil +} + +// timeSpec maps a time.Time field. An untagged time.Time becomes +// TIMESTAMP(ns, UTC), which arreflect's struct-field path does not currently +// produce; see the divergence note in README.md. +func timeSpec(opts tagOpts) (colSpec, error) { + if err := checkOpts(opts, "time.Time", optSupport{temporal: true}); err != nil { + return colSpec{}, err + } + switch opts.Temporal { + case "date32": + return colSpec{ + arrowType: "arrow.FixedWidthTypes.Date32", + builderType: "*array.Date32Builder", + appendStmt: func(bld, _, val string) string { + return fmt.Sprintf("%s.Append(arrow.Date32FromTime(%s))", bld, val) + }, + }, nil + case "date64": + return colSpec{ + arrowType: "arrow.FixedWidthTypes.Date64", + builderType: "*array.Date64Builder", + appendStmt: func(bld, _, val string) string { + return fmt.Sprintf("%s.Append(arrow.Date64FromTime(%s))", bld, val) + }, + }, nil + case "time32": + return timeOfDaySpec("&arrow.Time32Type{Unit: arrow.Millisecond}", "*array.Time32Builder", "arrow.Time32", " / 1e6"), nil + case "time64": + return timeOfDaySpec("&arrow.Time64Type{Unit: arrow.Nanosecond}", "*array.Time64Builder", "arrow.Time64", ""), nil + default: + return colSpec{ + arrowType: `&arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}`, + builderType: "*array.TimestampBuilder", + appendStmt: func(bld, recv, _ string) string { + return fmt.Sprintf("%s.Append(arrow.Timestamp(%s.UnixNano()))", bld, recv) + }, + }, nil + } +} + +// timeOfDaySpec renders arreflect's timeOfDayNanos inline, as a block rather +// than a package-level helper so two generated files in one package cannot +// collide over the helper's name. +func timeOfDaySpec(arrowType, builderType, cast, divisor string) colSpec { + return colSpec{ + arrowType: arrowType, + builderType: builderType, + needsTime: true, + appendStmt: func(bld, recv, _ string) string { + return fmt.Sprintf(`{ +tod := %s.UTC() +%s.Append(%s(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds()%s)) +}`, recv, bld, cast, divisor) + }, + } +} + +func byteSliceSpec(opts tagOpts) (colSpec, error) { + if err := checkOpts(opts, "[]byte", optSupport{largeView: true, dict: true}); err != nil { + return colSpec{}, err + } + // A nil []byte is null in arreflect even when the column is not nullable. + if opts.Dict { + spec := dictSpec("arrow.BinaryTypes.Binary", "*array.BinaryDictionaryBuilder", simpleAppend()) + spec.nilable = true + return spec, nil + } + switch { + case opts.View: + return colSpec{arrowType: "arrow.BinaryTypes.BinaryView", builderType: "*array.BinaryViewBuilder", appendStmt: simpleAppend(), nilable: true}, nil + case opts.Large: + return colSpec{arrowType: "arrow.BinaryTypes.LargeBinary", builderType: "*array.BinaryBuilder", appendStmt: simpleAppend(), nilable: true}, nil + default: + return colSpec{arrowType: "arrow.BinaryTypes.Binary", builderType: "*array.BinaryBuilder", appendStmt: simpleAppend(), nilable: true}, nil + } +} + +func stringSpec(opts tagOpts) (colSpec, error) { + if err := checkOpts(opts, "string", optSupport{largeView: true, dict: true}); err != nil { + return colSpec{}, err + } + if opts.Dict { + return dictSpec("arrow.BinaryTypes.String", "*array.BinaryDictionaryBuilder", func(bld, _, val string) string { + return fmt.Sprintf("%s.AppendString(%s)", bld, val) + }), nil + } + switch { + case opts.View: + return colSpec{arrowType: "arrow.BinaryTypes.StringView", builderType: "*array.StringViewBuilder", appendStmt: simpleAppend()}, nil + case opts.Large: + return colSpec{arrowType: "arrow.BinaryTypes.LargeString", builderType: "*array.LargeStringBuilder", appendStmt: simpleAppend()}, nil + default: + return colSpec{arrowType: "arrow.BinaryTypes.String", builderType: "*array.StringBuilder", appendStmt: simpleAppend()}, nil + } +} + +// dictSpec wraps a value type in Dictionary, matching +// arreflect.applyEncodingOpts. Dictionary appends can fail. +func dictSpec(valueType, builderType string, appendStmt func(bld, recv, val string) string) colSpec { + return colSpec{ + arrowType: fmt.Sprintf("&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: %s}", valueType), + builderType: builderType, + fallible: true, + appendStmt: appendStmt, + } +} + +// numeric describes one Go basic kind's Arrow column. +type numeric struct { + arrowType string // e.g. "arrow.PrimitiveTypes.Int64" + builder string // e.g. "Int64" + cast string // non-empty when the Go type is wider-named than the column ("int" -> int64) +} + +var numerics = map[types.BasicKind]numeric{ + types.Int8: {"arrow.PrimitiveTypes.Int8", "Int8", ""}, + types.Int16: {"arrow.PrimitiveTypes.Int16", "Int16", ""}, + types.Int32: {"arrow.PrimitiveTypes.Int32", "Int32", ""}, + types.Int64: {"arrow.PrimitiveTypes.Int64", "Int64", ""}, + types.Int: {"arrow.PrimitiveTypes.Int64", "Int64", "int64"}, + types.Uint8: {"arrow.PrimitiveTypes.Uint8", "Uint8", ""}, + types.Uint16: {"arrow.PrimitiveTypes.Uint16", "Uint16", ""}, + types.Uint32: {"arrow.PrimitiveTypes.Uint32", "Uint32", ""}, + types.Uint64: {"arrow.PrimitiveTypes.Uint64", "Uint64", ""}, + types.Uint: {"arrow.PrimitiveTypes.Uint64", "Uint64", "uint64"}, + types.Float32: {"arrow.PrimitiveTypes.Float32", "Float32", ""}, + types.Float64: {"arrow.PrimitiveTypes.Float64", "Float64", ""}, +} + +func basicSpec(b *types.Basic, opts tagOpts) (colSpec, error) { + switch b.Kind() { + case types.String: + return stringSpec(opts) + case types.Bool: + if err := checkOpts(opts, "bool", optSupport{}); err != nil { + return colSpec{}, err + } + return colSpec{arrowType: "arrow.FixedWidthTypes.Boolean", builderType: "*array.BooleanBuilder", appendStmt: simpleAppend()}, nil + } + + n, ok := numerics[b.Kind()] + if !ok { + return colSpec{}, fmt.Errorf("type %s is not supported", b) + } + if err := checkOpts(opts, b.Name(), optSupport{dict: true}); err != nil { + return colSpec{}, err + } + appendStmt := simpleAppend() + if n.cast != "" { + appendStmt = castAppend(n.cast) + } + if opts.Dict { + return dictSpec(n.arrowType, "*array."+n.builder+"DictionaryBuilder", appendStmt), nil + } + return colSpec{arrowType: n.arrowType, builderType: "*array." + n.builder + "Builder", appendStmt: appendStmt}, nil +} + +func simpleAppend() func(bld, recv, val string) string { + return func(bld, _, val string) string { return fmt.Sprintf("%s.Append(%s)", bld, val) } +} + +func castAppend(cast string) func(bld, recv, val string) string { + return func(bld, _, val string) string { return fmt.Sprintf("%s.Append(%s(%s))", bld, cast, val) } +} diff --git a/arrgen/arrgen/tag.go b/arrgen/arrgen/tag.go new file mode 100644 index 000000000..f375bcfca --- /dev/null +++ b/arrgen/arrgen/tag.go @@ -0,0 +1,140 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import ( + "fmt" + "strconv" + "strings" +) + +// tagOpts is the parsed form of one `arrow:"..."` struct tag. It mirrors the +// unexported tagOpts in arreflect; keep the two in sync. +type tagOpts struct { + Name string + Skip bool + Dict bool + View bool + REE bool + Large bool + DecimalPrecision int32 + DecimalScale int32 + HasDecimalOpts bool + Temporal string // "", "timestamp", "date32", "date64", "time32", "time64" +} + +// parseTag parses the value of an `arrow` struct tag. Unlike arreflect, which +// defers diagnostics, an unusable tag is an error. +func parseTag(tag string) (tagOpts, error) { + if tag == "-" { + return tagOpts{Skip: true}, nil + } + + name, rest, _ := strings.Cut(tag, ",") + opts := tagOpts{Name: name} + if rest == "" { + return opts, nil + } + for _, token := range splitTagTokens(rest) { + if err := applyTagToken(&opts, token); err != nil { + return tagOpts{}, err + } + } + return opts, nil +} + +// splitTagTokens splits the option list on commas outside parentheses, so +// decimal(18,2) stays one token. +func splitTagTokens(rest string) []string { + var tokens []string + depth, start := 0, 0 + for i := 0; i < len(rest); i++ { + switch rest[i] { + case '(': + depth++ + case ')': + depth-- + case ',': + if depth == 0 { + tokens = append(tokens, strings.TrimSpace(rest[start:i])) + start = i + 1 + } + } + } + if start < len(rest) { + tokens = append(tokens, strings.TrimSpace(rest[start:])) + } + return tokens +} + +func applyTagToken(opts *tagOpts, token string) error { + if strings.HasPrefix(token, "decimal(") && strings.HasSuffix(token, ")") { + return parseDecimalOpt(opts, token) + } + switch token { + case "dict": + opts.Dict = true + case "view": + opts.View = true + case "ree": + opts.REE = true + case "large": + opts.Large = true + case "date32", "date64", "time32", "time64", "timestamp": + opts.Temporal = token + default: + return fmt.Errorf("unknown option %q", token) + } + return nil +} + +func parseDecimalOpt(opts *tagOpts, token string) error { + inner := strings.TrimSuffix(strings.TrimPrefix(token, "decimal("), ")") + parts := strings.SplitN(inner, ",", 2) + if len(parts) != 2 { + return fmt.Errorf("invalid decimal tag %q: expected decimal(precision,scale)", token) + } + p, err := strconv.ParseInt(strings.TrimSpace(parts[0]), 10, 32) + if err != nil { + return fmt.Errorf("invalid decimal tag %q: precision %q is not an integer", token, strings.TrimSpace(parts[0])) + } + s, err := strconv.ParseInt(strings.TrimSpace(parts[1]), 10, 32) + if err != nil { + return fmt.Errorf("invalid decimal tag %q: scale %q is not an integer", token, strings.TrimSpace(parts[1])) + } + opts.HasDecimalOpts = true + opts.DecimalPrecision = int32(p) + opts.DecimalScale = int32(s) + return nil +} + +// validate rejects the option combinations arreflect also rejects. +func (o tagOpts) validate() error { + if o.REE { + return fmt.Errorf("ree is not supported on a struct field; use arreflect.FromSlice with WithREE at the top level") + } + n := 0 + for _, set := range []bool{o.Dict, o.View, o.REE} { + if set { + n++ + } + } + if n > 1 { + return fmt.Errorf("conflicting options: at most one of dict, view, ree may be set") + } + return nil +} diff --git a/arrgen/arrgen/template.go b/arrgen/arrgen/template.go new file mode 100644 index 000000000..9ab8f9941 --- /dev/null +++ b/arrgen/arrgen/template.go @@ -0,0 +1,138 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import "text/template" + +// fileTemplate renders the whole generated file. The emitted code imports only +// arrow, arrow/array and arrow/memory, plus time for time-of-day columns. +var fileTemplate = template.Must(template.New("arrgen").Parse( + `{{if .Header}}{{.Header}} + +{{end}}// Code generated by arrgen. DO NOT EDIT. + +package {{.Package}} + +import ( +{{- if .NeedTime}} + "time" +{{end}} + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) +{{range .Types}} +// {{.SchemaVar}} is the Arrow schema of {{.GoName}}, resolved from its arrow +// struct tags. +var {{.SchemaVar}} = arrow.NewSchema([]arrow.Field{ +{{- range .Fields}} + {Name: {{printf "%q" .Name}}, Type: {{.ArrowType}}, Nullable: {{.Nullable}}}, +{{- end}} +}, nil) + +// {{.SchemaFunc}} returns the Arrow schema encoded by {{.AppenderType}}. +func {{.SchemaFunc}}() *arrow.Schema { return {{.SchemaVar}} } + +// {{.AppenderType}} converts {{.GoName}} values into Arrow record batches +// without reflection. +// +// It is not safe for concurrent use. Release it when done. +type {{.AppenderType}} struct { + rb *array.RecordBuilder +{{- range .Fields}} + {{.BuilderVar}} {{.BuilderType}} +{{- end}} + err error +} + +// {{.CtorName}} returns an appender that builds {{.GoName}} batches with mem, +// or memory.DefaultAllocator when mem is nil. The typed field builders are +// resolved once here, so Append does no type assertions. +func {{.CtorName}}(mem memory.Allocator) *{{.AppenderType}} { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, {{.SchemaVar}}) + return &{{.AppenderType}}{ + rb: rb, +{{- range .Fields}} + {{.BuilderVar}}: rb.Field({{.Index}}).({{.BuilderType}}), +{{- end}} + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *{{.AppenderType}}) Schema() *arrow.Schema { return {{.SchemaVar}} } + +// RecordBuilder returns the underlying builder. Appending through it directly +// is allowed as long as every column stays the same length. +func (a *{{.AppenderType}}) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends keeps the append path free of allocations. +func (a *{{.AppenderType}}) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *{{.AppenderType}}) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller can reuse a single {{.GoName}} variable across calls. +func (a *{{.AppenderType}}) Append(v *{{.GoName}}) { +{{- range .Fields}} + {{.AppendStmt}} +{{- end}} +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *{{.AppenderType}}) AppendSlice(vs []{{.GoName}}) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch returns the appended rows as a record batch and resets the +// appender. The caller must Release the batch. +func (a *{{.AppenderType}}) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail. The error is sticky: NewRecordBatch does not clear it. +func (a *{{.AppenderType}}) Err() error { return a.err } +{{if .AnyFallible}} +// setErr keeps the first error. +func (a *{{.AppenderType}}) setErr(err error) { + if err != nil && a.err == nil { + a.err = err + } +} +{{end}} +// Release releases the appender's builders. +func (a *{{.AppenderType}}) Release() { a.rb.Release() } + +// {{.BatchFunc}} builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[{{.GoName}}](vs, mem). The +// caller must Release the batch. +func {{.BatchFunc}}(mem memory.Allocator, vs []{{.GoName}}) (arrow.RecordBatch, error) { + a := {{.CtorName}}(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} +{{end}}`)) diff --git a/arrgen/arrgen/testdata/basic.golden b/arrgen/arrgen/testdata/basic.golden new file mode 100644 index 000000000..1686ca0b8 --- /dev/null +++ b/arrgen/arrgen/testdata/basic.golden @@ -0,0 +1,226 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by arrgen. DO NOT EDIT. + +package basic + +import ( + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// metricArrowSchema is the Arrow schema of Metric, resolved from its arrow +// struct tags. +var metricArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "host", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: false}, + {Name: "count", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "Blob", Type: arrow.BinaryTypes.Binary, Nullable: false}, +}, nil) + +// MetricSchema returns the Arrow schema encoded by MetricAppender. +func MetricSchema() *arrow.Schema { return metricArrowSchema } + +// MetricAppender converts Metric values into Arrow record batches +// without reflection. +// +// It is not safe for concurrent use. Release it when done. +type MetricAppender struct { + rb *array.RecordBuilder + b0 *array.TimestampBuilder + b1 *array.BinaryDictionaryBuilder + b2 *array.Int64Builder + b3 *array.BinaryBuilder + err error +} + +// NewMetricAppender returns an appender that builds Metric batches with mem, +// or memory.DefaultAllocator when mem is nil. The typed field builders are +// resolved once here, so Append does no type assertions. +func NewMetricAppender(mem memory.Allocator) *MetricAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, metricArrowSchema) + return &MetricAppender{ + rb: rb, + b0: rb.Field(0).(*array.TimestampBuilder), + b1: rb.Field(1).(*array.BinaryDictionaryBuilder), + b2: rb.Field(2).(*array.Int64Builder), + b3: rb.Field(3).(*array.BinaryBuilder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *MetricAppender) Schema() *arrow.Schema { return metricArrowSchema } + +// RecordBuilder returns the underlying builder. Appending through it directly +// is allowed as long as every column stays the same length. +func (a *MetricAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends keeps the append path free of allocations. +func (a *MetricAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller can reuse a single Metric variable across calls. +func (a *MetricAppender) Append(v *Metric) { + a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) + a.setErr(a.b1.AppendString(v.Host)) + if v.Count == nil { + a.b2.AppendNull() + } else { + a.b2.Append(*v.Count) + } + if v.Blob == nil { + a.b3.AppendNull() + } else { + a.b3.Append(v.Blob) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *MetricAppender) AppendSlice(vs []Metric) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch returns the appended rows as a record batch and resets the +// appender. The caller must Release the batch. +func (a *MetricAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail. The error is sticky: NewRecordBatch does not clear it. +func (a *MetricAppender) Err() error { return a.err } + +// setErr keeps the first error. +func (a *MetricAppender) setErr(err error) { + if err != nil && a.err == nil { + a.err = err + } +} + +// Release releases the appender's builders. +func (a *MetricAppender) Release() { a.rb.Release() } + +// MetricRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Metric](vs, mem). The +// caller must Release the batch. +func MetricRecordBatch(mem memory.Allocator, vs []Metric) (arrow.RecordBatch, error) { + a := NewMetricAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} + +// readingArrowSchema is the Arrow schema of reading, resolved from its arrow +// struct tags. +var readingArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "at", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, + {Name: "value", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, +}, nil) + +// readingSchema returns the Arrow schema encoded by readingAppender. +func readingSchema() *arrow.Schema { return readingArrowSchema } + +// readingAppender converts reading values into Arrow record batches +// without reflection. +// +// It is not safe for concurrent use. Release it when done. +type readingAppender struct { + rb *array.RecordBuilder + b0 *array.Date32Builder + b1 *array.Float64Builder + err error +} + +// newReadingAppender returns an appender that builds reading batches with mem, +// or memory.DefaultAllocator when mem is nil. The typed field builders are +// resolved once here, so Append does no type assertions. +func newReadingAppender(mem memory.Allocator) *readingAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, readingArrowSchema) + return &readingAppender{ + rb: rb, + b0: rb.Field(0).(*array.Date32Builder), + b1: rb.Field(1).(*array.Float64Builder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *readingAppender) Schema() *arrow.Schema { return readingArrowSchema } + +// RecordBuilder returns the underlying builder. Appending through it directly +// is allowed as long as every column stays the same length. +func (a *readingAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends keeps the append path free of allocations. +func (a *readingAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *readingAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller can reuse a single reading variable across calls. +func (a *readingAppender) Append(v *reading) { + a.b0.Append(arrow.Date32FromTime(v.At)) + a.b1.Append(v.Value) +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *readingAppender) AppendSlice(vs []reading) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch returns the appended rows as a record batch and resets the +// appender. The caller must Release the batch. +func (a *readingAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail. The error is sticky: NewRecordBatch does not clear it. +func (a *readingAppender) Err() error { return a.err } + +// Release releases the appender's builders. +func (a *readingAppender) Release() { a.rb.Release() } + +// readingRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[reading](vs, mem). The +// caller must Release the batch. +func readingRecordBatch(mem memory.Allocator, vs []reading) (arrow.RecordBatch, error) { + a := newReadingAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} diff --git a/arrgen/arrgen/testdata/basic/basic.go b/arrgen/arrgen/testdata/basic/basic.go new file mode 100644 index 000000000..3bee85d9c --- /dev/null +++ b/arrgen/arrgen/testdata/basic/basic.go @@ -0,0 +1,37 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package basic is the fixture behind arrgen's golden test. It sits under +// testdata so the go tool does not build it; the generator loads it by +// directory. +package basic + +import "time" + +// Metric is an exported type, so its generated API is exported too. +type Metric struct { + Time time.Time `arrow:"ts"` + Host string `arrow:"host,dict"` + Count *int64 `arrow:"count"` + Blob []byte + Local string `arrow:"-"` +} + +// reading is unexported, so the generated appender is unexported as well. +type reading struct { + At time.Time `arrow:"at,date32"` + Value float64 `arrow:"value"` +} diff --git a/arrgen/arrgen/testdata/errors/errors.go b/arrgen/arrgen/testdata/errors/errors.go new file mode 100644 index 000000000..0a4472008 --- /dev/null +++ b/arrgen/arrgen/testdata/errors/errors.go @@ -0,0 +1,126 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package errors holds one struct per rejection the generator makes. +package errors + +import "time" + +// Inner is a nested struct used by Nested and Embedded. +type Inner struct { + A int64 +} + +// ID is a defined scalar type, which arreflect cannot map either. +type ID int64 + +type Embedded struct { + Inner + B int64 +} + +type Nested struct { + Field Inner +} + +type SliceField struct { + Field []int64 +} + +type ArrayField struct { + Field [3]int64 +} + +type MapField struct { + Field map[string]int64 +} + +type NamedScalar struct { + Field ID +} + +type UnknownOption struct { + Field int64 `arrow:"f,nope"` +} + +type BadDecimalTag struct { + Field int64 `arrow:"f,decimal(a,2)"` +} + +type ShortDecimalTag struct { + Field int64 `arrow:"f,decimal(4)"` +} + +type DuplicateNames struct { + A int64 `arrow:"same"` + B int64 `arrow:"same"` +} + +type TemporalOnInt struct { + Field int64 `arrow:"f,date32"` +} + +type DecimalOnInt struct { + Field int64 `arrow:"f,decimal(10,2)"` +} + +type DictOnBool struct { + Field bool `arrow:"f,dict"` +} + +type ViewOnInt struct { + Field int64 `arrow:"f,view"` +} + +type LargeOnInt struct { + Field int64 `arrow:"f,large"` +} + +type RunEndEncoded struct { + Field string `arrow:"f,ree"` +} + +type DictAndView struct { + Field string `arrow:"f,dict,view"` +} + +type DictAndLarge struct { + Field string `arrow:"f,dict,large"` +} + +type NoColumns struct { + Field string `arrow:"-"` + hidden int64 +} + +type DoublePointer struct { + Field **int64 +} + +type TimeSlice struct { + Field []time.Time +} + +// NotAStruct is a named type whose underlying type is not a struct. +type NotAStruct int64 + +type Collide struct { + Field int64 +} + +type collide struct { + Field int64 +} diff --git a/arrgen/arrgen/testdata/pregen/pregen.go b/arrgen/arrgen/testdata/pregen/pregen.go new file mode 100644 index 000000000..799ac0b68 --- /dev/null +++ b/arrgen/arrgen/testdata/pregen/pregen.go @@ -0,0 +1,37 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package pregen is the state a package is in the first time the generator runs +// against it: the call sites exist, the generated file does not, and so the +// package does not type-check. +package pregen + +import ( + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +type Sample struct { + At time.Time `arrow:"at"` + Value float64 `arrow:"value"` +} + +// Encode calls into the file arrgen has not written yet. +func Encode(mem memory.Allocator, samples []Sample) (arrow.RecordBatch, error) { + return SampleRecordBatch(mem, samples) +} diff --git a/arrgen/cmd/arrgen/main.go b/arrgen/cmd/arrgen/main.go new file mode 100644 index 000000000..a5052fd57 --- /dev/null +++ b/arrgen/cmd/arrgen/main.go @@ -0,0 +1,104 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Command arrgen generates zero-reflection Arrow appenders for Go struct types. +// +// It reads the same arrow:"..." struct tags that +// github.com/apache/arrow-go/v18/arrow/array/arreflect interprets at runtime, +// but reads them once, when you run it, and writes typed Go source instead. +// +// Usage: +// +// arrgen -type Metric [-type Other] [-output metric_arrow.go] [-header LICENSE.txt] +// +// Typically it is invoked through go:generate, next to the type: +// +// //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +// +// then "go generate ./..." writes metric_arrow.go beside the struct. Check the +// result in, as you would with stringer or easyjson, and regenerate when the +// struct changes. +package main + +import ( + "flag" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/apache/arrow-go/arrgen" +) + +func main() { + if err := run(os.Args[1:]); err != nil { + fmt.Fprintf(os.Stderr, "arrgen: %v\n", err) + os.Exit(1) + } +} + +// typeList collects -type values, so both "-type A,B" and "-type A -type B" work. +type typeList []string + +func (t *typeList) String() string { return strings.Join(*t, ",") } + +func (t *typeList) Set(v string) error { + for _, p := range strings.Split(v, ",") { + if p = strings.TrimSpace(p); p != "" { + *t = append(*t, p) + } + } + return nil +} + +func run(argv []string) error { + fs := flag.NewFlagSet("arrgen", flag.ContinueOnError) + var types typeList + fs.Var(&types, "type", "struct type name to generate for; repeatable, or comma-separated (required)") + dir := fs.String("dir", ".", "directory of the package holding the types") + output := fs.String("output", "", "output file, relative to -dir (default _arrow.go)") + header := fs.String("header", "", "file whose contents are copied above the generated-code marker, e.g. a license header") + if err := fs.Parse(argv); err != nil { + return err + } + if len(types) == 0 { + return fmt.Errorf("-type is required") + } + + cfg := arrgen.Config{Dir: *dir, Types: types} + if *header != "" { + b, err := os.ReadFile(*header) + if err != nil { + return fmt.Errorf("reading -header: %w", err) + } + cfg.Header = string(b) + } + + src, err := arrgen.Generate(cfg) + if err != nil { + return err + } + + out := *output + if out == "" { + out = strings.ToLower(types[0]) + "_arrow.go" + } + path := filepath.Join(*dir, out) + if err := os.WriteFile(path, src, 0o644); err != nil { + return fmt.Errorf("writing %s: %w", path, err) + } + return nil +} diff --git a/arrgen/doc.go b/arrgen/doc.go new file mode 100644 index 000000000..f8e27cd79 --- /dev/null +++ b/arrgen/doc.go @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package arrgen generates zero-reflection Arrow encoders for Go struct types. +// +// It is the code-generation counterpart to +// [github.com/apache/arrow-go/v18/arrow/array/arreflect]: arreflect interprets +// a struct's `arrow:"..."` tags with reflection on every value, while arrgen +// reads exactly the same tags once, at generate time, and emits typed Go source +// that appends struct fields straight into typed Arrow builders. +// +// The two paths are interchangeable. For every struct arrgen accepts, the +// generated schema equals arreflect.InferSchema and the generated record batch +// equals the one arreflect.RecordFromSlice builds, so adopting the generated +// code is a call-site change and nothing else. +// +// # Usage +// +// Put a go:generate directive next to the struct: +// +// //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +// +// then run `go generate ./...`, which writes metric_arrow.go next to the type. +// Nothing in the arrow-go module depends on arrgen, and the code it emits +// imports only arrow, arrow/array and arrow/memory, so the generator itself is +// a build-time-only dependency. +// +// # Supported field types +// +// bool, int8/16/32/64, int, uint8/16/32/64, uint, float32/64, string, []byte, +// time.Time, time.Duration, decimal.Decimal32, decimal.Decimal64, +// decimal128.Num and decimal256.Num, plus a pointer to any of those for a +// nullable column. +// +// Tag options mirror arreflect: a leading name, "-" to skip a field, the +// temporal overrides date32, date64, time32, time64 and timestamp, dict, view, +// large, and decimal(precision,scale). +// +// Anything arrgen cannot map exactly the way arreflect would - nested structs, +// slices other than []byte, arrays, maps, embedded fields, named scalar types - +// is a generate-time error naming the field, never a silently dropped column. +// Those structs still work with arreflect at runtime. +package arrgen diff --git a/arrgen/drift_test.go b/arrgen/drift_test.go new file mode 100644 index 000000000..40f840460 --- /dev/null +++ b/arrgen/drift_test.go @@ -0,0 +1,72 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen_test + +import ( + "bytes" + "os" + "path/filepath" + "testing" + + "github.com/apache/arrow-go/arrgen" +) + +// TestCheckedInFilesAreUpToDate regenerates every committed arrgen output in +// this module and fails if the result differs from what is on disk. +// +// Generated code that is checked in drifts from its source the first time +// somebody edits a struct and forgets to rerun go generate, and the failure +// then shows up as a mysteriously wrong column rather than as a broken build. +// This turns that into a test failure at the moment of the edit. +func TestCheckedInFilesAreUpToDate(t *testing.T) { + header, err := os.ReadFile("license_header.txt") + if err != nil { + t.Fatalf("reading license header: %v", err) + } + + // These mirror the go:generate directives next to each type. Add a line + // here whenever a new generated file is committed to this module. + targets := []struct { + dir string + types []string + file string + }{ + {"example", []string{"Metric"}, "metric_arrow.go"}, + {"internal/gentypes", []string{"Row", "Fixed"}, "row_arrow.go"}, + } + + for _, target := range targets { + t.Run(target.file, func(t *testing.T) { + got, err := arrgen.Generate(arrgen.Config{ + Dir: target.dir, + Types: target.types, + Header: string(header), + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + path := filepath.Join(target.dir, target.file) + want, err := os.ReadFile(path) + if err != nil { + t.Fatalf("reading %s: %v", path, err) + } + if !bytes.Equal(got, want) { + t.Errorf("%s is out of date; run go generate ./...", path) + } + }) + } +} diff --git a/arrgen/example/example_test.go b/arrgen/example/example_test.go new file mode 100644 index 000000000..dd2f3c644 --- /dev/null +++ b/arrgen/example/example_test.go @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package example_test + +import ( + "fmt" + "time" + + "github.com/apache/arrow-go/arrgen/example" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// Example encodes a slice of rows in one call. It is the drop-in replacement +// for arreflect.RecordFromSlice[Metric](metrics, mem). +func Example() { + cpuLoad := 0.75 + metrics := []example.Metric{ + {Time: time.Unix(0, 0).UTC(), Host: "web-1", CPU: 0.5, Value: &cpuLoad}, + {Time: time.Unix(1, 0).UTC(), Host: "web-2", CPU: 1.5}, // Value is nil: a null + } + + rec, err := example.MetricRecordBatch(memory.DefaultAllocator, metrics) + if err != nil { + fmt.Println("encode:", err) + return + } + defer rec.Release() + + fmt.Println("rows:", rec.NumRows()) + for i, col := range rec.Columns() { + fmt.Printf("%s: %v\n", rec.Schema().Field(i).Name, col) + } + // Output: + // rows: 2 + // ts: [0 1000000000] + // host: ["web-1" "web-2"] + // cpu: [0.5 1.5] + // value: [0.75 (null)] +} + +// ExampleMetricAppender streams rows in one at a time and cuts a batch at a +// size boundary, which is what an ingest loop actually looks like. The appender +// reads each row synchronously and never keeps a reference, so the same +// variable is reused for every row and the loop allocates nothing. +func ExampleMetricAppender() { + const batchSize = 2 + + a := example.NewMetricAppender(memory.DefaultAllocator) + defer a.Release() + a.Reserve(batchSize) + + var row example.Metric + for i := 0; i < 5; i++ { + row = example.Metric{ + Time: time.Unix(int64(i), 0).UTC(), + Host: fmt.Sprintf("web-%d", i), + CPU: float64(i), + } + a.Append(&row) + + if a.Len() == batchSize { + rec := a.NewRecordBatch() + fmt.Println("batch of", rec.NumRows()) + rec.Release() + a.Reserve(batchSize) + } + } + if a.Len() > 0 { + rec := a.NewRecordBatch() + fmt.Println("final batch of", rec.NumRows()) + rec.Release() + } + if err := a.Err(); err != nil { + fmt.Println("append:", err) + } + // Output: + // batch of 2 + // batch of 2 + // final batch of 1 +} + +// ExampleMetricSchema shows the schema the generator resolved from the tags. +// It is a package-level variable, so asking for it costs nothing. +func ExampleMetricSchema() { + for _, f := range example.MetricSchema().Fields() { + fmt.Printf("%s %s nullable=%t\n", f.Name, f.Type, f.Nullable) + } + // Output: + // ts timestamp[ns, tz=UTC] nullable=false + // host utf8 nullable=false + // cpu float64 nullable=false + // value float64 nullable=true +} diff --git a/arrgen/example/metric.go b/arrgen/example/metric.go new file mode 100644 index 000000000..b15679191 --- /dev/null +++ b/arrgen/example/metric.go @@ -0,0 +1,42 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package example is a worked, checked-in demonstration of arrgen. +// +// Metric carries the arrow struct tags; metric_arrow.go next to it is what +// arrgen emitted from them and is committed exactly as stringer or easyjson +// output would be. Regenerate it with: +// +// go generate ./... +// +// The package is importable rather than a main package so its tests can encode +// the same values through both the generated appender and arreflect and show +// they agree. +package example + +import "time" + +// Metric is a small, realistic row: a timestamp, two scalar columns, a nullable +// pointer column, and a field kept out of Arrow entirely. +type Metric struct { + Time time.Time `arrow:"ts"` + Host string `arrow:"host"` + CPU float64 `arrow:"cpu"` + Value *float64 `arrow:"value"` // nullable: a nil pointer appends null + Secret string `arrow:"-"` // never leaves the process +} + +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -header ../license_header.txt diff --git a/arrgen/example/metric_arrow.go b/arrgen/example/metric_arrow.go new file mode 100644 index 000000000..136666aa6 --- /dev/null +++ b/arrgen/example/metric_arrow.go @@ -0,0 +1,134 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by arrgen. DO NOT EDIT. + +package example + +import ( + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// metricArrowSchema is the Arrow schema of Metric, resolved from its arrow +// struct tags at generate time. +var metricArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "host", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: "cpu", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "value", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, +}, nil) + +// MetricSchema returns the Arrow schema encoded by MetricAppender. It is +// equal to the schema arreflect.InferSchema[Metric] infers at runtime. +func MetricSchema() *arrow.Schema { return metricArrowSchema } + +// MetricAppender converts Metric values into Arrow record batches +// without reflection: the struct tags were read when this file was generated, +// so appending a row is a handful of direct field loads into typed builders. +// +// It is not safe for concurrent use. Release it when done. +type MetricAppender struct { + rb *array.RecordBuilder + b0 *array.TimestampBuilder + b1 *array.StringBuilder + b2 *array.Float64Builder + b3 *array.Float64Builder + err error +} + +// NewMetricAppender returns an appender that builds Metric batches with mem, +// or memory.DefaultAllocator when mem is nil. +// +// The typed field builders are resolved once here rather than per row, so +// Append does no type assertions at all. +func NewMetricAppender(mem memory.Allocator) *MetricAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, metricArrowSchema) + return &MetricAppender{ + rb: rb, + b0: rb.Field(0).(*array.TimestampBuilder), + b1: rb.Field(1).(*array.StringBuilder), + b2: rb.Field(2).(*array.Float64Builder), + b3: rb.Field(3).(*array.Float64Builder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *MetricAppender) Schema() *arrow.Schema { return metricArrowSchema } + +// RecordBuilder returns the underlying builder, for callers that need an Arrow +// API this appender does not wrap. Appending through it directly is allowed as +// long as every column stays the same length. +func (a *MetricAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends is what keeps the append path free of allocations. +func (a *MetricAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller streaming rows can reuse a single Metric variable across calls. +func (a *MetricAppender) Append(v *Metric) { + a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) + a.b1.Append(v.Host) + a.b2.Append(v.CPU) + if v.Value == nil { + a.b3.AppendNull() + } else { + a.b3.Append(*v.Value) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *MetricAppender) AppendSlice(vs []Metric) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch snapshots the appended rows into a record batch and resets the +// appender for the next batch. The caller owns the batch and must Release it. +func (a *MetricAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail, so an appender over a schema without one always returns nil. +// +// The error is sticky: NewRecordBatch does not clear it, because a batch built +// after a failed append is missing a value in that column. +func (a *MetricAppender) Err() error { return a.err } + +// Release releases the appender's builders. +func (a *MetricAppender) Release() { a.rb.Release() } + +// MetricRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Metric](vs, mem) and +// produces an identical batch. The caller owns the batch and must Release it. +func MetricRecordBatch(mem memory.Allocator, vs []Metric) (arrow.RecordBatch, error) { + a := NewMetricAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} diff --git a/arrgen/generate.go b/arrgen/generate.go new file mode 100644 index 000000000..530bf1150 --- /dev/null +++ b/arrgen/generate.go @@ -0,0 +1,350 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import ( + "bytes" + "fmt" + "go/format" + "go/token" + "go/types" + "reflect" + "strings" + "unicode" + + "golang.org/x/tools/go/packages" +) + +// Config describes one generator run. +type Config struct { + // Dir is the directory of the package holding the struct types. go:generate + // runs in the package directory, so the command leaves this as ".". + Dir string + // Types are the struct type names to generate for, in output order. + Types []string + // Header, when non-empty, is emitted verbatim above the generated-code + // marker. Use it for a license header; it is expected to be Go comments. + Header string +} + +// Generate resolves cfg.Types in the package at cfg.Dir and returns the +// formatted source of their Arrow appenders. It never writes to disk, which +// keeps it usable from tests that only want to compare against a golden file. +// +// Output depends on nothing but the package's source, so regenerating an +// unchanged package reproduces the file byte for byte - the property the drift +// test in this module relies on. +func Generate(cfg Config) ([]byte, error) { + if len(cfg.Types) == 0 { + return nil, fmt.Errorf("arrgen: no types requested") + } + dir := cfg.Dir + if dir == "" { + dir = "." + } + pkg, loadErrs, err := loadPackage(dir) + if err != nil { + return nil, err + } + + file := genFile{Header: strings.TrimRight(cfg.Header, "\n"), Package: pkg.Name} + declared := make(map[string]string) // generated identifier -> type that claimed it + for _, name := range cfg.Types { + gt, err := genForType(pkg, name) + if err != nil { + return nil, withLoadContext(fmt.Errorf("arrgen: %s: %w", name, err), loadErrs) + } + if err := claimIdents(declared, name, gt); err != nil { + return nil, fmt.Errorf("arrgen: %s: %w", name, err) + } + file.NeedTime = file.NeedTime || gt.needTime + file.Types = append(file.Types, gt) + } + + var buf bytes.Buffer + if err := fileTemplate.Execute(&buf, file); err != nil { + return nil, fmt.Errorf("arrgen: rendering template: %w", err) + } + src, err := format.Source(buf.Bytes()) + if err != nil { + return nil, fmt.Errorf("arrgen: formatting generated source: %w\n--- raw ---\n%s", err, buf.String()) + } + return src, nil +} + +// loadPackage type-checks the package in dir and returns it along with any +// errors the type checker reported. +// +// Those errors are returned rather than raised, because the package a generator +// runs in usually does not compile yet: the code calling MetricRecordBatch is +// written before the file defining it exists. Refusing to run then would make +// the generator unusable exactly when it is needed. go/packages still resolves +// everything it can, so the struct being generated for is available; a load +// error only becomes the answer when resolving that struct fails, and +// withLoadContext then reports it as the likely cause. +func loadPackage(dir string) (*packages.Package, []string, error) { + cfg := &packages.Config{ + Mode: packages.NeedName | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax, + Dir: dir, + } + pkgs, err := packages.Load(cfg, ".") + if err != nil { + return nil, nil, fmt.Errorf("arrgen: loading package in %s: %w", dir, err) + } + if len(pkgs) != 1 { + return nil, nil, fmt.Errorf("arrgen: expected exactly one package in %s, found %d", dir, len(pkgs)) + } + var errs []string + for _, e := range pkgs[0].Errors { + errs = append(errs, e.Error()) + } + if pkgs[0].Types == nil { + return nil, nil, fmt.Errorf("arrgen: package in %s could not be type-checked: %s", dir, strings.Join(errs, "; ")) + } + return pkgs[0], errs, nil +} + +// withLoadContext appends the package's type errors to err, since a failure to +// resolve a type in a package that does not compile is usually a symptom of +// that rather than of the struct itself. +func withLoadContext(err error, loadErrs []string) error { + if len(loadErrs) == 0 { + return err + } + const maxShown = 5 + shown := loadErrs + suffix := "" + if len(shown) > maxShown { + shown, suffix = shown[:maxShown], fmt.Sprintf("\n ... and %d more", len(loadErrs)-maxShown) + } + return fmt.Errorf("%w\nthe package does not type-check, which may be the cause:\n %s%s", + err, strings.Join(shown, "\n "), suffix) +} + +func genForType(pkg *packages.Package, name string) (genType, error) { + st, err := findStruct(pkg, name) + if err != nil { + return genType{}, err + } + cols, err := collectColumns(st) + if err != nil { + return genType{}, err + } + if len(cols) == 0 { + return genType{}, fmt.Errorf("struct has no Arrow columns") + } + + exported := token.IsExported(name) + gt := genType{ + GoName: name, + SchemaVar: lowerFirst(name) + "ArrowSchema", + SchemaFunc: caseAs(exported, name+"Schema"), + AppenderType: caseAs(exported, name+"Appender"), + CtorName: caseAs(exported, "New"+upperFirst(name)+"Appender"), + BatchFunc: caseAs(exported, name+"RecordBatch"), + } + for i, c := range cols { + gt.needTime = gt.needTime || c.spec.needsTime + gt.AnyFallible = gt.AnyFallible || c.spec.fallible + gt.Fields = append(gt.Fields, genField{ + Index: i, + Name: c.name, + GoField: c.goField, + Nullable: c.nullable, + ArrowType: c.spec.arrowType, + BuilderType: c.spec.builderType, + BuilderVar: fmt.Sprintf("b%d", i), + AppendStmt: renderAppend(i, c), + }) + } + return gt, nil +} + +// claimIdents records the package-level names a type's generated code will +// declare, rejecting a second type that would declare the same one. Two types +// differing only in the case of their first letter collide this way; without +// this check the collision surfaces as a compile error in generated code the +// user did not write. +func claimIdents(declared map[string]string, name string, gt genType) error { + for _, ident := range []string{gt.SchemaVar, gt.SchemaFunc, gt.AppenderType, gt.CtorName, gt.BatchFunc} { + if prev, taken := declared[ident]; taken { + return fmt.Errorf("generated name %s collides with the one generated for %s; generate the two types into separate files", ident, prev) + } + declared[ident] = name + } + return nil +} + +func findStruct(pkg *packages.Package, name string) (*types.Struct, error) { + obj := pkg.Types.Scope().Lookup(name) + if obj == nil { + return nil, fmt.Errorf("type not found in package %s", pkg.Name) + } + named, ok := obj.Type().(*types.Named) + if !ok { + return nil, fmt.Errorf("not a named type") + } + st, ok := named.Underlying().(*types.Struct) + if !ok { + return nil, fmt.Errorf("not a struct type") + } + return st, nil +} + +// column is one resolved struct field. +type column struct { + name string + goField string + nullable bool + spec colSpec +} + +// collectColumns walks the struct's fields in declaration order, which is the +// order arreflect settles on for a flat struct. Embedded fields are rejected +// rather than promoted: arreflect resolves promoted names by breadth and tag, +// and quietly reimplementing those rules is how a generator ends up emitting a +// different schema than the runtime it claims to match. +func collectColumns(st *types.Struct) ([]column, error) { + var cols []column + seen := make(map[string]string, st.NumFields()) + + for i := 0; i < st.NumFields(); i++ { + f := st.Field(i) + if f.Anonymous() { + return nil, fmt.Errorf("field %s: embedded fields are not supported; give the field a name or encode the struct with arreflect", f.Name()) + } + if !f.Exported() { + continue + } + + tag, hasTag := reflect.StructTag(st.Tag(i)).Lookup("arrow") + var opts tagOpts + if hasTag { + var err error + if opts, err = parseTag(tag); err != nil { + return nil, fmt.Errorf("field %s: %w", f.Name(), err) + } + } + if opts.Skip { + continue + } + if err := opts.validate(); err != nil { + return nil, fmt.Errorf("field %s: %w", f.Name(), err) + } + + name := opts.Name + if name == "" { + name = f.Name() + } + if prev, dup := seen[name]; dup { + return nil, fmt.Errorf("fields %s and %s both map to column %q", prev, f.Name(), name) + } + seen[name] = f.Name() + + spec, nullable, err := resolveColumn(f.Type(), opts) + if err != nil { + return nil, fmt.Errorf("field %s: %w", f.Name(), err) + } + cols = append(cols, column{name: name, goField: f.Name(), nullable: nullable, spec: spec}) + } + return cols, nil +} + +// renderAppend emits the append statement for one column, wrapping it in the +// null guards the reflection path applies: a nil pointer is a null, and so is a +// nil []byte even when the column itself is not nullable. +func renderAppend(idx int, c column) string { + bld := fmt.Sprintf("a.b%d", idx) + val, recv := "v."+c.goField, "v."+c.goField + if c.nullable { + val, recv = "*v."+c.goField, "(*v."+c.goField+")" + } + + stmt := c.spec.appendStmt(bld, recv, val) + if c.spec.fallible { + stmt = "a.setErr(" + stmt + ")" + } + + var guard string + switch { + case c.nullable && c.spec.nilable: + guard = fmt.Sprintf("v.%s == nil || *v.%s == nil", c.goField, c.goField) + case c.nullable, c.spec.nilable: + guard = fmt.Sprintf("v.%s == nil", c.goField) + default: + return stmt + } + + // A statement that is already a block (a time-of-day column scopes a local) + // sheds its braces on the way into the guard's else, which would otherwise + // nest two sets for no reason. + inner := strings.TrimSuffix(strings.TrimPrefix(stmt, "{\n"), "\n}") + return fmt.Sprintf("if %s {\n%s.AppendNull()\n} else {\n%s\n}", guard, bld, inner) +} + +type genFile struct { + Header string + Package string + NeedTime bool + Types []genType +} + +type genType struct { + GoName string + SchemaVar string + SchemaFunc string + AppenderType string + CtorName string + BatchFunc string + Fields []genField + AnyFallible bool + + needTime bool +} + +type genField struct { + Index int + Name string + GoField string + Nullable bool + ArrowType string + BuilderType string + BuilderVar string + AppendStmt string +} + +func lowerFirst(s string) string { return mapFirst(s, unicode.ToLower) } +func upperFirst(s string) string { return mapFirst(s, unicode.ToUpper) } + +func mapFirst(s string, f func(rune) rune) string { + if s == "" { + return s + } + r := []rune(s) + r[0] = f(r[0]) + return string(r) +} + +// caseAs gives a generated identifier the same visibility as the struct it was +// derived from, so generating for an unexported type does not silently widen +// the package's API. +func caseAs(exported bool, s string) string { + if exported { + return upperFirst(s) + } + return lowerFirst(s) +} diff --git a/arrgen/generate_test.go b/arrgen/generate_test.go new file mode 100644 index 000000000..d2594b7d0 --- /dev/null +++ b/arrgen/generate_test.go @@ -0,0 +1,228 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen_test + +import ( + "bytes" + "flag" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/apache/arrow-go/arrgen" +) + +var update = flag.Bool("update", false, "rewrite the golden files instead of comparing against them") + +const goldenPath = "testdata/basic.golden" + +// TestGenerateGolden compares the whole emitted file against a checked-in +// copy. A golden file is worth more than a set of assertions about fragments +// here: what ships to users is the file, and a reviewer reading the diff of +// this one can see exactly what a generator change does to their code. +func TestGenerateGolden(t *testing.T) { + header, err := os.ReadFile("license_header.txt") + if err != nil { + t.Fatalf("reading license header: %v", err) + } + got, err := arrgen.Generate(arrgen.Config{ + Dir: "testdata/basic", + Types: []string{"Metric", "reading"}, + Header: string(header), + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + + if *update { + if err := os.WriteFile(goldenPath, got, 0o644); err != nil { + t.Fatalf("writing golden: %v", err) + } + return + } + + want, err := os.ReadFile(goldenPath) + if err != nil { + t.Fatalf("reading golden: %v (run: go test ./... -update)", err) + } + if !bytes.Equal(got, want) { + t.Errorf("generated output differs from %s; rerun with -update to accept\n--- got ---\n%s", goldenPath, got) + } +} + +// TestGenerateIsDeterministic guards the property the drift check in +// internal/gentypes depends on: the same package in, the same bytes out. +func TestGenerateIsDeterministic(t *testing.T) { + cfg := arrgen.Config{Dir: "testdata/basic", Types: []string{"Metric", "reading"}} + first, err := arrgen.Generate(cfg) + if err != nil { + t.Fatalf("Generate: %v", err) + } + second, err := arrgen.Generate(cfg) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if !bytes.Equal(first, second) { + t.Error("two runs over the same package produced different output") + } +} + +// TestGenerateHeaderOptional checks that the generated-code marker leads the +// file when no header is configured, which is what tools look for. +func TestGenerateHeaderOptional(t *testing.T) { + got, err := arrgen.Generate(arrgen.Config{Dir: "testdata/basic", Types: []string{"Metric"}}) + if err != nil { + t.Fatalf("Generate: %v", err) + } + const marker = "// Code generated by arrgen. DO NOT EDIT." + if !bytes.HasPrefix(got, []byte(marker)) { + t.Errorf("output does not start with %q:\n%s", marker, first(got, 200)) + } +} + +// TestGenerateRejects is the contract that arrgen never drops a column it does +// not understand. Every entry here is a struct that would otherwise be encoded +// wrongly, silently, or not at all. +func TestGenerateRejects(t *testing.T) { + tests := []struct { + typ string + want string + }{ + {"Embedded", "embedded fields are not supported"}, + {"Nested", "nested struct column"}, + {"SliceField", "list column for []int64 is not supported"}, + {"ArrayField", "is not supported"}, + {"MapField", "is not supported"}, + {"NamedScalar", "named type"}, + {"UnknownOption", `unknown option "nope"`}, + {"BadDecimalTag", "is not an integer"}, + {"ShortDecimalTag", "expected decimal(precision,scale)"}, + {"DuplicateNames", `both map to column "same"`}, + {"TemporalOnInt", "only valid on a time.Time field"}, + {"DecimalOnInt", "only valid on a decimal field"}, + {"DictOnBool", "dict is not supported on bool"}, + {"ViewOnInt", "view has no effect"}, + {"LargeOnInt", "large has no effect"}, + {"RunEndEncoded", "ree is not supported"}, + {"DictAndView", "at most one of dict, view, ree"}, + {"DictAndLarge", "cannot be combined with large"}, + {"NoColumns", "no Arrow columns"}, + {"DoublePointer", "multi-level pointer"}, + {"TimeSlice", "list column"}, + {"NotAStruct", "not a struct type"}, + {"Absent", "type not found"}, + } + + for _, tt := range tests { + t.Run(tt.typ, func(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{tt.typ}}) + if err == nil { + t.Fatalf("Generate(%s) succeeded, want an error mentioning %q", tt.typ, tt.want) + } + if !strings.Contains(err.Error(), tt.want) { + t.Errorf("Generate(%s) error = %q, want it to mention %q", tt.typ, err, tt.want) + } + if !strings.Contains(err.Error(), tt.typ) { + t.Errorf("Generate(%s) error = %q, want it to name the type", tt.typ, err) + } + }) + } +} + +// TestGenerateNamesTheField checks that a rejection points at the field, not +// just the struct: on a wide struct that is the difference between a fix and a +// hunt. +func TestGenerateNamesTheField(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Nested"}}) + if err == nil { + t.Fatal("Generate succeeded, want an error") + } + if !strings.Contains(err.Error(), "field Field") { + t.Errorf("error = %q, want it to name the offending field", err) + } +} + +// TestGenerateRejectsCollidingTypes covers two types in one file whose +// generated names would clash - the sort of thing that would otherwise surface +// as a compile error in code the user never wrote. +func TestGenerateRejectsCollidingTypes(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Collide", "collide"}}) + if err == nil { + t.Fatal("Generate succeeded, want a collision error") + } + if !strings.Contains(err.Error(), "collides with") { + t.Errorf("error = %q, want it to report the collision", err) + } +} + +func TestGenerateRejectsRepeatedType(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Collide", "Collide"}}) + if err == nil { + t.Fatal("Generate succeeded, want a collision error") + } +} + +// TestGenerateBeforePackageCompiles is the chicken-and-egg case: the call sites +// for the generated API are written first, so the package does not type-check +// until the generator has run. Refusing to run then would make the generator +// unusable exactly when it is needed. +func TestGenerateBeforePackageCompiles(t *testing.T) { + got, err := arrgen.Generate(arrgen.Config{Dir: "testdata/pregen", Types: []string{"Sample"}}) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if !bytes.Contains(got, []byte("func SampleRecordBatch(")) { + t.Error("generated output does not declare SampleRecordBatch") + } +} + +// TestGenerateReportsLoadErrors checks that when a type genuinely cannot be +// resolved, the package's own compile errors come along as the likely cause +// rather than leaving the user with a bare "type not found". +func TestGenerateReportsLoadErrors(t *testing.T) { + _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/pregen", Types: []string{"Absent"}}) + if err == nil { + t.Fatal("Generate succeeded, want an error") + } + if !strings.Contains(err.Error(), "does not type-check") { + t.Errorf("error = %q, want it to surface the package's compile errors", err) + } +} + +func TestGenerateNoTypes(t *testing.T) { + if _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/basic"}); err == nil { + t.Error("Generate with no types succeeded, want an error") + } +} + +func TestGenerateMissingPackage(t *testing.T) { + dir := filepath.Join(t.TempDir(), "empty") + if err := os.Mkdir(dir, 0o755); err != nil { + t.Fatal(err) + } + if _, err := arrgen.Generate(arrgen.Config{Dir: dir, Types: []string{"Metric"}}); err == nil { + t.Error("Generate in a directory with no Go files succeeded, want an error") + } +} + +func first(b []byte, n int) []byte { + if len(b) < n { + return b + } + return b[:n] +} diff --git a/arrgen/go.mod b/arrgen/go.mod new file mode 100644 index 000000000..ad12f8825 --- /dev/null +++ b/arrgen/go.mod @@ -0,0 +1,35 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +module github.com/apache/arrow-go/arrgen + +go 1.25.0 + +require ( + github.com/apache/arrow-go/v18 v18.7.0 + golang.org/x/tools v0.45.0 +) + +require ( + github.com/goccy/go-json v0.10.6 // indirect + github.com/google/flatbuffers v25.12.19+incompatible // indirect + github.com/klauspost/cpuid/v2 v2.4.0 // indirect + github.com/zeebo/xxh3 v1.1.0 // indirect + golang.org/x/exp v0.0.0-20260112195511-716be5621a96 // indirect + golang.org/x/mod v0.36.0 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/sys v0.47.0 // indirect +) diff --git a/arrgen/go.sum b/arrgen/go.sum new file mode 100644 index 000000000..963f34aaf --- /dev/null +++ b/arrgen/go.sum @@ -0,0 +1,44 @@ +github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM= +github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/apache/arrow-go/v18 v18.7.0 h1:Vw/i+cJyebUofT7JlqFpe65LrmwxULn166jjwStM4HY= +github.com/apache/arrow-go/v18 v18.7.0/go.mod h1:PM6IigLJkdMwIpeHXnymo+xZ52f42a9EYiLtRel4p/A= +github.com/apache/thrift v0.24.0 h1:zy31L1a49QTNB2bG1BBfMXol3yJrTH975G3pPubQVLQ= +github.com/apache/thrift v0.24.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU= +github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= +github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/klauspost/compress v1.19.0 h1:sXLILfc9jV2QYWkzFOPWStmcUVH2RHEB1JCdY2oVvCQ= +github.com/klauspost/compress v1.19.0/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw= +github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU= +github.com/pierrec/lz4/v4 v4.1.27 h1:+PhzhWDrjRj89TH2sw43nE3+4+W8lSxIuQadEHZyjUk= +github.com/pierrec/lz4/v4 v4.1.27/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= +github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96 h1:Z/6YuSHTLOHfNFdb8zVZomZr7cqNgTJvA8+Qz75D8gU= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96/go.mod h1:nzimsREAkjBCIEFtHiYkrJyT+2uy9YZJB7H1k68CXZU= +golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= +golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= +golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/arrgen/internal/gentypes/alloc_test.go b/arrgen/internal/gentypes/alloc_test.go new file mode 100644 index 000000000..2c6e8f10a --- /dev/null +++ b/arrgen/internal/gentypes/alloc_test.go @@ -0,0 +1,167 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gentypes_test + +import ( + "testing" + + "github.com/apache/arrow-go/arrgen/internal/gentypes" + "github.com/apache/arrow-go/v18/arrow/array/arreflect" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// TestStreamingAppendIsAllocationFree is the strongest allocation claim arrgen +// makes, so it is asserted rather than left to a benchmark to hint at: with +// room reserved up front, appending a row of fixed-width columns allocates +// nothing at all. The reflection path has no equivalent - it cannot encode +// anything until the caller has materialized the whole slice. +func TestStreamingAppendIsAllocationFree(t *testing.T) { + skipIfInstrumented(t) + const runs = 2000 + + a := gentypes.NewFixedAppender(memory.DefaultAllocator) + defer a.Release() + a.Reserve(runs + 8) // AllocsPerRun calls the function runs+1 times + + rows := makeFixed(4) + i := 0 + // The row is read through a single reused variable, the way a caller + // draining a stream would do it. + var row gentypes.Fixed + got := testing.AllocsPerRun(runs, func() { + row = rows[i%len(rows)] + i++ + a.Append(&row) + }) + if got != 0 { + t.Errorf("Append allocated %.2f times per row, want 0", got) + } + a.NewRecordBatch().Release() +} + +// TestSchemaIsAllocationFree pins the other cost generation removes outright: +// arreflect walks the struct and builds a fresh schema on every call, while the +// generated schema is a package-level variable. +func TestSchemaIsAllocationFree(t *testing.T) { + skipIfInstrumented(t) + if got := testing.AllocsPerRun(1000, func() { _ = gentypes.RowSchema() }); got != 0 { + t.Errorf("RowSchema allocated %.2f times per call, want 0", got) + } + + reflected := testing.AllocsPerRun(100, func() { + if _, err := arreflect.InferSchema[gentypes.Row](); err != nil { + t.Fatal(err) + } + }) + if reflected == 0 { + t.Skip("arreflect.InferSchema no longer allocates; the comparison is moot") + } + t.Logf("schema per call: arrgen 0 allocs, arreflect %.0f allocs", reflected) +} + +// TestBatchEncodingAllocatesLessThanReflection compares whole-batch encoding. +// The margin here is smaller than for streaming because both paths pay for the +// same Arrow buffers; what the generated path saves is the per-value reflection +// machinery on top of them. The threshold is deliberately loose - the point is +// to catch a regression that erases the advantage, not to pin an exact number +// that Arrow's own buffer growth policy is free to change. +func TestBatchEncodingAllocatesLessThanReflection(t *testing.T) { + skipIfInstrumented(t) + const rows = 1024 + fixtures := makeFixed(rows) + mem := memory.DefaultAllocator + + reflected := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := arreflect.RecordFromSlice(fixtures, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + generated := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := gentypes.FixedRecordBatch(mem, fixtures) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + + t.Logf("%d-row batch: arrgen %d allocs / %d ns, arreflect %d allocs / %d ns", + rows, generated.AllocsPerOp(), generated.NsPerOp(), reflected.AllocsPerOp(), reflected.NsPerOp()) + + if generated.AllocsPerOp() >= reflected.AllocsPerOp() { + t.Errorf("generated encoder allocated %d times per batch, want fewer than arreflect's %d", + generated.AllocsPerOp(), reflected.AllocsPerOp()) + } + if got, limit := generated.AllocsPerOp(), reflected.AllocsPerOp()*3/4; got > limit { + t.Errorf("generated encoder allocated %d times per batch, want at most %d (three quarters of arreflect's %d)", + got, limit, reflected.AllocsPerOp()) + } +} + +// TestBatchEncodingIsFasterThanReflection guards the headline claim. Wall-clock +// assertions are noisy on shared CI, so the bar is set at a fraction of the +// margin actually measured (around 3.5x on the fixed-width fixture) and the +// test steps aside for -short runs. +func TestBatchEncodingIsFasterThanReflection(t *testing.T) { + skipIfInstrumented(t) + if testing.Short() { + t.Skip("timing comparison skipped in short mode") + } + const rows = 1024 + fixtures := makeFixed(rows) + mem := memory.DefaultAllocator + + reflected := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := arreflect.RecordFromSlice(fixtures, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + generated := testing.Benchmark(func(b *testing.B) { + for b.Loop() { + rec, err := gentypes.FixedRecordBatch(mem, fixtures) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + }) + if generated.NsPerOp() == 0 || reflected.NsPerOp() == 0 { + t.Skip("benchmark produced no timing data") + } + + speedup := float64(reflected.NsPerOp()) / float64(generated.NsPerOp()) + t.Logf("%d-row batch: arrgen %d ns, arreflect %d ns (%.2fx)", rows, generated.NsPerOp(), reflected.NsPerOp(), speedup) + if speedup < 1.5 { + t.Errorf("generated encoder is only %.2fx faster than arreflect, want at least 1.5x", speedup) + } +} + +func skipIfInstrumented(t *testing.T) { + t.Helper() + if instrumented { + t.Skip("allocation counts and timings are not meaningful under -race or -asan") + } +} diff --git a/arrgen/internal/gentypes/bench_test.go b/arrgen/internal/gentypes/bench_test.go new file mode 100644 index 000000000..94a0bbb49 --- /dev/null +++ b/arrgen/internal/gentypes/bench_test.go @@ -0,0 +1,221 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gentypes_test + +import ( + "fmt" + "testing" + "time" + + "github.com/apache/arrow-go/arrgen/example" + "github.com/apache/arrow-go/arrgen/internal/gentypes" + "github.com/apache/arrow-go/v18/arrow/array/arreflect" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// benchRows is the batch size the encoder benchmarks build. It is large enough +// that per-batch setup does not dominate and small enough to stay in cache. +const benchRows = 1024 + +func makeMetrics(n int) []example.Metric { + metrics := make([]example.Metric, n) + v := 42.5 + for i := range metrics { + metrics[i] = example.Metric{ + Time: base.Add(time.Duration(i) * time.Second), + Host: fmt.Sprintf("host-%d", i%16), + CPU: float64(i) / 100, + } + if i%4 != 0 { + metrics[i].Value = &v + } + } + return metrics +} + +func makeFixed(n int) []gentypes.Fixed { + rows := make([]gentypes.Fixed, n) + v := 1.5 + for i := range rows { + rows[i] = gentypes.Fixed{ + Timestamp: base.Add(time.Duration(i) * time.Second), + ID: int64(i), + Value: float64(i) * 1.5, + OK: i%2 == 0, + } + if i%3 != 0 { + rows[i].Optional = &v + } + } + return rows +} + +// BenchmarkMetricBatch is the headline comparison: the same slice of rows +// encoded into a record batch by the reflection path and by the generated one. +func BenchmarkMetricBatch(b *testing.B) { + metrics := makeMetrics(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := arreflect.RecordFromSlice(metrics, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) + + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := example.MetricRecordBatch(mem, metrics) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) +} + +// BenchmarkRowBatch runs the same comparison over the wide fixture, where a +// row costs 46 columns of reflection instead of four. +func BenchmarkRowBatch(b *testing.B) { + rows := makeRows(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) + + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := gentypes.RowRecordBatch(mem, rows) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) +} + +// BenchmarkFixedBatch isolates the encoder from Arrow's variable-width buffers: +// every column is fixed width, so what is left in the allocation column is the +// cost of the encoding strategy itself rather than of growing a data buffer. +func BenchmarkFixedBatch(b *testing.B) { + rows := makeFixed(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) + + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + rec, err := gentypes.FixedRecordBatch(mem, rows) + if err != nil { + b.Fatal(err) + } + rec.Release() + } + b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") + }) +} + +// BenchmarkStreamAppend measures the per-row cost of the streaming entry point, +// where a caller feeds rows in one at a time from a reused variable. The +// reflection path has no single-row API, so its column is the amortized cost of +// encoding a batch of the same rows. +func BenchmarkStreamAppend(b *testing.B) { + rows := makeFixed(benchRows) + mem := memory.DefaultAllocator + + b.Run("arreflect/batched", func(b *testing.B) { + b.ReportAllocs() + n := 0 + for b.Loop() { + rec, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + b.Fatal(err) + } + rec.Release() + n += benchRows + } + b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(n), "ns/row") + }) + + b.Run("arrgen/streamed", func(b *testing.B) { + a := gentypes.NewFixedAppender(mem) + defer a.Release() + b.ReportAllocs() + var row gentypes.Fixed + i := 0 + for b.Loop() { + row = rows[i%benchRows] + i++ + if i%benchRows == 0 { + a.NewRecordBatch().Release() + a.Reserve(benchRows) + } + a.Append(&row) + } + a.NewRecordBatch().Release() + b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(b.N), "ns/row") + }) +} + +// BenchmarkSchema shows the other cost code generation removes: arreflect walks +// the struct to infer a schema, while the generated schema is a package-level +// variable built once at init. +func BenchmarkSchema(b *testing.B) { + b.Run("arreflect", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + if _, err := arreflect.InferSchema[gentypes.Row](); err != nil { + b.Fatal(err) + } + } + }) + b.Run("arrgen", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + _ = gentypes.RowSchema() + } + }) +} diff --git a/arrgen/internal/gentypes/equivalence_test.go b/arrgen/internal/gentypes/equivalence_test.go new file mode 100644 index 000000000..b86a3ff8e --- /dev/null +++ b/arrgen/internal/gentypes/equivalence_test.go @@ -0,0 +1,470 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gentypes_test + +import ( + "fmt" + "math" + "testing" + "time" + + "github.com/apache/arrow-go/arrgen/example" + "github.com/apache/arrow-go/arrgen/internal/gentypes" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/array/arreflect" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/decimal256" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// arreflectStructFieldGaps names the Row columns whose Go type arreflect cannot +// encode when it appears as a struct field, and what it emits instead. +// +// Its inferArrowType dispatches on reflect.Kind before it checks for the types +// it special-cases, so any Go struct that Arrow models as a scalar - time.Time, +// decimal128.Num, decimal256.Num - reaches inferStructType, which finds only +// unexported fields and returns an empty struct<>. The value is then dropped on +// the floor. Its own tests only ever pass these types as the top-level element +// of a slice, where a different branch handles them correctly, which is why the +// gap has gone unnoticed. +// +// arrgen emits the column arreflect's own InferType and FromSlice paths agree +// on (TIMESTAMP(ns, UTC), DECIMAL128, DECIMAL256), so these columns are the one +// place the two paths disagree and are compared separately, by value, in +// TestGapColumnValues. TestArreflectStructFieldGaps pins the current arreflect +// behavior so that fixing it upstream turns into a failure here telling us to +// shorten this list. +// A tag that names an explicit Arrow type - date32, decimal(20,3) - rescues the +// field, because arreflect overrides the inferred type from the tag afterwards. +// Only the untagged spellings are listed here. +var arreflectStructFieldGaps = map[string]string{ + "ts": "time.Time", + "pts": "*time.Time", + "dec256": "decimal256.Num", + "pdec128": "*decimal128.Num", +} + +// compareColumns asserts that got and want hold the same columns, skipping the +// ones arreflect cannot encode as struct fields. +// +// dictsByValue relaxes dictionary columns to a comparison of the values they +// decode to. An appender reused across batches carries its memo table forward, +// so its second batch encodes the same values against a longer dictionary than +// a freshly built one - a different encoding of identical data, and the reason +// Arrow has delta dictionaries at all. +func compareColumns(t *testing.T, label string, got, want arrow.RecordBatch, dictsByValue bool) { + t.Helper() + if got.NumRows() != want.NumRows() { + t.Fatalf("%s: NumRows = %d, want %d", label, got.NumRows(), want.NumRows()) + } + if got.NumCols() != want.NumCols() { + t.Fatalf("%s: NumCols = %d, want %d", label, got.NumCols(), want.NumCols()) + } + for i := 0; i < int(got.NumCols()); i++ { + name := got.Schema().Field(i).Name + if _, gap := arreflectStructFieldGaps[name]; gap { + continue + } + gc, wc := got.Column(i), want.Column(i) + if dictsByValue && gc.DataType().ID() == arrow.DICTIONARY { + compareByValueStr(t, label, name, gc, wc) + continue + } + if !array.Equal(gc, wc) { + t.Errorf("%s: column %q differs\n got: %s\nwant: %s", label, name, gc, wc) + } + } +} + +// compareByValueStr compares two arrays by the value each row decodes to, +// ignoring how that value is physically encoded. +func compareByValueStr(t *testing.T, label, name string, got, want arrow.Array) { + t.Helper() + if got.Len() != want.Len() { + t.Errorf("%s: column %q: len = %d, want %d", label, name, got.Len(), want.Len()) + return + } + for i := 0; i < got.Len(); i++ { + if got.IsNull(i) != want.IsNull(i) { + t.Errorf("%s: column %q row %d: null = %t, want %t", label, name, i, got.IsNull(i), want.IsNull(i)) + return + } + if got.IsNull(i) { + continue + } + if got.ValueStr(i) != want.ValueStr(i) { + t.Errorf("%s: column %q row %d = %s, want %s", label, name, i, got.ValueStr(i), want.ValueStr(i)) + return + } + } +} + +// base is a fixed instant with a non-zero time of day and sub-second part, so +// the date32/date64/time32/time64 columns each exercise a different truncation. +var base = time.Date(2024, time.March, 17, 13, 45, 12, 123456789, time.UTC) + +// makeRows builds n deterministic rows. Every third row nulls its pointer +// fields and empties its variable-width ones, so the comparison covers null +// handling and not just the happy path. +func makeRows(n int) []gentypes.Row { + rows := make([]gentypes.Row, n) + for i := range rows { + null := i%3 == 0 + r := gentypes.Row{ + Bool: i%2 == 0, + Int8: int8(i), + Int16: int16(i * 3), + Int32: int32(i * 7), + Int64: int64(i) * 1e6, + Int: i, + Uint8: uint8(i), + Uint16: uint16(i * 5), + Uint32: uint32(i * 11), + Uint64: uint64(i) * 1e9, + Uint: uint(i), + Float32: float32(i) + 0.25, + Float64: float64(i) * math.Pi, + Str: fmt.Sprintf("host-%d", i%7), + Bin: []byte{byte(i), byte(i >> 8)}, + Timestamp: base.Add(time.Duration(i) * time.Second), + Date32: base.AddDate(0, 0, i), + Date64: base.AddDate(0, 0, i), + Time32: base.Add(time.Duration(i) * time.Minute), + Time64: base.Add(time.Duration(i) * time.Minute), + Duration: time.Duration(i) * time.Millisecond, + Dec32: decimal.Decimal32(i * 100), + Dec64: decimal.Decimal64(i * 10000), + Dec128: decimal128.FromU64(uint64(i) * 1000), + Dec256: decimal256.FromU64(uint64(i) * 2000), + LargeStr: fmt.Sprintf("large-%d", i), + ViewStr: fmt.Sprintf("view-%d", i), + LargeBin: []byte(fmt.Sprintf("lb-%d", i)), + ViewBin: []byte(fmt.Sprintf("vb-%d", i)), + DictStr: fmt.Sprintf("region-%d", i%3), + DictBin: []byte(fmt.Sprintf("k%d", i%4)), + DictInt: int32(i % 5), + DictF64: float64(i % 6), + Untagged: int64(i), + Secret: "never encoded", + } + if null { + // Leave every pointer nil, and null out the two shapes that map to + // null without a pointer: a nil []byte column. + r.Bin = nil + r.LargeBin = nil + r.ViewBin = nil + r.DictBin = nil + rows[i] = r + continue + } + b := i%2 == 1 + i64 := int64(i) * 3 + f64 := float64(i) / 3 + s := fmt.Sprintf("p-%d", i) + bin := []byte(fmt.Sprintf("pb-%d", i)) + ts := base.Add(time.Duration(i) * time.Hour) + dur := time.Duration(i) * time.Microsecond + dec := decimal.Decimal32(i) + ds := fmt.Sprintf("pd-%d", i%2) + r.PBool, r.PInt64, r.PF64, r.PStr, r.PBin = &b, &i64, &f64, &s, &bin + dec128 := decimal128.FromU64(uint64(i)) + r.PTS, r.PDate32, r.PTime64, r.PDur, r.PDec32, r.PDictS = &ts, &ts, &ts, &dur, &dec, &ds + r.PDec128 = &dec128 + rows[i] = r + } + return rows +} + +// TestRowSchemaMatchesArreflect is the cheaper half of the equivalence +// guarantee: the generated schema is what arreflect would infer, column for +// column, including names, types and nullability. +func TestRowSchemaMatchesArreflect(t *testing.T) { + want, err := arreflect.InferSchema[gentypes.Row]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + got := gentypes.RowSchema() + if got.NumFields() != want.NumFields() { + t.Fatalf("field count = %d, want %d\n got: %s\nwant: %s", got.NumFields(), want.NumFields(), got, want) + } + for i := 0; i < got.NumFields(); i++ { + gf, wf := got.Field(i), want.Field(i) + if gf.Name != wf.Name { + t.Errorf("field %d: name = %q, want %q", i, gf.Name, wf.Name) + continue + } + if _, gap := arreflectStructFieldGaps[gf.Name]; gap { + continue + } + if !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { + t.Errorf("field %q: got %s nullable=%t, want %s nullable=%t", gf.Name, gf.Type, gf.Nullable, wf.Type, wf.Nullable) + } + } +} + +func TestMetricSchemaMatchesArreflect(t *testing.T) { + want, err := arreflect.InferSchema[example.Metric]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + got := example.MetricSchema() + // "ts" is a bare time.Time: see arreflectStructFieldGaps. + for i := 0; i < got.NumFields(); i++ { + gf, wf := got.Field(i), want.Field(i) + if gf.Name == "ts" { + continue + } + if gf.Name != wf.Name || !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { + t.Errorf("field %d: got %s, want %s", i, gf, wf) + } + } +} + +// TestRowRecordMatchesArreflect is the guarantee that matters: for the same +// input, the generated encoder and the reflection encoder produce equal +// columns. Everything else in this package - the benchmarks, the allocation +// assertions - is only interesting because this holds. +// +// It compares column by column rather than with array.RecordEqual so that the +// handful of columns arreflect cannot encode as struct fields can be skipped by +// name instead of poisoning the whole comparison; see arreflectStructFieldGaps. +func TestRowRecordMatchesArreflect(t *testing.T) { + for _, n := range []int{0, 1, 2, 3, 64} { + t.Run(fmt.Sprintf("rows=%d", n), func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + rows := makeRows(n) + + want, err := arreflect.RecordFromSlice(rows, mem) + if err != nil { + t.Fatalf("arreflect.RecordFromSlice: %v", err) + } + defer want.Release() + + got, err := gentypes.RowRecordBatch(mem, rows) + if err != nil { + t.Fatalf("RowRecordBatch: %v", err) + } + defer got.Release() + + compareColumns(t, fmt.Sprintf("rows=%d", n), got, want, false) + }) + } +} + +// TestArreflectStructFieldGaps pins the upstream behavior that +// arreflectStructFieldGaps exists to work around. If arreflect learns to encode +// these struct fields, this test fails and the skip list should shrink to match +// - a failure here is good news, not a regression. +func TestArreflectStructFieldGaps(t *testing.T) { + schema, err := arreflect.InferSchema[gentypes.Row]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + for name, goType := range arreflectStructFieldGaps { + idx := schema.FieldIndices(name) + if len(idx) != 1 { + t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) + } + if dt := schema.Field(idx[0]).Type; dt.ID() != arrow.STRUCT { + t.Errorf("arreflect now infers %s column %q as %s rather than an empty struct; "+ + "remove it from arreflectStructFieldGaps so the strict comparison covers it", goType, name, dt) + } + } +} + +// TestGapColumnValues covers the columns TestRowRecordMatchesArreflect has to +// skip. There is no reflection-built batch to compare them against, so the +// expected Arrow values are spelled out here directly. +func TestGapColumnValues(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + ts := base + dec128 := decimal128.FromU64(12345) + dec256 := decimal256.FromU64(67890) + rows := []gentypes.Row{ + {Timestamp: ts, Dec128: dec128, Dec256: dec256, PTS: &ts, PDec128: &dec128}, + {Timestamp: ts.Add(time.Second), Dec128: dec128, Dec256: dec256}, // pointers nil + } + + rec, err := gentypes.RowRecordBatch(mem, rows) + if err != nil { + t.Fatalf("RowRecordBatch: %v", err) + } + defer rec.Release() + + col := func(name string) arrow.Array { + idx := rec.Schema().FieldIndices(name) + if len(idx) != 1 { + t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) + } + return rec.Column(idx[0]) + } + + tsCol := col("ts").(*array.Timestamp) + if got, want := tsCol.Value(0), arrow.Timestamp(ts.UnixNano()); got != want { + t.Errorf("ts[0] = %d, want %d", got, want) + } + if got, want := tsCol.Value(1), arrow.Timestamp(ts.Add(time.Second).UnixNano()); got != want { + t.Errorf("ts[1] = %d, want %d", got, want) + } + if dt := tsCol.DataType().(*arrow.TimestampType); dt.Unit != arrow.Nanosecond || dt.TimeZone != "UTC" { + t.Errorf("ts type = %s, want timestamp[ns, tz=UTC]", dt) + } + + ptsCol := col("pts").(*array.Timestamp) + if ptsCol.IsNull(0) { + t.Error("pts[0] is null, want a value") + } + if !ptsCol.IsNull(1) { + t.Error("pts[1] is not null, want null") + } + + if got := col("dec128").(*array.Decimal128).Value(0); got != dec128 { + t.Errorf("dec128[0] = %v, want %v", got, dec128) + } + if got := col("dec256").(*array.Decimal256).Value(0); got != dec256 { + t.Errorf("dec256[0] = %v, want %v", got, dec256) + } + if !col("pdec128").(*array.Decimal128).IsNull(1) { + t.Error("pdec128[1] is not null, want null") + } +} + +func TestMetricRecordMatchesArreflect(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + v := 3.5 + metrics := []example.Metric{ + {Time: base, Host: "a", CPU: 0.5, Value: &v, Secret: "hidden"}, + {Time: base.Add(time.Second), Host: "b", CPU: 1.5, Value: nil}, + } + + want, err := arreflect.RecordFromSlice(metrics, mem) + if err != nil { + t.Fatalf("arreflect.RecordFromSlice: %v", err) + } + defer want.Release() + + got, err := example.MetricRecordBatch(mem, metrics) + if err != nil { + t.Fatalf("MetricRecordBatch: %v", err) + } + defer got.Release() + + // "ts" is a bare time.Time, one of the columns arreflect drops as a struct + // field; the rest must match it exactly. See arreflectStructFieldGaps. + for i := 0; i < int(got.NumCols()); i++ { + name := got.Schema().Field(i).Name + if name == "ts" { + continue + } + if !array.Equal(got.Column(i), want.Column(i)) { + t.Errorf("column %q differs\n got: %s\nwant: %s", name, got.Column(i), want.Column(i)) + } + } + if got, want := got.Column(0).(*array.Timestamp).Value(0), arrow.Timestamp(base.UnixNano()); got != want { + t.Errorf("ts[0] = %d, want %d", got, want) + } + if secrets := got.Schema().FieldIndices("Secret"); len(secrets) != 0 { + t.Errorf(`a field tagged arrow:"-" was encoded as column %v`, secrets) + } +} + +// TestAppendMatchesAppendSlice checks the streaming entry point against the +// bulk one: a caller feeding rows one at a time must land in the same place as +// a caller handing over a slice. +func TestAppendMatchesAppendSlice(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + rows := makeRows(16) + + bulk := gentypes.NewRowAppender(mem) + defer bulk.Release() + bulk.AppendSlice(rows) + bulkRec := bulk.NewRecordBatch() + defer bulkRec.Release() + + streamed := gentypes.NewRowAppender(mem) + defer streamed.Release() + streamed.Reserve(len(rows)) + var row gentypes.Row + for i := range rows { + row = rows[i] // a caller reusing one row variable must be safe + streamed.Append(&row) + } + if got, want := streamed.Len(), len(rows); got != want { + t.Errorf("Len() = %d, want %d", got, want) + } + streamedRec := streamed.NewRecordBatch() + defer streamedRec.Release() + + if !array.RecordEqual(streamedRec, bulkRec) { + t.Errorf("streamed batch differs from bulk batch\n got: %s\nwant: %s", streamedRec, bulkRec) + } + if err := streamed.Err(); err != nil { + t.Errorf("Err() = %v, want nil", err) + } +} + +// TestAppenderReusableAcrossBatches covers the roll boundary: NewRecordBatch +// hands over the rows so far and leaves the appender ready for the next batch. +func TestAppenderReusableAcrossBatches(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + rows := makeRows(9) + a := gentypes.NewRowAppender(mem) + defer a.Release() + + for i := 0; i < len(rows); i += 3 { + a.AppendSlice(rows[i : i+3]) + rec := a.NewRecordBatch() + if got := rec.NumRows(); got != 3 { + t.Errorf("batch %d: NumRows() = %d, want 3", i/3, got) + } + want, err := arreflect.RecordFromSlice(rows[i:i+3], mem) + if err != nil { + t.Fatalf("arreflect.RecordFromSlice: %v", err) + } + compareColumns(t, fmt.Sprintf("batch %d", i/3), rec, want, true) + want.Release() + rec.Release() + if got := a.Len(); got != 0 { + t.Errorf("batch %d: Len() after NewRecordBatch = %d, want 0", i/3, got) + } + } +} + +// TestNilAllocatorUsesDefault documents the constructor's contract rather than +// leaving callers to discover it by crashing. +func TestNilAllocatorUsesDefault(t *testing.T) { + a := gentypes.NewRowAppender(nil) + defer a.Release() + a.AppendSlice(makeRows(2)) + rec := a.NewRecordBatch() + defer rec.Release() + if got := rec.NumRows(); got != 2 { + t.Errorf("NumRows() = %d, want 2", got) + } +} diff --git a/arrgen/internal/gentypes/instrumented_race_test.go b/arrgen/internal/gentypes/instrumented_race_test.go new file mode 100644 index 000000000..dbdd4fe22 --- /dev/null +++ b/arrgen/internal/gentypes/instrumented_race_test.go @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build race || asan + +package gentypes_test + +const instrumented = true diff --git a/arrgen/internal/gentypes/instrumented_test.go b/arrgen/internal/gentypes/instrumented_test.go new file mode 100644 index 000000000..cff9316c5 --- /dev/null +++ b/arrgen/internal/gentypes/instrumented_test.go @@ -0,0 +1,24 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build !race && !asan + +package gentypes_test + +// instrumented reports whether the binary was built with a sanitizer that +// perturbs allocation counts and timings. The allocation and speed assertions +// step aside when it is set: they would be measuring the instrumentation. +const instrumented = false diff --git a/arrgen/internal/gentypes/row_arrow.go b/arrgen/internal/gentypes/row_arrow.go new file mode 100644 index 000000000..0a0483f2f --- /dev/null +++ b/arrgen/internal/gentypes/row_arrow.go @@ -0,0 +1,491 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by arrgen. DO NOT EDIT. + +package gentypes + +import ( + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// rowArrowSchema is the Arrow schema of Row, resolved from its arrow +// struct tags at generate time. +var rowArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "b", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, + {Name: "i8", Type: arrow.PrimitiveTypes.Int8, Nullable: false}, + {Name: "i16", Type: arrow.PrimitiveTypes.Int16, Nullable: false}, + {Name: "i32", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "i64", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "i", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "u8", Type: arrow.PrimitiveTypes.Uint8, Nullable: false}, + {Name: "u16", Type: arrow.PrimitiveTypes.Uint16, Nullable: false}, + {Name: "u32", Type: arrow.PrimitiveTypes.Uint32, Nullable: false}, + {Name: "u64", Type: arrow.PrimitiveTypes.Uint64, Nullable: false}, + {Name: "u", Type: arrow.PrimitiveTypes.Uint64, Nullable: false}, + {Name: "f32", Type: arrow.PrimitiveTypes.Float32, Nullable: false}, + {Name: "f64", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "s", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: "bin", Type: arrow.BinaryTypes.Binary, Nullable: false}, + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "d32", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, + {Name: "d64", Type: arrow.FixedWidthTypes.Date64, Nullable: false}, + {Name: "t32", Type: &arrow.Time32Type{Unit: arrow.Millisecond}, Nullable: false}, + {Name: "t64", Type: &arrow.Time64Type{Unit: arrow.Nanosecond}, Nullable: false}, + {Name: "dur", Type: &arrow.DurationType{Unit: arrow.Nanosecond}, Nullable: false}, + {Name: "dec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: false}, + {Name: "dec64", Type: &arrow.Decimal64Type{Precision: 18, Scale: 4}, Nullable: false}, + {Name: "dec128", Type: &arrow.Decimal128Type{Precision: 20, Scale: 3}, Nullable: false}, + {Name: "dec256", Type: &arrow.Decimal256Type{Precision: 76, Scale: 0}, Nullable: false}, + {Name: "ls", Type: arrow.BinaryTypes.LargeString, Nullable: false}, + {Name: "vs", Type: arrow.BinaryTypes.StringView, Nullable: false}, + {Name: "lb", Type: arrow.BinaryTypes.LargeBinary, Nullable: false}, + {Name: "vb", Type: arrow.BinaryTypes.BinaryView, Nullable: false}, + {Name: "ds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: false}, + {Name: "db", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.Binary}, Nullable: false}, + {Name: "di", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Int32}, Nullable: false}, + {Name: "df", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Float64}, Nullable: false}, + {Name: "Untagged", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "pb", Type: arrow.FixedWidthTypes.Boolean, Nullable: true}, + {Name: "pi64", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "pf64", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + {Name: "ps", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "pbin", Type: arrow.BinaryTypes.Binary, Nullable: true}, + {Name: "pts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: true}, + {Name: "pd32", Type: arrow.FixedWidthTypes.Date32, Nullable: true}, + {Name: "pt64", Type: &arrow.Time64Type{Unit: arrow.Nanosecond}, Nullable: true}, + {Name: "pdur", Type: &arrow.DurationType{Unit: arrow.Nanosecond}, Nullable: true}, + {Name: "pdec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: true}, + {Name: "pdec128", Type: &arrow.Decimal128Type{Precision: 38, Scale: 0}, Nullable: true}, + {Name: "pds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: true}, +}, nil) + +// RowSchema returns the Arrow schema encoded by RowAppender. It is +// equal to the schema arreflect.InferSchema[Row] infers at runtime. +func RowSchema() *arrow.Schema { return rowArrowSchema } + +// RowAppender converts Row values into Arrow record batches +// without reflection: the struct tags were read when this file was generated, +// so appending a row is a handful of direct field loads into typed builders. +// +// It is not safe for concurrent use. Release it when done. +type RowAppender struct { + rb *array.RecordBuilder + b0 *array.BooleanBuilder + b1 *array.Int8Builder + b2 *array.Int16Builder + b3 *array.Int32Builder + b4 *array.Int64Builder + b5 *array.Int64Builder + b6 *array.Uint8Builder + b7 *array.Uint16Builder + b8 *array.Uint32Builder + b9 *array.Uint64Builder + b10 *array.Uint64Builder + b11 *array.Float32Builder + b12 *array.Float64Builder + b13 *array.StringBuilder + b14 *array.BinaryBuilder + b15 *array.TimestampBuilder + b16 *array.Date32Builder + b17 *array.Date64Builder + b18 *array.Time32Builder + b19 *array.Time64Builder + b20 *array.DurationBuilder + b21 *array.Decimal32Builder + b22 *array.Decimal64Builder + b23 *array.Decimal128Builder + b24 *array.Decimal256Builder + b25 *array.LargeStringBuilder + b26 *array.StringViewBuilder + b27 *array.BinaryBuilder + b28 *array.BinaryViewBuilder + b29 *array.BinaryDictionaryBuilder + b30 *array.BinaryDictionaryBuilder + b31 *array.Int32DictionaryBuilder + b32 *array.Float64DictionaryBuilder + b33 *array.Int64Builder + b34 *array.BooleanBuilder + b35 *array.Int64Builder + b36 *array.Float64Builder + b37 *array.StringBuilder + b38 *array.BinaryBuilder + b39 *array.TimestampBuilder + b40 *array.Date32Builder + b41 *array.Time64Builder + b42 *array.DurationBuilder + b43 *array.Decimal32Builder + b44 *array.Decimal128Builder + b45 *array.BinaryDictionaryBuilder + err error +} + +// NewRowAppender returns an appender that builds Row batches with mem, +// or memory.DefaultAllocator when mem is nil. +// +// The typed field builders are resolved once here rather than per row, so +// Append does no type assertions at all. +func NewRowAppender(mem memory.Allocator) *RowAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, rowArrowSchema) + return &RowAppender{ + rb: rb, + b0: rb.Field(0).(*array.BooleanBuilder), + b1: rb.Field(1).(*array.Int8Builder), + b2: rb.Field(2).(*array.Int16Builder), + b3: rb.Field(3).(*array.Int32Builder), + b4: rb.Field(4).(*array.Int64Builder), + b5: rb.Field(5).(*array.Int64Builder), + b6: rb.Field(6).(*array.Uint8Builder), + b7: rb.Field(7).(*array.Uint16Builder), + b8: rb.Field(8).(*array.Uint32Builder), + b9: rb.Field(9).(*array.Uint64Builder), + b10: rb.Field(10).(*array.Uint64Builder), + b11: rb.Field(11).(*array.Float32Builder), + b12: rb.Field(12).(*array.Float64Builder), + b13: rb.Field(13).(*array.StringBuilder), + b14: rb.Field(14).(*array.BinaryBuilder), + b15: rb.Field(15).(*array.TimestampBuilder), + b16: rb.Field(16).(*array.Date32Builder), + b17: rb.Field(17).(*array.Date64Builder), + b18: rb.Field(18).(*array.Time32Builder), + b19: rb.Field(19).(*array.Time64Builder), + b20: rb.Field(20).(*array.DurationBuilder), + b21: rb.Field(21).(*array.Decimal32Builder), + b22: rb.Field(22).(*array.Decimal64Builder), + b23: rb.Field(23).(*array.Decimal128Builder), + b24: rb.Field(24).(*array.Decimal256Builder), + b25: rb.Field(25).(*array.LargeStringBuilder), + b26: rb.Field(26).(*array.StringViewBuilder), + b27: rb.Field(27).(*array.BinaryBuilder), + b28: rb.Field(28).(*array.BinaryViewBuilder), + b29: rb.Field(29).(*array.BinaryDictionaryBuilder), + b30: rb.Field(30).(*array.BinaryDictionaryBuilder), + b31: rb.Field(31).(*array.Int32DictionaryBuilder), + b32: rb.Field(32).(*array.Float64DictionaryBuilder), + b33: rb.Field(33).(*array.Int64Builder), + b34: rb.Field(34).(*array.BooleanBuilder), + b35: rb.Field(35).(*array.Int64Builder), + b36: rb.Field(36).(*array.Float64Builder), + b37: rb.Field(37).(*array.StringBuilder), + b38: rb.Field(38).(*array.BinaryBuilder), + b39: rb.Field(39).(*array.TimestampBuilder), + b40: rb.Field(40).(*array.Date32Builder), + b41: rb.Field(41).(*array.Time64Builder), + b42: rb.Field(42).(*array.DurationBuilder), + b43: rb.Field(43).(*array.Decimal32Builder), + b44: rb.Field(44).(*array.Decimal128Builder), + b45: rb.Field(45).(*array.BinaryDictionaryBuilder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *RowAppender) Schema() *arrow.Schema { return rowArrowSchema } + +// RecordBuilder returns the underlying builder, for callers that need an Arrow +// API this appender does not wrap. Appending through it directly is allowed as +// long as every column stays the same length. +func (a *RowAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends is what keeps the append path free of allocations. +func (a *RowAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *RowAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller streaming rows can reuse a single Row variable across calls. +func (a *RowAppender) Append(v *Row) { + a.b0.Append(v.Bool) + a.b1.Append(v.Int8) + a.b2.Append(v.Int16) + a.b3.Append(v.Int32) + a.b4.Append(v.Int64) + a.b5.Append(int64(v.Int)) + a.b6.Append(v.Uint8) + a.b7.Append(v.Uint16) + a.b8.Append(v.Uint32) + a.b9.Append(v.Uint64) + a.b10.Append(uint64(v.Uint)) + a.b11.Append(v.Float32) + a.b12.Append(v.Float64) + a.b13.Append(v.Str) + if v.Bin == nil { + a.b14.AppendNull() + } else { + a.b14.Append(v.Bin) + } + a.b15.Append(arrow.Timestamp(v.Timestamp.UnixNano())) + a.b16.Append(arrow.Date32FromTime(v.Date32)) + a.b17.Append(arrow.Date64FromTime(v.Date64)) + { + tod := v.Time32.UTC() + a.b18.Append(arrow.Time32(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds() / 1e6)) + } + { + tod := v.Time64.UTC() + a.b19.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) + } + a.b20.Append(arrow.Duration(v.Duration.Nanoseconds())) + a.b21.Append(v.Dec32) + a.b22.Append(v.Dec64) + a.b23.Append(v.Dec128) + a.b24.Append(v.Dec256) + a.b25.Append(v.LargeStr) + a.b26.Append(v.ViewStr) + if v.LargeBin == nil { + a.b27.AppendNull() + } else { + a.b27.Append(v.LargeBin) + } + if v.ViewBin == nil { + a.b28.AppendNull() + } else { + a.b28.Append(v.ViewBin) + } + a.setErr(a.b29.AppendString(v.DictStr)) + if v.DictBin == nil { + a.b30.AppendNull() + } else { + a.setErr(a.b30.Append(v.DictBin)) + } + a.setErr(a.b31.Append(v.DictInt)) + a.setErr(a.b32.Append(v.DictF64)) + a.b33.Append(v.Untagged) + if v.PBool == nil { + a.b34.AppendNull() + } else { + a.b34.Append(*v.PBool) + } + if v.PInt64 == nil { + a.b35.AppendNull() + } else { + a.b35.Append(*v.PInt64) + } + if v.PF64 == nil { + a.b36.AppendNull() + } else { + a.b36.Append(*v.PF64) + } + if v.PStr == nil { + a.b37.AppendNull() + } else { + a.b37.Append(*v.PStr) + } + if v.PBin == nil || *v.PBin == nil { + a.b38.AppendNull() + } else { + a.b38.Append(*v.PBin) + } + if v.PTS == nil { + a.b39.AppendNull() + } else { + a.b39.Append(arrow.Timestamp((*v.PTS).UnixNano())) + } + if v.PDate32 == nil { + a.b40.AppendNull() + } else { + a.b40.Append(arrow.Date32FromTime(*v.PDate32)) + } + if v.PTime64 == nil { + a.b41.AppendNull() + } else { + tod := (*v.PTime64).UTC() + a.b41.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) + } + if v.PDur == nil { + a.b42.AppendNull() + } else { + a.b42.Append(arrow.Duration((*v.PDur).Nanoseconds())) + } + if v.PDec32 == nil { + a.b43.AppendNull() + } else { + a.b43.Append(*v.PDec32) + } + if v.PDec128 == nil { + a.b44.AppendNull() + } else { + a.b44.Append(*v.PDec128) + } + if v.PDictS == nil { + a.b45.AppendNull() + } else { + a.setErr(a.b45.AppendString(*v.PDictS)) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *RowAppender) AppendSlice(vs []Row) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch snapshots the appended rows into a record batch and resets the +// appender for the next batch. The caller owns the batch and must Release it. +func (a *RowAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail, so an appender over a schema without one always returns nil. +// +// The error is sticky: NewRecordBatch does not clear it, because a batch built +// after a failed append is missing a value in that column. +func (a *RowAppender) Err() error { return a.err } + +// setErr keeps the first error so the hot path stays a single branch. +func (a *RowAppender) setErr(err error) { + if err != nil && a.err == nil { + a.err = err + } +} + +// Release releases the appender's builders. +func (a *RowAppender) Release() { a.rb.Release() } + +// RowRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Row](vs, mem) and +// produces an identical batch. The caller owns the batch and must Release it. +func RowRecordBatch(mem memory.Allocator, vs []Row) (arrow.RecordBatch, error) { + a := NewRowAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} + +// fixedArrowSchema is the Arrow schema of Fixed, resolved from its arrow +// struct tags at generate time. +var fixedArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "val", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "ok", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, + {Name: "opt", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, +}, nil) + +// FixedSchema returns the Arrow schema encoded by FixedAppender. It is +// equal to the schema arreflect.InferSchema[Fixed] infers at runtime. +func FixedSchema() *arrow.Schema { return fixedArrowSchema } + +// FixedAppender converts Fixed values into Arrow record batches +// without reflection: the struct tags were read when this file was generated, +// so appending a row is a handful of direct field loads into typed builders. +// +// It is not safe for concurrent use. Release it when done. +type FixedAppender struct { + rb *array.RecordBuilder + b0 *array.TimestampBuilder + b1 *array.Int64Builder + b2 *array.Float64Builder + b3 *array.BooleanBuilder + b4 *array.Float64Builder + err error +} + +// NewFixedAppender returns an appender that builds Fixed batches with mem, +// or memory.DefaultAllocator when mem is nil. +// +// The typed field builders are resolved once here rather than per row, so +// Append does no type assertions at all. +func NewFixedAppender(mem memory.Allocator) *FixedAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, fixedArrowSchema) + return &FixedAppender{ + rb: rb, + b0: rb.Field(0).(*array.TimestampBuilder), + b1: rb.Field(1).(*array.Int64Builder), + b2: rb.Field(2).(*array.Float64Builder), + b3: rb.Field(3).(*array.BooleanBuilder), + b4: rb.Field(4).(*array.Float64Builder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *FixedAppender) Schema() *arrow.Schema { return fixedArrowSchema } + +// RecordBuilder returns the underlying builder, for callers that need an Arrow +// API this appender does not wrap. Appending through it directly is allowed as +// long as every column stays the same length. +func (a *FixedAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends is what keeps the append path free of allocations. +func (a *FixedAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *FixedAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller streaming rows can reuse a single Fixed variable across calls. +func (a *FixedAppender) Append(v *Fixed) { + a.b0.Append(arrow.Timestamp(v.Timestamp.UnixNano())) + a.b1.Append(v.ID) + a.b2.Append(v.Value) + a.b3.Append(v.OK) + if v.Optional == nil { + a.b4.AppendNull() + } else { + a.b4.Append(*v.Optional) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *FixedAppender) AppendSlice(vs []Fixed) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch snapshots the appended rows into a record batch and resets the +// appender for the next batch. The caller owns the batch and must Release it. +func (a *FixedAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail, so an appender over a schema without one always returns nil. +// +// The error is sticky: NewRecordBatch does not clear it, because a batch built +// after a failed append is missing a value in that column. +func (a *FixedAppender) Err() error { return a.err } + +// Release releases the appender's builders. +func (a *FixedAppender) Release() { a.rb.Release() } + +// FixedRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Fixed](vs, mem) and +// produces an identical batch. The caller owns the batch and must Release it. +func FixedRecordBatch(mem memory.Allocator, vs []Fixed) (arrow.RecordBatch, error) { + a := NewFixedAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} diff --git a/arrgen/internal/gentypes/types.go b/arrgen/internal/gentypes/types.go new file mode 100644 index 000000000..f66f34fbe --- /dev/null +++ b/arrgen/internal/gentypes/types.go @@ -0,0 +1,103 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package gentypes holds the fixtures the arrgen test suite generates from. +// +// Row is deliberately exhaustive: it carries one field for every column shape +// arrgen claims to support, so the equivalence tests in this package compare +// the generated encoder against arreflect over the whole supported surface +// rather than over a convenient subset. +package gentypes + +import ( + "time" + + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/decimal256" +) + +// Row covers every supported field type, tag option and nullable variant. +type Row struct { + Bool bool `arrow:"b"` + Int8 int8 `arrow:"i8"` + Int16 int16 `arrow:"i16"` + Int32 int32 `arrow:"i32"` + Int64 int64 `arrow:"i64"` + Int int `arrow:"i"` + Uint8 uint8 `arrow:"u8"` + Uint16 uint16 `arrow:"u16"` + Uint32 uint32 `arrow:"u32"` + Uint64 uint64 `arrow:"u64"` + Uint uint `arrow:"u"` + Float32 float32 `arrow:"f32"` + Float64 float64 `arrow:"f64"` + Str string `arrow:"s"` + Bin []byte `arrow:"bin"` + + Timestamp time.Time `arrow:"ts"` + Date32 time.Time `arrow:"d32,date32"` + Date64 time.Time `arrow:"d64,date64"` + Time32 time.Time `arrow:"t32,time32"` + Time64 time.Time `arrow:"t64,time64"` + Duration time.Duration `arrow:"dur"` + + Dec32 decimal.Decimal32 `arrow:"dec32"` + Dec64 decimal.Decimal64 `arrow:"dec64,decimal(18,4)"` + Dec128 decimal128.Num `arrow:"dec128,decimal(20,3)"` + Dec256 decimal256.Num `arrow:"dec256"` + + LargeStr string `arrow:"ls,large"` + ViewStr string `arrow:"vs,view"` + LargeBin []byte `arrow:"lb,large"` + ViewBin []byte `arrow:"vb,view"` + + DictStr string `arrow:"ds,dict"` + DictBin []byte `arrow:"db,dict"` + DictInt int32 `arrow:"di,dict"` + DictF64 float64 `arrow:"df,dict"` + Untagged int64 // no tag: the column takes the Go field name + + PBool *bool `arrow:"pb"` + PInt64 *int64 `arrow:"pi64"` + PF64 *float64 `arrow:"pf64"` + PStr *string `arrow:"ps"` + PBin *[]byte `arrow:"pbin"` + PTS *time.Time `arrow:"pts"` + PDate32 *time.Time `arrow:"pd32,date32"` + PTime64 *time.Time `arrow:"pt64,time64"` + PDur *time.Duration `arrow:"pdur"` + PDec32 *decimal.Decimal32 `arrow:"pdec32"` + PDec128 *decimal128.Num `arrow:"pdec128"` + PDictS *string `arrow:"pds,dict"` + + Secret string `arrow:"-"` // excluded from Arrow entirely +} + +// Fixed holds only fixed-width columns, so appending a row never has to grow a +// variable-length data buffer. It is what the zero-allocation assertion in +// alloc_test.go measures: with space reserved up front, Append does no +// allocating at all, which is not something a string or []byte column can +// promise once its data buffer needs to double. +type Fixed struct { + Timestamp time.Time `arrow:"ts"` + ID int64 `arrow:"id"` + Value float64 `arrow:"val"` + OK bool `arrow:"ok"` + Optional *float64 `arrow:"opt"` +} + +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Row,Fixed -header ../../license_header.txt -output row_arrow.go diff --git a/arrgen/license_header.txt b/arrgen/license_header.txt new file mode 100644 index 000000000..f1fc31831 --- /dev/null +++ b/arrgen/license_header.txt @@ -0,0 +1,15 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. diff --git a/arrgen/mapping.go b/arrgen/mapping.go new file mode 100644 index 000000000..fc8c9a475 --- /dev/null +++ b/arrgen/mapping.go @@ -0,0 +1,330 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import ( + "fmt" + "go/types" +) + +// colSpec is everything the template needs to emit one Arrow column: the type +// expression for the schema, the concrete builder the appender caches, and a +// renderer for the append statement. +// +// The three are resolved together because they are three views of the same +// decision - Int64 columns get an *array.Int64Builder and an Append(int64(...)) +// - and splitting them across independent switches is how they drift apart. +type colSpec struct { + arrowType string // Go expression for the arrow.DataType + builderType string // Go type of the cached typed builder + fallible bool // the builder's Append returns an error (dictionaries) + needsTime bool // the append statement references the time package + nilable bool // a nil Go value maps to null even in a non-nullable column + + // appendStmt renders the append of a known-non-nil value. bld is the cached + // builder expression ("a.b3"), val is the value ("v.F" or "*v.F"), and recv + // is val parenthesized where needed so a method call binds to the value and + // not to the pointer. + appendStmt func(bld, recv, val string) string +} + +// optSupport records which tag options a column kind can honor. Options that +// would be silently ignored are rejected instead: arreflect drops a date32 tag +// on an int field on the floor, but at generate time there is a human waiting +// to be told about the typo. +type optSupport struct { + temporal bool + decimal bool + largeView bool + dict bool +} + +func checkOpts(o tagOpts, goType string, s optSupport) error { + switch { + case o.Temporal != "" && !s.temporal: + return fmt.Errorf("%q is only valid on a time.Time field, not %s", o.Temporal, goType) + case o.HasDecimalOpts && !s.decimal: + return fmt.Errorf("decimal(precision,scale) is only valid on a decimal field, not %s", goType) + case o.Large && !s.largeView: + return fmt.Errorf("large has no effect on %s; it is only valid on string and []byte fields", goType) + case o.View && !s.largeView: + return fmt.Errorf("view has no effect on %s; it is only valid on string and []byte fields", goType) + case o.Dict && !s.dict: + return fmt.Errorf("dict is not supported on %s; it is valid on string, []byte, integer and float fields", goType) + case o.Dict && o.Large: + return fmt.Errorf("dict cannot be combined with large: Dictionary is not implemented by arrow-go") + } + return nil +} + +// resolveColumn maps a Go field type and its parsed tag to an Arrow column, +// reporting nullable=true for a pointer field (matching arreflect, which marks +// a column nullable from the outermost pointer). +func resolveColumn(t types.Type, opts tagOpts) (spec colSpec, nullable bool, err error) { + t = types.Unalias(t) + if p, ok := t.(*types.Pointer); ok { + nullable = true + t = types.Unalias(p.Elem()) + if _, ok := t.(*types.Pointer); ok { + return colSpec{}, false, fmt.Errorf("multi-level pointer type %s is not supported", t) + } + } + spec, err = baseSpec(t, opts) + return spec, nullable, err +} + +func baseSpec(t types.Type, opts tagOpts) (colSpec, error) { + if named, ok := t.(*types.Named); ok { + return namedSpec(named, opts) + } + if sl, ok := t.(*types.Slice); ok { + if b, ok := types.Unalias(sl.Elem()).(*types.Basic); ok && b.Kind() == types.Uint8 { + return byteSliceSpec(opts) + } + return colSpec{}, fmt.Errorf("list column for %s is not supported; encode it with arreflect", t) + } + if b, ok := t.(*types.Basic); ok { + return basicSpec(b, opts) + } + return colSpec{}, fmt.Errorf("type %s is not supported", t) +} + +// namedSpec handles the handful of named types arreflect recognizes by identity. +// Every other named type - including a defined scalar such as `type ID int64` - +// is unsupported there too, so rejecting it here keeps the two paths aligned. +func namedSpec(named *types.Named, opts tagOpts) (colSpec, error) { + obj := named.Obj() + pkg := "" + if obj.Pkg() != nil { + pkg = obj.Pkg().Path() + } + switch { + case pkg == "time" && obj.Name() == "Time": + return timeSpec(opts) + case pkg == "time" && obj.Name() == "Duration": + if err := checkOpts(opts, "time.Duration", optSupport{}); err != nil { + return colSpec{}, err + } + return colSpec{ + arrowType: "&arrow.DurationType{Unit: arrow.Nanosecond}", + builderType: "*array.DurationBuilder", + appendStmt: func(bld, recv, _ string) string { + return fmt.Sprintf("%s.Append(arrow.Duration(%s.Nanoseconds()))", bld, recv) + }, + }, nil + case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal32": + return decimalSpec(opts, "Decimal32", 9) + case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal64": + return decimalSpec(opts, "Decimal64", 18) + case isDecimalPkg(pkg, "decimal128") && obj.Name() == "Num": + return decimalSpec(opts, "Decimal128", 38) + case isDecimalPkg(pkg, "decimal256") && obj.Name() == "Num": + return decimalSpec(opts, "Decimal256", 76) + } + if _, ok := named.Underlying().(*types.Struct); ok { + return colSpec{}, fmt.Errorf("nested struct column for %s is not supported; encode it with arreflect", named) + } + return colSpec{}, fmt.Errorf("named type %s is not supported; arreflect matches scalar fields by exact type, so a defined type such as this one has no Arrow mapping in either path", named) +} + +func isDecimalPkg(path, name string) bool { + const prefix = "github.com/apache/arrow-go/v18/arrow/" + return path == prefix+name +} + +func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32) (colSpec, error) { + goType := "decimal." + arrowName + if err := checkOpts(opts, goType, optSupport{decimal: true}); err != nil { + return colSpec{}, err + } + precision, scale := defaultPrecision, int32(0) + if opts.HasDecimalOpts { + precision, scale = opts.DecimalPrecision, opts.DecimalScale + } + return colSpec{ + arrowType: fmt.Sprintf("&arrow.%sType{Precision: %d, Scale: %d}", arrowName, precision, scale), + builderType: "*array." + arrowName + "Builder", + appendStmt: simpleAppend(), + }, nil +} + +// timeSpec maps a time.Time field. Note that an untagged time.Time is a +// TIMESTAMP(ns, UTC) column here, which is what arreflect.InferType[time.Time] +// and its top-level FromSlice path both produce - but not what its struct-field +// path produces today. See the divergence note in the package README. +func timeSpec(opts tagOpts) (colSpec, error) { + if err := checkOpts(opts, "time.Time", optSupport{temporal: true}); err != nil { + return colSpec{}, err + } + switch opts.Temporal { + case "date32": + return colSpec{ + arrowType: "arrow.FixedWidthTypes.Date32", + builderType: "*array.Date32Builder", + appendStmt: func(bld, _, val string) string { + return fmt.Sprintf("%s.Append(arrow.Date32FromTime(%s))", bld, val) + }, + }, nil + case "date64": + return colSpec{ + arrowType: "arrow.FixedWidthTypes.Date64", + builderType: "*array.Date64Builder", + appendStmt: func(bld, _, val string) string { + return fmt.Sprintf("%s.Append(arrow.Date64FromTime(%s))", bld, val) + }, + }, nil + case "time32": + return timeOfDaySpec("&arrow.Time32Type{Unit: arrow.Millisecond}", "*array.Time32Builder", "arrow.Time32", " / 1e6"), nil + case "time64": + return timeOfDaySpec("&arrow.Time64Type{Unit: arrow.Nanosecond}", "*array.Time64Builder", "arrow.Time64", ""), nil + default: + return colSpec{ + arrowType: `&arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}`, + builderType: "*array.TimestampBuilder", + appendStmt: func(bld, recv, _ string) string { + return fmt.Sprintf("%s.Append(arrow.Timestamp(%s.UnixNano()))", bld, recv) + }, + }, nil + } +} + +// timeOfDaySpec renders arreflect's timeOfDayNanos inline. It is emitted as a +// block rather than a package-level helper so that two generated files in the +// same package cannot collide over the helper's name. +func timeOfDaySpec(arrowType, builderType, cast, divisor string) colSpec { + return colSpec{ + arrowType: arrowType, + builderType: builderType, + needsTime: true, + appendStmt: func(bld, recv, _ string) string { + return fmt.Sprintf(`{ +tod := %s.UTC() +%s.Append(%s(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds()%s)) +}`, recv, bld, cast, divisor) + }, + } +} + +func byteSliceSpec(opts tagOpts) (colSpec, error) { + if err := checkOpts(opts, "[]byte", optSupport{largeView: true, dict: true}); err != nil { + return colSpec{}, err + } + // A nil []byte is a null in arreflect even when the column is not nullable, + // because its binary case tests the slice itself rather than the pointer. + if opts.Dict { + spec := dictSpec("arrow.BinaryTypes.Binary", "*array.BinaryDictionaryBuilder", simpleAppend()) + spec.nilable = true + return spec, nil + } + switch { + case opts.View: + return colSpec{arrowType: "arrow.BinaryTypes.BinaryView", builderType: "*array.BinaryViewBuilder", appendStmt: simpleAppend(), nilable: true}, nil + case opts.Large: + return colSpec{arrowType: "arrow.BinaryTypes.LargeBinary", builderType: "*array.BinaryBuilder", appendStmt: simpleAppend(), nilable: true}, nil + default: + return colSpec{arrowType: "arrow.BinaryTypes.Binary", builderType: "*array.BinaryBuilder", appendStmt: simpleAppend(), nilable: true}, nil + } +} + +func stringSpec(opts tagOpts) (colSpec, error) { + if err := checkOpts(opts, "string", optSupport{largeView: true, dict: true}); err != nil { + return colSpec{}, err + } + if opts.Dict { + return dictSpec("arrow.BinaryTypes.String", "*array.BinaryDictionaryBuilder", func(bld, _, val string) string { + return fmt.Sprintf("%s.AppendString(%s)", bld, val) + }), nil + } + switch { + case opts.View: + return colSpec{arrowType: "arrow.BinaryTypes.StringView", builderType: "*array.StringViewBuilder", appendStmt: simpleAppend()}, nil + case opts.Large: + return colSpec{arrowType: "arrow.BinaryTypes.LargeString", builderType: "*array.LargeStringBuilder", appendStmt: simpleAppend()}, nil + default: + return colSpec{arrowType: "arrow.BinaryTypes.String", builderType: "*array.StringBuilder", appendStmt: simpleAppend()}, nil + } +} + +// dictSpec wraps a value type in Dictionary, matching +// arreflect.applyEncodingOpts. Dictionary appends can fail, so the caller +// routes the result through the appender's first-error field. +func dictSpec(valueType, builderType string, appendStmt func(bld, recv, val string) string) colSpec { + return colSpec{ + arrowType: fmt.Sprintf("&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: %s}", valueType), + builderType: builderType, + fallible: true, + appendStmt: appendStmt, + } +} + +// numeric describes one Go basic kind's Arrow column. +type numeric struct { + arrowType string // e.g. "arrow.PrimitiveTypes.Int64" + builder string // e.g. "Int64" + cast string // non-empty when the Go type is wider-named than the column ("int" -> int64) +} + +var numerics = map[types.BasicKind]numeric{ + types.Int8: {"arrow.PrimitiveTypes.Int8", "Int8", ""}, + types.Int16: {"arrow.PrimitiveTypes.Int16", "Int16", ""}, + types.Int32: {"arrow.PrimitiveTypes.Int32", "Int32", ""}, + types.Int64: {"arrow.PrimitiveTypes.Int64", "Int64", ""}, + types.Int: {"arrow.PrimitiveTypes.Int64", "Int64", "int64"}, + types.Uint8: {"arrow.PrimitiveTypes.Uint8", "Uint8", ""}, + types.Uint16: {"arrow.PrimitiveTypes.Uint16", "Uint16", ""}, + types.Uint32: {"arrow.PrimitiveTypes.Uint32", "Uint32", ""}, + types.Uint64: {"arrow.PrimitiveTypes.Uint64", "Uint64", ""}, + types.Uint: {"arrow.PrimitiveTypes.Uint64", "Uint64", "uint64"}, + types.Float32: {"arrow.PrimitiveTypes.Float32", "Float32", ""}, + types.Float64: {"arrow.PrimitiveTypes.Float64", "Float64", ""}, +} + +func basicSpec(b *types.Basic, opts tagOpts) (colSpec, error) { + switch b.Kind() { + case types.String: + return stringSpec(opts) + case types.Bool: + if err := checkOpts(opts, "bool", optSupport{}); err != nil { + return colSpec{}, err + } + return colSpec{arrowType: "arrow.FixedWidthTypes.Boolean", builderType: "*array.BooleanBuilder", appendStmt: simpleAppend()}, nil + } + + n, ok := numerics[b.Kind()] + if !ok { + return colSpec{}, fmt.Errorf("type %s is not supported", b) + } + if err := checkOpts(opts, b.Name(), optSupport{dict: true}); err != nil { + return colSpec{}, err + } + appendStmt := simpleAppend() + if n.cast != "" { + appendStmt = castAppend(n.cast) + } + if opts.Dict { + return dictSpec(n.arrowType, "*array."+n.builder+"DictionaryBuilder", appendStmt), nil + } + return colSpec{arrowType: n.arrowType, builderType: "*array." + n.builder + "Builder", appendStmt: appendStmt}, nil +} + +func simpleAppend() func(bld, recv, val string) string { + return func(bld, _, val string) string { return fmt.Sprintf("%s.Append(%s)", bld, val) } +} + +func castAppend(cast string) func(bld, recv, val string) string { + return func(bld, _, val string) string { return fmt.Sprintf("%s.Append(%s(%s))", bld, cast, val) } +} diff --git a/arrgen/tag.go b/arrgen/tag.go new file mode 100644 index 000000000..a4c5b4813 --- /dev/null +++ b/arrgen/tag.go @@ -0,0 +1,146 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import ( + "fmt" + "strconv" + "strings" +) + +// tagOpts is the parsed form of one `arrow:"..."` struct tag. +// +// It is deliberately a mirror of the unexported tagOpts in arreflect: the two +// packages have to agree about what a tag means or generated code would encode +// a different schema than the reflection path, which the equivalence tests +// would then catch. Keep the two in sync when either grows an option. +type tagOpts struct { + Name string + Skip bool + Dict bool + View bool + REE bool + Large bool + DecimalPrecision int32 + DecimalScale int32 + HasDecimalOpts bool + Temporal string // "", "timestamp", "date32", "date64", "time32", "time64" +} + +// parseTag parses the value of an `arrow` struct tag. Unlike arreflect, which +// records a diagnostic and surfaces it later, an unusable tag is an error here: +// the generator has a file to write and no reason to defer the complaint. +func parseTag(tag string) (tagOpts, error) { + if tag == "-" { + return tagOpts{Skip: true}, nil + } + + name, rest, _ := strings.Cut(tag, ",") + opts := tagOpts{Name: name} + if rest == "" { + return opts, nil + } + for _, token := range splitTagTokens(rest) { + if err := applyTagToken(&opts, token); err != nil { + return tagOpts{}, err + } + } + return opts, nil +} + +// splitTagTokens splits the option list on commas that are not nested inside +// parentheses, so decimal(18,2) survives as a single token. +func splitTagTokens(rest string) []string { + var tokens []string + depth, start := 0, 0 + for i := 0; i < len(rest); i++ { + switch rest[i] { + case '(': + depth++ + case ')': + depth-- + case ',': + if depth == 0 { + tokens = append(tokens, strings.TrimSpace(rest[start:i])) + start = i + 1 + } + } + } + if start < len(rest) { + tokens = append(tokens, strings.TrimSpace(rest[start:])) + } + return tokens +} + +func applyTagToken(opts *tagOpts, token string) error { + if strings.HasPrefix(token, "decimal(") && strings.HasSuffix(token, ")") { + return parseDecimalOpt(opts, token) + } + switch token { + case "dict": + opts.Dict = true + case "view": + opts.View = true + case "ree": + opts.REE = true + case "large": + opts.Large = true + case "date32", "date64", "time32", "time64", "timestamp": + opts.Temporal = token + default: + return fmt.Errorf("unknown option %q", token) + } + return nil +} + +func parseDecimalOpt(opts *tagOpts, token string) error { + inner := strings.TrimSuffix(strings.TrimPrefix(token, "decimal("), ")") + parts := strings.SplitN(inner, ",", 2) + if len(parts) != 2 { + return fmt.Errorf("invalid decimal tag %q: expected decimal(precision,scale)", token) + } + p, err := strconv.ParseInt(strings.TrimSpace(parts[0]), 10, 32) + if err != nil { + return fmt.Errorf("invalid decimal tag %q: precision %q is not an integer", token, strings.TrimSpace(parts[0])) + } + s, err := strconv.ParseInt(strings.TrimSpace(parts[1]), 10, 32) + if err != nil { + return fmt.Errorf("invalid decimal tag %q: scale %q is not an integer", token, strings.TrimSpace(parts[1])) + } + opts.HasDecimalOpts = true + opts.DecimalPrecision = int32(p) + opts.DecimalScale = int32(s) + return nil +} + +// validate rejects option combinations that arreflect also rejects, so a struct +// arrgen accepts is always one arreflect can encode. +func (o tagOpts) validate() error { + if o.REE { + return fmt.Errorf("ree is not supported on a struct field; use arreflect.FromSlice with WithREE at the top level") + } + n := 0 + for _, set := range []bool{o.Dict, o.View, o.REE} { + if set { + n++ + } + } + if n > 1 { + return fmt.Errorf("conflicting options: at most one of dict, view, ree may be set") + } + return nil +} diff --git a/arrgen/template.go b/arrgen/template.go new file mode 100644 index 000000000..c44db14fc --- /dev/null +++ b/arrgen/template.go @@ -0,0 +1,147 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen + +import "text/template" + +// fileTemplate renders the whole generated file. The emitted code imports only +// arrow, arrow/array and arrow/memory (plus time, for time-of-day columns), so +// nothing a caller builds with arrgen ends up depending on arrgen itself. +var fileTemplate = template.Must(template.New("arrgen").Parse( + `{{if .Header}}{{.Header}} + +{{end}}// Code generated by arrgen. DO NOT EDIT. + +package {{.Package}} + +import ( +{{- if .NeedTime}} + "time" +{{end}} + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) +{{range .Types}} +// {{.SchemaVar}} is the Arrow schema of {{.GoName}}, resolved from its arrow +// struct tags at generate time. +var {{.SchemaVar}} = arrow.NewSchema([]arrow.Field{ +{{- range .Fields}} + {Name: {{printf "%q" .Name}}, Type: {{.ArrowType}}, Nullable: {{.Nullable}}}, +{{- end}} +}, nil) + +// {{.SchemaFunc}} returns the Arrow schema encoded by {{.AppenderType}}. It is +// equal to the schema arreflect.InferSchema[{{.GoName}}] infers at runtime. +func {{.SchemaFunc}}() *arrow.Schema { return {{.SchemaVar}} } + +// {{.AppenderType}} converts {{.GoName}} values into Arrow record batches +// without reflection: the struct tags were read when this file was generated, +// so appending a row is a handful of direct field loads into typed builders. +// +// It is not safe for concurrent use. Release it when done. +type {{.AppenderType}} struct { + rb *array.RecordBuilder +{{- range .Fields}} + {{.BuilderVar}} {{.BuilderType}} +{{- end}} + err error +} + +// {{.CtorName}} returns an appender that builds {{.GoName}} batches with mem, +// or memory.DefaultAllocator when mem is nil. +// +// The typed field builders are resolved once here rather than per row, so +// Append does no type assertions at all. +func {{.CtorName}}(mem memory.Allocator) *{{.AppenderType}} { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, {{.SchemaVar}}) + return &{{.AppenderType}}{ + rb: rb, +{{- range .Fields}} + {{.BuilderVar}}: rb.Field({{.Index}}).({{.BuilderType}}), +{{- end}} + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *{{.AppenderType}}) Schema() *arrow.Schema { return {{.SchemaVar}} } + +// RecordBuilder returns the underlying builder, for callers that need an Arrow +// API this appender does not wrap. Appending through it directly is allowed as +// long as every column stays the same length. +func (a *{{.AppenderType}}) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends is what keeps the append path free of allocations. +func (a *{{.AppenderType}}) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *{{.AppenderType}}) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller streaming rows can reuse a single {{.GoName}} variable across calls. +func (a *{{.AppenderType}}) Append(v *{{.GoName}}) { +{{- range .Fields}} + {{.AppendStmt}} +{{- end}} +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *{{.AppenderType}}) AppendSlice(vs []{{.GoName}}) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch snapshots the appended rows into a record batch and resets the +// appender for the next batch. The caller owns the batch and must Release it. +func (a *{{.AppenderType}}) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail, so an appender over a schema without one always returns nil. +// +// The error is sticky: NewRecordBatch does not clear it, because a batch built +// after a failed append is missing a value in that column. +func (a *{{.AppenderType}}) Err() error { return a.err } +{{if .AnyFallible}} +// setErr keeps the first error so the hot path stays a single branch. +func (a *{{.AppenderType}}) setErr(err error) { + if err != nil && a.err == nil { + a.err = err + } +} +{{end}} +// Release releases the appender's builders. +func (a *{{.AppenderType}}) Release() { a.rb.Release() } + +// {{.BatchFunc}} builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[{{.GoName}}](vs, mem) and +// produces an identical batch. The caller owns the batch and must Release it. +func {{.BatchFunc}}(mem memory.Allocator, vs []{{.GoName}}) (arrow.RecordBatch, error) { + a := {{.CtorName}}(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} +{{end}}`)) diff --git a/arrgen/testdata/basic.golden b/arrgen/testdata/basic.golden new file mode 100644 index 000000000..05fa66891 --- /dev/null +++ b/arrgen/testdata/basic.golden @@ -0,0 +1,242 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by arrgen. DO NOT EDIT. + +package basic + +import ( + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// metricArrowSchema is the Arrow schema of Metric, resolved from its arrow +// struct tags at generate time. +var metricArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "host", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: false}, + {Name: "count", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "Blob", Type: arrow.BinaryTypes.Binary, Nullable: false}, +}, nil) + +// MetricSchema returns the Arrow schema encoded by MetricAppender. It is +// equal to the schema arreflect.InferSchema[Metric] infers at runtime. +func MetricSchema() *arrow.Schema { return metricArrowSchema } + +// MetricAppender converts Metric values into Arrow record batches +// without reflection: the struct tags were read when this file was generated, +// so appending a row is a handful of direct field loads into typed builders. +// +// It is not safe for concurrent use. Release it when done. +type MetricAppender struct { + rb *array.RecordBuilder + b0 *array.TimestampBuilder + b1 *array.BinaryDictionaryBuilder + b2 *array.Int64Builder + b3 *array.BinaryBuilder + err error +} + +// NewMetricAppender returns an appender that builds Metric batches with mem, +// or memory.DefaultAllocator when mem is nil. +// +// The typed field builders are resolved once here rather than per row, so +// Append does no type assertions at all. +func NewMetricAppender(mem memory.Allocator) *MetricAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, metricArrowSchema) + return &MetricAppender{ + rb: rb, + b0: rb.Field(0).(*array.TimestampBuilder), + b1: rb.Field(1).(*array.BinaryDictionaryBuilder), + b2: rb.Field(2).(*array.Int64Builder), + b3: rb.Field(3).(*array.BinaryBuilder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *MetricAppender) Schema() *arrow.Schema { return metricArrowSchema } + +// RecordBuilder returns the underlying builder, for callers that need an Arrow +// API this appender does not wrap. Appending through it directly is allowed as +// long as every column stays the same length. +func (a *MetricAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends is what keeps the append path free of allocations. +func (a *MetricAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller streaming rows can reuse a single Metric variable across calls. +func (a *MetricAppender) Append(v *Metric) { + a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) + a.setErr(a.b1.AppendString(v.Host)) + if v.Count == nil { + a.b2.AppendNull() + } else { + a.b2.Append(*v.Count) + } + if v.Blob == nil { + a.b3.AppendNull() + } else { + a.b3.Append(v.Blob) + } +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *MetricAppender) AppendSlice(vs []Metric) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch snapshots the appended rows into a record batch and resets the +// appender for the next batch. The caller owns the batch and must Release it. +func (a *MetricAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail, so an appender over a schema without one always returns nil. +// +// The error is sticky: NewRecordBatch does not clear it, because a batch built +// after a failed append is missing a value in that column. +func (a *MetricAppender) Err() error { return a.err } + +// setErr keeps the first error so the hot path stays a single branch. +func (a *MetricAppender) setErr(err error) { + if err != nil && a.err == nil { + a.err = err + } +} + +// Release releases the appender's builders. +func (a *MetricAppender) Release() { a.rb.Release() } + +// MetricRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[Metric](vs, mem) and +// produces an identical batch. The caller owns the batch and must Release it. +func MetricRecordBatch(mem memory.Allocator, vs []Metric) (arrow.RecordBatch, error) { + a := NewMetricAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} + +// readingArrowSchema is the Arrow schema of reading, resolved from its arrow +// struct tags at generate time. +var readingArrowSchema = arrow.NewSchema([]arrow.Field{ + {Name: "at", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, + {Name: "value", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, +}, nil) + +// readingSchema returns the Arrow schema encoded by readingAppender. It is +// equal to the schema arreflect.InferSchema[reading] infers at runtime. +func readingSchema() *arrow.Schema { return readingArrowSchema } + +// readingAppender converts reading values into Arrow record batches +// without reflection: the struct tags were read when this file was generated, +// so appending a row is a handful of direct field loads into typed builders. +// +// It is not safe for concurrent use. Release it when done. +type readingAppender struct { + rb *array.RecordBuilder + b0 *array.Date32Builder + b1 *array.Float64Builder + err error +} + +// newReadingAppender returns an appender that builds reading batches with mem, +// or memory.DefaultAllocator when mem is nil. +// +// The typed field builders are resolved once here rather than per row, so +// Append does no type assertions at all. +func newReadingAppender(mem memory.Allocator) *readingAppender { + if mem == nil { + mem = memory.DefaultAllocator + } + rb := array.NewRecordBuilder(mem, readingArrowSchema) + return &readingAppender{ + rb: rb, + b0: rb.Field(0).(*array.Date32Builder), + b1: rb.Field(1).(*array.Float64Builder), + } +} + +// Schema returns the schema of the batches this appender builds. +func (a *readingAppender) Schema() *arrow.Schema { return readingArrowSchema } + +// RecordBuilder returns the underlying builder, for callers that need an Arrow +// API this appender does not wrap. Appending through it directly is allowed as +// long as every column stays the same length. +func (a *readingAppender) RecordBuilder() *array.RecordBuilder { return a.rb } + +// Reserve pre-allocates room for n more rows in every column. Reserving before +// a run of Appends is what keeps the append path free of allocations. +func (a *readingAppender) Reserve(n int) { a.rb.Reserve(n) } + +// Len reports the number of rows appended since the last NewRecordBatch. +func (a *readingAppender) Len() int { return a.rb.Field(0).Len() } + +// Append appends one row. v is read synchronously and never retained, so a +// caller streaming rows can reuse a single reading variable across calls. +func (a *readingAppender) Append(v *reading) { + a.b0.Append(arrow.Date32FromTime(v.At)) + a.b1.Append(v.Value) +} + +// AppendSlice appends every element of vs, reserving room for them up front. +func (a *readingAppender) AppendSlice(vs []reading) { + a.rb.Reserve(len(vs)) + for i := range vs { + a.Append(&vs[i]) + } +} + +// NewRecordBatch snapshots the appended rows into a record batch and resets the +// appender for the next batch. The caller owns the batch and must Release it. +func (a *readingAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } + +// Err returns the first error an Append hit, or nil. Only dictionary columns +// can fail, so an appender over a schema without one always returns nil. +// +// The error is sticky: NewRecordBatch does not clear it, because a batch built +// after a failed append is missing a value in that column. +func (a *readingAppender) Err() error { return a.err } + +// Release releases the appender's builders. +func (a *readingAppender) Release() { a.rb.Release() } + +// readingRecordBatch builds one record batch holding every element of vs. It is the +// generated equivalent of arreflect.RecordFromSlice[reading](vs, mem) and +// produces an identical batch. The caller owns the batch and must Release it. +func readingRecordBatch(mem memory.Allocator, vs []reading) (arrow.RecordBatch, error) { + a := newReadingAppender(mem) + defer a.Release() + a.AppendSlice(vs) + if err := a.Err(); err != nil { + return nil, err + } + return a.NewRecordBatch(), nil +} diff --git a/arrgen/testdata/basic/basic.go b/arrgen/testdata/basic/basic.go new file mode 100644 index 000000000..e46e0b9cf --- /dev/null +++ b/arrgen/testdata/basic/basic.go @@ -0,0 +1,37 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package basic is the fixture behind arrgen's golden test. It is under +// testdata so the go tool does not build it as part of the module; the +// generator loads it by directory instead. +package basic + +import "time" + +// Metric is an exported type, so its generated API is exported too. +type Metric struct { + Time time.Time `arrow:"ts"` + Host string `arrow:"host,dict"` + Count *int64 `arrow:"count"` + Blob []byte + Local string `arrow:"-"` +} + +// reading is unexported, so the generated appender is unexported as well. +type reading struct { + At time.Time `arrow:"at,date32"` + Value float64 `arrow:"value"` +} diff --git a/arrgen/testdata/errors/errors.go b/arrgen/testdata/errors/errors.go new file mode 100644 index 000000000..cc4e54f8d --- /dev/null +++ b/arrgen/testdata/errors/errors.go @@ -0,0 +1,128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package errors holds one struct per rejection the generator makes, so the +// error messages stay covered and stay readable. +package errors + +import "time" + +// Inner is a nested struct used by Nested and Embedded. +type Inner struct { + A int64 +} + +// ID is a defined scalar type, which arreflect matches by exact type and so +// cannot map either. +type ID int64 + +type Embedded struct { + Inner + B int64 +} + +type Nested struct { + Field Inner +} + +type SliceField struct { + Field []int64 +} + +type ArrayField struct { + Field [3]int64 +} + +type MapField struct { + Field map[string]int64 +} + +type NamedScalar struct { + Field ID +} + +type UnknownOption struct { + Field int64 `arrow:"f,nope"` +} + +type BadDecimalTag struct { + Field int64 `arrow:"f,decimal(a,2)"` +} + +type ShortDecimalTag struct { + Field int64 `arrow:"f,decimal(4)"` +} + +type DuplicateNames struct { + A int64 `arrow:"same"` + B int64 `arrow:"same"` +} + +type TemporalOnInt struct { + Field int64 `arrow:"f,date32"` +} + +type DecimalOnInt struct { + Field int64 `arrow:"f,decimal(10,2)"` +} + +type DictOnBool struct { + Field bool `arrow:"f,dict"` +} + +type ViewOnInt struct { + Field int64 `arrow:"f,view"` +} + +type LargeOnInt struct { + Field int64 `arrow:"f,large"` +} + +type RunEndEncoded struct { + Field string `arrow:"f,ree"` +} + +type DictAndView struct { + Field string `arrow:"f,dict,view"` +} + +type DictAndLarge struct { + Field string `arrow:"f,dict,large"` +} + +type NoColumns struct { + Field string `arrow:"-"` + hidden int64 +} + +type DoublePointer struct { + Field **int64 +} + +type TimeSlice struct { + Field []time.Time +} + +// NotAStruct is a named type whose underlying type is not a struct. +type NotAStruct int64 + +type Collide struct { + Field int64 +} + +type collide struct { + Field int64 +} diff --git a/arrgen/testdata/pregen/pregen.go b/arrgen/testdata/pregen/pregen.go new file mode 100644 index 000000000..009eecf33 --- /dev/null +++ b/arrgen/testdata/pregen/pregen.go @@ -0,0 +1,37 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package pregen is the state a package is in the first time the generator runs +// against it: the call sites exist, the generated file does not, and so the +// package does not type-check. The generator has to work anyway. +package pregen + +import ( + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +type Sample struct { + At time.Time `arrow:"at"` + Value float64 `arrow:"value"` +} + +// Encode calls into the file arrgen has not written yet. +func Encode(mem memory.Allocator, samples []Sample) (arrow.RecordBatch, error) { + return SampleRecordBatch(mem, samples) +} From 1d41aa7356a623ed25701a1a043aa63f7e2f52f8 Mon Sep 17 00:00:00 2001 From: jay <38236622+JayJamieson@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:37:52 +0000 Subject: [PATCH 2/4] fix(arrgen): address review feedback - Remove the duplicated arrgen/arrgen subtree. The two copies were functionally identical; they differed only in comment prose. - Reject the field spellings arreflect cannot infer as a struct field rather than generating a column that diverges from it: untagged (and ,timestamp) time.Time, and untagged decimal128.Num/decimal256.Num. arreflect's inferArrowType switches on reflect.Kind before the types it matches by identity, so all three infer an empty struct<> and drop the value. Every skip is gone from the equivalence tests, which now compare all 45 fixture columns and all three schemas strictly. TestArreflectCannotInferStructScalars pins the upstream behaviour so a fix there tells us to relax the rejections. - Replace the public example package with testable examples on arrgen itself, driven by the internal/gentypes fixture. - Give the nested module explicit CI coverage in ci/scripts/build.sh and ci/scripts/test.sh: build, vet, tests under the same -race/-asan args as the rest of the repo, and go generate followed by git diff --exit-code. A second pass runs the equivalence tests in a workspace over the local tree, so an arreflect change cannot break them silently. golangci-lint now covers the module in pre-commit too. - Document a go:generate workflow that resolves: go get -tool plus go tool arrgen. The unversioned go run spelling only works inside a module that already requires arrgen. --- .gitignore | 5 + .pre-commit-config.yaml | 3 + arrgen/README.md | 130 ++++- arrgen/arrgen/README.md | 209 -------- arrgen/arrgen/cmd/arrgen/main.go | 102 ---- arrgen/arrgen/doc.go | 51 -- arrgen/arrgen/drift_test.go | 66 --- arrgen/arrgen/example/example_test.go | 105 ---- arrgen/arrgen/example/metric_arrow.go | 126 ----- arrgen/arrgen/generate.go | 335 ------------ arrgen/arrgen/generate_test.go | 222 -------- arrgen/arrgen/go.mod | 35 -- arrgen/arrgen/go.sum | 44 -- arrgen/arrgen/internal/gentypes/alloc_test.go | 158 ------ arrgen/arrgen/internal/gentypes/bench_test.go | 219 -------- .../internal/gentypes/equivalence_test.go | 449 ----------------- .../gentypes/instrumented_race_test.go | 21 - .../internal/gentypes/instrumented_test.go | 24 - arrgen/arrgen/internal/gentypes/row_arrow.go | 475 ------------------ arrgen/arrgen/internal/gentypes/types.go | 98 ---- arrgen/arrgen/license_header.txt | 15 - arrgen/arrgen/mapping.go | 320 ------------ arrgen/arrgen/tag.go | 140 ------ arrgen/arrgen/template.go | 138 ----- arrgen/arrgen/testdata/basic.golden | 226 --------- arrgen/arrgen/testdata/basic/basic.go | 37 -- arrgen/arrgen/testdata/errors/errors.go | 126 ----- arrgen/arrgen/testdata/pregen/pregen.go | 37 -- arrgen/doc.go | 25 +- arrgen/drift_test.go | 2 +- arrgen/example/example_test.go | 107 ---- arrgen/example/metric.go | 42 -- arrgen/example_test.go | 166 ++++++ arrgen/generate_test.go | 43 ++ arrgen/internal/gentypes/bench_test.go | 22 +- arrgen/internal/gentypes/equivalence_test.go | 261 +++------- .../example => internal/gentypes}/metric.go | 17 +- .../gentypes}/metric_arrow.go | 10 +- arrgen/internal/gentypes/row_arrow.go | 228 +++++---- arrgen/internal/gentypes/types.go | 27 +- arrgen/mapping.go | 55 +- arrgen/testdata/basic.golden | 8 +- arrgen/testdata/basic/basic.go | 2 +- arrgen/testdata/errors/errors.go | 26 +- arrgen/testdata/pregen/pregen.go | 2 +- ci/scripts/build.sh | 7 + ci/scripts/test.sh | 31 ++ 47 files changed, 670 insertions(+), 4327 deletions(-) delete mode 100644 arrgen/arrgen/README.md delete mode 100644 arrgen/arrgen/cmd/arrgen/main.go delete mode 100644 arrgen/arrgen/doc.go delete mode 100644 arrgen/arrgen/drift_test.go delete mode 100644 arrgen/arrgen/example/example_test.go delete mode 100644 arrgen/arrgen/example/metric_arrow.go delete mode 100644 arrgen/arrgen/generate.go delete mode 100644 arrgen/arrgen/generate_test.go delete mode 100644 arrgen/arrgen/go.mod delete mode 100644 arrgen/arrgen/go.sum delete mode 100644 arrgen/arrgen/internal/gentypes/alloc_test.go delete mode 100644 arrgen/arrgen/internal/gentypes/bench_test.go delete mode 100644 arrgen/arrgen/internal/gentypes/equivalence_test.go delete mode 100644 arrgen/arrgen/internal/gentypes/instrumented_race_test.go delete mode 100644 arrgen/arrgen/internal/gentypes/instrumented_test.go delete mode 100644 arrgen/arrgen/internal/gentypes/row_arrow.go delete mode 100644 arrgen/arrgen/internal/gentypes/types.go delete mode 100644 arrgen/arrgen/license_header.txt delete mode 100644 arrgen/arrgen/mapping.go delete mode 100644 arrgen/arrgen/tag.go delete mode 100644 arrgen/arrgen/template.go delete mode 100644 arrgen/arrgen/testdata/basic.golden delete mode 100644 arrgen/arrgen/testdata/basic/basic.go delete mode 100644 arrgen/arrgen/testdata/errors/errors.go delete mode 100644 arrgen/arrgen/testdata/pregen/pregen.go delete mode 100644 arrgen/example/example_test.go delete mode 100644 arrgen/example/metric.go create mode 100644 arrgen/example_test.go rename arrgen/{arrgen/example => internal/gentypes}/metric.go (69%) rename arrgen/{example => internal/gentypes}/metric_arrow.go (95%) diff --git a/.gitignore b/.gitignore index ec2e3422b..57a189629 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,8 @@ /dev/release/apache-rat-*.jar /dev/release/filtered_rat.txt /dev/release/rat.xml + +# A go.work over the root module and arrgen is the local way to build the +# nested module against this tree rather than a released arrow-go. +/go.work +/go.work.sum diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5dd3a320a..9e06889a2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -30,6 +30,9 @@ repos: - id: golangci-lint-full name: golangci-lint-full-internal entry: bash -c 'cd internal && golangci-lint run' + - id: golangci-lint-full + name: golangci-lint-full-arrgen + entry: bash -c 'cd arrgen && golangci-lint run' - repo: local hooks: - id: rat diff --git a/arrgen/README.md b/arrgen/README.md index 6c231006a..f015ab56e 100644 --- a/arrgen/README.md +++ b/arrgen/README.md @@ -34,7 +34,14 @@ a build-time step; your binary does not grow a dependency for it. ## Quick start -Tag a struct as you would for `arreflect`, and add a `go:generate` line: +Add `arrgen` to your module as a tool dependency, which needs `go 1.24` or +later in your `go.mod`: + +```sh +go get -tool github.com/apache/arrow-go/arrgen/cmd/arrgen +``` + +Then tag a struct as you would for `arreflect` and add a `go:generate` line: ```go package telemetry @@ -42,14 +49,14 @@ package telemetry import "time" type Metric struct { - Time time.Time `arrow:"ts"` + Day time.Time `arrow:"day,date32"` Host string `arrow:"host"` CPU float64 `arrow:"cpu"` Value *float64 `arrow:"value"` // nullable Secret string `arrow:"-"` // not a column } -//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +//go:generate go tool arrgen -type Metric ``` ```sh @@ -57,13 +64,32 @@ go generate ./... # writes metric_arrow.go next to the type ``` Check the result in, as you would `stringer` or `easyjson` output, and -regenerate when the struct changes. If you would rather not add `arrgen` to your -module's dependencies at all, pin it in the directive instead: +regenerate when the struct changes. + +`go get -tool` records the generator in your `go.mod` and `go.sum` but not in +your build: nothing in your packages imports it, so it is not linked into your +binary. If you would rather not record it at all, name a version in the +directive instead, which resolves the generator per run without touching your +module files: ```go -//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen@v0.1.0 -type Metric +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen@arrgen/v0.1.0 -type Metric +``` + +Both forms need `arrgen` to be resolvable. As a nested module it is versioned +under its own `arrgen/vX.Y.Z` tags, and until the first of those is published +neither form resolves from a released version - use a `replace` directive +pointing at a checkout in the meantime: + +```sh +go mod edit -replace github.com/apache/arrow-go/arrgen=../arrow-go/arrgen ``` +The unversioned `go run github.com/apache/arrow-go/arrgen/cmd/arrgen` spelling +works only inside a module that already requires `arrgen`; anywhere else the go +tool refuses it with `no required module provides package`. The directives in +this module use that spelling because this module is `arrgen`. + ### What you get ```go @@ -91,6 +117,9 @@ a.Err() // first append error; only dictionary columns can fa `Append` reads `v` synchronously and never retains it, so a caller draining a stream can reuse one row variable for every row. +Runnable versions of all three are the testable examples in +[`example_test.go`](example_test.go). + ## Flags | Flag | Meaning | @@ -108,8 +137,16 @@ stream can reuse one row variable for every row. those for a nullable column. Tag options are `arreflect`'s: a leading column name, `-` to skip the field, the -temporal overrides `date32`, `date64`, `time32`, `time64` and `timestamp`, -`dict`, `view`, `large`, and `decimal(precision,scale)`. +temporal overrides `date32`, `date64`, `time32` and `time64`, `dict`, `view`, +`large`, and `decimal(precision,scale)`. + +Two of those types need an explicit tag, because the untagged spelling is one +`arreflect` cannot infer as a struct field - see [below](#the-two-types-that-need-a-tag): + +| Field | Required tag | +| --- | --- | +| `time.Time` | one of `date32`, `date64`, `time32`, `time64` | +| `decimal128.Num`, `decimal256.Num` | `decimal(precision,scale)` | Anything else such as nested structs, slices other than `[]byte`, arrays, maps, embedded fields, defined scalar types such as `type ID int64` is a @@ -124,20 +161,50 @@ The point of matching `arreflect`'s tag dialect is that switching between the two paths is a call-site change and nothing else. `internal/gentypes` holds a fixture with one field for every supported column shape and asserts, on every run, that the generated encoder and `arreflect` produce the same schema and the -same column data for the same input. +same column data for the same input - every column, with nothing excused. + +### The two types that need a tag + +`arreflect`'s `inferArrowType` switches on `reflect.Kind` before it reaches the +types it matches by identity, so a Go struct that Arrow models as a scalar +reaches `inferStructType` instead. That finds only unexported fields and yields +an empty `struct<>`, and the value is then dropped. Today that affects: + +| Field | `arreflect` infers | +| --- | --- | +| `time.Time`, untagged or `,timestamp` | `struct<>` | +| `decimal128.Num` / `decimal256.Num`, untagged | `struct<>` | + +A tag naming the Arrow type outright - `,date32`, `,decimal(20,3)` - rescues the +field, because `arreflect` applies the tag after the inferred type. +`arreflect`'s own tests only pass these types as the top-level element of a +slice, where `buildArray` special-cases them before the `reflect.Struct` branch +and they encode correctly. + +`arrgen` could emit the column Arrow plainly means here, and an earlier revision +did. But then generated code and `arreflect` would disagree about the schema, +which is the one thing this generator exists not to do, so the untaggable +spellings are a generate-time error naming the field and the tag that fixes it +instead. The practical consequence is that **`arrgen` cannot emit a `TIMESTAMP` +column**: `,timestamp` is rejected along with the untagged spelling, since +`arreflect` infers `struct<>` for both. + +`TestArreflectCannotInferStructScalars` pins that upstream behavior. If +`arrow/array/arreflect` is fixed, that test fails and tells us to drop these +rejections and generate the columns instead. ## Performance -`go test ./internal/gentypes/ -bench . -benchmem`, Go 1.25, linux/arm64. Batch -benchmarks encode 1024 rows per operation. +`go test ./internal/gentypes/ -bench . -benchmem`, Go 1.25, linux/arm64, 2 vCPU. +Batch benchmarks encode 1024 rows per operation. | Benchmark | arreflect | arrgen | | | --- | --- | --- | --- | -| `MetricBatch` (4 columns) | 56.9 µs, 98 allocs | 20.3 µs, 52 allocs | **2.8x faster** | -| `FixedBatch` (5 fixed-width columns) | 59.0 µs, 97 allocs | 17.4 µs, 46 allocs | **3.4x faster** | -| `RowBatch` (46 columns) | 843 µs, 3430 allocs | 413 µs, 3126 allocs | **2.0x faster** | -| `StreamAppend` (per row) | 58.3 ns | 18.6 ns, **0 allocs** | **3.1x faster** | -| `Schema` | 5.6 µs, 132 allocs | 1.5 ns, **0 allocs** | | +| `MetricBatch` (4 columns) | 69.6 µs, 97 allocs | 30.5 µs, 52 allocs | **2.3x faster** | +| `FixedBatch` (5 fixed-width columns) | 72.9 µs, 96 allocs | 26.7 µs, 46 allocs | **2.7x faster** | +| `RowBatch` (45 columns) | 985 µs, 3418 allocs | 515 µs, 3118 allocs | **1.9x faster** | +| `StreamAppend` (per row) | 73.2 ns | 27.5 ns, **0 allocs** | **2.7x faster** | +| `Schema` | 6.4 µs, 131 allocs | 1.8 ns, **0 allocs** | | Three things move the numbers: @@ -163,17 +230,31 @@ an appender can be fed a row at a time and cut into batches wherever you like. ## Working on arrgen -The repository root carries a `go.work` so this module builds against the -arrow-go tree next to it rather than the released version its `go.mod` names: +This module's `go.mod` names a released `github.com/apache/arrow-go/v18`, so by +default it builds against that rather than the tree it sits in: ```sh -go test ./arrgen/... # workspace: local arrow-go -cd arrgen && GOWORK=off go test ./... # standalone: released arrow-go +cd arrgen && go test ./... ``` -Both are expected to pass. `TestCheckedInFilesAreUpToDate` regenerates every -committed output in this module and fails if it differs, so a struct change -without a `go generate` is caught at the point of the edit. +To test it against the local arrow-go instead - which is what you want when +changing `arrow/array/arreflect`, since the equivalence tests are what catch a +divergence - put a workspace over the two modules: + +```sh +go work init . ./arrgen # from the repository root; go.work is gitignored +go test ./arrgen/... +``` + +`ci/scripts/build.sh` and `ci/scripts/test.sh` both descend into this module +explicitly, because `./...` in the root module stops at its `go.mod`. CI builds +it, vets it, runs its tests under the same `-race`/`-asan` args as the rest of +the repository, and runs `go generate ./...` followed by `git diff +--exit-code` so a committed output that drifted from its struct fails the build. + +`TestCheckedInFilesAreUpToDate` checks the same property from inside the test +binary, which is what fails first when you edit a struct and forget to +regenerate. The golden file in `testdata` is the review surface for generator changes: `go test ./arrgen/ -update` rewrites it, and the resulting diff is exactly what @@ -185,4 +266,5 @@ As a nested module, `arrgen` is versioned and tagged independently of the root module: its tags are `arrgen/vX.Y.Z`, not `vX.Y.Z`. It cannot share the root's `v18` line, because a module path without a `/vN` suffix is limited to v0 and v1. Nothing in the release scripts tags it yet - that is a deliberate omission -for maintainers to decide on, not an oversight. +for maintainers to decide on, not an oversight, and it is why the Quick Start +cannot yet name a version that resolves. diff --git a/arrgen/arrgen/README.md b/arrgen/arrgen/README.md deleted file mode 100644 index 41b88a024..000000000 --- a/arrgen/arrgen/README.md +++ /dev/null @@ -1,209 +0,0 @@ - - -# arrgen - -Zero-reflection Arrow encoders, generated from your Go structs. - -`arrgen` is the code-generation counterpart to -[`arrow/array/arreflect`](../arrow/array/arreflect). `arreflect` reads a -struct's `arrow:"..."` tags with reflection on every value; `arrgen` reads the -same tags once, when you run it, and writes typed Go source that appends struct -fields straight into typed Arrow builders. - -It lives in its own Go module: nothing in `github.com/apache/arrow-go/v18` -imports it, and the code it emits imports only `arrow`, `arrow/array` and -`arrow/memory` — never `arrgen` itself. Generating is a build-time step; your -binary does not grow a dependency for it. - -## Quick start - -Tag a struct as you would for `arreflect`, and add a `go:generate` line: - -```go -package telemetry - -import "time" - -type Metric struct { - Time time.Time `arrow:"ts"` - Host string `arrow:"host"` - CPU float64 `arrow:"cpu"` - Value *float64 `arrow:"value"` // nullable - Secret string `arrow:"-"` // not a column -} - -//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -``` - -```sh -go generate ./... # writes metric_arrow.go next to the type -``` - -Check the result in, as you would `stringer` or `easyjson` output, and -regenerate when the struct changes. If you would rather not add `arrgen` to your -module's dependencies at all, pin it in the directive instead: - -```go -//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen@v0.1.0 -type Metric -``` - -### What you get - -```go -// One batch from a slice - the drop-in for arreflect.RecordFromSlice. -rec, err := telemetry.MetricRecordBatch(mem, metrics) - -// Or stream rows in and cut batches where you want them. -a := telemetry.NewMetricAppender(mem) -defer a.Release() -a.Reserve(batchSize) -for row := range rows { - a.Append(&row) - if a.Len() == batchSize { - rec := a.NewRecordBatch() - ... // hand it off - rec.Release() - a.Reserve(batchSize) - } -} - -telemetry.MetricSchema() // the schema, built once at init -a.Err() // first append error; only dictionary columns can fail -``` - -`Append` reads `v` synchronously and never retains it, so a caller draining a -stream can reuse one row variable for every row. - -## Flags - -| Flag | Meaning | -| --- | --- | -| `-type` | Struct type name. Repeatable, or comma-separated. Required. | -| `-output` | Output file, relative to `-dir`. Defaults to `_arrow.go`. | -| `-dir` | Package directory. Defaults to `.`, which is where `go:generate` runs. | -| `-header` | File whose contents are copied above the generated-code marker, for a license header. | - -## Supported fields - -`bool`, `int8/16/32/64`, `int`, `uint8/16/32/64`, `uint`, `float32/64`, -`string`, `[]byte`, `time.Time`, `time.Duration`, `decimal.Decimal32`, -`decimal.Decimal64`, `decimal128.Num`, `decimal256.Num`, and a pointer to any of -those for a nullable column. - -Tag options are `arreflect`'s: a leading column name, `-` to skip the field, the -temporal overrides `date32`, `date64`, `time32`, `time64` and `timestamp`, -`dict`, `view`, `large`, and `decimal(precision,scale)`. - -Anything else — nested structs, slices other than `[]byte`, arrays, maps, -embedded fields, defined scalar types such as `type ID int64` — is a -**generate-time error naming the field**, never a silently dropped column. Those -structs still encode fine with `arreflect` at runtime. - -The generator is equally strict about options that would do nothing: `date32` on -an `int64` field is an error here even though `arreflect` ignores it. - -## Equivalence with arreflect - -The point of matching `arreflect`'s tag dialect is that switching between the -two paths is a call-site change and nothing else. `internal/gentypes` holds a -fixture with one field for every supported column shape and asserts, on every -run, that the generated encoder and `arreflect` produce the same schema and the -same column data for the same input. - -### Two columns where they differ, and why - -`arreflect`'s struct-field path dispatches on `reflect.Kind` before it checks -the types it special-cases, so a Go struct that Arrow models as a scalar reaches -its nested-struct branch, which finds only unexported fields and returns an -empty `struct<>`. The value is then dropped. This affects: - -| Field | `arreflect` infers | `arrgen` emits | -| --- | --- | --- | -| `time.Time` (untagged, or `,timestamp`) | `struct<>` | `timestamp[ns, tz=UTC]` | -| `decimal128.Num` / `decimal256.Num` (untagged) | `struct<>` | `decimal128` / `decimal256` | - -A tag that names an explicit Arrow type — `,date32`, `,decimal(20,3)` — rescues -the field, because `arreflect` overrides the inferred type from the tag -afterwards. `arreflect`'s own tests only pass these types as the top-level -element of a slice, where a different branch handles them correctly. - -`arrgen` emits what `arreflect.InferType[time.Time]` and its top-level -`FromSlice` path both agree on. `TestArreflectStructFieldGaps` pins the current -upstream behavior, so if `arrow/array/arreflect` is fixed, that test fails and -tells us to delete this section. - -## Performance - -`go test ./internal/gentypes/ -bench . -benchmem`, Go 1.25, linux/arm64. Batch -benchmarks encode 1024 rows per operation. - -| Benchmark | arreflect | arrgen | | -| --- | --- | --- | --- | -| `MetricBatch` (4 columns) | 56.9 µs, 98 allocs | 20.3 µs, 52 allocs | **2.8x faster** | -| `FixedBatch` (5 fixed-width columns) | 59.0 µs, 97 allocs | 17.4 µs, 46 allocs | **3.4x faster** | -| `RowBatch` (46 columns) | 843 µs, 3430 allocs | 413 µs, 3126 allocs | **2.0x faster** | -| `StreamAppend` (per row) | 58.3 ns | 18.6 ns, **0 allocs** | **3.1x faster** | -| `Schema` | 5.6 µs, 132 allocs | 1.5 ns, **0 allocs** | | - -Three things move the numbers: - -- **Field access.** The generated code loads `v.CPU` and calls - `Float64Builder.Append`. The reflection path resolves a field index, produces - a `reflect.Value`, and dispatches on the builder's dynamic type, per value. -- **Builder lookup.** Typed builders are resolved once, in the constructor, so - `Append` does no type assertions at all — not even the - `b.Field(i).(*array.Float64Builder)` a hand-written encoder usually repeats. -- **Schema construction.** The schema is a package-level variable rather than a - struct walk per call. - -Allocation counts converge on wide fixtures because both paths pay for the same -Arrow buffers, and because Arrow's dictionary builders allocate per appended -value regardless of who calls them. The place the difference is unambiguous is -the streaming path: with room reserved, `Append` allocates nothing, which -`TestStreamingAppendIsAllocationFree` asserts. - -The streaming path also removes a constraint: `arreflect` cannot encode anything -until the caller has materialized the whole `[]T`, while an appender can be fed -a row at a time and cut into batches wherever you like. - -## Working on arrgen - -The repository root carries a `go.work` so this module builds against the -arrow-go tree next to it rather than the released version its `go.mod` names: - -```sh -go test ./arrgen/... # workspace: local arrow-go -cd arrgen && GOWORK=off go test ./... # standalone: released arrow-go -``` - -Both are expected to pass. `TestCheckedInFilesAreUpToDate` regenerates every -committed output in this module and fails if it differs, so a struct change -without a `go generate` is caught at the point of the edit. - -The golden file in `testdata` is the review surface for generator changes: -`go test ./arrgen/ -update` rewrites it, and the resulting diff is exactly what -users would see in their own regenerated code. - -### Releasing - -As a nested module, `arrgen` is versioned and tagged independently of the root -module: its tags are `arrgen/vX.Y.Z`, not `vX.Y.Z`. It cannot share the root's -`v18` line, because a module path without a `/vN` suffix is limited to v0 and -v1. Nothing in the release scripts tags it yet. diff --git a/arrgen/arrgen/cmd/arrgen/main.go b/arrgen/arrgen/cmd/arrgen/main.go deleted file mode 100644 index 31d497b4f..000000000 --- a/arrgen/arrgen/cmd/arrgen/main.go +++ /dev/null @@ -1,102 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Command arrgen generates zero-reflection Arrow appenders for Go struct types. -// -// It reads the same arrow:"..." struct tags that -// github.com/apache/arrow-go/v18/arrow/array/arreflect interprets at runtime, -// but reads them once, when you run it, and writes typed Go source instead. -// -// Usage: -// -// arrgen -type Metric [-type Other] [-output metric_arrow.go] [-header LICENSE.txt] -// -// Typically it is invoked through go:generate, next to the type: -// -// //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -// -// Check the generated file in, as you would stringer or easyjson output. -package main - -import ( - "flag" - "fmt" - "os" - "path/filepath" - "strings" - - "github.com/apache/arrow-go/arrgen" -) - -func main() { - if err := run(os.Args[1:]); err != nil { - fmt.Fprintf(os.Stderr, "arrgen: %v\n", err) - os.Exit(1) - } -} - -// typeList collects -type values: both "-type A,B" and "-type A -type B" work. -type typeList []string - -func (t *typeList) String() string { return strings.Join(*t, ",") } - -func (t *typeList) Set(v string) error { - for _, p := range strings.Split(v, ",") { - if p = strings.TrimSpace(p); p != "" { - *t = append(*t, p) - } - } - return nil -} - -func run(argv []string) error { - fs := flag.NewFlagSet("arrgen", flag.ContinueOnError) - var types typeList - fs.Var(&types, "type", "struct type name to generate for; repeatable, or comma-separated (required)") - dir := fs.String("dir", ".", "directory of the package holding the types") - output := fs.String("output", "", "output file, relative to -dir (default _arrow.go)") - header := fs.String("header", "", "file whose contents are copied above the generated-code marker, e.g. a license header") - if err := fs.Parse(argv); err != nil { - return err - } - if len(types) == 0 { - return fmt.Errorf("-type is required") - } - - cfg := arrgen.Config{Dir: *dir, Types: types} - if *header != "" { - b, err := os.ReadFile(*header) - if err != nil { - return fmt.Errorf("reading -header: %w", err) - } - cfg.Header = string(b) - } - - src, err := arrgen.Generate(cfg) - if err != nil { - return err - } - - out := *output - if out == "" { - out = strings.ToLower(types[0]) + "_arrow.go" - } - path := filepath.Join(*dir, out) - if err := os.WriteFile(path, src, 0o644); err != nil { - return fmt.Errorf("writing %s: %w", path, err) - } - return nil -} diff --git a/arrgen/arrgen/doc.go b/arrgen/arrgen/doc.go deleted file mode 100644 index 682a70f79..000000000 --- a/arrgen/arrgen/doc.go +++ /dev/null @@ -1,51 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package arrgen generates zero-reflection Arrow encoders for Go struct types. -// -// It is the code-generation counterpart to -// [github.com/apache/arrow-go/v18/arrow/array/arreflect]: arreflect interprets -// a struct's `arrow:"..."` tags with reflection on every value, while arrgen -// reads the same tags once, at generate time, and emits typed Go source that -// appends struct fields into typed Arrow builders. For every struct arrgen -// accepts, the generated schema equals arreflect.InferSchema and the generated -// record batch equals the one arreflect.RecordFromSlice builds. -// -// # Usage -// -// Put a go:generate directive next to the struct: -// -// //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -// -// then run `go generate ./...`, which writes metric_arrow.go next to the type. -// The emitted code imports only arrow, arrow/array and arrow/memory, never -// arrgen itself. -// -// # Supported field types -// -// bool, int8/16/32/64, int, uint8/16/32/64, uint, float32/64, string, []byte, -// time.Time, time.Duration, decimal.Decimal32, decimal.Decimal64, -// decimal128.Num and decimal256.Num, plus a pointer to any of those for a -// nullable column. -// -// Tag options mirror arreflect: a leading name, "-" to skip a field, the -// temporal overrides date32, date64, time32, time64 and timestamp, dict, view, -// large, and decimal(precision,scale). -// -// Anything else - nested structs, slices other than []byte, arrays, maps, -// embedded fields, named scalar types - is a generate-time error naming the -// field, never a silently dropped column. -package arrgen diff --git a/arrgen/arrgen/drift_test.go b/arrgen/arrgen/drift_test.go deleted file mode 100644 index 6c891684f..000000000 --- a/arrgen/arrgen/drift_test.go +++ /dev/null @@ -1,66 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package arrgen_test - -import ( - "bytes" - "os" - "path/filepath" - "testing" - - "github.com/apache/arrow-go/arrgen" -) - -// TestCheckedInFilesAreUpToDate regenerates every committed arrgen output in -// this module and fails if the result differs from what is on disk. -func TestCheckedInFilesAreUpToDate(t *testing.T) { - header, err := os.ReadFile("license_header.txt") - if err != nil { - t.Fatalf("reading license header: %v", err) - } - - // Mirrors the go:generate directives next to each type. - targets := []struct { - dir string - types []string - file string - }{ - {"example", []string{"Metric"}, "metric_arrow.go"}, - {"internal/gentypes", []string{"Row", "Fixed"}, "row_arrow.go"}, - } - - for _, target := range targets { - t.Run(target.file, func(t *testing.T) { - got, err := arrgen.Generate(arrgen.Config{ - Dir: target.dir, - Types: target.types, - Header: string(header), - }) - if err != nil { - t.Fatalf("Generate: %v", err) - } - path := filepath.Join(target.dir, target.file) - want, err := os.ReadFile(path) - if err != nil { - t.Fatalf("reading %s: %v", path, err) - } - if !bytes.Equal(got, want) { - t.Errorf("%s is out of date; run go generate ./...", path) - } - }) - } -} diff --git a/arrgen/arrgen/example/example_test.go b/arrgen/arrgen/example/example_test.go deleted file mode 100644 index 5c7cbcb8c..000000000 --- a/arrgen/arrgen/example/example_test.go +++ /dev/null @@ -1,105 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package example_test - -import ( - "fmt" - "time" - - "github.com/apache/arrow-go/arrgen/example" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// Example encodes a slice of rows in one call, replacing -// arreflect.RecordFromSlice[Metric](metrics, mem). -func Example() { - cpuLoad := 0.75 - metrics := []example.Metric{ - {Time: time.Unix(0, 0).UTC(), Host: "web-1", CPU: 0.5, Value: &cpuLoad}, - {Time: time.Unix(1, 0).UTC(), Host: "web-2", CPU: 1.5}, // Value is nil: a null - } - - rec, err := example.MetricRecordBatch(memory.DefaultAllocator, metrics) - if err != nil { - fmt.Println("encode:", err) - return - } - defer rec.Release() - - fmt.Println("rows:", rec.NumRows()) - for i, col := range rec.Columns() { - fmt.Printf("%s: %v\n", rec.Schema().Field(i).Name, col) - } - // Output: - // rows: 2 - // ts: [0 1000000000] - // host: ["web-1" "web-2"] - // cpu: [0.5 1.5] - // value: [0.75 (null)] -} - -// ExampleMetricAppender streams rows in one at a time and cuts a batch at a -// size boundary. The appender never retains a row, so the loop reuses a single -// variable and allocates nothing. -func ExampleMetricAppender() { - const batchSize = 2 - - a := example.NewMetricAppender(memory.DefaultAllocator) - defer a.Release() - a.Reserve(batchSize) - - var row example.Metric - for i := 0; i < 5; i++ { - row = example.Metric{ - Time: time.Unix(int64(i), 0).UTC(), - Host: fmt.Sprintf("web-%d", i), - CPU: float64(i), - } - a.Append(&row) - - if a.Len() == batchSize { - rec := a.NewRecordBatch() - fmt.Println("batch of", rec.NumRows()) - rec.Release() - a.Reserve(batchSize) - } - } - if a.Len() > 0 { - rec := a.NewRecordBatch() - fmt.Println("final batch of", rec.NumRows()) - rec.Release() - } - if err := a.Err(); err != nil { - fmt.Println("append:", err) - } - // Output: - // batch of 2 - // batch of 2 - // final batch of 1 -} - -// ExampleMetricSchema shows the schema the generator resolved from the tags. -func ExampleMetricSchema() { - for _, f := range example.MetricSchema().Fields() { - fmt.Printf("%s %s nullable=%t\n", f.Name, f.Type, f.Nullable) - } - // Output: - // ts timestamp[ns, tz=UTC] nullable=false - // host utf8 nullable=false - // cpu float64 nullable=false - // value float64 nullable=true -} diff --git a/arrgen/arrgen/example/metric_arrow.go b/arrgen/arrgen/example/metric_arrow.go deleted file mode 100644 index 1909fd3a5..000000000 --- a/arrgen/arrgen/example/metric_arrow.go +++ /dev/null @@ -1,126 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Code generated by arrgen. DO NOT EDIT. - -package example - -import ( - "github.com/apache/arrow-go/v18/arrow" - "github.com/apache/arrow-go/v18/arrow/array" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// metricArrowSchema is the Arrow schema of Metric, resolved from its arrow -// struct tags. -var metricArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, - {Name: "host", Type: arrow.BinaryTypes.String, Nullable: false}, - {Name: "cpu", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, - {Name: "value", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, -}, nil) - -// MetricSchema returns the Arrow schema encoded by MetricAppender. -func MetricSchema() *arrow.Schema { return metricArrowSchema } - -// MetricAppender converts Metric values into Arrow record batches -// without reflection. -// -// It is not safe for concurrent use. Release it when done. -type MetricAppender struct { - rb *array.RecordBuilder - b0 *array.TimestampBuilder - b1 *array.StringBuilder - b2 *array.Float64Builder - b3 *array.Float64Builder - err error -} - -// NewMetricAppender returns an appender that builds Metric batches with mem, -// or memory.DefaultAllocator when mem is nil. The typed field builders are -// resolved once here, so Append does no type assertions. -func NewMetricAppender(mem memory.Allocator) *MetricAppender { - if mem == nil { - mem = memory.DefaultAllocator - } - rb := array.NewRecordBuilder(mem, metricArrowSchema) - return &MetricAppender{ - rb: rb, - b0: rb.Field(0).(*array.TimestampBuilder), - b1: rb.Field(1).(*array.StringBuilder), - b2: rb.Field(2).(*array.Float64Builder), - b3: rb.Field(3).(*array.Float64Builder), - } -} - -// Schema returns the schema of the batches this appender builds. -func (a *MetricAppender) Schema() *arrow.Schema { return metricArrowSchema } - -// RecordBuilder returns the underlying builder. Appending through it directly -// is allowed as long as every column stays the same length. -func (a *MetricAppender) RecordBuilder() *array.RecordBuilder { return a.rb } - -// Reserve pre-allocates room for n more rows in every column. Reserving before -// a run of Appends keeps the append path free of allocations. -func (a *MetricAppender) Reserve(n int) { a.rb.Reserve(n) } - -// Len reports the number of rows appended since the last NewRecordBatch. -func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } - -// Append appends one row. v is read synchronously and never retained, so a -// caller can reuse a single Metric variable across calls. -func (a *MetricAppender) Append(v *Metric) { - a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) - a.b1.Append(v.Host) - a.b2.Append(v.CPU) - if v.Value == nil { - a.b3.AppendNull() - } else { - a.b3.Append(*v.Value) - } -} - -// AppendSlice appends every element of vs, reserving room for them up front. -func (a *MetricAppender) AppendSlice(vs []Metric) { - a.rb.Reserve(len(vs)) - for i := range vs { - a.Append(&vs[i]) - } -} - -// NewRecordBatch returns the appended rows as a record batch and resets the -// appender. The caller must Release the batch. -func (a *MetricAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } - -// Err returns the first error an Append hit, or nil. Only dictionary columns -// can fail. The error is sticky: NewRecordBatch does not clear it. -func (a *MetricAppender) Err() error { return a.err } - -// Release releases the appender's builders. -func (a *MetricAppender) Release() { a.rb.Release() } - -// MetricRecordBatch builds one record batch holding every element of vs. It is the -// generated equivalent of arreflect.RecordFromSlice[Metric](vs, mem). The -// caller must Release the batch. -func MetricRecordBatch(mem memory.Allocator, vs []Metric) (arrow.RecordBatch, error) { - a := NewMetricAppender(mem) - defer a.Release() - a.AppendSlice(vs) - if err := a.Err(); err != nil { - return nil, err - } - return a.NewRecordBatch(), nil -} diff --git a/arrgen/arrgen/generate.go b/arrgen/arrgen/generate.go deleted file mode 100644 index 91e63de6d..000000000 --- a/arrgen/arrgen/generate.go +++ /dev/null @@ -1,335 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package arrgen - -import ( - "bytes" - "fmt" - "go/format" - "go/token" - "go/types" - "reflect" - "strings" - "unicode" - - "golang.org/x/tools/go/packages" -) - -// Config describes one generator run. -type Config struct { - // Dir is the directory of the package holding the struct types. go:generate - // runs in the package directory, so the command leaves this as ".". - Dir string - // Types are the struct type names to generate for, in output order. - Types []string - // Header, when non-empty, is emitted verbatim above the generated-code - // marker. Use it for a license header; it is expected to be Go comments. - Header string -} - -// Generate resolves cfg.Types in the package at cfg.Dir and returns the -// formatted source of their Arrow appenders. It does not write to disk. -// -// Output depends on nothing but the package's source, so regenerating an -// unchanged package reproduces the file byte for byte. -func Generate(cfg Config) ([]byte, error) { - if len(cfg.Types) == 0 { - return nil, fmt.Errorf("arrgen: no types requested") - } - dir := cfg.Dir - if dir == "" { - dir = "." - } - pkg, loadErrs, err := loadPackage(dir) - if err != nil { - return nil, err - } - - file := genFile{Header: strings.TrimRight(cfg.Header, "\n"), Package: pkg.Name} - declared := make(map[string]string) // generated identifier -> type that claimed it - for _, name := range cfg.Types { - gt, err := genForType(pkg, name) - if err != nil { - return nil, withLoadContext(fmt.Errorf("arrgen: %s: %w", name, err), loadErrs) - } - if err := claimIdents(declared, name, gt); err != nil { - return nil, fmt.Errorf("arrgen: %s: %w", name, err) - } - file.NeedTime = file.NeedTime || gt.needTime - file.Types = append(file.Types, gt) - } - - var buf bytes.Buffer - if err := fileTemplate.Execute(&buf, file); err != nil { - return nil, fmt.Errorf("arrgen: rendering template: %w", err) - } - src, err := format.Source(buf.Bytes()) - if err != nil { - return nil, fmt.Errorf("arrgen: formatting generated source: %w\n--- raw ---\n%s", err, buf.String()) - } - return src, nil -} - -// loadPackage type-checks the package in dir and returns it along with any -// errors the type checker reported. Those errors are returned rather than -// raised: the package a generator runs in usually does not compile yet, since -// the call sites are written before the generated file exists. -func loadPackage(dir string) (*packages.Package, []string, error) { - cfg := &packages.Config{ - Mode: packages.NeedName | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax, - Dir: dir, - } - pkgs, err := packages.Load(cfg, ".") - if err != nil { - return nil, nil, fmt.Errorf("arrgen: loading package in %s: %w", dir, err) - } - if len(pkgs) != 1 { - return nil, nil, fmt.Errorf("arrgen: expected exactly one package in %s, found %d", dir, len(pkgs)) - } - var errs []string - for _, e := range pkgs[0].Errors { - errs = append(errs, e.Error()) - } - if pkgs[0].Types == nil { - return nil, nil, fmt.Errorf("arrgen: package in %s could not be type-checked: %s", dir, strings.Join(errs, "; ")) - } - return pkgs[0], errs, nil -} - -// withLoadContext appends the package's type errors to err as a likely cause. -func withLoadContext(err error, loadErrs []string) error { - if len(loadErrs) == 0 { - return err - } - const maxShown = 5 - shown := loadErrs - suffix := "" - if len(shown) > maxShown { - shown, suffix = shown[:maxShown], fmt.Sprintf("\n ... and %d more", len(loadErrs)-maxShown) - } - return fmt.Errorf("%w\nthe package does not type-check, which may be the cause:\n %s%s", - err, strings.Join(shown, "\n "), suffix) -} - -func genForType(pkg *packages.Package, name string) (genType, error) { - st, err := findStruct(pkg, name) - if err != nil { - return genType{}, err - } - cols, err := collectColumns(st) - if err != nil { - return genType{}, err - } - if len(cols) == 0 { - return genType{}, fmt.Errorf("struct has no Arrow columns") - } - - exported := token.IsExported(name) - gt := genType{ - GoName: name, - SchemaVar: lowerFirst(name) + "ArrowSchema", - SchemaFunc: caseAs(exported, name+"Schema"), - AppenderType: caseAs(exported, name+"Appender"), - CtorName: caseAs(exported, "New"+upperFirst(name)+"Appender"), - BatchFunc: caseAs(exported, name+"RecordBatch"), - } - for i, c := range cols { - gt.needTime = gt.needTime || c.spec.needsTime - gt.AnyFallible = gt.AnyFallible || c.spec.fallible - gt.Fields = append(gt.Fields, genField{ - Index: i, - Name: c.name, - GoField: c.goField, - Nullable: c.nullable, - ArrowType: c.spec.arrowType, - BuilderType: c.spec.builderType, - BuilderVar: fmt.Sprintf("b%d", i), - AppendStmt: renderAppend(i, c), - }) - } - return gt, nil -} - -// claimIdents records the package-level names a type's generated code will -// declare, rejecting a second type that would declare the same one. Two types -// differing only in the case of their first letter collide this way. -func claimIdents(declared map[string]string, name string, gt genType) error { - for _, ident := range []string{gt.SchemaVar, gt.SchemaFunc, gt.AppenderType, gt.CtorName, gt.BatchFunc} { - if prev, taken := declared[ident]; taken { - return fmt.Errorf("generated name %s collides with the one generated for %s; generate the two types into separate files", ident, prev) - } - declared[ident] = name - } - return nil -} - -func findStruct(pkg *packages.Package, name string) (*types.Struct, error) { - obj := pkg.Types.Scope().Lookup(name) - if obj == nil { - return nil, fmt.Errorf("type not found in package %s", pkg.Name) - } - named, ok := obj.Type().(*types.Named) - if !ok { - return nil, fmt.Errorf("not a named type") - } - st, ok := named.Underlying().(*types.Struct) - if !ok { - return nil, fmt.Errorf("not a struct type") - } - return st, nil -} - -// column is one resolved struct field. -type column struct { - name string - goField string - nullable bool - spec colSpec -} - -// collectColumns walks the struct's fields in declaration order, which is the -// order arreflect settles on for a flat struct. Embedded fields are rejected -// rather than promoted, since arreflect resolves promoted names by breadth and -// tag. -func collectColumns(st *types.Struct) ([]column, error) { - var cols []column - seen := make(map[string]string, st.NumFields()) - - for i := 0; i < st.NumFields(); i++ { - f := st.Field(i) - if f.Anonymous() { - return nil, fmt.Errorf("field %s: embedded fields are not supported; give the field a name or encode the struct with arreflect", f.Name()) - } - if !f.Exported() { - continue - } - - tag, hasTag := reflect.StructTag(st.Tag(i)).Lookup("arrow") - var opts tagOpts - if hasTag { - var err error - if opts, err = parseTag(tag); err != nil { - return nil, fmt.Errorf("field %s: %w", f.Name(), err) - } - } - if opts.Skip { - continue - } - if err := opts.validate(); err != nil { - return nil, fmt.Errorf("field %s: %w", f.Name(), err) - } - - name := opts.Name - if name == "" { - name = f.Name() - } - if prev, dup := seen[name]; dup { - return nil, fmt.Errorf("fields %s and %s both map to column %q", prev, f.Name(), name) - } - seen[name] = f.Name() - - spec, nullable, err := resolveColumn(f.Type(), opts) - if err != nil { - return nil, fmt.Errorf("field %s: %w", f.Name(), err) - } - cols = append(cols, column{name: name, goField: f.Name(), nullable: nullable, spec: spec}) - } - return cols, nil -} - -// renderAppend emits the append statement for one column, wrapping it in the -// null guards arreflect applies: a nil pointer is null, and so is a nil []byte -// even when the column itself is not nullable. -func renderAppend(idx int, c column) string { - bld := fmt.Sprintf("a.b%d", idx) - val, recv := "v."+c.goField, "v."+c.goField - if c.nullable { - val, recv = "*v."+c.goField, "(*v."+c.goField+")" - } - - stmt := c.spec.appendStmt(bld, recv, val) - if c.spec.fallible { - stmt = "a.setErr(" + stmt + ")" - } - - var guard string - switch { - case c.nullable && c.spec.nilable: - guard = fmt.Sprintf("v.%s == nil || *v.%s == nil", c.goField, c.goField) - case c.nullable, c.spec.nilable: - guard = fmt.Sprintf("v.%s == nil", c.goField) - default: - return stmt - } - - // A statement that is already a block (a time-of-day column scopes a local) - // sheds its braces so the guard's else does not nest two sets. - inner := strings.TrimSuffix(strings.TrimPrefix(stmt, "{\n"), "\n}") - return fmt.Sprintf("if %s {\n%s.AppendNull()\n} else {\n%s\n}", guard, bld, inner) -} - -type genFile struct { - Header string - Package string - NeedTime bool - Types []genType -} - -type genType struct { - GoName string - SchemaVar string - SchemaFunc string - AppenderType string - CtorName string - BatchFunc string - Fields []genField - AnyFallible bool - - needTime bool -} - -type genField struct { - Index int - Name string - GoField string - Nullable bool - ArrowType string - BuilderType string - BuilderVar string - AppendStmt string -} - -func lowerFirst(s string) string { return mapFirst(s, unicode.ToLower) } -func upperFirst(s string) string { return mapFirst(s, unicode.ToUpper) } - -func mapFirst(s string, f func(rune) rune) string { - if s == "" { - return s - } - r := []rune(s) - r[0] = f(r[0]) - return string(r) -} - -// caseAs gives a generated identifier the same visibility as the struct it was -// derived from. -func caseAs(exported bool, s string) string { - if exported { - return upperFirst(s) - } - return lowerFirst(s) -} diff --git a/arrgen/arrgen/generate_test.go b/arrgen/arrgen/generate_test.go deleted file mode 100644 index a864013a1..000000000 --- a/arrgen/arrgen/generate_test.go +++ /dev/null @@ -1,222 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package arrgen_test - -import ( - "bytes" - "flag" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/apache/arrow-go/arrgen" -) - -var update = flag.Bool("update", false, "rewrite the golden files instead of comparing against them") - -const goldenPath = "testdata/basic.golden" - -// TestGenerateGolden compares the whole emitted file against a checked-in copy, -// so the diff of a generator change shows what users would see in their own -// regenerated code. -func TestGenerateGolden(t *testing.T) { - header, err := os.ReadFile("license_header.txt") - if err != nil { - t.Fatalf("reading license header: %v", err) - } - got, err := arrgen.Generate(arrgen.Config{ - Dir: "testdata/basic", - Types: []string{"Metric", "reading"}, - Header: string(header), - }) - if err != nil { - t.Fatalf("Generate: %v", err) - } - - if *update { - if err := os.WriteFile(goldenPath, got, 0o644); err != nil { - t.Fatalf("writing golden: %v", err) - } - return - } - - want, err := os.ReadFile(goldenPath) - if err != nil { - t.Fatalf("reading golden: %v (run: go test ./... -update)", err) - } - if !bytes.Equal(got, want) { - t.Errorf("generated output differs from %s; rerun with -update to accept\n--- got ---\n%s", goldenPath, got) - } -} - -// TestGenerateIsDeterministic checks that the same package in produces the same -// bytes out, which TestCheckedInFilesAreUpToDate depends on. -func TestGenerateIsDeterministic(t *testing.T) { - cfg := arrgen.Config{Dir: "testdata/basic", Types: []string{"Metric", "reading"}} - first, err := arrgen.Generate(cfg) - if err != nil { - t.Fatalf("Generate: %v", err) - } - second, err := arrgen.Generate(cfg) - if err != nil { - t.Fatalf("Generate: %v", err) - } - if !bytes.Equal(first, second) { - t.Error("two runs over the same package produced different output") - } -} - -// TestGenerateHeaderOptional checks that the generated-code marker leads the -// file when no header is configured. -func TestGenerateHeaderOptional(t *testing.T) { - got, err := arrgen.Generate(arrgen.Config{Dir: "testdata/basic", Types: []string{"Metric"}}) - if err != nil { - t.Fatalf("Generate: %v", err) - } - const marker = "// Code generated by arrgen. DO NOT EDIT." - if !bytes.HasPrefix(got, []byte(marker)) { - t.Errorf("output does not start with %q:\n%s", marker, first(got, 200)) - } -} - -// TestGenerateRejects covers the structs arrgen refuses rather than encode -// wrongly or drop a column from. -func TestGenerateRejects(t *testing.T) { - tests := []struct { - typ string - want string - }{ - {"Embedded", "embedded fields are not supported"}, - {"Nested", "nested struct column"}, - {"SliceField", "list column for []int64 is not supported"}, - {"ArrayField", "is not supported"}, - {"MapField", "is not supported"}, - {"NamedScalar", "named type"}, - {"UnknownOption", `unknown option "nope"`}, - {"BadDecimalTag", "is not an integer"}, - {"ShortDecimalTag", "expected decimal(precision,scale)"}, - {"DuplicateNames", `both map to column "same"`}, - {"TemporalOnInt", "only valid on a time.Time field"}, - {"DecimalOnInt", "only valid on a decimal field"}, - {"DictOnBool", "dict is not supported on bool"}, - {"ViewOnInt", "view has no effect"}, - {"LargeOnInt", "large has no effect"}, - {"RunEndEncoded", "ree is not supported"}, - {"DictAndView", "at most one of dict, view, ree"}, - {"DictAndLarge", "cannot be combined with large"}, - {"NoColumns", "no Arrow columns"}, - {"DoublePointer", "multi-level pointer"}, - {"TimeSlice", "list column"}, - {"NotAStruct", "not a struct type"}, - {"Absent", "type not found"}, - } - - for _, tt := range tests { - t.Run(tt.typ, func(t *testing.T) { - _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{tt.typ}}) - if err == nil { - t.Fatalf("Generate(%s) succeeded, want an error mentioning %q", tt.typ, tt.want) - } - if !strings.Contains(err.Error(), tt.want) { - t.Errorf("Generate(%s) error = %q, want it to mention %q", tt.typ, err, tt.want) - } - if !strings.Contains(err.Error(), tt.typ) { - t.Errorf("Generate(%s) error = %q, want it to name the type", tt.typ, err) - } - }) - } -} - -// TestGenerateNamesTheField checks that a rejection points at the field, not -// just the struct. -func TestGenerateNamesTheField(t *testing.T) { - _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Nested"}}) - if err == nil { - t.Fatal("Generate succeeded, want an error") - } - if !strings.Contains(err.Error(), "field Field") { - t.Errorf("error = %q, want it to name the offending field", err) - } -} - -// TestGenerateRejectsCollidingTypes covers two types in one file whose -// generated names would clash. -func TestGenerateRejectsCollidingTypes(t *testing.T) { - _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Collide", "collide"}}) - if err == nil { - t.Fatal("Generate succeeded, want a collision error") - } - if !strings.Contains(err.Error(), "collides with") { - t.Errorf("error = %q, want it to report the collision", err) - } -} - -func TestGenerateRejectsRepeatedType(t *testing.T) { - _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/errors", Types: []string{"Collide", "Collide"}}) - if err == nil { - t.Fatal("Generate succeeded, want a collision error") - } -} - -// TestGenerateBeforePackageCompiles covers the chicken-and-egg case: the call -// sites are written first, so the package does not type-check until the -// generator has run. -func TestGenerateBeforePackageCompiles(t *testing.T) { - got, err := arrgen.Generate(arrgen.Config{Dir: "testdata/pregen", Types: []string{"Sample"}}) - if err != nil { - t.Fatalf("Generate: %v", err) - } - if !bytes.Contains(got, []byte("func SampleRecordBatch(")) { - t.Error("generated output does not declare SampleRecordBatch") - } -} - -// TestGenerateReportsLoadErrors checks that a type which cannot be resolved -// reports the package's own compile errors as the likely cause. -func TestGenerateReportsLoadErrors(t *testing.T) { - _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/pregen", Types: []string{"Absent"}}) - if err == nil { - t.Fatal("Generate succeeded, want an error") - } - if !strings.Contains(err.Error(), "does not type-check") { - t.Errorf("error = %q, want it to surface the package's compile errors", err) - } -} - -func TestGenerateNoTypes(t *testing.T) { - if _, err := arrgen.Generate(arrgen.Config{Dir: "testdata/basic"}); err == nil { - t.Error("Generate with no types succeeded, want an error") - } -} - -func TestGenerateMissingPackage(t *testing.T) { - dir := filepath.Join(t.TempDir(), "empty") - if err := os.Mkdir(dir, 0o755); err != nil { - t.Fatal(err) - } - if _, err := arrgen.Generate(arrgen.Config{Dir: dir, Types: []string{"Metric"}}); err == nil { - t.Error("Generate in a directory with no Go files succeeded, want an error") - } -} - -func first(b []byte, n int) []byte { - if len(b) < n { - return b - } - return b[:n] -} diff --git a/arrgen/arrgen/go.mod b/arrgen/arrgen/go.mod deleted file mode 100644 index ad12f8825..000000000 --- a/arrgen/arrgen/go.mod +++ /dev/null @@ -1,35 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -module github.com/apache/arrow-go/arrgen - -go 1.25.0 - -require ( - github.com/apache/arrow-go/v18 v18.7.0 - golang.org/x/tools v0.45.0 -) - -require ( - github.com/goccy/go-json v0.10.6 // indirect - github.com/google/flatbuffers v25.12.19+incompatible // indirect - github.com/klauspost/cpuid/v2 v2.4.0 // indirect - github.com/zeebo/xxh3 v1.1.0 // indirect - golang.org/x/exp v0.0.0-20260112195511-716be5621a96 // indirect - golang.org/x/mod v0.36.0 // indirect - golang.org/x/sync v0.22.0 // indirect - golang.org/x/sys v0.47.0 // indirect -) diff --git a/arrgen/arrgen/go.sum b/arrgen/arrgen/go.sum deleted file mode 100644 index 963f34aaf..000000000 --- a/arrgen/arrgen/go.sum +++ /dev/null @@ -1,44 +0,0 @@ -github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM= -github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= -github.com/apache/arrow-go/v18 v18.7.0 h1:Vw/i+cJyebUofT7JlqFpe65LrmwxULn166jjwStM4HY= -github.com/apache/arrow-go/v18 v18.7.0/go.mod h1:PM6IigLJkdMwIpeHXnymo+xZ52f42a9EYiLtRel4p/A= -github.com/apache/thrift v0.24.0 h1:zy31L1a49QTNB2bG1BBfMXol3yJrTH975G3pPubQVLQ= -github.com/apache/thrift v0.24.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= -github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= -github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU= -github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= -github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= -github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= -github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= -github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= -github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= -github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/klauspost/compress v1.19.0 h1:sXLILfc9jV2QYWkzFOPWStmcUVH2RHEB1JCdY2oVvCQ= -github.com/klauspost/compress v1.19.0/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= -github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw= -github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU= -github.com/pierrec/lz4/v4 v4.1.27 h1:+PhzhWDrjRj89TH2sw43nE3+4+W8lSxIuQadEHZyjUk= -github.com/pierrec/lz4/v4 v4.1.27/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= -github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= -github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= -github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= -github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= -github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= -github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= -golang.org/x/exp v0.0.0-20260112195511-716be5621a96 h1:Z/6YuSHTLOHfNFdb8zVZomZr7cqNgTJvA8+Qz75D8gU= -golang.org/x/exp v0.0.0-20260112195511-716be5621a96/go.mod h1:nzimsREAkjBCIEFtHiYkrJyT+2uy9YZJB7H1k68CXZU= -golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= -golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ= -golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= -golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= -golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= -golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0= -gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= -gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= -gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= -gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/arrgen/arrgen/internal/gentypes/alloc_test.go b/arrgen/arrgen/internal/gentypes/alloc_test.go deleted file mode 100644 index 010172a09..000000000 --- a/arrgen/arrgen/internal/gentypes/alloc_test.go +++ /dev/null @@ -1,158 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package gentypes_test - -import ( - "testing" - - "github.com/apache/arrow-go/arrgen/internal/gentypes" - "github.com/apache/arrow-go/v18/arrow/array/arreflect" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// TestStreamingAppendIsAllocationFree asserts that with room reserved up front, -// appending a row of fixed-width columns allocates nothing. -func TestStreamingAppendIsAllocationFree(t *testing.T) { - skipIfInstrumented(t) - const runs = 2000 - - a := gentypes.NewFixedAppender(memory.DefaultAllocator) - defer a.Release() - a.Reserve(runs + 8) // AllocsPerRun calls the function runs+1 times - - rows := makeFixed(4) - i := 0 - // A single reused variable, as a caller draining a stream would do it. - var row gentypes.Fixed - got := testing.AllocsPerRun(runs, func() { - row = rows[i%len(rows)] - i++ - a.Append(&row) - }) - if got != 0 { - t.Errorf("Append allocated %.2f times per row, want 0", got) - } - a.NewRecordBatch().Release() -} - -// TestSchemaIsAllocationFree checks that the generated schema is a package-level -// variable, where arreflect rebuilds one on every call. -func TestSchemaIsAllocationFree(t *testing.T) { - skipIfInstrumented(t) - if got := testing.AllocsPerRun(1000, func() { _ = gentypes.RowSchema() }); got != 0 { - t.Errorf("RowSchema allocated %.2f times per call, want 0", got) - } - - reflected := testing.AllocsPerRun(100, func() { - if _, err := arreflect.InferSchema[gentypes.Row](); err != nil { - t.Fatal(err) - } - }) - if reflected == 0 { - t.Skip("arreflect.InferSchema no longer allocates; the comparison is moot") - } - t.Logf("schema per call: arrgen 0 allocs, arreflect %.0f allocs", reflected) -} - -// TestBatchEncodingAllocatesLessThanReflection compares whole-batch encoding. -// The threshold is loose: both paths pay for the same Arrow buffers, and the -// point is to catch a regression that erases the advantage, not to pin a number. -func TestBatchEncodingAllocatesLessThanReflection(t *testing.T) { - skipIfInstrumented(t) - const rows = 1024 - fixtures := makeFixed(rows) - mem := memory.DefaultAllocator - - reflected := testing.Benchmark(func(b *testing.B) { - for b.Loop() { - rec, err := arreflect.RecordFromSlice(fixtures, mem) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - }) - generated := testing.Benchmark(func(b *testing.B) { - for b.Loop() { - rec, err := gentypes.FixedRecordBatch(mem, fixtures) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - }) - - t.Logf("%d-row batch: arrgen %d allocs / %d ns, arreflect %d allocs / %d ns", - rows, generated.AllocsPerOp(), generated.NsPerOp(), reflected.AllocsPerOp(), reflected.NsPerOp()) - - if generated.AllocsPerOp() >= reflected.AllocsPerOp() { - t.Errorf("generated encoder allocated %d times per batch, want fewer than arreflect's %d", - generated.AllocsPerOp(), reflected.AllocsPerOp()) - } - if got, limit := generated.AllocsPerOp(), reflected.AllocsPerOp()*3/4; got > limit { - t.Errorf("generated encoder allocated %d times per batch, want at most %d (three quarters of arreflect's %d)", - got, limit, reflected.AllocsPerOp()) - } -} - -// TestBatchEncodingIsFasterThanReflection guards the speed claim. Wall-clock -// assertions are noisy on shared CI, so the bar sits well below the ~3.5x -// actually measured on the fixed-width fixture. -func TestBatchEncodingIsFasterThanReflection(t *testing.T) { - skipIfInstrumented(t) - if testing.Short() { - t.Skip("timing comparison skipped in short mode") - } - const rows = 1024 - fixtures := makeFixed(rows) - mem := memory.DefaultAllocator - - reflected := testing.Benchmark(func(b *testing.B) { - for b.Loop() { - rec, err := arreflect.RecordFromSlice(fixtures, mem) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - }) - generated := testing.Benchmark(func(b *testing.B) { - for b.Loop() { - rec, err := gentypes.FixedRecordBatch(mem, fixtures) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - }) - if generated.NsPerOp() == 0 || reflected.NsPerOp() == 0 { - t.Skip("benchmark produced no timing data") - } - - speedup := float64(reflected.NsPerOp()) / float64(generated.NsPerOp()) - t.Logf("%d-row batch: arrgen %d ns, arreflect %d ns (%.2fx)", rows, generated.NsPerOp(), reflected.NsPerOp(), speedup) - if speedup < 1.5 { - t.Errorf("generated encoder is only %.2fx faster than arreflect, want at least 1.5x", speedup) - } -} - -func skipIfInstrumented(t *testing.T) { - t.Helper() - if instrumented { - t.Skip("allocation counts and timings are not meaningful under -race or -asan") - } -} diff --git a/arrgen/arrgen/internal/gentypes/bench_test.go b/arrgen/arrgen/internal/gentypes/bench_test.go deleted file mode 100644 index acf2634c3..000000000 --- a/arrgen/arrgen/internal/gentypes/bench_test.go +++ /dev/null @@ -1,219 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package gentypes_test - -import ( - "fmt" - "testing" - "time" - - "github.com/apache/arrow-go/arrgen/example" - "github.com/apache/arrow-go/arrgen/internal/gentypes" - "github.com/apache/arrow-go/v18/arrow/array/arreflect" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// benchRows is the batch size the encoder benchmarks build: large enough that -// per-batch setup does not dominate, small enough to stay in cache. -const benchRows = 1024 - -func makeMetrics(n int) []example.Metric { - metrics := make([]example.Metric, n) - v := 42.5 - for i := range metrics { - metrics[i] = example.Metric{ - Time: base.Add(time.Duration(i) * time.Second), - Host: fmt.Sprintf("host-%d", i%16), - CPU: float64(i) / 100, - } - if i%4 != 0 { - metrics[i].Value = &v - } - } - return metrics -} - -func makeFixed(n int) []gentypes.Fixed { - rows := make([]gentypes.Fixed, n) - v := 1.5 - for i := range rows { - rows[i] = gentypes.Fixed{ - Timestamp: base.Add(time.Duration(i) * time.Second), - ID: int64(i), - Value: float64(i) * 1.5, - OK: i%2 == 0, - } - if i%3 != 0 { - rows[i].Optional = &v - } - } - return rows -} - -// BenchmarkMetricBatch encodes the same slice of rows into a record batch by -// the reflection path and by the generated one. -func BenchmarkMetricBatch(b *testing.B) { - metrics := makeMetrics(benchRows) - mem := memory.DefaultAllocator - - b.Run("arreflect", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - rec, err := arreflect.RecordFromSlice(metrics, mem) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") - }) - - b.Run("arrgen", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - rec, err := example.MetricRecordBatch(mem, metrics) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") - }) -} - -// BenchmarkRowBatch runs the same comparison over the wide fixture: 46 columns -// per row instead of four. -func BenchmarkRowBatch(b *testing.B) { - rows := makeRows(benchRows) - mem := memory.DefaultAllocator - - b.Run("arreflect", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - rec, err := arreflect.RecordFromSlice(rows, mem) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") - }) - - b.Run("arrgen", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - rec, err := gentypes.RowRecordBatch(mem, rows) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") - }) -} - -// BenchmarkFixedBatch isolates the encoder from Arrow's variable-width buffers: -// every column is fixed width, so the allocations are the encoding strategy's -// own rather than a growing data buffer. -func BenchmarkFixedBatch(b *testing.B) { - rows := makeFixed(benchRows) - mem := memory.DefaultAllocator - - b.Run("arreflect", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - rec, err := arreflect.RecordFromSlice(rows, mem) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") - }) - - b.Run("arrgen", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - rec, err := gentypes.FixedRecordBatch(mem, rows) - if err != nil { - b.Fatal(err) - } - rec.Release() - } - b.ReportMetric(float64(benchRows)*float64(b.N)/b.Elapsed().Seconds()/1e6, "Mrows/s") - }) -} - -// BenchmarkStreamAppend measures the per-row cost of the streaming entry point. -// The reflection path has no single-row API, so its column is the amortized -// cost of encoding a batch of the same rows. -func BenchmarkStreamAppend(b *testing.B) { - rows := makeFixed(benchRows) - mem := memory.DefaultAllocator - - b.Run("arreflect/batched", func(b *testing.B) { - b.ReportAllocs() - n := 0 - for b.Loop() { - rec, err := arreflect.RecordFromSlice(rows, mem) - if err != nil { - b.Fatal(err) - } - rec.Release() - n += benchRows - } - b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(n), "ns/row") - }) - - b.Run("arrgen/streamed", func(b *testing.B) { - a := gentypes.NewFixedAppender(mem) - defer a.Release() - b.ReportAllocs() - var row gentypes.Fixed - i := 0 - for b.Loop() { - row = rows[i%benchRows] - i++ - if i%benchRows == 0 { - a.NewRecordBatch().Release() - a.Reserve(benchRows) - } - a.Append(&row) - } - a.NewRecordBatch().Release() - b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(b.N), "ns/row") - }) -} - -// BenchmarkSchema compares arreflect's struct walk against the generated -// package-level schema variable. -func BenchmarkSchema(b *testing.B) { - b.Run("arreflect", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - if _, err := arreflect.InferSchema[gentypes.Row](); err != nil { - b.Fatal(err) - } - } - }) - b.Run("arrgen", func(b *testing.B) { - b.ReportAllocs() - for b.Loop() { - _ = gentypes.RowSchema() - } - }) -} diff --git a/arrgen/arrgen/internal/gentypes/equivalence_test.go b/arrgen/arrgen/internal/gentypes/equivalence_test.go deleted file mode 100644 index e99edb56f..000000000 --- a/arrgen/arrgen/internal/gentypes/equivalence_test.go +++ /dev/null @@ -1,449 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package gentypes_test - -import ( - "fmt" - "math" - "testing" - "time" - - "github.com/apache/arrow-go/arrgen/example" - "github.com/apache/arrow-go/arrgen/internal/gentypes" - "github.com/apache/arrow-go/v18/arrow" - "github.com/apache/arrow-go/v18/arrow/array" - "github.com/apache/arrow-go/v18/arrow/array/arreflect" - "github.com/apache/arrow-go/v18/arrow/decimal" - "github.com/apache/arrow-go/v18/arrow/decimal128" - "github.com/apache/arrow-go/v18/arrow/decimal256" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// arreflectStructFieldGaps names the Row columns arreflect cannot encode as a -// struct field, and the Go type of each. -// -// Its inferArrowType dispatches on reflect.Kind before checking the types it -// special-cases, so a Go struct that Arrow models as a scalar reaches -// inferStructType, which finds only unexported fields and returns an empty -// struct<>. A tag naming an explicit Arrow type rescues the field, so only the -// untagged spellings are listed. These columns are compared by value in -// TestGapColumnValues; see also README.md. -var arreflectStructFieldGaps = map[string]string{ - "ts": "time.Time", - "pts": "*time.Time", - "dec256": "decimal256.Num", - "pdec128": "*decimal128.Num", -} - -// compareColumns asserts that got and want hold the same columns, skipping the -// ones arreflect cannot encode as struct fields. -// -// dictsByValue relaxes dictionary columns to a comparison of the values they -// decode to: an appender reused across batches carries its memo table forward, -// so its later batches encode the same values against a longer dictionary. -func compareColumns(t *testing.T, label string, got, want arrow.RecordBatch, dictsByValue bool) { - t.Helper() - if got.NumRows() != want.NumRows() { - t.Fatalf("%s: NumRows = %d, want %d", label, got.NumRows(), want.NumRows()) - } - if got.NumCols() != want.NumCols() { - t.Fatalf("%s: NumCols = %d, want %d", label, got.NumCols(), want.NumCols()) - } - for i := 0; i < int(got.NumCols()); i++ { - name := got.Schema().Field(i).Name - if _, gap := arreflectStructFieldGaps[name]; gap { - continue - } - gc, wc := got.Column(i), want.Column(i) - if dictsByValue && gc.DataType().ID() == arrow.DICTIONARY { - compareByValueStr(t, label, name, gc, wc) - continue - } - if !array.Equal(gc, wc) { - t.Errorf("%s: column %q differs\n got: %s\nwant: %s", label, name, gc, wc) - } - } -} - -// compareByValueStr compares two arrays by the value each row decodes to, -// ignoring how that value is physically encoded. -func compareByValueStr(t *testing.T, label, name string, got, want arrow.Array) { - t.Helper() - if got.Len() != want.Len() { - t.Errorf("%s: column %q: len = %d, want %d", label, name, got.Len(), want.Len()) - return - } - for i := 0; i < got.Len(); i++ { - if got.IsNull(i) != want.IsNull(i) { - t.Errorf("%s: column %q row %d: null = %t, want %t", label, name, i, got.IsNull(i), want.IsNull(i)) - return - } - if got.IsNull(i) { - continue - } - if got.ValueStr(i) != want.ValueStr(i) { - t.Errorf("%s: column %q row %d = %s, want %s", label, name, i, got.ValueStr(i), want.ValueStr(i)) - return - } - } -} - -// base is a fixed instant with a non-zero time of day and sub-second part, so -// the date32/date64/time32/time64 columns each exercise a different truncation. -var base = time.Date(2024, time.March, 17, 13, 45, 12, 123456789, time.UTC) - -// makeRows builds n deterministic rows. Every third row nulls its pointer -// fields and empties its variable-width ones, so nulls are covered too. -func makeRows(n int) []gentypes.Row { - rows := make([]gentypes.Row, n) - for i := range rows { - null := i%3 == 0 - r := gentypes.Row{ - Bool: i%2 == 0, - Int8: int8(i), - Int16: int16(i * 3), - Int32: int32(i * 7), - Int64: int64(i) * 1e6, - Int: i, - Uint8: uint8(i), - Uint16: uint16(i * 5), - Uint32: uint32(i * 11), - Uint64: uint64(i) * 1e9, - Uint: uint(i), - Float32: float32(i) + 0.25, - Float64: float64(i) * math.Pi, - Str: fmt.Sprintf("host-%d", i%7), - Bin: []byte{byte(i), byte(i >> 8)}, - Timestamp: base.Add(time.Duration(i) * time.Second), - Date32: base.AddDate(0, 0, i), - Date64: base.AddDate(0, 0, i), - Time32: base.Add(time.Duration(i) * time.Minute), - Time64: base.Add(time.Duration(i) * time.Minute), - Duration: time.Duration(i) * time.Millisecond, - Dec32: decimal.Decimal32(i * 100), - Dec64: decimal.Decimal64(i * 10000), - Dec128: decimal128.FromU64(uint64(i) * 1000), - Dec256: decimal256.FromU64(uint64(i) * 2000), - LargeStr: fmt.Sprintf("large-%d", i), - ViewStr: fmt.Sprintf("view-%d", i), - LargeBin: []byte(fmt.Sprintf("lb-%d", i)), - ViewBin: []byte(fmt.Sprintf("vb-%d", i)), - DictStr: fmt.Sprintf("region-%d", i%3), - DictBin: []byte(fmt.Sprintf("k%d", i%4)), - DictInt: int32(i % 5), - DictF64: float64(i % 6), - Untagged: int64(i), - Secret: "never encoded", - } - if null { - // Leave every pointer nil, and null the []byte columns, which map to - // null without a pointer. - r.Bin = nil - r.LargeBin = nil - r.ViewBin = nil - r.DictBin = nil - rows[i] = r - continue - } - b := i%2 == 1 - i64 := int64(i) * 3 - f64 := float64(i) / 3 - s := fmt.Sprintf("p-%d", i) - bin := []byte(fmt.Sprintf("pb-%d", i)) - ts := base.Add(time.Duration(i) * time.Hour) - dur := time.Duration(i) * time.Microsecond - dec := decimal.Decimal32(i) - ds := fmt.Sprintf("pd-%d", i%2) - r.PBool, r.PInt64, r.PF64, r.PStr, r.PBin = &b, &i64, &f64, &s, &bin - dec128 := decimal128.FromU64(uint64(i)) - r.PTS, r.PDate32, r.PTime64, r.PDur, r.PDec32, r.PDictS = &ts, &ts, &ts, &dur, &dec, &ds - r.PDec128 = &dec128 - rows[i] = r - } - return rows -} - -// TestRowSchemaMatchesArreflect checks that the generated schema is what -// arreflect infers, column for column, including names, types and nullability. -func TestRowSchemaMatchesArreflect(t *testing.T) { - want, err := arreflect.InferSchema[gentypes.Row]() - if err != nil { - t.Fatalf("InferSchema: %v", err) - } - got := gentypes.RowSchema() - if got.NumFields() != want.NumFields() { - t.Fatalf("field count = %d, want %d\n got: %s\nwant: %s", got.NumFields(), want.NumFields(), got, want) - } - for i := 0; i < got.NumFields(); i++ { - gf, wf := got.Field(i), want.Field(i) - if gf.Name != wf.Name { - t.Errorf("field %d: name = %q, want %q", i, gf.Name, wf.Name) - continue - } - if _, gap := arreflectStructFieldGaps[gf.Name]; gap { - continue - } - if !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { - t.Errorf("field %q: got %s nullable=%t, want %s nullable=%t", gf.Name, gf.Type, gf.Nullable, wf.Type, wf.Nullable) - } - } -} - -func TestMetricSchemaMatchesArreflect(t *testing.T) { - want, err := arreflect.InferSchema[example.Metric]() - if err != nil { - t.Fatalf("InferSchema: %v", err) - } - got := example.MetricSchema() - // "ts" is a bare time.Time: see arreflectStructFieldGaps. - for i := 0; i < got.NumFields(); i++ { - gf, wf := got.Field(i), want.Field(i) - if gf.Name == "ts" { - continue - } - if gf.Name != wf.Name || !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { - t.Errorf("field %d: got %s, want %s", i, gf, wf) - } - } -} - -// TestRowRecordMatchesArreflect checks that for the same input the generated -// encoder and the reflection encoder produce equal columns. -// -// It compares column by column rather than with array.RecordEqual so the -// columns arreflect cannot encode as struct fields can be skipped by name; see -// arreflectStructFieldGaps. -func TestRowRecordMatchesArreflect(t *testing.T) { - for _, n := range []int{0, 1, 2, 3, 64} { - t.Run(fmt.Sprintf("rows=%d", n), func(t *testing.T) { - mem := memory.NewCheckedAllocator(memory.DefaultAllocator) - defer mem.AssertSize(t, 0) - - rows := makeRows(n) - - want, err := arreflect.RecordFromSlice(rows, mem) - if err != nil { - t.Fatalf("arreflect.RecordFromSlice: %v", err) - } - defer want.Release() - - got, err := gentypes.RowRecordBatch(mem, rows) - if err != nil { - t.Fatalf("RowRecordBatch: %v", err) - } - defer got.Release() - - compareColumns(t, fmt.Sprintf("rows=%d", n), got, want, false) - }) - } -} - -// TestArreflectStructFieldGaps pins the upstream behavior that -// arreflectStructFieldGaps works around: if arreflect learns to encode these -// struct fields, this fails and the skip list should shrink to match. -func TestArreflectStructFieldGaps(t *testing.T) { - schema, err := arreflect.InferSchema[gentypes.Row]() - if err != nil { - t.Fatalf("InferSchema: %v", err) - } - for name, goType := range arreflectStructFieldGaps { - idx := schema.FieldIndices(name) - if len(idx) != 1 { - t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) - } - if dt := schema.Field(idx[0]).Type; dt.ID() != arrow.STRUCT { - t.Errorf("arreflect now infers %s column %q as %s rather than an empty struct; "+ - "remove it from arreflectStructFieldGaps so the strict comparison covers it", goType, name, dt) - } - } -} - -// TestGapColumnValues covers the columns TestRowRecordMatchesArreflect skips. -// There is no reflection-built batch to compare against, so the expected Arrow -// values are spelled out directly. -func TestGapColumnValues(t *testing.T) { - mem := memory.NewCheckedAllocator(memory.DefaultAllocator) - defer mem.AssertSize(t, 0) - - ts := base - dec128 := decimal128.FromU64(12345) - dec256 := decimal256.FromU64(67890) - rows := []gentypes.Row{ - {Timestamp: ts, Dec128: dec128, Dec256: dec256, PTS: &ts, PDec128: &dec128}, - {Timestamp: ts.Add(time.Second), Dec128: dec128, Dec256: dec256}, // pointers nil - } - - rec, err := gentypes.RowRecordBatch(mem, rows) - if err != nil { - t.Fatalf("RowRecordBatch: %v", err) - } - defer rec.Release() - - col := func(name string) arrow.Array { - idx := rec.Schema().FieldIndices(name) - if len(idx) != 1 { - t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) - } - return rec.Column(idx[0]) - } - - tsCol := col("ts").(*array.Timestamp) - if got, want := tsCol.Value(0), arrow.Timestamp(ts.UnixNano()); got != want { - t.Errorf("ts[0] = %d, want %d", got, want) - } - if got, want := tsCol.Value(1), arrow.Timestamp(ts.Add(time.Second).UnixNano()); got != want { - t.Errorf("ts[1] = %d, want %d", got, want) - } - if dt := tsCol.DataType().(*arrow.TimestampType); dt.Unit != arrow.Nanosecond || dt.TimeZone != "UTC" { - t.Errorf("ts type = %s, want timestamp[ns, tz=UTC]", dt) - } - - ptsCol := col("pts").(*array.Timestamp) - if ptsCol.IsNull(0) { - t.Error("pts[0] is null, want a value") - } - if !ptsCol.IsNull(1) { - t.Error("pts[1] is not null, want null") - } - - if got := col("dec128").(*array.Decimal128).Value(0); got != dec128 { - t.Errorf("dec128[0] = %v, want %v", got, dec128) - } - if got := col("dec256").(*array.Decimal256).Value(0); got != dec256 { - t.Errorf("dec256[0] = %v, want %v", got, dec256) - } - if !col("pdec128").(*array.Decimal128).IsNull(1) { - t.Error("pdec128[1] is not null, want null") - } -} - -func TestMetricRecordMatchesArreflect(t *testing.T) { - mem := memory.NewCheckedAllocator(memory.DefaultAllocator) - defer mem.AssertSize(t, 0) - - v := 3.5 - metrics := []example.Metric{ - {Time: base, Host: "a", CPU: 0.5, Value: &v, Secret: "hidden"}, - {Time: base.Add(time.Second), Host: "b", CPU: 1.5, Value: nil}, - } - - want, err := arreflect.RecordFromSlice(metrics, mem) - if err != nil { - t.Fatalf("arreflect.RecordFromSlice: %v", err) - } - defer want.Release() - - got, err := example.MetricRecordBatch(mem, metrics) - if err != nil { - t.Fatalf("MetricRecordBatch: %v", err) - } - defer got.Release() - - // "ts" is a bare time.Time; see arreflectStructFieldGaps. - for i := 0; i < int(got.NumCols()); i++ { - name := got.Schema().Field(i).Name - if name == "ts" { - continue - } - if !array.Equal(got.Column(i), want.Column(i)) { - t.Errorf("column %q differs\n got: %s\nwant: %s", name, got.Column(i), want.Column(i)) - } - } - if got, want := got.Column(0).(*array.Timestamp).Value(0), arrow.Timestamp(base.UnixNano()); got != want { - t.Errorf("ts[0] = %d, want %d", got, want) - } - if secrets := got.Schema().FieldIndices("Secret"); len(secrets) != 0 { - t.Errorf(`a field tagged arrow:"-" was encoded as column %v`, secrets) - } -} - -// TestAppendMatchesAppendSlice checks that feeding rows one at a time lands in -// the same place as handing over a slice. -func TestAppendMatchesAppendSlice(t *testing.T) { - mem := memory.NewCheckedAllocator(memory.DefaultAllocator) - defer mem.AssertSize(t, 0) - - rows := makeRows(16) - - bulk := gentypes.NewRowAppender(mem) - defer bulk.Release() - bulk.AppendSlice(rows) - bulkRec := bulk.NewRecordBatch() - defer bulkRec.Release() - - streamed := gentypes.NewRowAppender(mem) - defer streamed.Release() - streamed.Reserve(len(rows)) - var row gentypes.Row - for i := range rows { - row = rows[i] // a caller reusing one row variable must be safe - streamed.Append(&row) - } - if got, want := streamed.Len(), len(rows); got != want { - t.Errorf("Len() = %d, want %d", got, want) - } - streamedRec := streamed.NewRecordBatch() - defer streamedRec.Release() - - if !array.RecordEqual(streamedRec, bulkRec) { - t.Errorf("streamed batch differs from bulk batch\n got: %s\nwant: %s", streamedRec, bulkRec) - } - if err := streamed.Err(); err != nil { - t.Errorf("Err() = %v, want nil", err) - } -} - -// TestAppenderReusableAcrossBatches checks that NewRecordBatch hands over the -// rows so far and leaves the appender ready for the next batch. -func TestAppenderReusableAcrossBatches(t *testing.T) { - mem := memory.NewCheckedAllocator(memory.DefaultAllocator) - defer mem.AssertSize(t, 0) - - rows := makeRows(9) - a := gentypes.NewRowAppender(mem) - defer a.Release() - - for i := 0; i < len(rows); i += 3 { - a.AppendSlice(rows[i : i+3]) - rec := a.NewRecordBatch() - if got := rec.NumRows(); got != 3 { - t.Errorf("batch %d: NumRows() = %d, want 3", i/3, got) - } - want, err := arreflect.RecordFromSlice(rows[i:i+3], mem) - if err != nil { - t.Fatalf("arreflect.RecordFromSlice: %v", err) - } - compareColumns(t, fmt.Sprintf("batch %d", i/3), rec, want, true) - want.Release() - rec.Release() - if got := a.Len(); got != 0 { - t.Errorf("batch %d: Len() after NewRecordBatch = %d, want 0", i/3, got) - } - } -} - -// TestNilAllocatorUsesDefault covers the constructor's nil-allocator contract. -func TestNilAllocatorUsesDefault(t *testing.T) { - a := gentypes.NewRowAppender(nil) - defer a.Release() - a.AppendSlice(makeRows(2)) - rec := a.NewRecordBatch() - defer rec.Release() - if got := rec.NumRows(); got != 2 { - t.Errorf("NumRows() = %d, want 2", got) - } -} diff --git a/arrgen/arrgen/internal/gentypes/instrumented_race_test.go b/arrgen/arrgen/internal/gentypes/instrumented_race_test.go deleted file mode 100644 index dbdd4fe22..000000000 --- a/arrgen/arrgen/internal/gentypes/instrumented_race_test.go +++ /dev/null @@ -1,21 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//go:build race || asan - -package gentypes_test - -const instrumented = true diff --git a/arrgen/arrgen/internal/gentypes/instrumented_test.go b/arrgen/arrgen/internal/gentypes/instrumented_test.go deleted file mode 100644 index f8d5afc8d..000000000 --- a/arrgen/arrgen/internal/gentypes/instrumented_test.go +++ /dev/null @@ -1,24 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//go:build !race && !asan - -package gentypes_test - -// instrumented reports whether the binary was built with a sanitizer that -// perturbs allocation counts and timings; the allocation and speed assertions -// skip when it is set. -const instrumented = false diff --git a/arrgen/arrgen/internal/gentypes/row_arrow.go b/arrgen/arrgen/internal/gentypes/row_arrow.go deleted file mode 100644 index 8a3a5332a..000000000 --- a/arrgen/arrgen/internal/gentypes/row_arrow.go +++ /dev/null @@ -1,475 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Code generated by arrgen. DO NOT EDIT. - -package gentypes - -import ( - "time" - - "github.com/apache/arrow-go/v18/arrow" - "github.com/apache/arrow-go/v18/arrow/array" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// rowArrowSchema is the Arrow schema of Row, resolved from its arrow -// struct tags. -var rowArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "b", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, - {Name: "i8", Type: arrow.PrimitiveTypes.Int8, Nullable: false}, - {Name: "i16", Type: arrow.PrimitiveTypes.Int16, Nullable: false}, - {Name: "i32", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, - {Name: "i64", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, - {Name: "i", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, - {Name: "u8", Type: arrow.PrimitiveTypes.Uint8, Nullable: false}, - {Name: "u16", Type: arrow.PrimitiveTypes.Uint16, Nullable: false}, - {Name: "u32", Type: arrow.PrimitiveTypes.Uint32, Nullable: false}, - {Name: "u64", Type: arrow.PrimitiveTypes.Uint64, Nullable: false}, - {Name: "u", Type: arrow.PrimitiveTypes.Uint64, Nullable: false}, - {Name: "f32", Type: arrow.PrimitiveTypes.Float32, Nullable: false}, - {Name: "f64", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, - {Name: "s", Type: arrow.BinaryTypes.String, Nullable: false}, - {Name: "bin", Type: arrow.BinaryTypes.Binary, Nullable: false}, - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, - {Name: "d32", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, - {Name: "d64", Type: arrow.FixedWidthTypes.Date64, Nullable: false}, - {Name: "t32", Type: &arrow.Time32Type{Unit: arrow.Millisecond}, Nullable: false}, - {Name: "t64", Type: &arrow.Time64Type{Unit: arrow.Nanosecond}, Nullable: false}, - {Name: "dur", Type: &arrow.DurationType{Unit: arrow.Nanosecond}, Nullable: false}, - {Name: "dec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: false}, - {Name: "dec64", Type: &arrow.Decimal64Type{Precision: 18, Scale: 4}, Nullable: false}, - {Name: "dec128", Type: &arrow.Decimal128Type{Precision: 20, Scale: 3}, Nullable: false}, - {Name: "dec256", Type: &arrow.Decimal256Type{Precision: 76, Scale: 0}, Nullable: false}, - {Name: "ls", Type: arrow.BinaryTypes.LargeString, Nullable: false}, - {Name: "vs", Type: arrow.BinaryTypes.StringView, Nullable: false}, - {Name: "lb", Type: arrow.BinaryTypes.LargeBinary, Nullable: false}, - {Name: "vb", Type: arrow.BinaryTypes.BinaryView, Nullable: false}, - {Name: "ds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: false}, - {Name: "db", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.Binary}, Nullable: false}, - {Name: "di", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Int32}, Nullable: false}, - {Name: "df", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Float64}, Nullable: false}, - {Name: "Untagged", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, - {Name: "pb", Type: arrow.FixedWidthTypes.Boolean, Nullable: true}, - {Name: "pi64", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, - {Name: "pf64", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, - {Name: "ps", Type: arrow.BinaryTypes.String, Nullable: true}, - {Name: "pbin", Type: arrow.BinaryTypes.Binary, Nullable: true}, - {Name: "pts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: true}, - {Name: "pd32", Type: arrow.FixedWidthTypes.Date32, Nullable: true}, - {Name: "pt64", Type: &arrow.Time64Type{Unit: arrow.Nanosecond}, Nullable: true}, - {Name: "pdur", Type: &arrow.DurationType{Unit: arrow.Nanosecond}, Nullable: true}, - {Name: "pdec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: true}, - {Name: "pdec128", Type: &arrow.Decimal128Type{Precision: 38, Scale: 0}, Nullable: true}, - {Name: "pds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: true}, -}, nil) - -// RowSchema returns the Arrow schema encoded by RowAppender. -func RowSchema() *arrow.Schema { return rowArrowSchema } - -// RowAppender converts Row values into Arrow record batches -// without reflection. -// -// It is not safe for concurrent use. Release it when done. -type RowAppender struct { - rb *array.RecordBuilder - b0 *array.BooleanBuilder - b1 *array.Int8Builder - b2 *array.Int16Builder - b3 *array.Int32Builder - b4 *array.Int64Builder - b5 *array.Int64Builder - b6 *array.Uint8Builder - b7 *array.Uint16Builder - b8 *array.Uint32Builder - b9 *array.Uint64Builder - b10 *array.Uint64Builder - b11 *array.Float32Builder - b12 *array.Float64Builder - b13 *array.StringBuilder - b14 *array.BinaryBuilder - b15 *array.TimestampBuilder - b16 *array.Date32Builder - b17 *array.Date64Builder - b18 *array.Time32Builder - b19 *array.Time64Builder - b20 *array.DurationBuilder - b21 *array.Decimal32Builder - b22 *array.Decimal64Builder - b23 *array.Decimal128Builder - b24 *array.Decimal256Builder - b25 *array.LargeStringBuilder - b26 *array.StringViewBuilder - b27 *array.BinaryBuilder - b28 *array.BinaryViewBuilder - b29 *array.BinaryDictionaryBuilder - b30 *array.BinaryDictionaryBuilder - b31 *array.Int32DictionaryBuilder - b32 *array.Float64DictionaryBuilder - b33 *array.Int64Builder - b34 *array.BooleanBuilder - b35 *array.Int64Builder - b36 *array.Float64Builder - b37 *array.StringBuilder - b38 *array.BinaryBuilder - b39 *array.TimestampBuilder - b40 *array.Date32Builder - b41 *array.Time64Builder - b42 *array.DurationBuilder - b43 *array.Decimal32Builder - b44 *array.Decimal128Builder - b45 *array.BinaryDictionaryBuilder - err error -} - -// NewRowAppender returns an appender that builds Row batches with mem, -// or memory.DefaultAllocator when mem is nil. The typed field builders are -// resolved once here, so Append does no type assertions. -func NewRowAppender(mem memory.Allocator) *RowAppender { - if mem == nil { - mem = memory.DefaultAllocator - } - rb := array.NewRecordBuilder(mem, rowArrowSchema) - return &RowAppender{ - rb: rb, - b0: rb.Field(0).(*array.BooleanBuilder), - b1: rb.Field(1).(*array.Int8Builder), - b2: rb.Field(2).(*array.Int16Builder), - b3: rb.Field(3).(*array.Int32Builder), - b4: rb.Field(4).(*array.Int64Builder), - b5: rb.Field(5).(*array.Int64Builder), - b6: rb.Field(6).(*array.Uint8Builder), - b7: rb.Field(7).(*array.Uint16Builder), - b8: rb.Field(8).(*array.Uint32Builder), - b9: rb.Field(9).(*array.Uint64Builder), - b10: rb.Field(10).(*array.Uint64Builder), - b11: rb.Field(11).(*array.Float32Builder), - b12: rb.Field(12).(*array.Float64Builder), - b13: rb.Field(13).(*array.StringBuilder), - b14: rb.Field(14).(*array.BinaryBuilder), - b15: rb.Field(15).(*array.TimestampBuilder), - b16: rb.Field(16).(*array.Date32Builder), - b17: rb.Field(17).(*array.Date64Builder), - b18: rb.Field(18).(*array.Time32Builder), - b19: rb.Field(19).(*array.Time64Builder), - b20: rb.Field(20).(*array.DurationBuilder), - b21: rb.Field(21).(*array.Decimal32Builder), - b22: rb.Field(22).(*array.Decimal64Builder), - b23: rb.Field(23).(*array.Decimal128Builder), - b24: rb.Field(24).(*array.Decimal256Builder), - b25: rb.Field(25).(*array.LargeStringBuilder), - b26: rb.Field(26).(*array.StringViewBuilder), - b27: rb.Field(27).(*array.BinaryBuilder), - b28: rb.Field(28).(*array.BinaryViewBuilder), - b29: rb.Field(29).(*array.BinaryDictionaryBuilder), - b30: rb.Field(30).(*array.BinaryDictionaryBuilder), - b31: rb.Field(31).(*array.Int32DictionaryBuilder), - b32: rb.Field(32).(*array.Float64DictionaryBuilder), - b33: rb.Field(33).(*array.Int64Builder), - b34: rb.Field(34).(*array.BooleanBuilder), - b35: rb.Field(35).(*array.Int64Builder), - b36: rb.Field(36).(*array.Float64Builder), - b37: rb.Field(37).(*array.StringBuilder), - b38: rb.Field(38).(*array.BinaryBuilder), - b39: rb.Field(39).(*array.TimestampBuilder), - b40: rb.Field(40).(*array.Date32Builder), - b41: rb.Field(41).(*array.Time64Builder), - b42: rb.Field(42).(*array.DurationBuilder), - b43: rb.Field(43).(*array.Decimal32Builder), - b44: rb.Field(44).(*array.Decimal128Builder), - b45: rb.Field(45).(*array.BinaryDictionaryBuilder), - } -} - -// Schema returns the schema of the batches this appender builds. -func (a *RowAppender) Schema() *arrow.Schema { return rowArrowSchema } - -// RecordBuilder returns the underlying builder. Appending through it directly -// is allowed as long as every column stays the same length. -func (a *RowAppender) RecordBuilder() *array.RecordBuilder { return a.rb } - -// Reserve pre-allocates room for n more rows in every column. Reserving before -// a run of Appends keeps the append path free of allocations. -func (a *RowAppender) Reserve(n int) { a.rb.Reserve(n) } - -// Len reports the number of rows appended since the last NewRecordBatch. -func (a *RowAppender) Len() int { return a.rb.Field(0).Len() } - -// Append appends one row. v is read synchronously and never retained, so a -// caller can reuse a single Row variable across calls. -func (a *RowAppender) Append(v *Row) { - a.b0.Append(v.Bool) - a.b1.Append(v.Int8) - a.b2.Append(v.Int16) - a.b3.Append(v.Int32) - a.b4.Append(v.Int64) - a.b5.Append(int64(v.Int)) - a.b6.Append(v.Uint8) - a.b7.Append(v.Uint16) - a.b8.Append(v.Uint32) - a.b9.Append(v.Uint64) - a.b10.Append(uint64(v.Uint)) - a.b11.Append(v.Float32) - a.b12.Append(v.Float64) - a.b13.Append(v.Str) - if v.Bin == nil { - a.b14.AppendNull() - } else { - a.b14.Append(v.Bin) - } - a.b15.Append(arrow.Timestamp(v.Timestamp.UnixNano())) - a.b16.Append(arrow.Date32FromTime(v.Date32)) - a.b17.Append(arrow.Date64FromTime(v.Date64)) - { - tod := v.Time32.UTC() - a.b18.Append(arrow.Time32(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds() / 1e6)) - } - { - tod := v.Time64.UTC() - a.b19.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) - } - a.b20.Append(arrow.Duration(v.Duration.Nanoseconds())) - a.b21.Append(v.Dec32) - a.b22.Append(v.Dec64) - a.b23.Append(v.Dec128) - a.b24.Append(v.Dec256) - a.b25.Append(v.LargeStr) - a.b26.Append(v.ViewStr) - if v.LargeBin == nil { - a.b27.AppendNull() - } else { - a.b27.Append(v.LargeBin) - } - if v.ViewBin == nil { - a.b28.AppendNull() - } else { - a.b28.Append(v.ViewBin) - } - a.setErr(a.b29.AppendString(v.DictStr)) - if v.DictBin == nil { - a.b30.AppendNull() - } else { - a.setErr(a.b30.Append(v.DictBin)) - } - a.setErr(a.b31.Append(v.DictInt)) - a.setErr(a.b32.Append(v.DictF64)) - a.b33.Append(v.Untagged) - if v.PBool == nil { - a.b34.AppendNull() - } else { - a.b34.Append(*v.PBool) - } - if v.PInt64 == nil { - a.b35.AppendNull() - } else { - a.b35.Append(*v.PInt64) - } - if v.PF64 == nil { - a.b36.AppendNull() - } else { - a.b36.Append(*v.PF64) - } - if v.PStr == nil { - a.b37.AppendNull() - } else { - a.b37.Append(*v.PStr) - } - if v.PBin == nil || *v.PBin == nil { - a.b38.AppendNull() - } else { - a.b38.Append(*v.PBin) - } - if v.PTS == nil { - a.b39.AppendNull() - } else { - a.b39.Append(arrow.Timestamp((*v.PTS).UnixNano())) - } - if v.PDate32 == nil { - a.b40.AppendNull() - } else { - a.b40.Append(arrow.Date32FromTime(*v.PDate32)) - } - if v.PTime64 == nil { - a.b41.AppendNull() - } else { - tod := (*v.PTime64).UTC() - a.b41.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) - } - if v.PDur == nil { - a.b42.AppendNull() - } else { - a.b42.Append(arrow.Duration((*v.PDur).Nanoseconds())) - } - if v.PDec32 == nil { - a.b43.AppendNull() - } else { - a.b43.Append(*v.PDec32) - } - if v.PDec128 == nil { - a.b44.AppendNull() - } else { - a.b44.Append(*v.PDec128) - } - if v.PDictS == nil { - a.b45.AppendNull() - } else { - a.setErr(a.b45.AppendString(*v.PDictS)) - } -} - -// AppendSlice appends every element of vs, reserving room for them up front. -func (a *RowAppender) AppendSlice(vs []Row) { - a.rb.Reserve(len(vs)) - for i := range vs { - a.Append(&vs[i]) - } -} - -// NewRecordBatch returns the appended rows as a record batch and resets the -// appender. The caller must Release the batch. -func (a *RowAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } - -// Err returns the first error an Append hit, or nil. Only dictionary columns -// can fail. The error is sticky: NewRecordBatch does not clear it. -func (a *RowAppender) Err() error { return a.err } - -// setErr keeps the first error. -func (a *RowAppender) setErr(err error) { - if err != nil && a.err == nil { - a.err = err - } -} - -// Release releases the appender's builders. -func (a *RowAppender) Release() { a.rb.Release() } - -// RowRecordBatch builds one record batch holding every element of vs. It is the -// generated equivalent of arreflect.RecordFromSlice[Row](vs, mem). The -// caller must Release the batch. -func RowRecordBatch(mem memory.Allocator, vs []Row) (arrow.RecordBatch, error) { - a := NewRowAppender(mem) - defer a.Release() - a.AppendSlice(vs) - if err := a.Err(); err != nil { - return nil, err - } - return a.NewRecordBatch(), nil -} - -// fixedArrowSchema is the Arrow schema of Fixed, resolved from its arrow -// struct tags. -var fixedArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, - {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, - {Name: "val", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, - {Name: "ok", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, - {Name: "opt", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, -}, nil) - -// FixedSchema returns the Arrow schema encoded by FixedAppender. -func FixedSchema() *arrow.Schema { return fixedArrowSchema } - -// FixedAppender converts Fixed values into Arrow record batches -// without reflection. -// -// It is not safe for concurrent use. Release it when done. -type FixedAppender struct { - rb *array.RecordBuilder - b0 *array.TimestampBuilder - b1 *array.Int64Builder - b2 *array.Float64Builder - b3 *array.BooleanBuilder - b4 *array.Float64Builder - err error -} - -// NewFixedAppender returns an appender that builds Fixed batches with mem, -// or memory.DefaultAllocator when mem is nil. The typed field builders are -// resolved once here, so Append does no type assertions. -func NewFixedAppender(mem memory.Allocator) *FixedAppender { - if mem == nil { - mem = memory.DefaultAllocator - } - rb := array.NewRecordBuilder(mem, fixedArrowSchema) - return &FixedAppender{ - rb: rb, - b0: rb.Field(0).(*array.TimestampBuilder), - b1: rb.Field(1).(*array.Int64Builder), - b2: rb.Field(2).(*array.Float64Builder), - b3: rb.Field(3).(*array.BooleanBuilder), - b4: rb.Field(4).(*array.Float64Builder), - } -} - -// Schema returns the schema of the batches this appender builds. -func (a *FixedAppender) Schema() *arrow.Schema { return fixedArrowSchema } - -// RecordBuilder returns the underlying builder. Appending through it directly -// is allowed as long as every column stays the same length. -func (a *FixedAppender) RecordBuilder() *array.RecordBuilder { return a.rb } - -// Reserve pre-allocates room for n more rows in every column. Reserving before -// a run of Appends keeps the append path free of allocations. -func (a *FixedAppender) Reserve(n int) { a.rb.Reserve(n) } - -// Len reports the number of rows appended since the last NewRecordBatch. -func (a *FixedAppender) Len() int { return a.rb.Field(0).Len() } - -// Append appends one row. v is read synchronously and never retained, so a -// caller can reuse a single Fixed variable across calls. -func (a *FixedAppender) Append(v *Fixed) { - a.b0.Append(arrow.Timestamp(v.Timestamp.UnixNano())) - a.b1.Append(v.ID) - a.b2.Append(v.Value) - a.b3.Append(v.OK) - if v.Optional == nil { - a.b4.AppendNull() - } else { - a.b4.Append(*v.Optional) - } -} - -// AppendSlice appends every element of vs, reserving room for them up front. -func (a *FixedAppender) AppendSlice(vs []Fixed) { - a.rb.Reserve(len(vs)) - for i := range vs { - a.Append(&vs[i]) - } -} - -// NewRecordBatch returns the appended rows as a record batch and resets the -// appender. The caller must Release the batch. -func (a *FixedAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } - -// Err returns the first error an Append hit, or nil. Only dictionary columns -// can fail. The error is sticky: NewRecordBatch does not clear it. -func (a *FixedAppender) Err() error { return a.err } - -// Release releases the appender's builders. -func (a *FixedAppender) Release() { a.rb.Release() } - -// FixedRecordBatch builds one record batch holding every element of vs. It is the -// generated equivalent of arreflect.RecordFromSlice[Fixed](vs, mem). The -// caller must Release the batch. -func FixedRecordBatch(mem memory.Allocator, vs []Fixed) (arrow.RecordBatch, error) { - a := NewFixedAppender(mem) - defer a.Release() - a.AppendSlice(vs) - if err := a.Err(); err != nil { - return nil, err - } - return a.NewRecordBatch(), nil -} diff --git a/arrgen/arrgen/internal/gentypes/types.go b/arrgen/arrgen/internal/gentypes/types.go deleted file mode 100644 index 318dade5e..000000000 --- a/arrgen/arrgen/internal/gentypes/types.go +++ /dev/null @@ -1,98 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package gentypes holds the fixtures the arrgen test suite generates from. -// Row carries one field for every column shape arrgen supports, so the -// equivalence tests cover the whole surface. -package gentypes - -import ( - "time" - - "github.com/apache/arrow-go/v18/arrow/decimal" - "github.com/apache/arrow-go/v18/arrow/decimal128" - "github.com/apache/arrow-go/v18/arrow/decimal256" -) - -// Row covers every supported field type, tag option and nullable variant. -type Row struct { - Bool bool `arrow:"b"` - Int8 int8 `arrow:"i8"` - Int16 int16 `arrow:"i16"` - Int32 int32 `arrow:"i32"` - Int64 int64 `arrow:"i64"` - Int int `arrow:"i"` - Uint8 uint8 `arrow:"u8"` - Uint16 uint16 `arrow:"u16"` - Uint32 uint32 `arrow:"u32"` - Uint64 uint64 `arrow:"u64"` - Uint uint `arrow:"u"` - Float32 float32 `arrow:"f32"` - Float64 float64 `arrow:"f64"` - Str string `arrow:"s"` - Bin []byte `arrow:"bin"` - - Timestamp time.Time `arrow:"ts"` - Date32 time.Time `arrow:"d32,date32"` - Date64 time.Time `arrow:"d64,date64"` - Time32 time.Time `arrow:"t32,time32"` - Time64 time.Time `arrow:"t64,time64"` - Duration time.Duration `arrow:"dur"` - - Dec32 decimal.Decimal32 `arrow:"dec32"` - Dec64 decimal.Decimal64 `arrow:"dec64,decimal(18,4)"` - Dec128 decimal128.Num `arrow:"dec128,decimal(20,3)"` - Dec256 decimal256.Num `arrow:"dec256"` - - LargeStr string `arrow:"ls,large"` - ViewStr string `arrow:"vs,view"` - LargeBin []byte `arrow:"lb,large"` - ViewBin []byte `arrow:"vb,view"` - - DictStr string `arrow:"ds,dict"` - DictBin []byte `arrow:"db,dict"` - DictInt int32 `arrow:"di,dict"` - DictF64 float64 `arrow:"df,dict"` - Untagged int64 // no tag: the column takes the Go field name - - PBool *bool `arrow:"pb"` - PInt64 *int64 `arrow:"pi64"` - PF64 *float64 `arrow:"pf64"` - PStr *string `arrow:"ps"` - PBin *[]byte `arrow:"pbin"` - PTS *time.Time `arrow:"pts"` - PDate32 *time.Time `arrow:"pd32,date32"` - PTime64 *time.Time `arrow:"pt64,time64"` - PDur *time.Duration `arrow:"pdur"` - PDec32 *decimal.Decimal32 `arrow:"pdec32"` - PDec128 *decimal128.Num `arrow:"pdec128"` - PDictS *string `arrow:"pds,dict"` - - Secret string `arrow:"-"` // excluded from Arrow entirely -} - -// Fixed holds only fixed-width columns, so appending a row never has to grow a -// variable-length data buffer. It is what the zero-allocation assertion in -// alloc_test.go measures. -type Fixed struct { - Timestamp time.Time `arrow:"ts"` - ID int64 `arrow:"id"` - Value float64 `arrow:"val"` - OK bool `arrow:"ok"` - Optional *float64 `arrow:"opt"` -} - -//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Row,Fixed -header ../../license_header.txt -output row_arrow.go diff --git a/arrgen/arrgen/license_header.txt b/arrgen/arrgen/license_header.txt deleted file mode 100644 index f1fc31831..000000000 --- a/arrgen/arrgen/license_header.txt +++ /dev/null @@ -1,15 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. diff --git a/arrgen/arrgen/mapping.go b/arrgen/arrgen/mapping.go deleted file mode 100644 index f00e6583b..000000000 --- a/arrgen/arrgen/mapping.go +++ /dev/null @@ -1,320 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package arrgen - -import ( - "fmt" - "go/types" -) - -// colSpec is everything the template needs to emit one Arrow column: the type -// expression for the schema, the concrete builder the appender caches, and a -// renderer for the append statement. The three are resolved together so they -// cannot drift apart. -type colSpec struct { - arrowType string // Go expression for the arrow.DataType - builderType string // Go type of the cached typed builder - fallible bool // the builder's Append returns an error (dictionaries) - needsTime bool // the append statement references the time package - nilable bool // a nil Go value maps to null even in a non-nullable column - - // appendStmt renders the append of a known-non-nil value. bld is the cached - // builder expression ("a.b3"), val is the value ("v.F" or "*v.F"), and recv - // is val parenthesized so a method call binds to the value, not the pointer. - appendStmt func(bld, recv, val string) string -} - -// optSupport records which tag options a column kind can honor. An option that -// would be ignored is rejected instead, unlike arreflect. -type optSupport struct { - temporal bool - decimal bool - largeView bool - dict bool -} - -func checkOpts(o tagOpts, goType string, s optSupport) error { - switch { - case o.Temporal != "" && !s.temporal: - return fmt.Errorf("%q is only valid on a time.Time field, not %s", o.Temporal, goType) - case o.HasDecimalOpts && !s.decimal: - return fmt.Errorf("decimal(precision,scale) is only valid on a decimal field, not %s", goType) - case o.Large && !s.largeView: - return fmt.Errorf("large has no effect on %s; it is only valid on string and []byte fields", goType) - case o.View && !s.largeView: - return fmt.Errorf("view has no effect on %s; it is only valid on string and []byte fields", goType) - case o.Dict && !s.dict: - return fmt.Errorf("dict is not supported on %s; it is valid on string, []byte, integer and float fields", goType) - case o.Dict && o.Large: - return fmt.Errorf("dict cannot be combined with large: Dictionary is not implemented by arrow-go") - } - return nil -} - -// resolveColumn maps a Go field type and its parsed tag to an Arrow column, -// reporting nullable=true for a pointer field, as arreflect does. -func resolveColumn(t types.Type, opts tagOpts) (spec colSpec, nullable bool, err error) { - t = types.Unalias(t) - if p, ok := t.(*types.Pointer); ok { - nullable = true - t = types.Unalias(p.Elem()) - if _, ok := t.(*types.Pointer); ok { - return colSpec{}, false, fmt.Errorf("multi-level pointer type %s is not supported", t) - } - } - spec, err = baseSpec(t, opts) - return spec, nullable, err -} - -func baseSpec(t types.Type, opts tagOpts) (colSpec, error) { - if named, ok := t.(*types.Named); ok { - return namedSpec(named, opts) - } - if sl, ok := t.(*types.Slice); ok { - if b, ok := types.Unalias(sl.Elem()).(*types.Basic); ok && b.Kind() == types.Uint8 { - return byteSliceSpec(opts) - } - return colSpec{}, fmt.Errorf("list column for %s is not supported; encode it with arreflect", t) - } - if b, ok := t.(*types.Basic); ok { - return basicSpec(b, opts) - } - return colSpec{}, fmt.Errorf("type %s is not supported", t) -} - -// namedSpec handles the named types arreflect recognizes by identity. Every -// other named type, including a defined scalar such as `type ID int64`, is -// unsupported there too. -func namedSpec(named *types.Named, opts tagOpts) (colSpec, error) { - obj := named.Obj() - pkg := "" - if obj.Pkg() != nil { - pkg = obj.Pkg().Path() - } - switch { - case pkg == "time" && obj.Name() == "Time": - return timeSpec(opts) - case pkg == "time" && obj.Name() == "Duration": - if err := checkOpts(opts, "time.Duration", optSupport{}); err != nil { - return colSpec{}, err - } - return colSpec{ - arrowType: "&arrow.DurationType{Unit: arrow.Nanosecond}", - builderType: "*array.DurationBuilder", - appendStmt: func(bld, recv, _ string) string { - return fmt.Sprintf("%s.Append(arrow.Duration(%s.Nanoseconds()))", bld, recv) - }, - }, nil - case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal32": - return decimalSpec(opts, "Decimal32", 9) - case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal64": - return decimalSpec(opts, "Decimal64", 18) - case isDecimalPkg(pkg, "decimal128") && obj.Name() == "Num": - return decimalSpec(opts, "Decimal128", 38) - case isDecimalPkg(pkg, "decimal256") && obj.Name() == "Num": - return decimalSpec(opts, "Decimal256", 76) - } - if _, ok := named.Underlying().(*types.Struct); ok { - return colSpec{}, fmt.Errorf("nested struct column for %s is not supported; encode it with arreflect", named) - } - return colSpec{}, fmt.Errorf("named type %s is not supported; arreflect matches scalar fields by exact type, so a defined type such as this one has no Arrow mapping in either path", named) -} - -func isDecimalPkg(path, name string) bool { - const prefix = "github.com/apache/arrow-go/v18/arrow/" - return path == prefix+name -} - -func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32) (colSpec, error) { - goType := "decimal." + arrowName - if err := checkOpts(opts, goType, optSupport{decimal: true}); err != nil { - return colSpec{}, err - } - precision, scale := defaultPrecision, int32(0) - if opts.HasDecimalOpts { - precision, scale = opts.DecimalPrecision, opts.DecimalScale - } - return colSpec{ - arrowType: fmt.Sprintf("&arrow.%sType{Precision: %d, Scale: %d}", arrowName, precision, scale), - builderType: "*array." + arrowName + "Builder", - appendStmt: simpleAppend(), - }, nil -} - -// timeSpec maps a time.Time field. An untagged time.Time becomes -// TIMESTAMP(ns, UTC), which arreflect's struct-field path does not currently -// produce; see the divergence note in README.md. -func timeSpec(opts tagOpts) (colSpec, error) { - if err := checkOpts(opts, "time.Time", optSupport{temporal: true}); err != nil { - return colSpec{}, err - } - switch opts.Temporal { - case "date32": - return colSpec{ - arrowType: "arrow.FixedWidthTypes.Date32", - builderType: "*array.Date32Builder", - appendStmt: func(bld, _, val string) string { - return fmt.Sprintf("%s.Append(arrow.Date32FromTime(%s))", bld, val) - }, - }, nil - case "date64": - return colSpec{ - arrowType: "arrow.FixedWidthTypes.Date64", - builderType: "*array.Date64Builder", - appendStmt: func(bld, _, val string) string { - return fmt.Sprintf("%s.Append(arrow.Date64FromTime(%s))", bld, val) - }, - }, nil - case "time32": - return timeOfDaySpec("&arrow.Time32Type{Unit: arrow.Millisecond}", "*array.Time32Builder", "arrow.Time32", " / 1e6"), nil - case "time64": - return timeOfDaySpec("&arrow.Time64Type{Unit: arrow.Nanosecond}", "*array.Time64Builder", "arrow.Time64", ""), nil - default: - return colSpec{ - arrowType: `&arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}`, - builderType: "*array.TimestampBuilder", - appendStmt: func(bld, recv, _ string) string { - return fmt.Sprintf("%s.Append(arrow.Timestamp(%s.UnixNano()))", bld, recv) - }, - }, nil - } -} - -// timeOfDaySpec renders arreflect's timeOfDayNanos inline, as a block rather -// than a package-level helper so two generated files in one package cannot -// collide over the helper's name. -func timeOfDaySpec(arrowType, builderType, cast, divisor string) colSpec { - return colSpec{ - arrowType: arrowType, - builderType: builderType, - needsTime: true, - appendStmt: func(bld, recv, _ string) string { - return fmt.Sprintf(`{ -tod := %s.UTC() -%s.Append(%s(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds()%s)) -}`, recv, bld, cast, divisor) - }, - } -} - -func byteSliceSpec(opts tagOpts) (colSpec, error) { - if err := checkOpts(opts, "[]byte", optSupport{largeView: true, dict: true}); err != nil { - return colSpec{}, err - } - // A nil []byte is null in arreflect even when the column is not nullable. - if opts.Dict { - spec := dictSpec("arrow.BinaryTypes.Binary", "*array.BinaryDictionaryBuilder", simpleAppend()) - spec.nilable = true - return spec, nil - } - switch { - case opts.View: - return colSpec{arrowType: "arrow.BinaryTypes.BinaryView", builderType: "*array.BinaryViewBuilder", appendStmt: simpleAppend(), nilable: true}, nil - case opts.Large: - return colSpec{arrowType: "arrow.BinaryTypes.LargeBinary", builderType: "*array.BinaryBuilder", appendStmt: simpleAppend(), nilable: true}, nil - default: - return colSpec{arrowType: "arrow.BinaryTypes.Binary", builderType: "*array.BinaryBuilder", appendStmt: simpleAppend(), nilable: true}, nil - } -} - -func stringSpec(opts tagOpts) (colSpec, error) { - if err := checkOpts(opts, "string", optSupport{largeView: true, dict: true}); err != nil { - return colSpec{}, err - } - if opts.Dict { - return dictSpec("arrow.BinaryTypes.String", "*array.BinaryDictionaryBuilder", func(bld, _, val string) string { - return fmt.Sprintf("%s.AppendString(%s)", bld, val) - }), nil - } - switch { - case opts.View: - return colSpec{arrowType: "arrow.BinaryTypes.StringView", builderType: "*array.StringViewBuilder", appendStmt: simpleAppend()}, nil - case opts.Large: - return colSpec{arrowType: "arrow.BinaryTypes.LargeString", builderType: "*array.LargeStringBuilder", appendStmt: simpleAppend()}, nil - default: - return colSpec{arrowType: "arrow.BinaryTypes.String", builderType: "*array.StringBuilder", appendStmt: simpleAppend()}, nil - } -} - -// dictSpec wraps a value type in Dictionary, matching -// arreflect.applyEncodingOpts. Dictionary appends can fail. -func dictSpec(valueType, builderType string, appendStmt func(bld, recv, val string) string) colSpec { - return colSpec{ - arrowType: fmt.Sprintf("&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: %s}", valueType), - builderType: builderType, - fallible: true, - appendStmt: appendStmt, - } -} - -// numeric describes one Go basic kind's Arrow column. -type numeric struct { - arrowType string // e.g. "arrow.PrimitiveTypes.Int64" - builder string // e.g. "Int64" - cast string // non-empty when the Go type is wider-named than the column ("int" -> int64) -} - -var numerics = map[types.BasicKind]numeric{ - types.Int8: {"arrow.PrimitiveTypes.Int8", "Int8", ""}, - types.Int16: {"arrow.PrimitiveTypes.Int16", "Int16", ""}, - types.Int32: {"arrow.PrimitiveTypes.Int32", "Int32", ""}, - types.Int64: {"arrow.PrimitiveTypes.Int64", "Int64", ""}, - types.Int: {"arrow.PrimitiveTypes.Int64", "Int64", "int64"}, - types.Uint8: {"arrow.PrimitiveTypes.Uint8", "Uint8", ""}, - types.Uint16: {"arrow.PrimitiveTypes.Uint16", "Uint16", ""}, - types.Uint32: {"arrow.PrimitiveTypes.Uint32", "Uint32", ""}, - types.Uint64: {"arrow.PrimitiveTypes.Uint64", "Uint64", ""}, - types.Uint: {"arrow.PrimitiveTypes.Uint64", "Uint64", "uint64"}, - types.Float32: {"arrow.PrimitiveTypes.Float32", "Float32", ""}, - types.Float64: {"arrow.PrimitiveTypes.Float64", "Float64", ""}, -} - -func basicSpec(b *types.Basic, opts tagOpts) (colSpec, error) { - switch b.Kind() { - case types.String: - return stringSpec(opts) - case types.Bool: - if err := checkOpts(opts, "bool", optSupport{}); err != nil { - return colSpec{}, err - } - return colSpec{arrowType: "arrow.FixedWidthTypes.Boolean", builderType: "*array.BooleanBuilder", appendStmt: simpleAppend()}, nil - } - - n, ok := numerics[b.Kind()] - if !ok { - return colSpec{}, fmt.Errorf("type %s is not supported", b) - } - if err := checkOpts(opts, b.Name(), optSupport{dict: true}); err != nil { - return colSpec{}, err - } - appendStmt := simpleAppend() - if n.cast != "" { - appendStmt = castAppend(n.cast) - } - if opts.Dict { - return dictSpec(n.arrowType, "*array."+n.builder+"DictionaryBuilder", appendStmt), nil - } - return colSpec{arrowType: n.arrowType, builderType: "*array." + n.builder + "Builder", appendStmt: appendStmt}, nil -} - -func simpleAppend() func(bld, recv, val string) string { - return func(bld, _, val string) string { return fmt.Sprintf("%s.Append(%s)", bld, val) } -} - -func castAppend(cast string) func(bld, recv, val string) string { - return func(bld, _, val string) string { return fmt.Sprintf("%s.Append(%s(%s))", bld, cast, val) } -} diff --git a/arrgen/arrgen/tag.go b/arrgen/arrgen/tag.go deleted file mode 100644 index f375bcfca..000000000 --- a/arrgen/arrgen/tag.go +++ /dev/null @@ -1,140 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package arrgen - -import ( - "fmt" - "strconv" - "strings" -) - -// tagOpts is the parsed form of one `arrow:"..."` struct tag. It mirrors the -// unexported tagOpts in arreflect; keep the two in sync. -type tagOpts struct { - Name string - Skip bool - Dict bool - View bool - REE bool - Large bool - DecimalPrecision int32 - DecimalScale int32 - HasDecimalOpts bool - Temporal string // "", "timestamp", "date32", "date64", "time32", "time64" -} - -// parseTag parses the value of an `arrow` struct tag. Unlike arreflect, which -// defers diagnostics, an unusable tag is an error. -func parseTag(tag string) (tagOpts, error) { - if tag == "-" { - return tagOpts{Skip: true}, nil - } - - name, rest, _ := strings.Cut(tag, ",") - opts := tagOpts{Name: name} - if rest == "" { - return opts, nil - } - for _, token := range splitTagTokens(rest) { - if err := applyTagToken(&opts, token); err != nil { - return tagOpts{}, err - } - } - return opts, nil -} - -// splitTagTokens splits the option list on commas outside parentheses, so -// decimal(18,2) stays one token. -func splitTagTokens(rest string) []string { - var tokens []string - depth, start := 0, 0 - for i := 0; i < len(rest); i++ { - switch rest[i] { - case '(': - depth++ - case ')': - depth-- - case ',': - if depth == 0 { - tokens = append(tokens, strings.TrimSpace(rest[start:i])) - start = i + 1 - } - } - } - if start < len(rest) { - tokens = append(tokens, strings.TrimSpace(rest[start:])) - } - return tokens -} - -func applyTagToken(opts *tagOpts, token string) error { - if strings.HasPrefix(token, "decimal(") && strings.HasSuffix(token, ")") { - return parseDecimalOpt(opts, token) - } - switch token { - case "dict": - opts.Dict = true - case "view": - opts.View = true - case "ree": - opts.REE = true - case "large": - opts.Large = true - case "date32", "date64", "time32", "time64", "timestamp": - opts.Temporal = token - default: - return fmt.Errorf("unknown option %q", token) - } - return nil -} - -func parseDecimalOpt(opts *tagOpts, token string) error { - inner := strings.TrimSuffix(strings.TrimPrefix(token, "decimal("), ")") - parts := strings.SplitN(inner, ",", 2) - if len(parts) != 2 { - return fmt.Errorf("invalid decimal tag %q: expected decimal(precision,scale)", token) - } - p, err := strconv.ParseInt(strings.TrimSpace(parts[0]), 10, 32) - if err != nil { - return fmt.Errorf("invalid decimal tag %q: precision %q is not an integer", token, strings.TrimSpace(parts[0])) - } - s, err := strconv.ParseInt(strings.TrimSpace(parts[1]), 10, 32) - if err != nil { - return fmt.Errorf("invalid decimal tag %q: scale %q is not an integer", token, strings.TrimSpace(parts[1])) - } - opts.HasDecimalOpts = true - opts.DecimalPrecision = int32(p) - opts.DecimalScale = int32(s) - return nil -} - -// validate rejects the option combinations arreflect also rejects. -func (o tagOpts) validate() error { - if o.REE { - return fmt.Errorf("ree is not supported on a struct field; use arreflect.FromSlice with WithREE at the top level") - } - n := 0 - for _, set := range []bool{o.Dict, o.View, o.REE} { - if set { - n++ - } - } - if n > 1 { - return fmt.Errorf("conflicting options: at most one of dict, view, ree may be set") - } - return nil -} diff --git a/arrgen/arrgen/template.go b/arrgen/arrgen/template.go deleted file mode 100644 index 9ab8f9941..000000000 --- a/arrgen/arrgen/template.go +++ /dev/null @@ -1,138 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package arrgen - -import "text/template" - -// fileTemplate renders the whole generated file. The emitted code imports only -// arrow, arrow/array and arrow/memory, plus time for time-of-day columns. -var fileTemplate = template.Must(template.New("arrgen").Parse( - `{{if .Header}}{{.Header}} - -{{end}}// Code generated by arrgen. DO NOT EDIT. - -package {{.Package}} - -import ( -{{- if .NeedTime}} - "time" -{{end}} - "github.com/apache/arrow-go/v18/arrow" - "github.com/apache/arrow-go/v18/arrow/array" - "github.com/apache/arrow-go/v18/arrow/memory" -) -{{range .Types}} -// {{.SchemaVar}} is the Arrow schema of {{.GoName}}, resolved from its arrow -// struct tags. -var {{.SchemaVar}} = arrow.NewSchema([]arrow.Field{ -{{- range .Fields}} - {Name: {{printf "%q" .Name}}, Type: {{.ArrowType}}, Nullable: {{.Nullable}}}, -{{- end}} -}, nil) - -// {{.SchemaFunc}} returns the Arrow schema encoded by {{.AppenderType}}. -func {{.SchemaFunc}}() *arrow.Schema { return {{.SchemaVar}} } - -// {{.AppenderType}} converts {{.GoName}} values into Arrow record batches -// without reflection. -// -// It is not safe for concurrent use. Release it when done. -type {{.AppenderType}} struct { - rb *array.RecordBuilder -{{- range .Fields}} - {{.BuilderVar}} {{.BuilderType}} -{{- end}} - err error -} - -// {{.CtorName}} returns an appender that builds {{.GoName}} batches with mem, -// or memory.DefaultAllocator when mem is nil. The typed field builders are -// resolved once here, so Append does no type assertions. -func {{.CtorName}}(mem memory.Allocator) *{{.AppenderType}} { - if mem == nil { - mem = memory.DefaultAllocator - } - rb := array.NewRecordBuilder(mem, {{.SchemaVar}}) - return &{{.AppenderType}}{ - rb: rb, -{{- range .Fields}} - {{.BuilderVar}}: rb.Field({{.Index}}).({{.BuilderType}}), -{{- end}} - } -} - -// Schema returns the schema of the batches this appender builds. -func (a *{{.AppenderType}}) Schema() *arrow.Schema { return {{.SchemaVar}} } - -// RecordBuilder returns the underlying builder. Appending through it directly -// is allowed as long as every column stays the same length. -func (a *{{.AppenderType}}) RecordBuilder() *array.RecordBuilder { return a.rb } - -// Reserve pre-allocates room for n more rows in every column. Reserving before -// a run of Appends keeps the append path free of allocations. -func (a *{{.AppenderType}}) Reserve(n int) { a.rb.Reserve(n) } - -// Len reports the number of rows appended since the last NewRecordBatch. -func (a *{{.AppenderType}}) Len() int { return a.rb.Field(0).Len() } - -// Append appends one row. v is read synchronously and never retained, so a -// caller can reuse a single {{.GoName}} variable across calls. -func (a *{{.AppenderType}}) Append(v *{{.GoName}}) { -{{- range .Fields}} - {{.AppendStmt}} -{{- end}} -} - -// AppendSlice appends every element of vs, reserving room for them up front. -func (a *{{.AppenderType}}) AppendSlice(vs []{{.GoName}}) { - a.rb.Reserve(len(vs)) - for i := range vs { - a.Append(&vs[i]) - } -} - -// NewRecordBatch returns the appended rows as a record batch and resets the -// appender. The caller must Release the batch. -func (a *{{.AppenderType}}) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } - -// Err returns the first error an Append hit, or nil. Only dictionary columns -// can fail. The error is sticky: NewRecordBatch does not clear it. -func (a *{{.AppenderType}}) Err() error { return a.err } -{{if .AnyFallible}} -// setErr keeps the first error. -func (a *{{.AppenderType}}) setErr(err error) { - if err != nil && a.err == nil { - a.err = err - } -} -{{end}} -// Release releases the appender's builders. -func (a *{{.AppenderType}}) Release() { a.rb.Release() } - -// {{.BatchFunc}} builds one record batch holding every element of vs. It is the -// generated equivalent of arreflect.RecordFromSlice[{{.GoName}}](vs, mem). The -// caller must Release the batch. -func {{.BatchFunc}}(mem memory.Allocator, vs []{{.GoName}}) (arrow.RecordBatch, error) { - a := {{.CtorName}}(mem) - defer a.Release() - a.AppendSlice(vs) - if err := a.Err(); err != nil { - return nil, err - } - return a.NewRecordBatch(), nil -} -{{end}}`)) diff --git a/arrgen/arrgen/testdata/basic.golden b/arrgen/arrgen/testdata/basic.golden deleted file mode 100644 index 1686ca0b8..000000000 --- a/arrgen/arrgen/testdata/basic.golden +++ /dev/null @@ -1,226 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Code generated by arrgen. DO NOT EDIT. - -package basic - -import ( - "github.com/apache/arrow-go/v18/arrow" - "github.com/apache/arrow-go/v18/arrow/array" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// metricArrowSchema is the Arrow schema of Metric, resolved from its arrow -// struct tags. -var metricArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, - {Name: "host", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: false}, - {Name: "count", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, - {Name: "Blob", Type: arrow.BinaryTypes.Binary, Nullable: false}, -}, nil) - -// MetricSchema returns the Arrow schema encoded by MetricAppender. -func MetricSchema() *arrow.Schema { return metricArrowSchema } - -// MetricAppender converts Metric values into Arrow record batches -// without reflection. -// -// It is not safe for concurrent use. Release it when done. -type MetricAppender struct { - rb *array.RecordBuilder - b0 *array.TimestampBuilder - b1 *array.BinaryDictionaryBuilder - b2 *array.Int64Builder - b3 *array.BinaryBuilder - err error -} - -// NewMetricAppender returns an appender that builds Metric batches with mem, -// or memory.DefaultAllocator when mem is nil. The typed field builders are -// resolved once here, so Append does no type assertions. -func NewMetricAppender(mem memory.Allocator) *MetricAppender { - if mem == nil { - mem = memory.DefaultAllocator - } - rb := array.NewRecordBuilder(mem, metricArrowSchema) - return &MetricAppender{ - rb: rb, - b0: rb.Field(0).(*array.TimestampBuilder), - b1: rb.Field(1).(*array.BinaryDictionaryBuilder), - b2: rb.Field(2).(*array.Int64Builder), - b3: rb.Field(3).(*array.BinaryBuilder), - } -} - -// Schema returns the schema of the batches this appender builds. -func (a *MetricAppender) Schema() *arrow.Schema { return metricArrowSchema } - -// RecordBuilder returns the underlying builder. Appending through it directly -// is allowed as long as every column stays the same length. -func (a *MetricAppender) RecordBuilder() *array.RecordBuilder { return a.rb } - -// Reserve pre-allocates room for n more rows in every column. Reserving before -// a run of Appends keeps the append path free of allocations. -func (a *MetricAppender) Reserve(n int) { a.rb.Reserve(n) } - -// Len reports the number of rows appended since the last NewRecordBatch. -func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } - -// Append appends one row. v is read synchronously and never retained, so a -// caller can reuse a single Metric variable across calls. -func (a *MetricAppender) Append(v *Metric) { - a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) - a.setErr(a.b1.AppendString(v.Host)) - if v.Count == nil { - a.b2.AppendNull() - } else { - a.b2.Append(*v.Count) - } - if v.Blob == nil { - a.b3.AppendNull() - } else { - a.b3.Append(v.Blob) - } -} - -// AppendSlice appends every element of vs, reserving room for them up front. -func (a *MetricAppender) AppendSlice(vs []Metric) { - a.rb.Reserve(len(vs)) - for i := range vs { - a.Append(&vs[i]) - } -} - -// NewRecordBatch returns the appended rows as a record batch and resets the -// appender. The caller must Release the batch. -func (a *MetricAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } - -// Err returns the first error an Append hit, or nil. Only dictionary columns -// can fail. The error is sticky: NewRecordBatch does not clear it. -func (a *MetricAppender) Err() error { return a.err } - -// setErr keeps the first error. -func (a *MetricAppender) setErr(err error) { - if err != nil && a.err == nil { - a.err = err - } -} - -// Release releases the appender's builders. -func (a *MetricAppender) Release() { a.rb.Release() } - -// MetricRecordBatch builds one record batch holding every element of vs. It is the -// generated equivalent of arreflect.RecordFromSlice[Metric](vs, mem). The -// caller must Release the batch. -func MetricRecordBatch(mem memory.Allocator, vs []Metric) (arrow.RecordBatch, error) { - a := NewMetricAppender(mem) - defer a.Release() - a.AppendSlice(vs) - if err := a.Err(); err != nil { - return nil, err - } - return a.NewRecordBatch(), nil -} - -// readingArrowSchema is the Arrow schema of reading, resolved from its arrow -// struct tags. -var readingArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "at", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, - {Name: "value", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, -}, nil) - -// readingSchema returns the Arrow schema encoded by readingAppender. -func readingSchema() *arrow.Schema { return readingArrowSchema } - -// readingAppender converts reading values into Arrow record batches -// without reflection. -// -// It is not safe for concurrent use. Release it when done. -type readingAppender struct { - rb *array.RecordBuilder - b0 *array.Date32Builder - b1 *array.Float64Builder - err error -} - -// newReadingAppender returns an appender that builds reading batches with mem, -// or memory.DefaultAllocator when mem is nil. The typed field builders are -// resolved once here, so Append does no type assertions. -func newReadingAppender(mem memory.Allocator) *readingAppender { - if mem == nil { - mem = memory.DefaultAllocator - } - rb := array.NewRecordBuilder(mem, readingArrowSchema) - return &readingAppender{ - rb: rb, - b0: rb.Field(0).(*array.Date32Builder), - b1: rb.Field(1).(*array.Float64Builder), - } -} - -// Schema returns the schema of the batches this appender builds. -func (a *readingAppender) Schema() *arrow.Schema { return readingArrowSchema } - -// RecordBuilder returns the underlying builder. Appending through it directly -// is allowed as long as every column stays the same length. -func (a *readingAppender) RecordBuilder() *array.RecordBuilder { return a.rb } - -// Reserve pre-allocates room for n more rows in every column. Reserving before -// a run of Appends keeps the append path free of allocations. -func (a *readingAppender) Reserve(n int) { a.rb.Reserve(n) } - -// Len reports the number of rows appended since the last NewRecordBatch. -func (a *readingAppender) Len() int { return a.rb.Field(0).Len() } - -// Append appends one row. v is read synchronously and never retained, so a -// caller can reuse a single reading variable across calls. -func (a *readingAppender) Append(v *reading) { - a.b0.Append(arrow.Date32FromTime(v.At)) - a.b1.Append(v.Value) -} - -// AppendSlice appends every element of vs, reserving room for them up front. -func (a *readingAppender) AppendSlice(vs []reading) { - a.rb.Reserve(len(vs)) - for i := range vs { - a.Append(&vs[i]) - } -} - -// NewRecordBatch returns the appended rows as a record batch and resets the -// appender. The caller must Release the batch. -func (a *readingAppender) NewRecordBatch() arrow.RecordBatch { return a.rb.NewRecordBatch() } - -// Err returns the first error an Append hit, or nil. Only dictionary columns -// can fail. The error is sticky: NewRecordBatch does not clear it. -func (a *readingAppender) Err() error { return a.err } - -// Release releases the appender's builders. -func (a *readingAppender) Release() { a.rb.Release() } - -// readingRecordBatch builds one record batch holding every element of vs. It is the -// generated equivalent of arreflect.RecordFromSlice[reading](vs, mem). The -// caller must Release the batch. -func readingRecordBatch(mem memory.Allocator, vs []reading) (arrow.RecordBatch, error) { - a := newReadingAppender(mem) - defer a.Release() - a.AppendSlice(vs) - if err := a.Err(); err != nil { - return nil, err - } - return a.NewRecordBatch(), nil -} diff --git a/arrgen/arrgen/testdata/basic/basic.go b/arrgen/arrgen/testdata/basic/basic.go deleted file mode 100644 index 3bee85d9c..000000000 --- a/arrgen/arrgen/testdata/basic/basic.go +++ /dev/null @@ -1,37 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package basic is the fixture behind arrgen's golden test. It sits under -// testdata so the go tool does not build it; the generator loads it by -// directory. -package basic - -import "time" - -// Metric is an exported type, so its generated API is exported too. -type Metric struct { - Time time.Time `arrow:"ts"` - Host string `arrow:"host,dict"` - Count *int64 `arrow:"count"` - Blob []byte - Local string `arrow:"-"` -} - -// reading is unexported, so the generated appender is unexported as well. -type reading struct { - At time.Time `arrow:"at,date32"` - Value float64 `arrow:"value"` -} diff --git a/arrgen/arrgen/testdata/errors/errors.go b/arrgen/arrgen/testdata/errors/errors.go deleted file mode 100644 index 0a4472008..000000000 --- a/arrgen/arrgen/testdata/errors/errors.go +++ /dev/null @@ -1,126 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package errors holds one struct per rejection the generator makes. -package errors - -import "time" - -// Inner is a nested struct used by Nested and Embedded. -type Inner struct { - A int64 -} - -// ID is a defined scalar type, which arreflect cannot map either. -type ID int64 - -type Embedded struct { - Inner - B int64 -} - -type Nested struct { - Field Inner -} - -type SliceField struct { - Field []int64 -} - -type ArrayField struct { - Field [3]int64 -} - -type MapField struct { - Field map[string]int64 -} - -type NamedScalar struct { - Field ID -} - -type UnknownOption struct { - Field int64 `arrow:"f,nope"` -} - -type BadDecimalTag struct { - Field int64 `arrow:"f,decimal(a,2)"` -} - -type ShortDecimalTag struct { - Field int64 `arrow:"f,decimal(4)"` -} - -type DuplicateNames struct { - A int64 `arrow:"same"` - B int64 `arrow:"same"` -} - -type TemporalOnInt struct { - Field int64 `arrow:"f,date32"` -} - -type DecimalOnInt struct { - Field int64 `arrow:"f,decimal(10,2)"` -} - -type DictOnBool struct { - Field bool `arrow:"f,dict"` -} - -type ViewOnInt struct { - Field int64 `arrow:"f,view"` -} - -type LargeOnInt struct { - Field int64 `arrow:"f,large"` -} - -type RunEndEncoded struct { - Field string `arrow:"f,ree"` -} - -type DictAndView struct { - Field string `arrow:"f,dict,view"` -} - -type DictAndLarge struct { - Field string `arrow:"f,dict,large"` -} - -type NoColumns struct { - Field string `arrow:"-"` - hidden int64 -} - -type DoublePointer struct { - Field **int64 -} - -type TimeSlice struct { - Field []time.Time -} - -// NotAStruct is a named type whose underlying type is not a struct. -type NotAStruct int64 - -type Collide struct { - Field int64 -} - -type collide struct { - Field int64 -} diff --git a/arrgen/arrgen/testdata/pregen/pregen.go b/arrgen/arrgen/testdata/pregen/pregen.go deleted file mode 100644 index 799ac0b68..000000000 --- a/arrgen/arrgen/testdata/pregen/pregen.go +++ /dev/null @@ -1,37 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package pregen is the state a package is in the first time the generator runs -// against it: the call sites exist, the generated file does not, and so the -// package does not type-check. -package pregen - -import ( - "time" - - "github.com/apache/arrow-go/v18/arrow" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -type Sample struct { - At time.Time `arrow:"at"` - Value float64 `arrow:"value"` -} - -// Encode calls into the file arrgen has not written yet. -func Encode(mem memory.Allocator, samples []Sample) (arrow.RecordBatch, error) { - return SampleRecordBatch(mem, samples) -} diff --git a/arrgen/doc.go b/arrgen/doc.go index f8e27cd79..df77b2ecc 100644 --- a/arrgen/doc.go +++ b/arrgen/doc.go @@ -29,14 +29,19 @@ // // # Usage // -// Put a go:generate directive next to the struct: +// Add the generator to your module as a tool dependency, then put a go:generate +// directive next to the struct: // -// //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric +// go get -tool github.com/apache/arrow-go/arrgen/cmd/arrgen +// +// //go:generate go tool arrgen -type Metric // // then run `go generate ./...`, which writes metric_arrow.go next to the type. // Nothing in the arrow-go module depends on arrgen, and the code it emits // imports only arrow, arrow/array and arrow/memory, so the generator itself is -// a build-time-only dependency. +// a build-time-only dependency. The unversioned `go run ` spelling this +// module's own directives use resolves only inside a module that already +// requires arrgen; see the README for the alternatives. // // # Supported field types // @@ -46,8 +51,18 @@ // nullable column. // // Tag options mirror arreflect: a leading name, "-" to skip a field, the -// temporal overrides date32, date64, time32, time64 and timestamp, dict, view, -// large, and decimal(precision,scale). +// temporal overrides date32, date64, time32 and time64, dict, view, large, and +// decimal(precision,scale). +// +// A time.Time field must carry one of the four temporal tags, and a +// decimal128.Num or decimal256.Num field must carry a decimal(precision,scale) +// tag. Untagged, all three are Go structs that arreflect's inferArrowType +// resolves through inferStructType - it switches on reflect.Kind before it +// reaches the types it matches by identity - so it infers an empty struct<> and +// drops the value. Generating the column Arrow means here would put the two +// paths out of step, so arrgen rejects those spellings instead. One consequence +// is that arrgen cannot emit a TIMESTAMP column at all: ",timestamp" is +// rejected along with the untagged spelling. // // Anything arrgen cannot map exactly the way arreflect would - nested structs, // slices other than []byte, arrays, maps, embedded fields, named scalar types - diff --git a/arrgen/drift_test.go b/arrgen/drift_test.go index 40f840460..724c0ad90 100644 --- a/arrgen/drift_test.go +++ b/arrgen/drift_test.go @@ -45,7 +45,7 @@ func TestCheckedInFilesAreUpToDate(t *testing.T) { types []string file string }{ - {"example", []string{"Metric"}, "metric_arrow.go"}, + {"internal/gentypes", []string{"Metric"}, "metric_arrow.go"}, {"internal/gentypes", []string{"Row", "Fixed"}, "row_arrow.go"}, } diff --git a/arrgen/example/example_test.go b/arrgen/example/example_test.go deleted file mode 100644 index dd2f3c644..000000000 --- a/arrgen/example/example_test.go +++ /dev/null @@ -1,107 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package example_test - -import ( - "fmt" - "time" - - "github.com/apache/arrow-go/arrgen/example" - "github.com/apache/arrow-go/v18/arrow/memory" -) - -// Example encodes a slice of rows in one call. It is the drop-in replacement -// for arreflect.RecordFromSlice[Metric](metrics, mem). -func Example() { - cpuLoad := 0.75 - metrics := []example.Metric{ - {Time: time.Unix(0, 0).UTC(), Host: "web-1", CPU: 0.5, Value: &cpuLoad}, - {Time: time.Unix(1, 0).UTC(), Host: "web-2", CPU: 1.5}, // Value is nil: a null - } - - rec, err := example.MetricRecordBatch(memory.DefaultAllocator, metrics) - if err != nil { - fmt.Println("encode:", err) - return - } - defer rec.Release() - - fmt.Println("rows:", rec.NumRows()) - for i, col := range rec.Columns() { - fmt.Printf("%s: %v\n", rec.Schema().Field(i).Name, col) - } - // Output: - // rows: 2 - // ts: [0 1000000000] - // host: ["web-1" "web-2"] - // cpu: [0.5 1.5] - // value: [0.75 (null)] -} - -// ExampleMetricAppender streams rows in one at a time and cuts a batch at a -// size boundary, which is what an ingest loop actually looks like. The appender -// reads each row synchronously and never keeps a reference, so the same -// variable is reused for every row and the loop allocates nothing. -func ExampleMetricAppender() { - const batchSize = 2 - - a := example.NewMetricAppender(memory.DefaultAllocator) - defer a.Release() - a.Reserve(batchSize) - - var row example.Metric - for i := 0; i < 5; i++ { - row = example.Metric{ - Time: time.Unix(int64(i), 0).UTC(), - Host: fmt.Sprintf("web-%d", i), - CPU: float64(i), - } - a.Append(&row) - - if a.Len() == batchSize { - rec := a.NewRecordBatch() - fmt.Println("batch of", rec.NumRows()) - rec.Release() - a.Reserve(batchSize) - } - } - if a.Len() > 0 { - rec := a.NewRecordBatch() - fmt.Println("final batch of", rec.NumRows()) - rec.Release() - } - if err := a.Err(); err != nil { - fmt.Println("append:", err) - } - // Output: - // batch of 2 - // batch of 2 - // final batch of 1 -} - -// ExampleMetricSchema shows the schema the generator resolved from the tags. -// It is a package-level variable, so asking for it costs nothing. -func ExampleMetricSchema() { - for _, f := range example.MetricSchema().Fields() { - fmt.Printf("%s %s nullable=%t\n", f.Name, f.Type, f.Nullable) - } - // Output: - // ts timestamp[ns, tz=UTC] nullable=false - // host utf8 nullable=false - // cpu float64 nullable=false - // value float64 nullable=true -} diff --git a/arrgen/example/metric.go b/arrgen/example/metric.go deleted file mode 100644 index b15679191..000000000 --- a/arrgen/example/metric.go +++ /dev/null @@ -1,42 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package example is a worked, checked-in demonstration of arrgen. -// -// Metric carries the arrow struct tags; metric_arrow.go next to it is what -// arrgen emitted from them and is committed exactly as stringer or easyjson -// output would be. Regenerate it with: -// -// go generate ./... -// -// The package is importable rather than a main package so its tests can encode -// the same values through both the generated appender and arreflect and show -// they agree. -package example - -import "time" - -// Metric is a small, realistic row: a timestamp, two scalar columns, a nullable -// pointer column, and a field kept out of Arrow entirely. -type Metric struct { - Time time.Time `arrow:"ts"` - Host string `arrow:"host"` - CPU float64 `arrow:"cpu"` - Value *float64 `arrow:"value"` // nullable: a nil pointer appends null - Secret string `arrow:"-"` // never leaves the process -} - -//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -header ../license_header.txt diff --git a/arrgen/example_test.go b/arrgen/example_test.go new file mode 100644 index 000000000..e39cacf9e --- /dev/null +++ b/arrgen/example_test.go @@ -0,0 +1,166 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package arrgen_test + +import ( + "fmt" + "go/ast" + "go/parser" + "go/token" + "time" + + "github.com/apache/arrow-go/arrgen" + "github.com/apache/arrow-go/arrgen/internal/gentypes" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// The examples below are driven by internal/gentypes.Metric, this module's +// checked-in fixture: +// +// type Metric struct { +// Day time.Time `arrow:"day,date32"` +// Host string `arrow:"host"` +// CPU float64 `arrow:"cpu"` +// Value *float64 `arrow:"value"` // nullable: a nil pointer appends null +// Secret string `arrow:"-"` // never leaves the process +// } +// +// Running arrgen over it emits MetricSchema, NewMetricAppender and +// MetricRecordBatch. In your own package the names are derived from your own +// struct the same way. + +// ExampleGenerate runs the generator the way the arrgen command does and lists +// the API it emitted. Generate returns formatted source and writes nothing, so +// a caller can diff or inspect the result before it lands on disk. +func ExampleGenerate() { + src, err := arrgen.Generate(arrgen.Config{ + Dir: "testdata/basic", + Types: []string{"Metric"}, + }) + if err != nil { + fmt.Println("generate:", err) + return + } + + file, err := parser.ParseFile(token.NewFileSet(), "metric_arrow.go", src, 0) + if err != nil { + fmt.Println("parse:", err) + return + } + for _, decl := range file.Decls { + switch d := decl.(type) { + case *ast.FuncDecl: + if d.Recv == nil { + fmt.Println("func", d.Name.Name) + } + case *ast.GenDecl: + if d.Tok == token.TYPE { + for _, spec := range d.Specs { + fmt.Println("type", spec.(*ast.TypeSpec).Name.Name) + } + } + } + } + // Output: + // func MetricSchema + // type MetricAppender + // func NewMetricAppender + // func MetricRecordBatch +} + +// Example encodes a slice of rows in one call, replacing +// arreflect.RecordFromSlice[Metric](metrics, mem). +func Example() { + cpuLoad := 0.75 + day := time.Date(2024, time.March, 17, 0, 0, 0, 0, time.UTC) + metrics := []gentypes.Metric{ + {Day: day, Host: "web-1", CPU: 0.5, Value: &cpuLoad}, + {Day: day.AddDate(0, 0, 1), Host: "web-2", CPU: 1.5}, // Value is nil: a null + } + + rec, err := gentypes.MetricRecordBatch(memory.DefaultAllocator, metrics) + if err != nil { + fmt.Println("encode:", err) + return + } + defer rec.Release() + + fmt.Println("rows:", rec.NumRows()) + for i, col := range rec.Columns() { + fmt.Printf("%s: %v\n", rec.Schema().Field(i).Name, col) + } + // Output: + // rows: 2 + // day: [19799 19800] + // host: ["web-1" "web-2"] + // cpu: [0.5 1.5] + // value: [0.75 (null)] +} + +// Example_streamingAppend streams rows in one at a time and cuts a batch at a +// size boundary. The appender never retains a row, so the loop reuses a single +// variable and allocates nothing. +func Example_streamingAppend() { + const batchSize = 2 + day := time.Date(2024, time.March, 17, 0, 0, 0, 0, time.UTC) + + a := gentypes.NewMetricAppender(memory.DefaultAllocator) + defer a.Release() + a.Reserve(batchSize) + + var row gentypes.Metric + for i := 0; i < 5; i++ { + row = gentypes.Metric{ + Day: day.AddDate(0, 0, i), + Host: fmt.Sprintf("web-%d", i), + CPU: float64(i), + } + a.Append(&row) + + if a.Len() == batchSize { + rec := a.NewRecordBatch() + fmt.Println("batch of", rec.NumRows()) + rec.Release() + a.Reserve(batchSize) + } + } + if a.Len() > 0 { + rec := a.NewRecordBatch() + fmt.Println("final batch of", rec.NumRows()) + rec.Release() + } + if err := a.Err(); err != nil { + fmt.Println("append:", err) + } + // Output: + // batch of 2 + // batch of 2 + // final batch of 1 +} + +// Example_schema shows the schema the generator resolved from the tags. It is +// a package-level value in the generated file, so reading it costs nothing. +func Example_schema() { + for _, f := range gentypes.MetricSchema().Fields() { + fmt.Printf("%s %s nullable=%t\n", f.Name, f.Type, f.Nullable) + } + // Output: + // day date32 nullable=false + // host utf8 nullable=false + // cpu float64 nullable=false + // value float64 nullable=true +} diff --git a/arrgen/generate_test.go b/arrgen/generate_test.go index d2594b7d0..26a229de5 100644 --- a/arrgen/generate_test.go +++ b/arrgen/generate_test.go @@ -23,8 +23,13 @@ import ( "path/filepath" "strings" "testing" + "time" "github.com/apache/arrow-go/arrgen" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array/arreflect" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/decimal256" ) var update = flag.Bool("update", false, "rewrite the golden files instead of comparing against them") @@ -121,6 +126,10 @@ func TestGenerateRejects(t *testing.T) { {"RunEndEncoded", "ree is not supported"}, {"DictAndView", "at most one of dict, view, ree"}, {"DictAndLarge", "cannot be combined with large"}, + {"BareTime", "one of the date32, date64, time32 or time64 tags"}, + {"TimestampTime", "one of the date32, date64, time32 or time64 tags"}, + {"BareDecimal128", "decimal128.Num needs a decimal(precision,scale) tag"}, + {"BareDecimal256", "decimal256.Num needs a decimal(precision,scale) tag"}, {"NoColumns", "no Arrow columns"}, {"DoublePointer", "multi-level pointer"}, {"TimeSlice", "list column"}, @@ -220,6 +229,40 @@ func TestGenerateMissingPackage(t *testing.T) { } } +// TestArreflectCannotInferStructScalars pins the upstream behavior that the +// rejections above exist for. +// +// time.Time, decimal128.Num and decimal256.Num are Go structs that Arrow models +// as scalars. arreflect's inferArrowType switches on reflect.Kind before it +// reaches the types it matches by identity, so as a struct field each one is +// resolved by inferStructType, which sees only unexported fields and yields an +// empty struct<>. arrgen refuses to generate the column Arrow means here +// because doing so would break the equivalence it promises. +// +// If arreflect learns to infer these, this test fails - which is good news, not +// a regression. It means the rejections in mapping.go can be dropped and the +// untagged spellings generated instead. +func TestArreflectCannotInferStructScalars(t *testing.T) { + type row struct { + Time time.Time `arrow:"t"` + TS time.Time `arrow:"ts,timestamp"` + D128 decimal128.Num `arrow:"d128"` + D256 decimal256.Num `arrow:"d256"` + } + + schema, err := arreflect.InferSchema[row]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + for i := 0; i < schema.NumFields(); i++ { + f := schema.Field(i) + if f.Type.ID() != arrow.STRUCT { + t.Errorf("arreflect now infers column %q as %s rather than an empty struct; "+ + "drop the matching rejection in mapping.go and generate the column instead", f.Name, f.Type) + } + } +} + func first(b []byte, n int) []byte { if len(b) < n { return b diff --git a/arrgen/internal/gentypes/bench_test.go b/arrgen/internal/gentypes/bench_test.go index 94a0bbb49..81b385e3b 100644 --- a/arrgen/internal/gentypes/bench_test.go +++ b/arrgen/internal/gentypes/bench_test.go @@ -19,9 +19,7 @@ package gentypes_test import ( "fmt" "testing" - "time" - "github.com/apache/arrow-go/arrgen/example" "github.com/apache/arrow-go/arrgen/internal/gentypes" "github.com/apache/arrow-go/v18/arrow/array/arreflect" "github.com/apache/arrow-go/v18/arrow/memory" @@ -31,12 +29,12 @@ import ( // that per-batch setup does not dominate and small enough to stay in cache. const benchRows = 1024 -func makeMetrics(n int) []example.Metric { - metrics := make([]example.Metric, n) +func makeMetrics(n int) []gentypes.Metric { + metrics := make([]gentypes.Metric, n) v := 42.5 for i := range metrics { - metrics[i] = example.Metric{ - Time: base.Add(time.Duration(i) * time.Second), + metrics[i] = gentypes.Metric{ + Day: base.AddDate(0, 0, i), Host: fmt.Sprintf("host-%d", i%16), CPU: float64(i) / 100, } @@ -52,10 +50,10 @@ func makeFixed(n int) []gentypes.Fixed { v := 1.5 for i := range rows { rows[i] = gentypes.Fixed{ - Timestamp: base.Add(time.Duration(i) * time.Second), - ID: int64(i), - Value: float64(i) * 1.5, - OK: i%2 == 0, + Day: base.AddDate(0, 0, i), + ID: int64(i), + Value: float64(i) * 1.5, + OK: i%2 == 0, } if i%3 != 0 { rows[i].Optional = &v @@ -85,7 +83,7 @@ func BenchmarkMetricBatch(b *testing.B) { b.Run("arrgen", func(b *testing.B) { b.ReportAllocs() for b.Loop() { - rec, err := example.MetricRecordBatch(mem, metrics) + rec, err := gentypes.MetricRecordBatch(mem, metrics) if err != nil { b.Fatal(err) } @@ -96,7 +94,7 @@ func BenchmarkMetricBatch(b *testing.B) { } // BenchmarkRowBatch runs the same comparison over the wide fixture, where a -// row costs 46 columns of reflection instead of four. +// row costs 45 columns of reflection instead of four. func BenchmarkRowBatch(b *testing.B) { rows := makeRows(benchRows) mem := memory.DefaultAllocator diff --git a/arrgen/internal/gentypes/equivalence_test.go b/arrgen/internal/gentypes/equivalence_test.go index b86a3ff8e..5de1eb21c 100644 --- a/arrgen/internal/gentypes/equivalence_test.go +++ b/arrgen/internal/gentypes/equivalence_test.go @@ -22,7 +22,6 @@ import ( "testing" "time" - "github.com/apache/arrow-go/arrgen/example" "github.com/apache/arrow-go/arrgen/internal/gentypes" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" @@ -33,35 +32,10 @@ import ( "github.com/apache/arrow-go/v18/arrow/memory" ) -// arreflectStructFieldGaps names the Row columns whose Go type arreflect cannot -// encode when it appears as a struct field, and what it emits instead. +// compareColumns asserts that got and want hold the same columns. // -// Its inferArrowType dispatches on reflect.Kind before it checks for the types -// it special-cases, so any Go struct that Arrow models as a scalar - time.Time, -// decimal128.Num, decimal256.Num - reaches inferStructType, which finds only -// unexported fields and returns an empty struct<>. The value is then dropped on -// the floor. Its own tests only ever pass these types as the top-level element -// of a slice, where a different branch handles them correctly, which is why the -// gap has gone unnoticed. -// -// arrgen emits the column arreflect's own InferType and FromSlice paths agree -// on (TIMESTAMP(ns, UTC), DECIMAL128, DECIMAL256), so these columns are the one -// place the two paths disagree and are compared separately, by value, in -// TestGapColumnValues. TestArreflectStructFieldGaps pins the current arreflect -// behavior so that fixing it upstream turns into a failure here telling us to -// shorten this list. -// A tag that names an explicit Arrow type - date32, decimal(20,3) - rescues the -// field, because arreflect overrides the inferred type from the tag afterwards. -// Only the untagged spellings are listed here. -var arreflectStructFieldGaps = map[string]string{ - "ts": "time.Time", - "pts": "*time.Time", - "dec256": "decimal256.Num", - "pdec128": "*decimal128.Num", -} - -// compareColumns asserts that got and want hold the same columns, skipping the -// ones arreflect cannot encode as struct fields. +// It walks the columns rather than calling array.RecordEqual so a failure names +// the column that differs instead of dumping two whole batches. // // dictsByValue relaxes dictionary columns to a comparison of the values they // decode to. An appender reused across batches carries its memo table forward, @@ -78,9 +52,6 @@ func compareColumns(t *testing.T, label string, got, want arrow.RecordBatch, dic } for i := 0; i < int(got.NumCols()); i++ { name := got.Schema().Field(i).Name - if _, gap := arreflectStructFieldGaps[name]; gap { - continue - } gc, wc := got.Column(i), want.Column(i) if dictsByValue && gc.DataType().ID() == arrow.DICTIONARY { compareByValueStr(t, label, name, gc, wc) @@ -127,41 +98,40 @@ func makeRows(n int) []gentypes.Row { for i := range rows { null := i%3 == 0 r := gentypes.Row{ - Bool: i%2 == 0, - Int8: int8(i), - Int16: int16(i * 3), - Int32: int32(i * 7), - Int64: int64(i) * 1e6, - Int: i, - Uint8: uint8(i), - Uint16: uint16(i * 5), - Uint32: uint32(i * 11), - Uint64: uint64(i) * 1e9, - Uint: uint(i), - Float32: float32(i) + 0.25, - Float64: float64(i) * math.Pi, - Str: fmt.Sprintf("host-%d", i%7), - Bin: []byte{byte(i), byte(i >> 8)}, - Timestamp: base.Add(time.Duration(i) * time.Second), - Date32: base.AddDate(0, 0, i), - Date64: base.AddDate(0, 0, i), - Time32: base.Add(time.Duration(i) * time.Minute), - Time64: base.Add(time.Duration(i) * time.Minute), - Duration: time.Duration(i) * time.Millisecond, - Dec32: decimal.Decimal32(i * 100), - Dec64: decimal.Decimal64(i * 10000), - Dec128: decimal128.FromU64(uint64(i) * 1000), - Dec256: decimal256.FromU64(uint64(i) * 2000), - LargeStr: fmt.Sprintf("large-%d", i), - ViewStr: fmt.Sprintf("view-%d", i), - LargeBin: []byte(fmt.Sprintf("lb-%d", i)), - ViewBin: []byte(fmt.Sprintf("vb-%d", i)), - DictStr: fmt.Sprintf("region-%d", i%3), - DictBin: []byte(fmt.Sprintf("k%d", i%4)), - DictInt: int32(i % 5), - DictF64: float64(i % 6), - Untagged: int64(i), - Secret: "never encoded", + Bool: i%2 == 0, + Int8: int8(i), + Int16: int16(i * 3), + Int32: int32(i * 7), + Int64: int64(i) * 1e6, + Int: i, + Uint8: uint8(i), + Uint16: uint16(i * 5), + Uint32: uint32(i * 11), + Uint64: uint64(i) * 1e9, + Uint: uint(i), + Float32: float32(i) + 0.25, + Float64: float64(i) * math.Pi, + Str: fmt.Sprintf("host-%d", i%7), + Bin: []byte{byte(i), byte(i >> 8)}, + Date32: base.AddDate(0, 0, i), + Date64: base.AddDate(0, 0, i), + Time32: base.Add(time.Duration(i) * time.Minute), + Time64: base.Add(time.Duration(i) * time.Minute), + Duration: time.Duration(i) * time.Millisecond, + Dec32: decimal.Decimal32(i * 100), + Dec64: decimal.Decimal64(i * 10000), + Dec128: decimal128.FromU64(uint64(i) * 1000), + Dec256: decimal256.FromU64(uint64(i) * 2000), + LargeStr: fmt.Sprintf("large-%d", i), + ViewStr: fmt.Sprintf("view-%d", i), + LargeBin: []byte(fmt.Sprintf("lb-%d", i)), + ViewBin: []byte(fmt.Sprintf("vb-%d", i)), + DictStr: fmt.Sprintf("region-%d", i%3), + DictBin: []byte(fmt.Sprintf("k%d", i%4)), + DictInt: int32(i % 5), + DictF64: float64(i % 6), + Untagged: int64(i), + Secret: "never encoded", } if null { // Leave every pointer nil, and null out the two shapes that map to @@ -184,7 +154,7 @@ func makeRows(n int) []gentypes.Row { ds := fmt.Sprintf("pd-%d", i%2) r.PBool, r.PInt64, r.PF64, r.PStr, r.PBin = &b, &i64, &f64, &s, &bin dec128 := decimal128.FromU64(uint64(i)) - r.PTS, r.PDate32, r.PTime64, r.PDur, r.PDec32, r.PDictS = &ts, &ts, &ts, &dur, &dec, &ds + r.PDate32, r.PDate64, r.PTime64, r.PDur, r.PDec32, r.PDictS = &ts, &ts, &ts, &dur, &dec, &ds r.PDec128 = &dec128 rows[i] = r } @@ -199,7 +169,30 @@ func TestRowSchemaMatchesArreflect(t *testing.T) { if err != nil { t.Fatalf("InferSchema: %v", err) } - got := gentypes.RowSchema() + compareSchemas(t, gentypes.RowSchema(), want) +} + +func TestMetricSchemaMatchesArreflect(t *testing.T) { + want, err := arreflect.InferSchema[gentypes.Metric]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + compareSchemas(t, gentypes.MetricSchema(), want) +} + +func TestFixedSchemaMatchesArreflect(t *testing.T) { + want, err := arreflect.InferSchema[gentypes.Fixed]() + if err != nil { + t.Fatalf("InferSchema: %v", err) + } + compareSchemas(t, gentypes.FixedSchema(), want) +} + +// compareSchemas asserts got and want are the same schema, field for field. +// Schema.Equal would answer the same question in one call, but a mismatch on a +// 45-column fixture is only actionable if the failure names the column. +func compareSchemas(t *testing.T, got, want *arrow.Schema) { + t.Helper() if got.NumFields() != want.NumFields() { t.Fatalf("field count = %d, want %d\n got: %s\nwant: %s", got.NumFields(), want.NumFields(), got, want) } @@ -209,30 +202,12 @@ func TestRowSchemaMatchesArreflect(t *testing.T) { t.Errorf("field %d: name = %q, want %q", i, gf.Name, wf.Name) continue } - if _, gap := arreflectStructFieldGaps[gf.Name]; gap { - continue - } if !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { t.Errorf("field %q: got %s nullable=%t, want %s nullable=%t", gf.Name, gf.Type, gf.Nullable, wf.Type, wf.Nullable) } } -} - -func TestMetricSchemaMatchesArreflect(t *testing.T) { - want, err := arreflect.InferSchema[example.Metric]() - if err != nil { - t.Fatalf("InferSchema: %v", err) - } - got := example.MetricSchema() - // "ts" is a bare time.Time: see arreflectStructFieldGaps. - for i := 0; i < got.NumFields(); i++ { - gf, wf := got.Field(i), want.Field(i) - if gf.Name == "ts" { - continue - } - if gf.Name != wf.Name || !arrow.TypeEqual(gf.Type, wf.Type) || gf.Nullable != wf.Nullable { - t.Errorf("field %d: got %s, want %s", i, gf, wf) - } + if !got.Equal(want) { + t.Errorf("schemas differ beyond their fields\n got: %s\nwant: %s", got, want) } } @@ -241,9 +216,8 @@ func TestMetricSchemaMatchesArreflect(t *testing.T) { // columns. Everything else in this package - the benchmarks, the allocation // assertions - is only interesting because this holds. // -// It compares column by column rather than with array.RecordEqual so that the -// handful of columns arreflect cannot encode as struct fields can be skipped by -// name instead of poisoning the whole comparison; see arreflectStructFieldGaps. +// Every column is compared: arrgen rejects at generate time any field it +// cannot map the way arreflect would, so there is nothing here to excuse. func TestRowRecordMatchesArreflect(t *testing.T) { for _, n := range []int{0, 1, 2, 3, 64} { t.Run(fmt.Sprintf("rows=%d", n), func(t *testing.T) { @@ -269,94 +243,14 @@ func TestRowRecordMatchesArreflect(t *testing.T) { } } -// TestArreflectStructFieldGaps pins the upstream behavior that -// arreflectStructFieldGaps exists to work around. If arreflect learns to encode -// these struct fields, this test fails and the skip list should shrink to match -// - a failure here is good news, not a regression. -func TestArreflectStructFieldGaps(t *testing.T) { - schema, err := arreflect.InferSchema[gentypes.Row]() - if err != nil { - t.Fatalf("InferSchema: %v", err) - } - for name, goType := range arreflectStructFieldGaps { - idx := schema.FieldIndices(name) - if len(idx) != 1 { - t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) - } - if dt := schema.Field(idx[0]).Type; dt.ID() != arrow.STRUCT { - t.Errorf("arreflect now infers %s column %q as %s rather than an empty struct; "+ - "remove it from arreflectStructFieldGaps so the strict comparison covers it", goType, name, dt) - } - } -} - -// TestGapColumnValues covers the columns TestRowRecordMatchesArreflect has to -// skip. There is no reflection-built batch to compare them against, so the -// expected Arrow values are spelled out here directly. -func TestGapColumnValues(t *testing.T) { - mem := memory.NewCheckedAllocator(memory.DefaultAllocator) - defer mem.AssertSize(t, 0) - - ts := base - dec128 := decimal128.FromU64(12345) - dec256 := decimal256.FromU64(67890) - rows := []gentypes.Row{ - {Timestamp: ts, Dec128: dec128, Dec256: dec256, PTS: &ts, PDec128: &dec128}, - {Timestamp: ts.Add(time.Second), Dec128: dec128, Dec256: dec256}, // pointers nil - } - - rec, err := gentypes.RowRecordBatch(mem, rows) - if err != nil { - t.Fatalf("RowRecordBatch: %v", err) - } - defer rec.Release() - - col := func(name string) arrow.Array { - idx := rec.Schema().FieldIndices(name) - if len(idx) != 1 { - t.Fatalf("column %q: found %d fields, want 1", name, len(idx)) - } - return rec.Column(idx[0]) - } - - tsCol := col("ts").(*array.Timestamp) - if got, want := tsCol.Value(0), arrow.Timestamp(ts.UnixNano()); got != want { - t.Errorf("ts[0] = %d, want %d", got, want) - } - if got, want := tsCol.Value(1), arrow.Timestamp(ts.Add(time.Second).UnixNano()); got != want { - t.Errorf("ts[1] = %d, want %d", got, want) - } - if dt := tsCol.DataType().(*arrow.TimestampType); dt.Unit != arrow.Nanosecond || dt.TimeZone != "UTC" { - t.Errorf("ts type = %s, want timestamp[ns, tz=UTC]", dt) - } - - ptsCol := col("pts").(*array.Timestamp) - if ptsCol.IsNull(0) { - t.Error("pts[0] is null, want a value") - } - if !ptsCol.IsNull(1) { - t.Error("pts[1] is not null, want null") - } - - if got := col("dec128").(*array.Decimal128).Value(0); got != dec128 { - t.Errorf("dec128[0] = %v, want %v", got, dec128) - } - if got := col("dec256").(*array.Decimal256).Value(0); got != dec256 { - t.Errorf("dec256[0] = %v, want %v", got, dec256) - } - if !col("pdec128").(*array.Decimal128).IsNull(1) { - t.Error("pdec128[1] is not null, want null") - } -} - func TestMetricRecordMatchesArreflect(t *testing.T) { mem := memory.NewCheckedAllocator(memory.DefaultAllocator) defer mem.AssertSize(t, 0) v := 3.5 - metrics := []example.Metric{ - {Time: base, Host: "a", CPU: 0.5, Value: &v, Secret: "hidden"}, - {Time: base.Add(time.Second), Host: "b", CPU: 1.5, Value: nil}, + metrics := []gentypes.Metric{ + {Day: base, Host: "a", CPU: 0.5, Value: &v, Secret: "hidden"}, + {Day: base.AddDate(0, 0, 1), Host: "b", CPU: 1.5, Value: nil}, } want, err := arreflect.RecordFromSlice(metrics, mem) @@ -365,26 +259,13 @@ func TestMetricRecordMatchesArreflect(t *testing.T) { } defer want.Release() - got, err := example.MetricRecordBatch(mem, metrics) + got, err := gentypes.MetricRecordBatch(mem, metrics) if err != nil { t.Fatalf("MetricRecordBatch: %v", err) } defer got.Release() - // "ts" is a bare time.Time, one of the columns arreflect drops as a struct - // field; the rest must match it exactly. See arreflectStructFieldGaps. - for i := 0; i < int(got.NumCols()); i++ { - name := got.Schema().Field(i).Name - if name == "ts" { - continue - } - if !array.Equal(got.Column(i), want.Column(i)) { - t.Errorf("column %q differs\n got: %s\nwant: %s", name, got.Column(i), want.Column(i)) - } - } - if got, want := got.Column(0).(*array.Timestamp).Value(0), arrow.Timestamp(base.UnixNano()); got != want { - t.Errorf("ts[0] = %d, want %d", got, want) - } + compareColumns(t, "Metric", got, want, false) if secrets := got.Schema().FieldIndices("Secret"); len(secrets) != 0 { t.Errorf(`a field tagged arrow:"-" was encoded as column %v`, secrets) } diff --git a/arrgen/arrgen/example/metric.go b/arrgen/internal/gentypes/metric.go similarity index 69% rename from arrgen/arrgen/example/metric.go rename to arrgen/internal/gentypes/metric.go index 60bf341f8..6a0ff79ee 100644 --- a/arrgen/arrgen/example/metric.go +++ b/arrgen/internal/gentypes/metric.go @@ -14,23 +14,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package example is a worked, checked-in demonstration of arrgen. -// -// Metric carries the arrow struct tags; metric_arrow.go next to it is what -// arrgen emitted from them, committed as stringer or easyjson output would be. -// Regenerate it with `go generate ./...`. -package example +package gentypes import "time" -// Metric is a small row: a timestamp, two scalar columns, a nullable pointer -// column, and a field kept out of Arrow entirely. +// Metric is the small, readable fixture behind arrgen's package examples: a +// date column, two scalar columns, a nullable pointer column, and a field kept +// out of Arrow entirely. metric_arrow.go next to it is what arrgen emitted from +// these tags, committed as stringer or easyjson output would be. type Metric struct { - Time time.Time `arrow:"ts"` + Day time.Time `arrow:"day,date32"` Host string `arrow:"host"` CPU float64 `arrow:"cpu"` Value *float64 `arrow:"value"` // nullable: a nil pointer appends null Secret string `arrow:"-"` // never leaves the process } -//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -header ../license_header.txt +//go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Metric -header ../../license_header.txt diff --git a/arrgen/example/metric_arrow.go b/arrgen/internal/gentypes/metric_arrow.go similarity index 95% rename from arrgen/example/metric_arrow.go rename to arrgen/internal/gentypes/metric_arrow.go index 136666aa6..b0d1aae5a 100644 --- a/arrgen/example/metric_arrow.go +++ b/arrgen/internal/gentypes/metric_arrow.go @@ -16,7 +16,7 @@ // Code generated by arrgen. DO NOT EDIT. -package example +package gentypes import ( "github.com/apache/arrow-go/v18/arrow" @@ -27,7 +27,7 @@ import ( // metricArrowSchema is the Arrow schema of Metric, resolved from its arrow // struct tags at generate time. var metricArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "day", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, {Name: "host", Type: arrow.BinaryTypes.String, Nullable: false}, {Name: "cpu", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, {Name: "value", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, @@ -44,7 +44,7 @@ func MetricSchema() *arrow.Schema { return metricArrowSchema } // It is not safe for concurrent use. Release it when done. type MetricAppender struct { rb *array.RecordBuilder - b0 *array.TimestampBuilder + b0 *array.Date32Builder b1 *array.StringBuilder b2 *array.Float64Builder b3 *array.Float64Builder @@ -63,7 +63,7 @@ func NewMetricAppender(mem memory.Allocator) *MetricAppender { rb := array.NewRecordBuilder(mem, metricArrowSchema) return &MetricAppender{ rb: rb, - b0: rb.Field(0).(*array.TimestampBuilder), + b0: rb.Field(0).(*array.Date32Builder), b1: rb.Field(1).(*array.StringBuilder), b2: rb.Field(2).(*array.Float64Builder), b3: rb.Field(3).(*array.Float64Builder), @@ -88,7 +88,7 @@ func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } // Append appends one row. v is read synchronously and never retained, so a // caller streaming rows can reuse a single Metric variable across calls. func (a *MetricAppender) Append(v *Metric) { - a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) + a.b0.Append(arrow.Date32FromTime(v.Day)) a.b1.Append(v.Host) a.b2.Append(v.CPU) if v.Value == nil { diff --git a/arrgen/internal/gentypes/row_arrow.go b/arrgen/internal/gentypes/row_arrow.go index 0a0483f2f..58821a692 100644 --- a/arrgen/internal/gentypes/row_arrow.go +++ b/arrgen/internal/gentypes/row_arrow.go @@ -44,7 +44,6 @@ var rowArrowSchema = arrow.NewSchema([]arrow.Field{ {Name: "f64", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, {Name: "s", Type: arrow.BinaryTypes.String, Nullable: false}, {Name: "bin", Type: arrow.BinaryTypes.Binary, Nullable: false}, - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, {Name: "d32", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, {Name: "d64", Type: arrow.FixedWidthTypes.Date64, Nullable: false}, {Name: "t32", Type: &arrow.Time32Type{Unit: arrow.Millisecond}, Nullable: false}, @@ -53,7 +52,7 @@ var rowArrowSchema = arrow.NewSchema([]arrow.Field{ {Name: "dec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: false}, {Name: "dec64", Type: &arrow.Decimal64Type{Precision: 18, Scale: 4}, Nullable: false}, {Name: "dec128", Type: &arrow.Decimal128Type{Precision: 20, Scale: 3}, Nullable: false}, - {Name: "dec256", Type: &arrow.Decimal256Type{Precision: 76, Scale: 0}, Nullable: false}, + {Name: "dec256", Type: &arrow.Decimal256Type{Precision: 40, Scale: 5}, Nullable: false}, {Name: "ls", Type: arrow.BinaryTypes.LargeString, Nullable: false}, {Name: "vs", Type: arrow.BinaryTypes.StringView, Nullable: false}, {Name: "lb", Type: arrow.BinaryTypes.LargeBinary, Nullable: false}, @@ -68,12 +67,12 @@ var rowArrowSchema = arrow.NewSchema([]arrow.Field{ {Name: "pf64", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, {Name: "ps", Type: arrow.BinaryTypes.String, Nullable: true}, {Name: "pbin", Type: arrow.BinaryTypes.Binary, Nullable: true}, - {Name: "pts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: true}, {Name: "pd32", Type: arrow.FixedWidthTypes.Date32, Nullable: true}, + {Name: "pd64", Type: arrow.FixedWidthTypes.Date64, Nullable: true}, {Name: "pt64", Type: &arrow.Time64Type{Unit: arrow.Nanosecond}, Nullable: true}, {Name: "pdur", Type: &arrow.DurationType{Unit: arrow.Nanosecond}, Nullable: true}, {Name: "pdec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: true}, - {Name: "pdec128", Type: &arrow.Decimal128Type{Precision: 38, Scale: 0}, Nullable: true}, + {Name: "pdec128", Type: &arrow.Decimal128Type{Precision: 20, Scale: 3}, Nullable: true}, {Name: "pds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: true}, }, nil) @@ -103,37 +102,36 @@ type RowAppender struct { b12 *array.Float64Builder b13 *array.StringBuilder b14 *array.BinaryBuilder - b15 *array.TimestampBuilder - b16 *array.Date32Builder - b17 *array.Date64Builder - b18 *array.Time32Builder - b19 *array.Time64Builder - b20 *array.DurationBuilder - b21 *array.Decimal32Builder - b22 *array.Decimal64Builder - b23 *array.Decimal128Builder - b24 *array.Decimal256Builder - b25 *array.LargeStringBuilder - b26 *array.StringViewBuilder - b27 *array.BinaryBuilder - b28 *array.BinaryViewBuilder + b15 *array.Date32Builder + b16 *array.Date64Builder + b17 *array.Time32Builder + b18 *array.Time64Builder + b19 *array.DurationBuilder + b20 *array.Decimal32Builder + b21 *array.Decimal64Builder + b22 *array.Decimal128Builder + b23 *array.Decimal256Builder + b24 *array.LargeStringBuilder + b25 *array.StringViewBuilder + b26 *array.BinaryBuilder + b27 *array.BinaryViewBuilder + b28 *array.BinaryDictionaryBuilder b29 *array.BinaryDictionaryBuilder - b30 *array.BinaryDictionaryBuilder - b31 *array.Int32DictionaryBuilder - b32 *array.Float64DictionaryBuilder - b33 *array.Int64Builder - b34 *array.BooleanBuilder - b35 *array.Int64Builder - b36 *array.Float64Builder - b37 *array.StringBuilder - b38 *array.BinaryBuilder - b39 *array.TimestampBuilder - b40 *array.Date32Builder - b41 *array.Time64Builder - b42 *array.DurationBuilder - b43 *array.Decimal32Builder - b44 *array.Decimal128Builder - b45 *array.BinaryDictionaryBuilder + b30 *array.Int32DictionaryBuilder + b31 *array.Float64DictionaryBuilder + b32 *array.Int64Builder + b33 *array.BooleanBuilder + b34 *array.Int64Builder + b35 *array.Float64Builder + b36 *array.StringBuilder + b37 *array.BinaryBuilder + b38 *array.Date32Builder + b39 *array.Date64Builder + b40 *array.Time64Builder + b41 *array.DurationBuilder + b42 *array.Decimal32Builder + b43 *array.Decimal128Builder + b44 *array.BinaryDictionaryBuilder err error } @@ -164,37 +162,36 @@ func NewRowAppender(mem memory.Allocator) *RowAppender { b12: rb.Field(12).(*array.Float64Builder), b13: rb.Field(13).(*array.StringBuilder), b14: rb.Field(14).(*array.BinaryBuilder), - b15: rb.Field(15).(*array.TimestampBuilder), - b16: rb.Field(16).(*array.Date32Builder), - b17: rb.Field(17).(*array.Date64Builder), - b18: rb.Field(18).(*array.Time32Builder), - b19: rb.Field(19).(*array.Time64Builder), - b20: rb.Field(20).(*array.DurationBuilder), - b21: rb.Field(21).(*array.Decimal32Builder), - b22: rb.Field(22).(*array.Decimal64Builder), - b23: rb.Field(23).(*array.Decimal128Builder), - b24: rb.Field(24).(*array.Decimal256Builder), - b25: rb.Field(25).(*array.LargeStringBuilder), - b26: rb.Field(26).(*array.StringViewBuilder), - b27: rb.Field(27).(*array.BinaryBuilder), - b28: rb.Field(28).(*array.BinaryViewBuilder), + b15: rb.Field(15).(*array.Date32Builder), + b16: rb.Field(16).(*array.Date64Builder), + b17: rb.Field(17).(*array.Time32Builder), + b18: rb.Field(18).(*array.Time64Builder), + b19: rb.Field(19).(*array.DurationBuilder), + b20: rb.Field(20).(*array.Decimal32Builder), + b21: rb.Field(21).(*array.Decimal64Builder), + b22: rb.Field(22).(*array.Decimal128Builder), + b23: rb.Field(23).(*array.Decimal256Builder), + b24: rb.Field(24).(*array.LargeStringBuilder), + b25: rb.Field(25).(*array.StringViewBuilder), + b26: rb.Field(26).(*array.BinaryBuilder), + b27: rb.Field(27).(*array.BinaryViewBuilder), + b28: rb.Field(28).(*array.BinaryDictionaryBuilder), b29: rb.Field(29).(*array.BinaryDictionaryBuilder), - b30: rb.Field(30).(*array.BinaryDictionaryBuilder), - b31: rb.Field(31).(*array.Int32DictionaryBuilder), - b32: rb.Field(32).(*array.Float64DictionaryBuilder), - b33: rb.Field(33).(*array.Int64Builder), - b34: rb.Field(34).(*array.BooleanBuilder), - b35: rb.Field(35).(*array.Int64Builder), - b36: rb.Field(36).(*array.Float64Builder), - b37: rb.Field(37).(*array.StringBuilder), - b38: rb.Field(38).(*array.BinaryBuilder), - b39: rb.Field(39).(*array.TimestampBuilder), - b40: rb.Field(40).(*array.Date32Builder), - b41: rb.Field(41).(*array.Time64Builder), - b42: rb.Field(42).(*array.DurationBuilder), - b43: rb.Field(43).(*array.Decimal32Builder), - b44: rb.Field(44).(*array.Decimal128Builder), - b45: rb.Field(45).(*array.BinaryDictionaryBuilder), + b30: rb.Field(30).(*array.Int32DictionaryBuilder), + b31: rb.Field(31).(*array.Float64DictionaryBuilder), + b32: rb.Field(32).(*array.Int64Builder), + b33: rb.Field(33).(*array.BooleanBuilder), + b34: rb.Field(34).(*array.Int64Builder), + b35: rb.Field(35).(*array.Float64Builder), + b36: rb.Field(36).(*array.StringBuilder), + b37: rb.Field(37).(*array.BinaryBuilder), + b38: rb.Field(38).(*array.Date32Builder), + b39: rb.Field(39).(*array.Date64Builder), + b40: rb.Field(40).(*array.Time64Builder), + b41: rb.Field(41).(*array.DurationBuilder), + b42: rb.Field(42).(*array.Decimal32Builder), + b43: rb.Field(43).(*array.Decimal128Builder), + b44: rb.Field(44).(*array.BinaryDictionaryBuilder), } } @@ -235,103 +232,102 @@ func (a *RowAppender) Append(v *Row) { } else { a.b14.Append(v.Bin) } - a.b15.Append(arrow.Timestamp(v.Timestamp.UnixNano())) - a.b16.Append(arrow.Date32FromTime(v.Date32)) - a.b17.Append(arrow.Date64FromTime(v.Date64)) + a.b15.Append(arrow.Date32FromTime(v.Date32)) + a.b16.Append(arrow.Date64FromTime(v.Date64)) { tod := v.Time32.UTC() - a.b18.Append(arrow.Time32(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds() / 1e6)) + a.b17.Append(arrow.Time32(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds() / 1e6)) } { tod := v.Time64.UTC() - a.b19.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) + a.b18.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) } - a.b20.Append(arrow.Duration(v.Duration.Nanoseconds())) - a.b21.Append(v.Dec32) - a.b22.Append(v.Dec64) - a.b23.Append(v.Dec128) - a.b24.Append(v.Dec256) - a.b25.Append(v.LargeStr) - a.b26.Append(v.ViewStr) + a.b19.Append(arrow.Duration(v.Duration.Nanoseconds())) + a.b20.Append(v.Dec32) + a.b21.Append(v.Dec64) + a.b22.Append(v.Dec128) + a.b23.Append(v.Dec256) + a.b24.Append(v.LargeStr) + a.b25.Append(v.ViewStr) if v.LargeBin == nil { - a.b27.AppendNull() + a.b26.AppendNull() } else { - a.b27.Append(v.LargeBin) + a.b26.Append(v.LargeBin) } if v.ViewBin == nil { - a.b28.AppendNull() + a.b27.AppendNull() } else { - a.b28.Append(v.ViewBin) + a.b27.Append(v.ViewBin) } - a.setErr(a.b29.AppendString(v.DictStr)) + a.setErr(a.b28.AppendString(v.DictStr)) if v.DictBin == nil { - a.b30.AppendNull() + a.b29.AppendNull() } else { - a.setErr(a.b30.Append(v.DictBin)) + a.setErr(a.b29.Append(v.DictBin)) } - a.setErr(a.b31.Append(v.DictInt)) - a.setErr(a.b32.Append(v.DictF64)) - a.b33.Append(v.Untagged) + a.setErr(a.b30.Append(v.DictInt)) + a.setErr(a.b31.Append(v.DictF64)) + a.b32.Append(v.Untagged) if v.PBool == nil { - a.b34.AppendNull() + a.b33.AppendNull() } else { - a.b34.Append(*v.PBool) + a.b33.Append(*v.PBool) } if v.PInt64 == nil { - a.b35.AppendNull() + a.b34.AppendNull() } else { - a.b35.Append(*v.PInt64) + a.b34.Append(*v.PInt64) } if v.PF64 == nil { - a.b36.AppendNull() + a.b35.AppendNull() } else { - a.b36.Append(*v.PF64) + a.b35.Append(*v.PF64) } if v.PStr == nil { - a.b37.AppendNull() + a.b36.AppendNull() } else { - a.b37.Append(*v.PStr) + a.b36.Append(*v.PStr) } if v.PBin == nil || *v.PBin == nil { - a.b38.AppendNull() + a.b37.AppendNull() } else { - a.b38.Append(*v.PBin) + a.b37.Append(*v.PBin) } - if v.PTS == nil { - a.b39.AppendNull() + if v.PDate32 == nil { + a.b38.AppendNull() } else { - a.b39.Append(arrow.Timestamp((*v.PTS).UnixNano())) + a.b38.Append(arrow.Date32FromTime(*v.PDate32)) } - if v.PDate32 == nil { - a.b40.AppendNull() + if v.PDate64 == nil { + a.b39.AppendNull() } else { - a.b40.Append(arrow.Date32FromTime(*v.PDate32)) + a.b39.Append(arrow.Date64FromTime(*v.PDate64)) } if v.PTime64 == nil { - a.b41.AppendNull() + a.b40.AppendNull() } else { tod := (*v.PTime64).UTC() - a.b41.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) + a.b40.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) } if v.PDur == nil { - a.b42.AppendNull() + a.b41.AppendNull() } else { - a.b42.Append(arrow.Duration((*v.PDur).Nanoseconds())) + a.b41.Append(arrow.Duration((*v.PDur).Nanoseconds())) } if v.PDec32 == nil { - a.b43.AppendNull() + a.b42.AppendNull() } else { - a.b43.Append(*v.PDec32) + a.b42.Append(*v.PDec32) } if v.PDec128 == nil { - a.b44.AppendNull() + a.b43.AppendNull() } else { - a.b44.Append(*v.PDec128) + a.b43.Append(*v.PDec128) } if v.PDictS == nil { - a.b45.AppendNull() + a.b44.AppendNull() } else { - a.setErr(a.b45.AppendString(*v.PDictS)) + a.setErr(a.b44.AppendString(*v.PDictS)) } } @@ -380,7 +376,7 @@ func RowRecordBatch(mem memory.Allocator, vs []Row) (arrow.RecordBatch, error) { // fixedArrowSchema is the Arrow schema of Fixed, resolved from its arrow // struct tags at generate time. var fixedArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "day", Type: arrow.FixedWidthTypes.Date32, Nullable: false}, {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, {Name: "val", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, {Name: "ok", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, @@ -398,7 +394,7 @@ func FixedSchema() *arrow.Schema { return fixedArrowSchema } // It is not safe for concurrent use. Release it when done. type FixedAppender struct { rb *array.RecordBuilder - b0 *array.TimestampBuilder + b0 *array.Date32Builder b1 *array.Int64Builder b2 *array.Float64Builder b3 *array.BooleanBuilder @@ -418,7 +414,7 @@ func NewFixedAppender(mem memory.Allocator) *FixedAppender { rb := array.NewRecordBuilder(mem, fixedArrowSchema) return &FixedAppender{ rb: rb, - b0: rb.Field(0).(*array.TimestampBuilder), + b0: rb.Field(0).(*array.Date32Builder), b1: rb.Field(1).(*array.Int64Builder), b2: rb.Field(2).(*array.Float64Builder), b3: rb.Field(3).(*array.BooleanBuilder), @@ -444,7 +440,7 @@ func (a *FixedAppender) Len() int { return a.rb.Field(0).Len() } // Append appends one row. v is read synchronously and never retained, so a // caller streaming rows can reuse a single Fixed variable across calls. func (a *FixedAppender) Append(v *Fixed) { - a.b0.Append(arrow.Timestamp(v.Timestamp.UnixNano())) + a.b0.Append(arrow.Date32FromTime(v.Day)) a.b1.Append(v.ID) a.b2.Append(v.Value) a.b3.Append(v.OK) diff --git a/arrgen/internal/gentypes/types.go b/arrgen/internal/gentypes/types.go index f66f34fbe..16a1ec063 100644 --- a/arrgen/internal/gentypes/types.go +++ b/arrgen/internal/gentypes/types.go @@ -48,17 +48,16 @@ type Row struct { Str string `arrow:"s"` Bin []byte `arrow:"bin"` - Timestamp time.Time `arrow:"ts"` - Date32 time.Time `arrow:"d32,date32"` - Date64 time.Time `arrow:"d64,date64"` - Time32 time.Time `arrow:"t32,time32"` - Time64 time.Time `arrow:"t64,time64"` - Duration time.Duration `arrow:"dur"` + Date32 time.Time `arrow:"d32,date32"` + Date64 time.Time `arrow:"d64,date64"` + Time32 time.Time `arrow:"t32,time32"` + Time64 time.Time `arrow:"t64,time64"` + Duration time.Duration `arrow:"dur"` Dec32 decimal.Decimal32 `arrow:"dec32"` Dec64 decimal.Decimal64 `arrow:"dec64,decimal(18,4)"` Dec128 decimal128.Num `arrow:"dec128,decimal(20,3)"` - Dec256 decimal256.Num `arrow:"dec256"` + Dec256 decimal256.Num `arrow:"dec256,decimal(40,5)"` LargeStr string `arrow:"ls,large"` ViewStr string `arrow:"vs,view"` @@ -76,12 +75,12 @@ type Row struct { PF64 *float64 `arrow:"pf64"` PStr *string `arrow:"ps"` PBin *[]byte `arrow:"pbin"` - PTS *time.Time `arrow:"pts"` PDate32 *time.Time `arrow:"pd32,date32"` + PDate64 *time.Time `arrow:"pd64,date64"` PTime64 *time.Time `arrow:"pt64,time64"` PDur *time.Duration `arrow:"pdur"` PDec32 *decimal.Decimal32 `arrow:"pdec32"` - PDec128 *decimal128.Num `arrow:"pdec128"` + PDec128 *decimal128.Num `arrow:"pdec128,decimal(20,3)"` PDictS *string `arrow:"pds,dict"` Secret string `arrow:"-"` // excluded from Arrow entirely @@ -93,11 +92,11 @@ type Row struct { // allocating at all, which is not something a string or []byte column can // promise once its data buffer needs to double. type Fixed struct { - Timestamp time.Time `arrow:"ts"` - ID int64 `arrow:"id"` - Value float64 `arrow:"val"` - OK bool `arrow:"ok"` - Optional *float64 `arrow:"opt"` + Day time.Time `arrow:"day,date32"` + ID int64 `arrow:"id"` + Value float64 `arrow:"val"` + OK bool `arrow:"ok"` + Optional *float64 `arrow:"opt"` } //go:generate go run github.com/apache/arrow-go/arrgen/cmd/arrgen -type Row,Fixed -header ../../license_header.txt -output row_arrow.go diff --git a/arrgen/mapping.go b/arrgen/mapping.go index fc8c9a475..6a4e31f36 100644 --- a/arrgen/mapping.go +++ b/arrgen/mapping.go @@ -127,13 +127,13 @@ func namedSpec(named *types.Named, opts tagOpts) (colSpec, error) { }, }, nil case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal32": - return decimalSpec(opts, "Decimal32", 9) + return decimalSpec(opts, "Decimal32", 9, "") case isDecimalPkg(pkg, "decimal") && obj.Name() == "Decimal64": - return decimalSpec(opts, "Decimal64", 18) + return decimalSpec(opts, "Decimal64", 18, "") case isDecimalPkg(pkg, "decimal128") && obj.Name() == "Num": - return decimalSpec(opts, "Decimal128", 38) + return decimalSpec(opts, "Decimal128", 38, "decimal128.Num") case isDecimalPkg(pkg, "decimal256") && obj.Name() == "Num": - return decimalSpec(opts, "Decimal256", 76) + return decimalSpec(opts, "Decimal256", 76, "decimal256.Num") } if _, ok := named.Underlying().(*types.Struct); ok { return colSpec{}, fmt.Errorf("nested struct column for %s is not supported; encode it with arreflect", named) @@ -146,11 +146,19 @@ func isDecimalPkg(path, name string) bool { return path == prefix+name } -func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32) (colSpec, error) { +// decimalSpec maps a decimal field. structGoType is empty for decimal.Decimal32 +// and decimal.Decimal64, whose Go types are defined integers, and the Go type's +// name for decimal128.Num and decimal256.Num, whose Go types are structs and so +// need an explicit decimal tag to stay in step with arreflect. See +// errUninferableStructScalar. +func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32, structGoType string) (colSpec, error) { goType := "decimal." + arrowName if err := checkOpts(opts, goType, optSupport{decimal: true}); err != nil { return colSpec{}, err } + if structGoType != "" && !opts.HasDecimalOpts { + return colSpec{}, errUninferableStructScalar(structGoType, "a decimal(precision,scale) tag") + } precision, scale := defaultPrecision, int32(0) if opts.HasDecimalOpts { precision, scale = opts.DecimalPrecision, opts.DecimalScale @@ -162,10 +170,9 @@ func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32) (colSpe }, nil } -// timeSpec maps a time.Time field. Note that an untagged time.Time is a -// TIMESTAMP(ns, UTC) column here, which is what arreflect.InferType[time.Time] -// and its top-level FromSlice path both produce - but not what its struct-field -// path produces today. See the divergence note in the package README. +// timeSpec maps a time.Time field, which needs one of the four temporal tags: +// the timestamp spelling arreflect cannot infer for a struct field is rejected +// rather than generated. See errUninferableStructScalar. func timeSpec(opts tagOpts) (colSpec, error) { if err := checkOpts(opts, "time.Time", optSupport{temporal: true}); err != nil { return colSpec{}, err @@ -192,16 +199,32 @@ func timeSpec(opts tagOpts) (colSpec, error) { case "time64": return timeOfDaySpec("&arrow.Time64Type{Unit: arrow.Nanosecond}", "*array.Time64Builder", "arrow.Time64", ""), nil default: - return colSpec{ - arrowType: `&arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}`, - builderType: "*array.TimestampBuilder", - appendStmt: func(bld, recv, _ string) string { - return fmt.Sprintf("%s.Append(arrow.Timestamp(%s.UnixNano()))", bld, recv) - }, - }, nil + // "" and "timestamp" are the two spellings that ask for a TIMESTAMP + // column, and neither is inferable as a struct field. + return colSpec{}, errUninferableStructScalar("time.Time", "one of the date32, date64, time32 or time64 tags") } } +// errUninferableStructScalar reports a field whose Go type is a struct that +// Arrow models as a scalar, tagged in a way arreflect cannot infer. +// +// arreflect's inferArrowType switches on reflect.Kind before it reaches the +// types it matches by identity, so a struct field of one of these types is +// resolved by inferStructType instead, which finds only unexported fields and +// yields an empty struct<> - and the value is then dropped. Only a tag that +// names the Arrow type outright survives, because arreflect applies the tag +// after the inferred type. +// +// arrgen could emit the column Arrow plainly means here, but generated code and +// arreflect would then disagree about the schema, which is the one thing this +// generator exists not to do. So the untaggable spellings are a generate-time +// error naming the field and the tag that fixes it. +func errUninferableStructScalar(goType, fix string) error { + return fmt.Errorf("%s needs %s: arreflect infers an empty struct<> for it as a struct field "+ + "and drops the value, so a column generated here would not match arreflect.InferSchema. "+ + "Tag it, or encode the struct with arreflect instead", goType, fix) +} + // timeOfDaySpec renders arreflect's timeOfDayNanos inline. It is emitted as a // block rather than a package-level helper so that two generated files in the // same package cannot collide over the helper's name. diff --git a/arrgen/testdata/basic.golden b/arrgen/testdata/basic.golden index 05fa66891..3890fa9e9 100644 --- a/arrgen/testdata/basic.golden +++ b/arrgen/testdata/basic.golden @@ -27,7 +27,7 @@ import ( // metricArrowSchema is the Arrow schema of Metric, resolved from its arrow // struct tags at generate time. var metricArrowSchema = arrow.NewSchema([]arrow.Field{ - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, Nullable: false}, + {Name: "day", Type: arrow.FixedWidthTypes.Date64, Nullable: false}, {Name: "host", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: false}, {Name: "count", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, {Name: "Blob", Type: arrow.BinaryTypes.Binary, Nullable: false}, @@ -44,7 +44,7 @@ func MetricSchema() *arrow.Schema { return metricArrowSchema } // It is not safe for concurrent use. Release it when done. type MetricAppender struct { rb *array.RecordBuilder - b0 *array.TimestampBuilder + b0 *array.Date64Builder b1 *array.BinaryDictionaryBuilder b2 *array.Int64Builder b3 *array.BinaryBuilder @@ -63,7 +63,7 @@ func NewMetricAppender(mem memory.Allocator) *MetricAppender { rb := array.NewRecordBuilder(mem, metricArrowSchema) return &MetricAppender{ rb: rb, - b0: rb.Field(0).(*array.TimestampBuilder), + b0: rb.Field(0).(*array.Date64Builder), b1: rb.Field(1).(*array.BinaryDictionaryBuilder), b2: rb.Field(2).(*array.Int64Builder), b3: rb.Field(3).(*array.BinaryBuilder), @@ -88,7 +88,7 @@ func (a *MetricAppender) Len() int { return a.rb.Field(0).Len() } // Append appends one row. v is read synchronously and never retained, so a // caller streaming rows can reuse a single Metric variable across calls. func (a *MetricAppender) Append(v *Metric) { - a.b0.Append(arrow.Timestamp(v.Time.UnixNano())) + a.b0.Append(arrow.Date64FromTime(v.Day)) a.setErr(a.b1.AppendString(v.Host)) if v.Count == nil { a.b2.AppendNull() diff --git a/arrgen/testdata/basic/basic.go b/arrgen/testdata/basic/basic.go index e46e0b9cf..39dd901b6 100644 --- a/arrgen/testdata/basic/basic.go +++ b/arrgen/testdata/basic/basic.go @@ -23,7 +23,7 @@ import "time" // Metric is an exported type, so its generated API is exported too. type Metric struct { - Time time.Time `arrow:"ts"` + Day time.Time `arrow:"day,date64"` Host string `arrow:"host,dict"` Count *int64 `arrow:"count"` Blob []byte diff --git a/arrgen/testdata/errors/errors.go b/arrgen/testdata/errors/errors.go index cc4e54f8d..5d414818b 100644 --- a/arrgen/testdata/errors/errors.go +++ b/arrgen/testdata/errors/errors.go @@ -18,7 +18,12 @@ // error messages stay covered and stay readable. package errors -import "time" +import ( + "time" + + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/decimal256" +) // Inner is a nested struct used by Nested and Embedded. type Inner struct { @@ -119,6 +124,25 @@ type TimeSlice struct { // NotAStruct is a named type whose underlying type is not a struct. type NotAStruct int64 +// BareTime, TimestampTime and the two bare decimals below are the spellings +// arreflect cannot infer as a struct field, so arrgen rejects them rather than +// generating a column the reflection path would not agree with. +type BareTime struct { + Field time.Time +} + +type TimestampTime struct { + Field time.Time `arrow:"f,timestamp"` +} + +type BareDecimal128 struct { + Field decimal128.Num +} + +type BareDecimal256 struct { + Field decimal256.Num +} + type Collide struct { Field int64 } diff --git a/arrgen/testdata/pregen/pregen.go b/arrgen/testdata/pregen/pregen.go index 009eecf33..d9fadce97 100644 --- a/arrgen/testdata/pregen/pregen.go +++ b/arrgen/testdata/pregen/pregen.go @@ -27,7 +27,7 @@ import ( ) type Sample struct { - At time.Time `arrow:"at"` + At time.Time `arrow:"at,date64"` Value float64 `arrow:"value"` } diff --git a/ci/scripts/build.sh b/ci/scripts/build.sh index 404d3e388..9ec9f540a 100755 --- a/ci/scripts/build.sh +++ b/ci/scripts/build.sh @@ -79,3 +79,10 @@ pushd "${source_dir}/parquet" go install -tags pqarrow_read_only ./... popd + +# arrgen is a nested module, so the traversals above never reach it. +pushd "${source_dir}/arrgen" + +go build -v ./... + +popd diff --git a/ci/scripts/test.sh b/ci/scripts/test.sh index 45064f0ae..35b3ed496 100755 --- a/ci/scripts/test.sh +++ b/ci/scripts/test.sh @@ -81,3 +81,34 @@ go test "${test_args[@]}" -tags assert ./... go test "${test_args[@]}" -tags assert,noasm ./... popd + +# arrgen is a nested module: "./..." above stops at its go.mod, so it needs its +# own vet, test and generation check. Its allocation assertions step aside under +# -race and -asan on their own, via build tags. +pushd "${source_dir}/arrgen" + +go vet ./... + +go test "${test_args[@]}" ./... + +# TestCheckedInFilesAreUpToDate covers the same ground through the arrgen +# package, but only running the committed go:generate directives covers the +# command wrapper and the flags they pass it. +go generate ./... +git diff --exit-code -- . + +popd + +# The run above resolves the released arrow-go that arrgen's go.mod names, which +# is what a consumer of the generator gets. Run it against this tree as well: the +# equivalence tests are what catch a change to arrow/array/arreflect that the +# generated encoders no longer match, and in module mode they would not see one +# until after the next release. +pushd "${source_dir}" + +rm -f go.work go.work.sum +go work init . ./arrgen +go test "${test_args[@]}" ./arrgen/... +rm -f go.work go.work.sum + +popd From eb8b976b84834361378709211ad2b70fb24a3152 Mon Sep 17 00:00:00 2001 From: jay <38236622+JayJamieson@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:21:07 +0000 Subject: [PATCH 3/4] feat(arrgen): support multi-level pointer fields arreflect's appendValue walks every pointer level to the value and nulls the column if any level along the way is nil, so **int64 and deeper were supported there while arrgen rejected them. Strip every level instead and emit one nil check per level, plus the existing check on a nil []byte value behind them. resolveColumn now reports pointer depth rather than a nullable flag; nullability is still read from the outermost pointer alone, which is arreflect's rule. The Row fixture grows **int64, **string, **[]byte and ***float64 columns, with rows that nil an inner level while the outer one stays set, so each generated nil check is exercised separately. --- arrgen/README.md | 7 ++-- arrgen/doc.go | 4 +- arrgen/generate.go | 43 ++++++++++++-------- arrgen/generate_test.go | 1 - arrgen/internal/gentypes/bench_test.go | 2 +- arrgen/internal/gentypes/equivalence_test.go | 17 +++++++- arrgen/internal/gentypes/row_arrow.go | 32 +++++++++++++++ arrgen/internal/gentypes/types.go | 8 ++++ arrgen/mapping.go | 23 +++++++---- arrgen/testdata/errors/errors.go | 4 -- 10 files changed, 104 insertions(+), 37 deletions(-) diff --git a/arrgen/README.md b/arrgen/README.md index f015ab56e..c1241e91e 100644 --- a/arrgen/README.md +++ b/arrgen/README.md @@ -133,8 +133,9 @@ Runnable versions of all three are the testable examples in `bool`, `int8/16/32/64`, `int`, `uint8/16/32/64`, `uint`, `float32/64`, `string`, `[]byte`, `time.Time`, `time.Duration`, `decimal.Decimal32`, -`decimal.Decimal64`, `decimal128.Num`, `decimal256.Num`, and a pointer to any of -those for a nullable column. +`decimal.Decimal64`, `decimal128.Num`, `decimal256.Num`, and one or more +pointers to any of those for a nullable column - a nil at any level is a null, +as it is in `arreflect`. Tag options are `arreflect`'s: a leading column name, `-` to skip the field, the temporal overrides `date32`, `date64`, `time32` and `time64`, `dict`, `view`, @@ -202,7 +203,7 @@ Batch benchmarks encode 1024 rows per operation. | --- | --- | --- | --- | | `MetricBatch` (4 columns) | 69.6 µs, 97 allocs | 30.5 µs, 52 allocs | **2.3x faster** | | `FixedBatch` (5 fixed-width columns) | 72.9 µs, 96 allocs | 26.7 µs, 46 allocs | **2.7x faster** | -| `RowBatch` (45 columns) | 985 µs, 3418 allocs | 515 µs, 3118 allocs | **1.9x faster** | +| `RowBatch` (49 columns) | 870 µs, 3502 allocs | 483 µs, 3174 allocs | **1.8x faster** | | `StreamAppend` (per row) | 73.2 ns | 27.5 ns, **0 allocs** | **2.7x faster** | | `Schema` | 6.4 µs, 131 allocs | 1.8 ns, **0 allocs** | | diff --git a/arrgen/doc.go b/arrgen/doc.go index df77b2ecc..a4bca0318 100644 --- a/arrgen/doc.go +++ b/arrgen/doc.go @@ -47,8 +47,8 @@ // // bool, int8/16/32/64, int, uint8/16/32/64, uint, float32/64, string, []byte, // time.Time, time.Duration, decimal.Decimal32, decimal.Decimal64, -// decimal128.Num and decimal256.Num, plus a pointer to any of those for a -// nullable column. +// decimal128.Num and decimal256.Num, plus one or more pointers to any of those +// for a nullable column - a nil at any level is a null, as it is in arreflect. // // Tag options mirror arreflect: a leading name, "-" to skip a field, the // temporal overrides date32, date64, time32 and time64, dict, view, large, and diff --git a/arrgen/generate.go b/arrgen/generate.go index 530bf1150..91b007999 100644 --- a/arrgen/generate.go +++ b/arrgen/generate.go @@ -164,7 +164,7 @@ func genForType(pkg *packages.Package, name string) (genType, error) { Index: i, Name: c.name, GoField: c.goField, - Nullable: c.nullable, + Nullable: c.nullable(), ArrowType: c.spec.arrowType, BuilderType: c.spec.builderType, BuilderVar: fmt.Sprintf("b%d", i), @@ -209,10 +209,14 @@ func findStruct(pkg *packages.Package, name string) (*types.Struct, error) { type column struct { name string goField string - nullable bool + ptrDepth int // pointer levels between the field and its value spec colSpec } +// nullable reports whether the column admits nulls, which for arreflect is a +// question about the field's outermost pointer and nothing else. +func (c column) nullable() bool { return c.ptrDepth > 0 } + // collectColumns walks the struct's fields in declaration order, which is the // order arreflect settles on for a flat struct. Embedded fields are rejected // rather than promoted: arreflect resolves promoted names by breadth and tag, @@ -255,23 +259,25 @@ func collectColumns(st *types.Struct) ([]column, error) { } seen[name] = f.Name() - spec, nullable, err := resolveColumn(f.Type(), opts) + spec, ptrDepth, err := resolveColumn(f.Type(), opts) if err != nil { return nil, fmt.Errorf("field %s: %w", f.Name(), err) } - cols = append(cols, column{name: name, goField: f.Name(), nullable: nullable, spec: spec}) + cols = append(cols, column{name: name, goField: f.Name(), ptrDepth: ptrDepth, spec: spec}) } return cols, nil } // renderAppend emits the append statement for one column, wrapping it in the -// null guards the reflection path applies: a nil pointer is a null, and so is a -// nil []byte even when the column itself is not nullable. +// null guards the reflection path applies: a nil pointer at any level is a +// null, and so is a nil []byte even when the column itself is not nullable. func renderAppend(idx int, c column) string { bld := fmt.Sprintf("a.b%d", idx) - val, recv := "v."+c.goField, "v."+c.goField - if c.nullable { - val, recv = "*v."+c.goField, "(*v."+c.goField+")" + val := strings.Repeat("*", c.ptrDepth) + "v." + c.goField + recv := val + if c.ptrDepth > 0 { + // A method on the value has to bind to the value, not to a pointer. + recv = "(" + val + ")" } stmt := c.spec.appendStmt(bld, recv, val) @@ -279,15 +285,20 @@ func renderAppend(idx int, c column) string { stmt = "a.setErr(" + stmt + ")" } - var guard string - switch { - case c.nullable && c.spec.nilable: - guard = fmt.Sprintf("v.%s == nil || *v.%s == nil", c.goField, c.goField) - case c.nullable, c.spec.nilable: - guard = fmt.Sprintf("v.%s == nil", c.goField) - default: + // One nil check per pointer level, plus one on the value itself when a nil + // value is a null in its own right. + levels := c.ptrDepth + if c.spec.nilable { + levels++ + } + if levels == 0 { return stmt } + checks := make([]string, levels) + for i := range checks { + checks[i] = strings.Repeat("*", i) + "v." + c.goField + " == nil" + } + guard := strings.Join(checks, " || ") // A statement that is already a block (a time-of-day column scopes a local) // sheds its braces on the way into the guard's else, which would otherwise diff --git a/arrgen/generate_test.go b/arrgen/generate_test.go index 26a229de5..b4aec8767 100644 --- a/arrgen/generate_test.go +++ b/arrgen/generate_test.go @@ -131,7 +131,6 @@ func TestGenerateRejects(t *testing.T) { {"BareDecimal128", "decimal128.Num needs a decimal(precision,scale) tag"}, {"BareDecimal256", "decimal256.Num needs a decimal(precision,scale) tag"}, {"NoColumns", "no Arrow columns"}, - {"DoublePointer", "multi-level pointer"}, {"TimeSlice", "list column"}, {"NotAStruct", "not a struct type"}, {"Absent", "type not found"}, diff --git a/arrgen/internal/gentypes/bench_test.go b/arrgen/internal/gentypes/bench_test.go index 81b385e3b..f850c3bdb 100644 --- a/arrgen/internal/gentypes/bench_test.go +++ b/arrgen/internal/gentypes/bench_test.go @@ -94,7 +94,7 @@ func BenchmarkMetricBatch(b *testing.B) { } // BenchmarkRowBatch runs the same comparison over the wide fixture, where a -// row costs 45 columns of reflection instead of four. +// row costs 49 columns of reflection instead of four. func BenchmarkRowBatch(b *testing.B) { rows := makeRows(benchRows) mem := memory.DefaultAllocator diff --git a/arrgen/internal/gentypes/equivalence_test.go b/arrgen/internal/gentypes/equivalence_test.go index 5de1eb21c..5b09939b9 100644 --- a/arrgen/internal/gentypes/equivalence_test.go +++ b/arrgen/internal/gentypes/equivalence_test.go @@ -156,6 +156,21 @@ func makeRows(n int) []gentypes.Row { dec128 := decimal128.FromU64(uint64(i)) r.PDate32, r.PDate64, r.PTime64, r.PDur, r.PDec32, r.PDictS = &ts, &ts, &ts, &dur, &dec, &ds r.PDec128 = &dec128 + + // Multi-level pointers. Every fourth row nils the inner level while the + // outer one stays set, so the per-level guards are covered separately + // rather than only ever failing at the first. + pi64, pf64, pstr := &i64, &f64, &s + ppf64 := &pf64 + r.PPInt64, r.PPStr, r.PPPF64 = &pi64, &pstr, &ppf64 + ppbin := &bin + if i%4 == 1 { + ppbin = nil // a nil inner pointer behind a set outer one + } else if i%4 == 2 { + var nilBin []byte + ppbin = &nilBin // a set pointer to a nil []byte + } + r.PPBin = &ppbin rows[i] = r } return rows @@ -190,7 +205,7 @@ func TestFixedSchemaMatchesArreflect(t *testing.T) { // compareSchemas asserts got and want are the same schema, field for field. // Schema.Equal would answer the same question in one call, but a mismatch on a -// 45-column fixture is only actionable if the failure names the column. +// 49-column fixture is only actionable if the failure names the column. func compareSchemas(t *testing.T, got, want *arrow.Schema) { t.Helper() if got.NumFields() != want.NumFields() { diff --git a/arrgen/internal/gentypes/row_arrow.go b/arrgen/internal/gentypes/row_arrow.go index 58821a692..9c8f1d2e3 100644 --- a/arrgen/internal/gentypes/row_arrow.go +++ b/arrgen/internal/gentypes/row_arrow.go @@ -74,6 +74,10 @@ var rowArrowSchema = arrow.NewSchema([]arrow.Field{ {Name: "pdec32", Type: &arrow.Decimal32Type{Precision: 9, Scale: 0}, Nullable: true}, {Name: "pdec128", Type: &arrow.Decimal128Type{Precision: 20, Scale: 3}, Nullable: true}, {Name: "pds", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, Nullable: true}, + {Name: "ppi64", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "pps", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "ppbin", Type: arrow.BinaryTypes.Binary, Nullable: true}, + {Name: "pppf64", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, }, nil) // RowSchema returns the Arrow schema encoded by RowAppender. It is @@ -132,6 +136,10 @@ type RowAppender struct { b42 *array.Decimal32Builder b43 *array.Decimal128Builder b44 *array.BinaryDictionaryBuilder + b45 *array.Int64Builder + b46 *array.StringBuilder + b47 *array.BinaryBuilder + b48 *array.Float64Builder err error } @@ -192,6 +200,10 @@ func NewRowAppender(mem memory.Allocator) *RowAppender { b42: rb.Field(42).(*array.Decimal32Builder), b43: rb.Field(43).(*array.Decimal128Builder), b44: rb.Field(44).(*array.BinaryDictionaryBuilder), + b45: rb.Field(45).(*array.Int64Builder), + b46: rb.Field(46).(*array.StringBuilder), + b47: rb.Field(47).(*array.BinaryBuilder), + b48: rb.Field(48).(*array.Float64Builder), } } @@ -329,6 +341,26 @@ func (a *RowAppender) Append(v *Row) { } else { a.setErr(a.b44.AppendString(*v.PDictS)) } + if v.PPInt64 == nil || *v.PPInt64 == nil { + a.b45.AppendNull() + } else { + a.b45.Append(**v.PPInt64) + } + if v.PPStr == nil || *v.PPStr == nil { + a.b46.AppendNull() + } else { + a.b46.Append(**v.PPStr) + } + if v.PPBin == nil || *v.PPBin == nil || **v.PPBin == nil { + a.b47.AppendNull() + } else { + a.b47.Append(**v.PPBin) + } + if v.PPPF64 == nil || *v.PPPF64 == nil || **v.PPPF64 == nil { + a.b48.AppendNull() + } else { + a.b48.Append(***v.PPPF64) + } } // AppendSlice appends every element of vs, reserving room for them up front. diff --git a/arrgen/internal/gentypes/types.go b/arrgen/internal/gentypes/types.go index 16a1ec063..0f8121468 100644 --- a/arrgen/internal/gentypes/types.go +++ b/arrgen/internal/gentypes/types.go @@ -83,6 +83,14 @@ type Row struct { PDec128 *decimal128.Num `arrow:"pdec128,decimal(20,3)"` PDictS *string `arrow:"pds,dict"` + // Multi-level pointers: arreflect walks to the value and nulls the column + // if any level is nil, so these cover the guard renderAppend emits per + // level - and, for ppbin, a nil []byte behind two of them. + PPInt64 **int64 `arrow:"ppi64"` + PPStr **string `arrow:"pps"` + PPBin **[]byte `arrow:"ppbin"` + PPPF64 ***float64 `arrow:"pppf64"` + Secret string `arrow:"-"` // excluded from Arrow entirely } diff --git a/arrgen/mapping.go b/arrgen/mapping.go index 6a4e31f36..8ac0018ce 100644 --- a/arrgen/mapping.go +++ b/arrgen/mapping.go @@ -72,19 +72,24 @@ func checkOpts(o tagOpts, goType string, s optSupport) error { } // resolveColumn maps a Go field type and its parsed tag to an Arrow column, -// reporting nullable=true for a pointer field (matching arreflect, which marks -// a column nullable from the outermost pointer). -func resolveColumn(t types.Type, opts tagOpts) (spec colSpec, nullable bool, err error) { +// also reporting how many pointers the field is behind. +// +// Every level is stripped, matching arreflect, whose appendValue walks pointers +// to the value and nulls the column if any level along the way is nil. The +// column is nullable when there is at least one, which is arreflect's rule too: +// it reads nullability from the outermost pointer alone. +func resolveColumn(t types.Type, opts tagOpts) (spec colSpec, ptrDepth int, err error) { t = types.Unalias(t) - if p, ok := t.(*types.Pointer); ok { - nullable = true - t = types.Unalias(p.Elem()) - if _, ok := t.(*types.Pointer); ok { - return colSpec{}, false, fmt.Errorf("multi-level pointer type %s is not supported", t) + for { + p, ok := t.(*types.Pointer) + if !ok { + break } + ptrDepth++ + t = types.Unalias(p.Elem()) } spec, err = baseSpec(t, opts) - return spec, nullable, err + return spec, ptrDepth, err } func baseSpec(t types.Type, opts tagOpts) (colSpec, error) { diff --git a/arrgen/testdata/errors/errors.go b/arrgen/testdata/errors/errors.go index 5d414818b..dd452f149 100644 --- a/arrgen/testdata/errors/errors.go +++ b/arrgen/testdata/errors/errors.go @@ -113,10 +113,6 @@ type NoColumns struct { hidden int64 } -type DoublePointer struct { - Field **int64 -} - type TimeSlice struct { Field []time.Time } From b2d5e73fe111321e9f20a23f9bcc09d3301b56ca Mon Sep 17 00:00:00 2001 From: jay <38236622+JayJamieson@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:30:12 +0000 Subject: [PATCH 4/4] docs(arrgen): correct and tighten the prose Two claims were wrong. The comment on timeSpec said arreflect's own tests only pass these types as the top-level element of a slice; they in fact assert the mapping against inferPrimitiveArrowType, which a struct field never reaches. The README said buildArray matches these types before it switches on reflect.Kind, which holds for the decimals but not for time.Time, matched inside the struct branch. Both now say what the code does, with the verified FromSlice output. --- arrgen/README.md | 67 ++++++++++---------- arrgen/doc.go | 24 +++---- arrgen/generate.go | 20 +++--- arrgen/generate_test.go | 20 +++--- arrgen/internal/gentypes/equivalence_test.go | 8 +-- arrgen/internal/gentypes/types.go | 6 +- arrgen/mapping.go | 47 ++++++-------- ci/scripts/test.sh | 9 ++- 8 files changed, 98 insertions(+), 103 deletions(-) diff --git a/arrgen/README.md b/arrgen/README.md index c1241e91e..b0ece14d4 100644 --- a/arrgen/README.md +++ b/arrgen/README.md @@ -29,7 +29,7 @@ fields straight into typed Arrow builders. It lives in its own Go module, so it is opt-in twice over: nothing in `github.com/apache/arrow-go/v18` imports it, and the code it emits imports only -`arrow`, `arrow/array` and `arrow/memory` - never `arrgen` itself. Generating is +`arrow`, `arrow/array` and `arrow/memory`, never `arrgen` itself. Generating is a build-time step; your binary does not grow a dependency for it. ## Quick start @@ -77,9 +77,8 @@ module files: ``` Both forms need `arrgen` to be resolvable. As a nested module it is versioned -under its own `arrgen/vX.Y.Z` tags, and until the first of those is published -neither form resolves from a released version - use a `replace` directive -pointing at a checkout in the meantime: +under its own `arrgen/vX.Y.Z` tags. Until the first of those is published, +neither form resolves from a released version, so point at a checkout instead: ```sh go mod edit -replace github.com/apache/arrow-go/arrgen=../arrow-go/arrgen @@ -93,7 +92,7 @@ this module use that spelling because this module is `arrgen`. ### What you get ```go -// One batch from a slice - the drop-in for arreflect.RecordFromSlice. +// One batch from a slice. The drop-in for arreflect.RecordFromSlice. rec, err := telemetry.MetricRecordBatch(mem, metrics) // Or stream rows in and cut batches where you want them. @@ -134,15 +133,16 @@ Runnable versions of all three are the testable examples in `bool`, `int8/16/32/64`, `int`, `uint8/16/32/64`, `uint`, `float32/64`, `string`, `[]byte`, `time.Time`, `time.Duration`, `decimal.Decimal32`, `decimal.Decimal64`, `decimal128.Num`, `decimal256.Num`, and one or more -pointers to any of those for a nullable column - a nil at any level is a null, +pointers to any of those for a nullable column. A nil at any level is a null, as it is in `arreflect`. Tag options are `arreflect`'s: a leading column name, `-` to skip the field, the temporal overrides `date32`, `date64`, `time32` and `time64`, `dict`, `view`, `large`, and `decimal(precision,scale)`. -Two of those types need an explicit tag, because the untagged spelling is one -`arreflect` cannot infer as a struct field - see [below](#the-two-types-that-need-a-tag): +Two of those types need an explicit tag, because `arreflect` cannot infer the +untagged spelling as a struct field. See +[below](#the-two-types-that-need-a-tag): | Field | Required tag | | --- | --- | @@ -162,7 +162,8 @@ The point of matching `arreflect`'s tag dialect is that switching between the two paths is a call-site change and nothing else. `internal/gentypes` holds a fixture with one field for every supported column shape and asserts, on every run, that the generated encoder and `arreflect` produce the same schema and the -same column data for the same input - every column, with nothing excused. +same column data for the same input. Every column is compared, with no +exceptions. ### The two types that need a tag @@ -176,23 +177,25 @@ an empty `struct<>`, and the value is then dropped. Today that affects: | `time.Time`, untagged or `,timestamp` | `struct<>` | | `decimal128.Num` / `decimal256.Num`, untagged | `struct<>` | -A tag naming the Arrow type outright - `,date32`, `,decimal(20,3)` - rescues the -field, because `arreflect` applies the tag after the inferred type. -`arreflect`'s own tests only pass these types as the top-level element of a -slice, where `buildArray` special-cases them before the `reflect.Struct` branch -and they encode correctly. +A tag naming the Arrow type survives, such as `,date32` or `,decimal(20,3)`, +because `arreflect` applies tags after inference. `arreflect.FromSlice` also +handles them correctly at the top level, returning `timestamp[ns, tz=UTC]` and +`decimal(38, 0)`, because `buildArray` matches them by type rather than reaching +`inferStructType`. -`arrgen` could emit the column Arrow plainly means here, and an earlier revision -did. But then generated code and `arreflect` would disagree about the schema, -which is the one thing this generator exists not to do, so the untaggable -spellings are a generate-time error naming the field and the tag that fixes it -instead. The practical consequence is that **`arrgen` cannot emit a `TIMESTAMP` -column**: `,timestamp` is rejected along with the untagged spelling, since -`arreflect` infers `struct<>` for both. +`arreflect`'s own tests assert the intended mapping against +`inferPrimitiveArrowType`, which a struct field never reaches, so they pass +either way. -`TestArreflectCannotInferStructScalars` pins that upstream behavior. If -`arrow/array/arreflect` is fixed, that test fails and tells us to drop these -rejections and generate the columns instead. +`arrgen` could emit the column Arrow means here, and an earlier revision did. +Generated code and `arreflect` would then disagree about the schema, so instead +these spellings are a generate-time error naming the field and the tag that +fixes it. One consequence: **`arrgen` cannot emit a `TIMESTAMP` column**, since +`,timestamp` is rejected along with the untagged spelling. + +`TestArreflectCannotInferStructScalars` pins the upstream behavior. If +`arrow/array/arreflect` is fixed, that test fails, which is the signal to drop +these rejections and generate the columns. ## Performance @@ -213,7 +216,7 @@ Three things move the numbers: `Float64Builder.Append`. The reflection path resolves a field index, produces a `reflect.Value`, and dispatches on the builder's dynamic type, per value. - **Builder lookup.** Typed builders are resolved once, in the constructor, so - `Append` does no type assertions at all - not even the + `Append` does no type assertions at all, not even the `b.Field(i).(*array.Float64Builder)` a hand-written encoder usually repeats. - **Schema construction.** The schema is a package-level variable rather than a struct walk per call. @@ -238,9 +241,9 @@ default it builds against that rather than the tree it sits in: cd arrgen && go test ./... ``` -To test it against the local arrow-go instead - which is what you want when -changing `arrow/array/arreflect`, since the equivalence tests are what catch a -divergence - put a workspace over the two modules: +To test it against the local arrow-go instead, put a workspace over the two +modules. This is what you want when changing `arrow/array/arreflect`, since the +equivalence tests are what catch a divergence: ```sh go work init . ./arrgen # from the repository root; go.work is gitignored @@ -257,7 +260,7 @@ the repository, and runs `go generate ./...` followed by `git diff binary, which is what fails first when you edit a struct and forget to regenerate. -The golden file in `testdata` is the review surface for generator changes: +The golden file in `testdata` is where generator changes get reviewed: `go test ./arrgen/ -update` rewrites it, and the resulting diff is exactly what users would see in their own regenerated code. @@ -266,6 +269,6 @@ users would see in their own regenerated code. As a nested module, `arrgen` is versioned and tagged independently of the root module: its tags are `arrgen/vX.Y.Z`, not `vX.Y.Z`. It cannot share the root's `v18` line, because a module path without a `/vN` suffix is limited to v0 and -v1. Nothing in the release scripts tags it yet - that is a deliberate omission -for maintainers to decide on, not an oversight, and it is why the Quick Start -cannot yet name a version that resolves. +v1. Nothing in the release scripts tags it yet. That is a deliberate omission +for maintainers to decide on, and it is why the Quick Start cannot yet name a +version that resolves. diff --git a/arrgen/doc.go b/arrgen/doc.go index a4bca0318..45c8e7176 100644 --- a/arrgen/doc.go +++ b/arrgen/doc.go @@ -48,7 +48,7 @@ // bool, int8/16/32/64, int, uint8/16/32/64, uint, float32/64, string, []byte, // time.Time, time.Duration, decimal.Decimal32, decimal.Decimal64, // decimal128.Num and decimal256.Num, plus one or more pointers to any of those -// for a nullable column - a nil at any level is a null, as it is in arreflect. +// for a nullable column. A nil at any level is a null, as it is in arreflect. // // Tag options mirror arreflect: a leading name, "-" to skip a field, the // temporal overrides date32, date64, time32 and time64, dict, view, large, and @@ -56,16 +56,16 @@ // // A time.Time field must carry one of the four temporal tags, and a // decimal128.Num or decimal256.Num field must carry a decimal(precision,scale) -// tag. Untagged, all three are Go structs that arreflect's inferArrowType -// resolves through inferStructType - it switches on reflect.Kind before it -// reaches the types it matches by identity - so it infers an empty struct<> and -// drops the value. Generating the column Arrow means here would put the two -// paths out of step, so arrgen rejects those spellings instead. One consequence -// is that arrgen cannot emit a TIMESTAMP column at all: ",timestamp" is -// rejected along with the untagged spelling. +// tag. All three are Go structs, and arreflect's inferArrowType switches on +// reflect.Kind before it reaches the types it matches by identity, so untagged +// it resolves them through inferStructType, infers an empty struct<>, and drops +// the value. Generating the column Arrow means here would put the two paths out +// of step, so arrgen rejects those spellings. One consequence: arrgen cannot +// emit a TIMESTAMP column at all, since ",timestamp" is rejected along with the +// untagged spelling. // -// Anything arrgen cannot map exactly the way arreflect would - nested structs, -// slices other than []byte, arrays, maps, embedded fields, named scalar types - -// is a generate-time error naming the field, never a silently dropped column. -// Those structs still work with arreflect at runtime. +// Anything arrgen cannot map exactly the way arreflect would is a generate-time +// error naming the field, never a silently dropped column: nested structs, +// slices other than []byte, arrays, maps, embedded fields and named scalar +// types. Those structs still work with arreflect at runtime. package arrgen diff --git a/arrgen/generate.go b/arrgen/generate.go index 91b007999..af960ea56 100644 --- a/arrgen/generate.go +++ b/arrgen/generate.go @@ -213,8 +213,8 @@ type column struct { spec colSpec } -// nullable reports whether the column admits nulls, which for arreflect is a -// question about the field's outermost pointer and nothing else. +// nullable reports whether the column admits nulls. arreflect decides this from +// the outermost pointer alone. func (c column) nullable() bool { return c.ptrDepth > 0 } // collectColumns walks the struct's fields in declaration order, which is the @@ -268,9 +268,9 @@ func collectColumns(st *types.Struct) ([]column, error) { return cols, nil } -// renderAppend emits the append statement for one column, wrapping it in the -// null guards the reflection path applies: a nil pointer at any level is a -// null, and so is a nil []byte even when the column itself is not nullable. +// renderAppend emits the append statement for one column, wrapped in the same +// nil checks the reflection path applies: a nil pointer at any level means a +// null, and so does a nil []byte even when the column is not nullable. func renderAppend(idx int, c column) string { bld := fmt.Sprintf("a.b%d", idx) val := strings.Repeat("*", c.ptrDepth) + "v." + c.goField @@ -286,7 +286,7 @@ func renderAppend(idx int, c column) string { } // One nil check per pointer level, plus one on the value itself when a nil - // value is a null in its own right. + // value counts as a null. levels := c.ptrDepth if c.spec.nilable { levels++ @@ -298,13 +298,13 @@ func renderAppend(idx int, c column) string { for i := range checks { checks[i] = strings.Repeat("*", i) + "v." + c.goField + " == nil" } - guard := strings.Join(checks, " || ") + cond := strings.Join(checks, " || ") // A statement that is already a block (a time-of-day column scopes a local) - // sheds its braces on the way into the guard's else, which would otherwise - // nest two sets for no reason. + // sheds its braces on the way into the else, which would otherwise nest two + // sets for no reason. inner := strings.TrimSuffix(strings.TrimPrefix(stmt, "{\n"), "\n}") - return fmt.Sprintf("if %s {\n%s.AppendNull()\n} else {\n%s\n}", guard, bld, inner) + return fmt.Sprintf("if %s {\n%s.AppendNull()\n} else {\n%s\n}", cond, bld, inner) } type genFile struct { diff --git a/arrgen/generate_test.go b/arrgen/generate_test.go index b4aec8767..2bfbbddcd 100644 --- a/arrgen/generate_test.go +++ b/arrgen/generate_test.go @@ -228,19 +228,17 @@ func TestGenerateMissingPackage(t *testing.T) { } } -// TestArreflectCannotInferStructScalars pins the upstream behavior that the -// rejections above exist for. +// TestArreflectCannotInferStructScalars pins the upstream behavior behind the +// rejections above. time.Time, decimal128.Num and decimal256.Num are Go structs +// that Arrow models as scalars, and as struct fields arreflect infers an empty +// struct<> for all three. // -// time.Time, decimal128.Num and decimal256.Num are Go structs that Arrow models -// as scalars. arreflect's inferArrowType switches on reflect.Kind before it -// reaches the types it matches by identity, so as a struct field each one is -// resolved by inferStructType, which sees only unexported fields and yields an -// empty struct<>. arrgen refuses to generate the column Arrow means here -// because doing so would break the equivalence it promises. +// arreflect's own tests assert the intended mapping against +// inferPrimitiveArrowType, which a struct field never reaches, so they pass +// either way. // -// If arreflect learns to infer these, this test fails - which is good news, not -// a regression. It means the rejections in mapping.go can be dropped and the -// untagged spellings generated instead. +// If arreflect learns to infer these, this test fails. That is the signal to +// drop the rejections in mapping.go and generate the columns instead. func TestArreflectCannotInferStructScalars(t *testing.T) { type row struct { Time time.Time `arrow:"t"` diff --git a/arrgen/internal/gentypes/equivalence_test.go b/arrgen/internal/gentypes/equivalence_test.go index 5b09939b9..c822de568 100644 --- a/arrgen/internal/gentypes/equivalence_test.go +++ b/arrgen/internal/gentypes/equivalence_test.go @@ -158,8 +158,8 @@ func makeRows(n int) []gentypes.Row { r.PDec128 = &dec128 // Multi-level pointers. Every fourth row nils the inner level while the - // outer one stays set, so the per-level guards are covered separately - // rather than only ever failing at the first. + // outer one stays set, so each per-level nil check is covered instead of + // only ever failing at the first. pi64, pf64, pstr := &i64, &f64, &s ppf64 := &pf64 r.PPInt64, r.PPStr, r.PPPF64 = &pi64, &pstr, &ppf64 @@ -231,8 +231,8 @@ func compareSchemas(t *testing.T, got, want *arrow.Schema) { // columns. Everything else in this package - the benchmarks, the allocation // assertions - is only interesting because this holds. // -// Every column is compared: arrgen rejects at generate time any field it -// cannot map the way arreflect would, so there is nothing here to excuse. +// Every column is compared. arrgen rejects at generate time any field it cannot +// map the way arreflect would, so nothing needs skipping here. func TestRowRecordMatchesArreflect(t *testing.T) { for _, n := range []int{0, 1, 2, 3, 64} { t.Run(fmt.Sprintf("rows=%d", n), func(t *testing.T) { diff --git a/arrgen/internal/gentypes/types.go b/arrgen/internal/gentypes/types.go index 0f8121468..d684481d7 100644 --- a/arrgen/internal/gentypes/types.go +++ b/arrgen/internal/gentypes/types.go @@ -83,9 +83,9 @@ type Row struct { PDec128 *decimal128.Num `arrow:"pdec128,decimal(20,3)"` PDictS *string `arrow:"pds,dict"` - // Multi-level pointers: arreflect walks to the value and nulls the column - // if any level is nil, so these cover the guard renderAppend emits per - // level - and, for ppbin, a nil []byte behind two of them. + // Multi-level pointers. arreflect walks down to the value and writes a null + // if any level is nil, so these cover the nil check renderAppend emits per + // level, and for ppbin a nil []byte behind two of them. PPInt64 **int64 `arrow:"ppi64"` PPStr **string `arrow:"pps"` PPBin **[]byte `arrow:"ppbin"` diff --git a/arrgen/mapping.go b/arrgen/mapping.go index 8ac0018ce..95480f4aa 100644 --- a/arrgen/mapping.go +++ b/arrgen/mapping.go @@ -71,13 +71,12 @@ func checkOpts(o tagOpts, goType string, s optSupport) error { return nil } -// resolveColumn maps a Go field type and its parsed tag to an Arrow column, -// also reporting how many pointers the field is behind. +// resolveColumn maps a Go field type and its parsed tag to an Arrow column, and +// reports how many pointers the field sits behind. // -// Every level is stripped, matching arreflect, whose appendValue walks pointers -// to the value and nulls the column if any level along the way is nil. The -// column is nullable when there is at least one, which is arreflect's rule too: -// it reads nullability from the outermost pointer alone. +// Every level is stripped, matching arreflect: its appendValue walks pointers +// down to the value and writes a null if any level is nil. The column is +// nullable when there is at least one pointer, which is arreflect's rule too. func resolveColumn(t types.Type, opts tagOpts) (spec colSpec, ptrDepth int, err error) { t = types.Unalias(t) for { @@ -152,9 +151,9 @@ func isDecimalPkg(path, name string) bool { } // decimalSpec maps a decimal field. structGoType is empty for decimal.Decimal32 -// and decimal.Decimal64, whose Go types are defined integers, and the Go type's -// name for decimal128.Num and decimal256.Num, whose Go types are structs and so -// need an explicit decimal tag to stay in step with arreflect. See +// and decimal.Decimal64, which are defined integers. For decimal128.Num and +// decimal256.Num it holds the Go type name: those are structs, and arreflect +// cannot infer them without an explicit decimal tag. See // errUninferableStructScalar. func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32, structGoType string) (colSpec, error) { goType := "decimal." + arrowName @@ -175,9 +174,9 @@ func decimalSpec(opts tagOpts, arrowName string, defaultPrecision int32, structG }, nil } -// timeSpec maps a time.Time field, which needs one of the four temporal tags: -// the timestamp spelling arreflect cannot infer for a struct field is rejected -// rather than generated. See errUninferableStructScalar. +// timeSpec maps a time.Time field. One of the four temporal tags is required, +// because arreflect cannot infer a timestamp column for a struct field. See +// errUninferableStructScalar. func timeSpec(opts tagOpts) (colSpec, error) { if err := checkOpts(opts, "time.Time", optSupport{temporal: true}); err != nil { return colSpec{}, err @@ -204,8 +203,8 @@ func timeSpec(opts tagOpts) (colSpec, error) { case "time64": return timeOfDaySpec("&arrow.Time64Type{Unit: arrow.Nanosecond}", "*array.Time64Builder", "arrow.Time64", ""), nil default: - // "" and "timestamp" are the two spellings that ask for a TIMESTAMP - // column, and neither is inferable as a struct field. + // "" and "timestamp" both ask for a TIMESTAMP column, which arreflect + // cannot infer for a struct field. return colSpec{}, errUninferableStructScalar("time.Time", "one of the date32, date64, time32 or time64 tags") } } @@ -214,20 +213,16 @@ func timeSpec(opts tagOpts) (colSpec, error) { // Arrow models as a scalar, tagged in a way arreflect cannot infer. // // arreflect's inferArrowType switches on reflect.Kind before it reaches the -// types it matches by identity, so a struct field of one of these types is -// resolved by inferStructType instead, which finds only unexported fields and -// yields an empty struct<> - and the value is then dropped. Only a tag that -// names the Arrow type outright survives, because arreflect applies the tag -// after the inferred type. +// types it matches by identity, so time.Time, decimal128.Num and decimal256.Num +// are resolved by inferStructType. That sees only unexported fields, returns an +// empty struct<>, and the value is dropped. Only a tag naming the Arrow type +// survives, since arreflect applies tags after inference. // -// arrgen could emit the column Arrow plainly means here, but generated code and -// arreflect would then disagree about the schema, which is the one thing this -// generator exists not to do. So the untaggable spellings are a generate-time -// error naming the field and the tag that fixes it. +// Generating the column Arrow means here would put the two paths out of step, +// so arrgen rejects the field and names the tag that fixes it. func errUninferableStructScalar(goType, fix string) error { - return fmt.Errorf("%s needs %s: arreflect infers an empty struct<> for it as a struct field "+ - "and drops the value, so a column generated here would not match arreflect.InferSchema. "+ - "Tag it, or encode the struct with arreflect instead", goType, fix) + return fmt.Errorf("%s needs %s: as a struct field arreflect infers an empty struct<> "+ + "and drops the value, so a generated column would not match arreflect.InferSchema", goType, fix) } // timeOfDaySpec renders arreflect's timeOfDayNanos inline. It is emitted as a diff --git a/ci/scripts/test.sh b/ci/scripts/test.sh index 35b3ed496..8fbf77b8b 100755 --- a/ci/scripts/test.sh +++ b/ci/scripts/test.sh @@ -99,11 +99,10 @@ git diff --exit-code -- . popd -# The run above resolves the released arrow-go that arrgen's go.mod names, which -# is what a consumer of the generator gets. Run it against this tree as well: the -# equivalence tests are what catch a change to arrow/array/arreflect that the -# generated encoders no longer match, and in module mode they would not see one -# until after the next release. +# The run above uses the released arrow-go that arrgen's go.mod names, which is +# what a consumer gets. Run it against this tree too: the equivalence tests are +# what catch an arrow/array/arreflect change the generated encoders no longer +# match, and in module mode they would not see one until the next release. pushd "${source_dir}" rm -f go.work go.work.sum