peter-toth commented on code in PR #58895:
URL: https://github.com/apache/spark/pull/58895#discussion_r4097172647
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -492,6 +938,535 @@ 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 #nonKeyColumns} 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;
+ }
+ // Splicing buffers one key value per surviving row of the whole row
group before it can emit
+ // the first batch, and that buffer is outside any MemoryConsumer, so
phase 1 counts what it
+ // holds and gives up past the cap. This row group is then read the
plain way: phase 2 takes
+ // every projected column under finalRanges, key columns included, so
nothing is buffered and
+ // the cost is one extra read of the key columns.
+ spliceCurrentRowGroup = true;
+ splicedBytes = 0L;
+
+ // Phase 0: rows allowed by the pushed data filter, at column-index
granularity. The full
+ // requestedSchema goes back on first, because phases 1 and 2 narrow it
and
+ // ParquetFileReader.getRowRanges computes ranges against the reader's
current paths.
+ lateMatReader.setRequestedSchema(requestedColumns);
+ // getRowRanges checks only whether a filter is pushed, not
options.useColumnIndexFilter(),
+ // so calling it unconditionally would keep applying column-index
filtering after a user
+ // turned it off -- the documented escape hatch for files whose column
index is wrong.
+ // Trusting a wrong column index here drops rows for good, since
finalRanges is a subset of
+ // these ranges and the post-scan Filter no longer holds the predicate.
Phase 2 is
+ // unaffected: it selects pages through the offset index, which this
conf says nothing
+ // about.
+ RowRanges pushedFilterRanges = useColumnIndexFilter
+ ? lateMatReader.getRowRanges(blockIdx)
+ : RowRanges.createSingle(blockRowCount);
+ // RowRanges.rowCount() walks every range, so resolve each range set's
count once.
+ long baselineRows = pushedFilterRanges.rowCount();
+ if (baselineRows == 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;
+ }
+
+ // What this feature can avoid reading is the non-key columns of the
rows the storage filter
+ // rejects, so that is the baseline both byte metrics are measured
against: the non-key bytes
+ // a plain read of this projection would transfer for every row the
pushed filter kept. The
+ // null checks only skip work for a caller that drives this reader
without a scan's metrics;
+ // FileSourceScanLike creates all five whenever storageFilters is
non-empty.
+ // compressedBytesForRowRanges never does IO of its own.
+ StorageFilterMetrics m = storageFilter.metrics();
+ SQLMetric bytesAvoidedRg = m.bytesAvoidedByRowGroup();
+ SQLMetric bytesAvoidedPf = m.bytesAvoidedByPageFiltering();
+ boolean needBytes = bytesAvoidedRg != null || bytesAvoidedPf != null;
+ Map<ColumnPath, ColumnChunkMetaData> blockChunks =
+ needBytes ? chunksByPath(lateMatReader, blockIdx) : null;
+ long nonKeyBaselineBytes = needBytes
+ ? compressedBytesForRowRanges(lateMatReader, blockIdx, blockChunks,
nonKeyColumns,
+ pushedFilterRanges, baselineRows)
+ : 0L;
+
+ // Phase 1: switch to key-only schema, read key columns under
pushedFilterRanges, evaluate the
+ // storage filter per row.
+ lateMatReader.setRequestedSchema(keyOnlyColumns);
+ 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 UnsupportedFileReadException(
+ "No key pages for row group " + blockIdx + " despite " +
baselineRows
+ + " rows selected by the pushed filter");
+ }
+ RowRanges finalRanges = evaluateStorageFilter(keyPages,
pushedFilterRanges);
+ long finalRowCount = finalRanges.rowCount();
+
+ if (finalRowCount == 0) {
+ // Every surviving row was rejected by the storage filter; skip block
entirely, which avoids
+ // the whole non-key baseline. Phase 1 still paid to read the key
columns, and that cost is
+ // not part of the baseline, so nothing has to be subtracted from it
here.
+ 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(nonKeyBaselineBytes);
+ continue;
+ }
+
+ // Phase 2: switch to non-key schema, read non-key columns under
finalRanges. Skipped entirely
+ // when the projection is all keys (nonKeyColumns is null); emit
reconstructs each batch from
+ // the key queues alone.
+ long keptRows;
+ long phase2Bytes;
+ PageReadStore dataPages = null;
+ // An all-keys projection has nothing for phase 2 to read -- but only if
the key values were
+ // buffered. A row group past the cap has to read them here like any
other column.
+ if (nonKeyColumns == null && spliceCurrentRowGroup) {
+ keptRows = finalRowCount;
+ phase2Bytes = 0L;
+ } else {
+ lateMatReader.setRequestedSchema(
+ spliceCurrentRowGroup ? nonKeyColumns : requestedColumns);
+ // Reading a strict subset of a block's rows needs a Parquet offset
index, and parquet
+ // enforces that itself: it resolves every requested column's offset
index before reading
+ // anything, and a column without one makes its column index store
throw
+ // MissingOffsetIndexException. Files written before parquet-mr 1.11,
or by a writer that
+ // omits the page index (pyarrow's `write_table` defaults to
`write_page_index=False`),
+ // have none. Degrading to a whole-block read is not an option,
because the key vectors
+ // hold only the survivors and the batch would misalign, and neither
is dropping the
+ // predicate, which extraction has removed from the post-scan Filter.
So the read fails,
+ // and all this adds is what the user can do about it.
+ //
+ // Nothing is checked up front: a row group the filter keeps whole
never needs the index
+ // (`readFilteredRowGroup` falls back to a plain read when the ranges
cover the block), and
+ // one it rejects whole is never read at all, so a file with no page
index still scans as
+ // long as the filter never has to prune inside a row group.
+ try {
+ dataPages = lateMatReader.readFilteredRowGroup(blockIdx,
finalRanges);
+ } catch (MissingOffsetIndexException e) {
+ throw new UnsupportedFileReadException(String.format(
+ "Storage-filter pushdown needs a Parquet offset index to read
the %d of %d rows its "
+ + "filter kept in row group %d of %s, but the file was
written without a page "
+ + "index. Set %s=false to read this file.",
+ finalRowCount, blockRowCount, blockIdx, lateMatReader.getFile(),
+
SQLConf$.MODULE$.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED().key()), e);
+ }
+ 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 UnsupportedFileReadException(
+ "No data pages for row group " + blockIdx + " despite " +
finalRowCount
+ + " surviving rows");
+ }
+ keptRows = dataPages.getRowCount();
+ // `needBytes`, not just `bytesAvoidedPf != null`: it is what built
`blockChunks`.
+ if (needBytes && bytesAvoidedPf != null) {
+ phase2Bytes = compressedBytesForRowRanges(lateMatReader, blockIdx,
blockChunks,
+ nonKeyColumns, finalRanges, finalRowCount);
+ if (!spliceCurrentRowGroup) {
+ // This row group gave splicing up, so phase 2 read the key
columns a second time. The
+ // baseline counts them once, in phase 1, so the extra read is a
cost against it --
+ // which can make the row group's contribution negative, and that
is the truth about it.
+ phase2Bytes += compressedBytesForRowRanges(lateMatReader,
blockIdx, blockChunks,
+ keyOnlyColumns, finalRanges, finalRowCount);
+ }
+ } else {
+ phase2Bytes = 0L;
+ }
+ }
+ long filteredRows = baselineRows - keptRows;
+ SQLMetric rowsExcludedWithinRg = m.rowsExcludedWithinRowGroup();
+ if (rowsExcludedWithinRg != null && filteredRows > 0)
rowsExcludedWithinRg.add(filteredRows);
+ if (bytesAvoidedPf != null) {
+ bytesAvoidedPf.add(nonKeyBaselineBytes - phase2Bytes);
+ }
+
+ if (dataPages != null) {
+ if (rowIndexGenerator != null) {
+ rowIndexGenerator.initFromPageReadStore(dataPages);
+ }
+ for (int i = 0; i < columnVectors.length; i++) {
+ if (spliceCurrentRowGroup && 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;
+ }
+
+
+ /**
+ * The block's column chunks by path, built once per row group and shared by
the byte-metric calls
+ * that consume it, since {@link BlockMetaData} offers no lookup of its own.
+ */
+ private static Map<ColumnPath, ColumnChunkMetaData> chunksByPath(
+ ParquetFileReader reader, int blockIndex) {
+ Map<ColumnPath, ColumnChunkMetaData> chunks = new HashMap<>();
+ for (ColumnChunkMetaData chunk :
reader.getRowGroups().get(blockIndex).getColumns()) {
+ chunks.put(chunk.getPath(), chunk);
+ }
+ return chunks;
+ }
+
+ /**
+ * Compressed bytes the reader transfers for the given leaf {@code columns}
when it reads exactly
+ * {@code rowRanges} of the given block. Page headers and the dictionary
page are included, since
+ * both are read whenever any page of a chunk is read. {@code rowRangeCount}
is
+ * {@code rowRanges.rowCount()}, passed in because that walks every range
and the caller has it.
+ *
+ * <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.
+ */
+ private static long compressedBytesForRowRanges(
+ ParquetFileReader reader,
+ int blockIndex,
+ Map<ColumnPath, ColumnChunkMetaData> chunks,
+ List<ColumnDescriptor> columns,
+ RowRanges rowRanges,
+ long rowRangeCount) {
+ if (columns == null || columns.isEmpty() || rowRangeCount == 0) {
+ return 0L;
+ }
+ long blockRowCount = reader.getRowGroups().get(blockIndex).getRowCount();
+ boolean wholeBlock = rowRangeCount == blockRowCount;
+ ColumnIndexStore ciStore = wholeBlock ? null :
reader.getColumnIndexStore(blockIndex);
+ long total = 0L;
+ for (ColumnDescriptor column : columns) {
+ 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;
+ }
+
+ /**
+ * Evaluates the storage filter over every row of a key-only {@link
PageReadStore}, in
+ * capacity-sized chunks, and returns the surviving rows as {@link
RowRanges} in block-row
+ * coordinates. The result is a subset of {@code pushedFilterRanges}: rows
outside it were never
+ * read.
+ *
+ * <p>Each survivor's key values are appended to {@link
#currentKeyAccumulators} for the emit path
+ * to splice, until the buffer passes its cap. From there the row group is
evaluated without
+ * buffering and {@link #spliceCurrentRowGroup} is false, so its phase 2
reads the key columns
+ * again along with everything else.
+ */
+ 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();
+
+ PrimitiveIterator.OfLong rowIndexIter = pushedFilterRanges.iterator();
+ RowRanges.Builder finalRangesBuilder = RowRanges.builder();
+ // Recomputed rather than taken from the caller: a count that disagreed
with this iterator would
+ // silently drop surviving rows, and no post-scan Filter is left to catch
that.
+ long remaining = pushedFilterRanges.rowCount();
+ boolean accumulate = true;
+ 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);
+ if (accumulate) {
+ accumulate = appendSurvivorRowToAccumulators(r);
+ }
+ }
+ }
+ remaining -= num;
+ }
+
+ if (accumulate) {
+ finalizePartialAccumulators();
+ }
+
+ return finalRangesBuilder.build();
+ }
+
+ private void ensureKeyScratchAllocated() {
+ if (keyScratchVectors != null) return;
+ // Assigned before the loop on purpose: an allocation failure part way
through then leaves the
+ // vectors allocated so far reachable for `close()`, which walks this
array element-wise.
+ keyScratchVectors = new WritableColumnVector[keyDescriptors.length];
+ boolean useOffHeap = MEMORY_MODE == MemoryMode.OFF_HEAP;
+ int[] keyIndices = storageFilter.keyColumnIndices();
+ for (int i = 0; i < keyDescriptors.length; i++) {
+ DataType dt = sparkRequestedSchema.fields()[keyIndices[i]].dataType();
+ keyScratchVectors[i] = useOffHeap
+ ? new OffHeapColumnVector(capacity, dt)
+ : new OnHeapColumnVector(capacity, dt);
+ }
+ keyScratchBatch = new ColumnarBatch(keyScratchVectors);
+ }
+
+ /**
+ * Allocates any accumulator slot left null by the last push to the queues,
{@link #capacity} rows
+ * each.
+ */
+ private void ensureCurrentKeyAccumulatorsAllocated() {
+ boolean useOffHeap = MEMORY_MODE == MemoryMode.OFF_HEAP;
+ int[] keyIndices = storageFilter.keyColumnIndices();
+ for (int i = 0; i < currentKeyAccumulators.length; i++) {
+ if (currentKeyAccumulators[i] == null) {
+ DataType dt = sparkRequestedSchema.fields()[keyIndices[i]].dataType();
+ currentKeyAccumulators[i] = useOffHeap
+ ? new OffHeapColumnVector(capacity, dt)
+ : new OnHeapColumnVector(capacity, dt);
+ }
+ }
+ currentKeyAccumulatorRowCount = 0;
+ }
+
+ /**
+ * Appends row {@code srcRow} of every key column to the accumulators,
pushing them onto their
+ * queues once full. All key columns advance in lockstep, which is what
keeps the queues aligned.
+ *
+ * <p>Returns false when the buffered survivors have passed
+ * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}
and this row group has
+ * given splicing up, in which case it has already released what it held and
the caller must stop
+ * calling this.
+ */
+ private boolean appendSurvivorRowToAccumulators(int srcRow) {
+ final int dstRow = currentKeyAccumulatorRowCount;
+ final WritableColumnVector[] accs = currentKeyAccumulators;
+ final WritableColumnVector[] srcs = keyScratchVectors;
+ final ValueCopier[] copiers = keyCopiers;
+ long valueBytes = 0L;
+ for (int i = 0, n = accs.length; i < n; i++) {
+ WritableColumnVector src = srcs[i];
+ WritableColumnVector dst = accs[i];
+ if (src.isNullAt(srcRow)) {
+ dst.putNull(dstRow);
+ } else {
+ copiers[i].copy(dst, dstRow, src, srcRow);
+ // Measured on the destination: a dictionary-encoded source has no
length of its own, since
+ // its values are read through the dictionary.
+ if (keyVariableLength[i]) valueBytes += dst.getArrayLength(dstRow);
+ }
+ }
+ splicedBytes += keyFixedBytesPerRow + valueBytes;
+ currentKeyAccumulatorRowCount = dstRow + 1;
+ if (currentKeyAccumulatorRowCount == capacity) {
+ for (int i = 0; i < currentKeyAccumulators.length; i++) {
+ keyVectorQueues[i].addLast(currentKeyAccumulators[i]);
+ currentKeyAccumulators[i] = null;
+ }
+ // Only here, so the cost is one comparison per capacity-sized vector
rather than per row. A
+ // row group whose survivors fit in a single accumulator is never
checked at all: it then
+ // holds one capacity-sized vector per key column, which is what a plain
read holds anyway.
+ if (splicedBytes > storageFilter.maxSplicedRowGroupBytes()) {
+ abandonSplicing();
+ return false;
+ }
+ ensureCurrentKeyAccumulatorsAllocated();
+ }
+ return true;
+ }
+
+ /**
+ * Gives up splicing for the row group being evaluated and releases every
survivor vector it has
+ * buffered. Phase 2 then reads the full projected schema and the emit path
takes the persistent
+ * batch, so the rows are unaffected.
+ */
+ private void abandonSplicing() {
+ for (java.util.ArrayDeque<WritableColumnVector> q : keyVectorQueues) {
+ for (WritableColumnVector v : q) v.close();
+ q.clear();
+ }
+ closeAll(currentKeyAccumulators);
+ Arrays.fill(currentKeyAccumulators, null);
+ currentKeyAccumulatorRowCount = 0;
+ spliceCurrentRowGroup = false;
+ }
+
+ /**
+ * Pushes any partially-filled accumulator into its queue at row-group end
so the emit path can
+ * dequeue it as the row group's final batch.
+ */
+ private void finalizePartialAccumulators() {
+ if (currentKeyAccumulatorRowCount == 0) return;
+ for (int i = 0; i < currentKeyAccumulators.length; i++) {
+ keyVectorQueues[i].addLast(currentKeyAccumulators[i]);
+ currentKeyAccumulators[i] = null;
+ }
+ currentKeyAccumulatorRowCount = 0;
+ }
+
+ /**
+ * Closes anything held by the splicing path: every vector still queued, the
published head
+ * included, and partially-filled accumulators. Called from {@link #close()}.
+ */
+ private void closeSplicingState() {
+ keyVectorsPublished = false;
+ if (keyVectorQueues != null) {
+ for (java.util.ArrayDeque<WritableColumnVector> q : keyVectorQueues) {
+ if (q != null) {
+ for (WritableColumnVector v : q) v.close();
+ q.clear();
+ }
+ }
+ }
+ closeAll(currentKeyAccumulators);
+ currentKeyAccumulators = null;
+ }
+
+ /** Closes every non-null vector of {@code vectors}; tolerates a null array.
*/
+ private static void closeAll(WritableColumnVector[] vectors) {
+ if (vectors == null) return;
+ for (WritableColumnVector v : vectors) {
+ if (v != null) v.close();
+ }
+ }
+
+ /**
+ * Copies one key value between column vectors. Picked per key column at
init time by
+ * {@link #copierFor(DataType)}; the caller handles null sources.
+ */
+ @FunctionalInterface
+ private interface ValueCopier {
+ void copy(WritableColumnVector dst, int dstRow, WritableColumnVector src,
int srcRow);
+ }
+
+ /**
+ * Returns a {@link ValueCopier} for the given key {@link DataType}. The set
of types handled here
+ * is the definition behind {@code ParquetStorageFilter.isSupportedKeyType},
which gates both
+ * planning-time extraction and {@code ParquetStorageFilter.create} -- so
the throw at the end is
+ * unreachable. Teach both sides at once when adding a type; a type admitted
there but missing
+ * here becomes a task failure instead of a planning-time rejection.
+ */
+ private static ValueCopier copierFor(DataType dt) {
Review Comment:
Kept as it is, for the reason you name as the condition. It is a
vector-to-vector copy with no row in between, which is the whole point of it in
that loop, and it only has to serve the types `isSupportedKeyType` admits.
The two lists cannot drift silently: a test asserts that
`isSupportedKeyType` covers exactly the types `copierFor` handles, and the
javadoc says to teach both sides at once.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -492,6 +938,535 @@ 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 #nonKeyColumns} 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;
+ }
+ // Splicing buffers one key value per surviving row of the whole row
group before it can emit
+ // the first batch, and that buffer is outside any MemoryConsumer, so
phase 1 counts what it
+ // holds and gives up past the cap. This row group is then read the
plain way: phase 2 takes
+ // every projected column under finalRanges, key columns included, so
nothing is buffered and
+ // the cost is one extra read of the key columns.
+ spliceCurrentRowGroup = true;
+ splicedBytes = 0L;
+
+ // Phase 0: rows allowed by the pushed data filter, at column-index
granularity. The full
+ // requestedSchema goes back on first, because phases 1 and 2 narrow it
and
+ // ParquetFileReader.getRowRanges computes ranges against the reader's
current paths.
+ lateMatReader.setRequestedSchema(requestedColumns);
+ // getRowRanges checks only whether a filter is pushed, not
options.useColumnIndexFilter(),
+ // so calling it unconditionally would keep applying column-index
filtering after a user
+ // turned it off -- the documented escape hatch for files whose column
index is wrong.
+ // Trusting a wrong column index here drops rows for good, since
finalRanges is a subset of
+ // these ranges and the post-scan Filter no longer holds the predicate.
Phase 2 is
+ // unaffected: it selects pages through the offset index, which this
conf says nothing
+ // about.
+ RowRanges pushedFilterRanges = useColumnIndexFilter
+ ? lateMatReader.getRowRanges(blockIdx)
+ : RowRanges.createSingle(blockRowCount);
+ // RowRanges.rowCount() walks every range, so resolve each range set's
count once.
+ long baselineRows = pushedFilterRanges.rowCount();
+ if (baselineRows == 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;
+ }
+
+ // What this feature can avoid reading is the non-key columns of the
rows the storage filter
+ // rejects, so that is the baseline both byte metrics are measured
against: the non-key bytes
+ // a plain read of this projection would transfer for every row the
pushed filter kept. The
+ // null checks only skip work for a caller that drives this reader
without a scan's metrics;
+ // FileSourceScanLike creates all five whenever storageFilters is
non-empty.
+ // compressedBytesForRowRanges never does IO of its own.
+ StorageFilterMetrics m = storageFilter.metrics();
+ SQLMetric bytesAvoidedRg = m.bytesAvoidedByRowGroup();
+ SQLMetric bytesAvoidedPf = m.bytesAvoidedByPageFiltering();
+ boolean needBytes = bytesAvoidedRg != null || bytesAvoidedPf != null;
+ Map<ColumnPath, ColumnChunkMetaData> blockChunks =
+ needBytes ? chunksByPath(lateMatReader, blockIdx) : null;
+ long nonKeyBaselineBytes = needBytes
+ ? compressedBytesForRowRanges(lateMatReader, blockIdx, blockChunks,
nonKeyColumns,
+ pushedFilterRanges, baselineRows)
+ : 0L;
+
+ // Phase 1: switch to key-only schema, read key columns under
pushedFilterRanges, evaluate the
+ // storage filter per row.
+ lateMatReader.setRequestedSchema(keyOnlyColumns);
+ 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 UnsupportedFileReadException(
+ "No key pages for row group " + blockIdx + " despite " +
baselineRows
+ + " rows selected by the pushed filter");
+ }
+ RowRanges finalRanges = evaluateStorageFilter(keyPages,
pushedFilterRanges);
+ long finalRowCount = finalRanges.rowCount();
+
+ if (finalRowCount == 0) {
+ // Every surviving row was rejected by the storage filter; skip block
entirely, which avoids
+ // the whole non-key baseline. Phase 1 still paid to read the key
columns, and that cost is
+ // not part of the baseline, so nothing has to be subtracted from it
here.
+ 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(nonKeyBaselineBytes);
+ continue;
+ }
+
+ // Phase 2: switch to non-key schema, read non-key columns under
finalRanges. Skipped entirely
+ // when the projection is all keys (nonKeyColumns is null); emit
reconstructs each batch from
+ // the key queues alone.
+ long keptRows;
+ long phase2Bytes;
+ PageReadStore dataPages = null;
+ // An all-keys projection has nothing for phase 2 to read -- but only if
the key values were
+ // buffered. A row group past the cap has to read them here like any
other column.
+ if (nonKeyColumns == null && spliceCurrentRowGroup) {
+ keptRows = finalRowCount;
+ phase2Bytes = 0L;
+ } else {
+ lateMatReader.setRequestedSchema(
+ spliceCurrentRowGroup ? nonKeyColumns : requestedColumns);
+ // Reading a strict subset of a block's rows needs a Parquet offset
index, and parquet
+ // enforces that itself: it resolves every requested column's offset
index before reading
+ // anything, and a column without one makes its column index store
throw
+ // MissingOffsetIndexException. Files written before parquet-mr 1.11,
or by a writer that
+ // omits the page index (pyarrow's `write_table` defaults to
`write_page_index=False`),
+ // have none. Degrading to a whole-block read is not an option,
because the key vectors
+ // hold only the survivors and the batch would misalign, and neither
is dropping the
+ // predicate, which extraction has removed from the post-scan Filter.
So the read fails,
+ // and all this adds is what the user can do about it.
+ //
+ // Nothing is checked up front: a row group the filter keeps whole
never needs the index
+ // (`readFilteredRowGroup` falls back to a plain read when the ranges
cover the block), and
+ // one it rejects whole is never read at all, so a file with no page
index still scans as
+ // long as the filter never has to prune inside a row group.
+ try {
+ dataPages = lateMatReader.readFilteredRowGroup(blockIdx,
finalRanges);
+ } catch (MissingOffsetIndexException e) {
+ throw new UnsupportedFileReadException(String.format(
+ "Storage-filter pushdown needs a Parquet offset index to read
the %d of %d rows its "
+ + "filter kept in row group %d of %s, but the file was
written without a page "
+ + "index. Set %s=false to read this file.",
+ finalRowCount, blockRowCount, blockIdx, lateMatReader.getFile(),
+
SQLConf$.MODULE$.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED().key()), e);
+ }
+ 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 UnsupportedFileReadException(
+ "No data pages for row group " + blockIdx + " despite " +
finalRowCount
+ + " surviving rows");
+ }
+ keptRows = dataPages.getRowCount();
+ // `needBytes`, not just `bytesAvoidedPf != null`: it is what built
`blockChunks`.
+ if (needBytes && bytesAvoidedPf != null) {
+ phase2Bytes = compressedBytesForRowRanges(lateMatReader, blockIdx,
blockChunks,
+ nonKeyColumns, finalRanges, finalRowCount);
+ if (!spliceCurrentRowGroup) {
+ // This row group gave splicing up, so phase 2 read the key
columns a second time. The
+ // baseline counts them once, in phase 1, so the extra read is a
cost against it --
+ // which can make the row group's contribution negative, and that
is the truth about it.
+ phase2Bytes += compressedBytesForRowRanges(lateMatReader,
blockIdx, blockChunks,
+ keyOnlyColumns, finalRanges, finalRowCount);
+ }
+ } else {
+ phase2Bytes = 0L;
+ }
+ }
+ long filteredRows = baselineRows - keptRows;
+ SQLMetric rowsExcludedWithinRg = m.rowsExcludedWithinRowGroup();
+ if (rowsExcludedWithinRg != null && filteredRows > 0)
rowsExcludedWithinRg.add(filteredRows);
+ if (bytesAvoidedPf != null) {
+ bytesAvoidedPf.add(nonKeyBaselineBytes - phase2Bytes);
+ }
+
+ if (dataPages != null) {
+ if (rowIndexGenerator != null) {
+ rowIndexGenerator.initFromPageReadStore(dataPages);
+ }
+ for (int i = 0; i < columnVectors.length; i++) {
+ if (spliceCurrentRowGroup && 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;
+ }
+
+
+ /**
+ * The block's column chunks by path, built once per row group and shared by
the byte-metric calls
+ * that consume it, since {@link BlockMetaData} offers no lookup of its own.
+ */
+ private static Map<ColumnPath, ColumnChunkMetaData> chunksByPath(
+ ParquetFileReader reader, int blockIndex) {
+ Map<ColumnPath, ColumnChunkMetaData> chunks = new HashMap<>();
+ for (ColumnChunkMetaData chunk :
reader.getRowGroups().get(blockIndex).getColumns()) {
+ chunks.put(chunk.getPath(), chunk);
+ }
+ return chunks;
+ }
+
+ /**
+ * Compressed bytes the reader transfers for the given leaf {@code columns}
when it reads exactly
+ * {@code rowRanges} of the given block. Page headers and the dictionary
page are included, since
+ * both are read whenever any page of a chunk is read. {@code rowRangeCount}
is
+ * {@code rowRanges.rowCount()}, passed in because that walks every range
and the caller has it.
+ *
+ * <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.
+ */
+ private static long compressedBytesForRowRanges(
+ ParquetFileReader reader,
+ int blockIndex,
+ Map<ColumnPath, ColumnChunkMetaData> chunks,
+ List<ColumnDescriptor> columns,
+ RowRanges rowRanges,
+ long rowRangeCount) {
+ if (columns == null || columns.isEmpty() || rowRangeCount == 0) {
+ return 0L;
+ }
+ long blockRowCount = reader.getRowGroups().get(blockIndex).getRowCount();
+ boolean wholeBlock = rowRangeCount == blockRowCount;
+ ColumnIndexStore ciStore = wholeBlock ? null :
reader.getColumnIndexStore(blockIndex);
+ long total = 0L;
+ for (ColumnDescriptor column : columns) {
+ 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;
+ }
+
+ /**
+ * Evaluates the storage filter over every row of a key-only {@link
PageReadStore}, in
+ * capacity-sized chunks, and returns the surviving rows as {@link
RowRanges} in block-row
+ * coordinates. The result is a subset of {@code pushedFilterRanges}: rows
outside it were never
+ * read.
+ *
+ * <p>Each survivor's key values are appended to {@link
#currentKeyAccumulators} for the emit path
+ * to splice, until the buffer passes its cap. From there the row group is
evaluated without
+ * buffering and {@link #spliceCurrentRowGroup} is false, so its phase 2
reads the key columns
+ * again along with everything else.
+ */
+ 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();
+
+ PrimitiveIterator.OfLong rowIndexIter = pushedFilterRanges.iterator();
+ RowRanges.Builder finalRangesBuilder = RowRanges.builder();
+ // Recomputed rather than taken from the caller: a count that disagreed
with this iterator would
+ // silently drop surviving rows, and no post-scan Filter is left to catch
that.
+ long remaining = pushedFilterRanges.rowCount();
+ boolean accumulate = true;
+ 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);
+ if (accumulate) {
+ accumulate = appendSurvivorRowToAccumulators(r);
+ }
+ }
+ }
+ remaining -= num;
+ }
+
+ if (accumulate) {
+ finalizePartialAccumulators();
+ }
+
+ return finalRangesBuilder.build();
+ }
+
+ private void ensureKeyScratchAllocated() {
+ if (keyScratchVectors != null) return;
+ // Assigned before the loop on purpose: an allocation failure part way
through then leaves the
+ // vectors allocated so far reachable for `close()`, which walks this
array element-wise.
+ keyScratchVectors = new WritableColumnVector[keyDescriptors.length];
+ boolean useOffHeap = MEMORY_MODE == MemoryMode.OFF_HEAP;
+ int[] keyIndices = storageFilter.keyColumnIndices();
+ for (int i = 0; i < keyDescriptors.length; i++) {
+ DataType dt = sparkRequestedSchema.fields()[keyIndices[i]].dataType();
+ keyScratchVectors[i] = useOffHeap
+ ? new OffHeapColumnVector(capacity, dt)
+ : new OnHeapColumnVector(capacity, dt);
+ }
+ keyScratchBatch = new ColumnarBatch(keyScratchVectors);
+ }
+
+ /**
+ * Allocates any accumulator slot left null by the last push to the queues,
{@link #capacity} rows
+ * each.
+ */
+ private void ensureCurrentKeyAccumulatorsAllocated() {
+ boolean useOffHeap = MEMORY_MODE == MemoryMode.OFF_HEAP;
+ int[] keyIndices = storageFilter.keyColumnIndices();
+ for (int i = 0; i < currentKeyAccumulators.length; i++) {
+ if (currentKeyAccumulators[i] == null) {
+ DataType dt = sparkRequestedSchema.fields()[keyIndices[i]].dataType();
+ currentKeyAccumulators[i] = useOffHeap
+ ? new OffHeapColumnVector(capacity, dt)
+ : new OnHeapColumnVector(capacity, dt);
+ }
+ }
+ currentKeyAccumulatorRowCount = 0;
+ }
+
+ /**
+ * Appends row {@code srcRow} of every key column to the accumulators,
pushing them onto their
+ * queues once full. All key columns advance in lockstep, which is what
keeps the queues aligned.
+ *
+ * <p>Returns false when the buffered survivors have passed
+ * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}
and this row group has
+ * given splicing up, in which case it has already released what it held and
the caller must stop
+ * calling this.
+ */
+ private boolean appendSurvivorRowToAccumulators(int srcRow) {
+ final int dstRow = currentKeyAccumulatorRowCount;
+ final WritableColumnVector[] accs = currentKeyAccumulators;
+ final WritableColumnVector[] srcs = keyScratchVectors;
+ final ValueCopier[] copiers = keyCopiers;
+ long valueBytes = 0L;
+ for (int i = 0, n = accs.length; i < n; i++) {
+ WritableColumnVector src = srcs[i];
+ WritableColumnVector dst = accs[i];
+ if (src.isNullAt(srcRow)) {
+ dst.putNull(dstRow);
+ } else {
+ copiers[i].copy(dst, dstRow, src, srcRow);
+ // Measured on the destination: a dictionary-encoded source has no
length of its own, since
+ // its values are read through the dictionary.
+ if (keyVariableLength[i]) valueBytes += dst.getArrayLength(dstRow);
+ }
+ }
+ splicedBytes += keyFixedBytesPerRow + valueBytes;
+ currentKeyAccumulatorRowCount = dstRow + 1;
+ if (currentKeyAccumulatorRowCount == capacity) {
+ for (int i = 0; i < currentKeyAccumulators.length; i++) {
+ keyVectorQueues[i].addLast(currentKeyAccumulators[i]);
+ currentKeyAccumulators[i] = null;
+ }
+ // Only here, so the cost is one comparison per capacity-sized vector
rather than per row. A
+ // row group whose survivors fit in a single accumulator is never
checked at all: it then
+ // holds one capacity-sized vector per key column, which is what a plain
read holds anyway.
+ if (splicedBytes > storageFilter.maxSplicedRowGroupBytes()) {
+ abandonSplicing();
+ return false;
+ }
+ ensureCurrentKeyAccumulatorsAllocated();
+ }
+ return true;
+ }
+
+ /**
+ * Gives up splicing for the row group being evaluated and releases every
survivor vector it has
+ * buffered. Phase 2 then reads the full projected schema and the emit path
takes the persistent
+ * batch, so the rows are unaffected.
+ */
+ private void abandonSplicing() {
+ for (java.util.ArrayDeque<WritableColumnVector> q : keyVectorQueues) {
+ for (WritableColumnVector v : q) v.close();
+ q.clear();
+ }
+ closeAll(currentKeyAccumulators);
+ Arrays.fill(currentKeyAccumulators, null);
+ currentKeyAccumulatorRowCount = 0;
+ spliceCurrentRowGroup = false;
+ }
+
+ /**
+ * Pushes any partially-filled accumulator into its queue at row-group end
so the emit path can
+ * dequeue it as the row group's final batch.
+ */
+ private void finalizePartialAccumulators() {
+ if (currentKeyAccumulatorRowCount == 0) return;
+ for (int i = 0; i < currentKeyAccumulators.length; i++) {
+ keyVectorQueues[i].addLast(currentKeyAccumulators[i]);
+ currentKeyAccumulators[i] = null;
+ }
+ currentKeyAccumulatorRowCount = 0;
+ }
+
+ /**
+ * Closes anything held by the splicing path: every vector still queued, the
published head
+ * included, and partially-filled accumulators. Called from {@link #close()}.
+ */
+ private void closeSplicingState() {
+ keyVectorsPublished = false;
+ if (keyVectorQueues != null) {
+ for (java.util.ArrayDeque<WritableColumnVector> q : keyVectorQueues) {
+ if (q != null) {
+ for (WritableColumnVector v : q) v.close();
+ q.clear();
+ }
+ }
+ }
+ closeAll(currentKeyAccumulators);
+ currentKeyAccumulators = null;
+ }
+
+ /** Closes every non-null vector of {@code vectors}; tolerates a null array.
*/
+ private static void closeAll(WritableColumnVector[] vectors) {
+ if (vectors == null) return;
+ for (WritableColumnVector v : vectors) {
+ if (v != null) v.close();
+ }
+ }
+
+ /**
+ * Copies one key value between column vectors. Picked per key column at
init time by
+ * {@link #copierFor(DataType)}; the caller handles null sources.
+ */
+ @FunctionalInterface
+ private interface ValueCopier {
+ void copy(WritableColumnVector dst, int dstRow, WritableColumnVector src,
int srcRow);
+ }
+
+ /**
+ * Returns a {@link ValueCopier} for the given key {@link DataType}. The set
of types handled here
+ * is the definition behind {@code ParquetStorageFilter.isSupportedKeyType},
which gates both
+ * planning-time extraction and {@code ParquetStorageFilter.create} -- so
the throw at the end is
+ * unreachable. Teach both sides at once when adding a type; a type admitted
there but missing
+ * here becomes a task failure instead of a planning-time rejection.
+ */
+ private static ValueCopier copierFor(DataType dt) {
+ if (dt instanceof BooleanType) {
+ return (dst, dRow, src, sRow) -> dst.putBoolean(dRow,
src.getBoolean(sRow));
+ }
+ if (dt instanceof ByteType) {
+ return (dst, dRow, src, sRow) -> dst.putByte(dRow, src.getByte(sRow));
+ }
+ if (dt instanceof ShortType) {
+ return (dst, dRow, src, sRow) -> dst.putShort(dRow, src.getShort(sRow));
+ }
+ if (dt instanceof IntegerType
+ || dt instanceof DateType
+ || dt instanceof YearMonthIntervalType) {
+ return (dst, dRow, src, sRow) -> dst.putInt(dRow, src.getInt(sRow));
+ }
+ if (dt instanceof LongType
+ || dt instanceof TimestampType
+ || dt instanceof TimestampNTZType
+ || dt instanceof TimeType
+ || dt instanceof DayTimeIntervalType) {
+ return (dst, dRow, src, sRow) -> dst.putLong(dRow, src.getLong(sRow));
+ }
+ if (dt instanceof FloatType) {
+ return (dst, dRow, src, sRow) -> dst.putFloat(dRow, src.getFloat(sRow));
+ }
+ if (dt instanceof DoubleType) {
+ return (dst, dRow, src, sRow) -> dst.putDouble(dRow,
src.getDouble(sRow));
+ }
+ if (dt instanceof DecimalType decimalType) {
+ int precision = decimalType.precision();
+ if (precision <= Decimal.MAX_INT_DIGITS()) {
+ return (dst, dRow, src, sRow) -> dst.putInt(dRow, src.getInt(sRow));
+ }
+ if (precision <= Decimal.MAX_LONG_DIGITS()) {
+ return (dst, dRow, src, sRow) -> dst.putLong(dRow, src.getLong(sRow));
+ }
+ return (dst, dRow, src, sRow) -> dst.putByteArray(dRow,
src.getBinary(sRow));
+ }
+ // StringType covers CHAR and VARCHAR: both extend it.
+ if (dt instanceof StringType || dt instanceof BinaryType) {
+ return (dst, dRow, src, sRow) -> dst.putByteArray(dRow,
src.getBinary(sRow));
Review Comment:
Follow-up in the description. The obvious shortcut does not work today:
`OffHeapColumnVector.getByteBuffer` allocates a `byte[]` of its own, a
dictionary-encoded source has to decode, and the two vector implementations
disagree on whether the buffer position is absolute. Removing the copy needs
the core column-vector classes, so it belongs in its own change.
Your point about plain-encoded high-cardinality keys is exactly right. The
key-type tests now cover both encodings for real, with an assertion on what
parquet actually wrote, which turned out to matter: the previous dictionary arm
was writing PLAIN files.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilter.scala:
##########
@@ -0,0 +1,217 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.parquet
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{And, BasePredicate,
BloomFilterMightContain, BoundReference, Expression, Literal, Predicate}
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType,
DataType, DateType, DayTimeIntervalType, DecimalType, DoubleType, FloatType,
IntegerType, LongType, ShortType, StringType, StructType, TimestampNTZType,
TimestampType, TimeType, YearMonthIntervalType}
+
+/**
+ * Optional SQL metrics the reader updates while applying a
[[ParquetStorageFilter]]. Every
+ * counter is scoped to what the storage filter added on top of a read of the
same projection
+ * without one. All fields are nullable; a null field disables that metric.
+ *
+ * - [[rowGroupsSkipped]] counts row groups whose data columns were never
read.
+ * - [[rowsExcludedByRowGroup]] sums the rows those skips excluded, per
skipped block the rows that
+ * survived the pushed data filter.
+ * - [[rowsExcludedWithinRowGroup]] sums rows excluded inside row groups that
were kept.
+ * - [[bytesAvoidedByRowGroup]] sums, per skipped row group, the non-key
bytes a plain read of this
+ * projection would have transferred for the rows that survived the pushed
data filter. Phase 1
+ * reads the key columns of every block, so key bytes are never part of it,
and it is zero on an
+ * all-keys projection, which can avoid nothing.
+ * - [[bytesAvoidedByPageFiltering]] sums, per kept row group, that same
non-key baseline minus the
+ * bytes phase 2 read, which is what `finalRanges` page selection pruned.
+ *
+ * The row counters' suffix says where a row was excluded, not what would have
saved it: a row
+ * inside a kept row group is read as part of its page and dropped during
decode, so page
+ * filtering did not save it. An all-keys projection has no page filtering at
all, and
+ * [[rowsExcludedWithinRowGroup]] still counts every row the filter dropped.
+ */
+case class StorageFilterMetrics(
+ rowGroupsSkipped: SQLMetric = null,
+ rowsExcludedByRowGroup: SQLMetric = null,
+ rowsExcludedWithinRowGroup: SQLMetric = null,
+ bytesAvoidedByRowGroup: SQLMetric = null,
+ bytesAvoidedByPageFiltering: SQLMetric = null)
+
+/**
+ * A runtime filter that the vectorized Parquet reader uses to drive late
materialization: read
+ * key-column pages first, evaluate this filter per row to decide which rows
survive, and skip
+ * data-column pages that do not overlap any surviving row range.
+ *
+ * [[keyColumnIndices]] are indices into the scan's requested data schema
identifying the leaf
+ * columns referenced by the filter. [[boundExpression]] has its references
rewritten to
+ * [[BoundReference]]s pointing at positions 0..(keyColumnIndices.length - 1);
the reader must
+ * evaluate it against rows whose fields correspond to those key columns in
that order.
+ */
+class ParquetStorageFilter private (
+ val keyColumnIndices: Array[Int],
+ val boundExpression: Expression,
+ val metrics: StorageFilterMetrics,
+ val maxSplicedRowGroupBytes: Long) extends Serializable {
+
+ // Codegen-produced predicates can be awkward to serialize from driver to
executor, so we defer
+ // construction to first use on the executor.
+ @transient private lazy val predicate: BasePredicate =
Predicate.create(boundExpression)
+
+ def test(keyRow: InternalRow): Boolean = predicate.eval(keyRow)
+
+ /**
+ * Returns a new filter for a physical file that is missing some key columns
(schema evolution).
+ * The [[BoundReference]]s at `missingKeyLocalPositions`, which are indices
into
+ * [[keyColumnIndices]], are replaced by `missingKeyValues`, and the
remaining references are
+ * renumbered onto the reduced key-row layout. [[keyColumnIndices]] keeps
the present columns in
+ * their original relative order, and SQL metrics are shared with `this`.
+ *
+ * `missingKeyValues(i)` must be the internal-format value the reader
produces for a missing
+ * column: its existence DEFAULT when it has one, else null.
`ParquetColumnVector` writes that
+ * default into the output vector, so substituting null instead would filter
on a value the scan
+ * never returns and could drop matching rows.
+ *
+ * The predicate has to be evaluated against the substitution rather than
skipped, because a null
+ * key does not always mean `false`: a `Coalesce`-wrapped reference still
produces a non-null
+ * result, and `XxHash64` is `nullable = false` and hashes a null input to
its seed.
+ *
+ * With every key position missing, the result holds no [[BoundReference]]
at all and
+ * [[evalAllMissing]] can read off its constant truth value.
+ */
+ def rewriteForMissingKeys(
+ missingKeyLocalPositions: Array[Int],
+ missingKeyValues: Array[Any]): ParquetStorageFilter = {
+ require(missingKeyLocalPositions.length == missingKeyValues.length,
+ "missingKeyLocalPositions and missingKeyValues must have the same
length")
+ val substitution = missingKeyLocalPositions.zip(missingKeyValues).toMap
+ val presentPositions =
keyColumnIndices.indices.filterNot(substitution.contains)
+ val newPosOf = presentPositions.zipWithIndex.toMap
+ val rewritten = boundExpression.transform {
Review Comment:
Fixed: the rewrite is memoized per set of missing key positions for the
task, so the bloom is deserialized once per scan rather than once per file. The
substituted values are a function of those positions, since they come from the
scan's schema, which is why the positions alone are a sound key. c255f88
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -409,23 +598,276 @@ public boolean nextBatch() throws IOException {
}
cv.assemble();
}
- // If needed, compute row indexes within a file.
- if (rowIndexGenerator != null) {
- rowIndexGenerator.populateRowIndex(columnVectors, num);
- }
+ }
- rowsReturned += num;
+ /** Publishes {@code num} rows as the current batch. */
+ private void finishBatch(int num) {
columnarBatch.setNumRows(num);
+ rowsReturned += num;
numBatched = num;
batchIdx = 0;
+ }
+
+ /**
+ * Splicing emit path. Publishes one survivor key vector per key column into
the batch, reads the
+ * non-key slots for {@code num} rows, and hands out the same {@link
ColumnarBatch} every time,
+ * its key slots rewritten in place. The batch is a view over vectors owned
elsewhere; see
+ * {@link #close()} and {@link #closeSplicingState()}.
+ */
+ private boolean nextBatchSplicing() throws IOException {
+ releasePublishedKeyVectors();
+ for (ParquetColumnVector cv : columnVectors) {
+ cv.reset();
+ }
+ // Zero the outgoing batch before the terminal checks below, as the plain
path does. Otherwise a
+ // terminal call leaves the batch pointing at key vectors that were just
closed -- off-heap,
+ // that is 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);
+
+ if (!spliceCurrentRowGroup) {
+ // This row group was read the plain way: phase 2 took every projected
column, so every slot
+ // of the batch comes from the persistent one.
+ readPersistentColumns(num, /* skipKeySlots= */ false);
+ if (rowIndexGenerator != null) {
+ rowIndexGenerator.populateRowIndex(columnVectors, num);
+ }
+ System.arraycopy(persistentBatchColumns, 0, spliceBatchColumns, 0,
spliceBatchColumns.length);
+ finishBatch(num);
+ return true;
+ }
+
+ for (int i = 0; i < keyVectorQueues.length; i++) {
+ if (keyVectorQueues[i].isEmpty()) {
+ // Unreachable: the queues hold exactly the survivors phase 1
accumulated, and the emit loop
+ // is driven by that same count. Named rather than left to
NoSuchElementException because
+ // this one must not be swallowed: under ignoreCorruptFiles that would
silently drop the
+ // rest of a healthy file's rows.
+ throw new UnsupportedFileReadException(String.format(
+ "Storage-filter survivor queue %d of row group %d in %s ran out
with %d rows still to "
+ + "emit", i, nextBlockIndex - 1, lateMatReader.getFile(),
num));
+ }
+ }
+ // The queues keep owning these until the next emit releases them, so a
phase-2 read below that
+ // throws leaves them reachable for `close()`.
+ keyVectorsPublished = true;
+ // Key slots are filled in ascending slot order while `keyIdx` walks the
queues in key-list
+ // order, so the pairing is the identity only because
`ParquetStorageFilter.create` sorts
+ // `keyColumnIndices` ascending. `isKeyTopLevel` says which slots are
keys, not where each
+ // sits in that list, so this loop cannot re-derive the pairing: an
unsorted list would swap
+ // key columns in the output batch.
+ int keyIdx = 0;
+ for (int i = 0; i < spliceBatchColumns.length; i++) {
+ if (i < isKeyTopLevel.length && isKeyTopLevel[i]) {
+ spliceBatchColumns[i] = keyVectorQueues[keyIdx++].peekFirst();
+ } else {
+ spliceBatchColumns[i] = persistentBatchColumns[i];
+ }
+ }
+
+ readPersistentColumns(num, /* skipKeySlots= */ true);
+ 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);
+ }
+ finishBatch(num);
return true;
}
+ /**
+ * Closes the key vectors the previous batch was built on. The queues own
them until here, so
+ * survivor memory drains as a row group is emitted rather than all at its
end.
+ */
+ private void releasePublishedKeyVectors() {
+ if (!keyVectorsPublished) return;
+ keyVectorsPublished = false;
+ for (java.util.ArrayDeque<WritableColumnVector> queue : keyVectorQueues) {
+ queue.removeFirst().close();
+ }
+ }
+
private void initializeInternal() throws IOException,
UnsupportedOperationException {
missingColumns = new HashSet<>();
for (ParquetColumn column :
CollectionConverters.asJava(parquetColumn.children())) {
checkColumn(column);
}
+ if (storageFilter != null) {
+ initializeLateMaterialization();
+ }
+ }
+
+ /**
+ * Sets the storage filter for late materialization. Must be called before
{@link #initialize};
+ * {@link #initializeLateMaterialization()} then inspects the per-file
schema and decides whether
+ * splicing engages.
+ *
+ * <p>It does not engage in one real case: every key column is missing from
this physical file
+ * under schema evolution. The predicate is then rewritten with each missing
key replaced by the
+ * constant the reader materializes for it (its existence DEFAULT, else
null) and evaluated once
+ * -- true keeps the file unfiltered, false or null skips it.
+ *
+ * <p>Every other precondition is guaranteed by
+ * {@code FileSourceStrategy.extractStorageFilters} and {@code
ParquetStorageFilter.create}, and a
+ * violation throws rather than falling back: extraction has already removed
the filter from the
+ * post-scan Filter, so not applying it would return rows the query rejected.
+ */
+ public void setStorageFilter(ParquetStorageFilter storageFilter) {
+ this.storageFilter = storageFilter;
+ }
+
+ private void initializeLateMaterialization() throws IOException {
+ lateMatReader = reader.getUnderlyingReader();
+ if (lateMatReader == null) {
+ // Unreachable: the only ParquetRowGroupReader ParquetFileFormat builds
is
+ // ParquetRowGroupReaderImpl, which exposes its reader.
+ throw new UnsupportedFileReadException(
+ "Storage-filter pushdown requires a reader backed by a
ParquetFileReader, but "
+ + reader.getClass().getName() + " does not expose one");
+ }
+ if (configuration != null) {
+ useColumnIndexFilter = configuration.getBoolean(
+ ParquetInputFormat.COLUMN_INDEX_FILTERING_ENABLED, true);
+ }
+
+ // Resolve each key column's top-level ParquetColumn. Partition into
present and missing (the
+ // latter can happen under schema evolution: a column is in the requested
schema but not in this
+ // physical parquet file). For any non-primitive key we still bail;
phase-1 reads only primitive
+ // leaves.
+ int[] keyIndices = storageFilter.keyColumnIndices();
+ List<ParquetColumn> presentKeyColumns = new ArrayList<>(keyIndices.length);
+ List<Integer> missingKeyLocalPositions = new ArrayList<>();
+ for (int i = 0; i < keyIndices.length; i++) {
+ int idx = keyIndices[i];
+ if (idx < 0 || idx >= parquetColumn.children().size()) {
+ // Unreachable: ParquetStorageFilter.create rejects out-of-range
ordinals.
+ throw new UnsupportedFileReadException(String.format(
+ "Storage-filter key ordinal %d is out of range for a %d-column
requested schema",
+ idx, parquetColumn.children().size()));
+ }
+ ParquetColumn column = parquetColumn.children().apply(idx);
+ if (!column.isPrimitive()) {
+ // Unreachable: ParquetStorageFilter.isSupportedKeyType admits only
types with a primitive
+ // Parquet leaf, and it gates both planning and
ParquetStorageFilter.create.
+ throw new UnsupportedFileReadException(
+ "Storage-filter key column is not a primitive Parquet column: " +
column.path());
+ }
+ if (missingColumns.contains(column)) {
+ missingKeyLocalPositions.add(i);
+ } else {
+ presentKeyColumns.add(column);
+ }
+ }
+
+ // If any key column is missing from this file, rewrite the predicate to
substitute the constant
+ // the reader will actually materialize for that column. That is the
column's existence DEFAULT
+ // when it has one (ParquetColumnVector writes it into the output vector
and marks the vector
+ // constant), otherwise null. Substituting null for a column that reads
back as its default
+ // would filter on a value the scan never returns. The predicate must be
evaluated against the
+ // substituted constant rather than skipped: null does not always mean
false in a filter.
+ if (!missingKeyLocalPositions.isEmpty()) {
+ int[] missing = new int[missingKeyLocalPositions.size()];
+ Object[] missingValues = new Object[missingKeyLocalPositions.size()];
+ Object[] existenceDefaults =
+ ResolveDefaultColumns.existenceDefaultValues(sparkRequestedSchema);
+ for (int i = 0; i < missing.length; i++) {
+ missing[i] = missingKeyLocalPositions.get(i);
+ missingValues[i] = existenceDefaults[keyIndices[missing[i]]];
+ }
+ storageFilter = storageFilter.rewriteForMissingKeys(missing,
missingValues);
+
+ if (presentKeyColumns.isEmpty()) {
Review Comment:
Fixed: the skip now walks every row group of the file and updates all three
counters.
It measures against the rows the pushed data filter kept, which is the
baseline every other skip path uses. Measuring against the whole row group
would have let `rowsExcludedByRowGroup` exceed the reader's own row count,
since that is `getFilteredRecordCount()`. A test asserts the numbers. c255f88
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormat.scala:
##########
@@ -165,6 +166,57 @@ 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. A format that supports storage-filter pushdown
overrides this, and must
+ * not let the two builders call each other: this default delegates one way,
so an override that
+ * delegates back recurses until the driver's stack runs out, `super`
included, since that call is
+ * virtual too. Route both to a private implementation instead, the way
`ParquetFileFormat` does.
+ *
+ * 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.
+ *
+ * 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,
Review Comment:
Resolved from the other end by the model change: honoring a storage filter
is optional, so this default is a legitimate implementation rather than dead
code. It ignores the filters and delegates, which is correct precisely because
the conjunct is still in the `Filter`, and the `require` is gone.
The recursion warning stays, with a corrected reason. The default delegates
on `this`, so an override of `buildReaderWithPartitionValues` that delegates
back here closes the loop. Reaching this default through `super` is part of
that loop rather than an escape from it, which is what the old wording got
wrong. c255f88
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1885,6 +1885,42 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED =
Review Comment:
Done, both confs and their accessors are after
`PARQUET_FILTER_PUSHDOWN_INFILTERTHRESHOLD` now, so the `filterPushdown` group
stays with its sub-confs. c255f88
--
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]
