Records

[tazarov]

It is easier to work with records than columnar data as the default API way of returning data.

Example

package tech.amikos.chromadb;

import java.util.List;

public class Record {
    private final String content;
    private final String id;
    private final String uri;
    private final Object metadata;
    private final List<Float> embedding;

    public Record(String content, String id, Object metadata, String uri, List<Float> embedding) {
        this.content = content;
        this.id = id;
        this.metadata = metadata;
        this.uri = uri;
        this.embedding = embedding;
    }

    public String getContent() {
        return content;
    }

    public String getId() {
        return id;
    }

    public Object getMetadata() {
        return metadata;
    }

    public String getUri() {
        return uri;
    }

    public List<Float> getEmbedding() {
        return embedding;
    }
}

Record builder example:

package tech.amikos.chromadb;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class RecordBuilder {

    private String content;
    private String id;

    private Map<String, Object> metadata;

    private String uri;

    private List<Float> embedding;

    private RecordBuilder() {
    }

    public static RecordBuilder forContent(String content) {
        RecordBuilder builder = new RecordBuilder();
        builder.content = content;
        return builder;
    }

    public static RecordBuilder forEmbedding(List<Float> embedding) {
        RecordBuilder builder = new RecordBuilder();
        builder.embedding = embedding;
        return builder;
    }

    public RecordBuilder withId(String id) {
        this.id = id;
        return this;
    }

    public RecordBuilder withMetadata(String key, Object value) {
        if (this.metadata == null) {
            this.metadata = new HashMap<>();
        }
        this.metadata.put(key, value);
        return this;
    }

    public RecordBuilder withMetadatas(Map<String, Object> metadata) {
        this.metadata = metadata;
        return this;
    }

    public RecordBuilder withUri(String uri) {
        this.uri = uri;
        return this;
    }

    public RecordBuilder withEmbedding(List<Float> embedding) {
        this.embedding = embedding;
        return this;
    }

    public Record build() {
        return new Record(content, id, metadata, uri, embedding);
    }
}

RecordSet example

package tech.amikos.chromadb;

import tech.amikos.chromadb.ids.IdGenerator;

import java.util.List;
import java.util.stream.Collectors;

public class RecordSet {
    private List<Record> records;
    private IdGenerator idGenerator;

    public List<String> getIds() {
        return records.stream().map(Record::getId).collect(Collectors.toList());
    }

    public List<String> getDocuments() {
        return records.stream().map(Record::getContent).collect(Collectors.toList());
    }

    public List<List<Float>> getEmbeddings() {
        return records.stream().map(Record::getEmbedding).collect(Collectors.toList());
    }

    public List<Object> getMetadatas() {
        return records.stream().map(Record::getMetadata).collect(Collectors.toList());
    }

    public List<String> getUris() {
        return records.stream().map(Record::getUri).collect(Collectors.toList());
    }
}