adelapena commented on code in PR #2310: URL: https://github.com/apache/cassandra/pull/2310#discussion_r1224161082
########## src/java/org/apache/cassandra/db/marshal/VectorType.java: ########## @@ -0,0 +1,613 @@ +/* + * 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.db.marshal; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; + +import org.apache.cassandra.cql3.CQL3Type; +import org.apache.cassandra.cql3.Term; +import org.apache.cassandra.cql3.Vectors; +import org.apache.cassandra.db.TypeSizes; +import org.apache.cassandra.exceptions.InvalidRequestException; +import org.apache.cassandra.serializers.MarshalException; +import org.apache.cassandra.serializers.TypeSerializer; +import org.apache.cassandra.transport.ProtocolVersion; +import org.apache.cassandra.utils.ByteBufferUtil; +import org.apache.cassandra.utils.JsonUtils; +import org.apache.cassandra.utils.bytecomparable.ByteComparable; +import org.apache.cassandra.utils.bytecomparable.ByteSource; + +public final class VectorType<T> extends AbstractType<List<T>> +{ + private static class Key + { + private final AbstractType<?> type; + private final int dimension; + + private Key(AbstractType<?> type, int dimension) + { + this.type = type; + this.dimension = dimension; + } + + private VectorType<?> create() + { + return new VectorType<>(type, dimension); + } + + @Override + public boolean equals(Object o) + { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + Key key = (Key) o; + return dimension == key.dimension && Objects.equals(type, key.type); + } + + @Override + public int hashCode() + { + return Objects.hash(type, dimension); + } + } + @SuppressWarnings("rawtypes") + private static final ConcurrentHashMap<Key, VectorType> instances = new ConcurrentHashMap<>(); + + public final AbstractType<T> elementType; + public final int dimension; + private final TypeSerializer<T> elementSerializer; + private final int valueLengthIfFixed; + private final VectorSerializer serializer; + + private VectorType(AbstractType<T> elementType, int dimension) + { + super(ComparisonType.CUSTOM); + if (dimension <= 0) + throw new InvalidRequestException(String.format("vectors may only have positive dimentions; given %d", dimension)); + this.elementType = elementType; + this.dimension = dimension; + this.elementSerializer = elementType.getSerializer(); + this.valueLengthIfFixed = elementType.isValueLengthFixed() ? + elementType.valueLengthIfFixed() * dimension : + super.valueLengthIfFixed(); + this.serializer = elementType.isValueLengthFixed() ? + new FixedLengthSerializer() : + new VariableLengthSerializer(); + } + + public static <T> VectorType<T> getInstance(AbstractType<T> elements, int dimension) + { + Key key = new Key(elements, dimension); + return instances.computeIfAbsent(key, Key::create); + } + + public static VectorType<?> getInstance(TypeParser parser) + { + TypeParser.Vector v = parser.getVectorParameters(); + return new VectorType<>(v.type, v.dimension); + } + + @Override + public boolean isVector() + { + return true; + } + + @Override + public <VL, VR> int compareCustom(VL left, ValueAccessor<VL> accessorL, VR right, ValueAccessor<VR> accessorR) + { + return getSerializer().compareCustom(left, accessorL, right, accessorR); + } + + @Override + public int valueLengthIfFixed() + { + return valueLengthIfFixed; + } + + @Override + public VectorSerializer getSerializer() + { + return serializer; + } + + public List<ByteBuffer> split(ByteBuffer buffer) + { + return split(buffer, ByteBufferAccessor.instance); + } + + public <V> List<V> split(V buffer, ValueAccessor<V> accessor) + { + return getSerializer().split(buffer, accessor); + } + + public float[] composeAsFloat(ByteBuffer input) + { + return composeAsFloat(input, ByteBufferAccessor.instance); + } + + public <V> float[] composeAsFloat(V input, ValueAccessor<V> accessor) + { + if (!(elementType instanceof FloatType)) + throw new IllegalStateException("Attempted to read as float, but element type is " + elementType.asCQL3Type()); + + if (isNull(input, accessor)) + return null; + float[] array = new float[dimension]; + int offset = 0; + for (int i = 0; i < dimension; i++) + { + array[i] = accessor.getFloat(input, offset); + offset += Float.BYTES; + } + return array; + } + + public ByteBuffer decomposeRaw(List<ByteBuffer> elements) + { + return decomposeRaw(elements, ByteBufferAccessor.instance); + } + + public <V> V decomposeRaw(List<V> elements, ValueAccessor<V> accessor) + { + return getSerializer().serializeRaw(elements, accessor); + } + + @Override + public <V> ByteSource asComparableBytes(ValueAccessor<V> accessor, V value, ByteComparable.Version version) + { + if (isNull(value, accessor)) + return null; + ByteSource[] srcs = new ByteSource[dimension]; + List<V> split = split(value, accessor); + for (int i = 0; i < dimension; i++) + srcs[i] = elementType.asComparableBytes(accessor, split.get(i), version); + return ByteSource.withTerminatorMaybeLegacy(version, 0x00, srcs); + } + + @Override + public <V> V fromComparableBytes(ValueAccessor<V> accessor, ByteSource.Peekable comparableBytes, ByteComparable.Version version) + { + if (comparableBytes == null) + return accessor.empty(); + assert version != ByteComparable.Version.LEGACY; // legacy translation is not reversible + + List<V> buffers = new ArrayList<>(); + int separator = comparableBytes.next(); + while (separator != ByteSource.TERMINATOR) + { + buffers.add(elementType.fromComparableBytes(accessor, comparableBytes, version)); + separator = comparableBytes.next(); + } + return decomposeRaw(buffers, accessor); + } + + @Override + public CQL3Type asCQL3Type() + { + return new CQL3Type.Vector(this); + } + + public AbstractType<T> getElementsType() + { + return elementType; + } + + // vector of nested types is hard to parse, so fall back to bytes string matching ListType + @Override + public <V> String getString(V value, ValueAccessor<V> accessor) + { + return BytesType.instance.getString(value, accessor); + } + + @Override + public ByteBuffer fromString(String source) throws MarshalException + { + try + { + return ByteBufferUtil.hexToBytes(source); + } + catch (NumberFormatException e) + { + throw new MarshalException(String.format("cannot parse '%s' as hex bytes", source), e); + } + } + + @Override + public List<AbstractType<?>> subTypes() + { + return Collections.singletonList(elementType); + } + + @Override + public String toJSONString(ByteBuffer buffer, ProtocolVersion protocolVersion) + { + return toJSONString(buffer, ByteBufferAccessor.instance, protocolVersion); + } + + @Override + public <V> String toJSONString(V value, ValueAccessor<V> accessor, ProtocolVersion protocolVersion) + { + StringBuilder sb = new StringBuilder(); + sb.append('['); + List<V> split = split(value, accessor); + for (int i = 0; i < dimension; i++) + { + if (i > 0) + sb.append(", "); + sb.append(elementType.toJSONString(split.get(i), accessor, protocolVersion)); + } + sb.append(']'); + return sb.toString(); + } + + @Override + public Term fromJSONObject(Object parsed) throws MarshalException + { + if (parsed instanceof String) + parsed = JsonUtils.decodeJson((String) parsed); + + if (!(parsed instanceof List)) + throw new MarshalException(String.format( + "Expected a list, but got a %s: %s", parsed.getClass().getSimpleName(), parsed)); + + List<?> list = (List<?>) parsed; + if (list.size() != dimension) + throw new MarshalException(String.format("List had incorrect size: expected %d but given %d; %s", dimension, list.size(), list)); + List<Term> terms = new ArrayList<>(list.size()); + for (Object element : list) + { + if (element == null) + throw new MarshalException("Invalid null element in list"); + terms.add(elementType.fromJSONObject(element)); + } + + return new Vectors.DelayedValue<>(this, terms); + } + + @Override + public boolean equals(Object o) + { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + VectorType<?> that = (VectorType<?>) o; + return dimension == that.dimension && Objects.equals(elementType, that.elementType); + } + + @Override + public int hashCode() + { + return Objects.hash(elementType, dimension); + } + + @Override + public String toString() + { + return toString(false); + } + + @Override + public String toString(boolean ignoreFreezing) + { + return getClass().getName() + TypeParser.stringifyVectorParameters(elementType, ignoreFreezing, dimension); + } + + private void check(List<?> values) + { + if (values.size() != dimension) + throw new MarshalException(String.format("Required %d elements, but saw %d", dimension, values.size())); + + // This code base always works with a list that is RandomAccess, so can use .get to avoid allocation + for (int i = 0; i < dimension; i++) + { + Object value = values.get(i); + if (value == null || (value instanceof ByteBuffer && elementSerializer.isNull((ByteBuffer) value))) + throw new MarshalException(String.format("Element at index %d is null (expected type %s); given %s", i, elementType.asCQL3Type(), values)); + } + } + + public abstract class VectorSerializer extends TypeSerializer<List<T>> Review Comment: Makes sense to me. -- 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]
