added support for long type datasets

This commit is contained in:
Mark Wolters 2023-12-05 15:27:12 -04:00
parent 796c786d1a
commit 6793828278
2 changed files with 84 additions and 0 deletions

View File

@ -43,6 +43,12 @@ public class EmbeddingGeneratorFactory {
}
return generators.get(type);
}
case "long" -> {
if (!generators.containsKey(type)) {
generators.put(type, new LongEmbeddingGenerator());
}
return generators.get(type);
}
default -> throw new RuntimeException("Unknown embedding type: " + type);
}
}

View File

@ -0,0 +1,78 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed 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 io.nosqlbench.virtdata.library.hdf5.helpers;
import java.util.List;
public class LongEmbeddingGenerator implements EmbeddingGenerator {
@Override
public List<Float> generateFloatListEmbeddingFrom(Object o, int[] dims) {
// in this case o will always be long[1][x]
long[] vector = ((long[][]) o)[0];
Float[] vector2 = new Float[vector.length];
for (int i = 0; i < vector.length; i++) {
vector2[i] = (float) vector[i];
}
return List.of(vector2);
}
@Override
public float[] generateFloatArrayEmbeddingFrom(Object o, int[] dims) {
long[] vector = ((long[][]) o)[0];
float[] vector2 = new float[vector.length];
for (int i = 0; i < vector.length; i++) {
vector2[i] = (float) vector[i];
}
return vector2;
}
@Override
public List<Long> generateLongListEmbeddingFrom(Object o, int[] dims) {
long[] vector = ((long[][]) o)[0];
Long[] vector2 = new Long[vector.length];
for (int i = 0; i < vector.length; i++) {
vector2[i] = vector[i];
}
return List.of(vector2);
}
@Override
public long[] generateLongArrayEmbeddingFrom(Object o, int[] dims) {
return ((long[][]) o)[0];
}
@Override
public List<Integer> generateIntListEmbeddingFrom(Object o, int[] dims) {
long[] vector = ((long[][]) o)[0];
Integer[] vector2 = new Integer[vector.length];
for (int i = 0; i < vector.length; i++) {
vector2[i] = Math.toIntExact(vector[i]);
}
return List.of(vector2);
}
@Override
public int[] generateIntArrayEmbeddingFrom(Object o, int[] dims) {
long[] vector = ((long[][]) o)[0];
int[] vector2 = new int[vector.length];
for (int i = 0; i < vector.length; i++) {
vector2[i] = Math.toIntExact(vector[i]);
}
return vector2;
}
}