diff --git a/purego/NEXT_CHANGELOG.md b/purego/NEXT_CHANGELOG.md index fad2f3e7..d40506bb 100644 --- a/purego/NEXT_CHANGELOG.md +++ b/purego/NEXT_CHANGELOG.md @@ -38,6 +38,19 @@ and encoder, ack-model, and opener seams let a protocol instantiate the core over its own payload type. Proto and JSON behavior is unchanged, and nothing is exposed through a public API yet. +- Add `internal/arrowproto`, the Arrow IPC payload for the upcoming Arrow + ingestion path. A payload is a canonical self-contained IPC stream materialized + when it is built, so the core can hold it across a reconnect without pinning the + caller's `RecordBatch`, and it can be sliced by row so a partially acknowledged + batch replays only its unacknowledged suffix. Admission sizes a batch from the + rows it covers rather than the whole buffers it points at, because a slice + shares its parent's buffers and would otherwise be charged for the entire + parent. Nested children are sized from the same row window, since slicing + rebases only the top-level node and leaves a struct's fields and a list's + values spanning the whole parent. Nothing is exposed through a public API yet. +- Add the `github.com/apache/arrow-go/v18` dependency. arrow-go requires + `google.golang.org/grpc` v1.82.0, which raises this module's grpc minimum from + v1.81.1. ### Breaking Changes diff --git a/purego/go.mod b/purego/go.mod index 17015b55..1e746cdc 100644 --- a/purego/go.mod +++ b/purego/go.mod @@ -3,13 +3,21 @@ module github.com/databricks/zerobus-sdk/purego go 1.25.0 require ( - google.golang.org/grpc v1.81.1 + github.com/apache/arrow-go/v18 v18.7.0 + google.golang.org/grpc v1.82.0 google.golang.org/protobuf v1.36.11 ) require ( - golang.org/x/net v0.51.0 // indirect - golang.org/x/sys v0.42.0 // indirect - golang.org/x/text v0.34.0 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 // indirect + github.com/goccy/go-json v0.10.6 // indirect + github.com/google/flatbuffers v25.12.19+incompatible // indirect + github.com/klauspost/compress v1.19.0 // indirect + github.com/klauspost/cpuid/v2 v2.4.0 // indirect + github.com/pierrec/lz4/v4 v4.1.27 // indirect + github.com/zeebo/xxh3 v1.1.0 // indirect + golang.org/x/exp v0.0.0-20260112195511-716be5621a96 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.47.0 // indirect + golang.org/x/text v0.37.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect ) diff --git a/purego/go.sum b/purego/go.sum index 44c671d1..6b517da9 100644 --- a/purego/go.sum +++ b/purego/go.sum @@ -1,15 +1,41 @@ +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/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +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/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +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/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +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.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= 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= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= @@ -22,17 +48,21 @@ go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfC go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= -golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo= -golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y= -golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= -golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk= -golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA= +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/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 h1:ggcbiqK8WWh6l1dnltU4BgWGIGo+EVYxCaAPih/zQXQ= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= -google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ= -google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.82.0 h1:vguDnZUPjE26w09A63VoxZPnvPjB5Riyc0mkXPFmAIU= +google.golang.org/grpc v1.82.0/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +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/purego/internal/arrowproto/encoder.go b/purego/internal/arrowproto/encoder.go new file mode 100644 index 00000000..623ac3ed --- /dev/null +++ b/purego/internal/arrowproto/encoder.go @@ -0,0 +1,630 @@ +// Package arrowproto implements the Arrow IPC payload for the Arrow ingestion +// path. A payload holds only self-contained IPC bytes, so caller-owned Arrow +// arrays are never retained by the stream core. +package arrowproto + +import ( + "bytes" + "fmt" + "math" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// payloadOverheadBytes charges the Payload container so tiny batches still +// consume byte capacity. +const payloadOverheadBytes = int64(64) + +// admissionSlopBytes is headroom for IPC framing a buffer sum cannot see: batch +// metadata, buffer padding, end markers. It bounds the estimate, not occupancy — +// a smaller MaxBufferedPayloadBytes rejects every batch. +const admissionSlopBytes = int64(64 * 1024) + +// Compression selects the compression used when serializing Arrow IPC batches. +type Compression uint8 + +const ( + CompressionNone Compression = iota + CompressionLZ4 + CompressionZstd +) + +// Options configures Arrow IPC materialization. +type Options struct { + Compression Compression + + // Allocator is mainly a test seam. Nil uses Arrow's default. + Allocator memory.Allocator +} + +// Payload is one eagerly materialized, non-empty Arrow RecordBatch held as a +// canonical self-contained IPC stream: schema, dictionaries, batch, end marker. +type Payload struct { + ipcBytes []byte + rows uint64 +} + +// UnitCount returns the number of row durability units in the payload. +func (p *Payload) UnitCount() uint64 { + if p == nil { + return 0 + } + return p.rows +} + +// RetainedSize returns a conservative heap charge for the payload. +func (p *Payload) RetainedSize() int64 { + if p == nil { + return 0 + } + return payloadOverheadBytes + int64(cap(p.ipcBytes)) +} + +// IPCBytes returns a caller-owned copy of the self-contained IPC stream. +func (p *Payload) IPCBytes() []byte { + if p == nil { + return nil + } + return bytes.Clone(p.ipcBytes) +} + +// Protocol owns one exact Arrow schema and the IPC encoding policy. +type Protocol struct { + schema *arrow.Schema + compression Compression + allocator memory.Allocator + admissionBaseBytes int64 +} + +// New constructs an Arrow protocol encoder, copying schema so the caller's +// Schema object is not retained. +func New(schema *arrow.Schema, options Options) (*Protocol, error) { + if schema == nil { + return nil, fmt.Errorf("arrow protocol: schema is required") + } + if _, err := compressionOption(options.Compression); err != nil { + return nil, err + } + allocator := options.Allocator + if allocator == nil { + allocator = memory.DefaultAllocator + } + ownedSchema, err := cloneSchemaThroughIPC(schema, allocator) + if err != nil { + return nil, err + } + schemaIPC, err := encodeSchemaIPC(ownedSchema) + if err != nil { + return nil, err + } + return &Protocol{ + schema: ownedSchema, + compression: options.Compression, + allocator: allocator, + // Every payload repeats the schema, and materialization holds a second + // copy of it while writing, so both are charged up front. + admissionBaseBytes: admissionSlopBytes + 2*int64(len(schemaIPC)), + }, nil +} + +// EncodeRecordBatch serializes batch immediately. The returned payload owns no +// references to batch, its columns, or their buffers. +func (p *Protocol) EncodeRecordBatch(batch arrow.RecordBatch) (*Payload, error) { + if batch == nil { + return nil, fmt.Errorf("arrow protocol: RecordBatch is required") + } + if err := p.validateBatch(batch); err != nil { + return nil, err + } + serialized, err := p.serialize(batch) + if err != nil { + return nil, fmt.Errorf("arrow protocol: serialize RecordBatch: %w", err) + } + return p.payloadFromCanonicalIPC(serialized, uint64(batch.NumRows())) +} + +// EstimateRecordBatchRetainedSize returns a conservative pre-materialization +// reservation: the batch sized from the rows it covers, doubled to cover +// framing, compression, and buffer growth, plus the per-payload schema base. +// The core admits on this value before encoding, so an over-estimate rejects a +// batch that would have fit. +func (p *Protocol) EstimateRecordBatchRetainedSize( + batch arrow.RecordBatch, +) (int64, error) { + if batch == nil { + return 0, fmt.Errorf("arrow protocol: RecordBatch is required") + } + if err := p.validateBatch(batch); err != nil { + return 0, err + } + return p.admissionEstimate(totalRecordBufferSize(batch)), nil +} + +func (p *Protocol) admissionEstimate(inputBytes int64) int64 { + if inputBytes < 0 || + inputBytes > math.MaxInt64/2 || + p.admissionBaseBytes > math.MaxInt64-payloadOverheadBytes { + return math.MaxInt64 + } + total := payloadOverheadBytes + p.admissionBaseBytes + if scaled := inputBytes * 2; scaled <= math.MaxInt64-total { + return total + scaled + } + return math.MaxInt64 +} + +// Slice drops an acknowledged row prefix and reserializes the remaining suffix +// as a standalone IPC payload. +func (p *Protocol) Slice(payload *Payload, acknowledgedPrefix uint64) (*Payload, error) { + if payload == nil { + return nil, fmt.Errorf("arrow protocol: payload is nil") + } + // The core only slices a partially acknowledged payload, so a prefix outside + // (0,rows) means its accounting is wrong. Say so rather than invent a payload. + if acknowledgedPrefix == 0 || acknowledgedPrefix >= payload.rows { + return nil, fmt.Errorf( + "arrow protocol: acknowledged prefix %d is invalid for %d rows", + acknowledgedPrefix, + payload.rows, + ) + } + + batch, err := p.decodeOne(payload.ipcBytes) + if err != nil { + return nil, fmt.Errorf("arrow protocol: decode payload for slicing: %w", err) + } + defer batch.Release() + // The prefix was checked against the header count, so a decoded batch that + // disagrees would make NewSlice panic on an out-of-range bound. + if uint64(batch.NumRows()) != payload.rows { + return nil, fmt.Errorf( + "arrow protocol: payload row count changed: header=%d decoded=%d", + payload.rows, + batch.NumRows(), + ) + } + + suffix := batch.NewSlice(int64(acknowledgedPrefix), batch.NumRows()) + defer suffix.Release() + serialized, err := p.serialize(suffix) + if err != nil { + return nil, fmt.Errorf("arrow protocol: serialize sliced suffix: %w", err) + } + return p.payloadFromCanonicalIPC(serialized, uint64(suffix.NumRows())) +} + +func (p *Protocol) payloadFromCanonicalIPC(data []byte, rows uint64) (*Payload, error) { + if len(data) == 0 || rows == 0 { + return nil, fmt.Errorf("arrow protocol: canonical IPC payload is empty") + } + // data is the serializer's own buffer, which nothing else aliases, so the + // payload can adopt it without a second full copy. + return &Payload{ipcBytes: data, rows: rows}, nil +} + +func (p *Protocol) validateBatch(batch arrow.RecordBatch) error { + if !exactSchemaEqual(batch.Schema(), p.schema) { + return fmt.Errorf( + "arrow protocol: RecordBatch schema does not exactly match stream schema", + ) + } + if batch.NumRows() <= 0 { + return fmt.Errorf("arrow protocol: RecordBatch must contain at least one row") + } + return nil +} + +func exactSchemaEqual(left, right *arrow.Schema) bool { + return left != nil && right != nil && + left.Equal(right) && + left.Metadata().Equal(right.Metadata()) +} + +// rowWindow is the absolute row range of one array node that a batch covers. +// Slicing rebases only the top-level nodes: a struct's fields stay parallel to +// the parent's offset and a list's values are reached through its offsets +// buffer, so a child still spans the parent's whole extent and cannot report +// its own covered rows. +type rowWindow struct { + offset int64 + length int64 +} + +func windowOf(data arrow.ArrayData) rowWindow { + return rowWindow{offset: int64(data.Offset()), length: int64(data.Len())} +} + +// absentArrayData reports whether data is missing. array.Data returns a nil +// *Data as a non-nil interface for a node without a dictionary, so the typed +// nil has to be caught before any method call on it. +func absentArrayData(data arrow.ArrayData) bool { + if data == nil { + return true + } + concrete, ok := data.(*array.Data) + return ok && concrete == nil +} + +// totalRecordBufferSize sums the bytes a batch's columns hold. A slice shares +// its parent's buffers, so a buffer's own length reports the parent's whole +// extent: charging that would reject a ten-row slice of a million-row batch. +// Every node — nested children included — is therefore charged from the row +// window the batch covers rather than from buffer lengths. Layouts with no +// per-row rule (list-view, union, run-end-encoded) still fall back to whole +// buffers, an over-estimate that reconciliation corrects once the payload +// exists. +func totalRecordBufferSize(batch arrow.RecordBatch) int64 { + seen := make(map[*memory.Buffer]struct{}) + var total int64 + for _, column := range batch.Columns() { + data := column.Data() + size, err := addInt64Saturating( + total, + arrayDataSize(data, windowOf(data), seen), + ) + if err != nil { + return math.MaxInt64 + } + total = size + } + return total +} + +func arrayDataSize( + data arrow.ArrayData, + window rowWindow, + seen map[*memory.Buffer]struct{}, +) int64 { + if absentArrayData(data) { + return 0 + } + total := ownedBufferSize(data, window, seen) + for _, child := range data.Children() { + if absentArrayData(child) { + continue + } + size, err := addInt64Saturating( + total, + arrayDataSize(child, childWindow(data, window, child), seen), + ) + if err != nil { + return math.MaxInt64 + } + total = size + } + // A dictionary is shared whole rather than sliced per row, so it is charged + // in full through the same recursion. + dictionary := data.Dictionary() + if absentArrayData(dictionary) { + return total + } + size, err := addInt64Saturating( + total, + arrayDataSize(dictionary, windowOf(dictionary), seen), + ) + if err != nil { + return math.MaxInt64 + } + return size +} + +// childWindow maps a parent's covered rows onto one of its children. A child +// carries its own offset as well, so the parent's range is applied on top of +// it. An unrecognized nesting falls back to the child's whole extent. +func childWindow( + parent arrow.ArrayData, + window rowWindow, + child arrow.ArrayData, +) rowWindow { + base := int64(child.Offset()) + switch dataType := parent.DataType().(type) { + case *arrow.StructType: + return rowWindow{offset: base + window.offset, length: window.length} + case *arrow.ListType, *arrow.MapType: + if first, last, ok := offsetRange(parent, window, false); ok { + return rowWindow{offset: base + first, length: last - first} + } + case *arrow.LargeListType: + if first, last, ok := offsetRange(parent, window, true); ok { + return rowWindow{offset: base + first, length: last - first} + } + case *arrow.FixedSizeListType: + width := int64(dataType.Len()) + if width > 0 && window.offset <= math.MaxInt64/width && + window.length <= math.MaxInt64/width { + return rowWindow{ + offset: base + window.offset*width, + length: window.length * width, + } + } + } + return windowOf(child) +} + +// ownedBufferSize charges one array node for the rows it actually covers. +func ownedBufferSize( + data arrow.ArrayData, + window rowWindow, + seen map[*memory.Buffer]struct{}, +) int64 { + rows := window.length + buffers := data.Buffers() + if rows < 0 { + return wholeBufferSize(buffers, seen) + } + var validityBytes int64 + if len(buffers) > 0 && buffers[0] != nil { + validityBytes = bitmapBytes(rows) + } + switch dataType := data.DataType().(type) { + case *arrow.StringType, *arrow.BinaryType: + if valueBytes, ok := variableWidthValueBytes(data, window, false); ok { + return validityBytes + (rows+1)*4 + valueBytes + } + case *arrow.LargeStringType, *arrow.LargeBinaryType: + if valueBytes, ok := variableWidthValueBytes(data, window, true); ok { + return validityBytes + (rows+1)*8 + valueBytes + } + case *arrow.ListType, *arrow.MapType: + // The values are a child node, so only the offsets are charged here. + if len(buffers) == 2 { + return validityBytes + (rows+1)*4 + } + case *arrow.LargeListType: + if len(buffers) == 2 { + return validityBytes + (rows+1)*8 + } + case *arrow.StructType, *arrow.FixedSizeListType: + // Both own a validity bitmap only; every value lives in a child. + if len(buffers) <= 1 { + return validityBytes + } + case arrow.FixedWidthDataType: + // A wider layout would put row data in a buffer this arithmetic does not + // know about, so only the canonical validity+values shape is derived. + if len(buffers) <= 2 { + return validityBytes + fixedWidthValueBytes(rows, dataType.BitWidth()) + } + } + return wholeBufferSize(buffers, seen) +} + +// variableWidthValueBytes sizes exactly the values the covered rows reference. +func variableWidthValueBytes( + data arrow.ArrayData, + window rowWindow, + large bool, +) (int64, bool) { + if len(data.Buffers()) != 3 { + return 0, false + } + first, last, ok := offsetRange(data, window, large) + if !ok { + return 0, false + } + return last - first, true +} + +// offsetRange reads an offsets-buffer layout to find the first and last value +// index the covered rows reference. +func offsetRange( + data arrow.ArrayData, + window rowWindow, + large bool, +) (first, last int64, ok bool) { + buffers := data.Buffers() + if len(buffers) < 2 || buffers[1] == nil { + return 0, 0, false + } + start := window.offset + end := start + window.length + if start < 0 || end < start { + return 0, 0, false + } + if large { + offsets := arrow.Int64Traits.CastFromBytes(buffers[1].Bytes()) + if int64(len(offsets)) <= end { + return 0, 0, false + } + first, last = offsets[start], offsets[end] + } else { + offsets := arrow.Int32Traits.CastFromBytes(buffers[1].Bytes()) + if int64(len(offsets)) <= end { + return 0, 0, false + } + first, last = int64(offsets[start]), int64(offsets[end]) + } + if first < 0 || last < first { + return 0, 0, false + } + return first, last, true +} + +func fixedWidthValueBytes(rows int64, bitWidth int) int64 { + if bitWidth <= 0 { + return 0 + } + if bitWidth < 8 { + return bitmapBytes(rows) + } + width := int64(bitWidth / 8) + if rows > math.MaxInt64/width { + return math.MaxInt64 + } + return rows * width +} + +func bitmapBytes(rows int64) int64 { + return (rows + 7) / 8 +} + +func wholeBufferSize( + buffers []*memory.Buffer, + seen map[*memory.Buffer]struct{}, +) int64 { + var total int64 + for _, buffer := range buffers { + if buffer == nil { + continue + } + if _, exists := seen[buffer]; exists { + continue + } + seen[buffer] = struct{}{} + size, err := addInt64Saturating(total, int64(buffer.Len())) + if err != nil { + return math.MaxInt64 + } + total = size + } + return total +} + +func addInt64Saturating(left, right int64) (int64, error) { + if left < 0 || right < 0 || right > math.MaxInt64-left { + return math.MaxInt64, fmt.Errorf("int64 size overflow") + } + return left + right, nil +} + +// encodeSchemaIPC serializes schema as a schema-only Arrow IPC stream. +func encodeSchemaIPC(schema *arrow.Schema) ([]byte, error) { + var output bytes.Buffer + writer := ipc.NewWriter(&output, ipc.WithSchema(schema)) + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("arrow protocol: serialize schema: %w", err) + } + return output.Bytes(), nil +} + +// cloneSchemaThroughIPC rebuilds nested and dictionary DataTypes independently. A +// shallow field copy would still alias mutable pointer-backed types such as +// DictionaryType. +func cloneSchemaThroughIPC( + schema *arrow.Schema, + allocator memory.Allocator, +) (*arrow.Schema, error) { + var output bytes.Buffer + writer := ipc.NewWriter( + &output, + ipc.WithSchema(schema), + ipc.WithAllocator(allocator), + ) + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("arrow protocol: serialize schema: %w", err) + } + reader, err := ipc.NewReader( + bytes.NewReader(output.Bytes()), + ipc.WithAllocator(allocator), + ) + if err != nil { + return nil, fmt.Errorf("arrow protocol: deserialize schema: %w", err) + } + defer reader.Release() + decoded := reader.Schema() + metadata := decoded.Metadata() + return arrow.NewSchemaWithEndian( + decoded.Fields(), + &metadata, + decoded.Endianness(), + ), nil +} + +// decodeOne parses exactly one non-empty RecordBatch carrying the protocol +// schema, rejecting anything the payload contract does not allow. +func (p *Protocol) decodeOne(data []byte) (arrow.RecordBatch, error) { + if len(data) == 0 { + return nil, fmt.Errorf("arrow protocol: IPC input is empty") + } + source := bytes.NewReader(data) + reader, err := ipc.NewReader( + source, + ipc.WithAllocator(p.allocator), + ) + if err != nil { + return nil, fmt.Errorf("arrow protocol: invalid IPC stream: %w", err) + } + defer reader.Release() + if !exactSchemaEqual(reader.Schema(), p.schema) { + return nil, fmt.Errorf("arrow protocol: IPC schema does not exactly match stream schema") + } + if !reader.Next() { + if err := reader.Err(); err != nil { + return nil, fmt.Errorf("arrow protocol: read IPC RecordBatch: %w", err) + } + return nil, fmt.Errorf("arrow protocol: IPC stream contains no RecordBatch") + } + batch := reader.RecordBatch() + batch.Retain() + if batch.NumRows() <= 0 { + batch.Release() + return nil, fmt.Errorf("arrow protocol: IPC RecordBatch must contain at least one row") + } + if reader.Next() { + batch.Release() + return nil, fmt.Errorf( + "arrow protocol: IPC stream must contain exactly one RecordBatch", + ) + } + if err := reader.Err(); err != nil { + batch.Release() + return nil, fmt.Errorf("arrow protocol: read trailing IPC data: %w", err) + } + if source.Len() != 0 { + batch.Release() + return nil, fmt.Errorf( + "arrow protocol: IPC stream contains %d trailing bytes", + source.Len(), + ) + } + return batch, nil +} + +func (p *Protocol) serialize(batch arrow.RecordBatch) ([]byte, error) { + var output bytes.Buffer + options, err := p.ipcOptions() + if err != nil { + return nil, err + } + options = append(options, ipc.WithSchema(p.schema)) + writer := ipc.NewWriter(&output, options...) + if err := writer.Write(batch); err != nil { + _ = writer.Close() + return nil, err + } + if err := writer.Close(); err != nil { + return nil, err + } + return output.Bytes(), nil +} + +func (p *Protocol) ipcOptions() ([]ipc.Option, error) { + compression, err := compressionOption(p.compression) + if err != nil { + return nil, err + } + options := []ipc.Option{ipc.WithAllocator(p.allocator)} + if compression != nil { + options = append(options, compression) + } + return options, nil +} + +func compressionOption(compression Compression) (ipc.Option, error) { + switch compression { + case CompressionNone: + return nil, nil + case CompressionLZ4: + return ipc.WithLZ4(), nil + case CompressionZstd: + return ipc.WithZstd(), nil + default: + return nil, fmt.Errorf( + "arrow protocol: unsupported IPC compression %d", + compression, + ) + } +} diff --git a/purego/internal/arrowproto/encoder_test.go b/purego/internal/arrowproto/encoder_test.go new file mode 100644 index 00000000..a267c571 --- /dev/null +++ b/purego/internal/arrowproto/encoder_test.go @@ -0,0 +1,491 @@ +package arrowproto + +import ( + "bytes" + "strings" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +func idSchema(metadata *arrow.Metadata) *arrow.Schema { + return arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: false, + }}, metadata) +} + +func idBatch( + t *testing.T, + allocator memory.Allocator, + schema *arrow.Schema, + values []int32, +) arrow.RecordBatch { + t.Helper() + builder := array.NewInt32Builder(allocator) + builder.AppendValues(values, nil) + column := builder.NewArray() + builder.Release() + record := array.NewRecordBatch(schema, []arrow.Array{column}, int64(len(values))) + column.Release() + return record +} + +func binaryBatch( + t *testing.T, + allocator memory.Allocator, + rows int, + valueBytes int, +) (*arrow.Schema, arrow.RecordBatch) { + t.Helper() + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "value", Type: arrow.BinaryTypes.String, Nullable: false}, + }, nil) + idBuilder := array.NewInt32Builder(allocator) + valueBuilder := array.NewStringBuilder(allocator) + value := strings.Repeat("x", valueBytes) + for i := range rows { + idBuilder.Append(int32(i)) + valueBuilder.Append(value) + } + ids := idBuilder.NewArray() + values := valueBuilder.NewArray() + idBuilder.Release() + valueBuilder.Release() + record := array.NewRecordBatch( + schema, + []arrow.Array{ids, values}, + int64(rows), + ) + ids.Release() + values.Release() + return schema, record +} + +func listBatch( + t *testing.T, + allocator memory.Allocator, + rows int, + valuesPerRow int, +) (*arrow.Schema, arrow.RecordBatch) { + t.Helper() + schema := arrow.NewSchema([]arrow.Field{{ + Name: "value", + Type: arrow.ListOf(arrow.PrimitiveTypes.Int64), + }}, nil) + builder := array.NewListBuilder(allocator, arrow.PrimitiveTypes.Int64) + values, ok := builder.ValueBuilder().(*array.Int64Builder) + if !ok { + t.Fatalf("list value builder = %T, want *array.Int64Builder", + builder.ValueBuilder()) + } + for i := range rows { + builder.Append(true) + for j := range valuesPerRow { + values.Append(int64(i*valuesPerRow + j)) + } + } + column := builder.NewArray() + builder.Release() + record := array.NewRecordBatch(schema, []arrow.Array{column}, int64(rows)) + column.Release() + return schema, record +} + +// structBatch builds one struct column whose string field holds valueBytes(row) +// bytes, so a slice's cost is dominated by a nested variable-width child. +func structBatch( + t *testing.T, + allocator memory.Allocator, + rows int, + valueBytes func(row int) int, +) (*arrow.Schema, arrow.RecordBatch) { + t.Helper() + structType := arrow.StructOf( + arrow.Field{Name: "id", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.String}, + ) + schema := arrow.NewSchema( + []arrow.Field{{Name: "entry", Type: structType}}, + nil, + ) + builder := array.NewStructBuilder(allocator, structType) + ids, ok := builder.FieldBuilder(0).(*array.Int64Builder) + if !ok { + t.Fatalf("field 0 builder = %T, want *array.Int64Builder", + builder.FieldBuilder(0)) + } + values, ok := builder.FieldBuilder(1).(*array.StringBuilder) + if !ok { + t.Fatalf("field 1 builder = %T, want *array.StringBuilder", + builder.FieldBuilder(1)) + } + for i := range rows { + builder.Append(true) + ids.Append(int64(i)) + values.Append(strings.Repeat("x", valueBytes(i))) + } + column := builder.NewArray() + builder.Release() + record := array.NewRecordBatch(schema, []arrow.Array{column}, int64(rows)) + column.Release() + return schema, record +} + +func readIDs(t *testing.T, data []byte) []int32 { + t.Helper() + reader, err := ipc.NewReader(bytes.NewReader(data)) + if err != nil { + t.Fatalf("new IPC reader: %v", err) + } + defer reader.Release() + if !reader.Next() { + t.Fatalf("read IPC record: %v", reader.Err()) + } + column, ok := reader.RecordBatch().Column(0).(*array.Int32) + if !ok { + t.Fatalf("first column type = %T, want *array.Int32", reader.RecordBatch().Column(0)) + } + values := append([]int32(nil), column.Int32Values()...) + if reader.Next() { + t.Fatal("decoded IPC contains more than one record") + } + if err := reader.Err(); err != nil { + t.Fatalf("finish IPC reader: %v", err) + } + return values +} + +func equalInt32(left, right []int32) bool { + if len(left) != len(right) { + return false + } + for i := range left { + if left[i] != right[i] { + return false + } + } + return true +} + +func TestTypedRecordBatchRoundTripOwnsIPC(t *testing.T) { + allocator := memory.NewCheckedAllocator(memory.DefaultAllocator) + schema := idSchema(nil) + protocol, err := New(schema, Options{Allocator: allocator}) + if err != nil { + t.Fatalf("New: %v", err) + } + if allocator.CurrentAlloc() != 0 { + t.Fatalf("schema construction retained %d allocator bytes", allocator.CurrentAlloc()) + } + + record := idBatch(t, allocator, schema, []int32{10, 20, 30}) + estimate, err := protocol.EstimateRecordBatchRetainedSize(record) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize: %v", err) + } + payload, err := protocol.EncodeRecordBatch(record) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + record.Release() + + if got := payload.UnitCount(); got != 3 { + t.Fatalf("UnitCount = %d, want 3", got) + } + if got := readIDs(t, payload.IPCBytes()); !equalInt32(got, []int32{10, 20, 30}) { + t.Fatalf("round-trip ids = %v", got) + } + if payload.RetainedSize() < int64(len(payload.ipcBytes)) { + t.Fatalf( + "RetainedSize = %d, smaller than IPC bytes %d", + payload.RetainedSize(), + len(payload.ipcBytes), + ) + } + if payload.RetainedSize() > estimate { + t.Fatalf( + "actual retained size %d exceeds admission estimate %d", + payload.RetainedSize(), + estimate, + ) + } + + exported := payload.IPCBytes() + exported[0] ^= 0xff + if bytes.Equal(exported, payload.IPCBytes()) { + t.Fatal("IPCBytes returned bytes aliasing the retained payload") + } + allocator.AssertSize(t, 0) +} + +func TestCanonicalPayloadTakesSerializerOwnership(t *testing.T) { + protocol, err := New(idSchema(nil), Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + canonical := []byte{1, 2, 3, 4} + payload, err := protocol.payloadFromCanonicalIPC(canonical, 1) + if err != nil { + t.Fatalf("payloadFromCanonicalIPC: %v", err) + } + if &payload.ipcBytes[0] != &canonical[0] { + t.Fatal("canonical IPC was copied instead of ownership being transferred") + } +} + +func TestExactSchemaMatchIncludesMetadata(t *testing.T) { + expectedMetadata := arrow.NewMetadata([]string{"owner"}, []string{"expected"}) + actualMetadata := arrow.NewMetadata([]string{"owner"}, []string{"different"}) + protocol, err := New(idSchema(&expectedMetadata), Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, idSchema(&actualMetadata), []int32{1}) + defer record.Release() + + if _, err := protocol.EncodeRecordBatch(record); err == nil { + t.Fatal("typed RecordBatch with different schema metadata accepted") + } +} + +func TestProtocolOwnsPointerBackedSchemaTypes(t *testing.T) { + dictionaryType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int8, + ValueType: arrow.BinaryTypes.String, + } + schema := arrow.NewSchema([]arrow.Field{{ + Name: "category", Type: dictionaryType, Nullable: false, + }}, nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + + dictionaryType.ValueType = arrow.BinaryTypes.Binary + owned := protocol.schema.Field(0).Type.(*arrow.DictionaryType) + if owned.ValueType.ID() != arrow.STRING { + t.Fatalf("owned dictionary value type = %v, want string", owned.ValueType) + } +} + +func TestSliceReserializesUnacknowledgedSuffix(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, schema, []int32{1, 2, 3, 4, 5}) + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + + suffix, err := protocol.Slice(payload, 2) + if err != nil { + t.Fatalf("Slice: %v", err) + } + if suffix.UnitCount() != 3 { + t.Fatalf("suffix units = %d, want 3", suffix.UnitCount()) + } + if got := readIDs(t, suffix.IPCBytes()); !equalInt32(got, []int32{3, 4, 5}) { + t.Fatalf("suffix ids = %v", got) + } + if got := readIDs(t, payload.IPCBytes()); !equalInt32(got, []int32{1, 2, 3, 4, 5}) { + t.Fatalf("original payload changed after Slice: %v", got) + } + if _, err := protocol.Slice(payload, 5); err == nil { + t.Fatal("Slice accepted a fully acknowledged prefix") + } + // The core never slices an unacknowledged payload, so a zero prefix is a bug + // to surface, not a payload to pass through. + if _, err := protocol.Slice(payload, 0); err == nil { + t.Fatal("Slice accepted an unacknowledged prefix") + } +} + +// TestSliceRejectsPayloadWhoseRowCountDrifted keeps a corrupted invariant on the +// error path. Slicing against a bound past the decoded batch panics, and that +// panic would escape through recovery into the caller's goroutine. +func TestSliceRejectsPayloadWhoseRowCountDrifted(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, schema, []int32{1, 2, 3, 4}) + defer record.Release() + payload, err := protocol.EncodeRecordBatch(record) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + + drifted := &Payload{ipcBytes: payload.ipcBytes, rows: 64} + _, err = protocol.Slice(drifted, 32) + if err == nil || !strings.Contains(err.Error(), "row count changed") { + t.Fatalf("Slice of a drifted payload = %v, want a row-count error", err) + } +} + +func TestIPCCompressionOptionsRoundTrip(t *testing.T) { + schema, record := binaryBatch(t, memory.DefaultAllocator, 2_000, 200) + defer record.Release() + + sizes := make(map[Compression]int) + for _, compression := range []Compression{ + CompressionNone, + CompressionLZ4, + CompressionZstd, + } { + protocol, err := New(schema, Options{Compression: compression}) + if err != nil { + t.Fatalf("New compression %d: %v", compression, err) + } + payload, err := protocol.EncodeRecordBatch(record) + if err != nil { + t.Fatalf("EncodeRecordBatch compression %d: %v", compression, err) + } + sizes[compression] = len(payload.ipcBytes) + reader, err := ipc.NewReader(bytes.NewReader(payload.IPCBytes())) + if err != nil { + t.Fatalf("read compression %d: %v", compression, err) + } + if !reader.Next() || reader.RecordBatch().NumRows() != 2_000 { + t.Fatalf("compression %d did not round-trip: %v", compression, reader.Err()) + } + reader.Release() + } + if sizes[CompressionLZ4] >= sizes[CompressionNone] { + t.Errorf("LZ4 size %d >= uncompressed %d", sizes[CompressionLZ4], sizes[CompressionNone]) + } + if sizes[CompressionZstd] >= sizes[CompressionNone] { + t.Errorf("Zstd size %d >= uncompressed %d", sizes[CompressionZstd], sizes[CompressionNone]) + } + if _, err := New(schema, Options{Compression: Compression(99)}); err == nil { + t.Fatal("unsupported compression accepted") + } +} + +// TestTypedAdmissionChargesSliceNotParentBuffers pins the reason admission sizes +// a batch from the rows it covers: a slice shares its parent's buffers, so a +// whole-buffer measurement charges a small slice for the entire parent. Nested +// layouts are covered because slicing rebases only the top-level node, leaving +// a struct's fields and a list's values spanning the whole parent. +func TestTypedAdmissionChargesSliceNotParentBuffers(t *testing.T) { + const ( + rows = 20_000 + valueBytes = 512 + sliceRows = 10 + ) + cases := []struct { + name string + build func(*testing.T) (*arrow.Schema, arrow.RecordBatch) + }{{ + name: "string", + build: func(t *testing.T) (*arrow.Schema, arrow.RecordBatch) { + return binaryBatch(t, memory.DefaultAllocator, rows, valueBytes) + }, + }, { + name: "list", + build: func(t *testing.T) (*arrow.Schema, arrow.RecordBatch) { + return listBatch(t, memory.DefaultAllocator, rows, valueBytes/8) + }, + }, { + name: "struct", + build: func(t *testing.T) (*arrow.Schema, arrow.RecordBatch) { + return structBatch(t, memory.DefaultAllocator, rows, + func(int) int { return valueBytes }) + }, + }} + + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + schema, record := testCase.build(t) + defer record.Release() + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + + slice := record.NewSlice(rows-sliceRows, rows) + defer slice.Release() + estimate, err := protocol.EstimateRecordBatchRetainedSize(slice) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize: %v", err) + } + payload, err := protocol.EncodeRecordBatch(slice) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + if estimate < payload.RetainedSize() { + t.Fatalf( + "slice estimate %d under-reserves its %d-byte payload", + estimate, + payload.RetainedSize(), + ) + } + parentEstimate, err := protocol.EstimateRecordBatchRetainedSize(record) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize(parent): %v", err) + } + if estimate > parentEstimate/100 { + t.Fatalf( + "slice estimate %d tracks the %d-byte parent rather than its own rows", + estimate, + parentEstimate, + ) + } + }) + } +} + +// TestNestedSliceAdmissionCoversHeavyTailRows pins that a nested child's window +// follows its parent's offset. Sizing the wrong rows of a variable-width child +// looks plausible on uniform data, but under-reserves once the sliced rows are +// the large ones. +func TestNestedSliceAdmissionCoversHeavyTailRows(t *testing.T) { + const ( + rows = 2_000 + sliceRows = 4 + smallLen = 16 + largeLen = 256 * 1024 + ) + schema, record := structBatch( + t, + memory.DefaultAllocator, + rows, + func(row int) int { + if row >= rows-sliceRows { + return largeLen + } + return smallLen + }, + ) + defer record.Release() + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + + slice := record.NewSlice(rows-sliceRows, rows) + defer slice.Release() + estimate, err := protocol.EstimateRecordBatchRetainedSize(slice) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize: %v", err) + } + payload, err := protocol.EncodeRecordBatch(slice) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + if estimate < payload.RetainedSize() { + t.Fatalf( + "estimate %d under-reserves the %d-byte heavy-tail payload", + estimate, + payload.RetainedSize(), + ) + } +}