peter-toth commented on code in PR #58895:
URL: https://github.com/apache/spark/pull/58895#discussion_r4084280609
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -421,11 +552,313 @@ public boolean nextBatch() throws IOException {
return true;
}
+ /**
+ * Splicing emit path. Closes the previous emit's dequeued key vectors
(releasing survivor memory
+ * incrementally), advances to the next row group if needed via {@link
#checkEndOfRowGroup()},
+ * dequeues one survivor key vector per key column, drives non-key column
readers for {@code num}
+ * rows, and assembles a fresh {@link ColumnarBatch} interleaving key
(dequeued) and non-key
+ * (persistent value-vector) slots in the original projection order.
+ * The per-emit batch is a transient view over vectors owned elsewhere; see
{@link #close()} and
+ * {@link #closeSplicingState()}.
+ */
+ private boolean nextBatchSplicing() throws IOException {
+ if (pendingCloseKeyVectors != null) {
+ for (WritableColumnVector v : pendingCloseKeyVectors) {
+ if (v != null) v.close();
+ }
+ pendingCloseKeyVectors = null;
+ }
+ for (int i = 0; i < columnVectors.length; i++) {
+ if (isKeyTopLevel[i]) continue;
+ columnVectors[i].reset();
+ }
+ // Match the eager path and zero the outgoing batch before the terminal
checks below. Without
+ // this, a terminal call leaves `columnarBatch` pointing at the previous
emit whose key slots
+ // were just closed above -- with off-heap vectors those buffers are
already freed, so a
+ // consumer that read the batch after nextBatch() returned false would see
freed memory.
+ if (columnarBatch != null) columnarBatch.setNumRows(0);
+ if (hitEndOfData) return false;
+ if (rowsReturned >= totalRowCount) return false;
+ checkEndOfRowGroup();
+ if (hitEndOfData) return false;
+
+ int num = (int) Math.min(capacity, totalCountLoadedSoFar - rowsReturned);
+
+ WritableColumnVector[] dequeued = new
WritableColumnVector[keyVectorQueues.length];
+ for (int i = 0; i < keyVectorQueues.length; i++) {
+ dequeued[i] = keyVectorQueues[i].removeFirst();
Review Comment:
Fixed, and the shape changed rather than the assignment moving: a survivor
vector now stays owned by its queue while the batch is built on it, and the
next emit removes and closes it. So there is no window where a vector is out of
the queues and not yet reachable from a field, and `pendingCloseKeyVectors` is
gone.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -492,6 +929,450 @@ private void checkEndOfRowGroup() throws IOException {
totalCountLoadedSoFar += pages.getRowCount();
}
+ /**
+ * Loads the next row group using the three-phase late-materialization
pattern, all driven by the
+ * single {@link #lateMatReader} with its requested schema mutated per phase:
+ * - Phase 0 (full schema): compute {@code pushedFilterRanges} from the
pushed data filter via
+ * column index (metadata-only) using {@link
ParquetFileReader#getRowRanges}.
+ * - Phase 1 (key-only schema): read key-column pages restricted to {@code
pushedFilterRanges},
+ * evaluate the storage filter per row, build {@code finalRanges}.
+ * - Phase 2 (non-key schema): read non-key columns restricted to {@code
finalRanges}.
+ * Skipped entirely when {@link #nonKeyRequestedSchema} is null
(all-keys projection).
+ *
+ * Row groups for which {@code finalRanges} is empty are skipped entirely
(no phase-2 IO).
+ * Sets {@link #hitEndOfData} when all row groups have been processed.
+ */
+ private void loadNextRowGroupWithLateMaterialization() throws IOException {
+ while (nextBlockIndex < totalBlockCount) {
+ int blockIdx = nextBlockIndex++;
+ long blockRowCount =
lateMatReader.getRowGroups().get(blockIdx).getRowCount();
+ if (blockRowCount == 0) {
+ // parquet-mr never writes these, but RowRanges.createSingle(0) would
build Range(0, -1) and
+ // trip parquet's own `from <= to` assertion. The plain read path
skips them too.
+ continue;
+ }
+
+ // Phase 0: rows allowed by the pushed data filter (column-index
granularity). Restore the
+ // full requestedSchema first: phases 1 and 2 below narrow the reader's
schema, and
+ // ParquetFileReader.getRowRanges computes ranges against whatever
schema is set (it passes
+ // the reader's current `paths` to ColumnIndexFilter). This is defensive
-- with column-index
+ // filtering on, getFilteredRecordCount() in initialize() has already
memoized every block's
+ // ranges under the full schema, and with it off we do not call
getRowRanges at all.
+ lateMatReader.setRequestedSchema(requestedSchema);
+ // ParquetFileReader.getRowRanges only checks whether a filter is
pushed, NOT
+ // options.useColumnIndexFilter(), so calling it unconditionally would
keep applying
+ // column-index filtering after a user turned it off. That conf is the
documented escape hatch
+ // for files whose column index is wrong, and trusting a wrong column
index here would drop
+ // rows for good: finalRanges is a subset of pushedFilterRanges, and the
post-scan Filter no
+ // longer holds this predicate.
+ RowRanges pushedFilterRanges = useColumnIndexFilter
+ ? lateMatReader.getRowRanges(blockIdx)
+ : RowRanges.createSingle(blockRowCount);
+ if (pushedFilterRanges.rowCount() == 0) {
+ // Pushed data filter rejects this block entirely via column index.
Not a storage-filter
+ // skip, so we don't increment storage-filter metrics.
+ continue;
+ }
+
+ // Both byte metrics are always wired in production (FileSourceScanLike
creates all five
+ // whenever storageFilters is non-empty), so this only skips the work on
the test-only path
+ // that drives the reader directly. compressedBytesForRowRanges never
does IO of its own, so
+ // there is nothing here to avoid on the production path.
+ StorageFilterMetrics m = storageFilter.metrics();
+ SQLMetric bytesAvoidedRg = m.bytesAvoidedByRowGroup();
+ SQLMetric bytesAvoidedPf = m.bytesAvoidedByPageFiltering();
+ boolean needBytes = bytesAvoidedRg != null || bytesAvoidedPf != null;
+ long baselineRows = pushedFilterRanges.rowCount();
+ long baselineBytes = needBytes
+ ? compressedBytesForRowRanges(lateMatReader, blockIdx,
requestedSchema,
+ pushedFilterRanges)
+ : 0L;
+
+ // Phase 1: switch to key-only schema, read key columns under
pushedFilterRanges, evaluate the
+ // storage filter per row.
+ lateMatReader.setRequestedSchema(keyOnlyRequestedSchema);
+ long phase1Bytes = needBytes
+ ? compressedBytesForRowRanges(lateMatReader, blockIdx,
keyOnlyRequestedSchema,
+ pushedFilterRanges)
+ : 0L;
+ PageReadStore keyPages = lateMatReader.readFilteredRowGroup(blockIdx,
pushedFilterRanges);
+ if (keyPages == null) {
+ // Unreachable: readFilteredRowGroup returns null only for an empty
block, and we already
+ // know pushedFilterRanges selects at least one row. Skipping the
block here would drop its
+ // surviving rows from the output, so assert rather than `continue`.
+ throw new IllegalStateException(
+ "No key pages for row group " + blockIdx + " despite "
+ + pushedFilterRanges.rowCount() + " rows selected by the
pushed filter");
+ }
+ RowRanges finalRanges = evaluateStorageFilter(keyPages,
pushedFilterRanges);
+
+ if (finalRanges.rowCount() == 0) {
+ // Every surviving row was rejected by the storage filter; skip block
entirely. We still
+ // paid phase-1 to read the key column, so the bytes avoided vs a
no-storage-filter read
+ // are baseline - phase1 (the non-key bytes the no-filter path would
have read).
+ SQLMetric rgSkipped = m.rowGroupsSkipped();
+ if (rgSkipped != null) rgSkipped.add(1L);
+ SQLMetric rowsExcludedRg = m.rowsExcludedByRowGroup();
+ if (rowsExcludedRg != null) rowsExcludedRg.add(baselineRows);
+ if (bytesAvoidedRg != null) bytesAvoidedRg.add(baselineBytes -
phase1Bytes);
+ continue;
+ }
+
+ // Phase 2: switch to non-key schema, read non-key columns under
finalRanges. Skipped entirely
+ // when the projection is all keys (nonKeyRequestedSchema is null); emit
reconstructs each
+ // batch from the key queues alone.
+ long keptRows;
+ long phase2Bytes;
+ PageReadStore dataPages = null;
+ if (nonKeyRequestedSchema == null) {
+ keptRows = finalRanges.rowCount();
+ phase2Bytes = 0L;
+ } else {
+ lateMatReader.setRequestedSchema(nonKeyRequestedSchema);
+ // requireOffsetIndexesForPhase2() already established that every
projected column of every
+ // row group has an offset index, so this page-filtering read cannot
fail for want of one.
+ dataPages = lateMatReader.readFilteredRowGroup(blockIdx, finalRanges);
+ if (dataPages == null) {
+ // Unreachable: readFilteredRowGroup returns null only for an empty
block or empty ranges,
+ // both excluded above. Match phase 1 and fail with a message rather
than an NPE.
+ throw new IllegalStateException(
+ "No data pages for row group " + blockIdx + " despite " +
finalRanges.rowCount()
+ + " surviving rows");
+ }
+ keptRows = dataPages.getRowCount();
+ phase2Bytes = bytesAvoidedPf != null
+ ? compressedBytesForRowRanges(lateMatReader, blockIdx,
nonKeyRequestedSchema,
+ finalRanges)
+ : 0L;
+ }
+ long filteredRows = baselineRows - keptRows;
+ SQLMetric rowsExcludedWithinRg = m.rowsExcludedWithinRowGroup();
+ if (rowsExcludedWithinRg != null && filteredRows > 0)
rowsExcludedWithinRg.add(filteredRows);
+ if (bytesAvoidedPf != null) {
+ bytesAvoidedPf.add(baselineBytes - phase1Bytes - phase2Bytes);
+ }
+
+ if (dataPages != null) {
+ if (rowIndexGenerator != null) {
+ rowIndexGenerator.initFromPageReadStore(dataPages);
+ }
+ for (int i = 0; i < columnVectors.length; i++) {
+ if (isKeyTopLevel[i]) {
+ // Key columns are sourced from the queues during emit; skip
phase-2 reader init.
+ continue;
+ }
+ initColumnReader(dataPages, columnVectors[i]);
+ }
+ }
+ totalCountLoadedSoFar += keptRows;
+ return;
+ }
+ hitEndOfData = true;
+ }
+
+ /**
+ * Compressed bytes the reader transfers for the leaf columns of {@code
schema} when it reads
+ * exactly {@code rowRanges} of the given block. Page headers and the
dictionary page are
+ * included, because both are read whenever any page of a chunk is read.
+ *
+ * <p>Two sources, chosen so this never causes IO of its own:
+ * <ul>
+ * <li>{@code rowRanges} covers the whole block: the answer is the sum of
the chunks'
+ * {@code getTotalSize()}, which is already in the footer. This is the
case that matters --
+ * whenever nothing else has built the block's {@link
ColumnIndexStore}, {@code rowRanges}
+ * is necessarily the whole block, because a narrower range can only
come from column-index
+ * filtering, which builds the store as a side effect.
+ * <li>{@code rowRanges} is a strict subset: walk the offset index, as
parquet's own read path
+ * does, and add the dictionary page the way {@code
calculateOffsetRanges} does. The store
+ * is guaranteed to exist here, so the walk is pure metadata
arithmetic.
+ * </ul>
+ *
+ * <p>Columns absent from this physical file (schema evolution) contribute
nothing, which is
+ * correct: the reader transfers nothing for them. Every caller for a given
block walks the same
+ * metadata, so a skipped column drops out of the baseline and the per-phase
totals alike.
+ */
+ private static long compressedBytesForRowRanges(
+ ParquetFileReader reader,
+ int blockIndex,
+ MessageType schema,
+ RowRanges rowRanges) {
+ if (schema == null || rowRanges.rowCount() == 0 ||
schema.getColumns().isEmpty()) {
+ return 0L;
+ }
+ BlockMetaData block = reader.getRowGroups().get(blockIndex);
+ long blockRowCount = block.getRowCount();
+ Map<ColumnPath, ColumnChunkMetaData> chunks = new HashMap<>();
+ for (ColumnChunkMetaData chunk : block.getColumns()) {
+ chunks.put(chunk.getPath(), chunk);
+ }
+ boolean wholeBlock = rowRanges.rowCount() == blockRowCount;
+ ColumnIndexStore ciStore = wholeBlock ? null :
reader.getColumnIndexStore(blockIndex);
+ long total = 0L;
+ for (ColumnDescriptor column : schema.getColumns()) {
+ ColumnPath path = ColumnPath.get(column.getPath());
+ ColumnChunkMetaData chunk = chunks.get(path);
+ if (chunk == null) {
+ // Column is in the (clipped) requested schema but not in this file.
+ continue;
+ }
+ if (wholeBlock) {
+ total += chunk.getTotalSize();
+ continue;
+ }
+ OffsetIndex offsetIndex;
+ try {
+ offsetIndex = ciStore.getOffsetIndex(path);
+ } catch (MissingOffsetIndexException e) {
+ continue;
+ }
+ if (offsetIndex == null) {
+ continue;
+ }
+ // The dictionary page is read whenever any data page of the chunk is,
so count it here the
+ // same way parquet's ColumnIndexFilterUtils.calculateOffsetRanges does.
+ total += dictionaryPageSize(chunk);
+ int pageCount = offsetIndex.getPageCount();
+ for (int i = 0; i < pageCount; i++) {
+ long from = offsetIndex.getFirstRowIndex(i);
+ long to = offsetIndex.getLastRowIndex(i, blockRowCount);
+ if (rowRanges.isOverlapping(from, to)) {
+ total += offsetIndex.getCompressedPageSize(i);
+ }
+ }
+ }
+ return total;
+ }
+
+ /**
+ * Compressed size of a chunk's dictionary page, or 0 if it has none.
+ * {@link ColumnChunkMetaData#getStartingPos()} already resolves to the
dictionary page offset
+ * when there is a valid one, so the gap up to the first data page is
exactly the dictionary page.
+ */
+ private static long dictionaryPageSize(ColumnChunkMetaData chunk) {
+ long startingPos = chunk.getStartingPos();
+ long firstDataPageOffset = chunk.getFirstDataPageOffset();
+ return startingPos < firstDataPageOffset ? firstDataPageOffset -
startingPos : 0L;
+ }
+
+ /**
+ * Reads all rows of the given key-only {@link PageReadStore} (which
contains only rows in
+ * {@code pushedFilterRanges}) in capacity-sized chunks, evaluates the
storage filter on each row,
+ * builds a {@link RowRanges} of surviving rows in original block-row
coordinates, and appends
+ * survivor key values into the per-key-column accumulators ({@link
#currentKeyAccumulators}).
+ * When an accumulator hits {@link #capacity}, it's pushed into {@link
#keyVectorQueues} and a
+ * fresh one is allocated. After all rows have been examined, any partial
trailing accumulator is
+ * pushed too.
+ *
+ * <p>The result is a subset of {@code pushedFilterRanges}: rows not in
{@code
+ * pushedFilterRanges} were never read and are implicitly excluded.
+ */
+ private RowRanges evaluateStorageFilter(
+ PageReadStore keyPages,
+ RowRanges pushedFilterRanges) throws IOException {
+ ensureKeyScratchAllocated();
+ VectorizedColumnReader[] readers = new
VectorizedColumnReader[keyDescriptors.length];
+ for (int i = 0; i < readers.length; i++) {
+ readers[i] = new VectorizedColumnReader(
+ keyDescriptors[i], keyRequired[i], keyPages, convertTz,
datetimeRebaseMode,
+ datetimeRebaseTz, int96RebaseMode, int96RebaseTz, writerVersion);
+ }
+ ensureCurrentKeyAccumulatorsAllocated();
+
+ long keyRowsTotal = pushedFilterRanges.rowCount();
+ PrimitiveIterator.OfLong rowIndexIter = pushedFilterRanges.iterator();
+ RowRanges.Builder finalRangesBuilder = RowRanges.builder();
+ long remaining = keyRowsTotal;
+ while (remaining > 0) {
+ int num = (int) Math.min((long) capacity, remaining);
+ for (int i = 0; i < keyScratchVectors.length; i++) {
+ keyScratchVectors[i].reset();
+ readers[i].readBatch(num, keyScratchVectors[i], null, null);
+ }
+ keyScratchBatch.setNumRows(num);
+ for (int r = 0; r < num; r++) {
+ long blockRow = rowIndexIter.nextLong();
+ if (storageFilter.test(keyScratchBatch.getRow(r))) {
+ finalRangesBuilder.addSelectedRow(blockRow);
+ appendSurvivorRowToAccumulators(r);
+ }
+ }
+ remaining -= num;
+ }
+
+ finalizePartialAccumulators();
+
+ return finalRangesBuilder.build();
+ }
+
+ private void ensureKeyScratchAllocated() {
Review Comment:
Fixed: the array is assigned before the loop and `close()` frees the scratch
vectors by walking it element-wise, so `keyScratchBatch` is no longer what
reaches them.
The guard stays, and it cannot be reached in the state you describe: an
allocation failure here propagates out of the reader, the task fails and
`close()` runs, so nothing calls this again on a half-filled array.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -421,11 +552,313 @@ public boolean nextBatch() throws IOException {
return true;
}
+ /**
+ * Splicing emit path. Closes the previous emit's dequeued key vectors
(releasing survivor memory
+ * incrementally), advances to the next row group if needed via {@link
#checkEndOfRowGroup()},
+ * dequeues one survivor key vector per key column, drives non-key column
readers for {@code num}
+ * rows, and assembles a fresh {@link ColumnarBatch} interleaving key
(dequeued) and non-key
+ * (persistent value-vector) slots in the original projection order.
+ * The per-emit batch is a transient view over vectors owned elsewhere; see
{@link #close()} and
+ * {@link #closeSplicingState()}.
+ */
+ private boolean nextBatchSplicing() throws IOException {
+ if (pendingCloseKeyVectors != null) {
+ for (WritableColumnVector v : pendingCloseKeyVectors) {
+ if (v != null) v.close();
+ }
+ pendingCloseKeyVectors = null;
+ }
+ for (int i = 0; i < columnVectors.length; i++) {
+ if (isKeyTopLevel[i]) continue;
+ columnVectors[i].reset();
+ }
+ // Match the eager path and zero the outgoing batch before the terminal
checks below. Without
+ // this, a terminal call leaves `columnarBatch` pointing at the previous
emit whose key slots
+ // were just closed above -- with off-heap vectors those buffers are
already freed, so a
+ // consumer that read the batch after nextBatch() returned false would see
freed memory.
+ if (columnarBatch != null) columnarBatch.setNumRows(0);
+ if (hitEndOfData) return false;
+ if (rowsReturned >= totalRowCount) return false;
+ checkEndOfRowGroup();
+ if (hitEndOfData) return false;
+
+ int num = (int) Math.min(capacity, totalCountLoadedSoFar - rowsReturned);
+
+ WritableColumnVector[] dequeued = new
WritableColumnVector[keyVectorQueues.length];
+ for (int i = 0; i < keyVectorQueues.length; i++) {
+ dequeued[i] = keyVectorQueues[i].removeFirst();
+ }
+
+ for (int i = 0; i < columnVectors.length; i++) {
+ if (isKeyTopLevel[i]) continue;
+ ParquetColumnVector cv = columnVectors[i];
+ for (ParquetColumnVector leafCv : cv.getLeaves()) {
+ VectorizedColumnReader columnReader = leafCv.getColumnReader();
+ if (columnReader != null) {
+ columnReader.readBatch(num, leafCv.getValueVector(),
+ leafCv.getRepetitionLevelVector(),
leafCv.getDefinitionLevelVector());
+ }
+ }
+ cv.assemble();
+ }
+ if (rowIndexGenerator != null) {
+ // Row-index column is identified by name
(ROW_INDEX_TEMPORARY_COLUMN_NAME), which is a
+ // synthetic metadata column never referenced by a storage filter, so
its slot is a non-key
+ // slot with a persistent ParquetColumnVector.
+ rowIndexGenerator.populateRowIndex(columnVectors, num);
+ }
+
+ ColumnVector[] cols = new ColumnVector[persistentBatchColumns.length];
+ // This walks batch slots in ascending order while `keyIdx` walks the
survivor queues in
+ // key-row-position order, so it pairs the k-th smallest key slot with
key-row position k.
+ // That is only the identity because `ParquetStorageFilter.create` sorts
`keyColumnIndices`
+ // ascending -- see the comment there. `isKeyTopLevel` marks which slots
are keys but not their
+ // position in that list, so this loop cannot reconstruct the pairing on
its own: if the list
+ // ever stops being sorted, key columns silently swap places in the output
batch.
+ int keyIdx = 0;
+ for (int i = 0; i < persistentBatchColumns.length; i++) {
+ if (i < isKeyTopLevel.length && isKeyTopLevel[i]) {
+ cols[i] = dequeued[keyIdx++];
+ } else {
+ cols[i] = persistentBatchColumns[i];
+ }
+ }
+ columnarBatch = new ColumnarBatch(cols);
Review Comment:
Took the second option, since documenting an exception to a public method's
contract is the worse half of the trade. There is one `ColumnarBatch` for the
whole read again, and the emit path rewrites its key slots in place --
`ColumnarBatch` holds the column array by reference, its staging row included,
so writing a slot publishes it. Two allocations per batch go away with it, and
the hoist-once pattern works on this path.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormat.scala:
##########
@@ -165,6 +166,47 @@ trait FileFormat {
}
}
+ /**
+ * Like [[buildReaderWithPartitionValues]] but additionally accepts a
sequence of storage filters:
+ * Catalyst expressions that the storage layer may evaluate to drive
value-column IO pruning based
+ * on key-column evaluation (e.g., late materialization with a runtime bloom
filter).
+ *
+ * The default implementation delegates to
[[buildReaderWithPartitionValues]] and accepts no
+ * storage filter at all. File formats that support storage-filter pushdown
(e.g., Parquet) should
+ * override this method.
+ *
+ * A non-empty `storageFilters` here is a planner bug, so the default body
rejects it rather than
+ * dropping it. The planner removes an extracted conjunct from the post-scan
`Filter`, so a reader
+ * that ignores it returns rows the filter rejects. Every other layer of the
feature fails loudly
+ * for the same reason. Unreachable today, because
`FileSourceStrategy.extractStorageFilters` only
+ * extracts for `ParquetFileFormat` itself, but that keeps the invariant one
method away from the
+ * code that relies on it.
+ *
+ * Scalar subqueries inside `storageFilters` are expected to have been
materialized before this
+ * method is called, so that the returned reader can be safely serialized to
executors.
+ *
+ * `storageFilterMetrics` is an optional map of SQL metrics the reader can
update during execution
+ * (e.g. number of row groups skipped). The scan is expected to expose these
metrics via its
+ * `metrics` field so they show up in the SQL UI.
+ */
+ def buildReaderWithStorageFilters(
+ sparkSession: SparkSession,
+ dataSchema: StructType,
+ partitionSchema: StructType,
+ requiredSchema: StructType,
+ filters: Seq[Filter],
+ storageFilters: Seq[Expression],
+ options: Map[String, String],
+ hadoopConf: Configuration,
+ storageFilterMetrics: Map[String, SQLMetric] = Map.empty
+ ): PartitionedFile => Iterator[InternalRow] = {
+ require(storageFilters.isEmpty,
+ s"${getClass.getSimpleName} does not support storage-filter pushdown,
but was given " +
+ storageFilters.mkString("[", ", ", "]"))
+ buildReaderWithPartitionValues(
Review Comment:
Closed by routing both public builders through a private
`buildParquetReader`, and the scaladoc on `buildReaderWithStorageFilters` now
says an override must do the same, `super` included, since that call is virtual
too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
