Skip to content

Commit 27af406

Browse files
authored
Add tsblock tsfile writer (#753)
1 parent ed91d2a commit 27af406

3 files changed

Lines changed: 492 additions & 12 deletions

File tree

java/tsfile/src/main/java/org/apache/tsfile/write/chunk/AlignedChunkGroupWriterImpl.java

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -51,20 +51,20 @@
5151
import java.util.stream.Collectors;
5252

5353
public class AlignedChunkGroupWriterImpl implements IChunkGroupWriter {
54-
private static final Logger LOG = LoggerFactory.getLogger(AlignedChunkGroupWriterImpl.class);
54+
protected static final Logger LOG = LoggerFactory.getLogger(AlignedChunkGroupWriterImpl.class);
5555

56-
private final IDeviceID deviceId;
56+
protected final IDeviceID deviceId;
5757

5858
// measurementID -> ValueChunkWriter
59-
private final Map<String, ValueChunkWriter> valueChunkWriterMap = new LinkedHashMap<>();
59+
protected final Map<String, ValueChunkWriter> valueChunkWriterMap = new LinkedHashMap<>();
6060

61-
private final TimeChunkWriter timeChunkWriter;
61+
protected final TimeChunkWriter timeChunkWriter;
6262

63-
private final EncryptParameter encryprParam;
63+
protected final EncryptParameter encryprParam;
6464

65-
private long lastTime = Long.MIN_VALUE;
66-
private boolean isInitLastTime = false;
67-
private boolean convertColumnNameToLowerCase = false;
65+
protected long lastTime = Long.MIN_VALUE;
66+
protected boolean isInitLastTime = false;
67+
protected boolean convertColumnNameToLowerCase = false;
6868

6969
public AlignedChunkGroupWriterImpl(IDeviceID deviceId) {
7070
this.deviceId = deviceId;
@@ -392,7 +392,7 @@ private void writeEmptyDataInOneRow(List<ValueChunkWriter> valueChunkWriterList)
392392
* check occupied memory size, if it exceeds the PageSize threshold, construct a page and put it
393393
* to pageBuffer
394394
*/
395-
private boolean checkPageSizeAndMayOpenANewPage() {
395+
protected boolean checkPageSizeAndMayOpenANewPage() {
396396
if (timeChunkWriter.checkPageSizeAndMayOpenANewPage()) {
397397
return true;
398398
}
@@ -404,21 +404,21 @@ private boolean checkPageSizeAndMayOpenANewPage() {
404404
return false;
405405
}
406406

407-
private void writePageToPageBuffer() {
407+
protected void writePageToPageBuffer() {
408408
timeChunkWriter.writePageToPageBuffer();
409409
for (ValueChunkWriter valueChunkWriter : valueChunkWriterMap.values()) {
410410
valueChunkWriter.writePageToPageBuffer();
411411
}
412412
}
413413

414-
private void sealAllChunks() {
414+
protected void sealAllChunks() {
415415
timeChunkWriter.sealCurrentPage();
416416
for (ValueChunkWriter valueChunkWriter : valueChunkWriterMap.values()) {
417417
valueChunkWriter.sealCurrentPage();
418418
}
419419
}
420420

421-
private void checkIsHistoryData(long time) throws WriteProcessException {
421+
protected void checkIsHistoryData(long time) throws WriteProcessException {
422422
if (isInitLastTime && time <= lastTime) {
423423
throw new WriteProcessException(
424424
"Not allowed to write out-of-order data in timeseries "
Lines changed: 270 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,270 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.tsfile.write.v4;
21+
22+
import org.apache.tsfile.block.column.Column;
23+
import org.apache.tsfile.common.conf.TSFileConfig;
24+
import org.apache.tsfile.exception.write.WriteProcessException;
25+
import org.apache.tsfile.file.metadata.IDeviceID;
26+
import org.apache.tsfile.file.metadata.StringArrayDeviceID;
27+
import org.apache.tsfile.file.metadata.TableSchema;
28+
import org.apache.tsfile.read.common.block.TsBlock;
29+
import org.apache.tsfile.read.common.block.column.TimeColumn;
30+
import org.apache.tsfile.utils.Pair;
31+
import org.apache.tsfile.write.UnSupportedDataTypeException;
32+
import org.apache.tsfile.write.chunk.AlignedChunkGroupWriterImpl;
33+
import org.apache.tsfile.write.chunk.IChunkGroupWriter;
34+
import org.apache.tsfile.write.chunk.TableChunkGroupWriterImpl;
35+
import org.apache.tsfile.write.chunk.ValueChunkWriter;
36+
import org.apache.tsfile.write.schema.IMeasurementSchema;
37+
38+
import java.io.File;
39+
import java.io.IOException;
40+
import java.util.ArrayList;
41+
import java.util.HashMap;
42+
import java.util.List;
43+
import java.util.Map;
44+
45+
/**
46+
* Convert TsBlock (table model) into TsFile format. Core responsibilities: 1. Split TsBlock by
47+
* device (based on tag columns) 2. Optionally generate a new time column per device 3. Dispatch
48+
* rows to corresponding ChunkGroupWriter
49+
*/
50+
public class TableTsBlock2TsFileWriter extends DeviceTableModelWriter {
51+
52+
private final boolean generateNewTimeColumn;
53+
private final int timeColumnIndexInTsBlock;
54+
private final int[] tagColumnIndexInTsBlock;
55+
private final int[] fieldColumnIndexInTsBlock;
56+
private final IMeasurementSchema[] fieldColumnSchemas;
57+
58+
private final String tableName;
59+
private final Map<IDeviceID, Long> deviceRowCountMap;
60+
61+
private int rowCount = 0;
62+
63+
public TableTsBlock2TsFileWriter(
64+
File file,
65+
TableSchema tableSchema,
66+
long memoryThreshold,
67+
boolean generateNewTimeColumn,
68+
int timeColumnIndexInTsBlock,
69+
int[] tagColumnIndexesInTsBlock,
70+
int[] fieldColumnIndexesInTsBlock,
71+
IMeasurementSchema[] fieldColumnSchemas)
72+
throws IOException {
73+
super(file, tableSchema, memoryThreshold);
74+
this.tableName = tableSchema.getTableName();
75+
this.generateNewTimeColumn = generateNewTimeColumn;
76+
this.timeColumnIndexInTsBlock = timeColumnIndexInTsBlock;
77+
this.tagColumnIndexInTsBlock = tagColumnIndexesInTsBlock;
78+
this.fieldColumnIndexInTsBlock = fieldColumnIndexesInTsBlock;
79+
this.deviceRowCountMap = generateNewTimeColumn ? new HashMap<>() : null;
80+
this.fieldColumnSchemas = fieldColumnSchemas;
81+
}
82+
83+
public void write(TsBlock tsBlock) throws IOException, WriteProcessException {
84+
if (tsBlock == null || tsBlock.isEmpty()) {
85+
return;
86+
}
87+
// Split TsBlock into device partitions and prepare time column
88+
Pair<Column, List<Pair<IDeviceID, Integer>>> timeColumnAndDeviceIdEndIndexPairs =
89+
splitTsBlockByDeviceAndGetTimeColumn(tsBlock);
90+
Column timeColumn = timeColumnAndDeviceIdEndIndexPairs.left;
91+
// Extract value columns according to schema mapping
92+
Column[] valueColumns = new Column[fieldColumnIndexInTsBlock.length];
93+
for (int i = 0; i < valueColumns.length; i++) {
94+
valueColumns[i] = tsBlock.getColumn(fieldColumnIndexInTsBlock[i]);
95+
}
96+
List<Pair<IDeviceID, Integer>> deviceIdEndIndexPairs = timeColumnAndDeviceIdEndIndexPairs.right;
97+
int startIndex = 0;
98+
99+
// Iterate each device segment and write data into its ChunkGroup
100+
for (Pair<IDeviceID, Integer> pair : deviceIdEndIndexPairs) {
101+
TableTsBlockChunkGroupWriterImpl chunkGroupWriter =
102+
(TableTsBlockChunkGroupWriterImpl) tryToInitialGroupWriter(pair.left, true, true);
103+
int writeCount = chunkGroupWriter.write(timeColumn, valueColumns, startIndex, pair.right);
104+
rowCount += writeCount;
105+
recordCount += writeCount;
106+
startIndex = pair.right;
107+
}
108+
109+
this.checkMemorySizeAndMayFlushChunks();
110+
}
111+
112+
/**
113+
* Split TsBlock by device boundary. If generateNewTimeColumn is true, generate a monotonically
114+
* increasing time column per device using deviceRowCountMap.
115+
*
116+
* @return Pair of (time column, device -> end index list)
117+
*/
118+
private Pair<Column, List<Pair<IDeviceID, Integer>>> splitTsBlockByDeviceAndGetTimeColumn(
119+
TsBlock tsBlock) {
120+
long[] timestamps = null;
121+
if (generateNewTimeColumn) {
122+
timestamps = new long[tsBlock.getPositionCount()];
123+
}
124+
List<Pair<IDeviceID, Integer>> deviceSplitResult = new ArrayList<>();
125+
IDeviceID lastDeviceID = null;
126+
long lastDeviceCount = 0;
127+
128+
// Iterate rows and detect device boundary changes
129+
for (int i = 0; i < tsBlock.getPositionCount(); i++) {
130+
IDeviceID currDeviceID = getDeviceId(tsBlock, i);
131+
// Device changed, flush previous segment
132+
if (!currDeviceID.equals(lastDeviceID)) {
133+
if (lastDeviceID != null) {
134+
deviceSplitResult.add(new Pair(lastDeviceID, i));
135+
if (generateNewTimeColumn) {
136+
deviceRowCountMap.put(lastDeviceID, lastDeviceCount);
137+
}
138+
}
139+
lastDeviceID = currDeviceID;
140+
if (generateNewTimeColumn) {
141+
lastDeviceCount = deviceRowCountMap.getOrDefault(lastDeviceID, 0L);
142+
}
143+
}
144+
// Generate synthetic time if required
145+
if (generateNewTimeColumn) {
146+
timestamps[i] = lastDeviceCount++;
147+
}
148+
}
149+
150+
deviceSplitResult.add(new Pair(lastDeviceID, tsBlock.getPositionCount()));
151+
if (generateNewTimeColumn) {
152+
deviceRowCountMap.put(lastDeviceID, lastDeviceCount);
153+
return new Pair<>(new TimeColumn(timestamps.length, timestamps), deviceSplitResult);
154+
} else {
155+
return new Pair<>(tsBlock.getColumn(timeColumnIndexInTsBlock), deviceSplitResult);
156+
}
157+
}
158+
159+
private IDeviceID getDeviceId(TsBlock tsBlock, int rowIdx) {
160+
String[] segments = new String[tagColumnIndexInTsBlock.length + 1];
161+
segments[0] = tableName;
162+
for (int i = 0; i < tagColumnIndexInTsBlock.length; i++) {
163+
Column tagColumn = tsBlock.getColumn(tagColumnIndexInTsBlock[i]);
164+
if (tagColumn.isNull(rowIdx)) {
165+
segments[i + 1] = null;
166+
} else {
167+
segments[i + 1] = tagColumn.getBinary(rowIdx).getStringValue(TSFileConfig.STRING_CHARSET);
168+
}
169+
}
170+
return new StringArrayDeviceID(segments);
171+
}
172+
173+
@Override
174+
protected IChunkGroupWriter tryToInitialGroupWriter(
175+
IDeviceID deviceId, boolean isAligned, boolean isTableModel) throws IOException {
176+
IChunkGroupWriter groupWriter = groupWriters.get(deviceId);
177+
if (groupWriter == null) {
178+
groupWriter = new TableTsBlockChunkGroupWriterImpl(deviceId);
179+
((AlignedChunkGroupWriterImpl) groupWriter)
180+
.setLastTime(alignedDeviceLastTimeMap.get(deviceId));
181+
groupWriters.put(deviceId, groupWriter);
182+
}
183+
return groupWriter;
184+
}
185+
186+
public int getDeviceCount() {
187+
return alignedDeviceLastTimeMap.size();
188+
}
189+
190+
public int getRowCount() {
191+
return rowCount;
192+
}
193+
194+
/**
195+
* ChunkGroup writer for a single device. Responsible for writing time column and multiple value
196+
* columns.
197+
*/
198+
private class TableTsBlockChunkGroupWriterImpl extends TableChunkGroupWriterImpl {
199+
private final ValueChunkWriter[] valueChunkWriters;
200+
201+
public TableTsBlockChunkGroupWriterImpl(IDeviceID deviceId) throws IOException {
202+
super(deviceId);
203+
// Initialize ValueChunkWriter for each measurement
204+
this.valueChunkWriters = new ValueChunkWriter[fieldColumnSchemas.length];
205+
for (int i = 0; i < fieldColumnSchemas.length; i++) {
206+
valueChunkWriters[i] = tryToAddSeriesWriterInternal(fieldColumnSchemas[i]);
207+
}
208+
}
209+
210+
/**
211+
* Write a range of rows into chunk group.
212+
*
213+
* @param startRowIndex inclusive
214+
* @param endRowIndex exclusive
215+
*/
216+
public int write(Column timeColumn, Column[] valueColumns, int startRowIndex, int endRowIndex)
217+
throws WriteProcessException {
218+
int pointCount = 0;
219+
for (int rowIndex = startRowIndex; rowIndex < endRowIndex; rowIndex++) {
220+
if (timeColumn.isNull(rowIndex)) {
221+
throw new WriteProcessException("All values in time column should not be null");
222+
}
223+
long time = timeColumn.getLong(rowIndex);
224+
checkIsHistoryData(time);
225+
for (int valueColumnIndex = 0; valueColumnIndex < valueColumns.length; valueColumnIndex++) {
226+
Column valueColumn = valueColumns[valueColumnIndex];
227+
ValueChunkWriter valueChunkWriter = valueChunkWriters[valueColumnIndex];
228+
boolean isNull = valueColumn.isNull(rowIndex);
229+
switch (valueChunkWriter.getDataType()) {
230+
case BOOLEAN:
231+
valueChunkWriter.write(
232+
time, isNull ? false : valueColumn.getBoolean(rowIndex), isNull);
233+
break;
234+
case INT32:
235+
case DATE:
236+
valueChunkWriter.write(time, isNull ? 0 : valueColumn.getInt(rowIndex), isNull);
237+
break;
238+
case INT64:
239+
case TIMESTAMP:
240+
valueChunkWriter.write(time, isNull ? 0 : valueColumn.getLong(rowIndex), isNull);
241+
break;
242+
case FLOAT:
243+
valueChunkWriter.write(time, isNull ? 0 : valueColumn.getFloat(rowIndex), isNull);
244+
break;
245+
case DOUBLE:
246+
valueChunkWriter.write(time, isNull ? 0 : valueColumn.getDouble(rowIndex), isNull);
247+
break;
248+
case TEXT:
249+
case BLOB:
250+
case STRING:
251+
valueChunkWriter.write(time, isNull ? null : valueColumn.getBinary(rowIndex), isNull);
252+
break;
253+
default:
254+
throw new UnSupportedDataTypeException(
255+
String.format(
256+
"Data type %s is not supported.", valueChunkWriter.getDataType().getType()));
257+
}
258+
}
259+
timeChunkWriter.write(time);
260+
lastTime = time;
261+
isInitLastTime = true;
262+
if (checkPageSizeAndMayOpenANewPage()) {
263+
writePageToPageBuffer();
264+
}
265+
pointCount++;
266+
}
267+
return pointCount;
268+
}
269+
}
270+
}

0 commit comments

Comments
 (0)