Add a sample factor to mongodb planning - #336
Conversation
WalkthroughAdds a SampleFactor to Mongo connector settings and integrates a new CLI flag to configure it. Updates sampling logic to scale sample count, cap excessive sampling, and introduce a skip/cycle pattern during aggregation. Sets default SampleFactor=1. No other connectors or systems are changed. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant U as User/CLI
participant App as App Options
participant Conn as Mongo Connector
participant Plan as Plan Generator
participant Agg as Aggregation/Sampler
participant DB as MongoDB
U->>App: Run with --sample-factor (optional)
App->>Conn: Build ConnectorSettings{SampleFactor or default=1}
Conn->>Plan: Initialize planning with settings
Plan->>DB: count = aggregateCount()
Note over Plan,DB: Determine partition doc counts
Plan->>Agg: Compute numSamples = (count / targetPerPartition) * SampleFactor
alt Excessive sampling
Agg->>Agg: if numSamples*20 >= count → numSamples = count/20
Note over Agg: Warn about cap
end
rect rgb(240,248,255)
note right of Agg: New skip/cycle behavior
Agg->>Agg: factorCount = floor(SampleFactor/2) skips
loop Cycle through SampleFactor
Agg->>DB: fetch next sample batch
DB-->>Agg: sample docs
end
end
Agg-->>Plan: sampled docs
Plan-->>Conn: partition plan
Conn-->>App: ready
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
connectors/mongo/conn.go (2)
39-39: Clarify the SampleFactor field documentation.The comment "a factor to determine how many extra samples per partition are used" is vague and doesn't accurately describe the behavior. The actual implementation requests
SampleFactor * base_samplesfrom MongoDB, then selects a subset during iteration, effectively maintaining the target partition count while potentially improving boundary distribution.Consider expanding the comment to explain the full behavior, for example:
SampleFactor int // factor for oversampling partition boundaries; requests SampleFactor × base samples from MongoDB, then subsamples to maintain target partition count while improving boundary distribution
243-250: Add explanatory comments for the skip/cycle pattern.The skip/cycle logic is not immediately clear. With
SampleFactor = 10, the pattern is: skip 5, use 1, skip 9, use 1, skip 9, ... This effectively selects ~10% of the oversampled results to maintain the target partition count.Consider adding an inline comment explaining the pattern and its purpose:
+// Skip/cycle pattern: by requesting SampleFactor × samples and subsampling, +// we get better partition boundaries (more representative of _id distribution) +// while maintaining the target partition count. +// Pattern: skip SampleFactor/2 initially, then use 1, skip SampleFactor-1, repeat. var factorCount = c.settings.SampleFactor / 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
connectors/mongo/conn.go(3 hunks)internal/app/options/connectorflags.go(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (1)
connectors/mongo/conn.go (1)
986-986: Verify the default SampleFactor value (1 vs 10).The default here is
1, but the CLI flag ininternal/app/options/connectorflags.go(line 363) sets it to10. This means:
- Direct calls to
NewConn()getSampleFactor = 1(no oversampling)- CLI usage gets
SampleFactor = 10(10× oversampling)Is this intentional? If so, consider adding a comment explaining why different defaults are appropriate for different use cases.
| numSamples := (count / c.settings.TargetDocCountPerPartition) * int64(c.settings.SampleFactor) | ||
| if numSamples*20 >= count { | ||
| slog.Warn("Too many samples requested, adjusting...", "original", numSamples, "new", count/20) | ||
| numSamples = count / 20 | ||
| } |
There was a problem hiding this comment.
Fix integer overflow risk and add input validation.
Two issues:
-
Integer overflow: Line 235's guard condition
numSamples*20 >= countcan overflow ifnumSamplesis large (e.g., for a collection with billions of documents and a high SampleFactor). In Go, integer overflow wraps silently, potentially causing the guard to fail when it should trigger. -
Missing validation: No check ensures
SampleFactoris positive. IfSampleFactor ≤ 0, the sampling logic behaves unexpectedly (e.g., withSampleFactor = 0,numSamples = 0, and the skip/cycle logic uses every sample instead of none).
Apply this diff to fix both issues:
+if c.settings.SampleFactor <= 0 {
+ return fmt.Errorf("SampleFactor must be positive, got %d", c.settings.SampleFactor)
+}
numSamples := (count / c.settings.TargetDocCountPerPartition) * int64(c.settings.SampleFactor)
-if numSamples*20 >= count {
+// Guard against excessive sampling: cap at 5% of collection
+// Use division to avoid overflow: numSamples*20 >= count ⟺ numSamples >= count/20
+if count < 20 || numSamples >= count/20 {
slog.Warn("Too many samples requested, adjusting...", "original", numSamples, "new", count/20)
numSamples = count / 20
+ if numSamples < 1 {
+ numSamples = 1
+ }
}Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In connectors/mongo/conn.go around lines 234-238, the multiplication
numSamples*20 can overflow and SampleFactor isn't validated; first validate
c.settings.SampleFactor > 0 (return an error or default it to 1) before
computing numSamples, cast/compute all values as int64, compute numSamples =
(count / c.settings.TargetDocCountPerPartition) * int64(c.settings.SampleFactor)
but avoid a direct multiply check by computing a safe threshold (e.g., threshold
:= count / 20) and then compare if numSamples >= threshold (or use
math.Ceil-style safe logic) to decide adjustment, and if adjustment is needed
set numSamples = threshold; ensure all arithmetic uses int64 to prevent wrapping
and add a clear error or panic path for invalid SampleFactor.
| altsrc.NewIntFlag(&cli.IntFlag{ | ||
| Name: "sample-factor", | ||
| Destination: &settings.SampleFactor, | ||
| Usage: "Number of extra samples per partition", | ||
| Value: 10, | ||
| }), |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Clarify the usage text for the sample-factor flag.
The usage text "Number of extra samples per partition" is misleading. The flag doesn't directly increase the number of samples used per partition; rather, it controls oversampling: MongoDB is asked for SampleFactor × base_samples, which are then subsampled during iteration to maintain the target partition count while improving boundary distribution.
Consider revising the usage text to be more accurate:
altsrc.NewIntFlag(&cli.IntFlag{
Name: "sample-factor",
Destination: &settings.SampleFactor,
- Usage: "Number of extra samples per partition",
+ Usage: "Oversampling factor for partition boundaries (requests N× samples, then subsamples to improve distribution)",
Value: 10,
}),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| altsrc.NewIntFlag(&cli.IntFlag{ | |
| Name: "sample-factor", | |
| Destination: &settings.SampleFactor, | |
| Usage: "Number of extra samples per partition", | |
| Value: 10, | |
| }), | |
| altsrc.NewIntFlag(&cli.IntFlag{ | |
| Name: "sample-factor", | |
| Destination: &settings.SampleFactor, | |
| Usage: "Oversampling factor for partition boundaries (requests N× samples, then subsamples to improve distribution)", | |
| Value: 10, | |
| }), |
🤖 Prompt for AI Agents
In internal/app/options/connectorflags.go around lines 359 to 364, the usage
text for the "sample-factor" flag is misleading; update the Usage string to
explain that this factor controls oversampling (MongoDB is requested for
SampleFactor × base_samples, which are then subsampled to achieve the target
partition count and improve boundary distribution). Change the Usage to a
concise, accurate description reflecting oversampling and subsampling semantics
and include the default value where helpful.
| if err != nil { | ||
| return err | ||
| } | ||
| var factorCount = c.settings.SampleFactor / 2 |
There was a problem hiding this comment.
what's the meaning of this?
2 participants
Summary by CodeRabbit