Skip to content

Commit

Permalink
[ARROW-7073][Java] Rewrite tests with vector populator and provide st…
Browse files Browse the repository at this point in the history
…atic utility
  • Loading branch information
liyafan82 committed Feb 3, 2020
1 parent 7c13ede commit c89211a
Show file tree
Hide file tree
Showing 5 changed files with 294 additions and 159 deletions.
Expand Up @@ -37,7 +37,7 @@
/**
* Utility to append two vectors together.
*/
public class VectorAppender implements VectorVisitor<ValueVector, Void> {
class VectorAppender implements VectorVisitor<ValueVector, Void> {

/**
* The targetVector to be appended.
Expand All @@ -50,9 +50,9 @@ public class VectorAppender implements VectorVisitor<ValueVector, Void> {
* Constructs a new targetVector appender, with the given targetVector.
* @param targetVector the targetVector to be appended.
*/
public VectorAppender(ValueVector targetVector) {
VectorAppender(ValueVector targetVector) {
this.targetVector = targetVector;
typeVisitor = new TypeEqualsVisitor(targetVector);
typeVisitor = new TypeEqualsVisitor(targetVector, false, true);
}

@Override
Expand Down Expand Up @@ -194,7 +194,7 @@ public ValueVector visit(FixedSizeListVector deltaVector, Void value) {
int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount();

int targetListSize = targetListVector.getValueCount() * targetListVector.getListSize();
int deltaListSize = deltaVector.getValueCount() * deltaVector.getValueCount();
int deltaListSize = deltaVector.getValueCount() * deltaVector.getListSize();

// make sure the underlying vector has value count set
targetListVector.getDataVector().setValueCount(targetListSize);
Expand Down
@@ -0,0 +1,39 @@
/*
* 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.arrow.vector.util;

import org.apache.arrow.vector.ValueVector;

/**
* Utility to add vector values in batch.
*/
public class VectorBatchAppender {

/**
* Add value vectors in batch.
* @param targetVector the target vector.
* @param vectorsToAppend the vectors to append.
* @param <V> the vector type.
*/
public static <V extends ValueVector> void batchAppend(V targetVector, V... vectorsToAppend) {
VectorAppender appender = new VectorAppender(targetVector);
for (V delta : vectorsToAppend) {
delta.accept(appender, null);
}
}
}
Expand Up @@ -17,10 +17,14 @@

package org.apache.arrow.vector.testing;

import static org.junit.Assert.assertEquals;

import java.nio.charset.StandardCharsets;
import java.util.List;

import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.BitVectorHelper;
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.DateMilliVector;
import org.apache.arrow.vector.DecimalVector;
Expand Down Expand Up @@ -51,7 +55,12 @@
import org.apache.arrow.vector.UInt8Vector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.complex.BaseRepeatedValueVector;
import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.holders.IntervalDayHolder;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.FieldType;

/**
* Utility for populating {@link org.apache.arrow.vector.ValueVector}.
Expand Down Expand Up @@ -533,10 +542,62 @@ public static void setVector(VarCharVector vector, String... values) {
vector.allocateNewSafe();
for (int i = 0; i < length; i++) {
if (values[i] != null) {
vector.set(i, values[i].getBytes(StandardCharsets.UTF_8));
vector.setSafe(i, values[i].getBytes(StandardCharsets.UTF_8));
}
}
vector.setValueCount(length);
}

/**
* Populate values for {@link ListVector}.
*/
public static void setVector(ListVector vector, List<Integer>... values) {
Types.MinorType type = Types.MinorType.INT;
vector.addOrGetVector(FieldType.nullable(type.getType()));

IntVector dataVector = (IntVector) vector.getDataVector();
dataVector.allocateNew();

// set underlying vectors
int curPos = 0;
vector.getOffsetBuffer().setInt(0, curPos);
for (int i = 0; i < values.length; i++) {
BitVectorHelper.setBit(vector.getValidityBuffer(), i);
for (int value : values[i]) {
dataVector.set(curPos, value);
curPos += 1;
}
vector.getOffsetBuffer().setInt((i + 1) * BaseRepeatedValueVector.OFFSET_WIDTH, curPos);
}
dataVector.setValueCount(curPos);
vector.setLastSet(values.length - 1);
vector.setValueCount(values.length);
}

/**
* Populate values for {@link FixedSizeListVector}.
*/
public static void setVector(FixedSizeListVector vector, List<Integer>... values) {
for (int i = 0; i < values.length; i++) {
assertEquals(vector.getListSize(), values[i].size());
}

Types.MinorType type = Types.MinorType.INT;
vector.addOrGetVector(FieldType.nullable(type.getType()));

IntVector dataVector = (IntVector) vector.getDataVector();
dataVector.allocateNew();

// set underlying vectors
int curPos = 0;
for (int i = 0; i < values.length; i++) {
BitVectorHelper.setBit(vector.getValidityBuffer(), i);
for (int value : values[i]) {
dataVector.set(curPos, value);
curPos += 1;
}
}
dataVector.setValueCount(curPos);
vector.setValueCount(values.length);
}
}

0 comments on commit c89211a

Please sign in to comment.