ceph/src/arrow/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/JdbcToArrowTestHelper.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.adapter.jdbc;
  19
  20 import static org.junit.Assert.assertArrayEquals;
  21 import static org.junit.Assert.assertEquals;
  22 import static org.junit.Assert.assertNotNull;
  23 import static org.junit.Assert.assertTrue;
  24
  25 import java.math.BigDecimal;
  26 import java.nio.charset.Charset;
  27 import java.sql.ResultSetMetaData;
  28 import java.sql.SQLException;
  29 import java.util.List;
  30 import java.util.Map;
  31
  32 import org.apache.arrow.vector.BaseValueVector;
  33 import org.apache.arrow.vector.BigIntVector;
  34 import org.apache.arrow.vector.BitVector;
  35 import org.apache.arrow.vector.DateDayVector;
  36 import org.apache.arrow.vector.DecimalVector;
  37 import org.apache.arrow.vector.Float4Vector;
  38 import org.apache.arrow.vector.Float8Vector;
  39 import org.apache.arrow.vector.IntVector;
  40 import org.apache.arrow.vector.NullVector;
  41 import org.apache.arrow.vector.SmallIntVector;
  42 import org.apache.arrow.vector.TimeMilliVector;
  43 import org.apache.arrow.vector.TimeStampVector;
  44 import org.apache.arrow.vector.TinyIntVector;
  45 import org.apache.arrow.vector.VarBinaryVector;
  46 import org.apache.arrow.vector.VarCharVector;
  47 import org.apache.arrow.vector.VectorSchemaRoot;
  48 import org.apache.arrow.vector.types.pojo.Field;
  49 import org.apache.arrow.vector.types.pojo.Schema;
  50
  51 /**
  52  * This is a Helper class which has functionalities to read and assert the values from the given FieldVector object.
  53  */
  54 public class JdbcToArrowTestHelper {
  55
  56   public static void assertIntVectorValues(IntVector intVector, int rowCount, Integer[] values) {
  57     assertEquals(rowCount, intVector.getValueCount());
  58
  59     for (int j = 0; j < intVector.getValueCount(); j++) {
  60       if (values[j] == null) {
  61         assertTrue(intVector.isNull(j));
  62       } else {
  63         assertEquals(values[j].intValue(), intVector.get(j));
  64       }
  65     }
  66   }
  67
  68   public static void assertBooleanVectorValues(BitVector bitVector, int rowCount, Boolean[] values) {
  69     assertEquals(rowCount, bitVector.getValueCount());
  70
  71     for (int j = 0; j < bitVector.getValueCount(); j++) {
  72       if (values[j] == null) {
  73         assertTrue(bitVector.isNull(j));
  74       } else {
  75         assertEquals(values[j].booleanValue(), bitVector.get(j) == 1);
  76       }
  77     }
  78   }
  79
  80   public static void assertBitVectorValues(BitVector bitVector, int rowCount, Integer[] values) {
  81     assertEquals(rowCount, bitVector.getValueCount());
  82
  83     for (int j = 0; j < bitVector.getValueCount(); j++) {
  84       if (values[j] == null) {
  85         assertTrue(bitVector.isNull(j));
  86       } else {
  87         assertEquals(values[j].intValue(), bitVector.get(j));
  88       }
  89     }
  90   }
  91
  92   public static void assertTinyIntVectorValues(TinyIntVector tinyIntVector, int rowCount, Integer[] values) {
  93     assertEquals(rowCount, tinyIntVector.getValueCount());
  94
  95     for (int j = 0; j < tinyIntVector.getValueCount(); j++) {
  96       if (values[j] == null) {
  97         assertTrue(tinyIntVector.isNull(j));
  98       } else {
  99         assertEquals(values[j].intValue(), tinyIntVector.get(j));
 100       }
 101     }
 102   }
 103
 104   public static void assertSmallIntVectorValues(SmallIntVector smallIntVector, int rowCount, Integer[] values) {
 105     assertEquals(rowCount, smallIntVector.getValueCount());
 106
 107     for (int j = 0; j < smallIntVector.getValueCount(); j++) {
 108       if (values[j] == null) {
 109         assertTrue(smallIntVector.isNull(j));
 110       } else {
 111         assertEquals(values[j].intValue(), smallIntVector.get(j));
 112       }
 113     }
 114   }
 115
 116   public static void assertBigIntVectorValues(BigIntVector bigIntVector, int rowCount, Long[] values) {
 117     assertEquals(rowCount, bigIntVector.getValueCount());
 118
 119     for (int j = 0; j < bigIntVector.getValueCount(); j++) {
 120       if (values[j] == null) {
 121         assertTrue(bigIntVector.isNull(j));
 122       } else {
 123         assertEquals(values[j].longValue(), bigIntVector.get(j));
 124       }
 125     }
 126   }
 127
 128   public static void assertDecimalVectorValues(DecimalVector decimalVector, int rowCount, BigDecimal[] values) {
 129     assertEquals(rowCount, decimalVector.getValueCount());
 130
 131     for (int j = 0; j < decimalVector.getValueCount(); j++) {
 132       if (values[j] == null) {
 133         assertTrue(decimalVector.isNull(j));
 134       } else {
 135         assertEquals(values[j].doubleValue(), decimalVector.getObject(j).doubleValue(), 0);
 136       }
 137     }
 138   }
 139
 140   public static void assertFloat8VectorValues(Float8Vector float8Vector, int rowCount, Double[] values) {
 141     assertEquals(rowCount, float8Vector.getValueCount());
 142
 143     for (int j = 0; j < float8Vector.getValueCount(); j++) {
 144       if (values[j] == null) {
 145         assertTrue(float8Vector.isNull(j));
 146       } else {
 147         assertEquals(values[j], float8Vector.get(j), 0.01);
 148       }
 149     }
 150   }
 151
 152   public static void assertFloat4VectorValues(Float4Vector float4Vector, int rowCount, Float[] values) {
 153     assertEquals(rowCount, float4Vector.getValueCount());
 154
 155     for (int j = 0; j < float4Vector.getValueCount(); j++) {
 156       if (values[j] == null) {
 157         assertTrue(float4Vector.isNull(j));
 158       } else {
 159         assertEquals(values[j], float4Vector.get(j), 0.01);
 160       }
 161     }
 162   }
 163
 164   public static void assertTimeVectorValues(TimeMilliVector timeMilliVector, int rowCount, Long[] values) {
 165     assertEquals(rowCount, timeMilliVector.getValueCount());
 166
 167     for (int j = 0; j < timeMilliVector.getValueCount(); j++) {
 168       if (values[j] == null) {
 169         assertTrue(timeMilliVector.isNull(j));
 170       } else {
 171         assertEquals(values[j].longValue(), timeMilliVector.get(j));
 172       }
 173     }
 174   }
 175
 176   public static void assertDateVectorValues(DateDayVector dateDayVector, int rowCount, Integer[] values) {
 177     assertEquals(rowCount, dateDayVector.getValueCount());
 178
 179     for (int j = 0; j < dateDayVector.getValueCount(); j++) {
 180       if (values[j] == null) {
 181         assertTrue(dateDayVector.isNull(j));
 182       } else {
 183         assertEquals(values[j].longValue(), dateDayVector.get(j));
 184       }
 185     }
 186   }
 187
 188   public static void assertTimeStampVectorValues(TimeStampVector timeStampVector, int rowCount, Long[] values) {
 189     assertEquals(rowCount, timeStampVector.getValueCount());
 190
 191     for (int j = 0; j < timeStampVector.getValueCount(); j++) {
 192       if (values[j] == null) {
 193         assertTrue(timeStampVector.isNull(j));
 194       } else {
 195         assertEquals(values[j].longValue(), timeStampVector.get(j));
 196       }
 197     }
 198   }
 199
 200   public static void assertVarBinaryVectorValues(VarBinaryVector varBinaryVector, int rowCount, byte[][] values) {
 201     assertEquals(rowCount, varBinaryVector.getValueCount());
 202
 203     for (int j = 0; j < varBinaryVector.getValueCount(); j++) {
 204       if (values[j] == null) {
 205         assertTrue(varBinaryVector.isNull(j));
 206       } else {
 207         assertArrayEquals(values[j], varBinaryVector.get(j));
 208       }
 209     }
 210   }
 211
 212   public static void assertVarcharVectorValues(VarCharVector varCharVector, int rowCount, byte[][] values) {
 213     assertEquals(rowCount, varCharVector.getValueCount());
 214
 215     for (int j = 0; j < varCharVector.getValueCount(); j++) {
 216       if (values[j] == null) {
 217         assertTrue(varCharVector.isNull(j));
 218       } else {
 219         assertArrayEquals(values[j], varCharVector.get(j));
 220       }
 221     }
 222   }
 223
 224   public static void assertNullVectorValues(NullVector vector, int rowCount) {
 225     assertEquals(rowCount, vector.getValueCount());
 226   }
 227
 228   public static void assertNullValues(BaseValueVector vector, int rowCount) {
 229     assertEquals(rowCount, vector.getValueCount());
 230
 231     for (int j = 0; j < vector.getValueCount(); j++) {
 232       assertTrue(vector.isNull(j));
 233     }
 234   }
 235
 236   public static void assertFieldMetadataIsEmpty(VectorSchemaRoot schema) {
 237     assertNotNull(schema);
 238     assertNotNull(schema.getSchema());
 239     assertNotNull(schema.getSchema().getFields());
 240
 241     for (Field field : schema.getSchema().getFields()) {
 242       assertNotNull(field.getMetadata());
 243       assertEquals(0, field.getMetadata().size());
 244     }
 245   }
 246
 247   public static void assertFieldMetadataMatchesResultSetMetadata(ResultSetMetaData rsmd, Schema schema)
 248       throws SQLException {
 249     assertNotNull(schema);
 250     assertNotNull(schema.getFields());
 251     assertNotNull(rsmd);
 252
 253     List<Field> fields = schema.getFields();
 254
 255     assertEquals(rsmd.getColumnCount(), fields.size());
 256
 257     // Vector columns are created in the same order as ResultSet columns.
 258     for (int i = 1; i <= rsmd.getColumnCount(); ++i) {
 259       Map<String, String> metadata = fields.get(i - 1).getMetadata();
 260
 261       assertNotNull(metadata);
 262       assertEquals(4, metadata.size());
 263
 264       assertEquals(rsmd.getCatalogName(i), metadata.get(Constants.SQL_CATALOG_NAME_KEY));
 265       assertEquals(rsmd.getTableName(i), metadata.get(Constants.SQL_TABLE_NAME_KEY));
 266       assertEquals(rsmd.getColumnLabel(i), metadata.get(Constants.SQL_COLUMN_NAME_KEY));
 267       assertEquals(rsmd.getColumnTypeName(i), metadata.get(Constants.SQL_TYPE_KEY));
 268     }
 269   }
 270
 271   public static byte[] hexStringToByteArray(String s) {
 272     int len = s.length();
 273     byte[] data = new byte[len / 2];
 274     for (int i = 0; i < len; i += 2) {
 275       data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4) +
 276               Character.digit(s.charAt(i + 1), 16));
 277     }
 278     return data;
 279   }
 280
 281   public static Integer[] getIntValues(String[] values, String dataType) {
 282     String[] dataArr = getValues(values, dataType);
 283     Integer[] valueArr = new Integer[dataArr.length];
 284     int i = 0;
 285     for (String data : dataArr) {
 286       valueArr[i++] = "null".equals(data.trim()) ? null : Integer.parseInt(data);
 287     }
 288     return valueArr;
 289   }
 290
 291   public static Boolean[] getBooleanValues(String[] values, String dataType) {
 292     String[] dataArr = getValues(values, dataType);
 293     Boolean[] valueArr = new Boolean[dataArr.length];
 294     int i = 0;
 295     for (String data : dataArr) {
 296       valueArr[i++] = "null".equals(data.trim()) ? null : data.trim().equals("1");
 297     }
 298     return valueArr;
 299   }
 300
 301   public static BigDecimal[] getDecimalValues(String[] values, String dataType) {
 302     String[] dataArr = getValues(values, dataType);
 303     BigDecimal[] valueArr = new BigDecimal[dataArr.length];
 304     int i = 0;
 305     for (String data : dataArr) {
 306       valueArr[i++] = "null".equals(data.trim()) ? null : new BigDecimal(data);
 307     }
 308     return valueArr;
 309   }
 310
 311   public static Double[] getDoubleValues(String[] values, String dataType) {
 312     String[] dataArr = getValues(values, dataType);
 313     Double[] valueArr = new Double[dataArr.length];
 314     int i = 0;
 315     for (String data : dataArr) {
 316       valueArr[i++] = "null".equals(data.trim()) ? null : Double.parseDouble(data);
 317     }
 318     return valueArr;
 319   }
 320
 321   public static Float[] getFloatValues(String[] values, String dataType) {
 322     String[] dataArr = getValues(values, dataType);
 323     Float[] valueArr = new Float[dataArr.length];
 324     int i = 0;
 325     for (String data : dataArr) {
 326       valueArr[i++] = "null".equals(data.trim()) ? null : Float.parseFloat(data);
 327     }
 328     return valueArr;
 329   }
 330
 331   public static Long[] getLongValues(String[] values, String dataType) {
 332     String[] dataArr = getValues(values, dataType);
 333     Long[] valueArr = new Long[dataArr.length];
 334     int i = 0;
 335     for (String data : dataArr) {
 336       valueArr[i++] = "null".equals(data.trim()) ? null : Long.parseLong(data);
 337     }
 338     return valueArr;
 339   }
 340
 341   public static byte[][] getCharArray(String[] values, String dataType) {
 342     String[] dataArr = getValues(values, dataType);
 343     byte[][] valueArr = new byte[dataArr.length][];
 344     int i = 0;
 345     for (String data : dataArr) {
 346       valueArr[i++] = "null".equals(data.trim()) ? null : data.trim().getBytes();
 347     }
 348     return valueArr;
 349   }
 350
 351   public static byte[][] getCharArrayWithCharSet(String[] values, String dataType, Charset charSet) {
 352     String[] dataArr = getValues(values, dataType);
 353     byte[][] valueArr = new byte[dataArr.length][];
 354     int i = 0;
 355     for (String data : dataArr) {
 356       valueArr[i++] = "null".equals(data.trim()) ? null : data.trim().getBytes(charSet);
 357     }
 358     return valueArr;
 359   }
 360
 361   public static byte[][] getBinaryValues(String[] values, String dataType) {
 362     String[] dataArr = getValues(values, dataType);
 363     byte[][] valueArr = new byte[dataArr.length][];
 364     int i = 0;
 365     for (String data : dataArr) {
 366       valueArr[i++] = "null".equals(data.trim()) ? null : hexStringToByteArray(data.trim());
 367     }
 368     return valueArr;
 369   }
 370
 371   public static String[] getValues(String[] values, String dataType) {
 372     String value = "";
 373     for (String val : values) {
 374       if (val.startsWith(dataType)) {
 375         value = val.split("=")[1];
 376         break;
 377       }
 378     }
 379     return value.split(",");
 380   }
 381 }