twuebi opened a new pull request, #1506: URL: https://github.com/apache/arrow-go/pull/1506
### Rationale for this change When the sink fails while a column chunk is being flushed, the Parquet writer leaves allocator memory charged after `Close` returns. With `memory.CheckedAllocator` the bytes stay charged, and with the cgo mallocator they leak. The default Go allocator hides the problem because the GC collects the buffers. Four early returns cause this: - `columnWriter.Close` registers its release defer only after `WriteDictionaryPage` and `FlushBufferedDataPages`, so a column whose flush fails releases nothing. - `rowGroupWriter.Close` returns at the first column writer whose `Close` fails and caches that error. The remaining column writers are never closed, and a second `Close` only returns the cached error. - `arrowColumnWriter.Write` defers the level builder releases inside the loop, after `NextColumn`. If `NextColumn` fails because closing the previous column failed, the builders and the arrays they retain are never released. - `pqarrow.FileWriter.Close` returns when the row group fails to close, without calling `file.Writer.Close`, so the sink is never closed and the encryption keys are never wiped. ### What changes are included in this PR? - `columnWriter.Close` registers the release defer before the flushes, so it runs on every return path. The defer is unchanged and runs once, so the success path behaves as before. - `rowGroupWriter.Close` closes every non-nil column writer, keeps the first error, and caches it as before. - `arrowColumnWriter.Write` defers each builder release before the loop. - `pqarrow.FileWriter.Close` always closes the underlying `file.Writer` and returns the row group error first. The early return in `rowGroupWriter.Close` when `checkRowsWritten` fails is left alone. Closing the column writers there would write pages for a row group already known to be invalid, and fixing it would need a separate discard path. ### Are these changes tested? Yes. `TestWriterReleasesBuffersWhenSinkFails` in `parquet/file` writes a plain string column and a dictionary-encoded string column through serial and buffered row groups. The writes use `memory.NewCheckedAllocator` through `parquet.WithAllocator`, and the sink fails at every byte offset after the magic bytes. Each case asserts that the error wraps the sink error and that `mem.AssertSize(t, 0)` holds. `TestFileWriterCloseReleasesBuffersWhenSinkFails` in `parquet/pqarrow` does the same through `pqarrow.FileWriter`, using `Write` and `WriteBuffered` on an int8 column, a dictionary column and a plain string column. It also asserts that the sink is closed. Both tests fail on main and pass with this change. Reverting each of the four fixes on its own makes at least one subtest fail. On macOS arm64 with Go 1.26.10, `go test ./parquet/...` passes except `TestFileEncryptionDecryption`, which fails the same way on main because the local parquet-testing checkout is missing a file. `go test -tags assert,test` on `parquet/file` and `parquet/pqarrow`, `-race` on the new and adjacent tests, `go vet` and `golangci-lint` all pass. ### Are there any user-facing changes? Yes. A failed file write no longer leaves allocator memory charged. When the row group close fails, `pqarrow.FileWriter.Close` now also closes the underlying writer and sink. This includes the internal `Close` that `Write` makes after a column write fails. There is no API change. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
