package com.aliyun.tablestore.grid.core; import com.alicloud.openservices.tablestore.AsyncClientInterface; import com.alicloud.openservices.tablestore.TableStoreCallback; import com.alicloud.openservices.tablestore.model.GetRowRequest; import com.alicloud.openservices.tablestore.model.GetRowResponse; import com.aliyun.tablestore.grid.GridDataFetcher; import com.aliyun.tablestore.grid.model.GetDataParam; import com.aliyun.tablestore.grid.model.GridDataSet; import com.aliyun.tablestore.grid.model.GridDataSetMeta; import com.aliyun.tablestore.grid.model.StoreOptions; import com.aliyun.tablestore.grid.model.grid.*; import com.aliyun.tablestore.grid.utils.BlockUtil; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicInteger; import static com.aliyun.tablestore.grid.consts.Constants.DATA_BLOCK_COL_NAME_FORMAT; /** * not thread-safe */ public class TableStoreDataFetcher implements GridDataFetcher { private AsyncClientInterface asyncClient; private GridDataSetMeta meta; private String tableName; private long dataSizeLimitForFetch; private Collection<String> variables; private Range tRange; private Range zRange; private Range xRange; private Range yRange; public TableStoreDataFetcher(AsyncClientInterface asyncClient, String tableName, GridDataSetMeta meta, long dataSizeLimitForFetch) { this.asyncClient = asyncClient; this.tableName = tableName; this.meta = meta; this.dataSizeLimitForFetch = dataSizeLimitForFetch; this.variables = this.meta.getVariables(); this.tRange = new Range(0, this.meta.gettSize()); this.zRange = new Range(0, this.meta.getzSize()); this.xRange = new Range(0, this.meta.getxSize()); this.yRange = new Range(0, this.meta.getySize()); } public GridDataFetcher setVariablesToGet(Collection<String> variables) { this.variables = variables; return this; } @Override public GridDataFetcher setT(int t) { return setTRange(new Range(t, t+1)); } @Override public GridDataFetcher setTRange(Range range) { if (range.getStart() < 0 || range.getEnd() > meta.gettSize()) { throw new IllegalArgumentException("range invalid"); } this.tRange = range; return this; } @Override public GridDataFetcher setZ(int z) { return setZRange(new Range(z, z+1)); } @Override public GridDataFetcher setZRange(Range range) { if (range.getStart() < 0 || range.getEnd() > meta.getzSize()) { throw new IllegalArgumentException("range invalid"); } this.zRange = range; return this; } @Override public GridDataFetcher setX(int x) { return setXRange(new Range(x, x+1)); } @Override public GridDataFetcher setXRange(Range range) { if (range.getStart() < 0 || range.getEnd() > meta.getxSize()) { throw new IllegalArgumentException("range invalid"); } this.xRange = range; return this; } @Override public GridDataFetcher setY(int y) { return setYRange(new Range(y, y+1)); } @Override public GridDataFetcher setYRange(Range range) { if (range.getStart() < 0 || range.getEnd() > meta.getySize()) { throw new IllegalArgumentException("range invalid"); } this.yRange = range; return this; } @Override public GridDataFetcher setOriginShape(int[] origin, int[] shape) { if (origin.length != 4 || shape.length != 4) { throw new IllegalArgumentException("the length of origin and shape must be 4"); } setTRange(new Range(origin[0], origin[0] + shape[0])); setZRange(new Range(origin[1], origin[1] + shape[1])); setXRange(new Range(origin[2], origin[2] + shape[2])); setYRange(new Range(origin[3], origin[3] + shape[3])); return this; } public int[] getOrigin() { return new int[] {tRange.getStart(), zRange.getStart(), xRange.getStart(), yRange.getStart()}; } public int[] getShape() { return new int[] {tRange.getSize(), zRange.getSize(), xRange.getSize(), yRange.getSize()}; } private long calcDataSize(int variableCount) { long dataSize = variableCount; return dataSize * meta.getDataType().getSize() * tRange.getSize() * zRange.getSize() * xRange.getSize() * yRange.getSize(); } private List<String> getColumnsToGet() { if (!meta.getStoreOptions().getStoreType().equals(StoreOptions.StoreType.SLICE)) { throw new IllegalArgumentException("unsupported store type"); } Plane plane = new Plane(new Range(meta.getxSize()), new Range(meta.getySize())); Plane subPlane = new Plane(xRange, yRange); if (plane.equals(subPlane)) { return null; } List<String> columnsToGet = new ArrayList<String>(); List<Point> points = BlockUtil.calcBlockPointsCanCoverSubPlane(plane, subPlane, meta.getStoreOptions().getxSplitCount(), meta.getStoreOptions().getySplitCount()); for (Point point : points) { columnsToGet.add(String.format(DATA_BLOCK_COL_NAME_FORMAT, point.getX(), point.getY())); } return columnsToGet; } private void addTask(final AtomicInteger counter, final byte[] buffer, final int pos, String variable, int t, int z, final CountDownLatch latch, final Queue<Exception> exceptions) { GetDataParam param = new GetDataParam(tableName, meta.getGridDataSetId(), variable, t, z, getColumnsToGet()); asyncClient.getRow(RequestBuilder.buildGetDataRequest(param), new TableStoreCallback<GetRowRequest, GetRowResponse>() { @Override public void onCompleted(GetRowRequest req, GetRowResponse res) { try { if (res.getRow() == null) { exceptions.add(new RuntimeException("the row in not exist, pk: " + req.getRowQueryCriteria().getPrimaryKey())); } RowParser.parseGridFromRow(res.getRow(), new Plane(xRange, yRange), meta, buffer, pos); counter.incrementAndGet(); } catch (Exception ex) { exceptions.add(ex); } finally { latch.countDown(); } } @Override public void onFailed(GetRowRequest req, Exception ex) { try { exceptions.add(ex); } finally { latch.countDown(); } } }); } public GridDataSet fetch() throws Exception { long totalFetchDataSize = calcDataSize(variables.size()); if (totalFetchDataSize == 0) { throw new RuntimeException("no data to fetch"); } if (totalFetchDataSize > dataSizeLimitForFetch) { throw new RuntimeException("exceed the max data limit for fetch"); } GridDataSet dataSet = new GridDataSet(meta); CountDownLatch latch = new CountDownLatch(variables.size() * tRange.getSize() * zRange.getSize()); Queue<Exception> exceptions = new ConcurrentLinkedQueue<Exception>(); AtomicInteger counter = new AtomicInteger(); int taskCount = 0; for (String variable : variables) { int dataSize = (int) calcDataSize(1); byte[] data = new byte[dataSize]; ByteBuffer buffer = ByteBuffer.wrap(data).asReadOnlyBuffer(); dataSet.addVariable(variable, new Grid4D(buffer, meta.getDataType(), getOrigin(), getShape())); int curPos = 0; for (int t = tRange.getStart(); t < tRange.getEnd(); t++) { for (int z = zRange.getStart(); z < zRange.getEnd(); z++) { addTask(counter, data, curPos, variable, t, z, latch, exceptions); curPos += xRange.getSize() * yRange.getSize() * meta.getDataType().getSize(); taskCount++; } } } latch.await(); if (!exceptions.isEmpty()) { throw exceptions.peek(); } if (counter.get() != taskCount) { throw new RuntimeException("not all task success"); } return dataSet; } }