lukasz-antoniak commented on code in PR #254: URL: https://github.com/apache/cassandra-analytics/pull/254#discussion_r4424253258
########## 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); Review Comment: Curious why `PostingListRangeIterator` does not respect the callback. Anyway, releasing resources in `finally` block. -- 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]
