Skip to content

[BUG] Pre-decode schema estimate is producing underfull batches #15694

Description

@kuhushukla

Describe the bug
MultiFileCoalescingPartitionReaderBase decides how many row groups to coalesce into one read batch by gating on GpuBatchUtils.estimateGpuMemory(schema, rows) against maxReadBatchSizeBytes. That estimate is a schema-derived worst case computed before decode — it sums
per-column fixed widths plus validity buffers from the type, with no knowledge of encoding, compression, dictionary reuse, or actual string lengths.

When the estimate over-counts, the reader can stop coalescing far below the real memory need and emits small batches for data that has nested and empty values for example. That is the suspected mechanism behind the underfull batches. On scan heavy queries the fullness drives better gpu time for those stages' tasks.

Steps/Code to reproduce bug

 // Excerpt--

  /**
   * When true, the coalesced read-batch is NOT capped by the schema-estimated GPU bytes
   * (maxReadBatchSizeBytes); it is bounded only by maxReadBatchSizeRows and the per-file
   * split boundary. Overridden by readers whose downstream decode already bounds the output
   * size (e.g. the cuDF chunked reader sizes output by actual decoded bytes).
   */
  protected def skipReadBatchByteLimit: Boolean = false

  -if (numBytes == 0 || numBytes + estimatedBytes <= maxReadBatchSizeBytes) {
  +if (numBytes == 0 || skipReadBatchByteLimit ||
  +    numBytes + estimatedBytes <= maxReadBatchSizeBytes) {

  Second, turn it on for the Parquet reader (GpuParquetScan.scala, in MultiFileParquetPartitionReader):

  override protected def skipReadBatchByteLimit: Boolean = true

  In words: the coalescing loop keeps accumulating row groups into one read batch until it hits the row limit or a file/split boundary, with the pre-decode GpuBatchUtils.estimateGpuMemory check bypassed entirely.

  Caveat : we enabled it unconditionally, which also covers the non-chunked Table.readParquet path where the estimate is the only GPU-memory guard. A real fix should scope the override to readers whose decode already bounds output size.

Expected behavior
For chunked reading, just let chunking logic take care of it. Tests show it works as intended.

Environment details (please complete the following information)

  • Environment location: [Local on A5000]
  • Use eventlogs for fullness calculation

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions