/* * Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.client.core.arrow; import net.snowflake.client.TestUtil; import net.snowflake.client.core.SFException; import net.snowflake.client.jdbc.ErrorCode; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.Test; import java.math.BigDecimal; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Random; import java.util.Set; import java.util.TimeZone; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; public class TinyIntToFixedConverterTest extends BaseConverterTest { /** * allocator for arrow */ private BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); /** * Random seed */ private Random random = new Random(); @Test public void testFixedNoScale() throws SFException { final int rowCount = 1000; List<Byte> expectedValues = new ArrayList<>(); Set<Integer> nullValIndex = new HashSet<>(); for (int i = 0; i < rowCount; i++) { expectedValues.add((byte) random.nextInt(1 << 8)); } Map<String, String> customFieldMeta = new HashMap<>(); customFieldMeta.put("logicalType", "FIXED"); customFieldMeta.put("precision", "10"); customFieldMeta.put("scale", "0"); FieldType fieldType = new FieldType(true, Types.MinorType.TINYINT.getType(), null, customFieldMeta); TinyIntVector vector = new TinyIntVector("col_one", fieldType, allocator); for (int i = 0; i < rowCount; i++) { boolean isNull = random.nextBoolean(); if (isNull) { vector.setNull(i); nullValIndex.add(i); } else { vector.setSafe(i, expectedValues.get(i)); } } ArrowVectorConverter converter = new TinyIntToFixedConverter(vector, 0, this); for (int i = 0; i < rowCount; i++) { byte byteVal = converter.toByte(i); Object longObject = converter.toObject(i); // the logical type is long String byteString = converter.toString(i); if (nullValIndex.contains(i)) { assertThat(byteVal, is((byte) 0)); assertThat(longObject, is(nullValue())); assertThat(byteString, is(nullValue())); } else { assertThat(byteVal, is(expectedValues.get(i))); assertEquals(longObject, (long) expectedValues.get(i)); assertThat(byteString, is(expectedValues.get(i).toString())); } } vector.clear(); } @Test public void testFixedWithScale() throws SFException { final int rowCount = 1000; List<Byte> expectedValues = new ArrayList<>(); Set<Integer> nullValIndex = new HashSet<>(); for (int i = 0; i < rowCount; i++) { expectedValues.add((byte) random.nextInt(1 << 8)); } Map<String, String> customFieldMeta = new HashMap<>(); customFieldMeta.put("logicalType", "FIXED"); customFieldMeta.put("precision", "10"); customFieldMeta.put("scale", "1"); FieldType fieldType = new FieldType(true, Types.MinorType.TINYINT.getType(), null, customFieldMeta); TinyIntVector vector = new TinyIntVector("col_one", fieldType, allocator); for (int i = 0; i < rowCount; i++) { boolean isNull = random.nextBoolean(); if (isNull) { vector.setNull(i); nullValIndex.add(i); } else { vector.setSafe(i, expectedValues.get(i)); } } ArrowVectorConverter converter = new TinyIntToScaledFixedConverter(vector, 0, this, 1); for (int i = 0; i < rowCount; i++) { BigDecimal bigDecimalVal = converter.toBigDecimal(i); Object objectVal = converter.toObject(i); String stringVal = converter.toString(i); if (nullValIndex.contains(i)) { assertThat(bigDecimalVal, nullValue()); assertThat(objectVal, nullValue()); assertThat(stringVal, nullValue()); } else { BigDecimal expectedVal = BigDecimal.valueOf(expectedValues.get(i), 1); assertThat(bigDecimalVal, is(expectedVal)); assertThat(objectVal, is(expectedVal)); assertThat(stringVal, is(expectedVal.toString())); } } vector.clear(); } @Test public void testInvalidConversion() { // try convert to int/long/byte/short with scale > 0 Map<String, String> customFieldMeta = new HashMap<>(); customFieldMeta.put("logicalType", "FIXED"); customFieldMeta.put("precision", "10"); customFieldMeta.put("scale", "1"); FieldType fieldType = new FieldType(true, Types.MinorType.TINYINT.getType(), null, customFieldMeta); TinyIntVector vector = new TinyIntVector("col_one", fieldType, allocator); vector.setSafe(0, 200); final ArrowVectorConverter converter = new TinyIntToScaledFixedConverter(vector, 0, this, 1); final int invalidConversionErrorCode = ErrorCode.INVALID_VALUE_CONVERT.getMessageCode(); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toBoolean(0)); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toLong(0)); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toInt(0)); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toShort(0)); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toDate(0)); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toTime(0)); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toTimestamp(0, TimeZone.getDefault())); vector.clear(); } @Test public void testGetSmallerIntegralType() throws SFException { // try convert to int/long/byte/short with scale > 0 Map<String, String> customFieldMeta = new HashMap<>(); customFieldMeta.put("logicalType", "FIXED"); customFieldMeta.put("precision", "10"); customFieldMeta.put("scale", "0"); FieldType fieldType = new FieldType(true, Types.MinorType.TINYINT.getType(), null, customFieldMeta); // test value which is in range of byte, all get method should return TinyIntVector vectorBar = new TinyIntVector("col_one", fieldType, allocator); // set value which is out of range of int, but falls in long vectorBar.setSafe(0, 10); vectorBar.setSafe(1, -10); final ArrowVectorConverter converterBar = new TinyIntToFixedConverter(vectorBar, 0, this); assertThat(converterBar.toShort(0), is((short) 10)); assertThat(converterBar.toShort(1), is((short) -10)); assertThat(converterBar.toInt(0), is(10)); assertThat(converterBar.toInt(1), is(-10)); assertThat(converterBar.toLong(0), is(10L)); assertThat(converterBar.toLong(1), is(-10L)); vectorBar.clear(); } @Test public void testGetBooleanNoScale() throws SFException { Map<String, String> customFieldMeta = new HashMap<>(); customFieldMeta.put("logicalType", "FIXED"); customFieldMeta.put("precision", "10"); customFieldMeta.put("scale", "0"); FieldType fieldType = new FieldType(true, Types.MinorType.TINYINT.getType(), null, customFieldMeta); TinyIntVector vector = new TinyIntVector("col_one", fieldType, allocator); vector.setSafe(0, 0); vector.setSafe(1, 1); vector.setNull(2); vector.setSafe(3, 5); ArrowVectorConverter converter = new TinyIntToFixedConverter(vector, 0, this); assertThat(false, is(converter.toBoolean(0))); assertThat(true, is(converter.toBoolean(1))); assertThat(false, is(converter.toBoolean(2))); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toBoolean(3)); vector.close(); } @Test public void testGetBooleanWithScale() throws SFException { Map<String, String> customFieldMeta = new HashMap<>(); customFieldMeta.put("logicalType", "FIXED"); customFieldMeta.put("precision", "10"); customFieldMeta.put("scale", "3"); FieldType fieldType = new FieldType(true, Types.MinorType.TINYINT.getType(), null, customFieldMeta); TinyIntVector vector = new TinyIntVector("col_one", fieldType, allocator); vector.setSafe(0, 0); vector.setSafe(1, 1); vector.setNull(2); vector.setSafe(3, 5); final ArrowVectorConverter converter = new TinyIntToScaledFixedConverter(vector, 0, this, 3); assertThat(false, is(converter.toBoolean(0))); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toBoolean(3)); assertThat(false, is(converter.toBoolean(2))); TestUtil.assertSFException(invalidConversionErrorCode, () -> converter.toBoolean(3)); vector.close(); } }