maedhroz commented on code in PR #254: URL: https://github.com/apache/cassandra-analytics/pull/254#discussion_r4406668623
########## cassandra-five-zero-bridge/src/main/java/org/apache/cassandra/spark/reader/sai/SaiIndexReader.java: ########## @@ -0,0 +1,596 @@ +/* + * 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.cassandra.spark.reader.sai; + +import java.io.Closeable; +import java.io.IOException; +import java.math.BigInteger; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.apache.cassandra.bridge.TokenRange; +import org.apache.cassandra.db.ClusteringComparator; +import org.apache.cassandra.db.ColumnFamilyStore; +import org.apache.cassandra.db.DecoratedKey; +import org.apache.cassandra.db.Keyspace; +import org.apache.cassandra.db.PartitionPosition; +import org.apache.cassandra.db.PartitionRangeReadCommand; +import org.apache.cassandra.db.filter.RowFilter; +import org.apache.cassandra.dht.Bounds; +import org.apache.cassandra.dht.IPartitioner; +import org.apache.cassandra.dht.Token; +import org.apache.cassandra.index.sai.QueryContext; +import org.apache.cassandra.index.sai.StorageAttachedIndex; +import org.apache.cassandra.index.sai.disk.PrimaryKeyMap; +import org.apache.cassandra.index.sai.disk.format.IndexDescriptor; +import org.apache.cassandra.index.sai.disk.format.SSTableIndexFileAccess; +import org.apache.cassandra.index.sai.disk.v1.MetadataSource; +import org.apache.cassandra.index.sai.disk.v1.PerColumnIndexFiles; +import org.apache.cassandra.index.sai.disk.v1.segment.IndexSegmentSearcher; +import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata; +import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator; +import org.apache.cassandra.index.sai.iterators.KeyRangeIterator; +import org.apache.cassandra.index.sai.iterators.KeyRangeUnionIterator; +import org.apache.cassandra.index.sai.plan.SaiPlanAccessor; +import org.apache.cassandra.index.sai.plan.Expression; +import org.apache.cassandra.index.sai.plan.QueryController; +import org.apache.cassandra.index.sai.plan.SaiQueryPlanner; +import org.apache.cassandra.index.sai.utils.PrimaryKey; +import org.apache.cassandra.io.sstable.Descriptor; +import org.apache.cassandra.io.util.File; +import org.apache.cassandra.schema.TableMetadata; +import org.apache.cassandra.spark.data.SSTable; +import org.apache.cassandra.spark.sparksql.filters.SaiFilter; +import org.apache.cassandra.spark.sparksql.filters.SparkRangeFilter; +import org.apache.cassandra.spark.utils.Preconditions; +import org.apache.cassandra.utils.TokenUtils; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +/** + * Opens Cassandra 5 SAI components, executes the query plan once, and materializes the resulting candidate tokens + * as an exact compact token set. + * + * <p>Cassandra's native SAI planner analyzes the conjunctive query, including same-column range folding and + * intersection. For every planned expression, matches are first UNIONed across all SSTables before Cassandra + * intersects different expressions. This ordering is important for correctness: values contributing to a logically + * reconciled row may live in different SSTables. Applying conjunctions inside each SSTable could therefore create + * false negatives.</p> + * + * <p>SAI is a read-planning phase only. The native iterator is fully consumed before {@code Data.db} scanning starts, + * then all SAI resources are closed. The resulting exact tokens are shared by every participating SSTable so the normal + * Cassandra compaction / reconciliation path can still apply newer values, TTLs and tombstones. Spark keeps the + * original predicates as residual filters, so stale index entries can only create false positives.</p> + */ +public final class SaiIndexReader +{ + private static final Logger LOGGER = LoggerFactory.getLogger(SaiIndexReader.class); + + private SaiIndexReader() + { + throw new IllegalStateException("Static utility class"); + } + + /** + * Executes all SAI predicates once and returns the sorted candidate tokens as an exact compact token set. + * + * @return empty {@link Optional} when SAI cannot safely be used and the caller should fall back to the ordinary SSTable + * scan; an empty {@link CandidateTokens} means SAI executed successfully and found no candidates + */ + @NotNull + public static Optional<CandidateTokens> findCandidateTokens(@NotNull TableMetadata metadata, + @NotNull Set<SSTable> sstables, + @NotNull List<SaiFilter> filters, + @Nullable SparkRangeFilter sparkRangeFilter, + int maxCandidateTokens) + { + if (sstables.isEmpty() || filters.isEmpty()) + { + return Optional.of(CandidateTokens.empty()); + } + + ResourceGroup resources = new ResourceGroup(); + try + { + SaiQueryPlanner.Plan plan = SaiQueryPlanner.plan(metadata, filters); + if (plan == null) + { + return Optional.empty(); + } + + List<StorageAttachedIndex> resourceIndexes = plan.resourceIndexes(); + + List<SSTableResources> sstableResources = new ArrayList<>(sstables.size()); + for (SSTable sstable : sstables) + { + SSTableResources opened = SSTableResources.open(metadata, sstable, resourceIndexes); + resources.add(opened); + sstableResources.add(opened); + } + + QueryContext queryContext = new QueryContext(null, Long.MAX_VALUE); + try (KeyRangeIterator finalMatches = executePlan(metadata, + plan, + sstableResources, + queryContext, + sparkRangeFilter)) + { + finalMatches.setOnClose(resources::close); + seekToSparkRangeStart(finalMatches, + sparkRangeFilter, + metadata.partitioner, + metadata.comparator); + return collectCandidateTokens(finalMatches, + sparkRangeFilter, + metadata.partitioner, + maxCandidateTokens); + } + } + catch (Exception exception) + { + resources.close(); + // No Data.db rows have been emitted yet, so every SAI failure can safely fail open to a normal scan. + LOGGER.warn("Unable to use SAI for SSTable pruning, falling back to normal SSTable scan", exception); + return Optional.empty(); + } + } + + /** + * Fully consumes the sorted native SAI result stream and converts it to the exact token-set representation + * used by the Data.db read-planning phase. + */ + @NotNull + static Optional<CandidateTokens> collectCandidateTokens(@NotNull Iterator<PrimaryKey> matches, + @Nullable SparkRangeFilter sparkRangeFilter, + @NotNull IPartitioner partitioner, + int maxCandidateTokens) + { + Preconditions.checkArgument(maxCandidateTokens >= 0, "maxCandidateTokens must be non-negative"); + BigInteger upperToken = sparkRangeFilter == null + ? null + : sparkRangeFilter.tokenRange().upperEndpoint(); + CandidateTokens.Builder candidates = CandidateTokens.builder(partitioner); + while (matches.hasNext()) + { + DecoratedKey partitionKey = matches.next().partitionKey(); + BigInteger token = TokenUtils.tokenToBigInteger(partitionKey.getToken()); + // SAI iterators are sorted by primary key/token. Once the worker's upper + // token is crossed there cannot be another candidate belonging to this task. + if (upperToken != null && token.compareTo(upperToken) > 0) + { + break; + } + if (sparkRangeFilter == null || !sparkRangeFilter.skipPartition(token)) + { + candidates.add(token); Review Comment: Two possible tests that might hit this, if you want to use them: ``` @Test void testCollectCandidatesSkipsRepeatedMatchesWithinSamePartition() { List<Candidate> distinctCandidates = sortedCandidates("one", "two", "three"); Candidate widePartition = distinctCandidates.get(0); Candidate nextPartition = distinctCandidates.get(1); PrimaryKey.Factory factory = new PrimaryKey.Factory(Murmur3Partitioner.instance, new ClusteringComparator()); List<PrimaryKey> keys = new ArrayList<>(); for (int i = 0; i < 20; i++) { keys.add(factory.create(widePartition.partitionKey)); } keys.add(factory.create(nextPartition.partitionKey)); TrackingKeyRangeIterator matches = new TrackingKeyRangeIterator(keys); Optional<CandidateTokens> result = SaiIndexReader.collectCandidateTokens(matches, null, Murmur3Partitioner.instance, new ClusteringComparator(), 100); assertThat(result).isPresent(); assertThat(result.get().size()).isEqualTo(2); assertThat(result.get().contains(widePartition.token)).isTrue(); assertThat(result.get().contains(nextPartition.token)).isTrue(); assertThat(matches.skipCalls()).isLessThanOrEqualTo(2); } @Test void testCollectCandidatesDoesNotSkipPastFollowingPartition() { List<Candidate> distinctCandidates = sortedCandidates("one", "two", "three"); Candidate first = distinctCandidates.get(0); Candidate second = distinctCandidates.get(1); Candidate third = distinctCandidates.get(2); PrimaryKey.Factory factory = new PrimaryKey.Factory(Murmur3Partitioner.instance, new ClusteringComparator()); List<PrimaryKey> keys = new ArrayList<>(); for (int i = 0; i < 5; i++) { keys.add(factory.create(first.partitionKey)); } for (int i = 0; i < 5; i++) { keys.add(factory.create(second.partitionKey)); } keys.add(factory.create(third.partitionKey)); TrackingKeyRangeIterator matches = new TrackingKeyRangeIterator(keys); Optional<CandidateTokens> result = SaiIndexReader.collectCandidateTokens(matches, null, Murmur3Partitioner.instance, new ClusteringComparator(), 100); assertThat(result).isPresent(); assertThat(result.get().size()).isEqualTo(3); assertThat(result.get().contains(first.token)).isTrue(); assertThat(result.get().contains(second.token)).isTrue(); assertThat(result.get().contains(third.token)).isTrue(); assertThat(matches.skipCalls()).isLessThanOrEqualTo(3); } ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
