> output = processAndAssertRetainedPageSize(pageProcessor, inputPage);
+ Page outputPage = getOnlyElement(ImmutableList.copyOf(output)).orElseThrow();
+ assertPageEquals(ImmutableList.of(BIGINT), outputPage, new Page(createLongSequenceBlock(101, 106).copyPositions(new int[] {0, 2, 4}, 0, 3)));
+ assertThat(outputPage.getBlock(0)).isInstanceOf(LongArrayBlock.class);
+ assertThat(inputPage.getPositionCount()).isEqualTo(10);
+ }
+
@Test
public void testSelectAllFilter()
{
diff --git a/core/trino-main/src/test/java/io/trino/sql/gen/columnar/TestColumnarFilterEvaluator.java b/core/trino-main/src/test/java/io/trino/sql/gen/columnar/TestColumnarFilterEvaluator.java
new file mode 100644
index 000000000000..ad7c009e40b6
--- /dev/null
+++ b/core/trino-main/src/test/java/io/trino/sql/gen/columnar/TestColumnarFilterEvaluator.java
@@ -0,0 +1,89 @@
+/*
+ * Licensed 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 io.trino.sql.gen.columnar;
+
+import io.trino.operator.TestingSourcePage;
+import io.trino.operator.project.InputChannels;
+import io.trino.operator.project.SelectedPositions;
+import io.trino.spi.block.Block;
+import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.connector.SourcePage;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import static io.trino.block.BlockAssertions.createLongSequenceBlock;
+import static io.trino.operator.project.SelectedPositions.positionsList;
+import static io.trino.operator.project.SelectedPositions.positionsRange;
+import static io.trino.testing.TestingConnectorSession.SESSION;
+import static java.lang.System.arraycopy;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class TestColumnarFilterEvaluator
+{
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ public void testMaterializesOnlyFilterInputs(boolean list)
+ {
+ Block first = createLongSequenceBlock(0, 4);
+ Block last = createLongSequenceBlock(100, 104);
+ TestingSourcePage source = new TestingSourcePage(4, first, null, last);
+ ColumnarFilter filter = new ColumnarFilter()
+ {
+ @Override
+ public InputChannels getInputChannels()
+ {
+ return new InputChannels(2, 0);
+ }
+
+ @Override
+ public int filterPositionsRange(ConnectorSession session, int[] outputPositions, int offset, int size, SourcePage loadedPage)
+ {
+ assertInputsLoaded(loadedPage);
+ for (int index = 0; index < size; index++) {
+ outputPositions[index] = offset + index;
+ }
+ return size;
+ }
+
+ @Override
+ public int filterPositionsList(ConnectorSession session, int[] outputPositions, int[] activePositions, int offset, int size, SourcePage loadedPage)
+ {
+ assertInputsLoaded(loadedPage);
+ arraycopy(activePositions, offset, outputPositions, 0, size);
+ return size;
+ }
+
+ private void assertInputsLoaded(SourcePage loadedPage)
+ {
+ assertThat(source.wasLoaded(0)).isTrue();
+ assertThat(source.wasLoaded(1)).isFalse();
+ assertThat(source.wasLoaded(2)).isTrue();
+ assertThat(loadedPage.getPositionCount()).isEqualTo(4);
+ assertThat(loadedPage.getChannelCount()).isEqualTo(2);
+ assertThat(loadedPage.getBlock(0)).isSameAs(last);
+ assertThat(loadedPage.getBlock(1)).isSameAs(first);
+ }
+ };
+ SelectedPositions activePositions = list ? positionsList(new int[] {0, 1, 3}, 1, 2) : positionsRange(1, 2);
+ SelectedPositions result = new ColumnarFilterEvaluator(filter).evaluate(SESSION, activePositions, source).selectedPositions();
+ assertThat(result.size()).isEqualTo(2);
+ assertThat(result.isList()).isEqualTo(list);
+ if (list) {
+ assertThat(result.getPositions()).containsExactly(1, 3);
+ }
+ else {
+ assertThat(result.getOffset()).isEqualTo(1);
+ }
+ }
+}
diff --git a/core/trino-main/src/test/java/io/trino/sql/query/TestAggregation.java b/core/trino-main/src/test/java/io/trino/sql/query/TestAggregation.java
index a8b9c6176d92..dd55c887d1b5 100644
--- a/core/trino-main/src/test/java/io/trino/sql/query/TestAggregation.java
+++ b/core/trino-main/src/test/java/io/trino/sql/query/TestAggregation.java
@@ -40,6 +40,33 @@ public void teardown()
assertions.close();
}
+ @Test
+ public void testDecimalAverageWithScaledOverflow()
+ {
+ assertThat(assertions.query("SELECT avg(x) FROM UNNEST(CAST(ARRAY[0.9, 0.9] AS array(decimal(38,38)))) t(x)"))
+ .matches("VALUES CAST(0.9 AS decimal(38,38))");
+ }
+
+ @Test
+ public void testDecimalAverageWithoutOverflow()
+ {
+ assertThat(assertions.query("SELECT avg(x) FROM UNNEST(CAST(ARRAY[0.8, 0.8] AS array(decimal(38,38)))) t(x)"))
+ .matches("VALUES CAST(0.8 AS decimal(38,38))");
+ }
+
+ @Test
+ public void testDecimalAverageWithScaledOverflowAndZeroValues()
+ {
+ assertThat(assertions.query(
+ """
+ SELECT avg(x)
+ FROM UNNEST(concat(
+ ARRAY[DECIMAL '9000000000000000000000000000000000000.0', DECIMAL '9000000000000000000000000000000000000.0'],
+ repeat(CAST(0 AS decimal(38,1)), 98))) t(x)
+ """))
+ .matches("VALUES CAST(DECIMAL '180000000000000000000000000000000000.0' AS decimal(38,1))");
+ }
+
@Test
public void testQuantifiedComparison()
{
diff --git a/core/trino-spi/src/main/java/io/trino/spi/StandardErrorCode.java b/core/trino-spi/src/main/java/io/trino/spi/StandardErrorCode.java
index 8f75da5fcd40..d0faec5bce5d 100644
--- a/core/trino-spi/src/main/java/io/trino/spi/StandardErrorCode.java
+++ b/core/trino-spi/src/main/java/io/trino/spi/StandardErrorCode.java
@@ -206,6 +206,7 @@ public enum StandardErrorCode
EXCEEDED_SCAN_LIMIT(131081, INSUFFICIENT_RESOURCES),
EXCEEDED_TASK_DESCRIPTOR_STORAGE_CAPACITY(131082, INSUFFICIENT_RESOURCES),
EXCEEDED_WRITE_LIMIT(131083, INSUFFICIENT_RESOURCES),
+ EXCEEDED_OUTPUT_LIMIT(131084, INSUFFICIENT_RESOURCES),
UNSUPPORTED_TABLE_TYPE(133001, EXTERNAL),
/**/;
diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/MemoryUsageReportingPageSource.java b/core/trino-spi/src/main/java/io/trino/spi/connector/MemoryUsageReportingPageSource.java
index 79793b15d04b..f291d807521a 100644
--- a/core/trino-spi/src/main/java/io/trino/spi/connector/MemoryUsageReportingPageSource.java
+++ b/core/trino-spi/src/main/java/io/trino/spi/connector/MemoryUsageReportingPageSource.java
@@ -155,6 +155,14 @@ public Page getPage()
return page;
}
+ @Override
+ public boolean trySelectPositions(int[] positions, int offset, int size)
+ {
+ boolean selected = sourcePage.trySelectPositions(positions, offset, size);
+ memoryContext.setBytes(pageSource.getMemoryUsage());
+ return selected;
+ }
+
@Override
public Page getColumns(int[] channels)
{
diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/SourcePage.java b/core/trino-spi/src/main/java/io/trino/spi/connector/SourcePage.java
index 37316e2331d5..2352317c00b4 100644
--- a/core/trino-spi/src/main/java/io/trino/spi/connector/SourcePage.java
+++ b/core/trino-spi/src/main/java/io/trino/spi/connector/SourcePage.java
@@ -86,6 +86,19 @@ static SourcePage create(Page page)
*/
Page getPage();
+ /**
+ * Attempts to apply the specified position selection when doing so can avoid loading or
+ * decoding data that has not been accessed yet. Returns whether the selection was applied.
+ * If this method returns false, the logical contents and position count of the page must remain unchanged.
+ *
+ * The positions array is owned by the caller and may only be accessed for the duration of
+ * this call. An implementation must copy any positions it retains.
+ */
+ default boolean trySelectPositions(int[] positions, int offset, int size)
+ {
+ return false;
+ }
+
/**
* Gets a projection of the page containing only the specified channels.
*/
diff --git a/docs/src/main/sphinx/admin/properties-query-management.md b/docs/src/main/sphinx/admin/properties-query-management.md
index 9a90d78f5fa6..e2561ea0c56a 100644
--- a/docs/src/main/sphinx/admin/properties-query-management.md
+++ b/docs/src/main/sphinx/admin/properties-query-management.md
@@ -168,6 +168,14 @@ The maximum physical size of data that can be written by a query during its exec
When this limit is reached, query processing is terminated to prevent excessive
resource usage.
+## `query.max-output-data-size`
+
+- **Type:** {ref}`prop-type-data-size`
+
+The maximum size of the output data that a query can produce during its execution.
+When this limit is exceeded, query processing is terminated. This limit protects
+against queries that return excessively large result sets to the client.
+
## `query.max-stage-count`
- **Type:** {ref}`prop-type-integer`
diff --git a/docs/src/main/sphinx/object-storage/file-formats.md b/docs/src/main/sphinx/object-storage/file-formats.md
index 1490e61bdc09..2879a544b58f 100644
--- a/docs/src/main/sphinx/object-storage/file-formats.md
+++ b/docs/src/main/sphinx/object-storage/file-formats.md
@@ -94,8 +94,7 @@ with Parquet files performed by supported object storage connectors:
- `true`
* - `parquet.use-column-index`
- Skip reading Parquet pages by using Parquet column indices. The equivalent
- catalog session property is `parquet_use_column_index`. Only supported by
- the Delta Lake and Hive connectors.
+ catalog session property is `parquet_use_column_index`.
- `true`
* - `parquet.ignore-statistics`
- Ignore statistics from Parquet to allow querying files with corrupted or
@@ -112,6 +111,11 @@ with Parquet files performed by supported object storage connectors:
entirely. The equivalent catalog session property is named
`parquet_small_file_threshold`.
- `3MB`
+* - `parquet.selected-positions-pushdown-enabled`
+ - Allow the Parquet reader to use rows selected by query filters to skip
+ decoding rejected values and decompressing entirely unselected data pages
+ when beneficial. Set to `false` to disable filter selection pushdown.
+ - `true`
* - `parquet.experimental.vectorized-decoding.enabled`
- Enable using Java Vector API (SIMD) for faster decoding of parquet files.
The equivalent catalog session property is
diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java
index 4ca8e2002540..9972e548fefc 100644
--- a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java
+++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java
@@ -39,6 +39,7 @@ public class ParquetReaderOptions
private final boolean useColumnIndex;
private final boolean useBloomFilter;
private final DataSize smallFileThreshold;
+ private final boolean selectedPositionsPushdownEnabled;
private final boolean vectorizedDecodingEnabled;
private final DataSize footerReadSize;
private final DataSize maxFooterReadSize;
@@ -54,6 +55,7 @@ private ParquetReaderOptions()
useColumnIndex = true;
useBloomFilter = true;
smallFileThreshold = DEFAULT_SMALL_FILE_THRESHOLD;
+ selectedPositionsPushdownEnabled = true;
vectorizedDecodingEnabled = true;
footerReadSize = DEFAULT_FOOTER_READ_SIZE;
maxFooterReadSize = DEFAULT_MAX_FOOTER_READ_SIZE;
@@ -69,6 +71,7 @@ private ParquetReaderOptions(
boolean useColumnIndex,
boolean useBloomFilter,
DataSize smallFileThreshold,
+ boolean selectedPositionsPushdownEnabled,
boolean vectorizedDecodingEnabled,
DataSize footerReadSize,
DataSize maxFooterReadSize,
@@ -83,6 +86,7 @@ private ParquetReaderOptions(
this.useColumnIndex = useColumnIndex;
this.useBloomFilter = useBloomFilter;
this.smallFileThreshold = requireNonNull(smallFileThreshold, "smallFileThreshold is null");
+ this.selectedPositionsPushdownEnabled = selectedPositionsPushdownEnabled;
this.vectorizedDecodingEnabled = vectorizedDecodingEnabled;
this.footerReadSize = requireNonNull(footerReadSize, "footerReadSize is null");
checkArgument(footerReadSize.toBytes() >= 8, "footerReadSize must be at least 8 bytes");
@@ -150,6 +154,11 @@ public DataSize getSmallFileThreshold()
return smallFileThreshold;
}
+ public boolean isSelectedPositionsPushdownEnabled()
+ {
+ return selectedPositionsPushdownEnabled;
+ }
+
public DataSize getMaxFooterReadSize()
{
return maxFooterReadSize;
@@ -175,6 +184,7 @@ public static class Builder
private boolean useColumnIndex;
private boolean useBloomFilter;
private DataSize smallFileThreshold;
+ private boolean selectedPositionsPushdownEnabled;
private boolean vectorizedDecodingEnabled;
private DataSize footerReadSize;
private DataSize maxFooterReadSize;
@@ -191,6 +201,7 @@ private Builder(ParquetReaderOptions parquetReaderOptions)
this.useColumnIndex = parquetReaderOptions.useColumnIndex;
this.useBloomFilter = parquetReaderOptions.useBloomFilter;
this.smallFileThreshold = parquetReaderOptions.smallFileThreshold;
+ this.selectedPositionsPushdownEnabled = parquetReaderOptions.selectedPositionsPushdownEnabled;
this.vectorizedDecodingEnabled = parquetReaderOptions.vectorizedDecodingEnabled;
this.footerReadSize = parquetReaderOptions.footerReadSize;
this.maxFooterReadSize = parquetReaderOptions.maxFooterReadSize;
@@ -245,6 +256,12 @@ public Builder withSmallFileThreshold(DataSize smallFileThreshold)
return this;
}
+ public Builder withSelectedPositionsPushdownEnabled(boolean selectedPositionsPushdownEnabled)
+ {
+ this.selectedPositionsPushdownEnabled = selectedPositionsPushdownEnabled;
+ return this;
+ }
+
public Builder withVectorizedDecodingEnabled(boolean vectorizedDecodingEnabled)
{
this.vectorizedDecodingEnabled = vectorizedDecodingEnabled;
@@ -280,6 +297,7 @@ public ParquetReaderOptions build()
useColumnIndex,
useBloomFilter,
smallFileThreshold,
+ selectedPositionsPushdownEnabled,
vectorizedDecodingEnabled,
footerReadSize,
maxFooterReadSize,
diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/AbstractColumnReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/AbstractColumnReader.java
index 18313d41c291..71b3a3c6556c 100644
--- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/AbstractColumnReader.java
+++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/AbstractColumnReader.java
@@ -84,12 +84,25 @@ public void setPageReader(PageReader pageReader, Optional row
// For dictionary based encodings - https://github.com/apache/parquet-format/blob/master/Encodings.md
if (dictionaryPage != null) {
log.debug("field %s, readDictionaryPage %s", field, dictionaryPage);
- dictionaryDecoder = dictionaryDecoderProvider.create(dictionaryPage, isNonNull());
- produceDictionaryBlock = shouldProduceDictionaryBlock(rowRanges);
+ try {
+ dictionaryDecoder = dictionaryDecoderProvider.create(dictionaryPage, isNonNull());
+ produceDictionaryBlock = shouldProduceDictionaryBlock(rowRanges);
+ }
+ finally {
+ pageReader.releaseCurrentPage();
+ }
}
this.rowRanges = createRowRangesIterator(rowRanges);
}
+ @Override
+ public void close()
+ {
+ if (pageReader != null) {
+ pageReader.close();
+ }
+ }
+
protected abstract boolean isNonNull();
protected boolean produceDictionaryBlock()
diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnChunk.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnChunk.java
index 2aac0cfcd49d..5db4db5444d5 100644
--- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnChunk.java
+++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnChunk.java
@@ -17,6 +17,7 @@
import java.util.OptionalLong;
+import static com.google.common.base.Preconditions.checkArgument;
import static java.util.Objects.requireNonNull;
public class ColumnChunk
@@ -24,6 +25,7 @@ public class ColumnChunk
private final Block block;
private final int[] definitionLevels;
private final int[] repetitionLevels;
+ private final int maxBlockPositionCount;
private OptionalLong maxBlockSize;
public ColumnChunk(Block block, int[] definitionLevels, int[] repetitionLevels)
@@ -32,11 +34,18 @@ public ColumnChunk(Block block, int[] definitionLevels, int[] repetitionLevels)
}
public ColumnChunk(Block block, int[] definitionLevels, int[] repetitionLevels, OptionalLong maxBlockSize)
+ {
+ this(block, definitionLevels, repetitionLevels, maxBlockSize, block.getPositionCount());
+ }
+
+ public ColumnChunk(Block block, int[] definitionLevels, int[] repetitionLevels, OptionalLong maxBlockSize, int maxBlockPositionCount)
{
this.block = requireNonNull(block, "block is null");
this.definitionLevels = requireNonNull(definitionLevels, "definitionLevels is null");
this.repetitionLevels = requireNonNull(repetitionLevels, "repetitionLevels is null");
this.maxBlockSize = maxBlockSize;
+ checkArgument(maxBlockPositionCount >= block.getPositionCount(), "maxBlockPositionCount is less than block position count");
+ this.maxBlockPositionCount = maxBlockPositionCount;
}
public Block getBlock()
@@ -61,4 +70,9 @@ public long getMaxBlockSize()
}
return maxBlockSize.orElseThrow();
}
+
+ public int getMaxBlockPositionCount()
+ {
+ return maxBlockPositionCount;
+ }
}
diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReader.java
index e06686ba7ec6..805462b2896d 100644
--- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReader.java
+++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReader.java
@@ -17,6 +17,8 @@
public interface ColumnReader
{
+ void close();
+
boolean hasPageReader();
void setPageReader(PageReader pageReader, Optional rowRanges);
@@ -24,4 +26,29 @@ public interface ColumnReader
void prepareNextRead(int batchSize);
ColumnChunk readPrimitive();
+
+ default boolean supportsSelectedPositions()
+ {
+ return false;
+ }
+
+ default ColumnChunk readPrimitive(int[] positions, int offset, int positionCount)
+ {
+ throw new UnsupportedOperationException("Selected positions are not supported");
+ }
+
+ default ColumnChunk readPrimitivePageFiltered(int[] positions, int offset, int positionCount)
+ {
+ throw new UnsupportedOperationException("Page-filtered positions are not supported");
+ }
+
+ default int getDataPageReadCount()
+ {
+ return 0;
+ }
+
+ default long preparePageFilteredRead(int[] positions, int offset, int positionCount, long maxBufferedBytes)
+ {
+ return 0;
+ }
}
diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/NestedColumnReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/NestedColumnReader.java
index a569e41df141..47932572f574 100644
--- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/NestedColumnReader.java
+++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/NestedColumnReader.java
@@ -490,14 +490,9 @@ else if (page instanceof DataPageV2 dataPageV2) {
remainingPageValueCount = pageValueCount;
rowRanges.resetForNewPage(page.getFirstRowIndex());
- // For a compressed data page, the memory used by the decompressed values data needs to be accounted
- // for separately as ParquetCompressionUtils#decompress allocates a new byte array for the decompressed result.
- // For an uncompressed data page, we read directly from input Slices whose memory usage is already accounted
- // for in AbstractParquetDataSource#ReferenceCountedReader.
- int dataPageSizeInBytes = pageReader.arePagesCompressed() ? page.getUncompressedSize() : 0;
long dictionarySizeInBytes = dictionaryDecoder == null ? 0 : dictionaryDecoder.getRetainedSizeInBytes();
long repetitionBufferSizeInBytes = sizeOf(repetitionBuffer);
- memoryContext.setBytes(dataPageSizeInBytes + dictionarySizeInBytes + repetitionBufferSizeInBytes);
+ memoryContext.setBytes(dictionarySizeInBytes + repetitionBufferSizeInBytes + pageReader.getRetainedPageBytes());
return true;
}
diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java
index a3e5fcda0bb1..37b6843721d3 100644
--- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java
+++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java
@@ -14,8 +14,7 @@
package io.trino.parquet.reader;
import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.Iterators;
-import com.google.common.collect.PeekingIterator;
+import com.google.common.collect.ImmutableList;
import io.airlift.slice.Slice;
import io.trino.parquet.DataPage;
import io.trino.parquet.DataPageV1;
@@ -37,7 +36,9 @@
import org.apache.parquet.internal.column.columnindex.OffsetIndex;
import java.io.IOException;
+import java.util.ArrayDeque;
import java.util.Iterator;
+import java.util.List;
import java.util.Optional;
import static com.google.common.base.Preconditions.checkArgument;
@@ -53,11 +54,18 @@ public final class PageReader
private final CompressionCodec codec;
private final boolean hasOnlyDictionaryEncodedPages;
private final boolean hasNoNulls;
- private final PeekingIterator compressedPages;
+ private final Iterator compressedPages;
+ private final ArrayDeque bufferedPages = new ArrayDeque<>();
private final Optional blockDecryptor;
private boolean dictionaryAlreadyRead;
private int dataPageReadCount;
+ private long bufferedPageBytes;
+ private long bufferedPageRetainedBytes;
+ private long currentPageRetainedBytes;
+ private boolean currentPageOwned;
+ @Nullable
+ private Page currentPage;
@Nullable
private byte[] dataPageAad;
@Nullable
@@ -114,7 +122,7 @@ public PageReader(
{
this.dataSourceId = requireNonNull(dataSourceId, "dataSourceId is null");
this.codec = codec;
- this.compressedPages = Iterators.peekingIterator(compressedPages);
+ this.compressedPages = requireNonNull(compressedPages, "compressedPages is null");
this.hasOnlyDictionaryEncodedPages = hasOnlyDictionaryEncodedPages;
this.hasNoNulls = hasNoNulls;
this.blockDecryptor = decryptionContext.map(ColumnDecryptionContext::dataDecryptor);
@@ -136,20 +144,21 @@ public boolean hasOnlyDictionaryEncodedPages()
public DataPage readPage()
{
- if (!compressedPages.hasNext()) {
+ if (!hasNext()) {
return null;
}
- Page compressedPage = compressedPages.next();
- checkState(compressedPage instanceof DataPage, "Found page %s instead of a DataPage", compressedPage);
- dataPageReadCount++;
+ Page compressedPage = nextCompressedPage();
try {
+ checkState(compressedPage instanceof DataPage, "Found page %s instead of a DataPage", compressedPage);
+ dataPageReadCount++;
if (blockDecryptor.isPresent()) {
AesCipherUtils.quickUpdatePageAAD(dataPageAad, ((DataPage) compressedPage).getPageIndex());
}
Slice slice = decryptSliceIfNeeded(compressedPage.getSlice(), dataPageAad);
if (compressedPage instanceof DataPageV1 dataPageV1) {
- return new DataPageV1(
- !arePagesCompressed() ? slice : decompress(dataSourceId, codec, slice, dataPageV1.getUncompressedSize()),
+ PageData pageData = getPageData(slice, dataPageV1.getUncompressedSize(), arePagesCompressed());
+ DataPage page = new DataPageV1(
+ pageData.data(),
dataPageV1.getValueCount(),
dataPageV1.getUncompressedSize(),
dataPageV1.getFirstRowIndex(),
@@ -157,31 +166,46 @@ public DataPage readPage()
dataPageV1.getDefinitionLevelEncoding(),
dataPageV1.getValueEncoding(),
dataPageV1.getPageIndex());
+ replaceCurrentPage(page, pageData.retainedBytes(), pageData.owned());
+ return page;
}
DataPageV2 dataPageV2 = (DataPageV2) compressedPage;
- if (!dataPageV2.isCompressed()) {
+ if (!dataPageV2.isCompressed() && blockDecryptor.isEmpty()) {
return dataPageV2;
}
- int uncompressedSize = dataPageV2.getUncompressedSize()
+ int uncompressedDataSize = dataPageV2.getUncompressedSize()
- dataPageV2.getDefinitionLevels().length()
- dataPageV2.getRepetitionLevels().length();
- return new DataPageV2(
+ PageData pageData = getPageData(slice, uncompressedDataSize, dataPageV2.isCompressed());
+ Slice repetitionLevels = copySlice(dataPageV2.getRepetitionLevels());
+ Slice definitionLevels = copySlice(dataPageV2.getDefinitionLevels());
+ DataPage page = new DataPageV2(
dataPageV2.getRowCount(),
dataPageV2.getNullCount(),
dataPageV2.getValueCount(),
- dataPageV2.getRepetitionLevels(),
- dataPageV2.getDefinitionLevels(),
+ repetitionLevels,
+ definitionLevels,
dataPageV2.getDataEncoding(),
- decompress(dataSourceId, codec, slice, uncompressedSize),
+ pageData.data(),
dataPageV2.getUncompressedSize(),
dataPageV2.getFirstRowIndex(),
dataPageV2.getStatistics(),
false,
dataPageV2.getPageIndex());
+ replaceCurrentPage(
+ page,
+ repetitionLevels.length() + definitionLevels.length() + pageData.retainedBytes(),
+ true);
+ return page;
}
catch (IOException e) {
+ releaseCurrentPage();
throw new RuntimeException("Could not decompress page", e);
}
+ catch (RuntimeException | Error e) {
+ releaseCurrentPage();
+ throw e;
+ }
}
public DictionaryPage readDictionaryPage()
@@ -189,41 +213,291 @@ public DictionaryPage readDictionaryPage()
checkState(!dictionaryAlreadyRead, "Dictionary was already read");
checkState(dataPageReadCount == 0, "Dictionary has to be read first but %s was read already", dataPageReadCount);
dictionaryAlreadyRead = true;
- if (!(compressedPages.peek() instanceof DictionaryPage)) {
+ if (!(peekCompressedPage() instanceof DictionaryPage)) {
return null;
}
try {
- DictionaryPage compressedDictionaryPage = (DictionaryPage) compressedPages.next();
+ DictionaryPage compressedDictionaryPage = (DictionaryPage) nextCompressedPage();
Slice slice = decryptSliceIfNeeded(compressedDictionaryPage.getSlice(), dictionaryPageAad);
- return new DictionaryPage(
- decompress(dataSourceId, codec, slice, compressedDictionaryPage.getUncompressedSize()),
+ PageData pageData = getPageData(slice, compressedDictionaryPage.getUncompressedSize(), arePagesCompressed());
+ DictionaryPage dictionaryPage = new DictionaryPage(
+ pageData.data(),
compressedDictionaryPage.getDictionarySize(),
compressedDictionaryPage.getEncoding());
+ replaceCurrentPage(dictionaryPage, pageData.retainedBytes(), pageData.owned());
+ return dictionaryPage;
}
catch (IOException e) {
+ releaseCurrentPage();
throw new RuntimeException("Error reading dictionary page", e);
}
+ catch (RuntimeException | Error e) {
+ releaseCurrentPage();
+ throw e;
+ }
}
public boolean hasNext()
{
- return compressedPages.hasNext();
+ return !bufferedPages.isEmpty() || compressedPages.hasNext();
}
public DataPage getNextPage()
{
verifyDictionaryPageRead();
- return (DataPage) compressedPages.peek();
+ return (DataPage) peekCompressedPage();
}
public void skipNextPage()
{
verifyDictionaryPageRead();
- compressedPages.next();
+ nextCompressedPage();
+ releaseCurrentPage();
+ }
+
+ public List getNextDataPages(int valueCount, int maxPageCount, long maxBufferedBytes)
+ {
+ verifyDictionaryPageRead();
+ checkArgument(valueCount >= 0, "valueCount is negative");
+ checkArgument(maxPageCount > 0, "maxPageCount must be positive");
+ checkArgument(maxBufferedBytes > 0, "maxBufferedBytes must be positive");
+
+ int bufferedValueCount = 0;
+ int pageCount = 0;
+ for (BufferedPage bufferedPage : bufferedPages) {
+ Page page = bufferedPage.page();
+ checkState(page instanceof DataPage, "Found page %s instead of a DataPage", page);
+ DataPage dataPage = (DataPage) page;
+ pageCount++;
+ bufferedValueCount += dataPage.getValueCount();
+ if (bufferedValueCount >= valueCount || pageCount >= maxPageCount) {
+ return getBufferedDataPages(valueCount, maxPageCount);
+ }
+ }
+
+ while (bufferedValueCount < valueCount && pageCount < maxPageCount && bufferedPageBytes < maxBufferedBytes && canAdvanceInput() && compressedPages.hasNext()) {
+ ownLastBufferedPage();
+ Page page = compressedPages.next();
+ checkState(page instanceof DataPage, "Found page %s instead of a DataPage", page);
+ addBufferedPage(page);
+ DataPage dataPage = (DataPage) page;
+ bufferedValueCount += dataPage.getValueCount();
+ pageCount++;
+ }
+ return getBufferedDataPages(valueCount, maxPageCount);
+ }
+
+ public long getRetainedPageBytes()
+ {
+ return bufferedPageRetainedBytes + currentPageRetainedBytes;
+ }
+
+ @VisibleForTesting
+ public int getDataPageReadCount()
+ {
+ return dataPageReadCount;
+ }
+
+ private Page peekCompressedPage()
+ {
+ if (bufferedPages.isEmpty()) {
+ addBufferedPage(compressedPages.next());
+ }
+ return bufferedPages.peekFirst().page();
+ }
+
+ private Page nextCompressedPage()
+ {
+ releaseCurrentPage();
+ if (!bufferedPages.isEmpty()) {
+ BufferedPage bufferedPage = bufferedPages.removeFirst();
+ bufferedPageBytes -= getPageBytes(bufferedPage.page());
+ bufferedPageRetainedBytes -= bufferedPage.retainedBytes();
+ currentPage = bufferedPage.page();
+ currentPageRetainedBytes = bufferedPage.retainedBytes();
+ currentPageOwned = bufferedPage.owned();
+ return currentPage;
+ }
+ currentPage = compressedPages.next();
+ currentPageOwned = false;
+ return currentPage;
+ }
+
+ public void releaseCurrentPage()
+ {
+ currentPage = null;
+ currentPageRetainedBytes = 0;
+ currentPageOwned = false;
+ }
+
+ public void close()
+ {
+ releaseCurrentPage();
+ bufferedPages.clear();
+ bufferedPageBytes = 0;
+ bufferedPageRetainedBytes = 0;
+ }
+
+ private void replaceCurrentPage(Page page, long retainedBytes, boolean owned)
+ {
+ currentPage = requireNonNull(page, "page is null");
+ currentPageRetainedBytes = retainedBytes;
+ currentPageOwned = owned;
+ }
+
+ private void addBufferedPage(Page page)
+ {
+ BufferedPage bufferedPage = new BufferedPage(page, 0, false);
+ bufferedPages.addLast(bufferedPage);
+ bufferedPageBytes += getPageBytes(page);
+ }
+
+ private void ownLastBufferedPage()
+ {
+ if (!bufferedPages.isEmpty()) {
+ BufferedPage bufferedPage = bufferedPages.removeLast();
+ if (!bufferedPage.owned()) {
+ BufferedPage ownedPage = copyPage(bufferedPage.page());
+ bufferedPageRetainedBytes += ownedPage.retainedBytes();
+ bufferedPage = ownedPage;
+ }
+ bufferedPages.addLast(bufferedPage);
+ }
+ }
+
+ private boolean canAdvanceInput()
+ {
+ return currentPage == null || currentPageOwned;
+ }
+
+ private List getBufferedDataPages(int valueCount, int maxPageCount)
+ {
+ ImmutableList.Builder pages = ImmutableList.builder();
+ int bufferedValueCount = 0;
+ int pageCount = 0;
+ for (BufferedPage bufferedPage : bufferedPages) {
+ Page page = bufferedPage.page();
+ checkState(page instanceof DataPage, "Found page %s instead of a DataPage", page);
+ DataPage dataPage = (DataPage) page;
+ pages.add(dataPage);
+ bufferedValueCount += dataPage.getValueCount();
+ pageCount++;
+ if (bufferedValueCount >= valueCount || pageCount >= maxPageCount) {
+ break;
+ }
+ }
+ return pages.build();
+ }
+
+ private static Slice copySlice(Slice slice)
+ {
+ return wrappedBuffer(slice.getBytes());
+ }
+
+ private PageData getPageData(Slice data, int uncompressedSize, boolean compressed)
+ throws IOException
+ {
+ if (compressed) {
+ Slice uncompressed = arePagesCompressed()
+ ? decompress(dataSourceId, codec, data, uncompressedSize)
+ : copySlice(data);
+ return new PageData(uncompressed, uncompressed.length(), true);
+ }
+ if (blockDecryptor.isPresent()) {
+ return new PageData(data, data.length(), true);
+ }
+ return new PageData(data, currentPageRetainedBytes, currentPageOwned);
+ }
+
+ private static BufferedPage copyPage(Page page)
+ {
+ if (page instanceof DataPageV1 dataPageV1) {
+ Slice data = copySlice(dataPageV1.getSlice());
+ return new BufferedPage(
+ new DataPageV1(
+ data,
+ dataPageV1.getValueCount(),
+ dataPageV1.getUncompressedSize(),
+ dataPageV1.getFirstRowIndex(),
+ dataPageV1.getRepetitionLevelEncoding(),
+ dataPageV1.getDefinitionLevelEncoding(),
+ dataPageV1.getValueEncoding(),
+ dataPageV1.getPageIndex()),
+ data.length(),
+ true);
+ }
+ if (page instanceof DataPageV2 dataPageV2) {
+ int repetitionLevelsLength = dataPageV2.getRepetitionLevels().length();
+ int definitionLevelsLength = dataPageV2.getDefinitionLevels().length();
+ int dataLength = dataPageV2.getSlice().length();
+ byte[] bytes = new byte[repetitionLevelsLength + definitionLevelsLength + dataLength];
+ dataPageV2.getRepetitionLevels().getBytes(0, bytes, 0, repetitionLevelsLength);
+ dataPageV2.getDefinitionLevels().getBytes(0, bytes, repetitionLevelsLength, definitionLevelsLength);
+ dataPageV2.getSlice().getBytes(0, bytes, repetitionLevelsLength + definitionLevelsLength, dataLength);
+ Slice data = wrappedBuffer(bytes);
+ return new BufferedPage(
+ new DataPageV2(
+ dataPageV2.getRowCount(),
+ dataPageV2.getNullCount(),
+ dataPageV2.getValueCount(),
+ data.slice(0, repetitionLevelsLength),
+ data.slice(repetitionLevelsLength, definitionLevelsLength),
+ dataPageV2.getDataEncoding(),
+ data.slice(repetitionLevelsLength + definitionLevelsLength, dataLength),
+ dataPageV2.getUncompressedSize(),
+ dataPageV2.getFirstRowIndex(),
+ dataPageV2.getStatistics(),
+ dataPageV2.isCompressed(),
+ dataPageV2.getPageIndex()),
+ bytes.length,
+ true);
+ }
+ if (page instanceof DictionaryPage dictionaryPage) {
+ Slice data = copySlice(dictionaryPage.getSlice());
+ return new BufferedPage(
+ new DictionaryPage(
+ data,
+ dictionaryPage.getUncompressedSize(),
+ dictionaryPage.getDictionarySize(),
+ dictionaryPage.getEncoding()),
+ data.length(),
+ true);
+ }
+ throw new IllegalArgumentException("Unsupported page: " + page);
+ }
+
+ private static long getPageBytes(Page page)
+ {
+ if (page instanceof DataPageV2 dataPageV2) {
+ return (long) dataPageV2.getRepetitionLevels().length()
+ + dataPageV2.getDefinitionLevels().length()
+ + dataPageV2.getSlice().length();
+ }
+ return page.getSlice().length();
+ }
+
+ private record BufferedPage(Page page, long retainedBytes, boolean owned)
+ {
+ private BufferedPage
+ {
+ requireNonNull(page, "page is null");
+ checkArgument(retainedBytes >= 0, "retainedBytes is negative");
+ checkArgument(owned || retainedBytes == 0, "borrowed page has retained bytes");
+ }
+ }
+
+ private record PageData(Slice data, long retainedBytes, boolean owned)
+ {
+ private PageData
+ {
+ requireNonNull(data, "data is null");
+ checkArgument(retainedBytes >= 0, "retainedBytes is negative");
+ checkArgument(owned || retainedBytes == 0, "borrowed data has retained bytes");
+ }
}
- public boolean arePagesCompressed()
+ private boolean arePagesCompressed()
{
return codec != CompressionCodec.UNCOMPRESSED;
}
diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java
index fcaeb964f446..5cbe3115e86b 100644
--- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java
+++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java
@@ -13,6 +13,7 @@
*/
package io.trino.parquet.reader;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
@@ -55,6 +56,7 @@
import io.trino.spi.metrics.Metric;
import io.trino.spi.metrics.Metrics;
import io.trino.spi.type.ArrayType;
+import io.trino.spi.type.DecimalType;
import io.trino.spi.type.MapType;
import io.trino.spi.type.RowType;
import jakarta.annotation.Nullable;
@@ -65,6 +67,7 @@
import org.apache.parquet.internal.column.columnindex.OffsetIndex;
import org.apache.parquet.internal.filter2.columnindex.ColumnIndexFilter;
import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore;
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName;
import org.joda.time.DateTimeZone;
import java.io.Closeable;
@@ -88,6 +91,7 @@
import static io.airlift.slice.SizeOf.instanceSize;
import static io.airlift.slice.SizeOf.sizeOf;
import static io.airlift.slice.Slices.utf8Slice;
+import static io.trino.parquet.ParquetReaderUtils.isOnlyDictionaryEncodingPages;
import static io.trino.parquet.ParquetValidationUtils.validateParquet;
import static io.trino.parquet.ParquetWriteValidation.StatisticsValidation.createStatisticsValidationBuilder;
import static io.trino.parquet.ParquetWriteValidation.WriteChecksumBuilder.createWriteChecksumBuilder;
@@ -101,8 +105,10 @@
import static java.lang.Math.min;
import static java.lang.Math.toIntExact;
import static java.lang.String.format;
+import static java.util.Objects.checkFromIndexSize;
import static java.util.Objects.checkIndex;
import static java.util.Objects.requireNonNull;
+import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.BINARY;
public class ParquetReader
implements Closeable
@@ -111,8 +117,19 @@ public class ParquetReader
private static final int INITIAL_BATCH_SIZE = 1;
private static final int BATCH_SIZE_GROWTH_FACTOR = 2;
+ private static final int MAX_FIXED_WIDTH_SELECTED_PERCENTAGE = 5;
+ private static final int MAX_NULLABLE_FIXED_WIDTH_SELECTED_PERCENTAGE = 20;
+ private static final int MAX_DICTIONARY_SELECTED_PERCENTAGE = 10;
+ private static final int MAX_BINARY_SELECTED_PERCENTAGE = 25;
+ private static final long VERY_WIDE_BINARY_BYTES = 1_024;
+ private static final int MAX_VERY_WIDE_BINARY_PROJECTED_COLUMNS = 4;
+ private static final int MIN_SKIPPED_POSITIONS_PER_RUN = 300;
+ private static final long MIN_SKIPPED_PAGE_BYTES = 4 * 1024;
+ private static final int MAX_PAGE_AWARE_SELECTED_PERCENTAGE = 75;
+ private static final int MAX_RETAINED_TO_LOGICAL_SIZE_RATIO = 2;
public static final String PARQUET_CODEC_METRIC_PREFIX = "ParquetReaderCompressionFormat_";
public static final String COLUMN_INDEX_ROWS_FILTERED = "ParquetColumnIndexRowsFiltered";
+ public static final String SELECTED_POSITIONS_PUSHDOWNS = "ParquetSelectedPositionsPushdowns";
private final Optional fileCreatedBy;
private final List rowGroups;
@@ -142,6 +159,7 @@ public class ParquetReader
private int batchSize;
private int nextBatchSize = INITIAL_BATCH_SIZE;
private final Map columnReaders;
+ private final Map selectedPositionsPushdownCharacteristics = new HashMap<>();
private final Map maxBytesPerCell;
private double maxCombinedBytesPerRow;
private final ParquetReaderOptions options;
@@ -157,9 +175,13 @@ public class ParquetReader
private final Map> codecMetrics;
private int currentPageId;
+ private int completedRowGroupDataPageReadCount;
+ private int selectedPositionsFallbackCount;
+ private int selectedPositionsPushdownCount;
private long columnIndexRowsFiltered = -1;
private final Optional decryptionContext;
+ private final boolean forceSelectedPositionsPushdown;
public ParquetReader(
Optional fileCreatedBy,
@@ -175,6 +197,38 @@ public ParquetReader(
Optional writeValidation,
Optional decryptionContext)
throws IOException
+ {
+ this(fileCreatedBy,
+ columnFields,
+ appendRowNumberColumn,
+ rowGroups,
+ dataSource,
+ timeZone,
+ memoryContext,
+ options,
+ exceptionTransform,
+ parquetPredicate,
+ writeValidation,
+ decryptionContext,
+ false);
+ }
+
+ @VisibleForTesting
+ public ParquetReader(
+ Optional fileCreatedBy,
+ List columnFields,
+ boolean appendRowNumberColumn,
+ List rowGroups,
+ ParquetDataSource dataSource,
+ DateTimeZone timeZone,
+ AggregatedMemoryContext memoryContext,
+ ParquetReaderOptions options,
+ Function exceptionTransform,
+ Optional parquetPredicate,
+ Optional writeValidation,
+ Optional decryptionContext,
+ boolean forceSelectedPositionsPushdown)
+ throws IOException
{
this.fileCreatedBy = requireNonNull(fileCreatedBy, "fileCreatedBy is null");
requireNonNull(columnFields, "columnFields is null");
@@ -192,6 +246,7 @@ public ParquetReader(
this.columnReaders = new HashMap<>();
this.maxBytesPerCell = new HashMap<>();
this.decryptionContext = requireNonNull(decryptionContext, "decryptionContext is null");
+ this.forceSelectedPositionsPushdown = forceSelectedPositionsPushdown;
this.writeValidation = requireNonNull(writeValidation, "writeValidation is null");
validateWrite(
@@ -259,12 +314,13 @@ public void close()
throws IOException
{
// Release memory usage from column readers
+ columnReaders.values().forEach(ColumnReader::close);
columnReaders.clear();
- currentRowGroupMemoryContext.close();
for (ChunkedInputStream chunkedInputStream : chunkReaders.values()) {
chunkedInputStream.close();
}
+ currentRowGroupMemoryContext.close();
dataSource.close();
if (writeChecksumBuilder.isPresent()) {
@@ -298,13 +354,17 @@ private class ParquetSourcePage
private final long batchStartRow = lastBatchStartRow();
private final long[] batchRowNumbers = currentBatchRowNumbers;
private SelectedPositions selectedPositions;
+ @Nullable
+ private SelectedPositionsReadMode[] selectedPositionsReadModes;
+ private int unloadedSelectedColumns;
+ private boolean selectedPositionsPushedDown;
private long sizeInBytes;
private long retainedSizeInBytes;
public ParquetSourcePage(int positionCount)
{
- selectedPositions = new SelectedPositions(positionCount, null);
+ selectedPositions = SelectedPositions.allPositions(positionCount);
retainedSizeInBytes = shallowRetainedSizeInBytes();
}
@@ -331,6 +391,7 @@ private long shallowRetainedSizeInBytes()
return INSTANCE_SIZE +
sizeOf(blocks) +
sizeOf(batchRowNumbers) +
+ sizeOf(selectedPositionsReadModes) +
selectedPositions.retainedSizeInBytes();
}
@@ -342,6 +403,9 @@ public void retainedBytesForEachPart(ObjLongConsumer