content: /*
   * 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.spark.sql.catalyst.expressions;
  
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  
  import org.apache.spark.SparkConf;
  import org.apache.spark.memory.TaskMemoryManager;
  import org.apache.spark.memory.TestMemoryManager;
  import org.apache.spark.sql.types.StructType;
  import org.apache.spark.sql.types.DataTypes;
  import org.apache.spark.sql.catalyst.expressions.codegen.BufferHolder;
  import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter;
  import org.apache.spark.unsafe.types.UTF8String;
  
  import java.util.Random;
  
  public class RowBasedKeyValueBatchSuite {
  
    private final Random rand = new Random(42);
  
    private TestMemoryManager memoryManager;
    private TaskMemoryManager taskMemoryManager;
    private StructType keySchema = new StructType().add("k1", DataTypes.LongType)
            .add("k2", DataTypes.StringType);
    private StructType fixedKeySchema = new StructType().add("k1", DataTypes.LongType)
            .add("k2", DataTypes.LongType);
    private StructType valueSchema = new StructType().add("count", DataTypes.LongType)
            .add("sum", DataTypes.LongType);
    private int DEFAULT_CAPACITY = 1 << 16;
  
    private String getRandomString(int length) {
      Assert.assertTrue(length >= 0);
      final byte[] bytes = new byte[length];
      rand.nextBytes(bytes);
      return new String(bytes);
    }
  
    private UnsafeRow makeKeyRow(long k1, String k2) {
      UnsafeRow row = new UnsafeRow(2);
      BufferHolder holder = new BufferHolder(row, 32);
      UnsafeRowWriter writer = new UnsafeRowWriter(holder, 2);
      holder.reset();
      writer.write(0, k1);
      writer.write(1, UTF8String.fromString(k2));
      row.setTotalSize(holder.totalSize());
      return row;
    }
  
    private UnsafeRow makeKeyRow(long k1, long k2) {
      UnsafeRow row = new UnsafeRow(2);
      BufferHolder holder = new BufferHolder(row, 0);
      UnsafeRowWriter writer = new UnsafeRowWriter(holder, 2);
      holder.reset();
      writer.write(0, k1);
      writer.write(1, k2);
      row.setTotalSize(holder.totalSize());
      return row;
    }
  
    private UnsafeRow makeValueRow(long v1, long v2) {
      UnsafeRow row = new UnsafeRow(2);
      BufferHolder holder = new BufferHolder(row, 0);
      UnsafeRowWriter writer = new UnsafeRowWriter(holder, 2);
      holder.reset();
      writer.write(0, v1);
      writer.write(1, v2);
      row.setTotalSize(holder.totalSize());
      return row;
    }
  
    private UnsafeRow appendRow(RowBasedKeyValueBatch batch, UnsafeRow key, UnsafeRow value) {
      return batch.appendRow(key.getBaseObject(), key.getBaseOffset(), key.getSizeInBytes(),
              value.getBaseObject(), value.getBaseOffset(), value.getSizeInBytes());
    }
  
    private void updateValueRow(UnsafeRow row, long v1, long v2) {
      row.setLong(0, v1);
      row.setLong(1, v2);
    }
  
    private boolean checkKey(UnsafeRow row, long k1, String k2) {
      return (row.getLong(0) == k1)
              && (row.getUTF8String(1).equals(UTF8String.fromString(k2)));
    }
  
    private boolean checkKey(UnsafeRow row, long k1, long k2) {
      return (row.getLong(0) == k1)
              && (row.getLong(1) == k2);
    }
  
    private boolean checkValue(UnsafeRow row, long v1, long v2) {
      return (row.getLong(0) == v1) && (row.getLong(1) == v2);
    }
  
    @Before
    public void setup() {
      memoryManager = new TestMemoryManager(new SparkConf()
              .set("spark.memory.offHeap.enabled", "false")
              .set("spark.shuffle.spill.compress", "false")
              .set("spark.shuffle.compress", "false"));
      taskMemoryManager = new TaskMemoryManager(memoryManager, 0);
    }
  
    @After
    public void tearDown() {
      if (taskMemoryManager != null) {
        Assert.assertEquals(0L, taskMemoryManager.cleanUpAllAllocatedMemory());
        long leakedMemory = taskMemoryManager.getMemoryConsumptionForThisTask();
        taskMemoryManager = null;
        Assert.assertEquals(0L, leakedMemory);
      }
    }
  
  
    @Test
    public void emptyBatch() throws Exception {
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      try {
        Assert.assertEquals(0, batch.numRows());
        try {
          batch.getKeyRow(-1);
          Assert.fail("Should not be able to get row -1");
        } catch (AssertionError e) {
          // Expected exception; do nothing.
        }
        try {
          batch.getValueRow(-1);
          Assert.fail("Should not be able to get row -1");
        } catch (AssertionError e) {
          // Expected exception; do nothing.
        }
        try {
          batch.getKeyRow(0);
          Assert.fail("Should not be able to get row 0 when batch is empty");
        } catch (AssertionError e) {
          // Expected exception; do nothing.
        }
        try {
          batch.getValueRow(0);
          Assert.fail("Should not be able to get row 0 when batch is empty");
        } catch (AssertionError e) {
          // Expected exception; do nothing.
        }
        Assert.assertFalse(batch.rowIterator().next());
      } finally {
        batch.close();
      }
    }
  
    @Test
    public void batchType() throws Exception {
      RowBasedKeyValueBatch batch1 = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      RowBasedKeyValueBatch batch2 = RowBasedKeyValueBatch.allocate(fixedKeySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      try {
        Assert.assertEquals(batch1.getClass(), VariableLengthRowBasedKeyValueBatch.class);
        Assert.assertEquals(batch2.getClass(), FixedLengthRowBasedKeyValueBatch.class);
      } finally {
        batch1.close();
        batch2.close();
      }
    }
  
    @Test
    public void setAndRetrieve() {
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      try {
        UnsafeRow ret1 = appendRow(batch, makeKeyRow(1, "A"), makeValueRow(1, 1));
        Assert.assertTrue(checkValue(ret1, 1, 1));
        UnsafeRow ret2 = appendRow(batch, makeKeyRow(2, "B"), makeValueRow(2, 2));
        Assert.assertTrue(checkValue(ret2, 2, 2));
        UnsafeRow ret3 = appendRow(batch, makeKeyRow(3, "C"), makeValueRow(3, 3));
        Assert.assertTrue(checkValue(ret3, 3, 3));
        Assert.assertEquals(3, batch.numRows());
        UnsafeRow retrievedKey1 = batch.getKeyRow(0);
        Assert.assertTrue(checkKey(retrievedKey1, 1, "A"));
        UnsafeRow retrievedKey2 = batch.getKeyRow(1);
        Assert.assertTrue(checkKey(retrievedKey2, 2, "B"));
        UnsafeRow retrievedValue1 = batch.getValueRow(1);
        Assert.assertTrue(checkValue(retrievedValue1, 2, 2));
        UnsafeRow retrievedValue2 = batch.getValueRow(2);
        Assert.assertTrue(checkValue(retrievedValue2, 3, 3));
        try {
          batch.getKeyRow(3);
          Assert.fail("Should not be able to get row 3");
        } catch (AssertionError e) {
          // Expected exception; do nothing.
        }
        try {
          batch.getValueRow(3);
          Assert.fail("Should not be able to get row 3");
        } catch (AssertionError e) {
          // Expected exception; do nothing.
        }
      } finally {
        batch.close();
      }
    }
  
    @Test
    public void setUpdateAndRetrieve() {
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      try {
        appendRow(batch, makeKeyRow(1, "A"), makeValueRow(1, 1));
        Assert.assertEquals(1, batch.numRows());
        UnsafeRow retrievedValue = batch.getValueRow(0);
        updateValueRow(retrievedValue, 2, 2);
        UnsafeRow retrievedValue2 = batch.getValueRow(0);
        Assert.assertTrue(checkValue(retrievedValue2, 2, 2));
      } finally {
        batch.close();
      }
    }
  
  
    @Test
    public void iteratorTest() throws Exception {
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      try {
        appendRow(batch, makeKeyRow(1, "A"), makeValueRow(1, 1));
        appendRow(batch, makeKeyRow(2, "B"), makeValueRow(2, 2));
        appendRow(batch, makeKeyRow(3, "C"), makeValueRow(3, 3));
        Assert.assertEquals(3, batch.numRows());
        org.apache.spark.unsafe.KVIterator<UnsafeRow, UnsafeRow> iterator
                = batch.rowIterator();
        Assert.assertTrue(iterator.next());
        UnsafeRow key1 = iterator.getKey();
        UnsafeRow value1 = iterator.getValue();
        Assert.assertTrue(checkKey(key1, 1, "A"));
        Assert.assertTrue(checkValue(value1, 1, 1));
        Assert.assertTrue(iterator.next());
        UnsafeRow key2 = iterator.getKey();
        UnsafeRow value2 = iterator.getValue();
        Assert.assertTrue(checkKey(key2, 2, "B"));
        Assert.assertTrue(checkValue(value2, 2, 2));
        Assert.assertTrue(iterator.next());
        UnsafeRow key3 = iterator.getKey();
        UnsafeRow value3 = iterator.getValue();
        Assert.assertTrue(checkKey(key3, 3, "C"));
        Assert.assertTrue(checkValue(value3, 3, 3));
        Assert.assertFalse(iterator.next());
      } finally {
        batch.close();
      }
    }
  
    @Test
    public void fixedLengthTest() throws Exception {
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(fixedKeySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      try {
        appendRow(batch, makeKeyRow(11, 11), makeValueRow(1, 1));
        appendRow(batch, makeKeyRow(22, 22), makeValueRow(2, 2));
        appendRow(batch, makeKeyRow(33, 33), makeValueRow(3, 3));
        UnsafeRow retrievedKey1 = batch.getKeyRow(0);
        Assert.assertTrue(checkKey(retrievedKey1, 11, 11));
        UnsafeRow retrievedKey2 = batch.getKeyRow(1);
        Assert.assertTrue(checkKey(retrievedKey2, 22, 22));
        UnsafeRow retrievedValue1 = batch.getValueRow(1);
        Assert.assertTrue(checkValue(retrievedValue1, 2, 2));
        UnsafeRow retrievedValue2 = batch.getValueRow(2);
        Assert.assertTrue(checkValue(retrievedValue2, 3, 3));
        Assert.assertEquals(3, batch.numRows());
        org.apache.spark.unsafe.KVIterator<UnsafeRow, UnsafeRow> iterator
                = batch.rowIterator();
        Assert.assertTrue(iterator.next());
        UnsafeRow key1 = iterator.getKey();
        UnsafeRow value1 = iterator.getValue();
        Assert.assertTrue(checkKey(key1, 11, 11));
        Assert.assertTrue(checkValue(value1, 1, 1));
        Assert.assertTrue(iterator.next());
        UnsafeRow key2 = iterator.getKey();
        UnsafeRow value2 = iterator.getValue();
        Assert.assertTrue(checkKey(key2, 22, 22));
        Assert.assertTrue(checkValue(value2, 2, 2));
        Assert.assertTrue(iterator.next());
        UnsafeRow key3 = iterator.getKey();
        UnsafeRow value3 = iterator.getValue();
        Assert.assertTrue(checkKey(key3, 33, 33));
        Assert.assertTrue(checkValue(value3, 3, 3));
        Assert.assertFalse(iterator.next());
      } finally {
        batch.close();
      }
    }
  
    @Test
    public void appendRowUntilExceedingCapacity() throws Exception {
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, 10);
      try {
        UnsafeRow key = makeKeyRow(1, "A");
        UnsafeRow value = makeValueRow(1, 1);
        for (int i = 0; i < 10; i++) {
          appendRow(batch, key, value);
        }
        UnsafeRow ret = appendRow(batch, key, value);
        Assert.assertEquals(batch.numRows(), 10);
        Assert.assertNull(ret);
        org.apache.spark.unsafe.KVIterator<UnsafeRow, UnsafeRow> iterator
                = batch.rowIterator();
        for (int i = 0; i < 10; i++) {
          Assert.assertTrue(iterator.next());
          UnsafeRow key1 = iterator.getKey();
          UnsafeRow value1 = iterator.getValue();
          Assert.assertTrue(checkKey(key1, 1, "A"));
          Assert.assertTrue(checkValue(value1, 1, 1));
        }
        Assert.assertFalse(iterator.next());
      } finally {
        batch.close();
      }
    }
  
    @Test
    public void appendRowUntilExceedingPageSize() throws Exception {
      // Use default size or spark.buffer.pageSize if specified
      int pageSizeToUse = (int) memoryManager.pageSizeBytes();
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, pageSizeToUse); //enough capacity
      try {
        UnsafeRow key = makeKeyRow(1, "A");
        UnsafeRow value = makeValueRow(1, 1);
        int recordLength = 8 + key.getSizeInBytes() + value.getSizeInBytes() + 8;
        int totalSize = 4;
        int numRows = 0;
        while (totalSize + recordLength < pageSizeToUse) {
          appendRow(batch, key, value);
          totalSize += recordLength;
          numRows++;
        }
        UnsafeRow ret = appendRow(batch, key, value);
        Assert.assertEquals(batch.numRows(), numRows);
        Assert.assertNull(ret);
        org.apache.spark.unsafe.KVIterator<UnsafeRow, UnsafeRow> iterator
                = batch.rowIterator();
        for (int i = 0; i < numRows; i++) {
          Assert.assertTrue(iterator.next());
          UnsafeRow key1 = iterator.getKey();
          UnsafeRow value1 = iterator.getValue();
          Assert.assertTrue(checkKey(key1, 1, "A"));
          Assert.assertTrue(checkValue(value1, 1, 1));
        }
        Assert.assertFalse(iterator.next());
      } finally {
        batch.close();
      }
    }
  
    @Test
    public void failureToAllocateFirstPage() throws Exception {
      memoryManager.limit(1024);
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      try {
        UnsafeRow key = makeKeyRow(1, "A");
        UnsafeRow value = makeValueRow(11, 11);
        UnsafeRow ret = appendRow(batch, key, value);
        Assert.assertNull(ret);
        Assert.assertFalse(batch.rowIterator().next());
      } finally {
        batch.close();
      }
    }
  
    @Test
    public void randomizedTest() {
      RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
              valueSchema, taskMemoryManager, DEFAULT_CAPACITY);
      int numEntry = 100;
      long[] expectedK1 = new long[numEntry];
      String[] expectedK2 = new String[numEntry];
      long[] expectedV1 = new long[numEntry];
      long[] expectedV2 = new long[numEntry];
  
      for (int i = 0; i < numEntry; i++) {
        long k1 = rand.nextLong();
        String k2 = getRandomString(rand.nextInt(256));
        long v1 = rand.nextLong();
        long v2 = rand.nextLong();
        appendRow(batch, makeKeyRow(k1, k2), makeValueRow(v1, v2));
        expectedK1[i] = k1;
        expectedK2[i] = k2;
        expectedV1[i] = v1;
        expectedV2[i] = v2;
      }
      try {
        for (int j = 0; j < 10000; j++) {
          int rowId = rand.nextInt(numEntry);
          if (rand.nextBoolean()) {
            UnsafeRow key = batch.getKeyRow(rowId);
            Assert.assertTrue(checkKey(key, expectedK1[rowId], expectedK2[rowId]));
          }
          if (rand.nextBoolean()) {
            UnsafeRow value = batch.getValueRow(rowId);
            Assert.assertTrue(checkValue(value, expectedV1[rowId], expectedV2[rowId]));
          }
        }
      } finally {
        batch.close();
      }
    }
  }
  