From 1b107ff69537cc69ef595d4c857a1d4300094a37 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Fri, 3 May 2024 12:58:40 -0700 Subject: [PATCH] QueryableIndex: Close columns after failed vector cursor setup. (#16365) * QueryableIndex: Close columns after failed vector cursor setup. If anything fails while setting up a vector cursor, the prior code in QueryableIndex would not close its ColumnCache and would therefore leak columns. Columns often contain references to buffers that must be closed. * Fix style. --- .../QueryableIndexCursorSequenceBuilder.java | 110 ++++++++++-------- 1 file changed, 59 insertions(+), 51 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorSequenceBuilder.java b/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorSequenceBuilder.java index 77bb40f946b6..86069068eaf9 100644 --- a/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorSequenceBuilder.java +++ b/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorSequenceBuilder.java @@ -52,6 +52,7 @@ import org.apache.druid.segment.vector.VectorColumnSelectorFactory; import org.apache.druid.segment.vector.VectorCursor; import org.apache.druid.segment.vector.VectorOffset; +import org.apache.druid.utils.CloseableUtils; import org.joda.time.DateTime; import org.joda.time.Interval; @@ -207,69 +208,76 @@ public VectorCursor buildVectorized(final int vectorSize) final Closer closer = Closer.create(); final ColumnCache columnCache = new ColumnCache(index, closer); - final ColumnSelectorColumnIndexSelector bitmapIndexSelector = new ColumnSelectorColumnIndexSelector( - index.getBitmapFactoryForDimensions(), - virtualColumns, - columnCache - ); - - final int numRows = index.getNumRows(); - final FilterBundle filterBundle = makeFilterBundle(bitmapIndexSelector, numRows); + // Wrap the remainder of cursor setup in a try, so if an error is encountered while setting it up, we don't + // leak columns in the ColumnCache. + try { + final ColumnSelectorColumnIndexSelector bitmapIndexSelector = new ColumnSelectorColumnIndexSelector( + index.getBitmapFactoryForDimensions(), + virtualColumns, + columnCache + ); - NumericColumn timestamps = null; + final int numRows = index.getNumRows(); + final FilterBundle filterBundle = makeFilterBundle(bitmapIndexSelector, numRows); - final int startOffset; - final int endOffset; + NumericColumn timestamps = null; - if (interval.getStartMillis() > minDataTimestamp) { - timestamps = (NumericColumn) columnCache.getColumn(ColumnHolder.TIME_COLUMN_NAME); + final int startOffset; + final int endOffset; - startOffset = timeSearch(timestamps, interval.getStartMillis(), 0, index.getNumRows()); - } else { - startOffset = 0; - } - - if (interval.getEndMillis() <= maxDataTimestamp) { - if (timestamps == null) { + if (interval.getStartMillis() > minDataTimestamp) { timestamps = (NumericColumn) columnCache.getColumn(ColumnHolder.TIME_COLUMN_NAME); - } - endOffset = timeSearch(timestamps, interval.getEndMillis(), startOffset, index.getNumRows()); - } else { - endOffset = index.getNumRows(); - } + startOffset = timeSearch(timestamps, interval.getStartMillis(), 0, index.getNumRows()); + } else { + startOffset = 0; + } - // filterBundle will only be null if the filter itself is null, otherwise check to see if the filter can use - // an index - final VectorOffset baseOffset = - filterBundle == null || filterBundle.getIndex() == null - ? new NoFilterVectorOffset(vectorSize, startOffset, endOffset) - : new BitmapVectorOffset(vectorSize, filterBundle.getIndex().getBitmap(), startOffset, endOffset); + if (interval.getEndMillis() <= maxDataTimestamp) { + if (timestamps == null) { + timestamps = (NumericColumn) columnCache.getColumn(ColumnHolder.TIME_COLUMN_NAME); + } - // baseColumnSelectorFactory using baseOffset is the column selector for filtering. - final VectorColumnSelectorFactory baseColumnSelectorFactory = makeVectorColumnSelectorFactoryForOffset( - columnCache, - baseOffset - ); + endOffset = timeSearch(timestamps, interval.getEndMillis(), startOffset, index.getNumRows()); + } else { + endOffset = index.getNumRows(); + } - // filterBundle will only be null if the filter itself is null, otherwise check to see if the filter needs to use - // a value matcher - if (filterBundle != null && filterBundle.getMatcherBundle() != null) { - final VectorValueMatcher vectorValueMatcher = filterBundle.getMatcherBundle() - .vectorMatcher(baseColumnSelectorFactory, baseOffset); - final VectorOffset filteredOffset = FilteredVectorOffset.create( - baseOffset, - vectorValueMatcher - ); + // filterBundle will only be null if the filter itself is null, otherwise check to see if the filter can use + // an index + final VectorOffset baseOffset = + filterBundle == null || filterBundle.getIndex() == null + ? new NoFilterVectorOffset(vectorSize, startOffset, endOffset) + : new BitmapVectorOffset(vectorSize, filterBundle.getIndex().getBitmap(), startOffset, endOffset); - // Now create the cursor and column selector that will be returned to the caller. - final VectorColumnSelectorFactory filteredColumnSelectorFactory = makeVectorColumnSelectorFactoryForOffset( + // baseColumnSelectorFactory using baseOffset is the column selector for filtering. + final VectorColumnSelectorFactory baseColumnSelectorFactory = makeVectorColumnSelectorFactoryForOffset( columnCache, - filteredOffset + baseOffset ); - return new QueryableIndexVectorCursor(filteredColumnSelectorFactory, filteredOffset, vectorSize, closer); - } else { - return new QueryableIndexVectorCursor(baseColumnSelectorFactory, baseOffset, vectorSize, closer); + + // filterBundle will only be null if the filter itself is null, otherwise check to see if the filter needs to use + // a value matcher + if (filterBundle != null && filterBundle.getMatcherBundle() != null) { + final VectorValueMatcher vectorValueMatcher = filterBundle.getMatcherBundle() + .vectorMatcher(baseColumnSelectorFactory, baseOffset); + final VectorOffset filteredOffset = FilteredVectorOffset.create( + baseOffset, + vectorValueMatcher + ); + + // Now create the cursor and column selector that will be returned to the caller. + final VectorColumnSelectorFactory filteredColumnSelectorFactory = makeVectorColumnSelectorFactoryForOffset( + columnCache, + filteredOffset + ); + return new QueryableIndexVectorCursor(filteredColumnSelectorFactory, filteredOffset, vectorSize, closer); + } else { + return new QueryableIndexVectorCursor(baseColumnSelectorFactory, baseOffset, vectorSize, closer); + } + } + catch (Throwable t) { + throw CloseableUtils.closeAndWrapInCatch(t, closer); } }