From 213dcfb673f406e7900cb939639f1c014542d25b Mon Sep 17 00:00:00 2001 From: Wes <169498386+wes-mil@users.noreply.github.com> Date: Fri, 14 Aug 2026 13:41:24 -0400 Subject: [PATCH 1/4] add parquet --- README.md | 10 +- cmd/retriever/README.md | 28 +++ cmd/retriever/main.go | 1 + cmd/retriever/main_test.go | 2 +- go.mod | 7 + go.sum | 12 + retriever/compression.go | 5 +- retriever/dump.go | 138 +++++++---- retriever/dump_checkpoint.go | 80 +++++-- retriever/dump_checkpoint_test.go | 332 +++++++++++++++++++++++++++ retriever/dump_test.go | 47 +++- retriever/fragment_writer.go | 150 ++++++++++++ retriever/fragment_writer_test.go | 365 ++++++++++++++++++++++++++++++ retriever/options.go | 1 + retriever/options_test.go | 7 + retriever/parquet.go | 111 +++++++++ retriever/parquet_test.go | 210 +++++++++++++++++ 17 files changed, 1431 insertions(+), 75 deletions(-) create mode 100644 retriever/fragment_writer.go create mode 100644 retriever/fragment_writer_test.go create mode 100644 retriever/parquet.go create mode 100644 retriever/parquet_test.go diff --git a/README.md b/README.md index 39fec353..db1f947f 100644 --- a/README.md +++ b/README.md @@ -81,10 +81,12 @@ against a previous JSONL baseline. manifest-based collections of compressed JSONL fragments. It supports PostgreSQL and Neo4j, uncompressed, gzip, and zstd fragments, bounded keyset scans, resumable dump checkpoints, checksum validation before load, optional -deterministic property scrubbing, and a read-throughput benchmark mode. It can -also package dumps as single HPKE/ML-KEM encrypted TAR archives. -See [cmd/retriever/README.md](cmd/retriever/README.md) for dump, encrypted -archive, load, scrubbed dump, metrics verification, and benchmark examples. +deterministic property scrubbing, optional write-only Parquet sidecars for +analytical readers, and a read-throughput benchmark mode. It can also package +dumps as single HPKE/ML-KEM encrypted TAR archives. See +[cmd/retriever/README.md](cmd/retriever/README.md) for dump, Parquet sidecar, +encrypted archive, load, scrubbed dump, metrics verification, and benchmark +examples. The same import/export functionality is available to library consumers from `github.com/specterops/dawgs/retriever`; callers provide an already-open `graph.Database`, and archive helpers support both path-based and stream-based diff --git a/cmd/retriever/README.md b/cmd/retriever/README.md index 67aba113..26122cb8 100644 --- a/cmd/retriever/README.md +++ b/cmd/retriever/README.md @@ -37,6 +37,34 @@ Use repeated `-graph` flags to dump multiple named graphs. For PostgreSQL, validates that expected node and edge partitions exist. For Neo4j, `-all-graphs` means the selected Neo4j database only. +Pass `-parquet` to write optional Parquet sidecars alongside the JSONL +fragments: + +```bash +retriever dump \ + -connection "$CONNECTION_STRING" \ + -out ./dumpdir \ + -graph default \ + -scrub none \ + -parquet +``` + +For example, `graphs/default/nodes-000001.jsonl.zst` is paired with +`graphs/default/nodes-000001.parquet`, and +`graphs/default/edges-000001.jsonl.zst` with +`graphs/default/edges-000001.parquet`. Node sidecar rows contain `id` (string), +`kinds` (list of strings), and `properties` (`VARIANT`). Edge sidecar rows +contain `start_id`, `end_id`, and `kind` (strings), plus `properties` +(`VARIANT`). The `VARIANT` property preserves JSON-like scalar, object, array, +and null values for analytical readers. + +Parquet sidecars are optional output and are not part of the production +collection path: loading, collection verification, manifest checksums and +accounting, and encrypted archives remain JSONL-only. Interrupted Parquet +dumps resume from committed JSONL/Parquet shard pairs. For each checkpointed +JSONL fragment, resume requires the Parquet partner to be present as a regular +file; it does not reopen or validate the sidecar's Parquet contents. + Existing non-empty output directories are refused unless `-force` is supplied. The manifest is written last as `manifest.json`; if a dump fails before that point, the directory is intentionally left for inspection without a success diff --git a/cmd/retriever/main.go b/cmd/retriever/main.go index 1b5ca896..30768f0f 100644 --- a/cmd/retriever/main.go +++ b/cmd/retriever/main.go @@ -91,6 +91,7 @@ func (s commandRuntime) runDump(ctx context.Context, args []string) error { flags.StringVar(&cfg.OutputDir, "out", "", "Output collection directory.") flags.BoolVar(&cfg.Force, "force", false, "Replace an existing non-empty output directory.") flags.BoolVar(&cfg.Resume, "resume", false, "Resume an interrupted dump from its validated checkpoint.") + flags.BoolVar(&cfg.Parquet, "parquet", false, "Also write Parquet sidecars for each JSONL fragment.") flags.StringVar(&archiveOut, "archive-out", "", "Optional encrypted archive output path.") flags.StringVar(&recipientPath, "recipient", "", "Recipient public key for -archive-out.") flags.StringVar(&scrubValue, "scrub", string(cfg.Scrub), "Scrub mode: none or full.") diff --git a/cmd/retriever/main_test.go b/cmd/retriever/main_test.go index 7e265278..310748a0 100644 --- a/cmd/retriever/main_test.go +++ b/cmd/retriever/main_test.go @@ -33,7 +33,7 @@ func TestCommandRuntimeHelpAndValidation(t *testing.T) { t.Fatalf("expected unknown command error, got %v", err) } - err = runtime.run(context.Background(), []string{"dump", "-out", t.TempDir(), "-scrub", "full"}) + err = runtime.run(context.Background(), []string{"dump", "-out", t.TempDir(), "-parquet", "-scrub", "full"}) if err == nil || !strings.Contains(err.Error(), "-scrub full requires") { t.Fatalf("expected scrub salt validation error, got %v", err) } diff --git a/go.mod b/go.mod index 1f380c05..1e8551d6 100644 --- a/go.mod +++ b/go.mod @@ -15,6 +15,7 @@ require ( github.com/jackc/pgx/v5 v5.10.0 github.com/klauspost/compress v1.19.0 github.com/neo4j/neo4j-go-driver/v5 v5.28.4 + github.com/parquet-go/parquet-go v0.32.0 github.com/pashagolub/pgxmock/v5 v5.1.0 github.com/pelletier/go-toml/v2 v2.4.3 github.com/stretchr/testify v1.11.1 @@ -59,6 +60,7 @@ require ( github.com/alfatraining/structtag v1.0.0 // indirect github.com/alingse/asasalint v0.0.11 // indirect github.com/alingse/nilnesserr v0.2.0 // indirect + github.com/andybalholm/brotli v1.1.1 // indirect github.com/ashanbrown/forbidigo/v2 v2.3.0 // indirect github.com/ashanbrown/makezero/v2 v2.1.0 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect @@ -120,6 +122,7 @@ require ( github.com/golangci/swaggoswag v0.0.0-20250504205917-77f2aca3143e // indirect github.com/golangci/unconvert v0.0.0-20250410112200-a129a6e6413e // indirect github.com/google/go-cmp v0.7.0 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/gordonklaus/ineffassign v0.2.0 // indirect github.com/gostaticanalysis/analysisutil v0.7.1 // indirect github.com/gostaticanalysis/comment v1.5.0 // indirect @@ -172,6 +175,9 @@ require ( github.com/nishanths/exhaustive v0.12.0 // indirect github.com/nishanths/predeclared v0.2.2 // indirect github.com/nunnatsa/ginkgolinter v0.23.0 // indirect + github.com/parquet-go/bitpack v1.0.0 // indirect + github.com/parquet-go/jsonlite v1.0.0 // indirect + github.com/pierrec/lz4/v4 v4.1.21 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/prometheus/client_golang v1.22.0 // indirect @@ -212,6 +218,7 @@ require ( github.com/timonwong/loggercheck v0.11.0 // indirect github.com/tomarrell/wrapcheck/v2 v2.12.0 // indirect github.com/tommy-muehle/go-mnd/v2 v2.5.1 // indirect + github.com/twpayne/go-geom v1.6.1 // indirect github.com/ultraware/funlen v0.2.0 // indirect github.com/ultraware/whitespace v0.2.0 // indirect github.com/uudashr/gocognit v1.2.1 // indirect diff --git a/go.sum b/go.sum index bfc4bb5c..805c2027 100644 --- a/go.sum +++ b/go.sum @@ -60,6 +60,8 @@ github.com/alingse/asasalint v0.0.11 h1:SFwnQXJ49Kx/1GghOFz1XGqHYKp21Kq1nHad/0WQ github.com/alingse/asasalint v0.0.11/go.mod h1:nCaoMhw7a9kSJObvQyVzNTPBDbNpdocqrSP7t/cW5+I= github.com/alingse/nilnesserr v0.2.0 h1:raLem5KG7EFVb4UIDAXgrv3N2JIaffeKNtcEXkEWd/w= github.com/alingse/nilnesserr v0.2.0/go.mod h1:1xJPrXonEtX7wyTq8Dytns5P2hNzoWymVUIaKm4HNFg= +github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= +github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA= github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ= github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw= github.com/ashanbrown/forbidigo/v2 v2.3.0 h1:OZZDOchCgsX5gvToVtEBoV2UWbFfI6RKQTir2UZzSxo= @@ -443,10 +445,18 @@ github.com/otiai10/curr v0.0.0-20150429015615-9b4961190c95/go.mod h1:9qAhocn7zKJ github.com/otiai10/curr v1.0.0/go.mod h1:LskTG5wDwr8Rs+nNQ+1LlxRjAtTZZjtJW4rMXl6j4vs= github.com/otiai10/mint v1.3.0/go.mod h1:F5AjcsTsWUqX+Na9fpHb52P8pcRX2CI6A3ctIT91xUo= github.com/otiai10/mint v1.3.1/go.mod h1:/yxELlJQ0ufhjUwhshSj+wFjZ78CnZ48/1wtmBH1OTc= +github.com/parquet-go/bitpack v1.0.0 h1:AUqzlKzPPXf2bCdjfj4sTeacrUwsT7NlcYDMUQxPcQA= +github.com/parquet-go/bitpack v1.0.0/go.mod h1:XnVk9TH+O40eOOmvpAVZ7K2ocQFrQwysLMnc6M/8lgs= +github.com/parquet-go/jsonlite v1.0.0 h1:87QNdi56wOfsE5bdgas0vRzHPxfJgzrXGml1zZdd7VU= +github.com/parquet-go/jsonlite v1.0.0/go.mod h1:nDjpkpL4EOtqs6NQugUsi0Rleq9sW/OtC1NnZEnxzF0= +github.com/parquet-go/parquet-go v0.32.0 h1:NWDqTUHfrCS4cJP/Fj2HlxvqsrVedWG3sayMkf+znzM= +github.com/parquet-go/parquet-go v0.32.0/go.mod h1:navtkAYr2LGoJVp141oXPlO/sxLvaOe3la2JEoD8+rg= github.com/pashagolub/pgxmock/v5 v5.1.0 h1:NZ4pl82b335sEGIbD/+tk2fVIgVs3yNWr1R42ukpUvU= github.com/pashagolub/pgxmock/v5 v5.1.0/go.mod h1:8IJct22b7+EuqecVmYb9aKiENJLLqTsbjFHXH/znAEg= github.com/pelletier/go-toml/v2 v2.4.3 h1:GTRvJQutkOSftxIFD5xw9aepkYNuPWmVJpffdDPYVpY= github.com/pelletier/go-toml/v2 v2.4.3/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +github.com/pierrec/lz4/v4 v4.1.21 h1:yOVMLb6qSIDP67pl/5F7RepeKYu/VmTyEXvuMI5d9mQ= +github.com/pierrec/lz4/v4 v4.1.21/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -568,6 +578,8 @@ github.com/tomarrell/wrapcheck/v2 v2.12.0 h1:H/qQ1aNWz/eeIhxKAFvkfIA+N7YDvq6TWVF github.com/tomarrell/wrapcheck/v2 v2.12.0/go.mod h1:AQhQuZd0p7b6rfW+vUwHm5OMCGgp63moQ9Qr/0BpIWo= github.com/tommy-muehle/go-mnd/v2 v2.5.1 h1:NowYhSdyE/1zwK9QCLeRb6USWdoif80Ie+v+yU8u1Zw= github.com/tommy-muehle/go-mnd/v2 v2.5.1/go.mod h1:WsUAkMJMYww6l/ufffCD3m+P7LEvr8TnZn9lwVDlgzw= +github.com/twpayne/go-geom v1.6.1 h1:iLE+Opv0Ihm/ABIcvQFGIiFBXd76oBIar9drAwHFhR4= +github.com/twpayne/go-geom v1.6.1/go.mod h1:Kr+Nly6BswFsKM5sd31YaoWS5PeDDH2NftJTK7Gd028= github.com/ultraware/funlen v0.2.0 h1:gCHmCn+d2/1SemTdYMiKLAHFYxTYz7z9VIDRaTGyLkI= github.com/ultraware/funlen v0.2.0/go.mod h1:ZE0q4TsJ8T1SQcjmkhN/w+MceuatI6pBFSxxyteHIJA= github.com/ultraware/whitespace v0.2.0 h1:TYowo2m9Nfj1baEQBjuHzvMRbp19i+RCcRYrSWoFa+g= diff --git a/retriever/compression.go b/retriever/compression.go index 1e1b07dc..932918e4 100644 --- a/retriever/compression.go +++ b/retriever/compression.go @@ -156,11 +156,14 @@ func newDecompressionReader(reader io.Reader, codec CompressionCodec) (io.ReadCl } func newCompressedJSONLinesWriter(path string, codec CompressionCodec, zstdLevel int) (*compressedJSONLinesWriter, error) { + return newCompressedJSONLinesWriterAtPaths(path, path+".tmp", codec, zstdLevel) +} + +func newCompressedJSONLinesWriterAtPaths(path, tempPath string, codec CompressionCodec, zstdLevel int) (*compressedJSONLinesWriter, error) { if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { return nil, fmt.Errorf("create fragment directory: %w", err) } - tempPath := path + ".tmp" file, err := os.OpenFile(tempPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) if err != nil { return nil, fmt.Errorf("open fragment temp file: %w", err) diff --git a/retriever/dump.go b/retriever/dump.go index b9baf10f..804c79b9 100644 --- a/retriever/dump.go +++ b/retriever/dump.go @@ -613,8 +613,9 @@ func dumpNodePhase(ctx context.Context, db graph.Database, targetGraph graph.Gra var ( files = append([]FileManifest(nil), committedFiles...) - fragmentWriter *compressedJSONLinesWriter + fragmentWriter *fragmentWriter[FragmentNode] fragmentRelativePath string + parquetRelativePath string shardActionCounts scrubActionCounts shardNumber = len(committedFiles) + 1 lastWrittenID graph.ID @@ -634,13 +635,13 @@ func dumpNodePhase(ctx context.Context, db graph.Database, targetGraph graph.Gra files = append(files, fileEntry) if !hasLastWrittenID { - _ = os.Remove(filepath.Join(options.OutputDir, filepath.FromSlash(fileEntry.Path))) + removePublishedFragmentOutputs(options.OutputDir, fileEntry.Path, parquetRelativePath, options.Parquet) files = files[:len(files)-1] return fmt.Errorf("node fragment %q closed without a committed source cursor", fileEntry.Path) } if onCommit != nil { if err := onCommit(fileEntry, lastWrittenID); err != nil { - _ = os.Remove(filepath.Join(options.OutputDir, filepath.FromSlash(fileEntry.Path))) + removePublishedFragmentOutputs(options.OutputDir, fileEntry.Path, parquetRelativePath, options.Parquet) files = files[:len(files)-1] return err } @@ -653,13 +654,14 @@ func dumpNodePhase(ctx context.Context, db graph.Database, targetGraph graph.Gra if _, err := scanDatabaseNodesFrom(ctx, db, targetGraph, entitySnapshot.NodeCount, options.BatchSize, options.ProgressInterval, afterID, hasAfterID, fileTotal(committedFiles), func(node *graph.Node) error { if fragmentWriter == nil { - nextWriter, nextRelativePath, err := openFragmentWriter(options.OutputDir, targetGraph.Name, PhaseNodes, shardNumber, options) + nextWriter, nextRelativePath, nextParquetRelativePath, err := openNodeFragmentWriter(options.OutputDir, targetGraph.Name, shardNumber, options) if err != nil { return err } fragmentWriter = nextWriter fragmentRelativePath = nextRelativePath + parquetRelativePath = nextParquetRelativePath } kinds := node.Kinds.Strings() @@ -720,8 +722,9 @@ func dumpEdgePhase(ctx context.Context, db graph.Database, targetGraph graph.Gra var ( files = append([]FileManifest(nil), committedFiles...) - fragmentWriter *compressedJSONLinesWriter + fragmentWriter *fragmentWriter[FragmentEdge] fragmentRelativePath string + parquetRelativePath string shardActionCounts scrubActionCounts shardNumber = len(committedFiles) + 1 lastWrittenID graph.ID @@ -741,13 +744,13 @@ func dumpEdgePhase(ctx context.Context, db graph.Database, targetGraph graph.Gra files = append(files, fileEntry) if !hasLastWrittenID { - _ = os.Remove(filepath.Join(options.OutputDir, filepath.FromSlash(fileEntry.Path))) + removePublishedFragmentOutputs(options.OutputDir, fileEntry.Path, parquetRelativePath, options.Parquet) files = files[:len(files)-1] return fmt.Errorf("edge fragment %q closed without a committed source cursor", fileEntry.Path) } if onCommit != nil { if err := onCommit(fileEntry, lastWrittenID); err != nil { - _ = os.Remove(filepath.Join(options.OutputDir, filepath.FromSlash(fileEntry.Path))) + removePublishedFragmentOutputs(options.OutputDir, fileEntry.Path, parquetRelativePath, options.Parquet) files = files[:len(files)-1] return err } @@ -760,13 +763,14 @@ func dumpEdgePhase(ctx context.Context, db graph.Database, targetGraph graph.Gra if _, err := scanDatabaseRelationshipsFrom(ctx, db, targetGraph, entitySnapshot.EdgeCount, options.BatchSize, options.ProgressInterval, afterID, hasAfterID, fileTotal(committedFiles), func(relationship *graph.Relationship) error { if fragmentWriter == nil { - nextWriter, nextRelativePath, err := openFragmentWriter(options.OutputDir, targetGraph.Name, PhaseEdges, shardNumber, options) + nextWriter, nextRelativePath, nextParquetRelativePath, err := openEdgeFragmentWriter(options.OutputDir, targetGraph.Name, shardNumber, options) if err != nil { return err } fragmentWriter = nextWriter fragmentRelativePath = nextRelativePath + parquetRelativePath = nextParquetRelativePath } kind := "" @@ -824,84 +828,117 @@ func dumpEdgePhase(ctx context.Context, db graph.Database, targetGraph graph.Gra } func writeNodeFragment(outputDir, graphName string, shardNumber int, options DumpOptions, items []FragmentNode, actionCounts map[string]int) (FileManifest, error) { - relativePath, err := fragmentPath(graphName, PhaseNodes, shardNumber, options.Compression) + writer, relativePath, _, err := openNodeFragmentWriter(outputDir, graphName, shardNumber, options) if err != nil { return FileManifest{}, err } - - absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) - fileEntry, err := writeCompressedJSONLines(absolutePath, options.Compression, options.ZstdLevel, items) - if err != nil { - return FileManifest{}, err + for _, item := range items { + if err := writer.Write(item); err != nil { + writer.Abort() + return FileManifest{}, err + } } - - fileEntry.Phase = PhaseNodes - fileEntry.Path = relativePath - fileEntry.Count = len(items) - fileEntry.ActionCounts = cloneActionCounts(actionCounts) - - return fileEntry, nil + return closeFragmentWriter(writer, relativePath, PhaseNodes, actionCounts) } func writeEdgeFragment(outputDir, graphName string, shardNumber int, options DumpOptions, items []FragmentEdge, actionCounts map[string]int) (FileManifest, error) { - relativePath, err := fragmentPath(graphName, PhaseEdges, shardNumber, options.Compression) + writer, relativePath, _, err := openEdgeFragmentWriter(outputDir, graphName, shardNumber, options) if err != nil { return FileManifest{}, err } + for _, item := range items { + if err := writer.Write(item); err != nil { + writer.Abort() + return FileManifest{}, err + } + } + return closeFragmentWriter(writer, relativePath, PhaseEdges, actionCounts) +} - absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) - fileEntry, err := writeCompressedJSONLines(absolutePath, options.Compression, options.ZstdLevel, items) +func fragmentPath(graphName string, fragmentPhase Phase, shardNumber int, codec CompressionCodec) (string, error) { + if shardNumber <= 0 { + return "", fmt.Errorf("shard number must be > 0") + } + extension, err := compressionExtension(codec) if err != nil { - return FileManifest{}, err + return "", err + } + prefix, err := fragmentPrefix(fragmentPhase) + if err != nil { + return "", err } - fileEntry.Phase = PhaseEdges - fileEntry.Path = relativePath - fileEntry.Count = len(items) - fileEntry.ActionCounts = cloneActionCounts(actionCounts) - - return fileEntry, nil + return path.Join("graphs", graphDirectoryName(graphName), fmt.Sprintf("%s-%06d.jsonl%s", prefix, shardNumber, extension)), nil } -func fragmentPath(graphName string, fragmentPhase Phase, shardNumber int, codec CompressionCodec) (string, error) { +func parquetFragmentPath(graphName string, fragmentPhase Phase, shardNumber int) (string, error) { if shardNumber <= 0 { return "", fmt.Errorf("shard number must be > 0") } - - extension, err := compressionExtension(codec) + prefix, err := fragmentPrefix(fragmentPhase) if err != nil { return "", err } - var prefix string + return path.Join("graphs", graphDirectoryName(graphName), fmt.Sprintf("%s-%06d.parquet", prefix, shardNumber)), nil +} + +func fragmentPrefix(fragmentPhase Phase) (string, error) { switch fragmentPhase { case PhaseNodes: - prefix = "nodes" + return "nodes", nil case PhaseEdges: - prefix = "edges" + return "edges", nil default: return "", fmt.Errorf("unsupported fragment phase %q", fragmentPhase) } - - return path.Join("graphs", graphDirectoryName(graphName), fmt.Sprintf("%s-%06d.jsonl%s", prefix, shardNumber, extension)), nil } -func openFragmentWriter(outputDir, graphName string, fragmentPhase Phase, shardNumber int, options DumpOptions) (*compressedJSONLinesWriter, string, error) { - relativePath, err := fragmentPath(graphName, fragmentPhase, shardNumber, options.Compression) +func openNodeFragmentWriter(outputDir, graphName string, shardNumber int, options DumpOptions) (*fragmentWriter[FragmentNode], string, string, error) { + relativePath, parquetRelativePath, err := fragmentPaths(graphName, PhaseNodes, shardNumber, options.Compression) + if err != nil { + return nil, "", "", err + } + writer, err := newNodeFragmentWriter( + filepath.Join(outputDir, filepath.FromSlash(relativePath)), + filepath.Join(outputDir, filepath.FromSlash(parquetRelativePath)), + options, + ) if err != nil { - return nil, "", err + return nil, "", "", err } + return writer, relativePath, parquetRelativePath, nil +} - absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) - writer, err := newCompressedJSONLinesWriter(absolutePath, options.Compression, options.ZstdLevel) +func openEdgeFragmentWriter(outputDir, graphName string, shardNumber int, options DumpOptions) (*fragmentWriter[FragmentEdge], string, string, error) { + relativePath, parquetRelativePath, err := fragmentPaths(graphName, PhaseEdges, shardNumber, options.Compression) + if err != nil { + return nil, "", "", err + } + writer, err := newEdgeFragmentWriter( + filepath.Join(outputDir, filepath.FromSlash(relativePath)), + filepath.Join(outputDir, filepath.FromSlash(parquetRelativePath)), + options, + ) if err != nil { - return nil, "", err + return nil, "", "", err } + return writer, relativePath, parquetRelativePath, nil +} - return writer, relativePath, nil +func fragmentPaths(graphName string, fragmentPhase Phase, shardNumber int, codec CompressionCodec) (string, string, error) { + relativePath, err := fragmentPath(graphName, fragmentPhase, shardNumber, codec) + if err != nil { + return "", "", err + } + parquetRelativePath, err := parquetFragmentPath(graphName, fragmentPhase, shardNumber) + if err != nil { + return "", "", err + } + return relativePath, parquetRelativePath, nil } -func closeFragmentWriter(writer *compressedJSONLinesWriter, relativePath string, fragmentPhase Phase, actionCounts map[string]int) (FileManifest, error) { +func closeFragmentWriter[T any](writer *fragmentWriter[T], relativePath string, fragmentPhase Phase, actionCounts map[string]int) (FileManifest, error) { fileEntry, err := writer.Close() if err != nil { return FileManifest{}, err @@ -914,6 +951,13 @@ func closeFragmentWriter(writer *compressedJSONLinesWriter, relativePath string, return fileEntry, nil } +func removePublishedFragmentOutputs(outputDir, relativePath, parquetRelativePath string, parquetEnabled bool) { + _ = os.Remove(filepath.Join(outputDir, filepath.FromSlash(relativePath))) + if parquetEnabled { + _ = os.Remove(filepath.Join(outputDir, filepath.FromSlash(parquetRelativePath))) + } +} + func fileTotal(files []FileManifest) int64 { var total int64 for _, fileEntry := range files { diff --git a/retriever/dump_checkpoint.go b/retriever/dump_checkpoint.go index ac59c40d..996d6c62 100644 --- a/retriever/dump_checkpoint.go +++ b/retriever/dump_checkpoint.go @@ -22,6 +22,7 @@ type dumpCheckpointIdentity struct { Graphs []string `json:"graphs"` Compression CompressionCodec `json:"compression"` CompressionLevel int `json:"compression_level"` + Parquet bool `json:"parquet"` Scrub ScrubMode `json:"scrub"` ScrubRulesVersion string `json:"scrub_rules_version,omitempty"` ScrubConfigSHA256 string `json:"scrub_config_sha256,omitempty"` @@ -54,6 +55,7 @@ func newDumpCheckpointIdentity(driverName string, targets []GraphTarget, options Graphs: make([]string, len(targets)), Compression: options.Compression, CompressionLevel: options.ZstdLevel, + Parquet: options.Parquet, Scrub: options.Scrub, ShardSize: options.ShardSize, BatchSize: options.BatchSize, @@ -142,7 +144,7 @@ func loadCompatibleDumpCheckpoint(outputDir string, expected dumpCheckpointIdent return dumpCheckpoint{}, err } if !reflect.DeepEqual(value.Identity, expected) { - return dumpCheckpoint{}, fmt.Errorf("dump checkpoint is incompatible with the requested driver, graphs, scrub, compression, shard, or batch options") + return dumpCheckpoint{}, fmt.Errorf("dump checkpoint is incompatible with the requested driver, graphs, scrub, compression, Parquet, shard, or batch options") } if err := validateDumpCheckpoint(value, targetCount); err != nil { return dumpCheckpoint{}, err @@ -271,11 +273,21 @@ func removeKnownDumpCheckpointTemps(outputDir string, value dumpCheckpoint) erro if err != nil { return err } - paths = append(paths, filepath.Join(outputDir, filepath.FromSlash(nextPath))+".tmp") + jsonlPath := filepath.Join(outputDir, filepath.FromSlash(nextPath)) + jsonlStagingPath := jsonlPath + ".tmp" + paths = append(paths, jsonlPath, jsonlStagingPath) + if value.Identity.Parquet { + nextParquetPath, err := parquetFragmentPath(value.Current.Name, value.Current.Phase, len(phaseFiles)+1) + if err != nil { + return err + } + parquetPath := filepath.Join(outputDir, filepath.FromSlash(nextParquetPath)) + paths = append(paths, jsonlStagingPath+".tmp", parquetPath, parquetPath+".tmp") + } } for _, candidate := range paths { if err := os.Remove(candidate); err != nil && !os.IsNotExist(err) { - return fmt.Errorf("remove stale dump temporary file %q: %w", candidate, err) + return fmt.Errorf("remove uncheckpointed dump file %q: %w", candidate, err) } } return nil @@ -284,25 +296,13 @@ func removeKnownDumpCheckpointTemps(outputDir string, value dumpCheckpoint) erro func validateDumpCheckpointFiles(outputDir string, value dumpCheckpoint) error { expected := map[string]struct{}{dumpCheckpointFileName: {}} for _, graphEntry := range value.Manifest.Graphs { - for _, fileEntry := range graphEntry.Files { - if _, found := expected[fileEntry.Path]; found { - return fmt.Errorf("dump checkpoint contains duplicate fragment path %q", fileEntry.Path) - } - expected[fileEntry.Path] = struct{}{} - if err := verifyDumpCheckpointFile(outputDir, fileEntry); err != nil { - return err - } + if err := addExpectedDumpCheckpointFiles(outputDir, expected, graphEntry.Name, graphEntry.Files, value.Identity.Parquet); err != nil { + return err } } if value.Current != nil { - for _, fileEntry := range value.Current.Files { - if _, found := expected[fileEntry.Path]; found { - return fmt.Errorf("dump checkpoint contains duplicate fragment path %q", fileEntry.Path) - } - expected[fileEntry.Path] = struct{}{} - if err := verifyDumpCheckpointFile(outputDir, fileEntry); err != nil { - return err - } + if err := addExpectedDumpCheckpointFiles(outputDir, expected, value.Current.Name, value.Current.Files, value.Identity.Parquet); err != nil { + return err } } @@ -325,6 +325,36 @@ func validateDumpCheckpointFiles(outputDir string, value dumpCheckpoint) error { }) } +func addExpectedDumpCheckpointFiles(outputDir string, expected map[string]struct{}, graphName string, files []FileManifest, parquetEnabled bool) error { + phaseShards := map[Phase]int{} + for _, fileEntry := range files { + if _, found := expected[fileEntry.Path]; found { + return fmt.Errorf("dump checkpoint contains duplicate fragment path %q", fileEntry.Path) + } + expected[fileEntry.Path] = struct{}{} + if err := verifyDumpCheckpointFile(outputDir, fileEntry); err != nil { + return err + } + + phaseShards[fileEntry.Phase]++ + if !parquetEnabled { + continue + } + parquetPath, err := parquetFragmentPath(graphName, fileEntry.Phase, phaseShards[fileEntry.Phase]) + if err != nil { + return err + } + if _, found := expected[parquetPath]; found { + return fmt.Errorf("dump checkpoint contains duplicate Parquet fragment path %q", parquetPath) + } + expected[parquetPath] = struct{}{} + if err := verifyDumpCheckpointParquetFile(outputDir, parquetPath); err != nil { + return err + } + } + return nil +} + func verifyDumpCheckpointFile(outputDir string, fileEntry FileManifest) error { absolutePath := filepath.Join(outputDir, filepath.FromSlash(fileEntry.Path)) info, err := os.Lstat(absolutePath) @@ -337,6 +367,18 @@ func verifyDumpCheckpointFile(outputDir string, fileEntry FileManifest) error { return verifyChecksum(absolutePath, fileEntry.SHA256, fileEntry.CompressedBytes) } +func verifyDumpCheckpointParquetFile(outputDir, relativePath string) error { + absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) + info, err := os.Lstat(absolutePath) + if err != nil { + return fmt.Errorf("inspect dump checkpoint Parquet fragment %q: %w", relativePath, err) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("dump checkpoint Parquet fragment %q is not a regular file", relativePath) + } + return nil +} + func filterPhaseFiles(files []FileManifest, phase Phase) []FileManifest { result := make([]FileManifest, 0, len(files)) for _, fileEntry := range files { diff --git a/retriever/dump_checkpoint_test.go b/retriever/dump_checkpoint_test.go index 43584205..2e424642 100644 --- a/retriever/dump_checkpoint_test.go +++ b/retriever/dump_checkpoint_test.go @@ -3,12 +3,14 @@ package retriever import ( "context" "errors" + "fmt" "os" "path/filepath" "reflect" "strings" "testing" + "github.com/parquet-go/parquet-go" cypherModel "github.com/specterops/dawgs/cypher/models/cypher" "github.com/specterops/dawgs/graph" ) @@ -154,6 +156,7 @@ func TestDumpResumesFromLastCommittedFragment(t *testing.T) { options.ShardSize = 2 options.Scrub = ScrubFull options.Salt = "resume-test-salt" + options.Parquet = true _, err := Dump(context.Background(), database, "test", []GraphTarget{{Name: "default"}}, options) if err == nil { @@ -170,6 +173,12 @@ func TestDumpResumesFromLastCommittedFragment(t *testing.T) { if checkpoint.Current == nil || checkpoint.Current.LastCommittedID != 2 || fileTotal(checkpoint.Current.Files) != 2 { t.Fatalf("unexpected interrupted checkpoint: %+v", checkpoint.Current) } + if !checkpoint.Identity.Parquet { + t.Fatal("interrupted checkpoint did not record Parquet enablement") + } + for _, fileEntry := range checkpoint.Current.Files { + assertCheckpointParquetPartnerReadable(t, outputDir, fileEntry) + } options.Resume = true database.nodes = append(database.nodes, graph.NewNode(6, graph.AsProperties(map[string]any{"name": "six"}), nodeKind)) @@ -192,6 +201,29 @@ func TestDumpResumesFromLastCommittedFragment(t *testing.T) { t.Fatalf("completed dump retained checkpoint: %v", err) } + parquetPartners := map[string]struct{}{} + for _, graphEntry := range result.Manifest.Graphs { + for _, fileEntry := range graphEntry.Files { + partnerPath := assertCheckpointParquetPartnerReadable(t, outputDir, fileEntry) + if _, found := parquetPartners[partnerPath]; found { + t.Fatalf("multiple JSONL manifest entries resolve to Parquet partner %q", partnerPath) + } + parquetPartners[partnerPath] = struct{}{} + } + } + actualParquetFiles := 0 + if err := filepath.WalkDir(outputDir, func(_ string, entry os.DirEntry, err error) error { + if err == nil && !entry.IsDir() && strings.HasSuffix(entry.Name(), ".parquet") { + actualParquetFiles++ + } + return err + }); err != nil { + t.Fatalf("walk completed dump: %v", err) + } + if actualParquetFiles != len(parquetPartners) { + t.Fatalf("completed dump has %d Parquet files for %d JSONL manifest entries", actualParquetFiles, len(parquetPartners)) + } + var nodeIDs []string for _, fileEntry := range result.Manifest.Graphs[0].Files { if fileEntry.Phase != PhaseNodes { @@ -235,6 +267,16 @@ func TestDumpResumeRejectsIncompatibleIdentityAndUnexpectedFiles(t *testing.T) { t.Fatal("expected incompatible checkpoint rejection") } + parquetOptions := options + parquetOptions.Parquet = true + parquetIdentity, err := newDumpCheckpointIdentity("test", []GraphTarget{{Name: "default"}}, parquetOptions, nil) + if err != nil { + t.Fatalf("Parquet checkpoint identity: %v", err) + } + if _, err := loadCompatibleDumpCheckpoint(outputDir, parquetIdentity, 1); err == nil { + t.Fatal("expected Parquet enablement mismatch rejection") + } + if err := os.WriteFile(filepath.Join(outputDir, "orphan.jsonl"), []byte("orphan"), 0o600); err != nil { t.Fatalf("write orphan: %v", err) } @@ -243,6 +285,296 @@ func TestDumpResumeRejectsIncompatibleIdentityAndUnexpectedFiles(t *testing.T) { } } +func TestDumpResumeRejectsMissingCommittedParquetPartner(t *testing.T) { + nodeKind := graph.StringKind("Node") + database := &checkpointTestDatabase{ + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{"name": "one"}), nodeKind), + graph.NewNode(2, graph.AsProperties(map[string]any{"name": "two"}), nodeKind), + }, + failNodeFetch: 2, + failNodeFetchOnce: true, + } + + outputDir := t.TempDir() + options := DefaultDumpOptions(outputDir) + options.Compression = CompressionNone + options.BatchSize = 1 + options.ShardSize = 1 + options.Parquet = true + + if _, err := Dump(context.Background(), database, "test", []GraphTarget{{Name: "default"}}, options); err == nil { + t.Fatal("expected injected dump failure") + } + checkpoint, err := readDumpCheckpoint(outputDir) + if err != nil { + t.Fatalf("read checkpoint: %v", err) + } + if checkpoint.Current == nil || len(checkpoint.Current.Files) != 1 { + t.Fatalf("unexpected interrupted checkpoint: %+v", checkpoint.Current) + } + partnerPath, err := parquetFragmentPath(checkpoint.Current.Name, checkpoint.Current.Files[0].Phase, 1) + if err != nil { + t.Fatalf("Parquet partner path: %v", err) + } + if err := os.Remove(filepath.Join(outputDir, filepath.FromSlash(partnerPath))); err != nil { + t.Fatalf("remove committed Parquet partner: %v", err) + } + + options.Resume = true + if _, err := Dump(context.Background(), database, "test", []GraphTarget{{Name: "default"}}, options); err == nil || !strings.Contains(err.Error(), "Parquet") { + t.Fatalf("expected missing committed Parquet partner rejection, got %v", err) + } +} + +func TestDumpResumeRemovesKnownNextShardTemps(t *testing.T) { + outputDir := t.TempDir() + options := DefaultDumpOptions(outputDir) + options.Compression = CompressionNone + options.Parquet = true + identity, err := newDumpCheckpointIdentity("test", []GraphTarget{{Name: "default"}}, options, nil) + if err != nil { + t.Fatalf("checkpoint identity: %v", err) + } + checkpoint := dumpCheckpoint{ + Identity: identity, + Current: &dumpGraphCheckpoint{ + Name: "default", + Phase: PhaseNodes, + }, + } + jsonlPath, err := fragmentPath("default", PhaseNodes, 1, CompressionNone) + if err != nil { + t.Fatalf("JSONL fragment path: %v", err) + } + parquetPath, err := parquetFragmentPath("default", PhaseNodes, 1) + if err != nil { + t.Fatalf("Parquet fragment path: %v", err) + } + tempPaths := []string{ + filepath.Join(outputDir, filepath.FromSlash(jsonlPath)) + ".tmp", + filepath.Join(outputDir, filepath.FromSlash(jsonlPath)) + ".tmp.tmp", + filepath.Join(outputDir, filepath.FromSlash(parquetPath)) + ".tmp", + } + for _, tempPath := range tempPaths { + if err := os.MkdirAll(filepath.Dir(tempPath), 0o755); err != nil { + t.Fatalf("create temp directory: %v", err) + } + if err := os.WriteFile(tempPath, []byte("stale"), 0o600); err != nil { + t.Fatalf("write stale temp %q: %v", tempPath, err) + } + } + + if err := removeKnownDumpCheckpointTemps(outputDir, checkpoint); err != nil { + t.Fatalf("remove known dump checkpoint temps: %v", err) + } + for _, tempPath := range tempPaths { + if _, err := os.Stat(tempPath); !os.IsNotExist(err) { + t.Fatalf("stale temp %q still exists: %v", tempPath, err) + } + } +} + +func TestDumpResumeRemovesUncheckpointedJSONLFinal(t *testing.T) { + outputDir := t.TempDir() + options := DefaultDumpOptions(outputDir) + options.Compression = CompressionNone + options.Parquet = true + identity, err := newDumpCheckpointIdentity("test", []GraphTarget{{Name: "default"}}, options, nil) + if err != nil { + t.Fatalf("checkpoint identity: %v", err) + } + manifest := newManifest("test", options.Compression, options.ZstdLevel, ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + }, 1) + metrics := newMetricsManifest(1) + manifest.Metrics = &metrics + checkpoint := dumpCheckpoint{ + Identity: identity, + Manifest: manifest, + Current: &dumpGraphCheckpoint{ + Index: 0, + Name: "default", + Snapshot: graphEntitySnapshot{NodeCount: 1}, + HasSnapshot: true, + Phase: PhaseNodes, + }, + } + if err := writeDumpCheckpoint(outputDir, checkpoint); err != nil { + t.Fatalf("write checkpoint: %v", err) + } + + jsonlPath, err := fragmentPath("default", PhaseNodes, 1, CompressionNone) + if err != nil { + t.Fatalf("JSONL fragment path: %v", err) + } + jsonlFinal := filepath.Join(outputDir, filepath.FromSlash(jsonlPath)) + if err := os.MkdirAll(filepath.Dir(jsonlFinal), 0o755); err != nil { + t.Fatalf("create fragment directory: %v", err) + } + if err := os.WriteFile(jsonlFinal, []byte("published before Parquet\n"), 0o600); err != nil { + t.Fatalf("write uncheckpointed JSONL final: %v", err) + } + + if _, err := loadCompatibleDumpCheckpoint(outputDir, identity, 1); err != nil { + t.Fatalf("load checkpoint with partially published shard: %v", err) + } + assertAbsent(t, jsonlFinal) +} + +func TestDumpResumeRemovesUncheckpointedPairWithoutRemovingCommittedPair(t *testing.T) { + outputDir := t.TempDir() + options := DefaultDumpOptions(outputDir) + options.Compression = CompressionNone + options.Parquet = true + identity, err := newDumpCheckpointIdentity("test", []GraphTarget{{Name: "default"}}, options, nil) + if err != nil { + t.Fatalf("checkpoint identity: %v", err) + } + + committedFile, err := writeNodeFragment(outputDir, "default", 1, options, []FragmentNode{{ + ID: "1", + Kinds: []string{"Person"}, + Properties: map[string]any{"name": "committed"}, + }}, nil) + if err != nil { + t.Fatalf("write committed fragment pair: %v", err) + } + committedParquetPath, err := parquetFragmentPath("default", PhaseNodes, 1) + if err != nil { + t.Fatalf("committed Parquet path: %v", err) + } + + manifest := newManifest("test", options.Compression, options.ZstdLevel, ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + }, 1) + metrics := newMetricsManifest(1) + manifest.Metrics = &metrics + checkpoint := dumpCheckpoint{ + Identity: identity, + Manifest: manifest, + Current: &dumpGraphCheckpoint{ + Index: 0, + Name: "default", + Snapshot: graphEntitySnapshot{NodeCount: 2}, + HasSnapshot: true, + Phase: PhaseNodes, + LastCommittedID: 1, + HasLastCommittedID: true, + Files: []FileManifest{committedFile}, + }, + } + if err := writeDumpCheckpoint(outputDir, checkpoint); err != nil { + t.Fatalf("write checkpoint: %v", err) + } + + uncheckpointedJSONLPath, err := fragmentPath("default", PhaseNodes, 2, CompressionNone) + if err != nil { + t.Fatalf("uncheckpointed JSONL path: %v", err) + } + uncheckpointedParquetPath, err := parquetFragmentPath("default", PhaseNodes, 2) + if err != nil { + t.Fatalf("uncheckpointed Parquet path: %v", err) + } + for _, relativePath := range []string{uncheckpointedJSONLPath, uncheckpointedParquetPath} { + absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) + if err := os.WriteFile(absolutePath, []byte("published before checkpoint"), 0o600); err != nil { + t.Fatalf("write uncheckpointed final %q: %v", relativePath, err) + } + } + + if _, err := loadCompatibleDumpCheckpoint(outputDir, identity, 1); err != nil { + t.Fatalf("load checkpoint with uncheckpointed pair: %v", err) + } + for _, relativePath := range []string{uncheckpointedJSONLPath, uncheckpointedParquetPath} { + assertAbsent(t, filepath.Join(outputDir, filepath.FromSlash(relativePath))) + } + for _, relativePath := range []string{committedFile.Path, committedParquetPath} { + if _, err := os.Stat(filepath.Join(outputDir, filepath.FromSlash(relativePath))); err != nil { + t.Fatalf("committed pair member %q was removed: %v", relativePath, err) + } + } +} + +func TestDumpResumeCommitCallbackFailureRemovesParquetPair(t *testing.T) { + nodeKind := graph.StringKind("Node") + database := &checkpointTestDatabase{ + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{"name": "one"}), nodeKind), + }, + } + outputDir := t.TempDir() + options := DefaultDumpOptions(outputDir) + options.Compression = CompressionNone + options.BatchSize = 1 + options.ShardSize = 1 + options.Parquet = true + jsonlPath, err := fragmentPath("default", PhaseNodes, 1, CompressionNone) + if err != nil { + t.Fatalf("JSONL fragment path: %v", err) + } + parquetPath, err := parquetFragmentPath("default", PhaseNodes, 1) + if err != nil { + t.Fatalf("Parquet fragment path: %v", err) + } + commitErr := errors.New("injected checkpoint commit failure") + observedPair := false + + _, err = dumpNodePhase(context.Background(), database, graph.Graph{Name: "default"}, options, nil, map[string]struct{}{}, &scrubActionCounts{}, graphEntitySnapshot{NodeCount: 1}, newMetricsBuilder("default", 1), nil, 0, false, func(FileManifest, graph.ID) error { + for _, relativePath := range []string{jsonlPath, parquetPath} { + if _, err := os.Stat(filepath.Join(outputDir, filepath.FromSlash(relativePath))); err != nil { + return fmt.Errorf("inspect published fragment pair member %q: %w", relativePath, err) + } + } + observedPair = true + return commitErr + }) + if !errors.Is(err, commitErr) { + t.Fatalf("dump node phase error = %v, want %v", err, commitErr) + } + if !observedPair { + t.Fatal("commit callback did not observe the published JSONL/Parquet pair") + } + for _, relativePath := range []string{jsonlPath, parquetPath} { + if _, err := os.Stat(filepath.Join(outputDir, filepath.FromSlash(relativePath))); !os.IsNotExist(err) { + t.Fatalf("callback failure retained %q: %v", relativePath, err) + } + } +} + +func assertCheckpointParquetPartnerReadable(t *testing.T, outputDir string, fileEntry FileManifest) string { + t.Helper() + + jsonlPath := strings.TrimSuffix(strings.TrimSuffix(fileEntry.Path, ".gz"), ".zst") + partnerPath := strings.TrimSuffix(jsonlPath, ".jsonl") + ".parquet" + absolutePath := filepath.Join(outputDir, filepath.FromSlash(partnerPath)) + switch fileEntry.Phase { + case PhaseNodes: + rows, err := parquet.ReadFile[parquetNodeRow](absolutePath) + if err != nil { + t.Fatalf("read node Parquet partner %q: %v", partnerPath, err) + } + if len(rows) != fileEntry.Count { + t.Fatalf("node Parquet partner %q has %d rows, want %d", partnerPath, len(rows), fileEntry.Count) + } + case PhaseEdges: + rows, err := parquet.ReadFile[parquetEdgeRow](absolutePath) + if err != nil { + t.Fatalf("read edge Parquet partner %q: %v", partnerPath, err) + } + if len(rows) != fileEntry.Count { + t.Fatalf("edge Parquet partner %q has %d rows, want %d", partnerPath, len(rows), fileEntry.Count) + } + default: + t.Fatalf("unsupported checkpoint fragment phase %q", fileEntry.Phase) + } + return partnerPath +} + func TestDumpRejectsInvalidGraphTargetsBeforeWritingOutput(t *testing.T) { for name, targets := range map[string][]GraphTarget{ "empty": {{Name: ""}}, diff --git a/retriever/dump_test.go b/retriever/dump_test.go index 011cae77..32913c37 100644 --- a/retriever/dump_test.go +++ b/retriever/dump_test.go @@ -2,7 +2,10 @@ package retriever import ( "path/filepath" + "reflect" "testing" + + "github.com/parquet-go/parquet-go" ) func TestFragmentPath(t *testing.T) { @@ -22,18 +25,41 @@ func TestFragmentPath(t *testing.T) { t.Fatalf("unexpected edge fragment path %q", edgePath) } + nodeParquetPath, err := parquetFragmentPath("graph/name", PhaseNodes, 7) + if err != nil { + t.Fatalf("node Parquet fragment path: %v", err) + } + if nodeParquetPath != "graphs/graph%2Fname/nodes-000007.parquet" { + t.Fatalf("unexpected node Parquet fragment path %q", nodeParquetPath) + } + + edgeParquetPath, err := parquetFragmentPath("default", PhaseEdges, 3) + if err != nil { + t.Fatalf("edge Parquet fragment path: %v", err) + } + if edgeParquetPath != "graphs/default/edges-000003.parquet" { + t.Fatalf("unexpected edge Parquet fragment path %q", edgeParquetPath) + } + if _, err := fragmentPath("default", Phase("bad"), 1, CompressionGzip); err == nil { t.Fatalf("expected unsupported Phase error") } if _, err := fragmentPath("default", PhaseNodes, 0, CompressionGzip); err == nil { t.Fatalf("expected invalid shard number error") } + if _, err := parquetFragmentPath("default", Phase("bad"), 1); err == nil { + t.Fatalf("expected unsupported Parquet Phase error") + } + if _, err := parquetFragmentPath("default", PhaseNodes, 0); err == nil { + t.Fatalf("expected invalid Parquet shard number error") + } } func TestWriteFragmentMetadata(t *testing.T) { options := DumpOptions{ OutputDir: t.TempDir(), Compression: CompressionGzip, + Parquet: true, ZstdLevel: DefaultZstdLevel, } @@ -54,11 +80,19 @@ func TestWriteFragmentMetadata(t *testing.T) { if _, err := readManifest(filepath.Join(options.OutputDir, "graphs")); err == nil { t.Fatalf("fragment write should not create Manifest") } + nodeRows, err := parquet.ReadFile[parquetNodeRow](filepath.Join(options.OutputDir, "graphs/default/nodes-000001.parquet")) + if err != nil { + t.Fatalf("read node Parquet sidecar: %v", err) + } + if len(nodeRows) != 1 || nodeRows[0].ID != "1" || !reflect.DeepEqual(nodeRows[0].Kinds, []string{"User"}) || !reflect.DeepEqual(nodeRows[0].Properties, map[string]any{"name": "alice"}) { + t.Fatalf("unexpected node Parquet rows: %#v", nodeRows) + } edgeEntry, err := writeEdgeFragment(options.OutputDir, "default", 2, options, []FragmentEdge{{ - StartID: "1", - EndID: "2", - Kind: "AdminTo", + StartID: "1", + EndID: "2", + Kind: "AdminTo", + Properties: map[string]any{"active": true}, }}, nil) if err != nil { t.Fatalf("write edge fragment: %v", err) @@ -66,6 +100,13 @@ func TestWriteFragmentMetadata(t *testing.T) { if edgeEntry.Phase != PhaseEdges || edgeEntry.Path != "graphs/default/edges-000002.jsonl.gz" || edgeEntry.Count != 1 { t.Fatalf("unexpected edge file Manifest: %+v", edgeEntry) } + edgeRows, err := parquet.ReadFile[parquetEdgeRow](filepath.Join(options.OutputDir, "graphs/default/edges-000002.parquet")) + if err != nil { + t.Fatalf("read edge Parquet sidecar: %v", err) + } + if !reflect.DeepEqual(edgeRows, []parquetEdgeRow{{StartID: "1", EndID: "2", Kind: "AdminTo", Properties: map[string]any{"active": true}}}) { + t.Fatalf("unexpected edge Parquet rows: %#v", edgeRows) + } } func TestKindAndActionHelpers(t *testing.T) { diff --git a/retriever/fragment_writer.go b/retriever/fragment_writer.go new file mode 100644 index 00000000..6c9c1286 --- /dev/null +++ b/retriever/fragment_writer.go @@ -0,0 +1,150 @@ +package retriever + +import ( + "fmt" + "os" +) + +var newNodeParquetSinkForFragmentWriter = newNodeParquetSink + +type fragmentWriter[T any] struct { + jsonl *compressedJSONLinesWriter + + parquet *parquetFragmentSink[T] + + jsonlPath string + jsonlStagingPath string + parquetPath string + parquetStagingPath string + + count int + closed bool +} + +func newNodeFragmentWriter(jsonlPath, parquetPath string, options DumpOptions) (*fragmentWriter[FragmentNode], error) { + return newFragmentWriter(jsonlPath, parquetPath, options, newNodeParquetSinkForFragmentWriter) +} + +func newEdgeFragmentWriter(jsonlPath, parquetPath string, options DumpOptions) (*fragmentWriter[FragmentEdge], error) { + return newFragmentWriter(jsonlPath, parquetPath, options, newEdgeParquetSink) +} + +func newFragmentWriter[T any](jsonlPath, parquetPath string, options DumpOptions, newParquetSink func(string) (*parquetFragmentSink[T], error)) (*fragmentWriter[T], error) { + if !options.Parquet { + jsonl, err := newCompressedJSONLinesWriter(jsonlPath, options.Compression, options.ZstdLevel) + if err != nil { + return nil, err + } + + return &fragmentWriter[T]{ + jsonl: jsonl, + }, nil + } + + jsonlStagingPath := jsonlPath + ".tmp" + jsonl, err := newCompressedJSONLinesWriterAtPaths(jsonlStagingPath, jsonlStagingPath+".tmp", options.Compression, options.ZstdLevel) + if err != nil { + return nil, err + } + + parquetStagingPath := parquetPath + ".tmp" + parquet, err := newParquetSink(parquetStagingPath) + if err != nil { + jsonl.Abort() + _ = os.Remove(jsonlStagingPath) + return nil, err + } + + return &fragmentWriter[T]{ + jsonl: jsonl, + parquet: parquet, + jsonlPath: jsonlPath, + jsonlStagingPath: jsonlStagingPath, + parquetPath: parquetPath, + parquetStagingPath: parquetStagingPath, + }, nil +} + +func (s *fragmentWriter[T]) Write(fragment T) error { + if s.closed { + return fmt.Errorf("write closed fragment") + } + + if err := s.jsonl.Write(fragment); err != nil { + s.Abort() + return err + } + if s.parquet != nil { + if err := s.parquet.Write(fragment); err != nil { + s.Abort() + return err + } + } + + s.count++ + return nil +} + +func (s *fragmentWriter[T]) Count() int { + return s.count +} + +func (s *fragmentWriter[T]) Close() (FileManifest, error) { + if s.closed { + return FileManifest{}, fmt.Errorf("close fragment more than once") + } + s.closed = true + + if s.parquet == nil { + return s.jsonl.Close() + } + + fileEntry, err := s.jsonl.Close() + if err != nil { + s.cleanup(false, false) + return FileManifest{}, err + } + if fileEntry.Count != s.count { + s.cleanup(false, false) + return FileManifest{}, fmt.Errorf("JSONL fragment count %d does not match logical count %d", fileEntry.Count, s.count) + } + if err := s.parquet.Close(); err != nil { + s.cleanup(false, false) + return FileManifest{}, fmt.Errorf("close Parquet fragment: %w", err) + } + if err := os.Rename(s.jsonlStagingPath, s.jsonlPath); err != nil { + s.cleanup(false, false) + return FileManifest{}, fmt.Errorf("publish JSONL fragment: %w", err) + } + if err := os.Rename(s.parquetStagingPath, s.parquetPath); err != nil { + s.cleanup(true, false) + return FileManifest{}, fmt.Errorf("publish Parquet fragment: %w", err) + } + + return fileEntry, nil +} + +func (s *fragmentWriter[T]) Abort() { + if s.closed { + return + } + s.closed = true + s.cleanup(false, false) +} + +func (s *fragmentWriter[T]) cleanup(jsonlPublished, parquetPublished bool) { + s.jsonl.Abort() + if s.parquet != nil { + s.parquet.Abort() + } + + _ = os.Remove(s.jsonlStagingPath) + _ = os.Remove(s.jsonlStagingPath + ".tmp") + _ = os.Remove(s.parquetStagingPath) + if jsonlPublished { + _ = os.Remove(s.jsonlPath) + } + if parquetPublished { + _ = os.Remove(s.parquetPath) + } +} diff --git a/retriever/fragment_writer_test.go b/retriever/fragment_writer_test.go new file mode 100644 index 00000000..4de476a6 --- /dev/null +++ b/retriever/fragment_writer_test.go @@ -0,0 +1,365 @@ +package retriever + +import ( + "encoding/json" + "errors" + "io" + "math" + "os" + "path/filepath" + "reflect" + "testing" + + "github.com/parquet-go/parquet-go" +) + +func TestNodeFragmentWriter(t *testing.T) { + t.Run("writes JSONL and Parquet when enabled", func(t *testing.T) { + jsonlPath := filepath.Join(t.TempDir(), "nodes.jsonl") + parquetPath := filepath.Join(t.TempDir(), "nodes.parquet") + options := fragmentWriterTestOptions(t, true) + first := FragmentNode{ID: "node-1", Kinds: []string{"Person", "Employee"}, Properties: representativeParquetProperties()} + second := FragmentNode{ID: "node-2", Kinds: []string{"Device"}, Properties: map[string]any{"name": "Laptop"}} + + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, options) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(first); err != nil { + t.Fatalf("write first node: %v", err) + } + if writer.Count() != 1 { + t.Fatalf("count after first node = %d, want 1", writer.Count()) + } + if err := writer.Write(second); err != nil { + t.Fatalf("write second node: %v", err) + } + if writer.Count() != 2 { + t.Fatalf("count after second node = %d, want 2", writer.Count()) + } + entry, err := writer.Close() + if err != nil { + t.Fatalf("close node writer: %v", err) + } + if entry.Count != 2 { + t.Fatalf("manifest count = %d, want 2", entry.Count) + } + + var jsonlRows []FragmentNode + readJSONL(t, jsonlPath, &jsonlRows) + if !reflect.DeepEqual(jsonlRows, []FragmentNode{first, second}) { + t.Fatalf("JSONL rows = %#v, want %#v", jsonlRows, []FragmentNode{first, second}) + } + + parquetRows, err := parquet.ReadFile[parquetNodeRow](parquetPath) + if err != nil { + t.Fatalf("read node Parquet file: %v", err) + } + if len(parquetRows) != 2 { + t.Fatalf("read %d node Parquet rows, want 2", len(parquetRows)) + } + if parquetRows[0].ID != first.ID || !reflect.DeepEqual(parquetRows[0].Kinds, first.Kinds) || !reflect.DeepEqual(parquetRows[0].Properties, first.Properties) { + t.Fatalf("first Parquet row = %#v, want node %#v", parquetRows[0], first) + } + if parquetRows[1].ID != second.ID || !reflect.DeepEqual(parquetRows[1].Kinds, second.Kinds) || !reflect.DeepEqual(parquetRows[1].Properties, second.Properties) { + t.Fatalf("second Parquet row = %#v, want node %#v", parquetRows[1], second) + } + }) + + t.Run("writes only JSONL when disabled", func(t *testing.T) { + jsonlPath := filepath.Join(t.TempDir(), "nodes.jsonl") + parquetPath := filepath.Join(t.TempDir(), "nodes.parquet") + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, false)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "node-1"}); err != nil { + t.Fatalf("write node: %v", err) + } + if _, err := writer.Close(); err != nil { + t.Fatalf("close node writer: %v", err) + } + + var rows []FragmentNode + readJSONL(t, jsonlPath, &rows) + if !reflect.DeepEqual(rows, []FragmentNode{{ID: "node-1"}}) { + t.Fatalf("JSONL rows = %#v, want node-1", rows) + } + assertAbsent(t, parquetPath) + assertAbsent(t, parquetPath+".tmp") + }) + + t.Run("abort removes both outputs", func(t *testing.T) { + jsonlPath := filepath.Join(t.TempDir(), "nodes.jsonl") + parquetPath := filepath.Join(t.TempDir(), "nodes.parquet") + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "node-1"}); err != nil { + t.Fatalf("write node: %v", err) + } + writer.Abort() + + assertAbsent(t, jsonlPath) + assertAbsent(t, jsonlPath+".tmp") + assertAbsent(t, parquetPath) + assertAbsent(t, parquetPath+".tmp") + }) + + t.Run("Parquet write failure removes staging and final outputs", func(t *testing.T) { + original := newNodeParquetSinkForFragmentWriter + newNodeParquetSinkForFragmentWriter = func(path string) (*parquetFragmentSink[FragmentNode], error) { + sink, err := newNodeParquetSink(path) + if err != nil { + return nil, err + } + sink.write = func(FragmentNode) error { return errors.New("injected parquet write failure") } + return sink, nil + } + t.Cleanup(func() { newNodeParquetSinkForFragmentWriter = original }) + + jsonlPath := filepath.Join(t.TempDir(), "nodes.jsonl") + parquetPath := filepath.Join(t.TempDir(), "nodes.parquet") + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "node-1"}); err == nil { + t.Fatal("expected Parquet write failure") + } + if writer.Count() != 0 { + t.Fatalf("count after failed node = %d, want 0", writer.Count()) + } + + assertAbsent(t, jsonlPath) + assertAbsent(t, jsonlPath+".tmp") + assertAbsent(t, parquetPath) + assertAbsent(t, parquetPath+".tmp") + }) + + t.Run("unsupported Parquet variant returns an error and removes every output", func(t *testing.T) { + outputDir := t.TempDir() + jsonlPath := filepath.Join(outputDir, "nodes.jsonl") + parquetPath := filepath.Join(outputDir, "nodes.parquet") + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + + err = writer.Write(FragmentNode{ + ID: "node-1", + Properties: map[string]any{"overflow": uint64(math.MaxUint64)}, + }) + if err == nil { + t.Fatal("expected unsupported Parquet VARIANT value to return an error") + } + if writer.Count() != 0 { + t.Fatalf("count after unsupported node = %d, want 0", writer.Count()) + } + + for _, path := range []string{ + jsonlPath, + jsonlPath + ".tmp", + jsonlPath + ".tmp.tmp", + parquetPath, + parquetPath + ".tmp", + } { + assertAbsent(t, path) + } + }) + + t.Run("abort panic does not escape or block output cleanup", func(t *testing.T) { + original := newNodeParquetSinkForFragmentWriter + writeErr := errors.New("injected parquet write failure") + newNodeParquetSinkForFragmentWriter = func(path string) (*parquetFragmentSink[FragmentNode], error) { + sink, err := newNodeParquetSink(path) + if err != nil { + return nil, err + } + originalAbort := sink.abort + sink.write = func(FragmentNode) error { return writeErr } + sink.abort = func() { + originalAbort() + panic("injected Parquet abort panic") + } + return sink, nil + } + t.Cleanup(func() { newNodeParquetSinkForFragmentWriter = original }) + + outputDir := t.TempDir() + jsonlPath := filepath.Join(outputDir, "nodes.jsonl") + parquetPath := filepath.Join(outputDir, "nodes.parquet") + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "node-1"}); !errors.Is(err, writeErr) { + t.Fatalf("write error = %v, want %v", err, writeErr) + } + + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath, parquetPath + ".tmp"} { + assertAbsent(t, path) + } + }) + + t.Run("JSONL close failure removes Parquet staging", func(t *testing.T) { + outputDir := t.TempDir() + jsonlPath := filepath.Join(outputDir, "nodes.jsonl") + parquetPath := filepath.Join(outputDir, "nodes.parquet") + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "node-1"}); err != nil { + t.Fatalf("write node: %v", err) + } + if err := writer.jsonl.file.Close(); err != nil { + t.Fatalf("inject JSONL close failure: %v", err) + } + if _, err := writer.Close(); err == nil { + t.Fatal("expected JSONL close failure") + } + + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath, parquetPath + ".tmp"} { + assertAbsent(t, path) + } + }) + + t.Run("Parquet close failure removes JSONL staging", func(t *testing.T) { + outputDir := t.TempDir() + jsonlPath := filepath.Join(outputDir, "nodes.jsonl") + parquetPath := filepath.Join(outputDir, "nodes.parquet") + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "node-1"}); err != nil { + t.Fatalf("write node: %v", err) + } + originalClose := writer.parquet.close + writer.parquet.close = func() error { + return errors.Join(originalClose(), errors.New("injected Parquet close failure")) + } + if _, err := writer.Close(); err == nil { + t.Fatal("expected Parquet close failure") + } + + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath, parquetPath + ".tmp"} { + assertAbsent(t, path) + } + }) + + t.Run("second rename failure removes published JSONL and Parquet staging", func(t *testing.T) { + outputDir := t.TempDir() + jsonlPath := filepath.Join(outputDir, "nodes.jsonl") + parquetPath := filepath.Join(outputDir, "nodes.parquet") + if err := os.Mkdir(parquetPath, 0o700); err != nil { + t.Fatalf("create Parquet rename blocker: %v", err) + } + writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) + if err != nil { + t.Fatalf("create node writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "node-1"}); err != nil { + t.Fatalf("write node: %v", err) + } + if _, err := writer.Close(); err == nil { + t.Fatal("expected Parquet publish rename failure") + } + + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath + ".tmp"} { + assertAbsent(t, path) + } + if info, err := os.Stat(parquetPath); err != nil || !info.IsDir() { + t.Fatalf("Parquet rename blocker was changed: info=%v err=%v", info, err) + } + }) +} + +func TestEdgeFragmentWriter(t *testing.T) { + jsonlPath := filepath.Join(t.TempDir(), "edges.jsonl") + parquetPath := filepath.Join(t.TempDir(), "edges.parquet") + options := fragmentWriterTestOptions(t, true) + first := FragmentEdge{StartID: "node-1", EndID: "node-2", Kind: "MemberOf", Properties: representativeParquetProperties()} + second := FragmentEdge{StartID: "node-2", EndID: "node-3", Kind: "AdminTo", Properties: map[string]any{"enabled": false}} + + writer, err := newEdgeFragmentWriter(jsonlPath, parquetPath, options) + if err != nil { + t.Fatalf("create edge writer: %v", err) + } + if err := writer.Write(first); err != nil { + t.Fatalf("write first edge: %v", err) + } + if err := writer.Write(second); err != nil { + t.Fatalf("write second edge: %v", err) + } + if writer.Count() != 2 { + t.Fatalf("count after edges = %d, want 2", writer.Count()) + } + entry, err := writer.Close() + if err != nil { + t.Fatalf("close edge writer: %v", err) + } + if entry.Count != 2 { + t.Fatalf("manifest count = %d, want 2", entry.Count) + } + + var jsonlRows []FragmentEdge + readJSONL(t, jsonlPath, &jsonlRows) + if !reflect.DeepEqual(jsonlRows, []FragmentEdge{first, second}) { + t.Fatalf("JSONL rows = %#v, want %#v", jsonlRows, []FragmentEdge{first, second}) + } + + parquetRows, err := parquet.ReadFile[parquetEdgeRow](parquetPath) + if err != nil { + t.Fatalf("read edge Parquet file: %v", err) + } + if len(parquetRows) != 2 { + t.Fatalf("read %d edge Parquet rows, want 2", len(parquetRows)) + } + if parquetRows[0].StartID != first.StartID || parquetRows[0].EndID != first.EndID || parquetRows[0].Kind != first.Kind || !reflect.DeepEqual(parquetRows[0].Properties, first.Properties) { + t.Fatalf("first Parquet row = %#v, want edge %#v", parquetRows[0], first) + } + if parquetRows[1].StartID != second.StartID || parquetRows[1].EndID != second.EndID || parquetRows[1].Kind != second.Kind || !reflect.DeepEqual(parquetRows[1].Properties, second.Properties) { + t.Fatalf("second Parquet row = %#v, want edge %#v", parquetRows[1], second) + } +} + +func fragmentWriterTestOptions(t *testing.T, parquetEnabled bool) DumpOptions { + t.Helper() + options := DefaultDumpOptions(t.TempDir()) + options.Compression = CompressionNone + options.Parquet = parquetEnabled + return options +} + +func readJSONL[T any](t *testing.T, path string, records *[]T) { + t.Helper() + + file, err := os.Open(path) + if err != nil { + t.Fatalf("open JSONL file: %v", err) + } + defer file.Close() + + decoder := json.NewDecoder(file) + for { + var record T + if err := decoder.Decode(&record); err != nil { + if errors.Is(err, io.EOF) { + return + } + t.Fatalf("decode JSONL record: %v", err) + } + *records = append(*records, record) + } +} + +func assertAbsent(t *testing.T, path string) { + t.Helper() + + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("expected %q to be absent, got %v", path, err) + } +} diff --git a/retriever/options.go b/retriever/options.go index 5e6af208..779d50ae 100644 --- a/retriever/options.go +++ b/retriever/options.go @@ -78,6 +78,7 @@ type DumpOptions struct { Salt string ScrubConfig io.Reader Compression CompressionCodec + Parquet bool ZstdLevel int ShardSize int BatchSize int diff --git a/retriever/options_test.go b/retriever/options_test.go index 7e0411e6..ae5fa978 100644 --- a/retriever/options_test.go +++ b/retriever/options_test.go @@ -132,6 +132,9 @@ func TestOptionsValidate(t *testing.T) { func TestDefaultOptions(t *testing.T) { dump := DefaultDumpOptions(t.TempDir()) + if dump.Parquet { + t.Fatal("expected Parquet output to be disabled by default") + } if dump.Scrub != ScrubNone || dump.Compression != CompressionZstd || dump.ZstdLevel != DefaultZstdLevel { t.Fatalf("unexpected dump defaults: %+v", dump) } @@ -142,6 +145,10 @@ func TestDefaultOptions(t *testing.T) { if err := dump.Validate(); err != nil { t.Fatalf("validate default dump options: %v", err) } + dump.Parquet = true + if err := dump.Validate(); err != nil { + t.Fatalf("validate Parquet dump options: %v", err) + } load := DefaultLoadOptions(t.TempDir()) if load.BatchSize != DefaultBatchSize || load.ProgressInterval != DefaultProgressInterval { diff --git a/retriever/parquet.go b/retriever/parquet.go new file mode 100644 index 00000000..2632a8b7 --- /dev/null +++ b/retriever/parquet.go @@ -0,0 +1,111 @@ +package retriever + +import ( + "errors" + "fmt" + "os" + + "github.com/parquet-go/parquet-go" +) + +type parquetNodeRow struct { + ID string `parquet:"id"` + Kinds []string `parquet:"kinds,list"` + Properties any `parquet:"properties,variant"` +} + +type parquetEdgeRow struct { + StartID string `parquet:"start_id"` + EndID string `parquet:"end_id"` + Kind string `parquet:"kind"` + Properties any `parquet:"properties,variant"` +} + +type parquetFragmentSink[T any] struct { + path string + write func(T) error + close func() error + abort func() + closed bool +} + +func newNodeParquetSink(path string) (*parquetFragmentSink[FragmentNode], error) { + return newParquetFragmentSink(path, func(fragment FragmentNode) parquetNodeRow { + return parquetNodeRow{ + ID: fragment.ID, + Kinds: fragment.Kinds, + Properties: fragment.Properties, + } + }) +} + +func newEdgeParquetSink(path string) (*parquetFragmentSink[FragmentEdge], error) { + return newParquetFragmentSink(path, func(fragment FragmentEdge) parquetEdgeRow { + return parquetEdgeRow{ + StartID: fragment.StartID, + EndID: fragment.EndID, + Kind: fragment.Kind, + Properties: fragment.Properties, + } + }) +} + +func newParquetFragmentSink[T, R any](path string, adapt func(T) R) (*parquetFragmentSink[T], error) { + file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return nil, fmt.Errorf("open parquet fragment: %w", err) + } + + writer := parquet.NewGenericWriter[R](file) + return &parquetFragmentSink[T]{ + path: path, + write: func(fragment T) error { + written, err := writer.Write([]R{adapt(fragment)}) + if err != nil { + return fmt.Errorf("write parquet row: %w", err) + } + if written != 1 { + return fmt.Errorf("write parquet row: wrote %d rows, want 1", written) + } + return nil + }, + close: func() error { + return errors.Join(writer.Close(), file.Close()) + }, + abort: func() { + defer func() { _ = os.Remove(path) }() + defer func() { _ = file.Close() }() + defer func() { _ = recover() }() + _ = writer.Close() + }, + }, nil +} + +func (s *parquetFragmentSink[T]) Write(fragment T) (err error) { + if s.closed { + return fmt.Errorf("write closed parquet fragment") + } + defer func() { + if value := recover(); value != nil { + err = fmt.Errorf("write parquet fragment: %v", value) + } + }() + return s.write(fragment) +} + +func (s *parquetFragmentSink[T]) Close() error { + if s.closed { + return fmt.Errorf("close parquet fragment more than once") + } + s.closed = true + return s.close() +} + +func (s *parquetFragmentSink[T]) Abort() { + if s.closed { + return + } + s.closed = true + defer func() { _ = recover() }() + s.abort() +} diff --git a/retriever/parquet_test.go b/retriever/parquet_test.go new file mode 100644 index 00000000..f5cb3532 --- /dev/null +++ b/retriever/parquet_test.go @@ -0,0 +1,210 @@ +package retriever + +import ( + "os" + "path/filepath" + "reflect" + "testing" + + "github.com/parquet-go/parquet-go" + "github.com/parquet-go/parquet-go/format" +) + +func TestNodeParquetSinkRoundTripsVariantProperties(t *testing.T) { + path := filepath.Join(t.TempDir(), "nodes.parquet") + properties := representativeParquetProperties() + + sink, err := newNodeParquetSink(path) + if err != nil { + t.Fatalf("create node parquet sink: %v", err) + } + if err := sink.Write(FragmentNode{ID: "node-1", Kinds: []string{"Person", "Employee"}, Properties: properties}); err != nil { + t.Fatalf("write first node: %v", err) + } + if err := sink.Write(FragmentNode{ID: "node-2", Kinds: []string{"Device"}, Properties: map[string]any{"name": "Laptop"}}); err != nil { + t.Fatalf("write second node: %v", err) + } + if err := sink.Close(); err != nil { + t.Fatalf("close node parquet sink: %v", err) + } + + rows, err := parquet.ReadFile[parquetNodeRow](path) + if err != nil { + t.Fatalf("read node parquet file: %v", err) + } + if len(rows) != 2 { + t.Fatalf("read %d node rows, want 2", len(rows)) + } + if rows[0].ID != "node-1" || !reflect.DeepEqual(rows[0].Kinds, []string{"Person", "Employee"}) { + t.Fatalf("unexpected first node identity: %+v", rows[0]) + } + if rows[1].ID != "node-2" || !reflect.DeepEqual(rows[1].Kinds, []string{"Device"}) { + t.Fatalf("unexpected second node identity: %+v", rows[1]) + } + if !reflect.DeepEqual(rows[0].Properties, properties) { + t.Fatalf("first node properties = %#v, want %#v", rows[0].Properties, properties) + } + if !reflect.DeepEqual(rows[1].Properties, map[string]any{"name": "Laptop"}) { + t.Fatalf("second node properties = %#v, want %#v", rows[1].Properties, map[string]any{"name": "Laptop"}) + } + assertNodeParquetFooterSchema(t, path) +} + +func TestEdgeParquetSinkRoundTripsVariantProperties(t *testing.T) { + path := filepath.Join(t.TempDir(), "edges.parquet") + properties := representativeParquetProperties() + + sink, err := newEdgeParquetSink(path) + if err != nil { + t.Fatalf("create edge parquet sink: %v", err) + } + if err := sink.Write(FragmentEdge{StartID: "node-1", EndID: "node-2", Kind: "MemberOf", Properties: properties}); err != nil { + t.Fatalf("write first edge: %v", err) + } + if err := sink.Write(FragmentEdge{StartID: "node-2", EndID: "node-3", Kind: "AdminTo", Properties: map[string]any{"enabled": false}}); err != nil { + t.Fatalf("write second edge: %v", err) + } + if err := sink.Close(); err != nil { + t.Fatalf("close edge parquet sink: %v", err) + } + + rows, err := parquet.ReadFile[parquetEdgeRow](path) + if err != nil { + t.Fatalf("read edge parquet file: %v", err) + } + if len(rows) != 2 { + t.Fatalf("read %d edge rows, want 2", len(rows)) + } + if rows[0].StartID != "node-1" || rows[0].EndID != "node-2" || rows[0].Kind != "MemberOf" { + t.Fatalf("unexpected first edge identity: %+v", rows[0]) + } + if rows[1].StartID != "node-2" || rows[1].EndID != "node-3" || rows[1].Kind != "AdminTo" { + t.Fatalf("unexpected second edge identity: %+v", rows[1]) + } + if !reflect.DeepEqual(rows[0].Properties, properties) { + t.Fatalf("first edge properties = %#v, want %#v", rows[0].Properties, properties) + } + if !reflect.DeepEqual(rows[1].Properties, map[string]any{"enabled": false}) { + t.Fatalf("second edge properties = %#v, want %#v", rows[1].Properties, map[string]any{"enabled": false}) + } + assertEdgeParquetFooterSchema(t, path) +} + +func representativeParquetProperties() map[string]any { + return map[string]any{ + "name": "Ada", + "enabled": true, + "score": float64(42.5), + "nested": map[string]any{"labels": []any{"a", float64(2)}}, + } +} + +func assertNodeParquetFooterSchema(t *testing.T, path string) { + t.Helper() + + schema := readParquetFooterSchema(t, path) + assertRequiredParquetStringField(t, schema, "id") + assertParquetVariantField(t, schema, "properties") + + kinds := parquetSchemaField(t, schema, "kinds") + if !kinds.Required() { + t.Fatalf("Parquet field %q is not required", kinds.Name()) + } + if logicalType := kinds.Type().LogicalType(); logicalType == nil { + t.Fatalf("Parquet field %q has no LIST logical annotation", kinds.Name()) + } else if _, ok := logicalType.Value.(*format.ListType); !ok { + t.Fatalf("Parquet field %q logical type = %T, want *format.ListType", kinds.Name(), logicalType.Value) + } + listFields := kinds.Fields() + if len(listFields) != 1 || listFields[0].Name() != "list" || !listFields[0].Repeated() { + t.Fatalf("Parquet field %q LIST group = %#v, want one repeated list field", kinds.Name(), listFields) + } + elementFields := listFields[0].Fields() + if len(elementFields) != 1 || elementFields[0].Name() != "element" || !elementFields[0].Required() { + t.Fatalf("Parquet field %q element group = %#v, want one required element field", kinds.Name(), elementFields) + } + assertParquetStringNode(t, elementFields[0]) +} + +func assertEdgeParquetFooterSchema(t *testing.T, path string) { + t.Helper() + + schema := readParquetFooterSchema(t, path) + for _, name := range []string{"start_id", "end_id", "kind"} { + assertRequiredParquetStringField(t, schema, name) + } + assertParquetVariantField(t, schema, "properties") +} + +func readParquetFooterSchema(t *testing.T, path string) *parquet.Schema { + t.Helper() + + file, err := os.Open(path) + if err != nil { + t.Fatalf("open Parquet file for footer schema: %v", err) + } + defer file.Close() + info, err := file.Stat() + if err != nil { + t.Fatalf("stat Parquet file for footer schema: %v", err) + } + parquetFile, err := parquet.OpenFile(file, info.Size()) + if err != nil { + t.Fatalf("open Parquet footer: %v", err) + } + return parquetFile.Schema() +} + +func parquetSchemaField(t *testing.T, schema *parquet.Schema, name string) parquet.Field { + t.Helper() + + for _, field := range schema.Fields() { + if field.Name() == name { + return field + } + } + t.Fatalf("Parquet footer schema does not contain field %q", name) + return nil +} + +func assertRequiredParquetStringField(t *testing.T, schema *parquet.Schema, name string) { + t.Helper() + + field := parquetSchemaField(t, schema, name) + if !field.Required() { + t.Fatalf("Parquet field %q is not required", name) + } + assertParquetStringNode(t, field) +} + +func assertParquetStringNode(t *testing.T, node parquet.Node) { + t.Helper() + + if !node.Leaf() { + t.Fatalf("Parquet string node is a group: %s", node) + } + if kind := node.Type().Kind(); kind != parquet.ByteArray { + t.Fatalf("Parquet string physical kind = %s, want BYTE_ARRAY", kind) + } + logicalType := node.Type().LogicalType() + if logicalType == nil { + t.Fatal("Parquet string node has no STRING logical annotation") + } + if _, ok := logicalType.Value.(*format.StringType); !ok { + t.Fatalf("Parquet string logical type = %T, want *format.StringType", logicalType.Value) + } +} + +func assertParquetVariantField(t *testing.T, schema *parquet.Schema, name string) { + t.Helper() + + field := parquetSchemaField(t, schema, name) + if !field.Required() { + t.Fatalf("Parquet field %q is not required", name) + } + if logicalType := field.Type().LogicalType(); logicalType == nil { + t.Fatalf("Parquet field %q has no VARIANT logical annotation", name) + } else if _, ok := logicalType.Value.(*format.VariantType); !ok { + t.Fatalf("Parquet field %q logical type = %T, want *format.VariantType", name, logicalType.Value) + } +} From 7269bd5cb97183f4df474b9d39676c4206320c56 Mon Sep 17 00:00:00 2001 From: Wes <169498386+wes-mil@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:09:01 -0400 Subject: [PATCH 2/4] clean up --- retriever/archive_envelope.go | 4 ---- retriever/archive_tar.go | 5 ----- retriever/dump.go | 9 --------- retriever/load.go | 8 -------- retriever/memory_benchmark_test.go | 8 ++++---- retriever/metrics.go | 9 --------- retriever/progress.go | 12 ------------ 7 files changed, 4 insertions(+), 51 deletions(-) diff --git a/retriever/archive_envelope.go b/retriever/archive_envelope.go index ae60d4a3..ac92ff0c 100644 --- a/retriever/archive_envelope.go +++ b/retriever/archive_envelope.go @@ -341,10 +341,6 @@ func UnpackEncryptedCollectionArchiveWithOptions(reader io.Reader, outputDir str return nil } -func unpackEncryptedCollectionArchiveToDirectory(reader io.Reader, outputDir string, identity hpke.PrivateKey) error { - return UnpackEncryptedCollectionArchive(reader, outputDir, identity) -} - func createUnpackStagingDirectory(outputDir string, force bool) (string, error) { outputDir = strings.TrimSpace(outputDir) if outputDir == "" { diff --git a/retriever/archive_tar.go b/retriever/archive_tar.go index 15276a77..0898dfc4 100644 --- a/retriever/archive_tar.go +++ b/retriever/archive_tar.go @@ -376,11 +376,6 @@ func unpackTar(reader io.Reader, outputDir string, force bool) error { return UnpackTar(reader, outputDir, force) } -func unpackTarFile(reader io.Reader, outputDir, relativePath string, expectedSize int64) error { - _, err := unpackTarFileTracked(reader, outputDir, relativePath, expectedSize, false) - return err -} - func unpackTarFileTracked(reader io.Reader, outputDir, relativePath string, expectedSize int64, trackIntegrity bool) (unpackedFileIntegrity, error) { absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) if err := os.MkdirAll(filepath.Dir(absolutePath), 0o755); err != nil { diff --git a/retriever/dump.go b/retriever/dump.go index 804c79b9..b19296cb 100644 --- a/retriever/dump.go +++ b/retriever/dump.go @@ -507,15 +507,6 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio return graphEntry, schemaEntry, metricsEntry, nil } -func countGraphEntities(ctx context.Context, db graph.Database, targetGraph graph.Graph) (int64, int64, error) { - entitySnapshot, err := countGraphEntitySnapshot(ctx, db, targetGraph) - if err != nil { - return 0, 0, err - } - - return entitySnapshot.NodeCount, entitySnapshot.EdgeCount, nil -} - func validateCompletedDumpSources(ctx context.Context, db graph.Database, graphEntries []GraphManifest) error { for _, graphEntry := range graphEntries { snapshot, err := countGraphEntitySnapshot(ctx, db, graph.Graph{Name: graphEntry.Name}) diff --git a/retriever/load.go b/retriever/load.go index fd3be4fa..6f5e72a3 100644 --- a/retriever/load.go +++ b/retriever/load.go @@ -750,10 +750,6 @@ func readNodeFragmentFile(inputDir string, codec CompressionCodec, fileEntry Fil return decodeNodeFragmentFile(inputDir, codec, fileEntry, false, handle) } -func verifyNodeFragmentFile(inputDir string, codec CompressionCodec, fileEntry FileManifest) (int, error) { - return decodeNodeFragmentFile(inputDir, codec, fileEntry, true, nil) -} - func decodeNodeFragmentFile(inputDir string, codec CompressionCodec, fileEntry FileManifest, verifyIntegrity bool, handle func(FragmentNode) error) (int, error) { if fileEntry.Phase != PhaseNodes { return 0, fmt.Errorf("fragment %s has phase %q, expected nodes", fileEntry.Path, fileEntry.Phase) @@ -794,10 +790,6 @@ func readEdgeFragmentFile(inputDir string, codec CompressionCodec, fileEntry Fil return decodeEdgeFragmentFile(inputDir, codec, fileEntry, false, handle) } -func verifyEdgeFragmentFile(inputDir string, codec CompressionCodec, fileEntry FileManifest) (int, error) { - return decodeEdgeFragmentFile(inputDir, codec, fileEntry, true, nil) -} - func decodeEdgeFragmentFile(inputDir string, codec CompressionCodec, fileEntry FileManifest, verifyIntegrity bool, handle func(FragmentEdge) error) (int, error) { if fileEntry.Phase != PhaseEdges { return 0, fmt.Errorf("fragment %s has phase %q, expected edges", fileEntry.Path, fileEntry.Phase) diff --git a/retriever/memory_benchmark_test.go b/retriever/memory_benchmark_test.go index e362b751..629a565b 100644 --- a/retriever/memory_benchmark_test.go +++ b/retriever/memory_benchmark_test.go @@ -101,8 +101,8 @@ func BenchmarkEdgeMetricsAllocation(b *testing.B) { } b.ReportAllocs() - b.ResetTimer() - for index := 0; index < b.N; index++ { + + for b.Loop() { if err := builder.observeDatabaseRelationship(1, 2, "Edge"); err != nil { b.Fatal(err) } @@ -121,8 +121,8 @@ func BenchmarkEdgeScrubAllocation(b *testing.B) { scrubber.scrubPropertiesWithCounts(properties) b.ReportAllocs() - b.ResetTimer() - for index := 0; index < b.N; index++ { + + for b.Loop() { scrubber.scrubPropertiesWithCounts(properties) } } diff --git a/retriever/metrics.go b/retriever/metrics.go index 91fed115..f129db92 100644 --- a/retriever/metrics.go +++ b/retriever/metrics.go @@ -596,15 +596,6 @@ func validateMetricHistogramSum(graphName string, histogramName string, histogra return nil } -func cloneMetricHistogram(source map[string]int64) map[string]int64 { - target := make(map[string]int64, len(source)) - for key, count := range source { - target[key] = count - } - - return target -} - func metricKindSetKey(kinds []string) string { seen := map[string]struct{}{} for _, kind := range kinds { diff --git a/retriever/progress.go b/retriever/progress.go index 98575ba9..25a5795e 100644 --- a/retriever/progress.go +++ b/retriever/progress.go @@ -7,10 +7,6 @@ import ( const retrieverProgressEntityInterval int64 = DefaultProgressInterval -func retrieverInitialProgressAt(planned int64) int64 { - return retrieverInitialProgressAtInterval(planned, retrieverProgressEntityInterval) -} - func retrieverInitialProgressAtInterval(planned int64, interval int64) int64 { interval = normalizedProgressInterval(interval) if planned <= interval { @@ -32,10 +28,6 @@ func retrieverBatchLimit(remaining int64, batchSize int) int { return batchSize } -func logRetrieverEntityProgress(message string, graphName string, phaseName Phase, processed int64, planned int64, startedAt time.Time, nextProgressAt int64, progress ProgressFunc) int64 { - return logRetrieverEntityProgressInterval(message, graphName, phaseName, processed, planned, startedAt, nextProgressAt, progress, retrieverProgressEntityInterval) -} - func logRetrieverEntityProgressInterval(message string, graphName string, phaseName Phase, processed int64, planned int64, startedAt time.Time, nextProgressAt int64, progress ProgressFunc, interval int64) int64 { if nextProgressAt == 0 || processed < nextProgressAt || processed >= planned { return nextProgressAt @@ -74,10 +66,6 @@ func logRetrieverEntityProgressInterval(message string, graphName string, phaseN return retrieverNextProgressAtInterval(processed, planned, nextProgressAt, interval) } -func retrieverNextProgressAt(processed int64, planned int64, nextProgressAt int64) int64 { - return retrieverNextProgressAtInterval(processed, planned, nextProgressAt, retrieverProgressEntityInterval) -} - func retrieverNextProgressAtInterval(processed int64, planned int64, nextProgressAt int64, interval int64) int64 { interval = normalizedProgressInterval(interval) if nextProgressAt <= processed { From 9ccefe06c15b96beb043475feb5979724367976e Mon Sep 17 00:00:00 2001 From: Wes <169498386+wes-mil@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:20:22 -0400 Subject: [PATCH 3/4] remove double tmp files --- retriever/compression.go | 21 +++++++++++++++------ retriever/dump_checkpoint.go | 2 +- retriever/dump_checkpoint_test.go | 1 - retriever/fragment_writer.go | 5 ++--- retriever/fragment_writer_test.go | 13 ++++++++----- 5 files changed, 26 insertions(+), 16 deletions(-) diff --git a/retriever/compression.go b/retriever/compression.go index 932918e4..642a2178 100644 --- a/retriever/compression.go +++ b/retriever/compression.go @@ -218,6 +218,21 @@ func (s *compressedJSONLinesWriter) Count() int { } func (s *compressedJSONLinesWriter) Close() (FileManifest, error) { + fileEntry, err := s.finalize() + if err != nil { + return FileManifest{}, err + } + + if err := os.Rename(s.tempPath, s.path); err != nil { + _ = os.Remove(s.tempPath) + + return FileManifest{}, fmt.Errorf("rename fragment: %w", err) + } + + return fileEntry, nil +} + +func (s *compressedJSONLinesWriter) finalize() (FileManifest, error) { if s.closed { return FileManifest{}, fmt.Errorf("close JSONL fragment more than once") } @@ -236,12 +251,6 @@ func (s *compressedJSONLinesWriter) Close() (FileManifest, error) { return FileManifest{}, fmt.Errorf("close fragment file: %w", err) } - if err := os.Rename(s.tempPath, s.path); err != nil { - _ = os.Remove(s.tempPath) - - return FileManifest{}, fmt.Errorf("rename fragment: %w", err) - } - return FileManifest{ Count: s.count, CompressedBytes: s.compressedCounter.count, diff --git a/retriever/dump_checkpoint.go b/retriever/dump_checkpoint.go index 996d6c62..484f090f 100644 --- a/retriever/dump_checkpoint.go +++ b/retriever/dump_checkpoint.go @@ -282,7 +282,7 @@ func removeKnownDumpCheckpointTemps(outputDir string, value dumpCheckpoint) erro return err } parquetPath := filepath.Join(outputDir, filepath.FromSlash(nextParquetPath)) - paths = append(paths, jsonlStagingPath+".tmp", parquetPath, parquetPath+".tmp") + paths = append(paths, parquetPath, parquetPath+".tmp") } } for _, candidate := range paths { diff --git a/retriever/dump_checkpoint_test.go b/retriever/dump_checkpoint_test.go index 2e424642..a52fd955 100644 --- a/retriever/dump_checkpoint_test.go +++ b/retriever/dump_checkpoint_test.go @@ -353,7 +353,6 @@ func TestDumpResumeRemovesKnownNextShardTemps(t *testing.T) { } tempPaths := []string{ filepath.Join(outputDir, filepath.FromSlash(jsonlPath)) + ".tmp", - filepath.Join(outputDir, filepath.FromSlash(jsonlPath)) + ".tmp.tmp", filepath.Join(outputDir, filepath.FromSlash(parquetPath)) + ".tmp", } for _, tempPath := range tempPaths { diff --git a/retriever/fragment_writer.go b/retriever/fragment_writer.go index 6c9c1286..9e64a570 100644 --- a/retriever/fragment_writer.go +++ b/retriever/fragment_writer.go @@ -42,7 +42,7 @@ func newFragmentWriter[T any](jsonlPath, parquetPath string, options DumpOptions } jsonlStagingPath := jsonlPath + ".tmp" - jsonl, err := newCompressedJSONLinesWriterAtPaths(jsonlStagingPath, jsonlStagingPath+".tmp", options.Compression, options.ZstdLevel) + jsonl, err := newCompressedJSONLinesWriterAtPaths(jsonlPath, jsonlStagingPath, options.Compression, options.ZstdLevel) if err != nil { return nil, err } @@ -99,7 +99,7 @@ func (s *fragmentWriter[T]) Close() (FileManifest, error) { return s.jsonl.Close() } - fileEntry, err := s.jsonl.Close() + fileEntry, err := s.jsonl.finalize() if err != nil { s.cleanup(false, false) return FileManifest{}, err @@ -139,7 +139,6 @@ func (s *fragmentWriter[T]) cleanup(jsonlPublished, parquetPublished bool) { } _ = os.Remove(s.jsonlStagingPath) - _ = os.Remove(s.jsonlStagingPath + ".tmp") _ = os.Remove(s.parquetStagingPath) if jsonlPublished { _ = os.Remove(s.jsonlPath) diff --git a/retriever/fragment_writer_test.go b/retriever/fragment_writer_test.go index 4de476a6..d4fd6255 100644 --- a/retriever/fragment_writer_test.go +++ b/retriever/fragment_writer_test.go @@ -25,6 +25,10 @@ func TestNodeFragmentWriter(t *testing.T) { if err != nil { t.Fatalf("create node writer: %v", err) } + if _, err := os.Stat(jsonlPath + ".tmp"); err != nil { + t.Fatalf("stat JSONL staging file: %v", err) + } + assertAbsent(t, jsonlPath+".tmp.tmp") if err := writer.Write(first); err != nil { t.Fatalf("write first node: %v", err) } @@ -161,7 +165,6 @@ func TestNodeFragmentWriter(t *testing.T) { for _, path := range []string{ jsonlPath, jsonlPath + ".tmp", - jsonlPath + ".tmp.tmp", parquetPath, parquetPath + ".tmp", } { @@ -198,7 +201,7 @@ func TestNodeFragmentWriter(t *testing.T) { t.Fatalf("write error = %v, want %v", err, writeErr) } - for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath, parquetPath + ".tmp"} { + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", parquetPath, parquetPath + ".tmp"} { assertAbsent(t, path) } }) @@ -221,7 +224,7 @@ func TestNodeFragmentWriter(t *testing.T) { t.Fatal("expected JSONL close failure") } - for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath, parquetPath + ".tmp"} { + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", parquetPath, parquetPath + ".tmp"} { assertAbsent(t, path) } }) @@ -245,7 +248,7 @@ func TestNodeFragmentWriter(t *testing.T) { t.Fatal("expected Parquet close failure") } - for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath, parquetPath + ".tmp"} { + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", parquetPath, parquetPath + ".tmp"} { assertAbsent(t, path) } }) @@ -268,7 +271,7 @@ func TestNodeFragmentWriter(t *testing.T) { t.Fatal("expected Parquet publish rename failure") } - for _, path := range []string{jsonlPath, jsonlPath + ".tmp", jsonlPath + ".tmp.tmp", parquetPath + ".tmp"} { + for _, path := range []string{jsonlPath, jsonlPath + ".tmp", parquetPath + ".tmp"} { assertAbsent(t, path) } if info, err := os.Stat(parquetPath); err != nil || !info.IsDir() { From 8cae37eee44a9024e29e76413050bf84d363f9f0 Mon Sep 17 00:00:00 2001 From: Wes <169498386+wes-mil@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:27:15 -0400 Subject: [PATCH 4/4] fix fragment writer --- retriever/fragment_writer.go | 4 +--- retriever/fragment_writer_test.go | 34 +++++++------------------------ 2 files changed, 8 insertions(+), 30 deletions(-) diff --git a/retriever/fragment_writer.go b/retriever/fragment_writer.go index 9e64a570..63155eb4 100644 --- a/retriever/fragment_writer.go +++ b/retriever/fragment_writer.go @@ -5,8 +5,6 @@ import ( "os" ) -var newNodeParquetSinkForFragmentWriter = newNodeParquetSink - type fragmentWriter[T any] struct { jsonl *compressedJSONLinesWriter @@ -22,7 +20,7 @@ type fragmentWriter[T any] struct { } func newNodeFragmentWriter(jsonlPath, parquetPath string, options DumpOptions) (*fragmentWriter[FragmentNode], error) { - return newFragmentWriter(jsonlPath, parquetPath, options, newNodeParquetSinkForFragmentWriter) + return newFragmentWriter(jsonlPath, parquetPath, options, newNodeParquetSink) } func newEdgeFragmentWriter(jsonlPath, parquetPath string, options DumpOptions) (*fragmentWriter[FragmentEdge], error) { diff --git a/retriever/fragment_writer_test.go b/retriever/fragment_writer_test.go index d4fd6255..3bdd43f5 100644 --- a/retriever/fragment_writer_test.go +++ b/retriever/fragment_writer_test.go @@ -112,23 +112,13 @@ func TestNodeFragmentWriter(t *testing.T) { }) t.Run("Parquet write failure removes staging and final outputs", func(t *testing.T) { - original := newNodeParquetSinkForFragmentWriter - newNodeParquetSinkForFragmentWriter = func(path string) (*parquetFragmentSink[FragmentNode], error) { - sink, err := newNodeParquetSink(path) - if err != nil { - return nil, err - } - sink.write = func(FragmentNode) error { return errors.New("injected parquet write failure") } - return sink, nil - } - t.Cleanup(func() { newNodeParquetSinkForFragmentWriter = original }) - jsonlPath := filepath.Join(t.TempDir(), "nodes.jsonl") parquetPath := filepath.Join(t.TempDir(), "nodes.parquet") writer, err := newNodeFragmentWriter(jsonlPath, parquetPath, fragmentWriterTestOptions(t, true)) if err != nil { t.Fatalf("create node writer: %v", err) } + writer.parquet.write = func(FragmentNode) error { return errors.New("injected parquet write failure") } if err := writer.Write(FragmentNode{ID: "node-1"}); err == nil { t.Fatal("expected Parquet write failure") } @@ -173,23 +163,7 @@ func TestNodeFragmentWriter(t *testing.T) { }) t.Run("abort panic does not escape or block output cleanup", func(t *testing.T) { - original := newNodeParquetSinkForFragmentWriter writeErr := errors.New("injected parquet write failure") - newNodeParquetSinkForFragmentWriter = func(path string) (*parquetFragmentSink[FragmentNode], error) { - sink, err := newNodeParquetSink(path) - if err != nil { - return nil, err - } - originalAbort := sink.abort - sink.write = func(FragmentNode) error { return writeErr } - sink.abort = func() { - originalAbort() - panic("injected Parquet abort panic") - } - return sink, nil - } - t.Cleanup(func() { newNodeParquetSinkForFragmentWriter = original }) - outputDir := t.TempDir() jsonlPath := filepath.Join(outputDir, "nodes.jsonl") parquetPath := filepath.Join(outputDir, "nodes.parquet") @@ -197,6 +171,12 @@ func TestNodeFragmentWriter(t *testing.T) { if err != nil { t.Fatalf("create node writer: %v", err) } + originalAbort := writer.parquet.abort + writer.parquet.write = func(FragmentNode) error { return writeErr } + writer.parquet.abort = func() { + originalAbort() + panic("injected Parquet abort panic") + } if err := writer.Write(FragmentNode{ID: "node-1"}); !errors.Is(err, writeErr) { t.Fatalf("write error = %v, want %v", err, writeErr) }