[ceph.git] / ceph / src / rocksdb / java / src / test / java / org / rocksdb / SstFileWriterTest.java

// Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).

package org.rocksdb;

import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.Assert.fail;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.rocksdb.util.BytewiseComparator;

public class SstFileWriterTest {
  private static final String SST_FILE_NAME = "test.sst";
  private static final String DB_DIRECTORY_NAME = "test_db";

  @ClassRule
  public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE
      = new RocksNativeLibraryResource();

  @Rule public TemporaryFolder parentFolder = new TemporaryFolder();

  enum OpType { PUT, PUT_BYTES, PUT_DIRECT, MERGE, MERGE_BYTES, DELETE, DELETE_BYTES }

  static class KeyValueWithOp {
    KeyValueWithOp(String key, String value, OpType opType) {
      this.key = key;
      this.value = value;
      this.opType = opType;
    }

    String getKey() {
      return key;
    }

    String getValue() {
      return value;
    }

    OpType getOpType() {
      return opType;
    }

    private final String key;
    private final String value;
    private final OpType opType;
  };

  private File newSstFile(final List<KeyValueWithOp> keyValues,
      boolean useJavaBytewiseComparator) throws IOException, RocksDBException {
    final EnvOptions envOptions = new EnvOptions();
    final StringAppendOperator stringAppendOperator = new StringAppendOperator();
    final Options options = new Options().setMergeOperator(stringAppendOperator);
    SstFileWriter sstFileWriter = null;
    ComparatorOptions comparatorOptions = null;
    BytewiseComparator comparator = null;
    if (useJavaBytewiseComparator) {
      comparatorOptions = new ComparatorOptions().setUseDirectBuffer(false);
      comparator = new BytewiseComparator(comparatorOptions);
      options.setComparator(comparator);
      sstFileWriter = new SstFileWriter(envOptions, options);
    } else {
      sstFileWriter = new SstFileWriter(envOptions, options);
    }

    final File sstFile = parentFolder.newFile(SST_FILE_NAME);
    try {
      sstFileWriter.open(sstFile.getAbsolutePath());
      assertThat(sstFileWriter.fileSize()).isEqualTo(0);
      for (KeyValueWithOp keyValue : keyValues) {
        Slice keySlice = new Slice(keyValue.getKey());
        Slice valueSlice = new Slice(keyValue.getValue());
        byte[] keyBytes = keyValue.getKey().getBytes();
        byte[] valueBytes = keyValue.getValue().getBytes();
        ByteBuffer keyDirect = ByteBuffer.allocateDirect(keyBytes.length);
        keyDirect.put(keyBytes);
        keyDirect.flip();
        ByteBuffer valueDirect = ByteBuffer.allocateDirect(valueBytes.length);
        valueDirect.put(valueBytes);
        valueDirect.flip();
        switch (keyValue.getOpType()) {
          case PUT:
            sstFileWriter.put(keySlice, valueSlice);
            break;
          case PUT_BYTES:
            sstFileWriter.put(keyBytes, valueBytes);
            break;
          case PUT_DIRECT:
            sstFileWriter.put(keyDirect, valueDirect);
            assertThat(keyDirect.position()).isEqualTo(keyBytes.length);
            assertThat(keyDirect.limit()).isEqualTo(keyBytes.length);
            assertThat(valueDirect.position()).isEqualTo(valueBytes.length);
            assertThat(valueDirect.limit()).isEqualTo(valueBytes.length);
            break;
          case MERGE:
            sstFileWriter.merge(keySlice, valueSlice);
            break;
          case MERGE_BYTES:
            sstFileWriter.merge(keyBytes, valueBytes);
            break;
          case DELETE:
            sstFileWriter.delete(keySlice);
            break;
          case DELETE_BYTES:
            sstFileWriter.delete(keyBytes);
            break;
          default:
            fail("Unsupported op type");
        }
        keySlice.close();
        valueSlice.close();
      }
      sstFileWriter.finish();
      assertThat(sstFileWriter.fileSize()).isGreaterThan(100);
    } finally {
      assertThat(sstFileWriter).isNotNull();
      sstFileWriter.close();
      options.close();
      envOptions.close();
      if (comparatorOptions != null) {
        comparatorOptions.close();
      }
      if (comparator != null) {
        comparator.close();
      }
    }
    return sstFile;
  }

  @Test
  public void generateSstFileWithJavaComparator()
      throws RocksDBException, IOException {
    final List<KeyValueWithOp> keyValues = new ArrayList<>();
    keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
    keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
    keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
    keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
    keyValues.add(new KeyValueWithOp("key5", "", OpType.DELETE));

    newSstFile(keyValues, true);
  }

  @Test
  public void generateSstFileWithNativeComparator()
      throws RocksDBException, IOException {
    final List<KeyValueWithOp> keyValues = new ArrayList<>();
    keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
    keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
    keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
    keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
    keyValues.add(new KeyValueWithOp("key5", "", OpType.DELETE));

    newSstFile(keyValues, false);
  }

  @Test
  public void ingestSstFile() throws RocksDBException, IOException {
    final List<KeyValueWithOp> keyValues = new ArrayList<>();
    keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
    keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT_DIRECT));
    keyValues.add(new KeyValueWithOp("key3", "value3", OpType.PUT_BYTES));
    keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
    keyValues.add(new KeyValueWithOp("key5", "value5", OpType.MERGE_BYTES));
    keyValues.add(new KeyValueWithOp("key6", "", OpType.DELETE));
    keyValues.add(new KeyValueWithOp("key7", "", OpType.DELETE));


    final File sstFile = newSstFile(keyValues, false);
    final File dbFolder = parentFolder.newFolder(DB_DIRECTORY_NAME);
    try(final StringAppendOperator stringAppendOperator =
            new StringAppendOperator();
        final Options options = new Options()
            .setCreateIfMissing(true)
            .setMergeOperator(stringAppendOperator);
        final RocksDB db = RocksDB.open(options, dbFolder.getAbsolutePath());
        final IngestExternalFileOptions ingestExternalFileOptions =
            new IngestExternalFileOptions()) {
      db.ingestExternalFile(Arrays.asList(sstFile.getAbsolutePath()),
          ingestExternalFileOptions);

      assertThat(db.get("key1".getBytes())).isEqualTo("value1".getBytes());
      assertThat(db.get("key2".getBytes())).isEqualTo("value2".getBytes());
      assertThat(db.get("key3".getBytes())).isEqualTo("value3".getBytes());
      assertThat(db.get("key4".getBytes())).isEqualTo("value4".getBytes());
      assertThat(db.get("key5".getBytes())).isEqualTo("value5".getBytes());
      assertThat(db.get("key6".getBytes())).isEqualTo(null);
      assertThat(db.get("key7".getBytes())).isEqualTo(null);
    }
  }

  @Test
  public void ingestSstFile_cf() throws RocksDBException, IOException {
    final List<KeyValueWithOp> keyValues = new ArrayList<>();
    keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
    keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
    keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
    keyValues.add(new KeyValueWithOp("key4", "", OpType.DELETE));

    final File sstFile = newSstFile(keyValues, false);
    final File dbFolder = parentFolder.newFolder(DB_DIRECTORY_NAME);
    try(final StringAppendOperator stringAppendOperator =
            new StringAppendOperator();
        final Options options = new Options()
            .setCreateIfMissing(true)
            .setCreateMissingColumnFamilies(true)
            .setMergeOperator(stringAppendOperator);
        final RocksDB db = RocksDB.open(options, dbFolder.getAbsolutePath());
        final IngestExternalFileOptions ingestExternalFileOptions =
            new IngestExternalFileOptions()) {

      try(final ColumnFamilyOptions cf_opts = new ColumnFamilyOptions()
              .setMergeOperator(stringAppendOperator);
          final ColumnFamilyHandle cf_handle = db.createColumnFamily(
              new ColumnFamilyDescriptor("new_cf".getBytes(), cf_opts))) {

        db.ingestExternalFile(cf_handle,
            Arrays.asList(sstFile.getAbsolutePath()),
            ingestExternalFileOptions);

        assertThat(db.get(cf_handle,
            "key1".getBytes())).isEqualTo("value1".getBytes());
        assertThat(db.get(cf_handle,
            "key2".getBytes())).isEqualTo("value2".getBytes());
        assertThat(db.get(cf_handle,
            "key3".getBytes())).isEqualTo("value3".getBytes());
        assertThat(db.get(cf_handle,
            "key4".getBytes())).isEqualTo(null);
      }
    }
  }
}
Commit	Line	Data
7c673cae	1	// Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
11fdf7f2 TL	2	// This source code is licensed under both the GPLv2 (found in the
	3	// COPYING file in the root directory) and Apache 2.0 License
	4	// (found in the LICENSE.Apache file in the root directory).
7c673cae FG	5
	6	package org.rocksdb;
	7
f67539c2 TL	8	import static org.assertj.core.api.Assertions.assertThat;
f67539c2 TL	9	import static org.junit.Assert.fail;
7c673cae FG	10
	11	import java.io.File;
	12	import java.io.IOException;
f67539c2	13	import java.nio.ByteBuffer;
11fdf7f2 TL	14	import java.util.ArrayList;
	15	import java.util.Arrays;
	16	import java.util.List;
f67539c2 TL	17	import org.junit.ClassRule;
	18	import org.junit.Rule;
	19	import org.junit.Test;
	20	import org.junit.rules.TemporaryFolder;
	21	import org.rocksdb.util.BytewiseComparator;
7c673cae FG	22
	23	public class SstFileWriterTest {
	24	private static final String SST_FILE_NAME = "test.sst";
	25	private static final String DB_DIRECTORY_NAME = "test_db";
	26
	27	@ClassRule
f67539c2 TL	28	public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE
f67539c2 TL	29	= new RocksNativeLibraryResource();
7c673cae FG	30
	31	@Rule public TemporaryFolder parentFolder = new TemporaryFolder();
	32
f67539c2	33	enum OpType { PUT, PUT_BYTES, PUT_DIRECT, MERGE, MERGE_BYTES, DELETE, DELETE_BYTES }
11fdf7f2	34
1e59de90	35	static class KeyValueWithOp {
11fdf7f2 TL	36	KeyValueWithOp(String key, String value, OpType opType) {
	37	this.key = key;
	38	this.value = value;
	39	this.opType = opType;
	40	}
	41
	42	String getKey() {
	43	return key;
	44	}
	45
	46	String getValue() {
	47	return value;
	48	}
	49
	50	OpType getOpType() {
	51	return opType;
	52	}
	53
1e59de90 TL	54	private final String key;
	55	private final String value;
	56	private final OpType opType;
11fdf7f2 TL	57	};
	58
	59	private File newSstFile(final List<KeyValueWithOp> keyValues,
	60	boolean useJavaBytewiseComparator) throws IOException, RocksDBException {
7c673cae	61	final EnvOptions envOptions = new EnvOptions();
11fdf7f2 TL	62	final StringAppendOperator stringAppendOperator = new StringAppendOperator();
11fdf7f2 TL	63	final Options options = new Options().setMergeOperator(stringAppendOperator);
7c673cae FG	64	SstFileWriter sstFileWriter = null;
	65	ComparatorOptions comparatorOptions = null;
	66	BytewiseComparator comparator = null;
	67	if (useJavaBytewiseComparator) {
f67539c2	68	comparatorOptions = new ComparatorOptions().setUseDirectBuffer(false);
7c673cae FG	69	comparator = new BytewiseComparator(comparatorOptions);
7c673cae FG	70	options.setComparator(comparator);
f67539c2	71	sstFileWriter = new SstFileWriter(envOptions, options);
7c673cae FG	72	} else {
	73	sstFileWriter = new SstFileWriter(envOptions, options);
	74	}
	75
	76	final File sstFile = parentFolder.newFile(SST_FILE_NAME);
	77	try {
	78	sstFileWriter.open(sstFile.getAbsolutePath());
f67539c2	79	assertThat(sstFileWriter.fileSize()).isEqualTo(0);
11fdf7f2	80	for (KeyValueWithOp keyValue : keyValues) {
7c673cae FG	81	Slice keySlice = new Slice(keyValue.getKey());
7c673cae FG	82	Slice valueSlice = new Slice(keyValue.getValue());
11fdf7f2 TL	83	byte[] keyBytes = keyValue.getKey().getBytes();
11fdf7f2 TL	84	byte[] valueBytes = keyValue.getValue().getBytes();
f67539c2 TL	85	ByteBuffer keyDirect = ByteBuffer.allocateDirect(keyBytes.length);
	86	keyDirect.put(keyBytes);
	87	keyDirect.flip();
	88	ByteBuffer valueDirect = ByteBuffer.allocateDirect(valueBytes.length);
	89	valueDirect.put(valueBytes);
	90	valueDirect.flip();
11fdf7f2 TL	91	switch (keyValue.getOpType()) {
	92	case PUT:
	93	sstFileWriter.put(keySlice, valueSlice);
	94	break;
	95	case PUT_BYTES:
	96	sstFileWriter.put(keyBytes, valueBytes);
	97	break;
f67539c2 TL	98	case PUT_DIRECT:
	99	sstFileWriter.put(keyDirect, valueDirect);
	100	assertThat(keyDirect.position()).isEqualTo(keyBytes.length);
	101	assertThat(keyDirect.limit()).isEqualTo(keyBytes.length);
	102	assertThat(valueDirect.position()).isEqualTo(valueBytes.length);
	103	assertThat(valueDirect.limit()).isEqualTo(valueBytes.length);
	104	break;
11fdf7f2 TL	105	case MERGE:
	106	sstFileWriter.merge(keySlice, valueSlice);
	107	break;
	108	case MERGE_BYTES:
	109	sstFileWriter.merge(keyBytes, valueBytes);
	110	break;
	111	case DELETE:
	112	sstFileWriter.delete(keySlice);
	113	break;
	114	case DELETE_BYTES:
	115	sstFileWriter.delete(keyBytes);
	116	break;
	117	default:
	118	fail("Unsupported op type");
	119	}
7c673cae FG	120	keySlice.close();
	121	valueSlice.close();
	122	}
	123	sstFileWriter.finish();
f67539c2	124	assertThat(sstFileWriter.fileSize()).isGreaterThan(100);
7c673cae FG	125	} finally {
	126	assertThat(sstFileWriter).isNotNull();
	127	sstFileWriter.close();
	128	options.close();
	129	envOptions.close();
	130	if (comparatorOptions != null) {
	131	comparatorOptions.close();
	132	}
	133	if (comparator != null) {
	134	comparator.close();
	135	}
	136	}
	137	return sstFile;
	138	}
	139
	140	@Test
11fdf7f2 TL	141	public void generateSstFileWithJavaComparator()
	142	throws RocksDBException, IOException {
	143	final List<KeyValueWithOp> keyValues = new ArrayList<>();
	144	keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
	145	keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
	146	keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
	147	keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
	148	keyValues.add(new KeyValueWithOp("key5", "", OpType.DELETE));
	149
7c673cae FG	150	newSstFile(keyValues, true);
	151	}
	152
	153	@Test
11fdf7f2 TL	154	public void generateSstFileWithNativeComparator()
	155	throws RocksDBException, IOException {
	156	final List<KeyValueWithOp> keyValues = new ArrayList<>();
	157	keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
	158	keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
	159	keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
	160	keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
	161	keyValues.add(new KeyValueWithOp("key5", "", OpType.DELETE));
	162
7c673cae FG	163	newSstFile(keyValues, false);
	164	}
	165
	166	@Test
	167	public void ingestSstFile() throws RocksDBException, IOException {
11fdf7f2 TL	168	final List<KeyValueWithOp> keyValues = new ArrayList<>();
11fdf7f2 TL	169	keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
f67539c2	170	keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT_DIRECT));
11fdf7f2 TL	171	keyValues.add(new KeyValueWithOp("key3", "value3", OpType.PUT_BYTES));
	172	keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
	173	keyValues.add(new KeyValueWithOp("key5", "value5", OpType.MERGE_BYTES));
	174	keyValues.add(new KeyValueWithOp("key6", "", OpType.DELETE));
	175	keyValues.add(new KeyValueWithOp("key7", "", OpType.DELETE));
	176
	177
	178	final File sstFile = newSstFile(keyValues, false);
	179	final File dbFolder = parentFolder.newFolder(DB_DIRECTORY_NAME);
	180	try(final StringAppendOperator stringAppendOperator =
	181	new StringAppendOperator();
	182	final Options options = new Options()
	183	.setCreateIfMissing(true)
	184	.setMergeOperator(stringAppendOperator);
	185	final RocksDB db = RocksDB.open(options, dbFolder.getAbsolutePath());
	186	final IngestExternalFileOptions ingestExternalFileOptions =
	187	new IngestExternalFileOptions()) {
	188	db.ingestExternalFile(Arrays.asList(sstFile.getAbsolutePath()),
	189	ingestExternalFileOptions);
	190
	191	assertThat(db.get("key1".getBytes())).isEqualTo("value1".getBytes());
	192	assertThat(db.get("key2".getBytes())).isEqualTo("value2".getBytes());
	193	assertThat(db.get("key3".getBytes())).isEqualTo("value3".getBytes());
	194	assertThat(db.get("key4".getBytes())).isEqualTo("value4".getBytes());
	195	assertThat(db.get("key5".getBytes())).isEqualTo("value5".getBytes());
	196	assertThat(db.get("key6".getBytes())).isEqualTo(null);
	197	assertThat(db.get("key7".getBytes())).isEqualTo(null);
	198	}
	199	}
	200
	201	@Test
	202	public void ingestSstFile_cf() throws RocksDBException, IOException {
	203	final List<KeyValueWithOp> keyValues = new ArrayList<>();
	204	keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
	205	keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
	206	keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
	207	keyValues.add(new KeyValueWithOp("key4", "", OpType.DELETE));
	208
7c673cae FG	209	final File sstFile = newSstFile(keyValues, false);
7c673cae FG	210	final File dbFolder = parentFolder.newFolder(DB_DIRECTORY_NAME);
11fdf7f2 TL	211	try(final StringAppendOperator stringAppendOperator =
	212	new StringAppendOperator();
	213	final Options options = new Options()
	214	.setCreateIfMissing(true)
	215	.setCreateMissingColumnFamilies(true)
	216	.setMergeOperator(stringAppendOperator);
	217	final RocksDB db = RocksDB.open(options, dbFolder.getAbsolutePath());
	218	final IngestExternalFileOptions ingestExternalFileOptions =
	219	new IngestExternalFileOptions()) {
7c673cae	220
11fdf7f2 TL	221	try(final ColumnFamilyOptions cf_opts = new ColumnFamilyOptions()
	222	.setMergeOperator(stringAppendOperator);
	223	final ColumnFamilyHandle cf_handle = db.createColumnFamily(
	224	new ColumnFamilyDescriptor("new_cf".getBytes(), cf_opts))) {
7c673cae	225
11fdf7f2 TL	226	db.ingestExternalFile(cf_handle,
	227	Arrays.asList(sstFile.getAbsolutePath()),
	228	ingestExternalFileOptions);
	229
	230	assertThat(db.get(cf_handle,
	231	"key1".getBytes())).isEqualTo("value1".getBytes());
	232	assertThat(db.get(cf_handle,
	233	"key2".getBytes())).isEqualTo("value2".getBytes());
	234	assertThat(db.get(cf_handle,
	235	"key3".getBytes())).isEqualTo("value3".getBytes());
	236	assertThat(db.get(cf_handle,
	237	"key4".getBytes())).isEqualTo(null);
	238	}
	239	}
7c673cae FG	240	}
7c673cae FG	241	}