Improve BigInteger handling in Parquet files (Closes #3125) (#3126)

* Allow unrecognized values in metadata JSON to be ignored.  Allow BigIntegers to be encoded as Deciomal types with a precision/scale of 1, instead of requiring a serializable codec

* Improve BigInteger handling

* Applied spotless

* Allow overriding default codecs of BigInteger and BigDecimal by ones set by users in ParquetInstructions.  Fix unit test

* add comment

* Missed a file

* Comments

* Spotless

* DO a table walk to determine precision, use a default valid value if everything is null

* spotless

* Renames

* moar spotless

Author: abaranec

engine/table/src/main/java/io/deephaven/engine/table/impl/CodecLookup.java

@@ -17,6 +17,7 @@
 
 import java.io.Externalizable;
 import java.math.BigDecimal;
+import java.math.BigInteger;
 
 /**
  * Utility class to concentrate {@link ObjectCodec} lookups.
@@ -76,7 +77,11 @@ private static boolean noCodecRequired(@NotNull final Class<?> dataType) {
                 // appropriate precision and scale calculated from column data,
                 // unless the user explicitly requested something else
                 // via instructions.
-                dataType == BigDecimal.class;
+                dataType == BigDecimal.class ||
+
+                // BigIntegers can be encoded as a DecimalLogicalType using a precision of 1 and scale of 0, which lets
+                // them be read by other parquet tools.
+                dataType == BigInteger.class;
     }
 
     /**

engine/table/src/main/java/io/deephaven/engine/util/BigDecimalUtils.java

@@ -24,6 +24,8 @@
  * refreshing tables, we need the user to tell us.
  */
 public class BigDecimalUtils {
+    private static final PrecisionAndScale EMPTY_TABLE_PRECISION_AND_SCALE = new PrecisionAndScale(1, 1);
+    private static final int TARGET_CHUNK_SIZE = 4096;
     public static final int INVALID_PRECISION_OR_SCALE = -1;
 
     /**
@@ -53,7 +55,8 @@ public static PrecisionAndScale computePrecisionAndScale(
     }
 
     /**
-     * Compute an overall precision and scale that would fit all existing values in a column source.
+     * Compute an overall precision and scale that would fit all existing values in a column source. Note that this
+     * requires a full table scan to ensure the correct values are determined.
      *
      * @param rowSet The rowset for the provided column
      * @param source a {@code ColumnSource} of {@code BigDecimal} type
@@ -62,29 +65,45 @@ public static PrecisionAndScale computePrecisionAndScale(
     public static PrecisionAndScale computePrecisionAndScale(
             final RowSet rowSet,
             final ColumnSource<BigDecimal> source) {
-        final int sz = 4096;
-        // we first compute max(precision - scale) and max(scale), which corresponds to
-        // max(digits left of the decimal point), max(digits right of the decimal point).
-        // Then we convert to (precision, scale) before returning.
-        int maxPrecisionMinusScale = 0;
-        int maxScale = 0;
-        try (final ChunkSource.GetContext context = source.makeGetContext(sz);
+        if (rowSet.isEmpty()) {
+            return EMPTY_TABLE_PRECISION_AND_SCALE;
+        }
+
+        // We will walk the entire table to determine the max(precision - scale) and
+        // max(scale), which corresponds to max(digits left of the decimal point), max(digits right of the decimal
+        // point). Then we convert to (precision, scale) before returning.
+        int maxPrecisionMinusScale = -1;
+        int maxScale = -1;
+        try (final ChunkSource.GetContext context = source.makeGetContext(TARGET_CHUNK_SIZE);
                 final RowSequence.Iterator it = rowSet.getRowSequenceIterator()) {
-            final RowSequence rowSeq = it.getNextRowSequenceWithLength(sz);
-            final ObjectChunk<BigDecimal, ? extends Values> chunk = source.getChunk(context, rowSeq).asObjectChunk();
-            for (int i = 0; i < chunk.size(); ++i) {
-                final BigDecimal x = chunk.get(i);
-                final int precision = x.precision();
-                final int scale = x.scale();
-                final int precisionMinusScale = precision - scale;
-                if (precisionMinusScale > maxPrecisionMinusScale) {
-                    maxPrecisionMinusScale = precisionMinusScale;
-                }
-                if (scale > maxScale) {
-                    maxScale = scale;
+            while (it.hasMore()) {
+                final RowSequence rowSeq = it.getNextRowSequenceWithLength(TARGET_CHUNK_SIZE);
+                final ObjectChunk<BigDecimal, ? extends Values> chunk =
+                        source.getChunk(context, rowSeq).asObjectChunk();
+                for (int i = 0; i < chunk.size(); ++i) {
+                    final BigDecimal x = chunk.get(i);
+                    if (x == null) {
+                        continue;
+                    }
+
+                    final int precision = x.precision();
+                    final int scale = x.scale();
+                    final int precisionMinusScale = precision - scale;
+                    if (precisionMinusScale > maxPrecisionMinusScale) {
+                        maxPrecisionMinusScale = precisionMinusScale;
+                    }
+                    if (scale > maxScale) {
+                        maxScale = scale;
+                    }
                 }
             }
         }
+
+        // If these are < 0, then every value we visited was null
+        if (maxPrecisionMinusScale < 0 && maxScale < 0) {
+            return EMPTY_TABLE_PRECISION_AND_SCALE;
+        }
+
         return new PrecisionAndScale(maxPrecisionMinusScale + maxScale, maxScale);
     }
 

engine/table/src/test/java/io/deephaven/engine/table/impl/TestCodecColumns.java

@@ -8,6 +8,7 @@
 import io.deephaven.engine.table.Table;
 import io.deephaven.engine.table.TableDefinition;
 import io.deephaven.engine.testutil.TstUtils;
+import io.deephaven.parquet.table.BigIntegerParquetBytesCodec;
 import io.deephaven.parquet.table.ParquetTools;
 import io.deephaven.engine.util.TableTools;
 import io.deephaven.parquet.table.ParquetInstructions;
@@ -51,16 +52,18 @@ public class TestCodecColumns {
                 ColumnDefinition.fromGenericType("VWBA", byte[].class, byte.class);
         writeBuilder.addColumnCodec("VWBA", SimpleByteArrayCodec.class.getName());
         readBuilder.addColumnCodec("VWBA", SimpleByteArrayCodec.class.getName());
+
         VARIABLE_WIDTH_COLUMN_DEFINITION_2 = ColumnDefinition.fromGenericType("VWCD", ArrayTuple.class);
         readBuilder.addColumnCodec("VWCD", ExternalizableCodec.class.getName(), ArrayTuple.class.getName());
         FIXED_WIDTH_BYTE_ARRAY_COLUMN_DEFINITION = ColumnDefinition.fromGenericType("FWBA", byte[].class, byte.class);
         writeBuilder.addColumnCodec("FWBA", SimpleByteArrayCodec.class.getName(), "9");
         readBuilder.addColumnCodec("FWBA", SimpleByteArrayCodec.class.getName(), "9");
+
         VARIABLE_WIDTH_BIG_INTEGER_COLUMN_DEFINITION = ColumnDefinition.fromGenericType("VWBI", BigInteger.class);
-        writeBuilder.addColumnCodec("VWBI", BigIntegerCodec.class.getName());
-        readBuilder.addColumnCodec("VWBI", BigIntegerCodec.class.getName());
+        writeBuilder.addColumnCodec("VWBI", SerializableCodec.class.getName());
+        readBuilder.addColumnCodec("VWBI", SerializableCodec.class.getName());
+
         VARIABLE_WIDTH_BIG_INTEGER_COLUMN_DEFINITION_S = ColumnDefinition.fromGenericType("VWBIS", BigInteger.class);
-        readBuilder.addColumnCodec("VWBIS", SerializableCodec.class.getName());
         expectedReadInstructions = readBuilder.build();
         writeInstructions = writeBuilder.build();
     }

extensions/parquet/table/src/main/java/io/deephaven/parquet/table/BigIntegerParquetBytesCodec.java

@@ -0,0 +1,90 @@
+package io.deephaven.parquet.table;
+
+import io.deephaven.datastructures.util.CollectionUtil;
+import io.deephaven.util.codec.ObjectCodec;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.math.BigInteger;
+import java.nio.ByteBuffer;
+
+public class BigIntegerParquetBytesCodec implements ObjectCodec<BigInteger> {
+    private final int encodedSizeInBytes;
+    private final byte[] nullBytes;
+
+    /**
+     *
+     * @param encodedSizeInBytes encoded size in bytes, if fixed size, or -1 if variable size. note that according to
+     *        the parquet spec, the minimum number of bytes required to represent the unscaled value should be used for
+     *        a variable sized (binary) encoding; in any case, the maximum encoded bytes is implicitly limited by
+     *        precision.
+     */
+    public BigIntegerParquetBytesCodec(final int encodedSizeInBytes) {
+        this.encodedSizeInBytes = encodedSizeInBytes;
+        if (encodedSizeInBytes > 0) {
+            nullBytes = new byte[encodedSizeInBytes];
+            for (int i = 0; i < encodedSizeInBytes; ++i) {
+                nullBytes[i] = (byte) 0xff;
+            }
+        } else {
+            nullBytes = CollectionUtil.ZERO_LENGTH_BYTE_ARRAY;
+        }
+    }
+
+    // Given how parquet encoding works for nulls, the actual value provided for a null is irrelevant.
+    @Override
+    public boolean isNullable() {
+        return true;
+    }
+
+    @Override
+    public int getPrecision() {
+        return 0;
+    }
+
+    @Override
+    public int getScale() {
+        return 1;
+    }
+
+    @Override
+    public int expectedObjectWidth() {
+        return encodedSizeInBytes <= 0 ? VARIABLE_WIDTH_SENTINEL : encodedSizeInBytes;
+    }
+
+    @NotNull
+    @Override
+    public byte[] encode(@Nullable final BigInteger input) {
+        if (input == null) {
+            return nullBytes;
+        }
+
+        return input.toByteArray();
+    }
+
+    @Nullable
+    @Override
+    public BigInteger decode(@NotNull final byte[] input, final int offset, final int length) {
+        if (length <= 0) {
+            return null;
+        }
+
+        if (length == encodedSizeInBytes) {
+            boolean allPreviousBitsSet = true;
+            for (int i = 0; i < encodedSizeInBytes; ++i) {
+                if (input[offset + i] != (byte) 0xff) {
+                    allPreviousBitsSet = false;
+                    break;
+                }
+            }
+            if (allPreviousBitsSet) {
+                return null;
+            }
+        }
+
+        final ByteBuffer buffer = ByteBuffer.wrap(input, offset, length);
+        final byte[] unscaledValueBytes = new byte[length];
+        buffer.get(unscaledValueBytes);
+        return new BigInteger(unscaledValueBytes);
+    }
+}

extensions/parquet/table/src/main/java/io/deephaven/parquet/table/ParquetSchemaReader.java

@@ -25,6 +25,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.math.BigInteger;
 import java.util.*;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
@@ -327,6 +328,11 @@ public Optional<Class<?>> visit(final LogicalTypeAnnotation.EnumLogicalTypeAnnot
             @Override
             public Optional<Class<?>> visit(
                     final LogicalTypeAnnotation.DecimalLogicalTypeAnnotation decimalLogicalType) {
+                // This pair of values (precision=1, scale=0) is set at write tiem as a marker so that we can recover
+                // the fact that the type is a BigInteger, not a BigDecimal when the fies are read.
+                if (decimalLogicalType.getPrecision() == 1 && decimalLogicalType.getScale() == 0) {
+                    return Optional.of(BigInteger.class);
+                }
                 return Optional.of(java.math.BigDecimal.class);
             }
 

extensions/parquet/table/src/main/java/io/deephaven/parquet/table/ParquetTableWriter.java

@@ -53,6 +53,7 @@
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.math.BigDecimal;
+import java.math.BigInteger;
 import java.nio.*;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -775,14 +776,23 @@ private static <DATA_TYPE> TransferObject<?> getDestinationBuffer(
             return new ByteTransfer(columnSource, maxValuesPerPage);
         } else if (String.class.equals(columnType)) {
             return new StringTransfer(columnSource, maxValuesPerPage);
-        } else if (BigDecimal.class.equals(columnType)) {
-            // noinspection unchecked
-            final ColumnSource<BigDecimal> bigDecimalColumnSource = (ColumnSource<BigDecimal>) columnSource;
-            final BigDecimalUtils.PrecisionAndScale precisionAndScale = TypeInfos.getPrecisionAndScale(
-                    computedCache, columnDefinition.getName(), tableRowSet, () -> bigDecimalColumnSource);
-            final ObjectCodec<BigDecimal> codec = new BigDecimalParquetBytesCodec(
-                    precisionAndScale.precision, precisionAndScale.scale, -1);
-            return new CodecTransfer<>(bigDecimalColumnSource, codec, maxValuesPerPage);
+        }
+
+        // If there's an explicit codec, we should disregard the defaults for these CodecLookup#lookup() will properly
+        // select the codec assigned by the instructions so we only need to check and redirect once.
+        if (!CodecLookup.explicitCodecPresent(instructions.getCodecName(columnDefinition.getName()))) {
+            if (BigDecimal.class.equals(columnType)) {
+                // noinspection unchecked
+                final ColumnSource<BigDecimal> bigDecimalColumnSource = (ColumnSource<BigDecimal>) columnSource;
+                final BigDecimalUtils.PrecisionAndScale precisionAndScale = TypeInfos.getPrecisionAndScale(
+                        computedCache, columnDefinition.getName(), tableRowSet, () -> bigDecimalColumnSource);
+                final ObjectCodec<BigDecimal> codec = new BigDecimalParquetBytesCodec(
+                        precisionAndScale.precision, precisionAndScale.scale, -1);
+                return new CodecTransfer<>(bigDecimalColumnSource, codec, maxValuesPerPage);
+            } else if (BigInteger.class.equals(columnType)) {
+                return new CodecTransfer<>((ColumnSource<BigInteger>) columnSource, new BigIntegerParquetBytesCodec(-1),
+                        maxValuesPerPage);
+            }
         }
 
         final ObjectCodec<? super DATA_TYPE> codec = CodecLookup.lookup(columnDefinition, instructions);

extensions/parquet/table/src/main/java/io/deephaven/parquet/table/TypeInfos.java

@@ -23,6 +23,7 @@
 
 import java.io.Externalizable;
 import java.math.BigDecimal;
+import java.math.BigInteger;
 import java.util.*;
 import java.util.function.Supplier;
 
@@ -34,7 +35,6 @@
  * and the data translation.
  */
 class TypeInfos {
-
     private static final TypeInfo[] TYPE_INFOS = new TypeInfo[] {
             IntType.INSTANCE,
             LongType.INSTANCE,
@@ -45,15 +45,16 @@ class TypeInfos {
             CharType.INSTANCE,
             ByteType.INSTANCE,
             StringType.INSTANCE,
-            DateTimeType.INSTANCE
+            DateTimeType.INSTANCE,
+            BigIntegerType.INSTANCE
     };
 
     private static final Map<Class<?>, TypeInfo> BY_CLASS;
 
     static {
         final Map<Class<?>, TypeInfo> fa = new HashMap<>();
-        for (TypeInfo typeInfo : TYPE_INFOS) {
-            for (Class<?> type : typeInfo.getTypes()) {
+        for (final TypeInfo typeInfo : TYPE_INFOS) {
+            for (final Class<?> type : typeInfo.getTypes()) {
                 fa.put(type, typeInfo);
             }
         }
@@ -95,6 +96,7 @@ static Pair<String, String> getCodecAndArgs(
         if (!CodecLookup.codecRequired(columnDefinition)) {
             return null;
         }
+
         // Impute an appropriate codec for the data type
         final Class<?> dataType = columnDefinition.getDataType();
         if (Externalizable.class.isAssignableFrom(dataType)) {
@@ -365,6 +367,28 @@ public PrimitiveBuilder<PrimitiveType> getBuilder(boolean required, boolean repe
         }
     }
 
+    /**
+     * We will encode BigIntegers as Decimal types. Parquet has no special type for BigIntegers, but we can maintain
+     * external compatibility by encoding them as fixed length decimals of scale 1. Internally, we'll record that we
+     * wrote this as a decimal, so we can properly decode it back to BigInteger.
+     */
+    private enum BigIntegerType implements TypeInfo {
+        INSTANCE;
+
+        private static final Set<Class<?>> clazzes = Collections.singleton(BigInteger.class);
+
+        @Override
+        public Set<Class<?>> getTypes() {
+            return clazzes;
+        }
+
+        @Override
+        public PrimitiveBuilder<PrimitiveType> getBuilder(boolean required, boolean repeating, Class<?> dataType) {
+            return type(PrimitiveTypeName.BINARY, required, repeating)
+                    .as(LogicalTypeAnnotation.decimalType(0, 1));
+        }
+    }
+
     interface TypeInfo {
 
         Set<Class<?>> getTypes();