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 new file mode 100644 index 000000000..b0ece14d4 --- /dev/null +++ b/arrgen/README.md @@ -0,0 +1,274 @@ + + +# 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 + +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 + +import "time" + +type Metric struct { + 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 tool 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. + +`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@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. 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 +``` + +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 +// 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. + +Runnable versions of all three are the testable examples in +[`example_test.go`](example_test.go). + +## 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 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`, +`large`, and `decimal(precision,scale)`. + +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 | +| --- | --- | +| `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 +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. Every column is compared, with no +exceptions. + +### 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 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`. + +`arreflect`'s own tests assert the intended mapping against +`inferPrimitiveArrowType`, which a struct field never reaches, so they pass +either way. + +`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 + +`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) | 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` (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** | | + +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 + +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 +cd arrgen && go test ./... +``` + +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 +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 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. + +### 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, and it is why the Quick Start cannot yet name a +version that resolves. 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..45c8e7176 --- /dev/null +++ b/arrgen/doc.go @@ -0,0 +1,71 @@ +// 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 +// +// Add the generator to your module as a tool dependency, then put a go:generate +// directive next to the struct: +// +// 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. 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 +// +// 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. +// +// Tag options mirror arreflect: a leading name, "-" to skip a field, the +// 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. 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 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/drift_test.go b/arrgen/drift_test.go new file mode 100644 index 000000000..724c0ad90 --- /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 + }{ + {"internal/gentypes", []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_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.go b/arrgen/generate.go new file mode 100644 index 000000000..af960ea56 --- /dev/null +++ b/arrgen/generate.go @@ -0,0 +1,361 @@ +// 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 + ptrDepth int // pointer levels between the field and its value + spec colSpec +} + +// 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 +// 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, 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(), ptrDepth: ptrDepth, spec: spec}) + } + return cols, nil +} + +// 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 + 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) + if c.spec.fallible { + stmt = "a.setErr(" + stmt + ")" + } + + // One nil check per pointer level, plus one on the value itself when a nil + // value counts as a null. + 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" + } + 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 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}", cond, 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..2bfbbddcd --- /dev/null +++ b/arrgen/generate_test.go @@ -0,0 +1,268 @@ +// 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" + "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") + +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"}, + {"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"}, + {"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") + } +} + +// 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. +// +// 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. 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"` + 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 + } + 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..f850c3bdb --- /dev/null +++ b/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" + + "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) []gentypes.Metric { + metrics := make([]gentypes.Metric, n) + v := 42.5 + for i := range metrics { + metrics[i] = gentypes.Metric{ + Day: base.AddDate(0, 0, i), + 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{ + 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 + } + } + 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 := gentypes.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 49 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..c822de568 --- /dev/null +++ b/arrgen/internal/gentypes/equivalence_test.go @@ -0,0 +1,366 @@ +// 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/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" +) + +// compareColumns asserts that got and want hold the same columns. +// +// 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, +// 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 + 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)}, + 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.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 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 + 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 +} + +// 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) + } + 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 +// 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() { + 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 !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) + } + } + if !got.Equal(want) { + t.Errorf("schemas differ beyond their fields\n got: %s\nwant: %s", got, want) + } +} + +// 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. +// +// 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) { + 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) + }) + } +} + +func TestMetricRecordMatchesArreflect(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + v := 3.5 + 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) + if err != nil { + t.Fatalf("arreflect.RecordFromSlice: %v", err) + } + defer want.Release() + + got, err := gentypes.MetricRecordBatch(mem, metrics) + if err != nil { + t.Fatalf("MetricRecordBatch: %v", err) + } + defer got.Release() + + 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) + } +} + +// 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/metric.go b/arrgen/internal/gentypes/metric.go new file mode 100644 index 000000000..6a0ff79ee --- /dev/null +++ b/arrgen/internal/gentypes/metric.go @@ -0,0 +1,33 @@ +// 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 + +import "time" + +// 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 { + 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 diff --git a/arrgen/internal/gentypes/metric_arrow.go b/arrgen/internal/gentypes/metric_arrow.go new file mode 100644 index 000000000..b0d1aae5a --- /dev/null +++ b/arrgen/internal/gentypes/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 gentypes + +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: "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}, +}, 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.Date32Builder + 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.Date32Builder), + 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.Date32FromTime(v.Day)) + 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/internal/gentypes/row_arrow.go b/arrgen/internal/gentypes/row_arrow.go new file mode 100644 index 000000000..9c8f1d2e3 --- /dev/null +++ b/arrgen/internal/gentypes/row_arrow.go @@ -0,0 +1,519 @@ +// 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: "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: 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}, + {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: "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: 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 +// 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.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.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 + b45 *array.Int64Builder + b46 *array.StringBuilder + b47 *array.BinaryBuilder + b48 *array.Float64Builder + 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.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.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), + b45: rb.Field(45).(*array.Int64Builder), + b46: rb.Field(46).(*array.StringBuilder), + b47: rb.Field(47).(*array.BinaryBuilder), + b48: rb.Field(48).(*array.Float64Builder), + } +} + +// 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.Date32FromTime(v.Date32)) + a.b16.Append(arrow.Date64FromTime(v.Date64)) + { + tod := v.Time32.UTC() + 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.b18.Append(arrow.Time64(tod.Sub(time.Date(tod.Year(), tod.Month(), tod.Day(), 0, 0, 0, 0, time.UTC)).Nanoseconds())) + } + 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.b26.AppendNull() + } else { + a.b26.Append(v.LargeBin) + } + if v.ViewBin == nil { + a.b27.AppendNull() + } else { + a.b27.Append(v.ViewBin) + } + a.setErr(a.b28.AppendString(v.DictStr)) + if v.DictBin == nil { + a.b29.AppendNull() + } else { + a.setErr(a.b29.Append(v.DictBin)) + } + a.setErr(a.b30.Append(v.DictInt)) + a.setErr(a.b31.Append(v.DictF64)) + a.b32.Append(v.Untagged) + if v.PBool == nil { + a.b33.AppendNull() + } else { + a.b33.Append(*v.PBool) + } + if v.PInt64 == nil { + a.b34.AppendNull() + } else { + a.b34.Append(*v.PInt64) + } + if v.PF64 == nil { + a.b35.AppendNull() + } else { + a.b35.Append(*v.PF64) + } + if v.PStr == nil { + a.b36.AppendNull() + } else { + a.b36.Append(*v.PStr) + } + if v.PBin == nil || *v.PBin == nil { + a.b37.AppendNull() + } else { + a.b37.Append(*v.PBin) + } + if v.PDate32 == nil { + a.b38.AppendNull() + } else { + a.b38.Append(arrow.Date32FromTime(*v.PDate32)) + } + if v.PDate64 == nil { + a.b39.AppendNull() + } else { + a.b39.Append(arrow.Date64FromTime(*v.PDate64)) + } + if v.PTime64 == nil { + a.b40.AppendNull() + } else { + tod := (*v.PTime64).UTC() + 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.b41.AppendNull() + } else { + a.b41.Append(arrow.Duration((*v.PDur).Nanoseconds())) + } + if v.PDec32 == nil { + a.b42.AppendNull() + } else { + a.b42.Append(*v.PDec32) + } + if v.PDec128 == nil { + a.b43.AppendNull() + } else { + a.b43.Append(*v.PDec128) + } + if v.PDictS == nil { + a.b44.AppendNull() + } 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. +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: "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}, + {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.Date32Builder + 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.Date32Builder), + 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.Date32FromTime(v.Day)) + 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..d684481d7 --- /dev/null +++ b/arrgen/internal/gentypes/types.go @@ -0,0 +1,110 @@ +// 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"` + + 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,decimal(40,5)"` + + 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"` + 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,decimal(20,3)"` + PDictS *string `arrow:"pds,dict"` + + // 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"` + PPPF64 ***float64 `arrow:"pppf64"` + + 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 { + 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/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..95480f4aa --- /dev/null +++ b/arrgen/mapping.go @@ -0,0 +1,353 @@ +// 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, and +// reports how many pointers the field sits behind. +// +// 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 { + p, ok := t.(*types.Pointer) + if !ok { + break + } + ptrDepth++ + t = types.Unalias(p.Elem()) + } + spec, err = baseSpec(t, opts) + return spec, ptrDepth, 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, "decimal128.Num") + case isDecimalPkg(pkg, "decimal256") && obj.Name() == "Num": + 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) + } + 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 +} + +// decimalSpec maps a decimal field. structGoType is empty for decimal.Decimal32 +// 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 + 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 + } + 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. 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 + } + 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: + // "" 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") + } +} + +// 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 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. +// +// 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: 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 +// 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..3890fa9e9 --- /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: "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}, +}, 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.Date64Builder + 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.Date64Builder), + 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.Date64FromTime(v.Day)) + 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..39dd901b6 --- /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 { + Day time.Time `arrow:"day,date64"` + 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..dd452f149 --- /dev/null +++ b/arrgen/testdata/errors/errors.go @@ -0,0 +1,148 @@ +// 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" + + "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 { + 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 TimeSlice struct { + Field []time.Time +} + +// 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 +} + +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..d9fadce97 --- /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,date64"` + 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/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..8fbf77b8b 100755 --- a/ci/scripts/test.sh +++ b/ci/scripts/test.sh @@ -81,3 +81,33 @@ 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 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 +go work init . ./arrgen +go test "${test_args[@]}" ./arrgen/... +rm -f go.work go.work.sum + +popd