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

The following examples show how to use org.apache.spark.sql.catalyst.util.ArrayData#numElements() . 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 File: SparkOrcWriter.java    From iceberg with 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 File: SparkOrcWriter.java    From iceberg with 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 3
Source File: SparkOrcWriter.java    From iceberg with 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 4
Source File: SparkValueWriters.java    From iceberg with 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 5
Source File: SparkOrcWriter.java    From iceberg with 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 6
Source File: TestHelpers.java    From iceberg with 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 7
Source File: SparkValueWriters.java    From iceberg with 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 8
Source File: SparkParquetWriters.java    From iceberg with Apache License 2.0 4 votes vote down vote up
private ElementIterator(ArrayData list) {
  this.list = list;
  size = list.numElements();
  index = 0;
}
 
Example 9
Source File: TestHelpers.java    From iceberg with Apache License 2.0 4 votes vote down vote up
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);
      }
    }
  }
}
 
Example 10
Source File: SparkOrcReader.java    From iceberg with Apache License 2.0 4 votes vote down vote up
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);
  }
}
 
Example 11
Source File: CodegenExamples.java    From iceberg with Apache License 2.0 4 votes vote down vote up
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;
}
 
Example 12
Source File: CodegenExamples.java    From iceberg with Apache License 2.0 4 votes vote down vote up
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;
}
 
Example 13
Source File: CodegenExamples.java    From iceberg with Apache License 2.0 4 votes vote down vote up
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;
}
 
Example 14
Source File: TestHelpers.java    From iceberg with Apache License 2.0 4 votes vote down vote up
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);
      }
    }
  }
}