[FEA] Use catalyst stats as rowsHint in RegexComplexityEstimator to avoid CPU fallback on small inputs

[wjxiz1992]

Describe the issue

RegexComplexityEstimator.isValid always computes numRows as gpuTargetBatchSizeBytes / StringType.defaultSize (≈ 53,687,091 with the default 1 GiB batch and Spark's 20-byte string default). This is intended as a per-batch upper bound, but it is applied even when the actual input is far smaller than one batch — small lookup tables, small joins, post-aggregation outputs — which are exactly the cases where regex over-rejection is the most painful (no real memory risk, but the planner still falls back to CPU).

After #14849 fixed the long-standing Int overflow in numStates * numRows * 2, some patterns that previously slipped through (because the overflow wrapped to a negative number that compared <= 2 GiB) are now correctly rejected. The over-rejection becomes user-visible.

Reproduction

70k-row table written to ORC, then queried with a 7-alternation pattern. The pattern has numStates = 64, and at the per-batch ceiling the estimator reports 64 × 53,687,091 × 2 ≈ 6.4 GiB, above the default 2 GiB spark.rapids.sql.regexp.maxStateMemoryBytes. The query falls back to CPU even though the table has only 70,000 rows and the actual NFA scratch needed is in the tens of KB.

Pattern: (quick brown|brown fox|fox jumps|jumps over|over the|the lazy|lazy dog)

import datetime, os
from pyspark.sql import SparkSession
from pyspark.sql.types import *
import pyspark.sql.functions as F

schema = StructType([
    StructField("rlike_multiple_match", StringType(), True),
])
data = [(f"row-{i} The quick brown fox jumps over the lazy dog.",) for i in range(70000)]
spark = SparkSession.builder.appName("regex-fallback-repro").getOrCreate()
spark.conf.set("spark.rapids.sql.enabled", "true")
spark.conf.set("spark.sql.adaptive.enabled", "false")
spark.createDataFrame(data, schema).write.orc("/tmp/df_rlike_orc", mode="overwrite")
df = spark.read.orc("/tmp/df_rlike_orc").withColumn("uniqueID", F.monotonically_increasing_id())
df.createOrReplaceTempView("df_rlike_orc_table")
spark.sql("""
  select rlike(rlike_multiple_match,
    "(quick brown|brown fox|fox jumps|jumps over|over the|the lazy|lazy dog)") as literal,
    uniqueID from df_rlike_orc_table order by literal, uniqueID
""").explain()

Observed: RLike runs on CPU with the message "estimated memory needed for regular expression exceeds the maximum. Set spark.rapids.sql.regexp.maxStateMemoryBytes to change it."

Root cause

RegexComplexityEstimator.estimateGpuMemory derives numRows purely from gpuTargetBatchSizeBytes / StringType.defaultSize. This is the worst-case rows-per-batch and never considers the actual input size. Catalyst already tracks Statistics.sizeInBytes (always populated) and Statistics.rowCount (when CBO or leaf relation provides it) on the LogicalPlan that the enclosing SparkPlan links to via logicalLink. We can read this through the meta hierarchy at planning time and pass it as a rowsHint to the estimator.

Proposed fix

Add an optional rowsHint: Option[Long] parameter to RegexComplexityEstimator.isValid, used as min(perBatchRows, rowsHint) — strictly tightens, never inflates, the estimate. Compute the hint in GpuRegExpUtils.validateRegExpComplexity by walking ExprMeta.parent up to the enclosing SparkPlanMeta, then reading wrapped.logicalLink.stats:

• rowCount if present (exact)
• sizeInBytes / StringType.defaultSize as upper bound otherwise

For the reproduction above, the catalyst-projected sizeInBytes is 5.5 KiB (Project output is boolean + bigint = 9 bytes per row), so bySize ≈ 281, and 64 × 281 × 2 ≈ 35 KB — comfortably below the 2 GiB budget. RLike (further rewritten to GpuContainsAny) runs on GPU.

Workaround in current releases

Either raise spark.rapids.sql.regexp.maxStateMemoryBytes (e.g. 8 GiB) or lower spark.rapids.sql.batchSizeBytes (e.g. 256 MiB) so the per-batch row ceiling shrinks proportionally. Both are safe for small inputs; neither is a great default.

[wjxiz1992]

Surfacing a performance / safety trade-off discovered while validating #14888 locally — looking for opinions before we lock in the design.

The trade-off

The hint shrinks the planning-time estimate by replacing the per-batch ceiling (gpuTargetBatchSizeBytes / StringType.defaultSize ≈ 53M rows at defaults) with min(perBatchRows, statsHint). That ceiling was historically used as a conservative gate: "assume each batch is full, never let GPU regex run if worst-case scratch exceeds maxRegExpStateMemory."

When catalyst stats are accurate (leaf relations from Parquet/ORC scans, or after ANALYZE TABLE), the hint is a strict tightening and there is no safety regression.

When catalyst stats are inaccurate — and this happens by design after Filter, Join, Aggregate, or any V2 source that doesn't report row counts — the hint can underestimate the rows that will actually flow through one batch at runtime. In that case the gate now lets through patterns that, in the worst case, would allocate up to the original per-batch ceiling (numStates × perBatchRows × 2) of NFA scratch. Concretely for a 64-state pattern at default conf, that's still bounded at ~6.4 GiB per batch.

The cuDF kernel allocates scratch per batch, not per table, so total table size does not amplify this. The risk is purely: a single batch of input grows large because the actual Filter selectivity (or join cardinality) deviated from the catalyst estimate, and the resulting numStates × actual_batch_rows × 2 exceeds GPU memory pressure tolerance.

Numbers from my local validation

• Box: RTX 5880 Ada, 48 GiB total memory
• Conf: spark.rapids.memory.gpu.allocFraction=0.3 → 14.4 GiB pool
• A 64-state alternation pattern on a Filter whose catalyst estimate is off by 100× could request ~6.4 GiB regex scratch + ~2.6 GiB input string column + concurrent operators in the same pool, leaving 3–5 GiB margin
• This is not a hypothetical — it's exactly the regime where the original conservative gate was protecting us

On A100 80 GB or H100 the margin is much larger, but on production deployments with smaller GPUs or higher tenancy, OOM under unfavorable stats is possible.

Failure mode

The spark-rapids retry framework catches GPU OOM and splits the batch, eventually falling back to CPU. So:

• ✅ No correctness regression
• ✅ No driver crash
• ⚠️ Latency: "GPU attempt → OOM → split-retry → maybe second OOM → eventual CPU" is much slower than "CPU from the start"

So the net effect on queries with bad stats can be slower than before this PR, even though the GPU path is faster on accurate-stats queries.

Options under consideration

A. Trust stats, keep current design. Argument: catalyst's stats are conservative by default (Spark uses defaultSizeInBytes = Long.MaxValue when stats are absent, and deriveRowsHint already returns None in that case); the cuDF retry framework is exactly the mechanism that handles this. PR description should call out the trade-off honestly.

B. Add a floor. Don't let the hint pull numRows below some fraction of perBatchRows, e.g.:

numRows = math.max(perBatchRows / 10, math.min(perBatchRows, hint))

The 10× floor preserves 90% of the small-table benefit (in #14887's example, 5.3M instead of 281 rows, estimate 670 MiB instead of 35 KiB, still well under the 2 GiB budget for the 64-state case) while preventing the worst-case underestimate.

C. Add an off switch. Introduce spark.rapids.sql.regexp.statsHintEnabled (default true) so deployments that prefer the original conservative behavior can opt out without a rebuild.

D. (B) + (C). Belt and suspenders.

What I'd like input on

1. Is the cuDF retry → split-and-CPU-fallback path actually graceful for the regex kernel specifically, or are there known sharp edges where regex OOM kills the executor / leaks GPU memory?
2. On real production stats accuracy — is the Filter/Join underestimation regime common enough to justify the floor in option B?
3. Anyone with strong opinions on conf-name precedent (option C)?

Happy to iterate on the PR with whichever direction we land on. CC people who'd have context on the cuDF regex memory model and on spark-rapids stats-handling history.

[thirtiseven]

Good point I have a few thoughts.

In ProjectExec we use the PreProjectSplitIterator to estimate memory usage, so the batch size input to a ProjectExec is unlikely to be the default number from batch size config. So the idea in this issue would also apply for large inputs with some code changes:

Currently we only estimate memory based on output size (see PreSplitOutSizeEstimator), because we cannot see scratch allocations in the estimator. Since we have RegexComplexityEstimator.estimateGpuMemory we can integrate the regex estimator to PreProjectSplitIterator. Once a batch can be pre-split smaller we can handle more complex expressions without retry or fallback.

Is the cuDF retry → split-and-CPU-fallback path actually graceful for the regex kernel specifically, or are there known sharp edges where regex OOM kills the executor / leaks GPU memory?

For the reliability, project split retry should be fine but I don't think we'll achieve good performance when either split retry or fallback to CPU occurs. Also I don't think we will fallback in the middle of a query, so with split retry, it will run on the GPU until the end. So, in theory, it still has a chance of beating the CPU even if split retry happens, but it's better to avoid it.

[wjxiz1992]

Following up on @thirtiseven's PreProjectSplitIterator suggestion — I prototyped it locally and benchmarked it against the CPU fallback it would replace. It holds up, with one caveat that turns into a concrete policy. (The baseline GPU-vs-CPU-vs-fallback numbers are in the PR perf comment.)

Prototype (local, ~90 lines, not pushed)

• Expose the regex numStates and invert the budget: maxRowsForBudget = maxRegExpStateMemory / (numStates * 2).
• Carry numStates into GpuRLike; GpuProjectExec computes the per-batch row budget and PreProjectSplitIterator caps rows/batch to it (alongside the existing output-size split).
• Relax the complexity gate only for a regex inside a ProjectExec (which pre-splits). Regex in Filter/Join/Aggregate still rejects to CPU exactly as today.

Validated: builds (buildver 330), the existing RegexComplexityEstimatorSuite is still green, GPU output == CPU output, and the NUM_PRE_SPLITS metric confirms batches are actually split to the budget. A genuine 35-state pattern ([a-z]{2,5} × 6, which is rejected today) now runs on GPU in a Project, while the same pattern in a Filter still falls back — i.e. the relaxation is scoped.

Perf

35-state regex, sum(cast(rlike(...) as int)), median ms, RTX 5880 Ada / local[8] / Spark 3.3.0.

Part A — default 2 GiB budget (the pattern is needlessly rejected today; no splitting needed at these sizes):

rows	cpu	fallback (today)	pre-split (b)
70k	79	155	143
1M	200	260	131
10M	1219	1556	214

Part B — 10M rows, shrinking the budget to force increasing split counts (the overhead curve):

pre-splits	pre-split median
0	208
8	205
96	336
992	1520

Read

In the common case (oversized pattern, few/no splits) pre-split runs on GPU at a flat ~130–214 ms — 2–7× faster than today's CPU fallback, which is the slowest option at every size. The only regime where it loses is heavy splitting: ~1.3 ms/split overhead, so it reaches break-even with CPU/fallback at roughly ~1000 splits on 10M rows.

So the policy the data points to: pre-split up to a split-count cap (a few hundred), and keep the existing CPU-fallback gate above it. That captures every win in the table and avoids the many-tiny-batches regime — which I think is the concrete form of @thirtiseven's "it still has a chance to beat CPU, but better to avoid it."

This is also stats-independent, so it works for path-based file reads — unlike the rowsHint approach in #14888, which I found is inert unless the source is an ANALYZE-d catalog table. So this could supersede that PR.

@revans2 @thirtiseven — does this direction look right? Two specific questions: (1) is the Project-scoped gate relaxation acceptable as a first step, or would you want pre-split extended to Filter/Join before changing the gate at all; (2) where would you set the split cap — a fixed constant, or derived from the GPU pool size? Happy to push the branch and open it as a PR if the direction is agreed.