org.apache.kudu.client.Operation Java Examples

private Operation createKuduOperation(OperationType operationType, Record record,
                                      List<String> fieldNames, Boolean ignoreNull,
                                      Boolean lowercaseFields, KuduTable kuduTable) {
    switch (operationType) {
        case DELETE:
            return deleteRecordFromKudu(kuduTable, record, fieldNames, ignoreNull, lowercaseFields);
        case INSERT:
        case INSERT_IGNORE:
            return insertRecordToKudu(kuduTable, record, fieldNames, ignoreNull, lowercaseFields);
        case UPSERT:
            return upsertRecordToKudu(kuduTable, record, fieldNames, ignoreNull, lowercaseFields);
        case UPDATE:
            return updateRecordToKudu(kuduTable, record, fieldNames, ignoreNull, lowercaseFields);
        default:
            throw new IllegalArgumentException(String.format("OperationType: %s not supported by Kudu", operationType));
    }
}
 

/**
 * Return Operation based on the operation code. If the code has a number
 * that Kudu destination doesn't support, it throws UnsupportedOperationException.
 * @param table
 * @param op
 * @return
 * @throws UnsupportedOperationException
 */
protected Operation getOperation(KuduTable table, int op) throws UnsupportedOperationException {
  Operation operation = null;
  switch (op) {
    case OperationType.INSERT_CODE:
      operation = table.newInsert();
      break;
    case OperationType.UPSERT_CODE:
      operation = table.newUpsert();
      break;
    case OperationType.UPDATE_CODE:
      operation = table.newUpdate();
      break;
    case OperationType.DELETE_CODE:
      operation = table.newDelete();
      break;
    default:
      LOG.error("Operation {} not supported", op);
      throw new UnsupportedOperationException(String.format("Unsupported Operation: %s", op));
  }
  return operation;
}
 

private static Supplier<Operation> getOperationCreater(String mode, KuduTable kuduTable)
{
    //INSERT OR UPSET OR UPDATE OR DELETE
    switch (mode.toUpperCase()) {
        case "INSERT":
            return () -> kuduTable.newInsert();
        case "UPSET":
            return () -> kuduTable.newUpsert();
        case "UPDATE":
            return () -> kuduTable.newUpdate();
        case "DELETE":
            return () -> kuduTable.newDelete();
        default:
            throw new IllegalArgumentException();
    }
}
 

@Override
public void process(Record record)
{
    Operation operation = operationCreater.get();
    try {
        for (int i = 0; i < fieldNames.size(); i++) {
            appendColumn(operation, fieldNames.get(i), record.getField(i));
        }

        kuduSession.apply(operation);
        // submit batch
        if (rowNumCnt++ > maxBatchSize) {
            this.flush();
        }
    }
    catch (IOException e) {
        throwsException(e);
    }
}
 

@Override
public void applyRandomMutations(List<Row> planned) throws Exception {
  KuduConnection connection = getConnection();
  KuduSession session = connection.getSession();
  KuduTable table = connection.getTable(config.getString(TABLE_CONFIG_NAME));

  List<Operation> operations = extractOperations(planned, table);

  for (Operation operation : operations) {
    session.apply(operation);
  }

  // Wait until all operations have completed before checking for errors.
  while (session.hasPendingOperations()) {
    Thread.sleep(1);
  }

  // Fail fast on any error applying mutations
  if (session.countPendingErrors() > 0) {
    RowError firstError = session.getPendingErrors().getRowErrors()[0];
    String errorMessage = String.format("Kudu output error '%s' during operation '%s' at tablet server '%s'",
        firstError.getErrorStatus(), firstError.getOperation(), firstError.getTsUUID());

    throw new RuntimeException(errorMessage);
  }
}
 

public Optional<Operation> createBaseOperation(T input, KuduTable table) {
    if (operation == null) {
        throw new UnsupportedOperationException("createBaseOperation must be overridden if no operation specified in constructor");
    }
    switch (operation) {
        case INSERT:
            return Optional.of(table.newInsert());
        case UPDATE:
            return Optional.of(table.newUpdate());
        case UPSERT:
            return Optional.of(table.newUpsert());
        case DELETE:
            return Optional.of(table.newDelete());
        default:
            throw new RuntimeException("Unknown operation " + operation);
    }
}
 

@Override
public List<Operation> createOperations(T input, KuduTable table) {
    Optional<Operation> operationOpt = createBaseOperation(input, table);
    if (!operationOpt.isPresent()) {
        return Collections.emptyList();
    }

    Operation operation = operationOpt.get();
    PartialRow partialRow = operation.getRow();

    for (int i = 0; i < columnNames.length; i++) {
        partialRow.addObject(columnNames[i], getField(input, i));
    }

    return Collections.singletonList(operation);
}
 

@Test
void testPojoMapper() {

    PojoOperationMapper<BookInfo> mapper = new PojoOperationMapper<>(BookInfo.class, KuduTestBase.columns, AbstractSingleOperationMapper.KuduOperation.INSERT);

    BookInfo bookInfo = KuduTestBase.booksDataPojo().get(0);

    assertEquals(bookInfo.id, mapper.getField(bookInfo, 0));
    assertEquals(bookInfo.title, mapper.getField(bookInfo, 1));
    assertEquals(bookInfo.author, mapper.getField(bookInfo, 2));
    assertEquals(bookInfo.price, mapper.getField(bookInfo, 3));
    assertEquals(bookInfo.quantity, mapper.getField(bookInfo, 4));

    List<Operation> operations = mapper.createOperations(bookInfo, mockTable);
    assertEquals(1, operations.size());

    PartialRow row = operations.get(0).getRow();
    Mockito.verify(row, Mockito.times(1)).addObject("id", bookInfo.id);
    Mockito.verify(row, Mockito.times(1)).addObject("quantity", bookInfo.quantity);

    Mockito.verify(row, Mockito.times(1)).addObject("title", bookInfo.title);
    Mockito.verify(row, Mockito.times(1)).addObject("author", bookInfo.author);

    Mockito.verify(row, Mockito.times(1)).addObject("price", bookInfo.price);
}
 

static Operation toOperation(KuduTable table, KuduConnector.WriteMode writeMode) {
    switch (writeMode) {
        case INSERT: return table.newInsert();
        case UPDATE: return table.newUpdate();
        case UPSERT: return table.newUpsert();
    }
    return table.newUpsert();
}
 

@Before
public void setup() throws Exception {
  // Sample table and schema
  List<ColumnSchema> columns = new ArrayList(2);
  columns.add(new ColumnSchema.ColumnSchemaBuilder("key", Type.INT32).key(true).build());
  columns.add(new ColumnSchema.ColumnSchemaBuilder("value", Type.STRING).build());
  columns.add(new ColumnSchema.ColumnSchemaBuilder("name", Type.STRING).build());
  final Schema schema = new Schema(columns);

  // Mock KuduTable class
  KuduTable table = PowerMockito.mock(KuduTable.class);
  PowerMockito.doReturn(schema).when(table).getSchema();

  // Mock KuduSession class
  PowerMockito.suppress(PowerMockito.method(
      AsyncKuduSession.class,
      "apply",
      Operation.class
  ));
  PowerMockito.suppress(PowerMockito.method(
      AsyncKuduSession.class,
      "flush"
  ));
  PowerMockito.suppress(PowerMockito.method(
      KuduLookupProcessor.class,
      "destroy"
  ));
}
 

@Override
public Operation apply(TableAndRecord<Long> input) {
  Upsert upsert = input.getTable().newUpsert();
  PartialRow row = upsert.getRow();
  row.addLong(COL_ID, input.getRecord());
  row.addString(COL_NAME, input.getRecord() + ": name");
  return upsert;
}
 

@Test
void testCorrectOperationUpsert() {
    RowOperationMapper mapper = new RowOperationMapper(KuduTestBase.columns, AbstractSingleOperationMapper.KuduOperation.UPSERT);
    Row inputRow = KuduTestBase.booksDataRow().get(0);

    List<Operation> operations = mapper.createOperations(inputRow, mockTable);

    assertEquals(1, operations.size());
    verify(mockTable).newUpsert();
}
 

@Test
void testCorrectOperationInsert() {
    RowOperationMapper mapper = new RowOperationMapper(KuduTestBase.columns, AbstractSingleOperationMapper.KuduOperation.INSERT);
    Row inputRow = KuduTestBase.booksDataRow().get(0);

    List<Operation> operations = mapper.createOperations(inputRow, mockTable);

    assertEquals(1, operations.size());
    verify(mockTable).newInsert();
}
 

@Test
void testCorrectOperationUpsert() {
    TupleOperationMapper<Tuple5<Integer, String, String, Double, Integer>> mapper =
            new TupleOperationMapper<>(KuduTestBase.columns, AbstractSingleOperationMapper.KuduOperation.UPSERT);
    Tuple5<Integer, String, String, Double, Integer> inputTuple = KuduTestBase.booksDataTuple().get(0);

    List<Operation> operations = mapper.createOperations(inputTuple, mockTable);

    assertEquals(1, operations.size());
    verify(mockTable).newUpsert();
}
 

@Test
void testCorrectOperationInsert() {
    TupleOperationMapper<Tuple5<Integer, String, String, Double, Integer>> mapper =
            new TupleOperationMapper<>(KuduTestBase.columns, AbstractSingleOperationMapper.KuduOperation.INSERT);
    Tuple5<Integer, String, String, Double, Integer> inputTuple = KuduTestBase.booksDataTuple().get(0);

    List<Operation> operations = mapper.createOperations(inputTuple, mockTable);

    assertEquals(1, operations.size());
    verify(mockTable).newInsert();
}
 

public void write(T input) throws IOException {
    checkAsyncErrors();

    for (Operation operation : operationMapper.createOperations(input, table)) {
        checkErrors(session.apply(operation));
    }
}
 

static Operation toOperation(KuduTable table, KuduConnector.WriteMode writeMode) {
    switch (writeMode) {
        case INSERT: return table.newInsert();
        case UPDATE: return table.newUpdate();
        case UPSERT: return table.newUpsert();
    }
    return table.newUpsert();
}
 

@Override
public Optional<Operation> createBaseOperation(Tuple2<Boolean, Row> input, KuduTable table) {
    return Optional.of(input.f0 ? table.newUpsert() : table.newDelete());
}
 

private void insert(KuduTable kuduTable, Operation operation, Record record, List<String> fieldNames){
    PartialRow row = operation.getRow();
    Schema colSchema = kuduTable.getSchema();

    for (String colName : fieldNames) {
        int colIdx = this.getColumnIndex(colSchema, colName);
        if (colIdx != -1) {
            Type colType = colSchema.getColumnByIndex(colIdx).getType();

            switch (colType.getDataType()) {
                case BOOL:
                    row.addBoolean(colIdx, record.getAsBoolean(colName));
                    break;
                case FLOAT:
                    row.addFloat(colIdx, record.getAsFloat(colName));
                    break;
                case DOUBLE:
                    row.addDouble(colIdx, record.getAsDouble(colName));
                    break;
                case BINARY:
                    row.addBinary(colIdx, record.getAsString(colName).getBytes());
                    break;
                case INT8:
                case INT16:
                    short temp = (short)record.getAsInt(colName).intValue();
                    row.addShort(colIdx, temp);
                case INT32:
                    row.addInt(colIdx, record.getAsInt(colName));
                    break;
                case INT64:
                    row.addLong(colIdx, record.getAsLong(colName));
                    break;
                case STRING:
                    row.addString(colIdx, record.getAsString(colName));
                    break;
                default:
                    throw new IllegalStateException(String.format("unknown column type %s", colType));
            }
        }
    }
}
 

private void appendColumn(Operation operation, String name, Object value)
{
    ColumnSchema columnSchema = kuduTable.getSchema().getColumn(name);

    if (value == null) {
        operation.getRow().setNull(name);
        return;
    }

    Type kuduType = columnSchema.getType();
    switch (kuduType) {
        case BINARY:
            operation.getRow().addBinary(name, (byte[]) value);
            break;

        case STRING:
            operation.getRow().addString(name, String.valueOf(value));
            break;
        case BOOL:
            operation.getRow().addBoolean(name, (Boolean) value);
            break;

        case INT8:
        case INT16:
            operation.getRow().addShort(name, (Short) value);
            break;

        case INT32:
            operation.getRow().addInt(name, (Integer) value);
            break;

        case INT64: {
            if (value instanceof Date) {
                operation.getRow().addLong(name, ((Date) value).getTime());
            }
            else if (value instanceof Time) {
                operation.getRow().addLong(name, ((Time) value).getTime());
            }
            else if (value instanceof Timestamp) {
                operation.getRow().addLong(name, ((Timestamp) value).getTime());
            }
            else {
                operation.getRow().addLong(name, (Long) value);
            }
            break;
        }
        case DOUBLE:
            operation.getRow().addDouble(name, (Double) value);
            break;
        case FLOAT:
            operation.getRow().addFloat(name, (Float) value);
            break;

        case DECIMAL:
            operation.getRow().addDecimal(name, (BigDecimal) value);
            break;

        default:
            throw new IllegalStateException("不受支持的kudu类型:" + kuduType);
    }
}
 

static Operation toOperation(KuduTable table, KuduConnector.WriteMode writeMode, KuduRow row) {
    final Operation operation = toOperation(table, writeMode);
    final PartialRow partialRow = operation.getRow();

    table.getSchema().getColumns().forEach(column -> {
        String columnName = column.getName();
        Object value = row.getField(column.getName());

        if (value == null) {
            partialRow.setNull(columnName);
        } else {
            Type type = column.getType();
            switch (type) {
                case STRING:
                    partialRow.addString(columnName, (String) value);
                    break;
                case FLOAT:
                    partialRow.addFloat(columnName, (Float) value);
                    break;
                case INT8:
                    partialRow.addByte(columnName, (Byte) value);
                    break;
                case INT16:
                    partialRow.addShort(columnName, (Short) value);
                    break;
                case INT32:
                    partialRow.addInt(columnName, (Integer) value);
                    break;
                case INT64:
                    partialRow.addLong(columnName, (Long) value);
                    break;
                case DOUBLE:
                    partialRow.addDouble(columnName, (Double) value);
                    break;
                case BOOL:
                    partialRow.addBoolean(columnName, (Boolean) value);
                    break;
                case UNIXTIME_MICROS:
                    //*1000 to correctly create date on kudu
                    partialRow.addLong(columnName, ((Long) value) * 1000);
                    break;
                case BINARY:
                    partialRow.addBinary(columnName, (byte[]) value);
                    break;
                default:
                    throw new IllegalArgumentException("Illegal var type: " + type);
            }
        }
    });
    return operation;
}
 

static Operation toOperation(KuduTable table, KuduConnector.WriteMode writeMode, KuduRow row) {
    final Operation operation = toOperation(table, writeMode);
    final PartialRow partialRow = operation.getRow();

    table.getSchema().getColumns().forEach(column -> {
        String columnName = column.getName();
        Object value = row.getField(column.getName());

        if (value == null) {
            partialRow.setNull(columnName);
        } else {
            Type type = column.getType();
            switch (type) {
                case STRING:
                    partialRow.addString(columnName, (String) value);
                    break;
                case FLOAT:
                    partialRow.addFloat(columnName, (Float) value);
                    break;
                case INT8:
                    partialRow.addByte(columnName, (Byte) value);
                    break;
                case INT16:
                    partialRow.addShort(columnName, (Short) value);
                    break;
                case INT32:
                    partialRow.addInt(columnName, (Integer) value);
                    break;
                case INT64:
                    partialRow.addLong(columnName, (Long) value);
                    break;
                case DOUBLE:
                    partialRow.addDouble(columnName, (Double) value);
                    break;
                case BOOL:
                    partialRow.addBoolean(columnName, (Boolean) value);
                    break;
                case UNIXTIME_MICROS:
                    //*1000 to correctly create date on kudu
                    partialRow.addLong(columnName, ((Long) value) * 1000);
                    break;
                case BINARY:
                    partialRow.addBinary(columnName, (byte[]) value);
                    break;
                default:
                    throw new IllegalArgumentException("Illegal var type: " + type);
            }
        }
    });
    return operation;
}
 

/**
 * Create a list of operations to be executed by the {@link KuduWriter} for the
 * current input
 *
 * @param input input element
 * @param table table for which the operations should be created
 * @return List of operations to be executed on the table
 */
List<Operation> createOperations(T input, KuduTable table);