From c8137d6b8c285ebcc7bd8881a3adbbdfd72f82ec Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 14:36:37 +0800 Subject: [PATCH 1/3] [core] Read primary-key vector results by position --- .../paimon/operation/RawFileSplitRead.java | 97 +++++++-- .../table/source/KeyValueTableRead.java | 2 + .../source/PrimaryKeyVectorDataSplit.java | 175 +++++++++++++++ .../PrimaryKeyVectorPositionReader.java | 106 ++++++++++ .../paimon/table/source/SplitSerializer.java | 8 +- .../PrimaryKeyVectorSplitReadProvider.java | 39 ++++ .../PrimaryKeyVectorPositionReaderTest.java | 199 ++++++++++++++++++ .../PrimaryKeyVectorRawFileSplitReadTest.java | 102 +++++++++ .../table/source/SplitSerializerTest.java | 43 ++++ 9 files changed, 749 insertions(+), 22 deletions(-) create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java index 464d9ea0e5d7..e89ecd05fd33 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java @@ -48,6 +48,8 @@ import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.DeletionFile; import org.apache.paimon.table.source.IncrementalSplit; +import org.apache.paimon.table.source.PrimaryKeyVectorDataSplit; +import org.apache.paimon.table.source.PrimaryKeyVectorPositionReader; import org.apache.paimon.table.source.Split; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.FileStorePathFactory; @@ -150,7 +152,9 @@ public SplitRead withLimit(@Nullable Integer limit) { @Override public RecordReader createReader(Split s) throws IOException { - if (s instanceof DataSplit) { + if (s instanceof PrimaryKeyVectorDataSplit) { + return createReader((PrimaryKeyVectorDataSplit) s); + } else if (s instanceof DataSplit) { DataSplit split = (DataSplit) s; return createReader( split.partition(), @@ -194,22 +198,7 @@ public RecordReader createReader( pathFactory.createDataFilePathFactory(partition, bucket); List> suppliers = new ArrayList<>(); - Builder formatReaderMappingBuilder = - new Builder( - formatDiscover, - readRowType.getFields(), - schema -> { - if (rowTrackingEnabled) { - // maybe file has no row id and sequence number, but in manifest - // entry - return rowTypeWithRowTracking(schema.logicalRowType(), true, true) - .getFields(); - } - return schema.fields(); - }, - filters, - topN, - limit); + Builder formatReaderMappingBuilder = createFormatReaderMappingBuilder(); for (DataFileMeta file : files) { suppliers.add( @@ -218,18 +207,69 @@ public RecordReader createReader( dataFilePathFactory, file, formatReaderMappingBuilder, - dvFactories)); + dvFactories, + null)); } return ConcatRecordReader.create(suppliers); } + private RecordReader createReader(PrimaryKeyVectorDataSplit split) + throws IOException { + DataSplit dataSplit = split.dataSplit(); + DataFileMeta dataFile = dataSplit.dataFiles().get(0); + DeletionVector.Factory dvFactory = + DeletionVector.factory( + fileIO, dataSplit.dataFiles(), dataSplit.deletionFiles().orElse(null)); + Map> dvFactories = new HashMap<>(); + dvFactories.put( + dataFile.fileName(), () -> dvFactory.create(dataFile.fileName()).orElse(null)); + DataFilePathFactory dataFilePathFactory = + pathFactory.createDataFilePathFactory(dataSplit.partition(), dataSplit.bucket()); + FileRecordReader reader = + (FileRecordReader) + createFileReader( + dataSplit.partition(), + dataFilePathFactory, + dataFile, + // ANN has already selected the rows. Applying a regular + // TopN or limit before position filtering can drop hits. + createFormatReaderMappingBuilder(null, null), + dvFactories, + split.rowPositions()) + .get(); + return new PrimaryKeyVectorPositionReader(reader, split.rowPositions(), split::score); + } + + private Builder createFormatReaderMappingBuilder() { + return createFormatReaderMappingBuilder(topN, limit); + } + + private Builder createFormatReaderMappingBuilder( + @Nullable TopN pushDownTopN, @Nullable Integer pushDownLimit) { + return new Builder( + formatDiscover, + readRowType.getFields(), + schema -> { + if (rowTrackingEnabled) { + // maybe file has no row id and sequence number, but in manifest entry + return rowTypeWithRowTracking(schema.logicalRowType(), true, true) + .getFields(); + } + return schema.fields(); + }, + filters, + pushDownTopN, + pushDownLimit); + } + private ReaderSupplier createFileReader( BinaryRow partition, DataFilePathFactory dataFilePathFactory, DataFileMeta file, Builder formatBuilder, - @Nullable Map> dvFactories) { + @Nullable Map> dvFactories, + @Nullable RoaringBitmap32 selectedPositions) { String formatIdentifier = DataFilePathFactory.formatIdentifier(file.fileName()); long schemaId = file.schemaId(); @@ -248,7 +288,12 @@ private ReaderSupplier createFileReader( dvFactories == null ? null : dvFactories.get(file.fileName()); return () -> createFileReader( - partition, file, dataFilePathFactory, formatReaderMapping, dvFactory); + partition, + file, + dataFilePathFactory, + formatReaderMapping, + dvFactory, + selectedPositions); } private FileRecordReader createFileReader( @@ -256,7 +301,8 @@ private FileRecordReader createFileReader( DataFileMeta file, DataFilePathFactory dataFilePathFactory, FormatReaderMapping formatReaderMapping, - IOExceptionSupplier dvFactory) + IOExceptionSupplier dvFactory, + @Nullable RoaringBitmap32 selectedPositions) throws IOException { FileIndexResult fileIndexResult = null; DeletionVector deletionVector = dvFactory == null ? null : dvFactory.get(); @@ -280,6 +326,15 @@ private FileRecordReader createFileReader( if (fileIndexResult instanceof BitmapIndexResult) { selection = ((BitmapIndexResult) fileIndexResult).get(); } + if (selectedPositions != null) { + selection = + selection == null + ? selectedPositions.clone() + : RoaringBitmap32.and(selection, selectedPositions); + if (selection.isEmpty()) { + return new EmptyFileRecordReader<>(); + } + } FormatReaderContext formatReaderContext = new FormatReaderContext( diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java index fda7d70ffdf6..072ed8c2276c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java @@ -34,6 +34,7 @@ import org.apache.paimon.table.source.splitread.IncrementalDiffReadProvider; import org.apache.paimon.table.source.splitread.MergeFileSplitReadProvider; import org.apache.paimon.table.source.splitread.PrimaryKeyTableRawFileSplitReadProvider; +import org.apache.paimon.table.source.splitread.PrimaryKeyVectorSplitReadProvider; import org.apache.paimon.table.source.splitread.SplitReadProvider; import org.apache.paimon.types.RowType; @@ -67,6 +68,7 @@ public KeyValueTableRead( super(schema); this.readProviders = Arrays.asList( + new PrimaryKeyVectorSplitReadProvider(batchRawReadSupplier, this::config), new PrimaryKeyTableRawFileSplitReadProvider( batchRawReadSupplier, this::config), new MergeFileSplitReadProvider(mergeReadSupplier, this::config), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java new file mode 100644 index 000000000000..bdc12a9a9dfc --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java @@ -0,0 +1,175 @@ +/* + * 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.paimon.table.source; + +import org.apache.paimon.io.DataInputView; +import org.apache.paimon.io.DataInputViewStreamWrapper; +import org.apache.paimon.io.DataOutputView; +import org.apache.paimon.io.DataOutputViewStreamWrapper; +import org.apache.paimon.utils.RoaringBitmap32; + +import java.io.IOException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.OptionalLong; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Selected physical row positions and vector scores in one primary-key table data file. */ +public class PrimaryKeyVectorDataSplit implements Split { + + private static final long serialVersionUID = 1L; + private static final long MAGIC = 0x504B5653504C4954L; // "PKVSPLIT" + private static final int VERSION = 1; + + private DataSplit dataSplit; + private RoaringBitmap32 rowPositions; + private Map scores; + + PrimaryKeyVectorDataSplit( + DataSplit dataSplit, RoaringBitmap32 rowPositions, Map scores) { + checkArgument( + dataSplit.dataFiles().size() == 1, + "Primary-key vector data split must contain exactly one data file."); + checkArgument( + !rowPositions.isEmpty(), + "Primary-key vector data split must select at least one row position."); + checkArgument( + rowPositions.getCardinality() == scores.size(), + "Primary-key vector positions and scores must have the same size."); + Iterator positionIterator = rowPositions.iterator(); + while (positionIterator.hasNext()) { + int rowPosition = positionIterator.next(); + checkArgument( + rowPosition >= 0 && rowPosition < dataSplit.dataFiles().get(0).rowCount(), + "Selected row position %s is outside data file %s row count %s.", + rowPosition, + dataSplit.dataFiles().get(0).fileName(), + dataSplit.dataFiles().get(0).rowCount()); + checkArgument( + scores.containsKey(rowPosition), + "Selected row position %s has no vector score.", + rowPosition); + } + this.dataSplit = dataSplit; + this.rowPositions = rowPositions.clone(); + this.scores = Collections.unmodifiableMap(new HashMap<>(scores)); + } + + public DataSplit dataSplit() { + return dataSplit; + } + + public RoaringBitmap32 rowPositions() { + return rowPositions.clone(); + } + + public float score(int rowPosition) { + Float score = scores.get(rowPosition); + checkArgument(score != null, "Row position %s is not selected.", rowPosition); + return score; + } + + public void serialize(DataOutputView out) throws IOException { + out.writeLong(MAGIC); + out.writeInt(VERSION); + dataSplit.serialize(out); + rowPositions.serialize(out); + out.writeInt(scores.size()); + Iterator positions = rowPositions.iterator(); + while (positions.hasNext()) { + int position = positions.next(); + out.writeInt(position); + out.writeFloat(scores.get(position)); + } + } + + public static PrimaryKeyVectorDataSplit deserialize(DataInputView in) throws IOException { + long magic = in.readLong(); + if (magic != MAGIC) { + throw new IOException( + "Corrupted PrimaryKeyVectorDataSplit: wrong magic number " + magic); + } + int version = in.readInt(); + if (version != VERSION) { + throw new IOException("Unsupported PrimaryKeyVectorDataSplit version: " + version); + } + DataSplit dataSplit = DataSplit.deserialize(in); + RoaringBitmap32 positions = new RoaringBitmap32(); + positions.deserialize(in); + int scoreCount = in.readInt(); + Map scores = new HashMap<>(scoreCount); + for (int i = 0; i < scoreCount; i++) { + scores.put(in.readInt(), in.readFloat()); + } + return new PrimaryKeyVectorDataSplit(dataSplit, positions, scores); + } + + private void writeObject(ObjectOutputStream out) throws IOException { + serialize(new DataOutputViewStreamWrapper(out)); + } + + private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException { + PrimaryKeyVectorDataSplit restored = deserialize(new DataInputViewStreamWrapper(in)); + this.dataSplit = restored.dataSplit; + this.rowPositions = restored.rowPositions; + this.scores = restored.scores; + } + + @Override + public long rowCount() { + return rowPositions.getCardinality(); + } + + @Override + public OptionalLong mergedRowCount() { + return OptionalLong.of(rowCount()); + } + + @Override + public Optional> deletionFiles() { + return dataSplit.deletionFiles(); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + PrimaryKeyVectorDataSplit that = (PrimaryKeyVectorDataSplit) o; + return dataSplit.equals(that.dataSplit) + && rowPositions.equals(that.rowPositions) + && scores.equals(that.scores); + } + + @Override + public int hashCode() { + return Objects.hash(dataSplit, rowPositions, scores); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java new file mode 100644 index 000000000000..6470fc26ebd7 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java @@ -0,0 +1,106 @@ +/* + * 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.paimon.table.source; + +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.reader.ScoreRecordIterator; +import org.apache.paimon.utils.RoaringBitmap32; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.function.IntFunction; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Reads selected physical file positions and exposes their vector-search scores. */ +public class PrimaryKeyVectorPositionReader implements RecordReader { + + private final FileRecordReader reader; + private final RoaringBitmap32 rowPositions; + private final IntFunction scoreGetter; + private final int lastPosition; + private boolean exhausted; + + public PrimaryKeyVectorPositionReader( + FileRecordReader reader, + RoaringBitmap32 rowPositions, + IntFunction scoreGetter) { + checkArgument(!rowPositions.isEmpty(), "Selected row positions must not be empty."); + this.reader = reader; + this.rowPositions = rowPositions.clone(); + this.scoreGetter = scoreGetter; + this.lastPosition = rowPositions.last(); + } + + @Nullable + @Override + public ScoreRecordIterator readBatch() throws IOException { + if (exhausted) { + return null; + } + FileRecordIterator batch = reader.readBatch(); + if (batch == null) { + return null; + } + FileRecordIterator selected = batch.selection(rowPositions); + return new ScoreRecordIterator() { + + private long returnedPosition = -1; + private float returnedScore = Float.NaN; + + @Override + public float returnedScore() { + return returnedScore; + } + + @Override + public long returnedRowId() { + return returnedPosition; + } + + @Nullable + @Override + public InternalRow next() throws IOException { + InternalRow row = selected.next(); + if (row != null) { + returnedPosition = selected.returnedPosition(); + returnedScore = scoreGetter.apply((int) returnedPosition); + if (returnedPosition >= lastPosition) { + exhausted = true; + } + } + return row; + } + + @Override + public void releaseBatch() { + selected.releaseBatch(); + } + }; + } + + @Override + public void close() throws IOException { + reader.close(); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java b/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java index d1f8ad625362..be2418081002 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java @@ -63,6 +63,7 @@ public class SplitSerializer { private static final int QUERY_AUTH_SPLIT = 5; private static final int FALLBACK_DATA_SPLIT = 6; private static final int FALLBACK_SPLIT = 7; + private static final int PRIMARY_KEY_VECTOR_DATA_SPLIT = 8; private SplitSerializer() {} @@ -76,7 +77,10 @@ public static void serialize(Split split, DataOutputView out) throws IOException out.writeLong(MAGIC); out.writeInt(VERSION); - if (split instanceof QueryAuthSplit) { + if (split instanceof PrimaryKeyVectorDataSplit) { + out.writeInt(PRIMARY_KEY_VECTOR_DATA_SPLIT); + ((PrimaryKeyVectorDataSplit) split).serialize(out); + } else if (split instanceof QueryAuthSplit) { out.writeInt(QUERY_AUTH_SPLIT); writeQueryAuthSplit((QueryAuthSplit) split, out); } else if (split instanceof FallbackReadFileStoreTable.FallbackDataSplit) { @@ -133,6 +137,8 @@ public static Split deserialize(DataInputView in) throws IOException { return FallbackReadFileStoreTable.FallbackDataSplit.deserialize(in); case FALLBACK_SPLIT: return readFallbackSplit(in); + case PRIMARY_KEY_VECTOR_DATA_SPLIT: + return PrimaryKeyVectorDataSplit.deserialize(in); default: throw new IOException("Unsupported split type: " + type); } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java new file mode 100644 index 000000000000..1aea9641f6b7 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java @@ -0,0 +1,39 @@ +/* + * 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.paimon.table.source.splitread; + +import org.apache.paimon.operation.RawFileSplitRead; +import org.apache.paimon.table.source.PrimaryKeyVectorDataSplit; +import org.apache.paimon.table.source.Split; + +import java.util.function.Supplier; + +/** Reads selected physical positions from one primary-key table data file. */ +public class PrimaryKeyVectorSplitReadProvider extends RawFileSplitReadProvider { + + public PrimaryKeyVectorSplitReadProvider( + Supplier supplier, SplitReadConfig splitReadConfig) { + super(supplier, splitReadConfig); + } + + @Override + public boolean match(Split split, Context context) { + return !context.forceKeepDelete() && split instanceof PrimaryKeyVectorDataSplit; + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java new file mode 100644 index 000000000000..dd10d797d48f --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java @@ -0,0 +1,199 @@ +/* + * 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.paimon.table.source; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.fs.Path; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.operation.MergeFileSplitRead; +import org.apache.paimon.operation.RawFileSplitRead; +import org.apache.paimon.reader.FileRecordIterator; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.reader.ScoreRecordIterator; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.utils.RoaringBitmap32; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Answers.RETURNS_SELF; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests physical row-position filtering with vector scores. */ +class PrimaryKeyVectorPositionReaderTest { + + @Test + void testSelectsPhysicalPositionsAndReturnsScores() throws Exception { + PrimaryKeyVectorPositionReader reader = + new PrimaryKeyVectorPositionReader( + new TestingFileReader(5), + RoaringBitmap32.bitmapOf(1, 3), + position -> position / 10F); + + ScoreRecordIterator batch = reader.readBatch(); + assertThat(batch.next().getInt(0)).isEqualTo(1); + assertThat(batch.returnedRowId()).isEqualTo(1); + assertThat(batch.returnedScore()).isEqualTo(0.1F); + assertThat(batch.next().getInt(0)).isEqualTo(3); + assertThat(batch.returnedRowId()).isEqualTo(3); + assertThat(batch.returnedScore()).isEqualTo(0.3F); + assertThat(batch.next()).isNull(); + assertThat(reader.readBatch()).isNull(); + } + + @Test + void testSelectedPositionsAcrossBatches() throws Exception { + PrimaryKeyVectorPositionReader reader = + new PrimaryKeyVectorPositionReader( + new TestingFileReader(5, 2), + RoaringBitmap32.bitmapOf(1, 3), + position -> position / 10F); + + ScoreRecordIterator firstBatch = reader.readBatch(); + assertThat(firstBatch.next().getInt(0)).isEqualTo(1); + assertThat(firstBatch.returnedRowId()).isEqualTo(1); + assertThat(firstBatch.next()).isNull(); + + ScoreRecordIterator secondBatch = reader.readBatch(); + assertThat(secondBatch.next().getInt(0)).isEqualTo(3); + assertThat(secondBatch.returnedScore()).isEqualTo(0.3F); + assertThat(secondBatch.next()).isNull(); + assertThat(reader.readBatch()).isNull(); + } + + @Test + void testKeyValueTableReadRoutesPhysicalSplitToRawReader() throws Exception { + RawFileSplitRead rawRead = mock(RawFileSplitRead.class, RETURNS_SELF); + RecordReader expectedReader = mock(RecordReader.class); + PrimaryKeyVectorDataSplit split = selectedSplit(); + when(rawRead.createReader(split)).thenReturn(expectedReader); + KeyValueTableRead tableRead = + new KeyValueTableRead( + () -> mock(MergeFileSplitRead.class), + () -> rawRead, + mock(TableSchema.class)); + + assertThat(tableRead.createReader(split)).isSameAs(expectedReader); + verify(rawRead).createReader(split); + } + + private static PrimaryKeyVectorDataSplit selectedSplit() { + DataFileMeta dataFile = + DataFileMeta.forAppend( + "data-file", + 100, + 5, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.APPEND, + null, + null, + null, + null); + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(3) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(1) + .withDataFiles(Collections.singletonList(dataFile)) + .build(); + HashMap scores = new HashMap<>(); + scores.put(1, 0.5F); + return new PrimaryKeyVectorDataSplit(dataSplit, RoaringBitmap32.bitmapOf(1), scores); + } + + private static class TestingFileReader implements FileRecordReader { + + private final int rowCount; + private final int batchSize; + private int nextPosition; + + private TestingFileReader(int rowCount) { + this(rowCount, rowCount); + } + + private TestingFileReader(int rowCount, int batchSize) { + this.rowCount = rowCount; + this.batchSize = batchSize; + } + + @Override + public FileRecordIterator readBatch() { + if (nextPosition == rowCount) { + return null; + } + int startPosition = nextPosition; + nextPosition = Math.min(rowCount, nextPosition + batchSize); + return new TestingFileIterator(startPosition, nextPosition); + } + + @Override + public void close() {} + } + + private static class TestingFileIterator implements FileRecordIterator { + + private final int endPosition; + private int nextPosition; + private int returnedPosition = -1; + + private TestingFileIterator(int startPosition, int endPosition) { + this.nextPosition = startPosition; + this.endPosition = endPosition; + } + + @Override + public long returnedPosition() { + return returnedPosition; + } + + @Override + public Path filePath() { + return new Path("/test"); + } + + @Override + public InternalRow next() { + if (nextPosition == endPosition) { + return null; + } + returnedPosition = nextPosition; + return GenericRow.of(nextPosition++); + } + + @Override + public void releaseBatch() {} + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java new file mode 100644 index 000000000000..75eda2f88f3b --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java @@ -0,0 +1,102 @@ +/* + * 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.paimon.table.source; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.options.Options; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.reader.ScoreRecordIterator; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.SchemaUtils; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.FileStoreTableFactory; +import org.apache.paimon.table.sink.BatchTableCommit; +import org.apache.paimon.table.sink.BatchTableWrite; +import org.apache.paimon.table.sink.BatchWriteBuilder; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.utils.RoaringBitmap32; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.util.HashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests raw-file materialization of primary-key vector results. */ +class PrimaryKeyVectorRawFileSplitReadTest { + + @TempDir java.nio.file.Path tempDir; + + @Test + void testRegularLimitDoesNotDropSelectedPhysicalPosition() throws Exception { + Path tablePath = new Path(tempDir.toUri()); + Options options = new Options(); + options.set(CoreOptions.PATH, tablePath.toString()); + options.set(CoreOptions.BUCKET, 1); + Schema schema = + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("value", DataTypes.INT()) + .primaryKey("id") + .options(options.toMap()) + .build(); + TableSchema tableSchema = + SchemaUtils.forceCommit(new SchemaManager(LocalFileIO.create(), tablePath), schema); + FileStoreTable table = + FileStoreTableFactory.create(LocalFileIO.create(), tablePath, tableSchema); + + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + for (int i = 0; i < 5; i++) { + write.write(GenericRow.of(i, i * 10)); + } + commit.commit(write.prepareCommit()); + } + + DataSplit dataSplit = table.newSnapshotReader().read().dataSplits().get(0); + assertThat(dataSplit.dataFiles()).hasSize(1); + Map scores = new HashMap<>(); + scores.put(3, 0.25F); + PrimaryKeyVectorDataSplit vectorSplit = + new PrimaryKeyVectorDataSplit(dataSplit, RoaringBitmap32.bitmapOf(3), scores); + + try (RecordReader reader = + table.newRead().withLimit(1).createReader(vectorSplit)) { + ScoreRecordIterator batch = + (ScoreRecordIterator) reader.readBatch(); + InternalRow row = batch.next(); + assertThat(row.getInt(0)).isEqualTo(3); + assertThat(row.getInt(1)).isEqualTo(30); + assertThat(batch.returnedRowId()).isEqualTo(3); + assertThat(batch.returnedScore()).isEqualTo(0.25F); + assertThat(batch.next()).isNull(); + batch.releaseBatch(); + assertThat(reader.readBatch()).isNull(); + } + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java index 7cc16771f97f..b58312e7de58 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java @@ -32,6 +32,7 @@ import org.apache.paimon.utils.IOUtils; import org.apache.paimon.utils.InstantiationUtil; import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringBitmap32; import org.junit.jupiter.api.Test; @@ -61,6 +62,48 @@ public void testRoundTrip() throws IOException { } } + @Test + public void testPrimaryKeyVectorDataSplitRoundTrip() throws IOException { + PrimaryKeyVectorDataSplit split = primaryKeyVectorDataSplit(); + + Split actual = SplitSerializer.deserialize(SplitSerializer.serialize(split)); + + assertThat(actual).isEqualTo(split); + PrimaryKeyVectorDataSplit vectorSplit = (PrimaryKeyVectorDataSplit) actual; + assertThat(vectorSplit.rowPositions()).isEqualTo(RoaringBitmap32.bitmapOf(1, 3)); + assertThat(vectorSplit.score(1)).isEqualTo(0.9F); + assertThat(vectorSplit.score(3)).isEqualTo(0.7F); + } + + @Test + public void testPrimaryKeyVectorDataSplitJavaSerialization() throws Exception { + PrimaryKeyVectorDataSplit split = primaryKeyVectorDataSplit(); + + PrimaryKeyVectorDataSplit actual = + InstantiationUtil.deserializeObject( + InstantiationUtil.serializeObject(split), getClass().getClassLoader()); + + assertThat(actual).isEqualTo(split); + } + + private static PrimaryKeyVectorDataSplit primaryKeyVectorDataSplit() { + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(45L) + .withPartition(DataFileTestUtils.row(2026, 10)) + .withBucket(6) + .withTotalBuckets(8) + .withBucketPath("dt=20260708/bucket-6") + .withDataFiles( + Collections.singletonList(dataFile("vector-file", 1, 0, 4, 400L))) + .rawConvertible(true) + .build(); + Map scores = new LinkedHashMap<>(); + scores.put(1, 0.9F); + scores.put(3, 0.7F); + return new PrimaryKeyVectorDataSplit(dataSplit, RoaringBitmap32.bitmapOf(1, 3), scores); + } + @Test public void testGoldenFiles() throws IOException { boolean generateGoldenFiles = From 2686538f48357d10b672dfb9f40a498bbd1f3a17 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 20:06:31 +0800 Subject: [PATCH 2/3] [core] Reuse IndexedSplit for vector positions --- .../paimon/globalindex/IndexedSplit.java | 8 +- .../paimon/operation/RawFileSplitRead.java | 52 +++++- .../table/source/KeyValueTableRead.java | 2 - .../source/PrimaryKeyVectorDataSplit.java | 175 ------------------ .../paimon/table/source/SplitSerializer.java | 8 +- ...imaryKeyTableRawFileSplitReadProvider.java | 4 + .../PrimaryKeyVectorSplitReadProvider.java | 39 ---- .../PrimaryKeyVectorPositionReaderTest.java | 12 +- .../PrimaryKeyVectorRawFileSplitReadTest.java | 24 ++- .../table/source/SplitSerializerTest.java | 43 ----- 10 files changed, 78 insertions(+), 289 deletions(-) delete mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java delete mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java diff --git a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java index d9b6fef28c46..49802d5d4c0c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java +++ b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplit.java @@ -37,7 +37,13 @@ import java.util.Objects; import java.util.OptionalLong; -/** Indexed split for global index. */ +/** + * Indexed split for global index. + * + *

Row ranges use the coordinate system of the table read path. Data Evolution reads them as + * stable row IDs, while primary-key raw reads use them as physical positions in the split's single + * data file. Scores, when present, follow the expanded row-range order. + */ public class IndexedSplit implements Split { private static final long serialVersionUID = 1L; diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java index e89ecd05fd33..f3b32054a057 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java @@ -31,6 +31,7 @@ import org.apache.paimon.format.FormatKey; import org.apache.paimon.format.FormatReaderContext; import org.apache.paimon.fs.FileIO; +import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.io.DataFilePathFactory; import org.apache.paimon.io.DataFileRecordReader; @@ -48,7 +49,6 @@ import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.DeletionFile; import org.apache.paimon.table.source.IncrementalSplit; -import org.apache.paimon.table.source.PrimaryKeyVectorDataSplit; import org.apache.paimon.table.source.PrimaryKeyVectorPositionReader; import org.apache.paimon.table.source.Split; import org.apache.paimon.types.RowType; @@ -56,6 +56,7 @@ import org.apache.paimon.utils.FormatReaderMapping; import org.apache.paimon.utils.FormatReaderMapping.Builder; import org.apache.paimon.utils.IOExceptionSupplier; +import org.apache.paimon.utils.Range; import org.apache.paimon.utils.RoaringBitmap32; import org.slf4j.Logger; @@ -68,9 +69,11 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.IntFunction; import static org.apache.paimon.predicate.PredicateBuilder.splitAnd; import static org.apache.paimon.table.SpecialFields.rowTypeWithRowTracking; +import static org.apache.paimon.utils.Preconditions.checkArgument; /** A {@link SplitRead} to read raw file directly from {@link DataSplit}. */ public class RawFileSplitRead implements SplitRead { @@ -152,8 +155,8 @@ public SplitRead withLimit(@Nullable Integer limit) { @Override public RecordReader createReader(Split s) throws IOException { - if (s instanceof PrimaryKeyVectorDataSplit) { - return createReader((PrimaryKeyVectorDataSplit) s); + if (s instanceof IndexedSplit) { + return createReader((IndexedSplit) s); } else if (s instanceof DataSplit) { DataSplit split = (DataSplit) s; return createReader( @@ -214,10 +217,45 @@ public RecordReader createReader( return ConcatRecordReader.create(suppliers); } - private RecordReader createReader(PrimaryKeyVectorDataSplit split) - throws IOException { + private RecordReader createReader(IndexedSplit split) throws IOException { DataSplit dataSplit = split.dataSplit(); + checkArgument( + dataSplit.dataFiles().size() == 1, + "Indexed split for a primary-key table must contain exactly one data file."); DataFileMeta dataFile = dataSplit.dataFiles().get(0); + RoaringBitmap32 rowPositions = new RoaringBitmap32(); + float[] scores = split.scores(); + Map scoreByPosition = scores == null ? null : new HashMap<>(); + int scoreIndex = 0; + for (Range range : split.rowRanges()) { + checkArgument( + range.from >= 0 && range.to < dataFile.rowCount(), + "Indexed row range %s is outside data file %s row count %s.", + range, + dataFile.fileName(), + dataFile.rowCount()); + checkArgument( + range.to <= Integer.MAX_VALUE, + "Indexed row range %s exceeds supported physical positions.", + range); + for (long position = range.from; position <= range.to; position++) { + rowPositions.add((int) position); + if (scoreByPosition != null) { + checkArgument( + scoreIndex < scores.length, + "Scores length does not match row ranges in indexed split."); + scoreByPosition.put((int) position, scores[scoreIndex++]); + } + } + } + checkArgument(!rowPositions.isEmpty(), "Indexed split must select at least one row."); + if (scores != null) { + checkArgument( + scoreIndex == scores.length, + "Scores length does not match row ranges in indexed split."); + } + IntFunction scoreGetter = + scoreByPosition == null ? position -> Float.NaN : scoreByPosition::get; DeletionVector.Factory dvFactory = DeletionVector.factory( fileIO, dataSplit.dataFiles(), dataSplit.deletionFiles().orElse(null)); @@ -236,9 +274,9 @@ private RecordReader createReader(PrimaryKeyVectorDataSplit split) // TopN or limit before position filtering can drop hits. createFormatReaderMappingBuilder(null, null), dvFactories, - split.rowPositions()) + rowPositions) .get(); - return new PrimaryKeyVectorPositionReader(reader, split.rowPositions(), split::score); + return new PrimaryKeyVectorPositionReader(reader, rowPositions, scoreGetter); } private Builder createFormatReaderMappingBuilder() { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java index 072ed8c2276c..fda7d70ffdf6 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java @@ -34,7 +34,6 @@ import org.apache.paimon.table.source.splitread.IncrementalDiffReadProvider; import org.apache.paimon.table.source.splitread.MergeFileSplitReadProvider; import org.apache.paimon.table.source.splitread.PrimaryKeyTableRawFileSplitReadProvider; -import org.apache.paimon.table.source.splitread.PrimaryKeyVectorSplitReadProvider; import org.apache.paimon.table.source.splitread.SplitReadProvider; import org.apache.paimon.types.RowType; @@ -68,7 +67,6 @@ public KeyValueTableRead( super(schema); this.readProviders = Arrays.asList( - new PrimaryKeyVectorSplitReadProvider(batchRawReadSupplier, this::config), new PrimaryKeyTableRawFileSplitReadProvider( batchRawReadSupplier, this::config), new MergeFileSplitReadProvider(mergeReadSupplier, this::config), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java deleted file mode 100644 index bdc12a9a9dfc..000000000000 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorDataSplit.java +++ /dev/null @@ -1,175 +0,0 @@ -/* - * 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.paimon.table.source; - -import org.apache.paimon.io.DataInputView; -import org.apache.paimon.io.DataInputViewStreamWrapper; -import org.apache.paimon.io.DataOutputView; -import org.apache.paimon.io.DataOutputViewStreamWrapper; -import org.apache.paimon.utils.RoaringBitmap32; - -import java.io.IOException; -import java.io.ObjectInputStream; -import java.io.ObjectOutputStream; -import java.util.Collections; -import java.util.HashMap; -import java.util.Iterator; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; -import java.util.OptionalLong; - -import static org.apache.paimon.utils.Preconditions.checkArgument; - -/** Selected physical row positions and vector scores in one primary-key table data file. */ -public class PrimaryKeyVectorDataSplit implements Split { - - private static final long serialVersionUID = 1L; - private static final long MAGIC = 0x504B5653504C4954L; // "PKVSPLIT" - private static final int VERSION = 1; - - private DataSplit dataSplit; - private RoaringBitmap32 rowPositions; - private Map scores; - - PrimaryKeyVectorDataSplit( - DataSplit dataSplit, RoaringBitmap32 rowPositions, Map scores) { - checkArgument( - dataSplit.dataFiles().size() == 1, - "Primary-key vector data split must contain exactly one data file."); - checkArgument( - !rowPositions.isEmpty(), - "Primary-key vector data split must select at least one row position."); - checkArgument( - rowPositions.getCardinality() == scores.size(), - "Primary-key vector positions and scores must have the same size."); - Iterator positionIterator = rowPositions.iterator(); - while (positionIterator.hasNext()) { - int rowPosition = positionIterator.next(); - checkArgument( - rowPosition >= 0 && rowPosition < dataSplit.dataFiles().get(0).rowCount(), - "Selected row position %s is outside data file %s row count %s.", - rowPosition, - dataSplit.dataFiles().get(0).fileName(), - dataSplit.dataFiles().get(0).rowCount()); - checkArgument( - scores.containsKey(rowPosition), - "Selected row position %s has no vector score.", - rowPosition); - } - this.dataSplit = dataSplit; - this.rowPositions = rowPositions.clone(); - this.scores = Collections.unmodifiableMap(new HashMap<>(scores)); - } - - public DataSplit dataSplit() { - return dataSplit; - } - - public RoaringBitmap32 rowPositions() { - return rowPositions.clone(); - } - - public float score(int rowPosition) { - Float score = scores.get(rowPosition); - checkArgument(score != null, "Row position %s is not selected.", rowPosition); - return score; - } - - public void serialize(DataOutputView out) throws IOException { - out.writeLong(MAGIC); - out.writeInt(VERSION); - dataSplit.serialize(out); - rowPositions.serialize(out); - out.writeInt(scores.size()); - Iterator positions = rowPositions.iterator(); - while (positions.hasNext()) { - int position = positions.next(); - out.writeInt(position); - out.writeFloat(scores.get(position)); - } - } - - public static PrimaryKeyVectorDataSplit deserialize(DataInputView in) throws IOException { - long magic = in.readLong(); - if (magic != MAGIC) { - throw new IOException( - "Corrupted PrimaryKeyVectorDataSplit: wrong magic number " + magic); - } - int version = in.readInt(); - if (version != VERSION) { - throw new IOException("Unsupported PrimaryKeyVectorDataSplit version: " + version); - } - DataSplit dataSplit = DataSplit.deserialize(in); - RoaringBitmap32 positions = new RoaringBitmap32(); - positions.deserialize(in); - int scoreCount = in.readInt(); - Map scores = new HashMap<>(scoreCount); - for (int i = 0; i < scoreCount; i++) { - scores.put(in.readInt(), in.readFloat()); - } - return new PrimaryKeyVectorDataSplit(dataSplit, positions, scores); - } - - private void writeObject(ObjectOutputStream out) throws IOException { - serialize(new DataOutputViewStreamWrapper(out)); - } - - private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException { - PrimaryKeyVectorDataSplit restored = deserialize(new DataInputViewStreamWrapper(in)); - this.dataSplit = restored.dataSplit; - this.rowPositions = restored.rowPositions; - this.scores = restored.scores; - } - - @Override - public long rowCount() { - return rowPositions.getCardinality(); - } - - @Override - public OptionalLong mergedRowCount() { - return OptionalLong.of(rowCount()); - } - - @Override - public Optional> deletionFiles() { - return dataSplit.deletionFiles(); - } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - PrimaryKeyVectorDataSplit that = (PrimaryKeyVectorDataSplit) o; - return dataSplit.equals(that.dataSplit) - && rowPositions.equals(that.rowPositions) - && scores.equals(that.scores); - } - - @Override - public int hashCode() { - return Objects.hash(dataSplit, rowPositions, scores); - } -} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java b/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java index be2418081002..d1f8ad625362 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/SplitSerializer.java @@ -63,7 +63,6 @@ public class SplitSerializer { private static final int QUERY_AUTH_SPLIT = 5; private static final int FALLBACK_DATA_SPLIT = 6; private static final int FALLBACK_SPLIT = 7; - private static final int PRIMARY_KEY_VECTOR_DATA_SPLIT = 8; private SplitSerializer() {} @@ -77,10 +76,7 @@ public static void serialize(Split split, DataOutputView out) throws IOException out.writeLong(MAGIC); out.writeInt(VERSION); - if (split instanceof PrimaryKeyVectorDataSplit) { - out.writeInt(PRIMARY_KEY_VECTOR_DATA_SPLIT); - ((PrimaryKeyVectorDataSplit) split).serialize(out); - } else if (split instanceof QueryAuthSplit) { + if (split instanceof QueryAuthSplit) { out.writeInt(QUERY_AUTH_SPLIT); writeQueryAuthSplit((QueryAuthSplit) split, out); } else if (split instanceof FallbackReadFileStoreTable.FallbackDataSplit) { @@ -137,8 +133,6 @@ public static Split deserialize(DataInputView in) throws IOException { return FallbackReadFileStoreTable.FallbackDataSplit.deserialize(in); case FALLBACK_SPLIT: return readFallbackSplit(in); - case PRIMARY_KEY_VECTOR_DATA_SPLIT: - return PrimaryKeyVectorDataSplit.deserialize(in); default: throw new IOException("Unsupported split type: " + type); } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java index 713977dc590d..3ae77dd74776 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java @@ -18,6 +18,7 @@ package org.apache.paimon.table.source.splitread; +import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.operation.RawFileSplitRead; import org.apache.paimon.table.source.DataSplit; @@ -35,6 +36,9 @@ public PrimaryKeyTableRawFileSplitReadProvider( @Override public boolean match(Split split, Context context) { + if (split instanceof IndexedSplit) { + return !context.forceKeepDelete(); + } if (!(split instanceof DataSplit)) { return false; } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java deleted file mode 100644 index 1aea9641f6b7..000000000000 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyVectorSplitReadProvider.java +++ /dev/null @@ -1,39 +0,0 @@ -/* - * 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.paimon.table.source.splitread; - -import org.apache.paimon.operation.RawFileSplitRead; -import org.apache.paimon.table.source.PrimaryKeyVectorDataSplit; -import org.apache.paimon.table.source.Split; - -import java.util.function.Supplier; - -/** Reads selected physical positions from one primary-key table data file. */ -public class PrimaryKeyVectorSplitReadProvider extends RawFileSplitReadProvider { - - public PrimaryKeyVectorSplitReadProvider( - Supplier supplier, SplitReadConfig splitReadConfig) { - super(supplier, splitReadConfig); - } - - @Override - public boolean match(Split split, Context context) { - return !context.forceKeepDelete() && split instanceof PrimaryKeyVectorDataSplit; - } -} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java index dd10d797d48f..b600f94abd0a 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java @@ -22,6 +22,7 @@ import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; import org.apache.paimon.fs.Path; +import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.manifest.FileSource; import org.apache.paimon.operation.MergeFileSplitRead; @@ -32,12 +33,12 @@ import org.apache.paimon.reader.ScoreRecordIterator; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.utils.Range; import org.apache.paimon.utils.RoaringBitmap32; import org.junit.jupiter.api.Test; import java.util.Collections; -import java.util.HashMap; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Answers.RETURNS_SELF; @@ -91,7 +92,7 @@ void testSelectedPositionsAcrossBatches() throws Exception { void testKeyValueTableReadRoutesPhysicalSplitToRawReader() throws Exception { RawFileSplitRead rawRead = mock(RawFileSplitRead.class, RETURNS_SELF); RecordReader expectedReader = mock(RecordReader.class); - PrimaryKeyVectorDataSplit split = selectedSplit(); + IndexedSplit split = selectedSplit(); when(rawRead.createReader(split)).thenReturn(expectedReader); KeyValueTableRead tableRead = new KeyValueTableRead( @@ -103,7 +104,7 @@ void testKeyValueTableReadRoutesPhysicalSplitToRawReader() throws Exception { verify(rawRead).createReader(split); } - private static PrimaryKeyVectorDataSplit selectedSplit() { + private static IndexedSplit selectedSplit() { DataFileMeta dataFile = DataFileMeta.forAppend( "data-file", @@ -129,9 +130,8 @@ private static PrimaryKeyVectorDataSplit selectedSplit() { .withTotalBuckets(1) .withDataFiles(Collections.singletonList(dataFile)) .build(); - HashMap scores = new HashMap<>(); - scores.put(1, 0.5F); - return new PrimaryKeyVectorDataSplit(dataSplit, RoaringBitmap32.bitmapOf(1), scores); + return new IndexedSplit( + dataSplit, Collections.singletonList(new Range(1, 1)), new float[] {0.5F}); } private static class TestingFileReader implements FileRecordReader { diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java index 75eda2f88f3b..73b75cfec576 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorRawFileSplitReadTest.java @@ -23,6 +23,7 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.fs.Path; import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.options.Options; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.reader.ScoreRecordIterator; @@ -36,13 +37,12 @@ import org.apache.paimon.table.sink.BatchTableWrite; import org.apache.paimon.table.sink.BatchWriteBuilder; import org.apache.paimon.types.DataTypes; -import org.apache.paimon.utils.RoaringBitmap32; +import org.apache.paimon.utils.Range; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; -import java.util.HashMap; -import java.util.Map; +import java.util.Arrays; import static org.assertj.core.api.Assertions.assertThat; @@ -52,7 +52,7 @@ class PrimaryKeyVectorRawFileSplitReadTest { @TempDir java.nio.file.Path tempDir; @Test - void testRegularLimitDoesNotDropSelectedPhysicalPosition() throws Exception { + void testRegularLimitDoesNotDropSelectedPhysicalPositions() throws Exception { Path tablePath = new Path(tempDir.toUri()); Options options = new Options(); options.set(CoreOptions.PATH, tablePath.toString()); @@ -80,16 +80,22 @@ void testRegularLimitDoesNotDropSelectedPhysicalPosition() throws Exception { DataSplit dataSplit = table.newSnapshotReader().read().dataSplits().get(0); assertThat(dataSplit.dataFiles()).hasSize(1); - Map scores = new HashMap<>(); - scores.put(3, 0.25F); - PrimaryKeyVectorDataSplit vectorSplit = - new PrimaryKeyVectorDataSplit(dataSplit, RoaringBitmap32.bitmapOf(3), scores); + IndexedSplit vectorSplit = + new IndexedSplit( + dataSplit, + Arrays.asList(new Range(1, 1), new Range(3, 3)), + new float[] {0.5F, 0.25F}); try (RecordReader reader = - table.newRead().withLimit(1).createReader(vectorSplit)) { + table.newRead().withLimit(2).createReader(vectorSplit)) { ScoreRecordIterator batch = (ScoreRecordIterator) reader.readBatch(); InternalRow row = batch.next(); + assertThat(row.getInt(0)).isEqualTo(1); + assertThat(row.getInt(1)).isEqualTo(10); + assertThat(batch.returnedRowId()).isEqualTo(1); + assertThat(batch.returnedScore()).isEqualTo(0.5F); + row = batch.next(); assertThat(row.getInt(0)).isEqualTo(3); assertThat(row.getInt(1)).isEqualTo(30); assertThat(batch.returnedRowId()).isEqualTo(3); diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java index b58312e7de58..7cc16771f97f 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/SplitSerializerTest.java @@ -32,7 +32,6 @@ import org.apache.paimon.utils.IOUtils; import org.apache.paimon.utils.InstantiationUtil; import org.apache.paimon.utils.Range; -import org.apache.paimon.utils.RoaringBitmap32; import org.junit.jupiter.api.Test; @@ -62,48 +61,6 @@ public void testRoundTrip() throws IOException { } } - @Test - public void testPrimaryKeyVectorDataSplitRoundTrip() throws IOException { - PrimaryKeyVectorDataSplit split = primaryKeyVectorDataSplit(); - - Split actual = SplitSerializer.deserialize(SplitSerializer.serialize(split)); - - assertThat(actual).isEqualTo(split); - PrimaryKeyVectorDataSplit vectorSplit = (PrimaryKeyVectorDataSplit) actual; - assertThat(vectorSplit.rowPositions()).isEqualTo(RoaringBitmap32.bitmapOf(1, 3)); - assertThat(vectorSplit.score(1)).isEqualTo(0.9F); - assertThat(vectorSplit.score(3)).isEqualTo(0.7F); - } - - @Test - public void testPrimaryKeyVectorDataSplitJavaSerialization() throws Exception { - PrimaryKeyVectorDataSplit split = primaryKeyVectorDataSplit(); - - PrimaryKeyVectorDataSplit actual = - InstantiationUtil.deserializeObject( - InstantiationUtil.serializeObject(split), getClass().getClassLoader()); - - assertThat(actual).isEqualTo(split); - } - - private static PrimaryKeyVectorDataSplit primaryKeyVectorDataSplit() { - DataSplit dataSplit = - DataSplit.builder() - .withSnapshot(45L) - .withPartition(DataFileTestUtils.row(2026, 10)) - .withBucket(6) - .withTotalBuckets(8) - .withBucketPath("dt=20260708/bucket-6") - .withDataFiles( - Collections.singletonList(dataFile("vector-file", 1, 0, 4, 400L))) - .rawConvertible(true) - .build(); - Map scores = new LinkedHashMap<>(); - scores.put(1, 0.9F); - scores.put(3, 0.7F); - return new PrimaryKeyVectorDataSplit(dataSplit, RoaringBitmap32.bitmapOf(1, 3), scores); - } - @Test public void testGoldenFiles() throws IOException { boolean generateGoldenFiles = From 17ed7d8a3d8315db4af453d3c7913f13eda21eb8 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 20:21:14 +0800 Subject: [PATCH 3/3] [core] Separate primary-key indexed split read --- .../operation/PrimaryKeyIndexedSplitRead.java | 121 ++++++++++++++++++ .../paimon/operation/RawFileSplitRead.java | 73 ++--------- .../table/source/KeyValueTableRead.java | 2 + .../PrimaryKeyIndexedSplitReadProvider.java | 55 ++++++++ ...imaryKeyTableRawFileSplitReadProvider.java | 4 - .../PrimaryKeyIndexedSplitReadTest.java | 98 ++++++++++++++ .../PrimaryKeyVectorPositionReaderTest.java | 10 +- 7 files changed, 294 insertions(+), 69 deletions(-) create mode 100644 paimon-core/src/main/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitRead.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyIndexedSplitReadProvider.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitReadTest.java diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitRead.java b/paimon-core/src/main/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitRead.java new file mode 100644 index 000000000000..bb8d82e32c42 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitRead.java @@ -0,0 +1,121 @@ +/* + * 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.paimon.operation; + +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.disk.IOManager; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.PrimaryKeyVectorPositionReader; +import org.apache.paimon.table.source.Split; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringBitmap32; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.function.IntFunction; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Reads an {@link IndexedSplit} as physical positions in one primary-key data file. */ +public class PrimaryKeyIndexedSplitRead implements SplitRead { + + private final RawFileSplitRead rawRead; + + public PrimaryKeyIndexedSplitRead(RawFileSplitRead rawRead) { + this.rawRead = rawRead; + } + + @Override + public SplitRead forceKeepDelete() { + rawRead.forceKeepDelete(); + return this; + } + + @Override + public SplitRead withIOManager(@Nullable IOManager ioManager) { + rawRead.withIOManager(ioManager); + return this; + } + + @Override + public SplitRead withReadType(RowType readType) { + rawRead.withReadType(readType); + return this; + } + + @Override + public SplitRead withFilter(@Nullable Predicate predicate) { + rawRead.withFilter(predicate); + return this; + } + + @Override + public RecordReader createReader(Split split) throws IOException { + IndexedSplit indexedSplit = (IndexedSplit) split; + DataSplit dataSplit = indexedSplit.dataSplit(); + checkArgument( + dataSplit.dataFiles().size() == 1, + "Indexed split for a primary-key table must contain exactly one data file."); + DataFileMeta dataFile = dataSplit.dataFiles().get(0); + RoaringBitmap32 rowPositions = new RoaringBitmap32(); + float[] scores = indexedSplit.scores(); + Map scoreByPosition = scores == null ? null : new HashMap<>(); + int scoreIndex = 0; + for (Range range : indexedSplit.rowRanges()) { + checkArgument( + range.from >= 0 && range.to < dataFile.rowCount(), + "Indexed row range %s is outside data file %s row count %s.", + range, + dataFile.fileName(), + dataFile.rowCount()); + checkArgument( + range.to <= Integer.MAX_VALUE, + "Indexed row range %s exceeds supported physical positions.", + range); + for (long position = range.from; position <= range.to; position++) { + rowPositions.add((int) position); + if (scoreByPosition != null) { + checkArgument( + scoreIndex < scores.length, + "Scores length does not match row ranges in indexed split."); + scoreByPosition.put((int) position, scores[scoreIndex++]); + } + } + } + checkArgument(!rowPositions.isEmpty(), "Indexed split must select at least one row."); + if (scores != null) { + checkArgument( + scoreIndex == scores.length, + "Scores length does not match row ranges in indexed split."); + } + IntFunction scoreGetter = + scoreByPosition == null ? position -> Float.NaN : scoreByPosition::get; + FileRecordReader reader = rawRead.createFileReader(dataSplit, rowPositions); + return new PrimaryKeyVectorPositionReader(reader, rowPositions, scoreGetter); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java index f3b32054a057..56680aeb835e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java @@ -31,7 +31,6 @@ import org.apache.paimon.format.FormatKey; import org.apache.paimon.format.FormatReaderContext; import org.apache.paimon.fs.FileIO; -import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.io.DataFilePathFactory; import org.apache.paimon.io.DataFileRecordReader; @@ -49,14 +48,12 @@ import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.DeletionFile; import org.apache.paimon.table.source.IncrementalSplit; -import org.apache.paimon.table.source.PrimaryKeyVectorPositionReader; import org.apache.paimon.table.source.Split; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.FileStorePathFactory; import org.apache.paimon.utils.FormatReaderMapping; import org.apache.paimon.utils.FormatReaderMapping.Builder; import org.apache.paimon.utils.IOExceptionSupplier; -import org.apache.paimon.utils.Range; import org.apache.paimon.utils.RoaringBitmap32; import org.slf4j.Logger; @@ -69,11 +66,9 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.function.IntFunction; import static org.apache.paimon.predicate.PredicateBuilder.splitAnd; import static org.apache.paimon.table.SpecialFields.rowTypeWithRowTracking; -import static org.apache.paimon.utils.Preconditions.checkArgument; /** A {@link SplitRead} to read raw file directly from {@link DataSplit}. */ public class RawFileSplitRead implements SplitRead { @@ -155,9 +150,7 @@ public SplitRead withLimit(@Nullable Integer limit) { @Override public RecordReader createReader(Split s) throws IOException { - if (s instanceof IndexedSplit) { - return createReader((IndexedSplit) s); - } else if (s instanceof DataSplit) { + if (s instanceof DataSplit) { DataSplit split = (DataSplit) s; return createReader( split.partition(), @@ -217,45 +210,9 @@ public RecordReader createReader( return ConcatRecordReader.create(suppliers); } - private RecordReader createReader(IndexedSplit split) throws IOException { - DataSplit dataSplit = split.dataSplit(); - checkArgument( - dataSplit.dataFiles().size() == 1, - "Indexed split for a primary-key table must contain exactly one data file."); + FileRecordReader createFileReader( + DataSplit dataSplit, RoaringBitmap32 selectedPositions) throws IOException { DataFileMeta dataFile = dataSplit.dataFiles().get(0); - RoaringBitmap32 rowPositions = new RoaringBitmap32(); - float[] scores = split.scores(); - Map scoreByPosition = scores == null ? null : new HashMap<>(); - int scoreIndex = 0; - for (Range range : split.rowRanges()) { - checkArgument( - range.from >= 0 && range.to < dataFile.rowCount(), - "Indexed row range %s is outside data file %s row count %s.", - range, - dataFile.fileName(), - dataFile.rowCount()); - checkArgument( - range.to <= Integer.MAX_VALUE, - "Indexed row range %s exceeds supported physical positions.", - range); - for (long position = range.from; position <= range.to; position++) { - rowPositions.add((int) position); - if (scoreByPosition != null) { - checkArgument( - scoreIndex < scores.length, - "Scores length does not match row ranges in indexed split."); - scoreByPosition.put((int) position, scores[scoreIndex++]); - } - } - } - checkArgument(!rowPositions.isEmpty(), "Indexed split must select at least one row."); - if (scores != null) { - checkArgument( - scoreIndex == scores.length, - "Scores length does not match row ranges in indexed split."); - } - IntFunction scoreGetter = - scoreByPosition == null ? position -> Float.NaN : scoreByPosition::get; DeletionVector.Factory dvFactory = DeletionVector.factory( fileIO, dataSplit.dataFiles(), dataSplit.deletionFiles().orElse(null)); @@ -264,19 +221,17 @@ private RecordReader createReader(IndexedSplit split) throws IOExce dataFile.fileName(), () -> dvFactory.create(dataFile.fileName()).orElse(null)); DataFilePathFactory dataFilePathFactory = pathFactory.createDataFilePathFactory(dataSplit.partition(), dataSplit.bucket()); - FileRecordReader reader = - (FileRecordReader) - createFileReader( - dataSplit.partition(), - dataFilePathFactory, - dataFile, - // ANN has already selected the rows. Applying a regular - // TopN or limit before position filtering can drop hits. - createFormatReaderMappingBuilder(null, null), - dvFactories, - rowPositions) - .get(); - return new PrimaryKeyVectorPositionReader(reader, rowPositions, scoreGetter); + return (FileRecordReader) + createFileReader( + dataSplit.partition(), + dataFilePathFactory, + dataFile, + // The caller has already selected the rows. Applying a regular + // TopN or limit before position filtering can drop hits. + createFormatReaderMappingBuilder(null, null), + dvFactories, + selectedPositions) + .get(); } private Builder createFormatReaderMappingBuilder() { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java index fda7d70ffdf6..48b08df0c1dc 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java @@ -33,6 +33,7 @@ import org.apache.paimon.table.source.splitread.IncrementalChangelogReadProvider; import org.apache.paimon.table.source.splitread.IncrementalDiffReadProvider; import org.apache.paimon.table.source.splitread.MergeFileSplitReadProvider; +import org.apache.paimon.table.source.splitread.PrimaryKeyIndexedSplitReadProvider; import org.apache.paimon.table.source.splitread.PrimaryKeyTableRawFileSplitReadProvider; import org.apache.paimon.table.source.splitread.SplitReadProvider; import org.apache.paimon.types.RowType; @@ -67,6 +68,7 @@ public KeyValueTableRead( super(schema); this.readProviders = Arrays.asList( + new PrimaryKeyIndexedSplitReadProvider(batchRawReadSupplier, this::config), new PrimaryKeyTableRawFileSplitReadProvider( batchRawReadSupplier, this::config), new MergeFileSplitReadProvider(mergeReadSupplier, this::config), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyIndexedSplitReadProvider.java b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyIndexedSplitReadProvider.java new file mode 100644 index 000000000000..c96e995a349d --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyIndexedSplitReadProvider.java @@ -0,0 +1,55 @@ +/* + * 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.paimon.table.source.splitread; + +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.operation.PrimaryKeyIndexedSplitRead; +import org.apache.paimon.operation.RawFileSplitRead; +import org.apache.paimon.table.source.Split; +import org.apache.paimon.utils.LazyField; + +import java.util.function.Supplier; + +/** Provides physical-position reads for primary-key {@link IndexedSplit}s. */ +public class PrimaryKeyIndexedSplitReadProvider implements SplitReadProvider { + + private final LazyField splitRead; + + public PrimaryKeyIndexedSplitReadProvider( + Supplier supplier, SplitReadConfig splitReadConfig) { + this.splitRead = + new LazyField<>( + () -> { + PrimaryKeyIndexedSplitRead read = + new PrimaryKeyIndexedSplitRead(supplier.get()); + splitReadConfig.config(read); + return read; + }); + } + + @Override + public boolean match(Split split, Context context) { + return !context.forceKeepDelete() && split instanceof IndexedSplit; + } + + @Override + public LazyField get() { + return splitRead; + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java index 3ae77dd74776..713977dc590d 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/splitread/PrimaryKeyTableRawFileSplitReadProvider.java @@ -18,7 +18,6 @@ package org.apache.paimon.table.source.splitread; -import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.operation.RawFileSplitRead; import org.apache.paimon.table.source.DataSplit; @@ -36,9 +35,6 @@ public PrimaryKeyTableRawFileSplitReadProvider( @Override public boolean match(Split split, Context context) { - if (split instanceof IndexedSplit) { - return !context.forceKeepDelete(); - } if (!(split instanceof DataSplit)) { return false; } diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitReadTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitReadTest.java new file mode 100644 index 000000000000..3d4ef2413f5e --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexedSplitReadTest.java @@ -0,0 +1,98 @@ +/* + * 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.paimon.operation; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.reader.FileRecordReader; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.PrimaryKeyVectorPositionReader; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringBitmap32; + +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import java.util.Arrays; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests for {@link PrimaryKeyIndexedSplitRead}. */ +class PrimaryKeyIndexedSplitReadTest { + + @Test + void testConvertsRangesToPhysicalPositionsAndDelegates() throws Exception { + DataSplit dataSplit = dataSplit(); + IndexedSplit split = + new IndexedSplit( + dataSplit, + Arrays.asList(new Range(1, 1), new Range(3, 3)), + new float[] {0.5F, 0.25F}); + RawFileSplitRead rawRead = mock(RawFileSplitRead.class); + FileRecordReader fileReader = mock(FileRecordReader.class); + when(rawRead.createFileReader(eq(dataSplit), any(RoaringBitmap32.class))) + .thenReturn(fileReader); + + RecordReader reader = + new PrimaryKeyIndexedSplitRead(rawRead).createReader(split); + + assertThat(reader).isInstanceOf(PrimaryKeyVectorPositionReader.class); + ArgumentCaptor positions = ArgumentCaptor.forClass(RoaringBitmap32.class); + verify(rawRead).createFileReader(eq(dataSplit), positions.capture()); + assertThat(positions.getValue()).isEqualTo(RoaringBitmap32.bitmapOf(1, 3)); + } + + private static DataSplit dataSplit() { + DataFileMeta dataFile = + DataFileMeta.forAppend( + "data-file", + 100, + 5, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.APPEND, + null, + null, + null, + null); + return DataSplit.builder() + .withSnapshot(3) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(1) + .withDataFiles(Collections.singletonList(dataFile)) + .build(); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java index b600f94abd0a..c99f1d0c41f2 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReaderTest.java @@ -29,7 +29,6 @@ import org.apache.paimon.operation.RawFileSplitRead; import org.apache.paimon.reader.FileRecordIterator; import org.apache.paimon.reader.FileRecordReader; -import org.apache.paimon.reader.RecordReader; import org.apache.paimon.reader.ScoreRecordIterator; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.stats.SimpleStats; @@ -43,8 +42,8 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Answers.RETURNS_SELF; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; /** Tests physical row-position filtering with vector scores. */ class PrimaryKeyVectorPositionReaderTest { @@ -91,17 +90,16 @@ void testSelectedPositionsAcrossBatches() throws Exception { @Test void testKeyValueTableReadRoutesPhysicalSplitToRawReader() throws Exception { RawFileSplitRead rawRead = mock(RawFileSplitRead.class, RETURNS_SELF); - RecordReader expectedReader = mock(RecordReader.class); IndexedSplit split = selectedSplit(); - when(rawRead.createReader(split)).thenReturn(expectedReader); KeyValueTableRead tableRead = new KeyValueTableRead( () -> mock(MergeFileSplitRead.class), () -> rawRead, mock(TableSchema.class)); - assertThat(tableRead.createReader(split)).isSameAs(expectedReader); - verify(rawRead).createReader(split); + assertThat(tableRead.createReader(split)) + .isInstanceOf(PrimaryKeyVectorPositionReader.class); + verify(rawRead, never()).createReader(split); } private static IndexedSplit selectedSplit() {