Add aws_s3_parquet_stream output for memory-efficient Parquet streaming to S3 - #661
Add aws_s3_parquet_stream output for memory-efficient Parquet streaming to S3#661triddell wants to merge 29 commits into
Conversation
Adds a new output plugin that writes Parquet files to S3 using multipart uploads with incremental row group streaming, reducing memory usage by 88% compared to the batch approach. Key features: - Streams row groups directly to S3 (no full file buffering) - partition_by parameter for correct partition routing - schema_file parameter for external schema definitions - Supports all Parquet compression types (snappy, zstd, gzip, etc.)
- Added S3API interface to enable mock-based testing - Updated StreamingParquetWriter to use S3API instead of concrete *s3.Client - Added 8 new unit tests covering S3 multipart upload lifecycle: - Initialization and double-init protection - Write lifecycle (before init, after close) - Part upload tracking - CompleteMultipartUpload on Close - Multiple parts with correct numbering - AbortMultipartUpload on errors
Fixes "insufficient definition levels" errors by: 1. Using consistent schema/message type generation (both with struct tags + pointers) 2. Preserving actual column metadata from temp files instead of approximating
Fixed Nested Optional StructsFound and fixed the "insufficient definition levels" bug for nested optional STRUCT fields. Root causes:
Fix:
|
Tests verify column metadata extraction from temp files and proper definition/repetition level encoding in streaming parquet output.
- Fix documentation typo: timestamp().format() → now().ts_format() - Replace fmt.Errorf with errors.New for static error strings - Update AWS SDK endpoint resolver to use BaseEndpoint - Update testify assertions to use specific assertion methods - Remove unused test struct fields - Regenerate documentation
CI FixesFixed all linting errors and CI check failures:
All tests and lint checks now pass. I will review the generated documentation output more closely for any additional improvements needed. |
Preserve ColumnIndex and OffsetIndex from temporary parquet files to enable query engine optimization through page-level pruning and predicate pushdown. - Extract and store page index data during row group flush - Write in correct Parquet format: all ColumnIndex, then all OffsetIndex - Track filePosition separately from uploadSize for accurate offset calculation - Add comprehensive unit tests for page index preservation
Page Index Preservation FixIssueFurther in-depth testing revealed that while page index metadata pointers (ColumnIndexOffset, OffsetIndexOffset, etc.) were being copied from temporary files, the actual ColumnIndex and OffsetIndex binary data was not preserved. This prevented query engines from utilizing page-level pruning. Solution
|
This commit addresses three critical issues with the parquet streaming writer and schema generation that were causing file corruption and compatibility issues: 1. Clear page index and bloom filter offsets in streaming writer The streaming writer doesn't include page indexes or bloom filters, so these offsets are invalid. If left set, parquet-go will panic with "slice bounds out of range" errors when reading the files. - Clear MetaData.IndexPageOffset and BloomFilterOffset - Clear ColumnChunk-level OffsetIndex and ColumnIndex offsets 2. Preserve LogicalType annotations for all field types Critical for systems like Apache Iceberg that validate types and rely on modern LogicalType annotations (STRING, INT, TIMESTAMP, LIST, MAP, etc.) 3. Fix optional LIST type handling in schema generation Don't wrap slice types in pointers when optional. The parquet-go library requires "list" tag on slice types, not pointer types. LIST types already handle nullability through 3-level encoding.
Critical Bug FixesThis update addresses three critical issues with the parquet streaming writer: 1. Page Index & Bloom Filter Offset Handling
|
Config naming consistency
Memory fix
Parquet schema modernization
Production validation
|
Expose column_index_enabled, column_index_size_limit, and data_page_statistics config fields to give users control over Parquet metadata generation for query optimization
…mplementation-specific docs
Docs Check Failing Despite Running
|
Signed-off-by: Jem Davies <[email protected]>
…dell/bento into feature/streaming-s3-parquet-output
Signed-off-by: Jem Davies <[email protected]>
Did a merge commit and a make docs and seems to be working now - taking a look at this PR 👀 |
| }), | ||
| ).Description("Parquet schema. Mutually exclusive with schema_file.").Optional(), | ||
| service.NewStringField(spsoFieldSchemaFile). | ||
| Description("Path to a YAML file containing a Parquet schema definition. The file should contain a parquet_encode processor resource with a schema section. Mutually exclusive with schema."). |
There was a problem hiding this comment.
I think that this is somewhat confusing - I think that it would be better to only add schemas via the schema field - and later we could add a PR that would enable loading schemas via a file to the relevant components, (i.e. this one, parquet_encode & parquet_decode).
There was a problem hiding this comment.
Re: schema_file Support
I understand the desire for consistency, but there's a strong practical need for schema_file in production use cases.
Real-world schema complexity:
Production schemas (OCSF, security events, etc.) are typically 200-300+ lines with deeply nested structures. Here's an example of a real OCSF Network Activity schema we use in production: [link to schema_example.yml - 279 lines, 4+ levels of nested STRUCTs]
Embedding schemas this large inline in pipeline configs is impractical because:
- Maintainability: 279 lines of schema embedded in config obscures the actual pipeline logic
- Reusability: Multiple pipelines (dev/staging/prod environments) share the same schema definition
- Version control: Teams need to track schema evolution separately from pipeline config changes
- Validation: External schema files can be validated/tested independently
Why this component specifically needs it:
Unlike parquet_encode (which processes batches), this streaming output:
- Needs the schema at initialization to create the Parquet writer pool
- Schema must be stable across multiple files/partitions during the entire output lifecycle
- Schema loading happens once per output initialization, not per message/batch
Proposed path forward:
I'd like to propose Option A: Keep schema_file in this PR, with a commitment to add it to parquet_encode and parquet_decode in a follow-up PR for consistency. This PR establishes the pattern that others can follow.
Alternatively, I can extract the schema loading logic into internal/impl/parquet/schema_loader.go right now, making it available to all parquet components immediately.
There was a problem hiding this comment.
Sorry I wasn't really clear here.
When I said it is somewhat confusing it was around the implementation of the schema_file field. For instance, I was thinking that because you are required to include it as a processor_resource & parquet_encode when actually a Bento stream config could not contain them, that would be confusing for a reader of config to understand.
I have added comments regarding this here: #661 (review)
Then also the function loadSchemaFromFile() does more yaml marshalling / un-marshalling than necessary & duplication of the schema fields, which again added to poor code readablity.
My suggestion to take out this, was more to reduce the volume of code that is being added in the PR, and potentially add this in a separate feature later - making code review more manageable.
There is a real need for a feature like this, i.e. the ability to have a yaml file that contains some bento stream config that can be referenced by other stream config files.
It seems that we have a few ways to do this but they all fall short of being able to do exactly what we need here:
- yaml anchors - it is possible to reduce config repetition via yaml anchors, however they cannot cross file boundries.
- templates - You could use Bento Templates to avoid having a large schema obscure pipeline logic, but you would need to duplicate the schema for each template
- resources - these are good but require a full & valid component configuration - so it doesn't allow supplying only part of it i.e.
schema.
We should be looking to add this functionality in such that it could be used by all components rather than just aws_s3_parquet_stream and the other parquet_encode & parquet_decode I think.
| Description("Forces path style URLs for S3 requests."). | ||
| Default(false). | ||
| Advanced(), | ||
| service.NewObjectListField("schema", |
There was a problem hiding this comment.
Dedupe these from other parquet components?
There was a problem hiding this comment.
Re: Schema Field Duplication
Good catch - there is duplication of schema field configuration (~50 lines) between this output and the parquet package.
I can refactor this to extract a shared parquetSchemaFieldConfig() function that both parquet_encode and this output can use. This would:
- Eliminate duplication
- Ensure schema field definitions stay consistent
- Make future schema enhancements easier
Should I include this refactoring in this PR, or would you prefer it as a follow-up?
| } | ||
|
|
||
| func (s *s3ParquetStreamOutput) Connect(ctx context.Context) error { | ||
| s.log.Infof("Streaming Parquet output configured for bucket: %s", s.conf.Bucket) |
There was a problem hiding this comment.
The bento engine will already produce a log for this so remove this one.
There was a problem hiding this comment.
Re: Redundant log in Connect method
Good point - removed the redundant log line. The Bento engine already logs output connections, so this was unnecessary noise.
| ) | ||
|
|
||
| // S3API defines the S3 operations needed for multipart uploads | ||
| type S3API interface { |
There was a problem hiding this comment.
Duplicate of s3StreamingAPI in output_s3_stream_writer.go
There was a problem hiding this comment.
Re: Duplicate of s3StreamingAPI in output_s3_stream_writer.go
Good catch! There is indeed duplication here. I have two identical interfaces defined:
- Line 18 in
output_s3_stream_writer.go:s3StreamingAPI(unexported) - Line 25 in
output_aws_s3_parquet_stream_writer.go:S3API(exported)
Both define the same 4 S3 multipart upload methods.
Proposed Solution
I'll extract this to a shared exported interface that both components can use. This provides:
- Eliminates duplication - single source of truth
- Better testability - exported interface can be mocked in tests
- Future reusability - any other S3 multipart upload components can use it
Implementation:
Create internal/impl/aws/s3_interfaces.go:
package aws
import (
"context"
"github.com/aws/aws-sdk-go-v2/service/s3"
)
// S3MultipartAPI defines the S3 operations required for multipart uploads
type S3MultipartAPI interface {
CreateMultipartUpload(ctx context.Context, input *s3.CreateMultipartUploadInput, opts ...func(*s3.Options)) (*s3.CreateMultipartUploadOutput, error)
UploadPart(ctx context.Context, input *s3.UploadPartInput, opts ...func(*s3.Options)) (*s3.UploadPartOutput, error)
CompleteMultipartUpload(ctx context.Context, input *s3.CompleteMultipartUploadInput, opts ...func(*s3.Options)) (*s3.CompleteMultipartUploadOutput, error)
AbortMultipartUpload(ctx context.Context, input *s3.AbortMultipartUploadInput, opts ...func(*s3.Options)) (*s3.AbortMultipartUploadOutput, error)
}Then update both files to use S3MultipartAPI and remove their local interface definitions.
Should I proceed with this refactoring?
|
Questions regarding this:
I can see that there is an identical interface included with
The PR description talks about issues relating to the
Perhaps it would be better to refactor some of the parquet fields such they are deduplicated rather than copied into the aws package? 🤔 |
Co-authored-by: Jem Davies <[email protected]>
Co-authored-by: Jem Davies <[email protected]>
Re: Do we need this after
|
| // loadSchemaFromFile loads a Parquet schema from an external YAML file. | ||
| // The file should contain a processor_resources section with a parquet_encode processor. | ||
| func loadSchemaFromFile(filePath string) (*service.ParsedConfig, error) { | ||
| // Read the file |
There was a problem hiding this comment.
Redundant comment, I suppose it's somewhat subjective but why include a comment "Read the file" for then the subsequent line os.ReadFile()? It is self-explanatory.
|
|
||
| // Parse into a ParsedConfig | ||
| // We use a dummy ConfigSpec that matches the schema structure | ||
| spec := service.NewConfigSpec().Field( |
There was a problem hiding this comment.
These will have been duplicated 3 times now, from the parquet package & above in this output definition.
We could move these into a function:
func getSchemaFields() *service.ConfigField {
return service.NewObjectListField("schema",
service.NewStringField("name").Description("The name of the column."),
service.NewStringEnumField("type", "BOOLEAN", "INT8", "INT16", "INT32", "INT64", "DECIMAL64", "DECIMAL32", "FLOAT", "DOUBLE", "BYTE_ARRAY", "UTF8", "MAP", "LIST", "STRUCT").
Description("The type of the column, only applicable for leaf columns with no child fields. STRUCT represents nested objects with defined field schemas. MAP supports only string keys, but can support values of all types. Some logical types can be specified here such as UTF8.").Optional(),
service.NewIntField("decimal_precision").Description("Precision to use for DECIMAL32/DECIMAL64 type").Default(0),
service.NewIntField("decimal_scale").Description("Scale to use for DECIMAL32/DECIMAL64 type").Default(0),
service.NewBoolField("repeated").Description("Whether the field is repeated.").Default(false),
service.NewBoolField("optional").Description("Whether the field is optional.").Default(false),
service.NewAnyListField("fields").Description("A list of child fields.").Optional().Example([]any{
map[string]any{
"name": "foo",
"type": "INT64",
},
map[string]any{
"name": "bar",
"type": "BYTE_ARRAY",
},
}),
).Description("Parquet schema. Mutually exclusive with schema_file.").Optional()
}|
|
||
| // loadSchemaFromFile loads a Parquet schema from an external YAML file. | ||
| // The file should contain a processor_resources section with a parquet_encode processor. | ||
| func loadSchemaFromFile(filePath string) (*service.ParsedConfig, error) { |
There was a problem hiding this comment.
This function could be made much simpler:
func loadSchemaFromFile(filePath string) (*service.ParsedConfig, error) {
data, err := os.ReadFile(filePath)
if err != nil {
return nil, err
}
spec := service.NewConfigSpec().Fields(getSchemaFields())
return spec.ParseYAML(string(data), nil)
}Then the contents of the file could be:
schema:
- name: id
type: INT64
- name: message
type: UTF8
- name: value
type: DOUBLE
optional: trueThere is no real need to have processor_resources & parquet_encode in there?
… API Updates the S3 Parquet streaming writer to use thrift.New() wrapper for nullable fields, required by parquet-go v0.29.0. Changes: - ColumnIndexSizeLimit now takes a function parameter - Schema element nullable fields (NumChildren, RepetitionType, Type, LogicalType) now use thrift.New() instead of raw pointers
…r writer Each partition_by key gets its own StreamingParquetWriter, and every one of them pre-allocated a rowBuffer sized to the full row_group_size up front — even a writer that only ever buffers a handful of rows. That's fine with a low writer count, but a partition_by with a high-cardinality key (e.g. account+webacl) can create dozens of concurrent writers per invocation, and in a short-lived Lambda that memory is paid entirely upfront, all at once, with no time for anything to settle or get reclaimed before the process exits. Cap the initial capacity at initialRowBufferCap (128) regardless of row_group_size; Go's normal slice growth handles writers that actually need more. Close() already flushes any partial buffer, so this only changes an allocation-size detail, not behavior — confirmed via the existing writer test suite (all passing except one pre-existing, unrelated thrift type-equality failure).
…l-file offsets Close() rebased ColumnChunk.MetaData.DataPageOffset/DictionaryPageOffset from temp-buffer-relative to final-file-absolute positions, but the OffsetIndex's PageLocation.Offset values were copied verbatim, so every row group after the first described page positions as if it started at byte 0. Strict readers (arrow-rs's SerializedPageReader) panic with an offset underflow on read.
2 participants
Overview
This PR adds a new
aws_s3_parquet_streamoutput plugin that streams Parquet files directly to S3 using multipart uploads, addressing memory limitations and partition routing issues in the current batch approach.Resolves: #660
Related: #659 (partition routing bug)
Motivation
Problem 1: High Memory Usage
The current approach (
aws_s3+parquet_encodeprocessor) buffers entire Parquet files in memory before uploading to S3:Memory usage for 100K events (minimal schema):
For production workloads with complex schemas (OCSF) or larger datasets, memory can exceed 10+ GB, causing OOM errors in constrained environments.
Problem 2: Partition Routing
As documented in #659, the batch approach evaluates path expressions once per batch, causing incorrect partition routing:
With 100K events across 2 accounts:
This breaks Iceberg/Hive/Delta table partitioning and causes silent data corruption.
Solution
This PR introduces
aws_s3_parquet_streamwith:partition_byexpressionsArchitecture
Memory profile per writer: ~30-60 MB (independent of dataset size)
Key Features
1. Partition Routing with
partition_bypartition_byexpressions2. Schema Loading
Supports both inline and external schema definitions:
3. Compression Support
All Parquet compression types supported:
uncompressedsnappy(default)gzipzstdbrotlilz4raw4. Configurable Row Groups
Implementation Details
Files Added
output_aws_s3_parquet_stream.go(543 lines)parquet_streaming_writer.go(623 lines)output_aws_s3_parquet_stream_test.go(327 lines)parquet_streaming_writer_test.go(216 lines)output_aws_s3_parquet_stream_localstack_test.go(266 lines)Files Modified
internal/impl/parquet/convert.go(+17 lines)json.Numbertype support for int32/int64 conversionjson.Numberinternal/impl/parquet/schema.go(+26 lines)SchemaOptsstruct for use by output pluginsTesting
Unit Tests
Integration Tests
partition_bywith multiple partitionsManual Testing
Extensive testing documented in project with:
Performance
Test: 100K events with minimal schema
Test: Multiple partitions (2 dates × 2 accounts)
Backwards Compatibility
parquetpackageConfiguration Example
Use Cases
This output is particularly valuable for: