org.apache.spark.sql.catalyst.util.ArrayData Java Examples

public void addValue(int rowId, int column, SpecializedGetters data,
                     ColumnVector output) {
  if (data.isNullAt(column)) {
    output.noNulls = false;
    output.isNull[rowId] = true;
  } else {
    output.isNull[rowId] = false;
    MapData map = data.getMap(column);
    ArrayData key = map.keyArray();
    ArrayData value = map.valueArray();
    MapColumnVector cv = (MapColumnVector) output;
    // record the length and start of the list elements
    cv.lengths[rowId] = value.numElements();
    cv.offsets[rowId] = cv.childCount;
    cv.childCount += cv.lengths[rowId];
    // make sure the child is big enough
    cv.keys.ensureSize(cv.childCount, true);
    cv.values.ensureSize(cv.childCount, true);
    // Add each element
    for(int e=0; e < cv.lengths[rowId]; ++e) {
      int pos = (int)(e + cv.offsets[rowId]);
      keyConverter.addValue(pos, e, key, cv.keys);
      valueConverter.addValue(pos, e, value, cv.values);
    }
  }
}
 

private static void assertEquals(String context, MapType map, MapData expected, MapData actual) {
  Assert.assertEquals("Should have the same number of elements",
      expected.numElements(), actual.numElements());

  DataType keyType = map.keyType();
  ArrayData expectedKeys = expected.keyArray();
  ArrayData expectedValues = expected.valueArray();

  DataType valueType = map.valueType();
  ArrayData actualKeys = actual.keyArray();
  ArrayData actualValues = actual.valueArray();

  for (int i = 0; i < actual.numElements(); i += 1) {
    assertEquals(context + ".key", keyType,
        expectedKeys.get(i, keyType), actualKeys.get(i, keyType));
    assertEquals(context + ".value", valueType,
        expectedValues.get(i, valueType), actualValues.get(i, valueType));
  }
}
 

public void addValue(int rowId, int column, SpecializedGetters data,
                     ColumnVector output) {
  if (data.isNullAt(column)) {
    output.noNulls = false;
    output.isNull[rowId] = true;
  } else {
    output.isNull[rowId] = false;
    ArrayData value = data.getArray(column);
    ListColumnVector cv = (ListColumnVector) output;
    // record the length and start of the list elements
    cv.lengths[rowId] = value.numElements();
    cv.offsets[rowId] = cv.childCount;
    cv.childCount += cv.lengths[rowId];
    // make sure the child is big enough
    cv.child.ensureSize(cv.childCount, true);
    // Add each element
    for(int e=0; e < cv.lengths[rowId]; ++e) {
      children.addValue((int) (e + cv.offsets[rowId]), e, value, cv.child);
    }
  }
}
 

private static void assertEqualsUnsafe(Types.MapType map, Map<?, ?> expected, MapData actual) {
  Type keyType = map.keyType();
  Type valueType = map.valueType();

  List<Map.Entry<?, ?>> expectedElements = Lists.newArrayList(expected.entrySet());
  ArrayData actualKeys = actual.keyArray();
  ArrayData actualValues = actual.valueArray();

  for (int i = 0; i < expectedElements.size(); i += 1) {
    Map.Entry<?, ?> expectedPair = expectedElements.get(i);
    Object actualKey = actualKeys.get(i, convert(keyType));
    Object actualValue = actualValues.get(i, convert(keyType));

    assertEqualsUnsafe(keyType, expectedPair.getKey(), actualKey);
    assertEqualsUnsafe(valueType, expectedPair.getValue(), actualValue);
  }
}
 

private static void assertEqualsUnsafe(Types.MapType map, Map<?, ?> expected, MapData actual) {
  Type keyType = map.keyType();
  Type valueType = map.valueType();

  List<Map.Entry<?, ?>> expectedElements = Lists.newArrayList(expected.entrySet());
  ArrayData actualKeys = actual.keyArray();
  ArrayData actualValues = actual.valueArray();

  for (int i = 0; i < expectedElements.size(); i += 1) {
    Map.Entry<?, ?> expectedPair = expectedElements.get(i);
    Object actualKey = actualKeys.get(i, convert(keyType));
    Object actualValue = actualValues.get(i, convert(keyType));

    assertEqualsUnsafe(keyType, expectedPair.getKey(), actualKey);
    assertEqualsUnsafe(valueType, expectedPair.getValue(), actualValue);
  }
}
 

private static void assertEqualsUnsafe(Types.MapType map, Map<?, ?> expected, MapData actual) {
  Type keyType = map.keyType();
  Type valueType = map.valueType();

  List<Map.Entry<?, ?>> expectedElements = Lists.newArrayList(expected.entrySet());
  ArrayData actualKeys = actual.keyArray();
  ArrayData actualValues = actual.valueArray();

  for (int i = 0; i < expectedElements.size(); i += 1) {
    Map.Entry<?, ?> expectedPair = expectedElements.get(i);
    Object actualKey = actualKeys.get(i, convert(keyType));
    Object actualValue = actualValues.get(i, convert(keyType));

    assertEqualsUnsafe(keyType, expectedPair.getKey(), actualKey);
    assertEqualsUnsafe(valueType, expectedPair.getValue(), actualValue);
  }
}
 

@Override
public void addValue(int rowId, int column, SpecializedGetters data,
                     ColumnVector output) {
  if (data.isNullAt(column)) {
    output.noNulls = false;
    output.isNull[rowId] = true;
  } else {
    output.isNull[rowId] = false;
    ArrayData value = data.getArray(column);
    ListColumnVector cv = (ListColumnVector) output;
    // record the length and start of the list elements
    cv.lengths[rowId] = value.numElements();
    cv.offsets[rowId] = cv.childCount;
    cv.childCount += cv.lengths[rowId];
    // make sure the child is big enough
    cv.child.ensureSize(cv.childCount, true);
    // Add each element
    for (int e = 0; e < cv.lengths[rowId]; ++e) {
      children.addValue((int) (e + cv.offsets[rowId]), e, value, cv.child);
    }
  }
}
 

private static void assertEqualsUnsafe(Types.ListType list, Collection<?> expected, ArrayData actual) {
  Type elementType = list.elementType();
  List<?> expectedElements = Lists.newArrayList(expected);
  for (int i = 0; i < expectedElements.size(); i += 1) {
    Object expectedValue = expectedElements.get(i);
    Object actualValue = actual.get(i, convert(elementType));

    assertEqualsUnsafe(elementType, expectedValue, actualValue);
  }
}
 

private static void assertEquals(String context, ArrayType array, ArrayData expected, ArrayData actual) {
  Assert.assertEquals("Should have the same number of elements",
      expected.numElements(), actual.numElements());
  DataType type = array.elementType();
  for (int i = 0; i < actual.numElements(); i += 1) {
    assertEquals(context + ".element", type, expected.get(i, type), actual.get(i, type));
  }
}
 

private static void assertEquals(String context, DataType type, Object expected, Object actual) {
  if (expected == null && actual == null) {
    return;
  }

  if (type instanceof StructType) {
    Assert.assertTrue("Expected should be an InternalRow: " + context,
        expected instanceof InternalRow);
    Assert.assertTrue("Actual should be an InternalRow: " + context,
        actual instanceof InternalRow);
    assertEquals(context, (StructType) type, (InternalRow) expected, (InternalRow) actual);

  } else if (type instanceof ArrayType) {
    Assert.assertTrue("Expected should be an ArrayData: " + context,
        expected instanceof ArrayData);
    Assert.assertTrue("Actual should be an ArrayData: " + context,
        actual instanceof ArrayData);
    assertEquals(context, (ArrayType) type, (ArrayData) expected, (ArrayData) actual);

  } else if (type instanceof MapType) {
    Assert.assertTrue("Expected should be a MapData: " + context,
        expected instanceof MapData);
    Assert.assertTrue("Actual should be a MapData: " + context,
        actual instanceof MapData);
    assertEquals(context, (MapType) type, (MapData) expected, (MapData) actual);

  } else if (type instanceof BinaryType) {
    assertEqualBytes(context, (byte[]) expected, (byte[]) actual);
  } else {
    Assert.assertEquals("Value should match expected: " + context, expected, actual);
  }
}
 

private static void assertEqualsUnsafe(Types.ListType list, Collection<?> expected, ArrayData actual) {
  Type elementType = list.elementType();
  List<?> expectedElements = Lists.newArrayList(expected);
  for (int i = 0; i < expectedElements.size(); i += 1) {
    Object expectedValue = expectedElements.get(i);
    Object actualValue = actual.get(i, convert(elementType));

    assertEqualsUnsafe(elementType, expectedValue, actualValue);
  }
}
 

private static void assertEqualsUnsafe(Types.ListType list, Collection<?> expected, ArrayData actual) {
  Type elementType = list.elementType();
  List<?> expectedElements = Lists.newArrayList(expected);
  for (int i = 0; i < expectedElements.size(); i += 1) {
    Object expectedValue = expectedElements.get(i);
    Object actualValue = actual.get(i, convert(elementType));

    assertEqualsUnsafe(elementType, expectedValue, actualValue);
  }
}
 

private ArrayData collectionToArrayData(Type elementType, Collection<?> values) {
  switch (elementType.typeId()) {
    case BOOLEAN:
    case INTEGER:
    case DATE:
    case TIME:
    case LONG:
    case TIMESTAMP:
    case FLOAT:
    case DOUBLE:
      return fillArray(values, array -> (pos, value) -> array[pos] = value);
    case STRING:
      return fillArray(values, array ->
          (BiConsumer<Integer, CharSequence>) (pos, seq) -> array[pos] = UTF8String.fromString(seq.toString()));
    case FIXED:
    case BINARY:
      return fillArray(values, array ->
          (BiConsumer<Integer, ByteBuffer>) (pos, buf) -> array[pos] = ByteBuffers.toByteArray(buf));
    case DECIMAL:
      return fillArray(values, array ->
          (BiConsumer<Integer, BigDecimal>) (pos, dec) -> array[pos] = Decimal.apply(dec));
    case STRUCT:
      return fillArray(values, array -> (BiConsumer<Integer, StructLike>) (pos, tuple) ->
          array[pos] = new StructInternalRow(elementType.asStructType(), tuple));
    case LIST:
      return fillArray(values, array -> (BiConsumer<Integer, Collection<?>>) (pos, list) ->
          array[pos] = collectionToArrayData(elementType.asListType(), list));
    case MAP:
      return fillArray(values, array -> (BiConsumer<Integer, Map<?, ?>>) (pos, map) ->
          array[pos] = mapToMapData(elementType.asMapType(), map));
    default:
      throw new UnsupportedOperationException("Unsupported array element type: " + elementType);
  }
}
 

@Override
@SuppressWarnings("unchecked")
public void write(ArrayData array, Encoder encoder) throws IOException {
  encoder.writeArrayStart();
  int numElements = array.numElements();
  encoder.setItemCount(numElements);
  for (int i = 0; i < numElements; i += 1) {
    encoder.startItem();
    elementWriter.write((T) array.get(i, elementType), encoder);
  }
  encoder.writeArrayEnd();
}
 

@Override
public void addValue(int rowId, int column, SpecializedGetters data,
                     ColumnVector output) {
  if (data.isNullAt(column)) {
    output.noNulls = false;
    output.isNull[rowId] = true;
  } else {
    output.isNull[rowId] = false;
    MapData map = data.getMap(column);
    ArrayData key = map.keyArray();
    ArrayData value = map.valueArray();
    MapColumnVector cv = (MapColumnVector) output;
    // record the length and start of the list elements
    cv.lengths[rowId] = value.numElements();
    cv.offsets[rowId] = cv.childCount;
    cv.childCount += cv.lengths[rowId];
    // make sure the child is big enough
    cv.keys.ensureSize(cv.childCount, true);
    cv.values.ensureSize(cv.childCount, true);
    // Add each element
    for (int e = 0; e < cv.lengths[rowId]; ++e) {
      int pos = (int) (e + cv.offsets[rowId]);
      keyConverter.addValue(pos, e, key, cv.keys);
      valueConverter.addValue(pos, e, value, cv.values);
    }
  }
}
 

@Override
public ArrayData nonNullRead(ColumnVector vector, int row) {
  ListColumnVector listVector = (ListColumnVector) vector;
  int offset = (int) listVector.offsets[row];
  int length = (int) listVector.lengths[row];
  List<Object> elements = Lists.newArrayListWithExpectedSize(length);
  for (int c = 0; c < length; ++c) {
    elements.add(elementReader.read(listVector.child, offset + c));
  }
  return new GenericArrayData(elements.toArray());
}
 

@Override
@SuppressWarnings("unchecked")
public void write(MapData map, Encoder encoder) throws IOException {
  encoder.writeArrayStart();
  int numElements = map.numElements();
  encoder.setItemCount(numElements);
  ArrayData keyArray = map.keyArray();
  ArrayData valueArray = map.valueArray();
  for (int i = 0; i < numElements; i += 1) {
    encoder.startItem();
    keyWriter.write((K) keyArray.get(i, keyType), encoder);
    valueWriter.write((V) valueArray.get(i, valueType), encoder);
  }
  encoder.writeArrayEnd();
}
 

@Override
@SuppressWarnings("unchecked")
public void write(MapData map, Encoder encoder) throws IOException {
  encoder.writeMapStart();
  int numElements = map.numElements();
  encoder.setItemCount(numElements);
  ArrayData keyArray = map.keyArray();
  ArrayData valueArray = map.valueArray();
  for (int i = 0; i < numElements; i += 1) {
    encoder.startItem();
    keyWriter.write((K) keyArray.get(i, keyType), encoder);
    valueWriter.write((V) valueArray.get(i, valueType), encoder);
  }
  encoder.writeMapEnd();
}
 

@Override
@SuppressWarnings("unchecked")
protected ReusableArrayData newListData(ArrayData reuse) {
  this.readPos = 0;
  this.writePos = 0;

  if (reuse instanceof ReusableArrayData) {
    return (ReusableArrayData) reuse;
  } else {
    return new ReusableArrayData();
  }
}
 

@Override
@SuppressWarnings("unchecked")
public void write(MapData map, Encoder encoder) throws IOException {
  encoder.writeMapStart();
  int numElements = map.numElements();
  encoder.setItemCount(numElements);
  ArrayData keyArray = map.keyArray();
  ArrayData valueArray = map.valueArray();
  for (int i = 0; i < numElements; i += 1) {
    encoder.startItem();
    keyWriter.write((K) keyArray.get(i, keyType), encoder);
    valueWriter.write((V) valueArray.get(i, valueType), encoder);
  }
  encoder.writeMapEnd();
}
 

@Override
@SuppressWarnings("unchecked")
public void write(MapData map, Encoder encoder) throws IOException {
  encoder.writeArrayStart();
  int numElements = map.numElements();
  encoder.setItemCount(numElements);
  ArrayData keyArray = map.keyArray();
  ArrayData valueArray = map.valueArray();
  for (int i = 0; i < numElements; i += 1) {
    encoder.startItem();
    keyWriter.write((K) keyArray.get(i, keyType), encoder);
    valueWriter.write((V) valueArray.get(i, valueType), encoder);
  }
  encoder.writeArrayEnd();
}
 

@Override
@SuppressWarnings("unchecked")
public void write(ArrayData array, Encoder encoder) throws IOException {
  encoder.writeArrayStart();
  int numElements = array.numElements();
  encoder.setItemCount(numElements);
  for (int i = 0; i < numElements; i += 1) {
    encoder.startItem();
    elementWriter.write((T) array.get(i, elementType), encoder);
  }
  encoder.writeArrayEnd();
}
 

@Override
@SuppressWarnings("unchecked")
protected ReusableArrayData newListData(ArrayData reuse) {
  this.readPos = 0;
  this.writePos = 0;

  if (reuse instanceof ReusableArrayData) {
    return (ReusableArrayData) reuse;
  } else {
    return new ReusableArrayData();
  }
}
 

@Override
public ArrayData copy() {
  return new GenericArrayData(array());
}
 

@Override
public ArrayData getArray(int ordinal) {
  return (ArrayData) values[ordinal];
}
 

@Override
public <T> List<T> getList(int i) {
    ArrayData arrayData = row.getArray(i); // Slow but functional
    return Arrays.<T>asList((T[])arrayData.array());
}
 

static ValueReader<ArrayData> array(ValueReader<?> elementReader) {
  return new ArrayReader(elementReader);
}
 

public UnsafeRow apply(InternalRow i) {
  holder.reset();

  rowWriter.zeroOutNullBytes();


  boolean isNull = i.isNullAt(0);
  MapData value = isNull ? null : (i.getMap(0));
  if (isNull) {
    rowWriter.setNullAt(0);
  } else {
    // Remember the current cursor so that we can calculate how many bytes are
    // written later.
    final int tmpCursor = holder.cursor;

    if (value instanceof UnsafeMapData) {

      final int sizeInBytes = ((UnsafeMapData) value).getSizeInBytes();
      // grow the global buffer before writing data.
      holder.grow(sizeInBytes);
      ((UnsafeMapData) value).writeToMemory(holder.buffer, holder.cursor);
      holder.cursor += sizeInBytes;

    } else {
      final ArrayData keys = value.keyArray();
      final ArrayData values = value.valueArray();

      // preserve 8 bytes to write the key array numBytes later.
      holder.grow(8);
      holder.cursor += 8;

      // Remember the current cursor so that we can write numBytes of key array later.
      final int tmpCursor1 = holder.cursor;


      if (keys instanceof UnsafeArrayData) {

        final int sizeInBytes1 = ((UnsafeArrayData) keys).getSizeInBytes();
        // grow the global buffer before writing data.
        holder.grow(sizeInBytes1);
        ((UnsafeArrayData) keys).writeToMemory(holder.buffer, holder.cursor);
        holder.cursor += sizeInBytes1;

      } else {
        final int numElements = keys.numElements();
        arrayWriter.initialize(holder, numElements, 8);

        for (int index = 0; index < numElements; index++) {
          if (keys.isNullAt(index)) {
            arrayWriter.setNull(index);
          } else {
            final UTF8String element = keys.getUTF8String(index);
            arrayWriter.write(index, element);
          }
        }
      }

      // Write the numBytes of key array into the first 8 bytes.
      Platform.putLong(holder.buffer, tmpCursor1 - 8, holder.cursor - tmpCursor1);


      if (values instanceof UnsafeArrayData) {

        final int sizeInBytes2 = ((UnsafeArrayData) values).getSizeInBytes();
        // grow the global buffer before writing data.
        holder.grow(sizeInBytes2);
        ((UnsafeArrayData) values).writeToMemory(holder.buffer, holder.cursor);
        holder.cursor += sizeInBytes2;

      } else {
        final int numElements1 = values.numElements();
        arrayWriter1.initialize(holder, numElements1, 8);

        for (int index1 = 0; index1 < numElements1; index1++) {
          if (values.isNullAt(index1)) {
            arrayWriter1.setNull(index1);
          } else {
            final UTF8String element1 = values.getUTF8String(index1);
            arrayWriter1.write(index1, element1);
          }
        }
      }

    }

    rowWriter.setOffsetAndSize(0, tmpCursor, holder.cursor - tmpCursor);
  }
  result.setTotalSize(holder.totalSize());
  return result;
}
 

public UnsafeRow apply(InternalRow i) {
  holder.reset();

  rowWriter.zeroOutNullBytes();


  boolean isNull = i.isNullAt(0);
  ArrayData value = isNull ? null : (i.getArray(0));
  if (isNull) {
    rowWriter.setNullAt(0);
  } else {
    // Remember the current cursor so that we can calculate how many bytes are
    // written later.
    final int tmpCursor = holder.cursor;

    if (value instanceof UnsafeArrayData) {

      final int sizeInBytes1 = ((UnsafeArrayData) value).getSizeInBytes();
      // grow the global buffer before writing data.
      holder.grow(sizeInBytes1);
      ((UnsafeArrayData) value).writeToMemory(holder.buffer, holder.cursor);
      holder.cursor += sizeInBytes1;

    } else {
      final int numElements = value.numElements();
      arrayWriter.initialize(holder, numElements, 8);

      for (int index = 0; index < numElements; index++) {
        if (value.isNullAt(index)) {
          arrayWriter.setNull(index);
        } else {
          final InternalRow element = value.getStruct(index, 2);

          final int tmpCursor1 = holder.cursor;

          if (element instanceof UnsafeRow) {

            final int sizeInBytes = ((UnsafeRow) element).getSizeInBytes();
            // grow the global buffer before writing data.
            holder.grow(sizeInBytes);
            ((UnsafeRow) element).writeToMemory(holder.buffer, holder.cursor);
            holder.cursor += sizeInBytes;

          } else {
            rowWriter1.reset();


            boolean isNull1 = element.isNullAt(0);
            int value1 = isNull1 ? -1 : element.getInt(0);

            if (isNull1) {
              rowWriter1.setNullAt(0);
            } else {
              rowWriter1.write(0, value1);
            }


            boolean isNull2 = element.isNullAt(1);
            int value2 = isNull2 ? -1 : element.getInt(1);

            if (isNull2) {
              rowWriter1.setNullAt(1);
            } else {
              rowWriter1.write(1, value2);
            }
          }

          arrayWriter.setOffsetAndSize(index, tmpCursor1, holder.cursor - tmpCursor1);

        }
      }
    }

    rowWriter.setOffsetAndSize(0, tmpCursor, holder.cursor - tmpCursor);
  }
  result.setTotalSize(holder.totalSize());
  return result;
}
 

public UnsafeRow apply(InternalRow i) {
  holder.reset();

  rowWriter.zeroOutNullBytes();


  boolean isNull = i.isNullAt(0);
  MapData value = isNull ? null : (i.getMap(0));
  if (isNull) {
    rowWriter.setNullAt(0);
  } else {
    // Remember the current cursor so that we can calculate how many bytes are
    // written later.
    final int tmpCursor = holder.cursor;

    if (value instanceof UnsafeMapData) {

      final int sizeInBytes = ((UnsafeMapData) value).getSizeInBytes();
      // grow the global buffer before writing data.
      holder.grow(sizeInBytes);
      ((UnsafeMapData) value).writeToMemory(holder.buffer, holder.cursor);
      holder.cursor += sizeInBytes;

    } else {
      final ArrayData keys = value.keyArray();
      final ArrayData values = value.valueArray();

      // preserve 8 bytes to write the key array numBytes later.
      holder.grow(8);
      holder.cursor += 8;

      // Remember the current cursor so that we can write numBytes of key array later.
      final int tmpCursor1 = holder.cursor;


      if (keys instanceof UnsafeArrayData) {

        final int sizeInBytes1 = ((UnsafeArrayData) keys).getSizeInBytes();
        // grow the global buffer before writing data.
        holder.grow(sizeInBytes1);
        ((UnsafeArrayData) keys).writeToMemory(holder.buffer, holder.cursor);
        holder.cursor += sizeInBytes1;

      } else {
        final int numElements = keys.numElements();
        arrayWriter.initialize(holder, numElements, 8);

        for (int index = 0; index < numElements; index++) {
          if (keys.isNullAt(index)) {
            arrayWriter.setNull(index);
          } else {
            final UTF8String element = keys.getUTF8String(index);
            arrayWriter.write(index, element);
          }
        }
      }

      // Write the numBytes of key array into the first 8 bytes.
      Platform.putLong(holder.buffer, tmpCursor1 - 8, holder.cursor - tmpCursor1);


      if (values instanceof UnsafeArrayData) {

        final int sizeInBytes2 = ((UnsafeArrayData) values).getSizeInBytes();
        // grow the global buffer before writing data.
        holder.grow(sizeInBytes2);
        ((UnsafeArrayData) values).writeToMemory(holder.buffer, holder.cursor);
        holder.cursor += sizeInBytes2;

      } else {
        final int numElements1 = values.numElements();
        arrayWriter1.initialize(holder, numElements1, 8);

        for (int index1 = 0; index1 < numElements1; index1++) {
          if (values.isNullAt(index1)) {
            arrayWriter1.setNull(index1);
          } else {
            final UTF8String element1 = values.getUTF8String(index1);
            arrayWriter1.write(index1, element1);
          }
        }
      }

    }

    rowWriter.setOffsetAndSize(0, tmpCursor, holder.cursor - tmpCursor);
  }
  result.setTotalSize(holder.totalSize());
  return result;
}