ru.yandex.clickhouse.util.ClickHouseRowBinaryStream Java Examples

static void writeRow(ClickHouseRowBinaryStream stream, TableSchema schema, Row row)
    throws IOException {
  for (TableSchema.Column column : schema.columns()) {
    if (!column.materializedOrAlias()) {
      Object value = row.getValue(column.name());

      if (column.columnType().nullable()) {
        writeNullableValue(stream, column.columnType(), value);
      } else {
        if (value == null) {
          value = column.defaultValue();
        }

        writeValue(stream, column.columnType(), value);
      }
    }
  }
}
 

@Override
protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {

    try (Connection connection = manager.get(req.getParameter("connection_string")); Statement sth = connection.createStatement()) {
        String schema = req.getParameter("schema");
        String table = req.getParameter("table");

        String quote = connection.getMetaData().getIdentifierQuoteString();
        String schemaAndTable = Stream.of(schema, table)
                .filter(s -> !StringUtil.isBlank(s))
                .map(s -> quote + s + quote)
                .collect(Collectors.joining("."));

        String queryRewrite = "SELECT * FROM " + schemaAndTable + " WHERE 1 = 0";
        log.info("Inferring schema by query {}", queryRewrite);

        ResultSet resultset = sth.executeQuery(queryRewrite);
        String ddl = converter.getColumnsDDL(resultset.getMetaData());
        resp.setContentType("application/octet-stream");
        ClickHouseRowBinaryStream stream = new ClickHouseRowBinaryStream(resp.getOutputStream(), null, new ClickHouseProperties());
        stream.writeString(ddl);
    } catch (Exception err) {
        log.error(err.getMessage(), err);
        resp.sendError(HttpStatus.INTERNAL_SERVER_ERROR_500, err.getMessage());
    }
}
 

@Test
public void testRowBinary() throws Exception {
    statement.write().send("INSERT INTO test.writer", new ClickHouseStreamCallback() {
        @Override
        public void writeTo(ClickHouseRowBinaryStream stream) throws IOException {
            for (int i = 0; i < 10; i++) {
                stream.writeInt32(i);
                stream.writeString("Имя " + i);
            }
        }
    }, ClickHouseFormat.RowBinary);

    assertTableRowCount(10);
    ResultSet rs = statement.executeQuery("SELECT count() FROM test.writer WHERE name = concat('Имя ', toString(id))");
    rs.next();
    assertEquals(rs.getInt(1), 10);
}
 

/**
 * (metric, value, timestamp, date, updated).
 *
 * @param metric
 * @param stream
 * @throws IOException
 */
private void writeMetric(Metric metric, ClickHouseRowBinaryStream stream) throws IOException {
    stream.writeUnsignedLeb128(metric.getMetricDescription().getNameLengthInBytes());
    metric.getMetricDescription().writeName(stream);
    stream.writeFloat64(metric.getValue());
    stream.writeUInt32(metric.getTimestampSeconds());
    stream.writeUInt16(getUnsignedDaysSinceEpoch(metric.getTimestampSeconds()));
    stream.writeUInt32((int) Integer.toUnsignedLong(metric.getUpdatedSeconds()));
}
 

@SuppressWarnings("unchecked")
static void writeNullableValue(
    ClickHouseRowBinaryStream stream, ColumnType columnType, Object value) throws IOException {

  if (value == null) {
    stream.markNextNullable(true);
  } else {
    stream.markNextNullable(false);
    writeValue(stream, columnType, value);
  }
}
 

@Test
public void testNative() throws Exception {
    statement.write().send("INSERT INTO test.writer", new ClickHouseStreamCallback() {
        @Override
        public void writeTo(ClickHouseRowBinaryStream stream) throws IOException {

            int numberOfRows = 1000;
            stream.writeUnsignedLeb128(2); // 2 columns
            stream.writeUnsignedLeb128(numberOfRows);

            stream.writeString("id");
            stream.writeString("Int32");

            for (int i = 0; i < numberOfRows; i++) {
                stream.writeInt32(i);
            }

            stream.writeString("name");
            stream.writeString("String");

            for (int i = 0; i < numberOfRows; i++) {
                stream.writeString("Имя " + i);
            }
        }
    }, ClickHouseFormat.Native);

    assertTableRowCount(1000);
    ResultSet rs = statement.executeQuery("SELECT count() FROM test.writer WHERE name = concat('Имя ', toString(id))");
    rs.next();
    assertEquals(rs.getInt(1), 1000);
}
 

@Test
public void testLowCardinality() throws Exception{
    final ClickHouseStatement statement = connection.createStatement();
    connection.createStatement().execute("DROP TABLE IF EXISTS test.low_cardinality");
    connection.createStatement().execute(
        "CREATE TABLE test.low_cardinality (date Date, lowCardinality LowCardinality(String), string String) ENGINE = MergeTree(date, (date), 8192)"
    );

    final Date date1 = new Date(1497474018000L);

    statement.sendNativeStream(
        "INSERT INTO test.low_cardinality (date, lowCardinality, string)",
        new ClickHouseStreamCallback() {
            @Override
            public void writeTo(ClickHouseRowBinaryStream stream) throws IOException {
                stream.writeUnsignedLeb128(3); // Columns number
                stream.writeUnsignedLeb128(1); // Rows number

                stream.writeString("date"); // Column name
                stream.writeString("Date");  // Column type
                stream.writeDate(date1);  // value

                stream.writeString("lowCardinality"); // Column name
                stream.writeString("String");  // Column type
                stream.writeString("string");  // value

                stream.writeString("string"); // Column name
                stream.writeString("String");  // Column type
                stream.writeString("string");  // value
            }
        }
    );

    ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM test.low_cardinality");

    Assert.assertTrue(rs.next());
    assertEquals(rs.getString("lowCardinality"), "string");
    assertEquals(rs.getString("string"), "string");
}
 

@Test
public void testTimeZone() throws Exception{
    final ClickHouseStatement statement = connection.createStatement();
    connection.createStatement().execute("DROP TABLE IF EXISTS test.binary_tz");
    connection.createStatement().execute(
        "CREATE TABLE test.binary_tz (date Date, dateTime DateTime) ENGINE = MergeTree(date, (date), 8192)"
    );

    final Date date1 = new Date(1497474018000L);

    statement.sendRowBinaryStream(
        "INSERT INTO test.binary_tz (date, dateTime)",
        new ClickHouseStreamCallback() {
            @Override
            public void writeTo(ClickHouseRowBinaryStream stream) throws IOException {
                stream.writeDate(date1);
                stream.writeDateTime(date1);
            }
        }
    );

    ResultSet rs = connection.createStatement().executeQuery(
        "SELECT date, dateTime from test.binary_tz"
    );

    Assert.assertTrue(rs.next());
    assertEquals(rs.getTime("dateTime"), new Time(date1.getTime()));
    Date expectedDate = withTimeAtStartOfDay(date1); // expected start of the day in local timezone
    assertEquals(rs.getDate("date"), expectedDate);
}
 

@Test
public void multiRowTest() throws SQLException {
    connection.createStatement().execute("DROP TABLE IF EXISTS test.big_data");
    connection.createStatement().execute(
            "CREATE TABLE test.big_data (" +
                    "value Int32" +
                    ") ENGINE = TinyLog()"
    );

    final int count = 1000000;
    final AtomicLong sum = new AtomicLong();

    connection.createStatement().sendRowBinaryStream(
            "INSERT INTO test.big_data (value)",
            new ClickHouseStreamCallback() {
                @Override
                public void writeTo(ClickHouseRowBinaryStream stream) throws IOException {
                    for (int i = 0; i < count; i++) {
                        stream.writeInt32(i);
                        sum.addAndGet(i);
                    }
                }
            }
    );

    ResultSet rs = connection.createStatement().executeQuery("SELECT count() AS cnt, sum(value) AS sum FROM test.big_data");

    Assert.assertTrue(rs.next());
    assertEquals(rs.getInt("cnt"), count);
    assertEquals(rs.getLong("sum"), sum.get());
}
 

@Test
public void writeMetric() throws Exception {

    MetricDir root = new InMemoryMetricDir(null, null, MetricStatus.SIMPLE);
    MetricDir a = new InMemoryMetricDir(root, "a", MetricStatus.SIMPLE);
    MetricDir b = new InMemoryMetricDir(a, "b", MetricStatus.SIMPLE);
    MetricName c = new MetricName(b, "c", MetricStatus.SIMPLE, new DefaultRetentionProvider());

    Metric metric = new Metric(
        c,
        1492342562,
        42.21,
        1492350000
    );

    MetricsStreamCallback entity = new MetricsStreamCallback(Collections.singletonList(metric), TimeZone.getTimeZone("Europe/Moscow"));

    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
    entity.writeTo(new ClickHouseRowBinaryStream(byteArrayOutputStream, null, new ClickHouseProperties()));

    byte[] actual = byteArrayOutputStream.toByteArray();

    //Result of
    //clickhouse-client -q "select 'a.b.c', toFloat64(42.21), toUInt32(1492342562), toDate('2017-04-16'), toUInt32(1492350000) format RowBinary" | od -vAn -td1
    byte[] expected = {
        5, 97, 46, 98, 46, 99, 123, 20, -82, 71, -31, 26, 69, 64, 34, 87, -13, 88, 120, 67, 48, 116, -13, 88
    };
    Assert.assertArrayEquals(expected, actual);
}
 

@Override
protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {

    try {
        String query = req.getParameter("query");
        if (StringUtil.isBlank(query)) {
            // a hack for wrong input from CH
            String requestBody = StreamUtils.toString(req.getInputStream());
            String[] parts = requestBody.split("query=", 2);
            if (parts.length == 2) {
                query = parts[1];
            }
        }

        if (StringUtil.isBlank(query)) {
            throw new IllegalArgumentException("Query is blank or empty");
        }

        try (Connection connection = manager.get(req.getParameter("connection_string")); Statement sth = connection.createStatement()) {
            ResultSet resultset = sth.executeQuery(query);
            ResultSetMetaData meta = resultset.getMetaData();

            ClickHouseRowSerializer serializer = ClickHouseRowSerializer.create(meta);
            ClickHouseRowBinaryStream stream = new ClickHouseRowBinaryStream(resp.getOutputStream(), null, new ClickHouseProperties());

            resp.setContentType("application/octet-stream");
            while (resultset.next()) {
                serializer.serialize(resultset, stream);
            }
        }
    } catch (Exception err) {
        log.error(err.getMessage(), err);
        resp.sendError(HttpStatus.INTERNAL_SERVER_ERROR_500, err.getMessage());
    }
}
 

@Override
public void writeName(ClickHouseRowBinaryStream stream) throws IOException {
    if (!parent.isRoot()) {
        parent.writeName(stream);
    }
    stream.writeBytes(name.getBytes());
    if (isDir()) {
        stream.writeByte(LEVEL_SPLITTER);
    }
}
 

public void serialize(ResultSet resultSet, int position, ClickHouseRowBinaryStream stream) throws SQLException, IOException {
    T value = pair.getExtractor().apply(resultSet, position);

    if (canBeNull) {
        final boolean isNull = resultSet.wasNull() || null == value;
        stream.writeByte((byte) (isNull ? 1 : 0));
        if (isNull) {
            return;
        }
    }
    pair.getSerializer().accept(value, stream);
}
 

@Override
protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
    try (Connection connection = manager.get(req.getParameter("connection_string"))) {
        ClickHouseRowBinaryStream stream = new ClickHouseRowBinaryStream(resp.getOutputStream(), null, new ClickHouseProperties());
        final String identifierQuoteString = connection.getMetaData().getIdentifierQuoteString();
        resp.setContentType("application/octet-stream");
        stream.writeString(identifierQuoteString);
    } catch (Exception err) {
        log.error(err.getMessage(), err);
        resp.sendError(HttpStatus.INTERNAL_SERVER_ERROR_500, err.getMessage());
    }
}
 

@Override
public void writeTo(ClickHouseRowBinaryStream stream) throws IOException {
    for (Metric metric : metrics) {
        writeMetric(metric, stream);
    }
}
 

@Test
public void testSerialize() throws Exception {
    try (ResultSet resultset = conn.createStatement()
            .executeQuery("SELECT * FROM " + quote + tableName + quote);
         ResultSet resultsetCopy = conn.createStatement()
                 .executeQuery("SELECT * FROM " + quote + tableName + quote)) {

        File tmp = File.createTempFile("pref", "suf");
        tmp.deleteOnExit();
        OutputStream outputStream = new FileOutputStream(tmp);

        ClickHouseRowSerializer ser = ClickHouseRowSerializer.create(resultset.getMetaData());
        ClickHouseRowBinaryStream stream = new ClickHouseRowBinaryStream(outputStream, null, new ClickHouseProperties());

        while (resultset.next()) {
            ser.serialize(resultset, stream);
        }

        outputStream.flush();
        outputStream.close();

        // If a path to binary clickhouse-local has been specified, try to validate file content via invocation
        String clickHouseLocalBinaryPath = System.getProperty("clickhouse.local.bin", null);
        assumeTrue(clickHouseLocalBinaryPath != null);

        ProcessBuilder builder = new ProcessBuilder()
                .command(clickHouseLocalBinaryPath,
                        "--structure",
                        ClickHouseConverter.getCLIStructure(resultsetCopy.getMetaData()),
                        "--input-format=RowBinary",
                        "--format=Vertical",
                        "--query",
                        "SELECT * FROM table",
                        "--file=" + tmp + ""
                );

        builder.inheritIO();
        Process process = builder.start();
        if (!process.waitFor(10, TimeUnit.SECONDS)) {
            process.destroy();
            fail("Failed to validate output via invocation of clickhouse-local");
        } else if (process.exitValue() != 0) {
            fail("clickhouse-local exited with non-zero code");
        }
    }
}
 

public void serialize(ResultSet resultSet, ClickHouseRowBinaryStream stream) throws IOException, SQLException {
    int i = 0;
    for (ClickHouseFieldSerializer serializer : serializers) {
        serializer.serialize(resultSet, ++i, stream);
    }
}
 

@SuppressWarnings("unchecked")
static void writeValue(ClickHouseRowBinaryStream stream, ColumnType columnType, Object value)
    throws IOException {

  switch (columnType.typeName()) {
    case FIXEDSTRING:
      byte[] bytes;

      if (value instanceof String) {
        bytes = ((String) value).getBytes(Charsets.UTF_8);
      } else {
        bytes = ((byte[]) value);
      }

      stream.writeBytes(bytes);
      break;

    case FLOAT32:
      stream.writeFloat32((Float) value);
      break;

    case FLOAT64:
      stream.writeFloat64((Double) value);
      break;

    case INT8:
      stream.writeInt8((Byte) value);
      break;

    case INT16:
      stream.writeInt16((Short) value);
      break;

    case INT32:
      stream.writeInt32((Integer) value);
      break;

    case INT64:
      stream.writeInt64((Long) value);
      break;

    case STRING:
      stream.writeString((String) value);
      break;

    case UINT8:
      stream.writeUInt8((Short) value);
      break;

    case UINT16:
      stream.writeUInt16((Integer) value);
      break;

    case UINT32:
      stream.writeUInt32((Long) value);
      break;

    case UINT64:
      stream.writeUInt64((Long) value);
      break;

    case ENUM8:
      Integer enum8 = columnType.enumValues().get((String) value);
      Preconditions.checkNotNull(
          enum8,
          "unknown enum value '" + value + "', possible values: " + columnType.enumValues());
      stream.writeInt8(enum8);
      break;

    case ENUM16:
      Integer enum16 = columnType.enumValues().get((String) value);
      Preconditions.checkNotNull(
          enum16,
          "unknown enum value '" + value + "', possible values: " + columnType.enumValues());
      stream.writeInt16(enum16);
      break;

    case DATE:
      Days epochDays = Days.daysBetween(EPOCH_INSTANT, (ReadableInstant) value);
      stream.writeUInt16(epochDays.getDays());
      break;

    case DATETIME:
      long epochSeconds = ((ReadableInstant) value).getMillis() / 1000L;
      stream.writeUInt32(epochSeconds);
      break;

    case ARRAY:
      List<Object> values = (List<Object>) value;
      stream.writeUnsignedLeb128(values.size());
      for (Object arrayValue : values) {
        writeValue(stream, columnType.arrayElementType(), arrayValue);
      }
      break;
  }
}
 

/**
 * Serialize non-null value into stream
 */
void accept(T t, ClickHouseRowBinaryStream stream) throws IOException;
 

void writeName(ClickHouseRowBinaryStream stream) throws IOException;