dcapwell commented on code in PR #3785: URL: https://github.com/apache/cassandra/pull/3785#discussion_r1927772812
########## test/harry/main/org/apache/cassandra/harry/model/BytesPartitionState.java: ########## @@ -0,0 +1,554 @@ +/* + * 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.harry.model; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.NavigableSet; +import java.util.Objects; +import java.util.Set; +import java.util.TreeSet; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.annotation.Nullable; + +import org.apache.cassandra.cql3.ast.Symbol; +import org.apache.cassandra.db.Clustering; +import org.apache.cassandra.db.ClusteringComparator; +import org.apache.cassandra.db.marshal.AbstractType; +import org.apache.cassandra.dht.Token; +import org.apache.cassandra.harry.MagicConstants; +import org.apache.cassandra.harry.gen.BijectionCache; +import org.apache.cassandra.harry.gen.Bijections; +import org.apache.cassandra.harry.gen.ValueGenerators; +import org.apache.cassandra.harry.util.BitSet; +import org.apache.cassandra.schema.ColumnMetadata; +import org.apache.cassandra.schema.TableMetadata; +import org.apache.cassandra.utils.FastByteOperations; +import org.apache.cassandra.utils.ImmutableUniqueList; +import org.apache.cassandra.utils.TableUtil; + +public class BytesPartitionState +{ + private final Factory factory; + public final Clustering<ByteBuffer> key; + private final Token token; + private final PartitionState state; + + private BytesPartitionState(Factory factory, Clustering<ByteBuffer> key) + { + this.factory = factory; + this.key = key; + this.token = factory.metadata.partitioner.getToken(key.serializeAsPartitionKey()); + this.state = factory.partitionState(key); + } + + public void deleteRow(Clustering<ByteBuffer> clustering) + { + long cd = factory.clusteringCache.deflateOrUndefined(clustering); + if (MagicConstants.UNSET_DESCR == cd) + return; + state.delete(cd, MagicConstants.NO_TIMESTAMP); + } + + public void deleteColumns(Clustering<ByteBuffer> clustering, Set<Symbol> columns) + { + long cd = factory.clusteringCache.deflateOrUndefined(clustering); + if (cd != MagicConstants.UNSET_DESCR) + { + BitSet regularColumns = bitset(columns, true); + if (!regularColumns.allUnset()) + state.deleteRegularColumns(MagicConstants.NO_TIMESTAMP, cd, regularColumns); + } + deleteStaticColumns(columns); + } + + public void deleteStaticColumns(Set<Symbol> columns) + { + BitSet staticColumns = bitset(columns, false); + if (!staticColumns.allUnset()) + state.deleteStaticColumns(MagicConstants.NO_TIMESTAMP, staticColumns); + } + + private BitSet bitset(Set<Symbol> columns, boolean regular) + { + ImmutableUniqueList<Symbol> positions = regular ? factory.regularPositions : factory.staticPositions; + BitSet bitSet = new BitSet.BitSet64Bit(positions.size()); + for (int i = 0; i < positions.size(); i++) + { + Symbol column = positions.get(i); + if (columns.contains(column)) + bitSet.set(i); + } + return bitSet; + } + + public Ref ref() + { + return new Ref(factory, key, token); + } + + public PrimaryKey partitionRowRef() + { + return new PrimaryKey(ref(), null); + } + + public void setStaticColumns(Map<Symbol, ByteBuffer> values) + { + if (factory.staticPositions.isEmpty() || values.isEmpty()) + throw new IllegalStateException("Attempt to write to static columns; but they do not exist"); + long[] sds = toDescriptor(factory.staticPositions, values); + state.writeStatic(sds, MagicConstants.NO_TIMESTAMP); + } + + public void setColumns(Clustering<ByteBuffer> clustering, Map<Symbol, ByteBuffer> values, boolean writePrimaryKeyLiveness) + { + long cd = factory.clusteringCache.deflate(clustering); + long[] vds = toDescriptor(factory.regularPositions, values); + state.writeRegular(cd, vds, MagicConstants.NO_TIMESTAMP, writePrimaryKeyLiveness); + } + + private long[] toDescriptor(ImmutableUniqueList<Symbol> positions, Map<Symbol, ByteBuffer> values) + { + long[] vds = new long[positions.size()]; + for (int i = 0; i < positions.size(); i++) + { + Symbol column = positions.get(i); + if (values.containsKey(column)) + { + long vd = factory.valueCache.deflate(new Value(column.type(), values.get(column))); + vds[i] = vd; + } + else + { + vds[i] = MagicConstants.UNSET_DESCR; + } + } + return vds; + } + + private ByteBuffer[] fromDescriptor(ImmutableUniqueList<Symbol> positions, long[] values) + { + if (positions.size() != values.length) + throw new IllegalArgumentException(String.format("Attempted to extract values but expected columns didn't match; expected %s, but given %d values", positions, values.length)); + ByteBuffer[] bbs = new ByteBuffer[values.length]; + for (int i = 0; i < bbs.length; i++) + { + long vd = values[i]; + if (vd == MagicConstants.NIL_DESCR) + { + bbs[i] = null; + } + else + { + var value = factory.valueCache.inflate(vd); + Symbol column = positions.get(i); + if (!value.type.equals(column.type())) + throw new IllegalStateException(String.format("Given value descriptor %d that maps to the wrong type; expected %s, given %s", vd, column.type().asCQL3Type(), value.type.asCQL3Type())); + bbs[i] = value.value; + } + } + return bbs; + } + + public int size() + { + return state.rows().size(); + } + + public boolean isEmpty() + { + return state.rows().isEmpty(); + } + + @Nullable + public Row get(Clustering<ByteBuffer> clustering) + { + long cd = factory.clusteringCache.deflateOrUndefined(clustering); + if (cd == MagicConstants.UNSET_DESCR) + return null; + PartitionState.RowState rowState = state.rows().get(cd); + if (rowState == null) + return null; + return toRow(rowState); + } + + private Row toRow(PartitionState.RowState rowState) + { + Clustering<ByteBuffer> clustering; + ByteBuffer[] values; + if (PartitionState.STATIC_CLUSTERING == rowState.cd) + { + clustering = Clustering.STATIC_CLUSTERING; + values = fromDescriptor(factory.staticPositions, rowState.vds); + } + else + { + clustering = factory.clusteringCache.inflate(rowState.cd); + values = fromDescriptor(factory.regularPositions, rowState.vds); + } + return new Row(clustering, values); + } + + public Collection<Row> rows() + { + return state.rows().values().stream().map(this::toRow).collect(Collectors.toList()); + } + + public NavigableSet<Clustering<ByteBuffer>> clusteringKeys() + { + NavigableSet<Clustering<ByteBuffer>> navigableSet = new TreeSet<>(factory.clusteringComparator); + state.rows().keySet().stream().map(factory.clusteringCache::inflate).forEach(navigableSet::add); + return navigableSet; + } + + public Row staticRow() + { + return toRow(state.staticRow()); + } + + public boolean shouldDelete() + { + return state.shouldDelete(); + } + + static List<String> asCQL(ImmutableUniqueList<Symbol> columns, ByteBuffer[] row) + { + List<String> cql = new ArrayList<>(row.length); + for (int i = 0; i < row.length; i++) + { + AbstractType type = columns.get(i).type(); + cql.add(type.toCQLString(row[i])); + } + return cql; + } + + public class PrimaryKey implements Comparable<PrimaryKey> + { + public final BytesPartitionState.Ref partition; + @Nullable + public final Clustering<ByteBuffer> clustering; + + public PrimaryKey(BytesPartitionState.Ref partition, @Nullable Clustering<ByteBuffer> clustering) + { + this.partition = partition; + this.clustering = clustering; + } + + @Override + public int compareTo(PrimaryKey o) + { + int rc = partition.compareTo(o.partition); + if (rc != 0) return rc; + if (clustering == null) return rc; // if the partition matches, and clustering is null (partition doesn't have rows) then it would be a bug if o.clustering was non-null + rc = factory.clusteringComparator.compare(clustering, o.clustering); + return rc; + } + + @Override + public boolean equals(Object o) + { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + PrimaryKey that = (PrimaryKey) o; + return partition.equals(that.partition) && Objects.equals(clustering, that.clustering); + } + + @Override + public int hashCode() + { + return Objects.hash(partition, clustering); + } + } + + public static class Ref implements Comparable<Ref> Review Comment: this logic has 2 different types of ref 1) partition - this is the `(partition columns, token)` 2) row - this is `(partition ref, clustering columns)` in this case its a reference to a partition, aka `(partition columns, token)`. This is basically the same as `org.apache.cassandra.db.DecoratedKey` but with 1 big change. I don't have a single `ByteBuffer` for the partition, i have `ByteBuffer[]` so I can directly access each column. -- 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]
