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

The following examples show how to use org.apache.spark.sql.catalyst.util.ArrayData. These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source Project: iceberg   Source File: SparkOrcWriter.java    License: Apache License 2.0 6 votes vote down vote up
@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);
    }
  }
}
 
Example 2
Source Project: iceberg   Source File: GenericsHelpers.java    License: Apache License 2.0 6 votes vote down vote up
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);
  }
}
 
Example 3
Source Project: iceberg   Source File: TestHelpers.java    License: Apache License 2.0 6 votes vote down vote up
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);
  }
}
 
Example 4
Source Project: iceberg   Source File: TestHelpers.java    License: Apache License 2.0 6 votes vote down vote up
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));
  }
}
 
Example 5
Source Project: iceberg   Source File: SparkOrcWriter.java    License: Apache License 2.0 6 votes vote down vote up
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);
    }
  }
}
 
Example 6
Source Project: iceberg   Source File: SparkOrcWriter.java    License: Apache License 2.0 6 votes vote down vote up
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);
    }
  }
}
 
Example 7
Source Project: iceberg   Source File: TestHelpers.java    License: Apache License 2.0 6 votes vote down vote up
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);
  }
}
 
Example 8
Source Project: iceberg   Source File: SparkValueWriters.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 9
Source Project: iceberg   Source File: SparkValueWriters.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 10
Source Project: iceberg   Source File: SparkValueWriters.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 11
Source Project: iceberg   Source File: SparkParquetReaders.java    License: Apache License 2.0 5 votes vote down vote up
@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();
  }
}
 
Example 12
Source Project: iceberg   Source File: SparkOrcValueReaders.java    License: Apache License 2.0 5 votes vote down vote up
@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());
}
 
Example 13
Source Project: iceberg   Source File: SparkOrcWriter.java    License: Apache License 2.0 5 votes vote down vote up
@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);
    }
  }
}
 
Example 14
Source Project: iceberg   Source File: StructInternalRow.java    License: Apache License 2.0 5 votes vote down vote up
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);
  }
}
 
Example 15
Source Project: iceberg   Source File: GenericsHelpers.java    License: Apache License 2.0 5 votes vote down vote up
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);
  }
}
 
Example 16
Source Project: iceberg   Source File: TestHelpers.java    License: Apache License 2.0 5 votes vote down vote up
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);
  }
}
 
Example 17
Source Project: iceberg   Source File: TestHelpers.java    License: Apache License 2.0 5 votes vote down vote up
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);
  }
}
 
Example 18
Source Project: iceberg   Source File: TestHelpers.java    License: Apache License 2.0 5 votes vote down vote up
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));
  }
}
 
Example 19
Source Project: iceberg   Source File: SparkValueWriters.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 20
Source Project: iceberg   Source File: SparkValueWriters.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 21
Source Project: iceberg   Source File: SparkValueWriters.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 22
Source Project: iceberg   Source File: SparkParquetReaders.java    License: Apache License 2.0 5 votes vote down vote up
@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();
  }
}
 
Example 23
Source Project: iceberg   Source File: TestHelpers.java    License: Apache License 2.0 5 votes vote down vote up
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);
  }
}
 
Example 24
Source Project: iceberg   Source File: SparkParquetWriters.java    License: Apache License 2.0 4 votes vote down vote up
@Override
protected Iterator<E> elements(ArrayData list) {
  return new ElementIterator<>(list);
}
 
Example 25
Source Project: iceberg   Source File: SparkParquetWriters.java    License: Apache License 2.0 4 votes vote down vote up
private ElementIterator(ArrayData list) {
  this.list = list;
  size = list.numElements();
  index = 0;
}
 
Example 26
Source Project: iceberg   Source File: SparkValueWriters.java    License: Apache License 2.0 4 votes vote down vote up
static <T> ValueWriter<ArrayData> array(ValueWriter<T> elementWriter, DataType elementType) {
  return new ArrayWriter<>(elementWriter, elementType);
}
 
Example 27
Source Project: iceberg   Source File: SparkParquetReaders.java    License: Apache License 2.0 4 votes vote down vote up
@Override
protected ArrayData buildList(ReusableArrayData list) {
  list.setNumElements(writePos);
  return list;
}
 
Example 28
Source Project: iceberg   Source File: SparkParquetReaders.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public ArrayData copy() {
  return new GenericArrayData(array());
}
 
Example 29
Source Project: iceberg   Source File: SparkParquetReaders.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public ArrayData getArray(int ordinal) {
  return (ArrayData) values[ordinal];
}
 
Example 30
Source Project: iceberg   Source File: SparkValueReaders.java    License: Apache License 2.0 4 votes vote down vote up
static ValueReader<ArrayData> array(ValueReader<?> elementReader) {
  return new ArrayReader(elementReader);
}