This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 7dffad77c18 Reduce memory footprint and make FST regexp matching
interruptible (#18754)
7dffad77c18 is described below
commit 7dffad77c18dfcdc3520a7d3e160059a4ed1d490
Author: Xiang Fu <[email protected]>
AuthorDate: Mon Jun 15 00:57:50 2026 -0700
Reduce memory footprint and make FST regexp matching interruptible (#18754)
* Reduce memory footprint and make FST regexp matching interruptible
REGEXP_LIKE backed by an FST index walks the query automaton over the FST
via RegexpMatcher / RegexpMatcherCaseInsensitive. On a high-cardinality,
long-string column a broad regexp (anything containing `.*`) visits a huge
number of paths, and the traversal had three problems that together turned
memory pressure into server-wide heap OOMs:
1. Every pushed path carried a full copy of its accumulated input prefix in
an
IntsRefBuilder (newInput.copyInts(...) per arc), an O(paths x key-length)
allocation. The prefix was write-only -- it was never read to produce the
result -- so it was pure garbage. Removed the field.
2. Matches were retained as Path objects in an `endNodes` list until the
whole
traversal finished, then materialized into a boxed List<Long>, then
copied
into the result bitmap -- holding the entire match set (and copied FST
arcs)
live on the heap. Matches are now emitted to an IntConsumer as soon as a
final state is reached, so the readers add straight into the
RoaringBitmap.
3. The traversal loop never checked for query termination, so the
OOM-protection
framework (and query deadline / interrupt) could not stop it. All query
threads stayed pinned in the loop and the JVM was killed instead of the
offending query. The loop now calls
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(...), and
the
readers let QueryException propagate unwrapped so the kill/timeout is
honored.
JMH (BenchmarkFSTRegexpMatcher, 500k keys x 40 chars), allocation per op:
FST `.*` 206 MB -> 67 MB (-67%), 62.7 ms -> 18.7 ms
FST `.*9999` 193 MB -> 67 MB (-66%), 27.2 ms -> 19.7 ms
IFST `.*` 299 MB -> 134 MB (-55%), 74.8 ms -> 30.0 ms
IFST `.*9999` 234 MB -> 108 MB (-54%)
Selective queries (e.g. `key0001.*`) are unchanged.
Added regression tests asserting full enumeration through the new consumer
path
and that an interrupt mid-match surfaces EarlyTerminationException both
from the
matcher and through LuceneFSTIndexReader/LuceneIFSTIndexReader#getDictIds.
* Close off-heap FST buffers in benchmark teardown
The readers' close() is a no-op, so the PinotDataBuffer instances returned
by
loadBigEndianFile own the mapped memory and must be closed explicitly;
otherwise the mapping leaks and can block temp-dir deletion.
---
.../pinot/perf/BenchmarkFSTRegexpMatcher.java | 160 +++++++++++++++++++++
.../index/readers/LuceneFSTIndexReader.java | 10 +-
.../index/readers/LuceneIFSTIndexReader.java | 10 +-
.../segment/local/utils/fst/RegexpMatcher.java | 121 +++++++---------
.../utils/fst/RegexpMatcherCaseInsensitive.java | 152 +++++++++-----------
.../segment/local/utils/fst/FSTBuilderTest.java | 40 ++++++
.../segment/local/utils/fst/IFSTBuilderTest.java | 38 +++++
7 files changed, 368 insertions(+), 163 deletions(-)
diff --git
a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkFSTRegexpMatcher.java
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkFSTRegexpMatcher.java
new file mode 100644
index 00000000000..15c900a1e2d
--- /dev/null
+++
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkFSTRegexpMatcher.java
@@ -0,0 +1,160 @@
+/**
+ * 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.pinot.perf;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.SortedMap;
+import java.util.TreeMap;
+import java.util.concurrent.TimeUnit;
+import org.apache.commons.io.FileUtils;
+import org.apache.lucene.store.OutputStreamDataOutput;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.fst.FST;
+import
org.apache.pinot.segment.local.segment.index.readers.LuceneFSTIndexReader;
+import
org.apache.pinot.segment.local.segment.index.readers.LuceneIFSTIndexReader;
+import org.apache.pinot.segment.local.utils.fst.FSTBuilder;
+import org.apache.pinot.segment.local.utils.fst.IFSTBuilder;
+import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.options.ChainedOptionsBuilder;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+
+/**
+ * Benchmark for FST/IFST regexp matching (REGEXP_LIKE backed by an FST index).
+ *
+ * This exercises {@code LuceneFSTIndexReader#getDictIds} / {@code
LuceneIFSTIndexReader#getDictIds}, the production
+ * path used by the FST-based regexp predicate evaluator. It is designed to
highlight the memory footprint of the
+ * automaton-over-FST traversal for high-cardinality, long-string columns (the
case that caused server heap OOMs).
+ *
+ * Run with the GC profiler to see the allocation reduction, e.g.:
+ * java -cp 'pinot-perf/target/pinot-perf-pkg/lib/*' org.openjdk.jmh.Main \
+ * org.apache.pinot.perf.BenchmarkFSTRegexpMatcher -prof gc -wi 2 -i 3 -f
1 -r 5s -w 5s
+ */
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Fork(1)
+@Warmup(iterations = 3, time = 5)
+@Measurement(iterations = 5, time = 5)
+@State(Scope.Benchmark)
+public class BenchmarkFSTRegexpMatcher {
+
+ /// Number of distinct keys (column cardinality).
+ @Param({"500000"})
+ public int _cardinality;
+
+ /// Length of the key prefix padding, simulating a "long string" column.
+ @Param({"40"})
+ public int _keyLength;
+
+ /// Regexp patterns: a broad `.*` (worst case, matches everything), a
selective prefix, and a selective suffix.
+ @Param({".*", "key0001.*", ".*9999"})
+ public String _regex;
+
+ private File _tempDir;
+ private PinotDataBuffer _fstBuffer;
+ private PinotDataBuffer _ifstBuffer;
+ private LuceneFSTIndexReader _fstReader;
+ private LuceneIFSTIndexReader _ifstReader;
+
+ @Setup(Level.Trial)
+ public void setUp()
+ throws IOException {
+ _tempDir = new File(FileUtils.getTempDirectory(),
"BenchmarkFSTRegexpMatcher");
+ FileUtils.deleteDirectory(_tempDir);
+ FileUtils.forceMkdir(_tempDir);
+
+ // Build a high-cardinality dictionary of long string keys: "key" +
zero-padded id, padded to _keyLength.
+ SortedMap<String, Integer> entries = new TreeMap<>();
+ for (int i = 0; i < _cardinality; i++) {
+ String key = String.format("key%0" + (_keyLength - 3) + "d", i);
+ entries.put(key, i);
+ }
+
+ FST<Long> fst = FSTBuilder.buildFST(entries);
+ File fstFile = new File(_tempDir, "fst.lucene");
+ try (FileOutputStream outputStream = new FileOutputStream(fstFile);
+ OutputStreamDataOutput dataOutput = new
OutputStreamDataOutput(outputStream)) {
+ fst.save(dataOutput, dataOutput);
+ }
+ _fstBuffer = PinotDataBuffer.loadBigEndianFile(fstFile);
+ _fstReader = new LuceneFSTIndexReader(_fstBuffer);
+
+ FST<BytesRef> ifst = IFSTBuilder.buildIFST(entries);
+ File ifstFile = new File(_tempDir, "ifst.lucene");
+ try (FileOutputStream outputStream = new FileOutputStream(ifstFile);
+ OutputStreamDataOutput dataOutput = new
OutputStreamDataOutput(outputStream)) {
+ ifst.save(dataOutput, dataOutput);
+ }
+ _ifstBuffer = PinotDataBuffer.loadBigEndianFile(ifstFile);
+ _ifstReader = new LuceneIFSTIndexReader(_ifstBuffer);
+ }
+
+ @TearDown(Level.Trial)
+ public void tearDown()
+ throws IOException {
+ if (_fstReader != null) {
+ _fstReader.close();
+ }
+ if (_ifstReader != null) {
+ _ifstReader.close();
+ }
+ // The readers' close() is a no-op; the off-heap buffers own the mapped
memory and must be closed explicitly,
+ // otherwise the mapping stays live and can prevent the temp dir from
being deleted.
+ if (_fstBuffer != null) {
+ _fstBuffer.close();
+ }
+ if (_ifstBuffer != null) {
+ _ifstBuffer.close();
+ }
+ FileUtils.deleteDirectory(_tempDir);
+ }
+
+ @Benchmark
+ public void fstGetDictIds(Blackhole blackhole) {
+ blackhole.consume(_fstReader.getDictIds(_regex));
+ }
+
+ @Benchmark
+ public void ifstGetDictIds(Blackhole blackhole) {
+ blackhole.consume(_ifstReader.getDictIds(_regex));
+ }
+
+ public static void main(String[] args)
+ throws Exception {
+ ChainedOptionsBuilder opt = new
OptionsBuilder().include(BenchmarkFSTRegexpMatcher.class.getSimpleName());
+ new Runner(opt.build()).run();
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneFSTIndexReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneFSTIndexReader.java
index 324d15a0f3e..1f81f0c25e6 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneFSTIndexReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneFSTIndexReader.java
@@ -19,7 +19,6 @@
package org.apache.pinot.segment.local.segment.index.readers;
import java.io.IOException;
-import java.util.List;
import org.apache.lucene.util.fst.FST;
import org.apache.lucene.util.fst.OffHeapFSTStore;
import org.apache.lucene.util.fst.PositiveIntOutputs;
@@ -27,6 +26,7 @@ import
org.apache.pinot.segment.local.utils.fst.PinotBufferIndexInput;
import org.apache.pinot.segment.local.utils.fst.RegexpMatcher;
import org.apache.pinot.segment.spi.index.reader.TextIndexReader;
import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
+import org.apache.pinot.spi.exception.QueryException;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.slf4j.Logger;
@@ -61,11 +61,11 @@ public class LuceneFSTIndexReader implements
TextIndexReader {
public ImmutableRoaringBitmap getDictIds(String searchQuery) {
try {
MutableRoaringBitmap dictIds = new MutableRoaringBitmap();
- List<Long> matchingIds = RegexpMatcher.regexMatch(searchQuery, _fst);
- for (Long matchingId : matchingIds) {
- dictIds.add(matchingId.intValue());
- }
+ RegexpMatcher.regexMatch(searchQuery, _fst, dictIds::add);
return dictIds.toImmutableRoaringBitmap();
+ } catch (QueryException ex) {
+ // Let query termination exceptions (timeout, OOM-protection kill)
propagate as-is.
+ throw ex;
} catch (Exception ex) {
LOGGER.error("Error getting matching Ids from FST", ex);
throw new RuntimeException(ex);
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneIFSTIndexReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneIFSTIndexReader.java
index 3edacf009c7..fa1114051bc 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneIFSTIndexReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/LuceneIFSTIndexReader.java
@@ -19,7 +19,6 @@
package org.apache.pinot.segment.local.segment.index.readers;
import java.io.IOException;
-import java.util.List;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.fst.ByteSequenceOutputs;
import org.apache.lucene.util.fst.FST;
@@ -28,6 +27,7 @@ import
org.apache.pinot.segment.local.utils.fst.PinotBufferIndexInput;
import org.apache.pinot.segment.local.utils.fst.RegexpMatcherCaseInsensitive;
import org.apache.pinot.segment.spi.index.reader.TextIndexReader;
import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
+import org.apache.pinot.spi.exception.QueryException;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.slf4j.Logger;
@@ -62,11 +62,11 @@ public class LuceneIFSTIndexReader implements
TextIndexReader {
public ImmutableRoaringBitmap getDictIds(String searchQuery) {
try {
MutableRoaringBitmap dictIds = new MutableRoaringBitmap();
- List<Long> matchingIds =
RegexpMatcherCaseInsensitive.regexMatch(searchQuery, _ifst);
- for (Long matchingId : matchingIds) {
- dictIds.add(matchingId.intValue());
- }
+ RegexpMatcherCaseInsensitive.regexMatch(searchQuery, _ifst,
dictIds::add);
return dictIds.toImmutableRoaringBitmap();
+ } catch (QueryException ex) {
+ // Let query termination exceptions (timeout, OOM-protection kill)
propagate as-is.
+ throw ex;
} catch (Exception ex) {
LOGGER.error("Error getting matching Ids from FST", ex);
throw new RuntimeException(ex);
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcher.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcher.java
index 8afb8ce6db6..938e52140e0 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcher.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcher.java
@@ -20,25 +20,25 @@ package org.apache.pinot.segment.local.utils.fst;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.Collections;
import java.util.List;
-import org.apache.lucene.util.IntsRefBuilder;
+import java.util.function.IntConsumer;
import org.apache.lucene.util.automaton.Automaton;
import org.apache.lucene.util.automaton.CharacterRunAutomaton;
import org.apache.lucene.util.automaton.RegExp;
import org.apache.lucene.util.automaton.Transition;
import org.apache.lucene.util.fst.FST;
import org.apache.lucene.util.fst.Util;
+import org.apache.pinot.spi.query.QueryThreadContext;
-/**
- * RegexpMatcher is a helper to retrieve matching values for a given regexp
query.
- * Regexp query is converted into an automaton and we run the matching
algorithm on FST.
- *
- * Two main functions of this class are
- * regexMatchOnFST() Function runs matching on FST (See function comments
for more details)
- * match(input) Function builds the automaton and matches given input.
- */
+/// RegexpMatcher is a helper to retrieve matching values for a given regexp
query.
+/// Regexp query is converted into an automaton, and we run the matching
algorithm on FST.
+///
+/// Two main functions of this class are
+/// regexMatchOnFST() Function runs matching on FST (See function comments
for more details)
+/// match(input) Function builds the automaton and matches given input.
+///
+/// This class is not thread-safe. Create a new instance (or use the static
helper) per query.
public class RegexpMatcher {
private final FST<Long> _fst;
private final Automaton _automaton;
@@ -48,9 +48,12 @@ public class RegexpMatcher {
_automaton = (new RegExp(regexQuery)).toAutomaton();
}
- public static List<Long> regexMatch(String regexQuery, FST<Long> fst)
+ /// Runs matching of the given regexp query on the FST, emitting the value
(dict id) of each matched entry to the
+ /// given consumer. Values are emitted as the traversal hits final states,
so the caller can collect them without
+ /// this class materializing the full result set in memory.
+ public static void regexMatch(String regexQuery, FST<Long> fst, IntConsumer
dictIdConsumer)
throws IOException {
- return new RegexpMatcher(regexQuery, fst).regexMatchOnFST();
+ new RegexpMatcher(regexQuery, fst).regexMatchOnFST(dictIdConsumer);
}
// Matches "input" string with _regexQuery Automaton.
@@ -59,97 +62,79 @@ public class RegexpMatcher {
return characterRunAutomaton.run(input);
}
- /**
- * This function runs matching on automaton built from regexQuery and the
FST.
- * FST stores key (string) to a value (Long). Both are state machines and
state transition is based on
- * a input character.
- *
- * This algorithm starts with Queue containing (Automaton Start Node, FST
Start Node).
- * Each step an entry is popped from the queue:
- * 1) if the automaton state is accept and the FST Node is final (i.e.
end node) then the value stored for that FST
- * is added to the set of result.
- * 2) Else next set of transitions on automaton are gathered and for each
transition target node for that character
- * is figured out in FST Node, resulting pair of (automaton state, fst
node) are added to the queue.
- * 3) This process is bound to complete since we are making progression
on the FST (which is a DAG) towards final
- * nodes.
- * @return
- * @throws IOException
- */
- public List<Long> regexMatchOnFST()
+ /// This function runs matching on automaton built from regexQuery and the
FST.
+ /// FST stores key (string) to a value (Long). Both are state machines and
state transition is based on an input
+ /// character.
+ ///
+ /// This algorithm starts with a stack containing (Automaton Start Node, FST
Start Node).
+ /// Each step an entry is popped from the stack (DFS order to bound the
frontier size):
+ /// 1) if the automaton state is accept and the FST Node is final (i.e.
end node) then the value stored for that
+ /// FST node is emitted to the consumer.
+ /// 2) Else next set of transitions on automaton are gathered and for each
transition target node for that
+ /// character is figured out in FST Node, resulting pair of (automaton
state, fst node) are added to the stack.
+ /// 3) This process is bound to complete since we are making progression
on the FST (which is a DAG) towards final
+ /// nodes.
+ ///
+ /// A broad regexp on a high-cardinality column can visit a huge number of
paths, so the loop periodically checks
+ /// for query termination (timeout or OOM-protection kill) to remain
interruptible.
+ public void regexMatchOnFST(IntConsumer dictIdConsumer)
throws IOException {
- final List<Path<Long>> queue = new ArrayList<>();
- final List<Path<Long>> endNodes = new ArrayList<>();
if (_automaton.getNumStates() == 0) {
- return Collections.emptyList();
+ return;
}
- // Automaton start state and FST start node is added to the queue.
- queue.add(new Path<>(0, _fst.getFirstArc(new FST.Arc<Long>()),
_fst.outputs.getNoOutput(), new IntsRefBuilder()));
+ // Automaton start state and FST start node is added to the stack.
+ List<Path<Long>> stack = new ArrayList<>();
+ stack.add(new Path<>(0, _fst.getFirstArc(new FST.Arc<>()),
_fst.outputs.getNoOutput()));
- final FST.Arc<Long> scratchArc = new FST.Arc<>();
- final FST.BytesReader fstReader = _fst.getBytesReader();
+ FST.Arc<Long> scratchArc = new FST.Arc<>();
+ FST.BytesReader fstReader = _fst.getBytesReader();
Transition t = new Transition();
- while (!queue.isEmpty()) {
- final Path<Long> path = queue.remove(queue.size() - 1);
+ int numPathsProcessed = 0;
+ while (!stack.isEmpty()) {
+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(numPathsProcessed++,
+ "RegexpMatcher#regexMatchOnFST");
+ Path<Long> path = stack.remove(stack.size() - 1);
- // If automaton is in accept state and the fstNode is final (i.e. end
node) then add the entry to endNodes which
- // contains the result set.
- if (_automaton.isAccept(path._state)) {
- if (path._fstNode.isFinal()) {
- endNodes.add(path);
- }
+ // If automaton is in accept state and the fstNode is final (i.e. end
node) then emit the matched value.
+ if (_automaton.isAccept(path._state) && path._fstNode.isFinal()) {
+ dictIdConsumer.accept(path._output.intValue());
}
// Gather next set of transitions on automaton and find target nodes in
FST.
- IntsRefBuilder currentInput = path._input;
int count = _automaton.initTransition(path._state, t);
for (int i = 0; i < count; i++) {
_automaton.getNextTransition(t);
- final int min = t.min;
- final int max = t.max;
+ int min = t.min;
+ int max = t.max;
if (min == max) {
- final FST.Arc<Long> nextArc = _fst.findTargetArc(t.min,
path._fstNode, scratchArc, fstReader);
+ FST.Arc<Long> nextArc = _fst.findTargetArc(t.min, path._fstNode,
scratchArc, fstReader);
if (nextArc != null) {
- final IntsRefBuilder newInput = new IntsRefBuilder();
- newInput.copyInts(currentInput.get());
- newInput.append(t.min);
- queue.add(new Path<Long>(t.dest, new
FST.Arc<Long>().copyFrom(nextArc),
- _fst.outputs.add(path._output, nextArc.output()), newInput));
+ stack.add(new Path<>(t.dest, new FST.Arc<Long>().copyFrom(nextArc),
+ _fst.outputs.add(path._output, nextArc.output())));
}
} else {
FST.Arc<Long> nextArc = Util.readCeilArc(min, _fst, path._fstNode,
scratchArc, fstReader);
while (nextArc != null && nextArc.label() <= max) {
- final IntsRefBuilder newInput = new IntsRefBuilder();
- newInput.copyInts(currentInput.get());
- newInput.append(nextArc.label());
- queue.add(new Path<>(t.dest, new FST.Arc<Long>().copyFrom(nextArc),
- _fst.outputs.add(path._output, nextArc.output()), newInput));
+ stack.add(new Path<>(t.dest, new FST.Arc<Long>().copyFrom(nextArc),
+ _fst.outputs.add(path._output, nextArc.output())));
nextArc = nextArc.isLast() ? null : _fst.readNextRealArc(nextArc,
fstReader);
}
}
}
}
-
- // From the result set of matched entries gather the values stored and
return.
- ArrayList<Long> matchedIds = new ArrayList<>();
- for (Path<Long> path : endNodes) {
- matchedIds.add(path._output);
- }
- return matchedIds;
}
public static final class Path<T> {
public final int _state;
public final FST.Arc<T> _fstNode;
public final T _output;
- public final IntsRefBuilder _input;
- public Path(int state, FST.Arc<T> fstNode, T output, IntsRefBuilder input)
{
+ public Path(int state, FST.Arc<T> fstNode, T output) {
_state = state;
_fstNode = fstNode;
_output = output;
- _input = input;
}
}
}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcherCaseInsensitive.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcherCaseInsensitive.java
index 03fb87135a2..6e4fb356aeb 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcherCaseInsensitive.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/fst/RegexpMatcherCaseInsensitive.java
@@ -20,25 +20,26 @@ package org.apache.pinot.segment.local.utils.fst;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.Collections;
import java.util.List;
+import java.util.function.IntConsumer;
import org.apache.lucene.util.BytesRef;
-import org.apache.lucene.util.IntsRefBuilder;
import org.apache.lucene.util.automaton.Automaton;
import org.apache.lucene.util.automaton.CharacterRunAutomaton;
import org.apache.lucene.util.automaton.RegExp;
import org.apache.lucene.util.automaton.Transition;
import org.apache.lucene.util.fst.FST;
+import org.apache.lucene.util.fst.Util;
+import org.apache.pinot.spi.query.QueryThreadContext;
-/**
- * RegexpMatcherCaseInsensitive is a specialized helper to retrieve matching
values for case-insensitive regexp queries.
- * This class handles only case-insensitive matching using FST<BytesRef> type.
- *
- * The regexp query is converted to lowercase and then into an automaton for
matching against the FST.
- * This class is separate from RegexpMatcher because case-insensitive logic is
fundamentally different
- * from case-sensitive logic and doesn't share code.
- */
+/// RegexpMatcherCaseInsensitive is a specialized helper to retrieve matching
values for case-insensitive regexp
+/// queries. This class handles only case-insensitive matching using
`FST<BytesRef>` type.
+///
+/// The regexp query is converted to lowercase and then into an automaton for
matching against the FST.
+/// This class is separate from RegexpMatcher because case-insensitive logic
is fundamentally different
+/// from case-sensitive logic and doesn't share code.
+///
+/// This class is not thread-safe. Create a new instance (or use the static
helper) per query.
public class RegexpMatcherCaseInsensitive {
private final FST<BytesRef> _ifst;
private final Automaton _automaton;
@@ -49,9 +50,12 @@ public class RegexpMatcherCaseInsensitive {
_automaton = (new RegExp(regexQuery.toLowerCase())).toAutomaton();
}
- public static List<Long> regexMatch(String regexQuery, FST<BytesRef> ifst)
+ /// Runs case-insensitive matching of the given regexp query on the FST,
emitting the values (dict ids) of each
+ /// matched entry to the given consumer. Values are emitted as the traversal
hits final states, so the caller can
+ /// collect them without this class materializing the full result set in
memory.
+ public static void regexMatch(String regexQuery, FST<BytesRef> ifst,
IntConsumer dictIdConsumer)
throws IOException {
- return new RegexpMatcherCaseInsensitive(regexQuery,
ifst).regexMatchOnFST();
+ new RegexpMatcherCaseInsensitive(regexQuery,
ifst).regexMatchOnFST(dictIdConsumer);
}
// Matches "input" string with _regexQuery Automaton (case-insensitive).
@@ -60,114 +64,92 @@ public class RegexpMatcherCaseInsensitive {
return characterRunAutomaton.run(input.toLowerCase());
}
- /**
- * This function runs case-insensitive matching on automaton built from
regexQuery and the FST.
- * FST stores key (string) to a value (BytesRef). Both are state machines
and state transition is based on
- * a input character.
- *
- * This algorithm starts with Queue containing (Automaton Start Node, FST
Start Node).
- * Each step an entry is popped from the queue:
- * 1) if the automaton state is accept and the FST Node is final (i.e.
end node) then the value stored for that FST
- * is added to the set of result.
- * 2) Else next set of transitions on automaton are gathered and for each
transition target node for that character
- * is figured out in FST Node, resulting pair of (automaton state, fst
node) are added to the queue.
- * 3) This process is bound to complete since we are making progression
on the FST (which is a DAG) towards final
- * nodes.
- * @return List of matched IDs
- * @throws IOException
- */
- public List<Long> regexMatchOnFST()
+ /// This function runs case-insensitive matching on automaton built from
regexQuery and the FST.
+ /// FST stores key (string) to a value (BytesRef). Both are state machines
and state transition is based on an input
+ /// character.
+ ///
+ /// This algorithm starts with a stack containing (Automaton Start Node, FST
Start Node).
+ /// Each step an entry is popped from the stack (DFS order to bound the
frontier size):
+ /// 1) if the automaton state is accept and the FST Node is final (i.e.
end node) then the values stored for that
+ /// FST node are emitted to the consumer.
+ /// 2) Else next set of transitions on automaton are gathered and for each
transition target node for that
+ /// character is figured out in FST Node, resulting pair of (automaton
state, fst node) are added to the stack.
+ /// 3) This process is bound to complete since we are making progression
on the FST (which is a DAG) towards final
+ /// nodes.
+ ///
+ /// A broad regexp on a high-cardinality column can visit a huge number of
paths, so the loop periodically checks
+ /// for query termination (timeout or OOM-protection kill) to remain
interruptible.
+ public void regexMatchOnFST(IntConsumer dictIdConsumer)
throws IOException {
- final List<Path<BytesRef>> queue = new ArrayList<>();
- final List<Path<BytesRef>> endNodes = new ArrayList<>();
if (_automaton.getNumStates() == 0) {
- return Collections.emptyList();
+ return;
}
- // Automaton start state and FST start node is added to the queue.
- queue.add(
- new Path<>(0, _ifst.getFirstArc(new FST.Arc<BytesRef>()),
_ifst.outputs.getNoOutput(), new IntsRefBuilder()));
+ // Automaton start state and FST start node is added to the stack.
+ List<Path<BytesRef>> stack = new ArrayList<>();
+ stack.add(new Path<>(0, _ifst.getFirstArc(new FST.Arc<>()),
_ifst.outputs.getNoOutput()));
- final FST.Arc<BytesRef> scratchArc = new FST.Arc<>();
- final FST.BytesReader fstReader = _ifst.getBytesReader();
+ FST.Arc<BytesRef> scratchArc = new FST.Arc<>();
+ FST.BytesReader fstReader = _ifst.getBytesReader();
Transition t = new Transition();
- while (!queue.isEmpty()) {
- final Path<BytesRef> path = queue.remove(queue.size() - 1);
+ int numPathsProcessed = 0;
+ while (!stack.isEmpty()) {
+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(numPathsProcessed++,
+ "RegexpMatcherCaseInsensitive#regexMatchOnFST");
+ Path<BytesRef> path = stack.remove(stack.size() - 1);
- // If automaton is in accept state and the fstNode is final (i.e. end
node) then add the entry to endNodes which
- // contains the result set.
- if (_automaton.isAccept(path._state)) {
- if (path._fstNode.isFinal()) {
- // For final nodes, we need to combine the accumulated output with
the final output
- BytesRef finalOutput = path._fstNode.nextFinalOutput();
- BytesRef completeOutput;
- if (finalOutput != null && finalOutput.length > 0) {
- // Combine accumulated output with final output
- completeOutput = _ifst.outputs.add(path._output, finalOutput);
- } else {
- // Use the accumulated output if no final output
- completeOutput = path._output;
- }
- // Create a new path with the complete output
- endNodes.add(new Path<>(path._state, path._fstNode, completeOutput,
path._input));
+ // If automaton is in accept state and the fstNode is final (i.e. end
node) then emit the matched values.
+ if (_automaton.isAccept(path._state) && path._fstNode.isFinal()) {
+ // For final nodes, we need to combine the accumulated output with the
final output
+ BytesRef finalOutput = path._fstNode.nextFinalOutput();
+ BytesRef completeOutput;
+ if (finalOutput != null && finalOutput.length > 0) {
+ // Combine accumulated output with final output
+ completeOutput = _ifst.outputs.add(path._output, finalOutput);
+ } else {
+ // Use the accumulated output if no final output
+ completeOutput = path._output;
+ }
+ // Deserialize BytesRef to the list of matched dict ids
+ for (Integer value :
IFSTBuilder.deserializeBytesRefToIntegerList(completeOutput)) {
+ dictIdConsumer.accept(value);
}
}
// Gather next set of transitions on automaton and find target nodes in
FST.
- IntsRefBuilder currentInput = path._input;
int count = _automaton.initTransition(path._state, t);
for (int i = 0; i < count; i++) {
_automaton.getNextTransition(t);
- final int min = t.min;
- final int max = t.max;
+ int min = t.min;
+ int max = t.max;
if (min == max) {
- final FST.Arc<BytesRef> nextArc = _ifst.findTargetArc(t.min,
path._fstNode, scratchArc, fstReader);
+ FST.Arc<BytesRef> nextArc = _ifst.findTargetArc(t.min,
path._fstNode, scratchArc, fstReader);
if (nextArc != null) {
- final IntsRefBuilder newInput = new IntsRefBuilder();
- newInput.copyInts(currentInput.get());
- newInput.append(t.min);
- queue.add(new Path<BytesRef>(t.dest, new
FST.Arc<BytesRef>().copyFrom(nextArc),
- _ifst.outputs.add(path._output, nextArc.output()), newInput));
+ stack.add(new Path<>(t.dest, new
FST.Arc<BytesRef>().copyFrom(nextArc),
+ _ifst.outputs.add(path._output, nextArc.output())));
}
} else {
- FST.Arc<BytesRef> nextArc =
- org.apache.lucene.util.fst.Util.readCeilArc(min, _ifst,
path._fstNode, scratchArc, fstReader);
+ FST.Arc<BytesRef> nextArc = Util.readCeilArc(min, _ifst,
path._fstNode, scratchArc, fstReader);
while (nextArc != null && nextArc.label() <= max) {
- final IntsRefBuilder newInput = new IntsRefBuilder();
- newInput.copyInts(currentInput.get());
- newInput.append(nextArc.label());
- queue.add(new Path<>(t.dest, new
FST.Arc<BytesRef>().copyFrom(nextArc),
- _ifst.outputs.add(path._output, nextArc.output()), newInput));
+ stack.add(new Path<>(t.dest, new
FST.Arc<BytesRef>().copyFrom(nextArc),
+ _ifst.outputs.add(path._output, nextArc.output())));
nextArc = nextArc.isLast() ? null : _ifst.readNextRealArc(nextArc,
fstReader);
}
}
}
}
-
- // From the result set of matched entries gather the values stored and
return.
- ArrayList<Long> matchedIds = new ArrayList<>();
- for (Path<BytesRef> path : endNodes) {
- // Deserialize BytesRef to List<Integer> and convert to List<Long>
- List<Integer> intValues =
IFSTBuilder.deserializeBytesRefToIntegerList(path._output);
- for (Integer value : intValues) {
- matchedIds.add(value.longValue());
- }
- }
- return matchedIds;
}
public static final class Path<T> {
public final int _state;
public final FST.Arc<T> _fstNode;
public final T _output;
- public final IntsRefBuilder _input;
- public Path(int state, FST.Arc<T> fstNode, T output, IntsRefBuilder input)
{
+ public Path(int state, FST.Arc<T> fstNode, T output) {
_state = state;
_fstNode = fstNode;
_output = output;
- _input = input;
}
}
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java
index cedbc0c0de0..e32cbb5e42e 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java
@@ -29,13 +29,17 @@ import org.apache.lucene.util.fst.FST;
import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule;
import
org.apache.pinot.segment.local.segment.index.readers.LuceneFSTIndexReader;
import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
+import org.apache.pinot.spi.exception.EarlyTerminationException;
+import org.apache.pinot.spi.query.QueryThreadContext;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
import static org.testng.Assert.assertTrue;
@@ -88,6 +92,42 @@ public class FSTBuilderTest implements
PinotBuffersAfterMethodCheckRule {
result = reader.getDictIds(".*world");
assertEquals(result.getCardinality(), 1);
assertTrue(result.contains(12));
+
+ result = reader.getDictIds(".*");
+ assertEquals(result.getCardinality(), 3);
+ assertTrue(result.contains(12));
+ assertTrue(result.contains(21));
+ assertTrue(result.contains(123));
+ }
+ }
+
+ @Test
+ public void testRegexMatchHonorsTermination()
+ throws IOException {
+ SortedMap<String, Integer> x = new TreeMap<>();
+ for (int i = 0; i < 1000; i++) {
+ x.put(String.format("key-%06d", i), i);
+ }
+ FST<Long> fst = FSTBuilder.buildFST(x);
+ File outputFile = new File(TEMP_DIR, "test_termination.lucene");
+ try (FileOutputStream outputStream = new FileOutputStream(outputFile);
+ OutputStreamDataOutput dataOutput = new
OutputStreamDataOutput(outputStream)) {
+ fst.save(dataOutput, dataOutput);
+ }
+
+ try (PinotDataBuffer dataBuffer =
PinotDataBuffer.loadBigEndianFile(outputFile);
+ LuceneFSTIndexReader reader = new LuceneFSTIndexReader(dataBuffer);
+ QueryThreadContext ignore = QueryThreadContext.openForSseTest()) {
+ Thread.currentThread().interrupt();
+ assertThrows(EarlyTerminationException.class,
+ () -> RegexpMatcher.regexMatch(".*", fst, new
MutableRoaringBitmap()::add));
+
+ // The reader should propagate the termination exception as-is instead
of wrapping it
+ Thread.currentThread().interrupt();
+ assertThrows(EarlyTerminationException.class, () ->
reader.getDictIds(".*"));
+ } finally {
+ // Clear the interrupt flag in case the matcher did not consume it
+ Thread.interrupted();
}
}
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/IFSTBuilderTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/IFSTBuilderTest.java
index 5e99a978113..9a02a5df456 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/IFSTBuilderTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/IFSTBuilderTest.java
@@ -32,13 +32,17 @@ import org.apache.lucene.util.fst.FST;
import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule;
import
org.apache.pinot.segment.local.segment.index.readers.LuceneIFSTIndexReader;
import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
+import org.apache.pinot.spi.exception.EarlyTerminationException;
+import org.apache.pinot.spi.query.QueryThreadContext;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
import static org.testng.Assert.assertTrue;
@@ -105,6 +109,40 @@ public class IFSTBuilderTest implements
PinotBuffersAfterMethodCheckRule {
assertEquals(result.getCardinality(), 2); // Should still match both
assertTrue(result.contains(12));
assertTrue(result.contains(13));
+
+ // Test full enumeration
+ result = reader.getDictIds(".*");
+ assertEquals(result.getCardinality(), 6); // Should match all entries
+ }
+ }
+
+ @Test
+ public void testRegexMatchHonorsTermination()
+ throws IOException {
+ SortedMap<String, Integer> x = new TreeMap<>();
+ for (int i = 0; i < 1000; i++) {
+ x.put(String.format("key-%06d", i), i);
+ }
+ FST<BytesRef> ifst = IFSTBuilder.buildIFST(x);
+ File outputFile = new File(TEMP_DIR,
"test_termination_case_insensitive.lucene");
+ try (FileOutputStream outputStream = new FileOutputStream(outputFile);
+ OutputStreamDataOutput dataOutput = new
OutputStreamDataOutput(outputStream)) {
+ ifst.save(dataOutput, dataOutput);
+ }
+
+ try (PinotDataBuffer dataBuffer =
PinotDataBuffer.loadBigEndianFile(outputFile);
+ LuceneIFSTIndexReader reader = new LuceneIFSTIndexReader(dataBuffer);
+ QueryThreadContext ignore = QueryThreadContext.openForSseTest()) {
+ Thread.currentThread().interrupt();
+ assertThrows(EarlyTerminationException.class,
+ () -> RegexpMatcherCaseInsensitive.regexMatch(".*", ifst, new
MutableRoaringBitmap()::add));
+
+ // The reader should propagate the termination exception as-is instead
of wrapping it
+ Thread.currentThread().interrupt();
+ assertThrows(EarlyTerminationException.class, () ->
reader.getDictIds(".*"));
+ } finally {
+ // Clear the interrupt flag in case the matcher did not consume it
+ Thread.interrupted();
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
