iotdb-core/tsfile/src/test/java/org/apache/iotdb/tsfile/common/block/TsBlockTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.iotdb.tsfile.common.block;

import org.apache.iotdb.tsfile.common.conf.TSFileConfig;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.apache.iotdb.tsfile.read.common.block.TsBlock;
import org.apache.iotdb.tsfile.read.common.block.TsBlock.TsBlockSingleColumnIterator;
import org.apache.iotdb.tsfile.read.common.block.TsBlockBuilder;
import org.apache.iotdb.tsfile.read.common.block.column.BinaryColumn;
import org.apache.iotdb.tsfile.read.common.block.column.BooleanColumn;
import org.apache.iotdb.tsfile.read.common.block.column.Column;
import org.apache.iotdb.tsfile.read.common.block.column.DoubleColumn;
import org.apache.iotdb.tsfile.read.common.block.column.FloatColumn;
import org.apache.iotdb.tsfile.read.common.block.column.IntColumn;
import org.apache.iotdb.tsfile.read.common.block.column.LongColumn;
import org.apache.iotdb.tsfile.read.common.block.column.RunLengthEncodedColumn;
import org.apache.iotdb.tsfile.utils.Binary;

import org.junit.Assert;
import org.junit.Test;

import java.util.Arrays;
import java.util.Collections;
import java.util.Optional;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

public class TsBlockTest {

  private static final double DELTA = 0.000001;

  @Test
  public void testBooleanTsBlock() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    boolean[] valueArray = {true, false, false, false, false};
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.BOOLEAN));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      builder.getColumnBuilder(0).writeBoolean(valueArray[i]);
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof BooleanColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertFalse(tsBlock.getColumn(0).isNull(i));
      assertEquals(valueArray[i], tsBlock.getColumn(0).getBoolean(i));
    }
  }

  @Test
  public void testIntTsBlock() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    int[] valueArray = {10, 20, 30, 40, 50};
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      builder.getColumnBuilder(0).writeInt(valueArray[i]);
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof IntColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertFalse(tsBlock.getColumn(0).isNull(i));
      assertEquals(valueArray[i], tsBlock.getColumn(0).getInt(i));
    }
  }

  @Test
  public void testLongTsBlock() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    long[] valueArray = {10L, 20L, 30L, 40L, 50L};
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT64));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      builder.getColumnBuilder(0).writeLong(valueArray[i]);
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof LongColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertFalse(tsBlock.getColumn(0).isNull(i));
      assertEquals(valueArray[i], tsBlock.getColumn(0).getLong(i));
    }
  }

  @Test
  public void testFloatTsBlock() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    float[] valueArray = {10.0f, 20.0f, 30.0f, 40.0f, 50.0f};
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.FLOAT));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      builder.getColumnBuilder(0).writeFloat(valueArray[i]);
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof FloatColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertFalse(tsBlock.getColumn(0).isNull(i));
      assertEquals(valueArray[i], tsBlock.getColumn(0).getFloat(i), DELTA);
    }
  }

  @Test
  public void testDoubleTsBlock() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    double[] valueArray = {10.0, 20.0, 30.0, 40.0, 50.0};
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.DOUBLE));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      builder.getColumnBuilder(0).writeDouble(valueArray[i]);
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof DoubleColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertFalse(tsBlock.getColumn(0).isNull(i));
      assertEquals(valueArray[i], tsBlock.getColumn(0).getDouble(i), DELTA);
    }
  }

  @Test
  public void testBinaryTsBlock() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    Binary[] valueArray = {
      new Binary("10", TSFileConfig.STRING_CHARSET),
      new Binary("20", TSFileConfig.STRING_CHARSET),
      new Binary("30", TSFileConfig.STRING_CHARSET),
      new Binary("40", TSFileConfig.STRING_CHARSET),
      new Binary("50", TSFileConfig.STRING_CHARSET)
    };
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.TEXT));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      builder.getColumnBuilder(0).writeBinary(valueArray[i]);
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof BinaryColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertFalse(tsBlock.getColumn(0).isNull(i));
      assertEquals(valueArray[i], tsBlock.getColumn(0).getBinary(i));
    }
  }

  @Test
  public void testIntTsBlockWithNull() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    int[] valueArray = {10, 20, 30, 40, 50};
    boolean[] isNull = {false, false, true, true, false};
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      if (isNull[i]) {
        builder.getColumnBuilder(0).appendNull();

      } else {
        builder.getColumnBuilder(0).writeInt(valueArray[i]);
      }
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof IntColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertEquals(isNull[i], tsBlock.getColumn(0).isNull(i));
      if (!isNull[i]) {
        assertEquals(valueArray[i], tsBlock.getColumn(0).getInt(i));
      }
    }
  }

  @Test
  public void testIntTsBlockWithAllNull() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32));
    for (long l : timeArray) {
      builder.getTimeColumnBuilder().writeLong(l);
      builder.getColumnBuilder(0).appendNull();
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(1, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof RunLengthEncodedColumn);
    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertTrue(tsBlock.getColumn(0).isNull(i));
    }
  }

  @Test
  public void testMultiColumnTsBlockWithNull() {
    long[] timeArray = {1L, 2L, 3L, 4L, 5L};
    boolean[] booleanValueArray = {true, false, false, false, true};
    boolean[] booleanIsNull = {true, true, false, true, false};
    int[] intValueArray = {10, 20, 30, 40, 50};
    boolean[] intIsNull = {false, true, false, false, true};
    long[] longValueArray = {100L, 200L, 300L, 400, 500L};
    boolean[] longIsNull = {true, false, false, true, true};
    float[] floatValueArray = {1000.0f, 2000.0f, 3000.0f, 4000.0f, 5000.0f};
    boolean[] floatIsNull = {false, false, true, true, false};
    double[] doubleValueArray = {10000.0, 20000.0, 30000.0, 40000.0, 50000.0};
    boolean[] doubleIsNull = {true, false, false, true, false};
    Binary[] binaryValueArray = {
      new Binary("19970909", TSFileConfig.STRING_CHARSET),
      new Binary("ty", TSFileConfig.STRING_CHARSET),
      new Binary("love", TSFileConfig.STRING_CHARSET),
      new Binary("zm", TSFileConfig.STRING_CHARSET),
      new Binary("19950421", TSFileConfig.STRING_CHARSET)
    };
    boolean[] binaryIsNull = {false, false, false, false, false};

    TsBlockBuilder builder =
        new TsBlockBuilder(
            Arrays.asList(
                TSDataType.BOOLEAN,
                TSDataType.INT32,
                TSDataType.INT64,
                TSDataType.FLOAT,
                TSDataType.DOUBLE,
                TSDataType.TEXT));
    for (int i = 0; i < timeArray.length; i++) {
      builder.getTimeColumnBuilder().writeLong(timeArray[i]);
      if (booleanIsNull[i]) {
        builder.getColumnBuilder(0).appendNull();
      } else {
        builder.getColumnBuilder(0).writeBoolean(booleanValueArray[i]);
      }
      if (intIsNull[i]) {
        builder.getColumnBuilder(1).appendNull();
      } else {
        builder.getColumnBuilder(1).writeInt(intValueArray[i]);
      }
      if (longIsNull[i]) {
        builder.getColumnBuilder(2).appendNull();
      } else {
        builder.getColumnBuilder(2).writeLong(longValueArray[i]);
      }
      if (floatIsNull[i]) {
        builder.getColumnBuilder(3).appendNull();
      } else {
        builder.getColumnBuilder(3).writeFloat(floatValueArray[i]);
      }
      if (doubleIsNull[i]) {
        builder.getColumnBuilder(4).appendNull();
      } else {
        builder.getColumnBuilder(4).writeDouble(doubleValueArray[i]);
      }
      if (binaryIsNull[i]) {
        builder.getColumnBuilder(5).appendNull();
      } else {
        builder.getColumnBuilder(5).writeBinary(binaryValueArray[i]);
      }
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    assertEquals(timeArray.length, tsBlock.getPositionCount());
    assertEquals(6, tsBlock.getValueColumnCount());
    assertTrue(tsBlock.getColumn(0) instanceof BooleanColumn);
    assertTrue(tsBlock.getColumn(1) instanceof IntColumn);
    assertTrue(tsBlock.getColumn(2) instanceof LongColumn);
    assertTrue(tsBlock.getColumn(3) instanceof FloatColumn);
    assertTrue(tsBlock.getColumn(4) instanceof DoubleColumn);
    assertTrue(tsBlock.getColumn(5) instanceof BinaryColumn);

    for (int i = 0; i < timeArray.length; i++) {
      assertEquals(timeArray[i], tsBlock.getTimeByIndex(i));
      assertEquals(booleanIsNull[i], tsBlock.getColumn(0).isNull(i));
      if (!booleanIsNull[i]) {
        assertEquals(booleanValueArray[i], tsBlock.getColumn(0).getBoolean(i));
      }
      assertEquals(intIsNull[i], tsBlock.getColumn(1).isNull(i));
      if (!intIsNull[i]) {
        assertEquals(intValueArray[i], tsBlock.getColumn(1).getInt(i));
      }
      assertEquals(longIsNull[i], tsBlock.getColumn(2).isNull(i));
      if (!longIsNull[i]) {
        assertEquals(longValueArray[i], tsBlock.getColumn(2).getLong(i));
      }
      assertEquals(floatIsNull[i], tsBlock.getColumn(3).isNull(i));
      if (!floatIsNull[i]) {
        assertEquals(floatValueArray[i], tsBlock.getColumn(3).getFloat(i), DELTA);
      }
      assertEquals(doubleIsNull[i], tsBlock.getColumn(4).isNull(i));
      if (!doubleIsNull[i]) {
        assertEquals(doubleValueArray[i], tsBlock.getColumn(4).getDouble(i), DELTA);
      }
      assertEquals(binaryIsNull[i], tsBlock.getColumn(5).isNull(i));
      if (!binaryIsNull[i]) {
        assertEquals(binaryValueArray[i], tsBlock.getColumn(5).getBinary(i));
      }
    }
  }

  @Test
  public void testSubTsBlock() {
    TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32));
    for (int i = 0; i < 10; i++) {
      builder.getTimeColumnBuilder().writeLong(i);
      builder.getColumnBuilder(0).writeInt(i);
      builder.declarePosition();
    }
    TsBlock tsBlock = builder.build();
    TsBlockSingleColumnIterator iterator = tsBlock.getTsBlockSingleColumnIterator();
    int index = 0;
    while (iterator.hasNext()) {
      Assert.assertEquals(index, iterator.currentTime());
      Assert.assertEquals(index, iterator.currentValue());
      iterator.next();
      index++;
    }
    // get subTsBlock from TsBlock, offset = 3
    int offset = 3;
    TsBlock subTsBlock = tsBlock.subTsBlock(offset);
    iterator = subTsBlock.getTsBlockSingleColumnIterator();
    index = offset;
    while (iterator.hasNext()) {
      Assert.assertEquals(index, iterator.currentTime());
      Assert.assertEquals(index, iterator.currentValue());
      iterator.next();
      index++;
    }
    // get subSubTsBlock from subTsBlock, offset = 2
    int nextOffset = 2;
    TsBlock subSubTsBlock = subTsBlock.subTsBlock(nextOffset);
    iterator = subSubTsBlock.getTsBlockSingleColumnIterator();
    index = offset + nextOffset;
    while (iterator.hasNext()) {
      Assert.assertEquals(index, iterator.currentTime());
      Assert.assertEquals(index, iterator.currentValue());
      iterator.next();
      index++;
    }
    try {
      subSubTsBlock.subTsBlock(3);
    } catch (IllegalArgumentException e) {
      Assert.assertTrue(
          e.getMessage().contains("FromIndex of subTsBlock cannot over positionCount."));
    }
  }

  private TsBlock getOriginalTsBlock() {
    TsBlockBuilder tsBlockBuilder =
        new TsBlockBuilder(Arrays.asList(TSDataType.INT64, TSDataType.DOUBLE));
    for (int i = 0; i < 19; i++) {
      tsBlockBuilder.getTimeColumnBuilder().writeLong(i);
      tsBlockBuilder.getColumnBuilder(0).writeLong(i);
      tsBlockBuilder.getColumnBuilder(1).writeDouble(i * 0.1);
      tsBlockBuilder.declarePosition();
    }
    return tsBlockBuilder.build();
  }

  @Test
  public void ColumnInsertTest() {
    TsBlock tsBlock = getOriginalTsBlock();
    int[] targetValue = new int[19];
    for (int i = 0; i < 19; i++) {
      targetValue[i] = i;
    }

    Column targetColumn = new IntColumn(19, Optional.empty(), targetValue);
    tsBlock = tsBlock.insertValueColumn(1, new Column[] {targetColumn});
    for (int i = 0; i < 19; i++) {
      assertEquals(i, tsBlock.getColumn(0).getLong(i));
      assertEquals(i, tsBlock.getColumn(1).getInt(i));
      assertEquals(i * 0.1, tsBlock.getColumn(2).getDouble(i), DELTA);
    }
  }

  @Test
  public void MultiColumnInsertTest() {
    TsBlock tsBlock = getOriginalTsBlock();
    int[] targetValue1 = new int[19];
    double[] targetValue2 = new double[19];
    for (int i = 0; i < 19; i++) {
      targetValue1[i] = i;
      targetValue2[i] = i * 0.2;
    }
    Column[] targetColumn =
        new Column[] {
          new IntColumn(19, Optional.empty(), targetValue1),
          new DoubleColumn(19, Optional.empty(), targetValue2)
        };
    tsBlock = tsBlock.insertValueColumn(1, targetColumn);
    for (int i = 0; i < 19; i++) {
      assertEquals(i, tsBlock.getColumn(0).getLong(i));
      assertEquals(i, tsBlock.getColumn(1).getInt(i));
      assertEquals(i * 0.2, tsBlock.getColumn(2).getDouble(i), DELTA);
      assertEquals(i * 0.1, tsBlock.getColumn(3).getDouble(i), DELTA);
    }
  }

  @Test
  public void ColumnAppendTest() {
    TsBlock tsBlock = getOriginalTsBlock();
    int[] targetValue = new int[19];
    long[] targetValue2 = new long[19];
    for (int i = 0; i < 19; i++) {
      targetValue[i] = i * 2;
      targetValue2[i] = i * 3;
    }
    Column[] targetColumn =
        new Column[] {
          new IntColumn(19, Optional.empty(), targetValue),
          new LongColumn(19, Optional.empty(), targetValue2)
        };
    tsBlock = tsBlock.appendValueColumns(targetColumn);
    for (int i = 0; i < 19; i++) {
      assertEquals(i, tsBlock.getColumn(0).getLong(i));
      assertEquals(i * 0.1, tsBlock.getColumn(1).getDouble(i), DELTA);
      assertEquals(i * 2, tsBlock.getColumn(2).getInt(i));
      assertEquals(i * 3, tsBlock.getColumn(3).getLong(i));
    }
  }
}