peter-toth commented on code in PR #58895:
URL: https://github.com/apache/spark/pull/58895#discussion_r4084285791
##########
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);
+ columnarBatch.setNumRows(num);
+
+ rowsReturned += num;
+ numBatched = num;
+ batchIdx = 0;
+ pendingCloseKeyVectors = dequeued;
+ return true;
+ }
+
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()} (run from {@link #initialize})
inspects the per-file
+ * schema and decides whether splicing actually engages.
+ *
+ * <p>Splicing does NOT engage in one real case: all key columns are missing
from this physical
+ * file under schema evolution. The predicate is rewritten with each missing
key replaced by the
+ * constant the reader materializes for it (its existence DEFAULT, else null
-- see
+ * {@code ParquetStorageFilter.rewriteForMissingKeys}) and evaluated as a
constant: a true result
+ * keeps the file with no filtering, a false/null result skips it entirely.
Either way the rows
+ * the scan returns are exactly the rows that satisfy the filter, so this is
safe.
+ *
+ * <p>Every OTHER precondition is guaranteed by planning-time checks in
+ * {@code FileSourceStrategy.extractStorageFilters} and {@code
ParquetStorageFilter.create}, and
+ * violations throw rather than fall back: once extraction has moved a bloom
filter onto the scan
+ * it is gone from the post-scan Filter, so quietly not applying it would
produce wrong rows.
+ */
+ public void setStorageFilter(ParquetStorageFilter storageFilter) {
+ this.storageFilter = storageFilter;
+ }
+
+ private void initializeLateMaterialization() throws IOException {
+ lateMatReader = reader.getUnderlyingReader();
+ if (lateMatReader == null) {
+ // Unreachable in production: the only ParquetRowGroupReader
ParquetFileFormat builds is
+ // ParquetRowGroupReaderImpl, which returns its reader. Dropping the
filter here would return
+ // rows it rejects, since extraction already removed it from the
post-scan Filter.
+ throw new IllegalStateException(
+ "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. Fail loudly
+ // rather than dropping the filter -- extractStorageFilters already
removed it from the
+ // post-scan Filter, so silently ignoring it here would return wrong
rows.
+ throw new IllegalStateException(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. Fail loudly for
+ // the same reason as above.
+ throw new IllegalStateException(
+ "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()) {
+ // All key columns missing: the rewritten predicate is fully constant.
Evaluate once and
+ // apply uniformly to the whole file.
+ boolean keepAll = storageFilter.evalAllMissing();
+ if (keepAll) {
+ // Predicate is constant-true for this file: no filtering to do.
+ storageFilter = null;
+ } else {
+ // Predicate is constant-false/null: no row from this file can pass
the filter.
+ storageFilter = null;
+ hitEndOfData = true;
+ }
+ return;
+ }
+ }
+
+ keyDescriptors = new ColumnDescriptor[presentKeyColumns.size()];
+ keyRequired = new boolean[presentKeyColumns.size()];
+ Types.MessageTypeBuilder keySchemaBuilder = Types.buildMessage();
+ Set<String> keyTopLevelNames = new HashSet<>();
+ for (int i = 0; i < presentKeyColumns.size(); i++) {
+ ParquetColumn column = presentKeyColumns.get(i);
+ keyDescriptors[i] = column.descriptor().get();
+ keyRequired[i] = column.required();
+ // Preserve the field name/type as it appears at the top of
requestedSchema.
+ String topLevelName = keyDescriptors[i].getPath()[0];
+ keySchemaBuilder.addField(requestedSchema.getType(topLevelName));
+ keyTopLevelNames.add(topLevelName);
+ }
+ keyOnlyRequestedSchema = keySchemaBuilder.named(requestedSchema.getName());
+
+ // Build the non-key (complement) schema. When the projection has at least
one non-key column,
+ // phase 2 will switch lateMatReader's schema to this and read those
columns under finalRanges.
+ // When the projection is *all* key columns (e.g. a scan whose only column
is the bloom probe),
+ // splicing still pays off: phase 2 has nothing useful to read, so we skip
it entirely (no read,
+ // no IO) and emit batches purely from the key queues. The
`nonKeyRequestedSchema != null` check
+ // downstream gates phase-2 IO.
+ Types.MessageTypeBuilder nonKeyBuilder = Types.buildMessage();
+ int nonKeyFieldCount = 0;
+ for (Type field : requestedSchema.getFields()) {
+ if (!keyTopLevelNames.contains(field.getName())) {
+ nonKeyBuilder.addField(field);
+ nonKeyFieldCount++;
+ }
+ }
+ if (nonKeyFieldCount > 0) {
+ nonKeyRequestedSchema = nonKeyBuilder.named(requestedSchema.getName());
+ requireOffsetIndexesForPhase2();
+ }
+ initializeSplicingState(presentKeyColumns);
+
+ totalBlockCount = lateMatReader.getRowGroups().size();
+ nextBlockIndex = 0;
+ }
+
+ /**
+ * Fails now if any projected column of any row group lacks a Parquet offset
index.
+ *
+ * <p>Phase 2 reads a strict subset of a row group's rows, which parquet can
only do via the
+ * offset index; files written before parquet-mr 1.11, or by writers that
omit it, have none. We
+ * cannot
+ * widen phase 2 to the whole block instead, because the key vectors already
hold only the
+ * survivors and the batch would misalign -- and we cannot skip the filter
either, since
+ * {@code extractStorageFilters} has already removed it from the post-scan
Filter.
+ *
+ * <p>Every projected column is checked, not just the non-key ones phase 2
reads, because parquet
+ * builds one column index store per row group and reuses it. Phase 0 asks
for the row ranges
+ * under the full requested schema, so {@code ColumnIndexStoreImpl.create}
is called with the key
+ * columns in its path set, and it returns its {@code EMPTY} singleton as
soon as any one of those
+ * paths has no offset index. {@code ParquetFileReader.getColumnIndexStore}
memoizes that store
+ * per block, and {@code EMPTY.getOffsetIndex} throws for *every* column. So
a key column with no
+ * offset index kills phase 2 too, with a raw {@code
MissingOffsetIndexException} naming some
+ * non-key column and none of the guidance below.
+ *
+ * <p>Checking up front rather than at the first partially-kept row group is
deliberate: whether
+ * phase 2 needs the offset index otherwise depends on how selective the
filter turns out to be on
+ * this particular file, so the same query would fail or not depending on
the data. This is
+ * conservative -- a filter that happens to keep every row of every block
would not have needed
+ * the offset index -- but such a filter also saves nothing, so failing
loudly loses nothing.
+ *
+ * <p>The check itself is free: {@code getOffsetIndexReference()} is a
footer field that
+ * {@link #initialize} has already read.
+ */
+ private void requireOffsetIndexesForPhase2() {
+ Set<ColumnPath> projectedPaths = new HashSet<>();
+ for (ColumnDescriptor column : requestedSchema.getColumns()) {
+ projectedPaths.add(ColumnPath.get(column.getPath()));
+ }
+ List<BlockMetaData> blocks = lateMatReader.getRowGroups();
+ for (int blockIdx = 0; blockIdx < blocks.size(); blockIdx++) {
+ for (ColumnChunkMetaData chunk : blocks.get(blockIdx).getColumns()) {
+ if (projectedPaths.contains(chunk.getPath()) &&
chunk.getOffsetIndexReference() == null) {
Review Comment:
Narrowed rather than turned into a per-file fallback, and it came out a net
deletion.
The up-front check is gone. Parquet resolves every requested column's offset
index before it reads a byte, so phase 2's read is simply wrapped and
`MissingOffsetIndexException` rethrown with the conf to set. That is narrower
than the check was -- a row group the filter keeps whole never needs the index,
since `readFilteredRowGroup` degrades to a plain read when the ranges cover the
block, and one it rejects whole is never read at all -- so a file written
without a page index still scans as long as the filter never prunes inside a
row group. It is also more complete: an IO error while reading an index
surfaces the same way, which a footer walk could not see.
The eager-read fallback does not work in this design: phase 1 has already
buffered only the survivors, so a whole-block phase 2 would misalign the batch,
and dropping the predicate is not an option because extraction has removed it
from the post-scan `Filter`. What would remove the failure class entirely is a
parquet-side change -- `readFilteredRowGroup` reading whole chunks while
keeping the row ranges for `getRowIndexes()` -- and that is now named as a
follow-up in the description.
##########
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();
Review Comment:
Bounded now, in the stopgap form you suggested, and the conf is
`spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes` (internal,
64MB). Phase 1 counts the bytes it has buffered and past the cap releases them
and gives that row group up: phase 2 then reads every projected column of the
surviving rows, so nothing is buffered and the cost is one extra read of the
key columns. The count is examined once per accumulator, so the hot loop pays
one add per key column and the check costs one comparison per batch worth of
survivors.
Worth recording why it counts rather than estimates, since the first version
did estimate from the footer. `ColumnChunkMetaData.getTotalUncompressedSize()`
is the size of the uncompressed *pages*, which still hold dictionary ids:
measured, a 2M-row BIGINT column with 1000 distinct values reports 1.26
bytes/row dictionary-encoded against the 9 the accumulator holds. Spark writes
dictionary-encoded by default, so that estimate underestimated by 7x on the
default path.
Chunk-at-a-time is discussed in the description. It would shorten what the
first batch waits for, but the IO of a `LIMIT` query is unchanged either way,
and chunking turns phase 2 into one read call per chunk and refetches any page
that straddles a boundary.
##########
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);
Review Comment:
It is a latency and CPU cost, not an IO one: `internalReadFilteredRowGroup`
reads all of the row group's requested chunks before it hands back a
`PageReadStore`, so `LIMIT 5` transfers the same bytes on the eager path too.
What the first batch waits for is one row group's key decode and predicate
evaluation, and the buffer cap now bounds that as well.
Chunk-at-a-time would shorten it, but at a price the description spells out:
one read call per chunk instead of one per row group, pages that straddle a
chunk boundary fetched more than once with nothing caching them (parquet caches
footers, not pages, and S3A's default input stream caches nothing either), and
the coalescing and parallelism of a single vectored read given up -- parquet
1.18 has vectored IO on by default. So it is not filed.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -0,0 +1,1965 @@
+/*
+ * 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 java.io.{ByteArrayOutputStream, File}
+import java.net.URI
+import java.time.LocalTime
+import java.util.concurrent.atomic.AtomicLong
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileStatus, FSDataInputStream, FSInputStream,
Path, RawLocalFileSystem}
+import org.apache.hadoop.mapreduce.Job
+import org.apache.parquet.hadoop.{ParquetInputFormat, ParquetOutputFormat}
+
+import org.apache.spark.paths.SparkPath
+import org.apache.spark.sql.{sources, QueryTest, Row, SparkSession}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute,
BloomFilterMightContain, BoundReference, Coalesce, Expression,
GreaterThanOrEqual, IsNull, LessThanOrEqual, Literal, Or, Predicate, Rand,
XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter => LogicalFilter}
+import org.apache.spark.sql.execution.{CollapseCodegenStages,
ColumnarToRowExec, FileSourceScanExec, FileSourceScanLike, FilterExec,
LocalLimitExec, SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.datasources.{FileFormat,
FileSourceStrategy, OutputWriterFactory, PartitionedFile}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.functions.col
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.util.sketch.BloomFilter
+
+/**
+ * Tests the late-materialization path of [[VectorizedParquetRecordReader]]
driven by a
+ * [[ParquetStorageFilter]]. Writes small multi-row-group parquet files, wires
a hand-built filter
+ * into the reader, and asserts correctness + the two storage-filter metrics.
+ */
+class ParquetStorageFilterSuite extends QueryTest with SharedSparkSession {
+ import testImplicits._
+
+ // Writes a parquet file with the given rows and row-group size; returns the
path.
+ private def writeParquetFile(
+ dir: File,
+ rows: Seq[(Long, String)],
+ rowGroupSize: Long = 1024L,
+ pageSize: Option[Long] = None): String = {
+ val outDir = new File(dir, s"test-${System.nanoTime()}").getAbsolutePath
+ val writer = rows.toDF("k", "v")
+ .repartition(1)
+ .write
+ .option(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize)
+ // Dictionary encoding off keeps row-group sizing predictable. Note this
does NOT disable the
+ // column index, despite what an earlier version of this comment claimed.
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, "false")
+ // A small page size gives each row group several pages per column, which
is what lets
+ // column-index filtering produce a row range narrower than the whole row
group.
+ pageSize.foreach(size => writer.option(ParquetOutputFormat.PAGE_SIZE,
size))
+ writer.parquet(outDir)
+ val files = new File(outDir).listFiles((_, name) =>
name.endsWith(".parquet"))
+ assert(files != null && files.length == 1, s"expected exactly one parquet
file under $outDir")
+ files(0).getAbsolutePath
+ }
+
+ // Collects all rows from a reader initialized with the given storage filter.
+ private def readAll(
+ filePath: String,
+ storageFilter: ParquetStorageFilter): (Seq[(Long, String)],
VectorizedParquetRecordReader) = {
+ val reader = new VectorizedParquetRecordReader(false, 4096)
+ reader.setStorageFilter(storageFilter)
+ reader.initialize(filePath, java.util.Arrays.asList("k", "v"))
+ reader.initBatch(new StructType(), null)
+ val collected = mutable.ArrayBuffer[(Long, String)]()
+ while (reader.nextBatch()) {
+ val batch = reader.resultBatch().asInstanceOf[ColumnarBatch]
+ val n = batch.numRows()
+ val kVec = batch.column(0)
+ val vVec = batch.column(1)
+ var i = 0
+ while (i < n) {
+ collected += ((kVec.getLong(i), vVec.getUTF8String(i).toString))
+ i += 1
+ }
+ }
+ (collected.toSeq, reader)
Review Comment:
Fixed in all four: the drain happens inside `Utils.tryInitializeResource`,
which closes the reader if anything in the loop throws and leaves it open
otherwise, so the caller's `finally` still owns the close on the passing path.
--
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]
