com.aliyun.openservices.paifeaturestore.dao.FeatureViewFeatureDBDao Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of paifeaturestore-sdk Show documentation
Show all versions of paifeaturestore-sdk Show documentation
SDK for PAI-FeatureStore service
package com.aliyun.openservices.paifeaturestore.dao;
import com.aliyun.openservices.paifeaturestore.constants.FSType;
import com.aliyun.openservices.paifeaturestore.datasource.FeatureDBClient;
import com.aliyun.openservices.paifeaturestore.datasource.FeatureDBFactory;
import com.aliyun.openservices.paifeaturestore.datasource.RecordBlock;
import com.aliyun.openservices.paifeaturestore.datasource.UInt8ValueColumn;
import com.aliyun.openservices.paifeaturestore.domain.FeatureResult;
import com.aliyun.openservices.paifeaturestore.domain.FeatureStoreResult;
import com.aliyun.openservices.paifeaturestore.model.FeatureViewSeqConfig;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;
import java.util.stream.Collectors;
public class FeatureViewFeatureDBDao implements FeatureViewDao {
private static Log log = LogFactory.getLog(FeatureViewFeatureDBDao.class);//日志工厂
private FeatureDBClient featureDBClient;
private String database;
private String schema;
private String table;
private String primaryKeyField;
public Map fieldTypeMap;
private List fields;
private final List> writeData = new ArrayList<>();
private final ReentrantLock lock = new ReentrantLock();
private final Condition condition = lock.newCondition();
private final ExecutorService executor = Executors.newFixedThreadPool(4);
public FeatureViewFeatureDBDao(DaoConfig daoConfig) {
this.database = daoConfig.featureDBDatabase;
this.schema = daoConfig.featureDBSchema;
this.table = daoConfig.featureDBTable;
FeatureDBClient client = FeatureDBFactory.get(daoConfig.featureDBName);
if (null == client) {
throw new RuntimeException(String.format("featuredbclient:%s not found", daoConfig.featureDBName));
}
this.featureDBClient = client;
this.fieldTypeMap = daoConfig.fieldTypeMap;
this.primaryKeyField = daoConfig.primaryKeyField;
this.fields = daoConfig.fields;
this.startAsyncWrite();
}
@Override
public FeatureResult getFeatures(String[] keys, String[] selectFields) {
Set selectFieldSet = new HashSet<>(Arrays.asList(selectFields));
List keyList = Arrays.asList(keys);
final int GROUP_SIZE = 200;
List> groups = new ArrayList<>();
for (int i = 0; i < keyList.size(); i+= GROUP_SIZE) {
int end = i + GROUP_SIZE;
if (end > keyList.size()) {
end = keyList.size();
}
groups.add(keyList.subList(i, end));
}
List> futures = groups.stream().map(group-> CompletableFuture.supplyAsync(()->{
FeatureStoreResult featureResult = new FeatureStoreResult();
List> featuresList = new ArrayList<>(group.size());
try {
byte[] content = this.featureDBClient.requestFeatureDB(group, this.database, this.schema, this.table);
RecordBlock recordBlock = RecordBlock.getRootAsRecordBlock(ByteBuffer.wrap(content));
for(int i= 0; i < recordBlock.valuesLength(); i++) {
UInt8ValueColumn valueColumn = new UInt8ValueColumn();
recordBlock.values(valueColumn, i);
if (valueColumn.valueLength() < 2 ) {
continue;
}
ByteBuffer byteBuffer = valueColumn.valueAsByteBuffer();
byteBuffer = byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
byte protoFlag = byteBuffer.get();
byte protoVersion = byteBuffer.get();
if (protoFlag != 'F' || protoVersion != '1') {
String errorMsg =String.format("invalid proto version, %d, %d", protoFlag, protoVersion);
log.error(errorMsg);
throw new RuntimeException(errorMsg);
}
Map featureMap = new HashMap<>(selectFields.length);
for (String featureName : this.fields) {
if (featureName.equals(this.primaryKeyField)) {
continue;
}
byte isNull = byteBuffer.get();
if (1 == isNull) {
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, null);
}
continue;
}
switch (this.fieldTypeMap.get(featureName)) {
case FS_INT32:
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, byteBuffer.getInt());
}
break;
case FS_INT64:
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, byteBuffer.getLong());
}
break;
case FS_DOUBLE:
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, byteBuffer.getDouble());
}
break;
case FS_BOOLEAN:
byte boolValue = byteBuffer.get();
if (boolValue == 0) {
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, false);
}
} else {
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, true);
}
}
break;
case FS_FLOAT:
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, byteBuffer.getFloat());
}
break;
default:
int len = byteBuffer.getInt();
byte[] bytes = new byte[len];
byteBuffer.get(bytes, 0, len);
if (selectFieldSet.contains(featureName)) {
featureMap.put(featureName, new String(bytes));
}
break;
}
}
featureMap.put(this.primaryKeyField, group.get(i));
featuresList.add(featureMap);
}
featureResult.setFeatureDataList(featuresList);
} catch (Exception e) {
log.error(String.format("request featuredb error:%s", e.getMessage()));
throw new RuntimeException(e);
}
return featureResult;
})).collect(Collectors.toList());
CompletableFuture allFutures = CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]));
CompletableFuture> allFutureResults = allFutures.thenApply(v ->
futures.stream()
.map(CompletableFuture::join)
.collect(Collectors.toList())
);
//List featureResultList = completableFutureStream.map(CompletableFuture::join).collect(Collectors.toList());
List featureResultList = null;
try {
featureResultList = allFutureResults.get();
} catch (Exception e) {
throw new RuntimeException(e);
}
FeatureStoreResult featureResult = new FeatureStoreResult();
featureResult.setFeatureFields(selectFields);
featureResult.setFeatureFieldTypeMap(this.fieldTypeMap);
List> featureDataList = new ArrayList<>(keys.length);
for (FeatureResult result : featureResultList) {
if (null != result.getFeatureData()) {
featureDataList.addAll(result.getFeatureData());
}
}
featureResult.setFeatureDataList(featureDataList);
return featureResult;
}
@Override
public FeatureResult getSequenceFeatures(String[] keys, String userIdField, FeatureViewSeqConfig featureViewSeqConfig) {
return null;
}
@Override
public void writeFeatures(List> data) {
lock.lock();
try {
writeData.addAll(data);
if (writeData.size() >= 20) {
condition.signal();
}
} finally {
lock.unlock();
}
}
@Override
public void writeFlush() {
lock.lock();
try {
if (writeData.size() > 0) {
try {
doWriteFeatures();
this.executor.shutdown();
try {
if (!this.executor.awaitTermination(60, TimeUnit.SECONDS)) {
this.executor.shutdownNow(); // 取消正在执行的任务
}
} catch (InterruptedException e) {
this.executor.shutdownNow();
}
} catch (Exception e) {
log.error(String.format("request featuredb error:%s", e.getMessage()));
}
}
} finally {
lock.unlock();
}
}
private void startAsyncWrite() {
new Thread(() -> {
while (true) {
lock.lock();
try {
condition.await(50, java.util.concurrent.TimeUnit.MILLISECONDS);
if (!writeData.isEmpty()) {
doWriteFeatures();
}
} catch (InterruptedException e) {
e.printStackTrace();
} finally {
lock.unlock();
}
}
}).start();
}
private void doWriteFeatures() {
List> tempList = new ArrayList<>(writeData);
writeData.clear();
// 异步处理 tempList
this.executor.submit(()->{
try {
this.featureDBClient.writeFeatureDB(tempList, this.database, this.schema, this.table);
} catch (Exception e) {
log.error(String.format("request featuredb error:%s", e.getMessage()));
}
});
}
}