Add TokenOnlyPrimaryKey class and replace isTokenOnly

k-rus · k-rus · commit 097761986d72 · 2025-12-08T11:29:32.000+01:00
It's part of the patch.
diff --git a/src/java/org/apache/cassandra/index/sai/disk/v1/PartitionAwarePrimaryKeyFactory.java b/src/java/org/apache/cassandra/index/sai/disk/v1/PartitionAwarePrimaryKeyFactory.java
@@ -35,13 +35,6 @@
  */
 public class PartitionAwarePrimaryKeyFactory implements PrimaryKey.Factory
 {
-    @Override
-    public PrimaryKey createTokenOnly(Token token)
-    {
-        assert token != null;
-        return new PartitionAwarePrimaryKey(token, null, null);
-    }
-
     @Override
     public PrimaryKey createDeferred(Token token, Supplier<PrimaryKey> primaryKeySupplier)
     {
@@ -76,6 +69,7 @@ public PrimaryKey loadDeferred()
             {
                 this.partitionKey = primaryKeySupplier.get().partitionKey();
                 primaryKeySupplier = null;
+                assert this.token == this.partitionKey.getToken() : "Deferred primary key must contain the same token";
             }
             return this;
         }
@@ -86,12 +80,6 @@ public PartitionAwarePrimaryKey forStaticRow()
             return this;
         }
 
-        @Override
-        public boolean isTokenOnly()
-        {
-            return partitionKey == null && primaryKeySupplier == null;
-        }
-
         @Override
         public Token token()
         {
diff --git a/src/java/org/apache/cassandra/index/sai/disk/v2/PrimaryKeyWithSource.java b/src/java/org/apache/cassandra/index/sai/disk/v2/PrimaryKeyWithSource.java
@@ -80,12 +80,6 @@ public Token token()
         return primaryKey().token();
     }
 
-    @Override
-    public boolean isTokenOnly()
-    {
-        return false;
-    }
-
     @Override
     public DecoratedKey partitionKey()
     {
diff --git a/src/java/org/apache/cassandra/index/sai/disk/v2/RowAwarePrimaryKeyFactory.java b/src/java/org/apache/cassandra/index/sai/disk/v2/RowAwarePrimaryKeyFactory.java
@@ -50,12 +50,6 @@ public RowAwarePrimaryKeyFactory(ClusteringComparator clusteringComparator)
         this.hasEmptyClustering = clusteringComparator.size() == 0;
     }
 
-    @Override
-    public PrimaryKey createTokenOnly(Token token)
-    {
-        return new RowAwarePrimaryKey(token, null, null, null);
-    }
-
     @Override
     public PrimaryKey createDeferred(Token token, Supplier<PrimaryKey> primaryKeySupplier)
     {
@@ -94,12 +88,6 @@ public RowAwarePrimaryKey forStaticRow()
             return new RowAwarePrimaryKey(token, partitionKey, Clustering.STATIC_CLUSTERING, primaryKeySupplier);
         }
 
-        @Override
-        public boolean isTokenOnly()
-        {
-            return partitionKey == null && clustering == null && primaryKeySupplier == null;
-        }
-
         @Override
         public Token token()
         {
@@ -129,6 +117,7 @@ public PrimaryKey loadDeferred()
                 this.partitionKey = deferredPrimaryKey.partitionKey();
                 this.clustering = deferredPrimaryKey.clustering();
                 primaryKeySupplier = null;
+                assert this.token == this.partitionKey.getToken() : "Deferred primary key must contain the same token";
             }
             return this;
         }
diff --git a/src/java/org/apache/cassandra/index/sai/disk/v2/TokenOnlyPrimaryKey.java b/src/java/org/apache/cassandra/index/sai/disk/v2/TokenOnlyPrimaryKey.java
@@ -0,0 +1,118 @@
+package org.apache.cassandra.index.sai.disk.v2;
+
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
+
+public class TokenOnlyPrimaryKey implements PrimaryKey
+{
+    protected final Token token;
+
+    public TokenOnlyPrimaryKey(Token token)
+    {
+        this.token = token;
+    }
+
+    @Override
+    public boolean isTokenOnly()
+    {
+        return true;
+    }
+
+    @Override
+    public Token token()
+    {
+        return token;
+    }
+
+    @Override
+    public DecoratedKey partitionKey()
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Clustering<?> clustering()
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public ByteSource asComparableBytes(Version version)
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int compareTo(PrimaryKey o)
+    {
+        return token().compareTo(o.token());
+    }
+
+    @Override
+    public long ramBytesUsed()
+    {
+        // TODO Auto-generated method stub
+        throw new UnsupportedOperationException("Unimplemented method 'ramBytesUsed'");
+    }
+
+    @Override
+    public PrimaryKey forStaticRow()
+    {
+        // TODO Auto-generated method stub
+        throw new UnsupportedOperationException("Unimplemented method 'forStaticRow'");
+    }
+
+    @Override
+    public PrimaryKey loadDeferred()
+    {
+        // TODO Auto-generated method stub
+        throw new UnsupportedOperationException("Unimplemented method 'loadDeferred'");
+    }
+
+    @Override
+    public ByteSource asComparableBytesMinPrefix(Version version)
+    {
+        // TODO Auto-generated method stub
+        throw new UnsupportedOperationException("Unimplemented method 'asComparableBytesMinPrefix'");
+    }
+
+    @Override
+    public ByteSource asComparableBytesMaxPrefix(Version version)
+    {
+        // TODO Auto-generated method stub
+        throw new UnsupportedOperationException("Unimplemented method 'asComparableBytesMaxPrefix'");
+    }
+
+    // @Override
+    // public int hashCode()
+    // {
+    //     return Objects.hash(token(), clusteringComparator);
+    // }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (o instanceof PrimaryKey)
+            return compareTo((PrimaryKey) o) == 0;
+        return false;
+    }
+
+    // @Override
+    // public boolean equals(Object o, boolean strict)
+    // {
+    //     if (o == null)
+    //         return false;
+    //     if (o instanceof PrimaryKey)
+    //         return compareTo((PrimaryKey) o, strict) == 0;
+    //     return false;
+    // }
+
+    @Override
+    public String toString()
+    {
+        return String.format("PrimaryKey: { token: %s }", token());
+    }
+}
diff --git a/src/java/org/apache/cassandra/index/sai/utils/PrimaryKey.java b/src/java/org/apache/cassandra/index/sai/utils/PrimaryKey.java
@@ -27,6 +27,7 @@
 import org.apache.cassandra.index.sai.disk.format.IndexFeatureSet;
 import org.apache.cassandra.index.sai.disk.v1.PartitionAwarePrimaryKeyFactory;
 import org.apache.cassandra.index.sai.disk.v2.RowAwarePrimaryKeyFactory;
+import org.apache.cassandra.index.sai.disk.v2.TokenOnlyPrimaryKey;
 import org.apache.cassandra.utils.bytecomparable.ByteComparable;
 import org.apache.cassandra.utils.bytecomparable.ByteSource;
 
@@ -55,7 +56,11 @@ interface Factory
          * @param token the {@link Token}
          * @return a {@link PrimaryKey} represented by a token only
          */
-        PrimaryKey createTokenOnly(Token token);
+        default PrimaryKey createTokenOnly(Token token)
+        {
+            assert token != null;
+            return new TokenOnlyPrimaryKey(token);
+        }
 
         /**
          * Creates a {@link PrimaryKey} that is represented by a {@link DecoratedKey}.
@@ -118,7 +123,10 @@ static Factory factory(ClusteringComparator clusteringComparator, IndexFeatureSe
      */
     PrimaryKey forStaticRow();
 
-    boolean isTokenOnly();
+    default boolean isTokenOnly()
+    {
+        return false;
+    }
 
     /**
      * Returns the {@link Token} associated with this primary key.
diff --git a/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeyWithByteComparable.java b/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeyWithByteComparable.java
@@ -76,12 +76,6 @@ protected boolean isIndexDataEqualToLiveData(ByteBuffer value)
         }
     }
 
-    @Override
-    public boolean isTokenOnly()
-    {
-        return false;
-    }
-
     @Override
     public int compareTo(PrimaryKey o)
     {
diff --git a/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeyWithScore.java b/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeyWithScore.java
@@ -62,12 +62,6 @@ protected boolean isIndexDataEqualToLiveData(ByteBuffer value)
         return true;
     }
 
-    @Override
-    public boolean isTokenOnly()
-    {
-        return false;
-    }
-
     @Override
     public int compareTo(PrimaryKey o)
     {

Original file line number	Diff line number	Diff line change
`@@ -35,13 +35,6 @@`
`35`	`35`	`*/`
`36`	`36`	`public class PartitionAwarePrimaryKeyFactory implements PrimaryKey.Factory`
`37`	`37`	`{`
`38`		`- @Override`
`39`		`- public PrimaryKey createTokenOnly(Token token)`
`40`		`- {`
`41`		`- assert token != null;`
`42`		`- return new PartitionAwarePrimaryKey(token, null, null);`
`43`		`- }`
`44`		`-`
`45`	`38`	`@Override`
`46`	`39`	`public PrimaryKey createDeferred(Token token, Supplier<PrimaryKey> primaryKeySupplier)`
`47`	`40`	`{`
`@@ -76,6 +69,7 @@ public PrimaryKey loadDeferred()`
`76`	`69`	`{`
`77`	`70`	`this.partitionKey = primaryKeySupplier.get().partitionKey();`
`78`	`71`	`primaryKeySupplier = null;`
	`72`	`+ assert this.token == this.partitionKey.getToken() : "Deferred primary key must contain the same token";`
`79`	`73`	`}`
`80`	`74`	`return this;`
`81`	`75`	`}`
`@@ -86,12 +80,6 @@ public PartitionAwarePrimaryKey forStaticRow()`
`86`	`80`	`return this;`
`87`	`81`	`}`
`88`	`82`
`89`		`- @Override`
`90`		`- public boolean isTokenOnly()`
`91`		`- {`
`92`		`- return partitionKey == null && primaryKeySupplier == null;`
`93`		`- }`
`94`		`-`
`95`	`83`	`@Override`
`96`	`84`	`public Token token()`
`97`	`85`	`{`
Original file line number	Diff line number	Diff line change
`@@ -80,12 +80,6 @@ public Token token()`
`80`	`80`	`return primaryKey().token();`
`81`	`81`	`}`
`82`	`82`
`83`		`- @Override`
`84`		`- public boolean isTokenOnly()`
`85`		`- {`
`86`		`- return false;`
`87`		`- }`
`88`		`-`
`89`	`83`	`@Override`
`90`	`84`	`public DecoratedKey partitionKey()`
`91`	`85`	`{`
Original file line number	Diff line number	Diff line change
`@@ -50,12 +50,6 @@ public RowAwarePrimaryKeyFactory(ClusteringComparator clusteringComparator)`
`50`	`50`	`this.hasEmptyClustering = clusteringComparator.size() == 0;`
`51`	`51`	`}`
`52`	`52`
`53`		`- @Override`
`54`		`- public PrimaryKey createTokenOnly(Token token)`
`55`		`- {`
`56`		`- return new RowAwarePrimaryKey(token, null, null, null);`
`57`		`- }`
`58`		`-`
`59`	`53`	`@Override`
`60`	`54`	`public PrimaryKey createDeferred(Token token, Supplier<PrimaryKey> primaryKeySupplier)`
`61`	`55`	`{`
`@@ -94,12 +88,6 @@ public RowAwarePrimaryKey forStaticRow()`
`94`	`88`	`return new RowAwarePrimaryKey(token, partitionKey, Clustering.STATIC_CLUSTERING, primaryKeySupplier);`
`95`	`89`	`}`
`96`	`90`
`97`		`- @Override`
`98`		`- public boolean isTokenOnly()`
`99`		`- {`
`100`		`- return partitionKey == null && clustering == null && primaryKeySupplier == null;`
`101`		`- }`
`102`		`-`
`103`	`91`	`@Override`
`104`	`92`	`public Token token()`
`105`	`93`	`{`
`@@ -129,6 +117,7 @@ public PrimaryKey loadDeferred()`
`129`	`117`	`this.partitionKey = deferredPrimaryKey.partitionKey();`
`130`	`118`	`this.clustering = deferredPrimaryKey.clustering();`
`131`	`119`	`primaryKeySupplier = null;`
	`120`	`+ assert this.token == this.partitionKey.getToken() : "Deferred primary key must contain the same token";`
`132`	`121`	`}`
`133`	`122`	`return this;`
`134`	`123`	`}`
Original file line number	Diff line number	Diff line change
`@@ -76,12 +76,6 @@ protected boolean isIndexDataEqualToLiveData(ByteBuffer value)`
`76`	`76`	`}`
`77`	`77`	`}`
`78`	`78`
`79`		`- @Override`
`80`		`- public boolean isTokenOnly()`
`81`		`- {`
`82`		`- return false;`
`83`		`- }`
`84`		`-`
`85`	`79`	`@Override`
`86`	`80`	`public int compareTo(PrimaryKey o)`
`87`	`81`	`{`
Original file line number	Diff line number	Diff line change
`@@ -62,12 +62,6 @@ protected boolean isIndexDataEqualToLiveData(ByteBuffer value)`
`62`	`62`	`return true;`
`63`	`63`	`}`
`64`	`64`
`65`		`- @Override`
`66`		`- public boolean isTokenOnly()`
`67`		`- {`
`68`		`- return false;`
`69`		`- }`
`70`		`-`
`71`	`65`	`@Override`
`72`	`66`	`public int compareTo(PrimaryKey o)`
`73`	`67`	`{`