Java Code Examples for org.apache.spark.sql.catalyst.util.ArrayData#numElements()

@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);
    }
  }
}
 

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);
    }
  }
}
 

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
@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);
    }
  }
}
 

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));
  }
}
 

@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();
}
 

private ElementIterator(ArrayData list) {
  this.list = list;
  size = list.numElements();
  index = 0;
}
 

private static void assertEqualsLists(String prefix, Types.ListType type,
                                      ArrayData expected, List actual) {
  if (expected == null || actual == null) {
    Assert.assertEquals(prefix, expected, actual);
  } else {
    Assert.assertEquals(prefix + " length", expected.numElements(), actual.size());
    Type childType = type.elementType();
    for (int e = 0; e < expected.numElements(); ++e) {
      switch (childType.typeId()) {
        case BOOLEAN:
        case INTEGER:
        case LONG:
        case FLOAT:
        case DOUBLE:
        case STRING:
        case DECIMAL:
        case DATE:
        case TIMESTAMP:
          Assert.assertEquals(prefix + ".elem " + e + " - " + childType,
              getValue(expected, e, childType),
              actual.get(e));
          break;
        case UUID:
        case FIXED:
        case BINARY:
          assertEqualBytes(prefix + ".elem " + e,
              (byte[]) getValue(expected, e, childType),
              (byte[]) actual.get(e));
          break;
        case STRUCT: {
          Types.StructType st = (Types.StructType) childType;
          assertEquals(prefix + ".elem " + e, st,
              expected.getStruct(e, st.fields().size()), (Row) actual.get(e));
          break;
        }
        case LIST:
          assertEqualsLists(prefix + ".elem " + e, childType.asListType(),
              expected.getArray(e),
              toList((Seq<?>) actual.get(e)));
          break;
        case MAP:
          assertEqualsMaps(prefix + ".elem " + e, childType.asMapType(),
              expected.getMap(e), toJavaMap((scala.collection.Map<?, ?>) actual.get(e)));
          break;
        default:
          throw new IllegalArgumentException("Unhandled type " + childType);
      }
    }
  }
}
 

private static void printRow(SpecializedGetters row, int ord, TypeDescription schema) {
  switch (schema.getCategory()) {
    case BOOLEAN:
      System.out.print(row.getBoolean(ord));
      break;
    case BYTE:
      System.out.print(row.getByte(ord));
      break;
    case SHORT:
      System.out.print(row.getShort(ord));
      break;
    case INT:
      System.out.print(row.getInt(ord));
      break;
    case LONG:
      System.out.print(row.getLong(ord));
      break;
    case FLOAT:
      System.out.print(row.getFloat(ord));
      break;
    case DOUBLE:
      System.out.print(row.getDouble(ord));
      break;
    case CHAR:
    case VARCHAR:
    case STRING:
      System.out.print("\"" + row.getUTF8String(ord) + "\"");
      break;
    case BINARY: {
      byte[] bin = row.getBinary(ord);
      if (bin == null) {
        System.out.print("null");
      } else {
        System.out.print("[");
        for (int i = 0; i < bin.length; ++i) {
          if (i != 0) {
            System.out.print(", ");
          }
          int v = bin[i] & 0xff;
          if (v < 16) {
            System.out.print("0" + Integer.toHexString(v));
          } else {
            System.out.print(Integer.toHexString(v));
          }
        }
        System.out.print("]");
      }
      break;
    }
    case DECIMAL:
      System.out.print(row.getDecimal(ord, schema.getPrecision(), schema.getScale()));
      break;
    case DATE:
      System.out.print("\"" + new DateWritable(row.getInt(ord)) + "\"");
      break;
    case TIMESTAMP:
      System.out.print("\"" + new Timestamp(row.getLong(ord)) + "\"");
      break;
    case STRUCT:
      printRow(row.getStruct(ord, schema.getChildren().size()), schema);
      break;
    case LIST: {
      TypeDescription child = schema.getChildren().get(0);
      System.out.print("[");
      ArrayData list = row.getArray(ord);
      for(int e=0; e < list.numElements(); ++e) {
        if (e != 0) {
          System.out.print(", ");
        }
        printRow(list, e, child);
      }
      System.out.print("]");
      break;
    }
    case MAP: {
      TypeDescription keyType = schema.getChildren().get(0);
      TypeDescription valueType = schema.getChildren().get(1);
      MapData map = row.getMap(ord);
      ArrayData keys = map.keyArray();
      ArrayData values = map.valueArray();
      System.out.print("[");
      for(int e=0; e < map.numElements(); ++e) {
        if (e != 0) {
          System.out.print(", ");
        }
        printRow(keys, e, keyType);
        System.out.print(": ");
        printRow(values, e, valueType);
      }
      System.out.print("]");
      break;
    }
    default:
      throw new IllegalArgumentException("Unhandled type " + schema);
  }
}
 

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;
}
 

private static void assertEqualsLists(String prefix, Types.ListType type,
                                      ArrayData expected, List actual) {
  if (expected == null || actual == null) {
    Assert.assertEquals(prefix, expected, actual);
  } else {
    Assert.assertEquals(prefix + " length", expected.numElements(), actual.size());
    Type childType = type.elementType();
    for (int e = 0; e < expected.numElements(); ++e) {
      switch (childType.typeId()) {
        case BOOLEAN:
        case INTEGER:
        case LONG:
        case FLOAT:
        case DOUBLE:
        case STRING:
        case DECIMAL:
        case DATE:
        case TIMESTAMP:
          Assert.assertEquals(prefix + ".elem " + e + " - " + childType,
              getValue(expected, e, childType),
              actual.get(e));
          break;
        case UUID:
        case FIXED:
        case BINARY:
          assertEqualBytes(prefix + ".elem " + e,
              (byte[]) getValue(expected, e, childType),
              (byte[]) actual.get(e));
          break;
        case STRUCT: {
          Types.StructType st = (Types.StructType) childType;
          assertEquals(prefix + ".elem " + e, st,
              expected.getStruct(e, st.fields().size()), (Row) actual.get(e));
          break;
        }
        case LIST:
          assertEqualsLists(prefix + ".elem " + e, childType.asListType(),
              expected.getArray(e),
              toList((Seq<?>) actual.get(e)));
          break;
        case MAP:
          assertEqualsMaps(prefix + ".elem " + e, childType.asMapType(),
              expected.getMap(e), toJavaMap((scala.collection.Map<?, ?>) actual.get(e)));
          break;
        default:
          throw new IllegalArgumentException("Unhandled type " + childType);
      }
    }
  }
}