Untie decimal conversion from arrow metadata

src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java

@@ -4,8 +4,7 @@
 import static com.databricks.jdbc.common.DatabricksJdbcConstants.VOLUME_OPERATION_STATUS_COLUMN_NAME;
 import static com.databricks.jdbc.common.MetadataResultConstants.LARGE_DISPLAY_COLUMNS;
 import static com.databricks.jdbc.common.MetadataResultConstants.REMARKS_COLUMN;
-import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getTypeFromTypeDesc;
-import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getTypeTextFromTypeDesc;
+import static com.databricks.jdbc.common.util.DatabricksThriftUtil.*;
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.TIMESTAMP;
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.TIMESTAMP_NTZ;
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.VARIANT;
@@ -53,7 +52,6 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData {
   /**
    * Constructs a {@code DatabricksResultSetMetaData} object for a SEA result set.
    *
-   * @param statementId the unique identifier of the SQL statement execution
    * @param resultManifest the manifest containing metadata about the result set, including column
    *     information and types
    * @param usesExternalLinks whether or not the resultData contains external links (cloud fetch is
@@ -179,38 +177,37 @@ public DatabricksResultSetMetaData(
         for (int columnIndex = 0;
             columnIndex < resultManifest.getSchema().getColumnsSize();
             columnIndex++) {
-          TColumnDesc columnInfo = resultManifest.getSchema().getColumns().get(columnIndex);
-          ColumnInfoTypeName columnTypeName = getTypeFromTypeDesc(columnInfo.getTypeDesc());
-          int columnType = DatabricksTypeUtil.getColumnType(columnTypeName);
-          int[] precisionAndScale = getPrecisionAndScale(columnInfo, columnType);
+          TColumnDesc columnDesc = resultManifest.getSchema().getColumns().get(columnIndex);
+          ColumnInfo columnInfo = getColumnInfoFromTColumnDesc(columnDesc);
+          int[] precisionAndScale = getPrecisionAndScale(columnInfo);
           int precision = precisionAndScale[0];
           int scale = precisionAndScale[1];
 
           ImmutableDatabricksColumn.Builder columnBuilder = getColumnBuilder();
           columnBuilder
-              .columnName(columnInfo.getColumnName())
-              .columnTypeClassName(DatabricksTypeUtil.getColumnTypeClassName(columnTypeName))
-              .columnType(columnType)
-              .columnTypeText(
-                  getTypeTextFromTypeDesc(
-                      columnInfo
-                          .getTypeDesc())) // columnInfoTypeName does not have BIGINT, SMALLINT.
-              // Extracting from thriftType in typeDesc
+              .columnName(columnInfo.getName())
+              .columnTypeClassName(
+                  DatabricksTypeUtil.getColumnTypeClassName(columnInfo.getTypeName()))
+              .columnType(DatabricksTypeUtil.getColumnType(columnInfo.getTypeName()))
+              .columnTypeText(getTypeTextFromTypeDesc(columnDesc.getTypeDesc()))
+              // columnInfoTypeName does not have BIGINT, SMALLINT. Extracting from thriftType in
+              // typeDesc
               .typePrecision(precision)
               .typeScale(scale)
-              .displaySize(DatabricksTypeUtil.getDisplaySize(columnTypeName, precision, scale))
+              .displaySize(
+                  DatabricksTypeUtil.getDisplaySize(columnInfo.getTypeName(), precision, scale))
               .isSearchable(true)
               .schemaName(null)
               .tableName(null)
-              .isSigned(DatabricksTypeUtil.isSigned(columnTypeName));
+              .isSigned(DatabricksTypeUtil.isSigned(columnInfo.getTypeName()));
           if (isVariantColumn(arrowMetadata, columnIndex)) {
             columnBuilder
                 .columnTypeClassName("java.lang.String")
                 .columnType(Types.OTHER)
                 .columnTypeText(VARIANT);
           }
           columnsBuilder.add(columnBuilder.build());
-          columnNameToIndexMap.putIfAbsent(columnInfo.getColumnName(), ++currIndex);
+          columnNameToIndexMap.putIfAbsent(columnInfo.getName(), ++currIndex);
         }
       }
     }
@@ -628,20 +625,9 @@ public int[] getPrecisionAndScale(ColumnInfo columnInfo, int columnType) {
     return result;
   }
 
-  public int[] getPrecisionAndScale(TColumnDesc columnInfo, int columnType) {
-    int[] result = getBasePrecisionAndScale(columnType, ctx);
-    if (columnInfo.getTypeDesc() != null && columnInfo.getTypeDesc().getTypesSize() > 0) {
-      TTypeEntry tTypeEntry = columnInfo.getTypeDesc().getTypes().get(0);
-      if (tTypeEntry.isSetPrimitiveEntry()
-          && tTypeEntry.getPrimitiveEntry().isSetTypeQualifiers()
-          && tTypeEntry.getPrimitiveEntry().getTypeQualifiers().isSetQualifiers()) {
-        Map<String, TTypeQualifierValue> qualifiers =
-            tTypeEntry.getPrimitiveEntry().getTypeQualifiers().getQualifiers();
-        result[0] = qualifiers.get("precision").getI32Value(); // precision
-        result[1] = qualifiers.get("scale").getI32Value(); // scale
-      }
-    }
-    return result;
+  public int[] getPrecisionAndScale(ColumnInfo columnInfo) {
+    return getPrecisionAndScale(
+        columnInfo, DatabricksTypeUtil.getColumnType(columnInfo.getTypeName()));
   }
 
   private boolean isLargeColumn(String columnName) {

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java

@@ -18,6 +18,7 @@
 import com.databricks.jdbc.model.core.ExternalLink;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.sdk.service.sql.BaseChunkInfo;
+import com.databricks.sdk.service.sql.ColumnInfo;
 import com.databricks.sdk.service.sql.ColumnInfoTypeName;
 import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
@@ -195,12 +196,15 @@ boolean hasNextRow() {
 
     /** Returns object in the current row at the specified columnIndex. */
     Object getColumnObjectAtCurrentRow(
-        int columnIndex, ColumnInfoTypeName requiredType, String arrowMetadata)
+        int columnIndex,
+        ColumnInfoTypeName requiredType,
+        String arrowMetadata,
+        ColumnInfo columnInfo)
         throws DatabricksSQLException {
       ValueVector columnVector =
           this.resultChunk.getColumnVector(this.recordBatchCursorInChunk, columnIndex);
       return ArrowToJavaObjectConverter.convert(
-          columnVector, this.rowCursorInRecordBatch, requiredType, arrowMetadata);
+          columnVector, this.rowCursorInRecordBatch, requiredType, arrowMetadata, columnInfo);
     }
 
     String getType(int columnIndex) {

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java

@@ -1,6 +1,6 @@
 package com.databricks.jdbc.api.impl.arrow;
 
-import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getTypeFromTypeDesc;
+import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getColumnInfoFromTColumnDesc;
 
 import com.databricks.jdbc.api.impl.ComplexDataTypeParser;
 import com.databricks.jdbc.api.impl.IExecutionResult;
@@ -142,12 +142,13 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
 
       Object result =
           chunkIterator.getColumnObjectAtCurrentRow(
-              columnIndex, ColumnInfoTypeName.STRING, "STRING");
+              columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfos.get(columnIndex));
       ComplexDataTypeParser parser = new ComplexDataTypeParser();
       return parser.formatComplexTypeString(result.toString(), requiredType.name(), arrowMetadata);
     }
 
-    return chunkIterator.getColumnObjectAtCurrentRow(columnIndex, requiredType, arrowMetadata);
+    return chunkIterator.getColumnObjectAtCurrentRow(
+        columnIndex, requiredType, arrowMetadata, columnInfos.get(columnIndex));
   }
 
   /**
@@ -224,8 +225,8 @@ private void setColumnInfo(TGetResultSetMetadataResp resultManifest) {
     if (resultManifest.getSchema() == null) {
       return;
     }
-    for (TColumnDesc columnInfo : resultManifest.getSchema().getColumns()) {
-      columnInfos.add(new ColumnInfo().setTypeName(getTypeFromTypeDesc(columnInfo.getTypeDesc())));
+    for (TColumnDesc tColumnDesc : resultManifest.getSchema().getColumns()) {
+      columnInfos.add(getColumnInfoFromTColumnDesc(tColumnDesc));
     }
   }
 }

src/main/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProvider.java

@@ -199,8 +199,8 @@ private static Schema hiveSchemaToArrowSchema(TTableSchema hiveSchema)
   }
 
   private static Field getArrowField(TColumnDesc columnDesc) throws SQLException {
-    TTypeId thriftType = getThriftTypeFromTypeDesc(columnDesc.getTypeDesc());
-    ArrowType arrowType = mapThriftToArrowType(thriftType);
+    TPrimitiveTypeEntry primitiveTypeEntry = getTPrimitiveTypeOrDefault(columnDesc.getTypeDesc());
+    ArrowType arrowType = mapThriftToArrowType(primitiveTypeEntry.getType());
     FieldType fieldType = new FieldType(true, arrowType, null);
     return new Field(columnDesc.getColumnName(), fieldType, null);
   }

src/main/java/com/databricks/jdbc/api/impl/converters/ArrowToJavaObjectConverter.java

@@ -12,6 +12,7 @@
 import com.databricks.jdbc.exception.DatabricksValidationException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
+import com.databricks.sdk.service.sql.ColumnInfo;
 import com.databricks.sdk.service.sql.ColumnInfoTypeName;
 import java.math.BigDecimal;
 import java.math.RoundingMode;
@@ -61,7 +62,8 @@ public static Object convert(
       ValueVector columnVector,
       int vectorIndex,
       ColumnInfoTypeName requiredType,
-      String arrowMetadata)
+      String arrowMetadata,
+      ColumnInfo columnInfo)
       throws DatabricksSQLException {
     // check isNull before getting the object from the vector
     if (columnVector.isNull(vectorIndex)) {
@@ -102,7 +104,7 @@ public static Object convert(
       case DOUBLE:
         return convertToNumber(object, Double::parseDouble, Number::doubleValue);
       case DECIMAL:
-        return convertToDecimal(object, arrowMetadata);
+        return convertToDecimal(object, columnInfo);
       case BINARY:
         return convertToByteArray(object);
       case BOOLEAN:
@@ -262,30 +264,22 @@ private static byte[] convertToByteArray(Object object) {
     return (byte[]) object;
   }
 
-  static BigDecimal convertToDecimal(Object object, String arrowMetadata)
+  static BigDecimal convertToDecimal(Object object, ColumnInfo columnInfo)
       throws DatabricksValidationException {
+    BigDecimal bigDecimal = null;
+    int bigDecimalScale = columnInfo.getTypePrecision().intValue();
     if (object instanceof Text) {
-      return new BigDecimal(object.toString());
+      bigDecimal = new BigDecimal(object.toString());
+    } else if (object instanceof Number) {
+      bigDecimal = new BigDecimal(object.toString());
     }
-    int scale;
-    try {
-      scale =
-          Integer.parseInt(
-              arrowMetadata
-                  .substring(arrowMetadata.indexOf(',') + 1, arrowMetadata.indexOf(')'))
-                  .trim());
-    } catch (Exception e) {
-      LOGGER.error(
-          e, "Failed to parse scale from arrow metadata: {}. Defaulting to scale 0", arrowMetadata);
-      scale = 0;
+    if (bigDecimal == null) {
+      String errorMessage =
+          String.format("Unsupported object type for decimal conversion: %s", object.getClass());
+      LOGGER.error(errorMessage);
+      throw new DatabricksValidationException(errorMessage);
     }
-    if (object instanceof Number) {
-      return new BigDecimal(object.toString()).setScale(scale, RoundingMode.HALF_UP);
-    }
-    String errorMessage =
-        String.format("Unsupported object type for decimal conversion: %s", object.getClass());
-    LOGGER.error(errorMessage);
-    throw new DatabricksValidationException(errorMessage);
+    return bigDecimal.setScale(bigDecimalScale, RoundingMode.HALF_UP);
   }
 
   private static <T extends Number> T convertToNumber(

src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java

@@ -2,6 +2,7 @@
 
 import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT;
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.*;
+import static com.databricks.jdbc.model.client.thrift.generated.TTypeId.*;
 
 import com.databricks.jdbc.api.internal.IDatabricksSession;
 import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
@@ -14,13 +15,37 @@
 import com.databricks.jdbc.model.core.ExternalLink;
 import com.databricks.jdbc.model.core.StatementStatus;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+import com.databricks.sdk.service.sql.ColumnInfo;
 import com.databricks.sdk.service.sql.ColumnInfoTypeName;
 import com.databricks.sdk.service.sql.StatementState;
 import java.nio.ByteBuffer;
 import java.util.*;
 
 public class DatabricksThriftUtil {
 
+  private static final Map<TTypeId, ColumnInfoTypeName> T_TYPE_ID_COLUMN_INFO_TYPE_NAME_MAP =
+      Map.ofEntries(
+          Map.entry(BOOLEAN_TYPE, ColumnInfoTypeName.BOOLEAN),
+          Map.entry(TINYINT_TYPE, ColumnInfoTypeName.BYTE),
+          Map.entry(SMALLINT_TYPE, ColumnInfoTypeName.SHORT),
+          Map.entry(INT_TYPE, ColumnInfoTypeName.INT),
+          Map.entry(BIGINT_TYPE, ColumnInfoTypeName.LONG),
+          Map.entry(FLOAT_TYPE, ColumnInfoTypeName.FLOAT),
+          Map.entry(DOUBLE_TYPE, ColumnInfoTypeName.DOUBLE),
+          Map.entry(STRING_TYPE, ColumnInfoTypeName.STRING),
+          Map.entry(VARCHAR_TYPE, ColumnInfoTypeName.STRING),
+          Map.entry(TIMESTAMP_TYPE, ColumnInfoTypeName.TIMESTAMP),
+          Map.entry(BINARY_TYPE, ColumnInfoTypeName.BINARY),
+          Map.entry(DECIMAL_TYPE, ColumnInfoTypeName.DECIMAL),
+          Map.entry(DATE_TYPE, ColumnInfoTypeName.DATE),
+          Map.entry(CHAR_TYPE, ColumnInfoTypeName.CHAR),
+          Map.entry(INTERVAL_YEAR_MONTH_TYPE, ColumnInfoTypeName.INTERVAL),
+          Map.entry(INTERVAL_DAY_TIME_TYPE, ColumnInfoTypeName.INTERVAL),
+          Map.entry(ARRAY_TYPE, ColumnInfoTypeName.ARRAY),
+          Map.entry(MAP_TYPE, ColumnInfoTypeName.MAP),
+          Map.entry(NULL_TYPE, ColumnInfoTypeName.STRING),
+          Map.entry(STRUCT_TYPE, ColumnInfoTypeName.STRUCT));
+
   private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksThriftUtil.class);
   private static final List<TStatusCode> SUCCESS_STATUS_LIST =
       List.of(TStatusCode.SUCCESS_STATUS, TStatusCode.SUCCESS_WITH_INFO_STATUS);
@@ -148,49 +173,33 @@ public static StatementStatus getAsyncStatus(TStatus status) {
   }
 
   public static String getTypeTextFromTypeDesc(TTypeDesc typeDesc) {
-    TTypeId type = getThriftTypeFromTypeDesc(typeDesc);
-    return type.name().replace("_TYPE", "");
+    TPrimitiveTypeEntry primitiveTypeEntry = getTPrimitiveTypeOrDefault(typeDesc);
+    return primitiveTypeEntry.getType().name().replace("_TYPE", "");
   }
 
-  public static ColumnInfoTypeName getTypeFromTypeDesc(TTypeDesc typeDesc) {
-    TTypeId type = getThriftTypeFromTypeDesc(typeDesc);
-    switch (type) {
-      case BOOLEAN_TYPE:
-        return ColumnInfoTypeName.BOOLEAN;
-      case TINYINT_TYPE:
-        return ColumnInfoTypeName.BYTE;
-      case SMALLINT_TYPE:
-        return ColumnInfoTypeName.SHORT;
-      case INT_TYPE:
-        return ColumnInfoTypeName.INT;
-      case BIGINT_TYPE:
-        return ColumnInfoTypeName.LONG;
-      case FLOAT_TYPE:
-        return ColumnInfoTypeName.FLOAT;
-      case DOUBLE_TYPE:
-        return ColumnInfoTypeName.DOUBLE;
-      case TIMESTAMP_TYPE:
-        return ColumnInfoTypeName.TIMESTAMP;
-      case BINARY_TYPE:
-        return ColumnInfoTypeName.BINARY;
-      case DECIMAL_TYPE:
-        return ColumnInfoTypeName.DECIMAL;
-      case DATE_TYPE:
-        return ColumnInfoTypeName.DATE;
-      case CHAR_TYPE:
-        return ColumnInfoTypeName.CHAR;
-      case INTERVAL_YEAR_MONTH_TYPE:
-      case INTERVAL_DAY_TIME_TYPE:
-        return ColumnInfoTypeName.INTERVAL;
-      case ARRAY_TYPE:
-        return ColumnInfoTypeName.ARRAY;
-      case MAP_TYPE:
-        return ColumnInfoTypeName.MAP;
-      case STRUCT_TYPE:
-        return ColumnInfoTypeName.STRUCT;
-      default:
-        return ColumnInfoTypeName.STRING;
+  public static ColumnInfo getColumnInfoFromTColumnDesc(TColumnDesc columnDesc) {
+    TPrimitiveTypeEntry primitiveTypeEntry = getTPrimitiveTypeOrDefault(columnDesc.getTypeDesc());
+    ColumnInfoTypeName columnInfoTypeName =
+        T_TYPE_ID_COLUMN_INFO_TYPE_NAME_MAP.get(primitiveTypeEntry.getType());
+    ColumnInfo columnInfo =
+        new ColumnInfo()
+            .setName(columnDesc.getColumnName())
+            .setPosition((long) columnDesc.getPosition())
+            .setTypeName(columnInfoTypeName);
+    if (primitiveTypeEntry.isSetTypeQualifiers()) {
+      TTypeQualifiers typeQualifiers = primitiveTypeEntry.getTypeQualifiers();
+      String scaleQualifierKey = TCLIServiceConstants.SCALE,
+          precisionQualifierKey = TCLIServiceConstants.PRECISION;
+      if (typeQualifiers.getQualifiers().get(scaleQualifierKey) != null) {
+        columnInfo.setTypeScale(
+            (long) typeQualifiers.getQualifiers().get(scaleQualifierKey).getI32Value());
+      }
+      if (typeQualifiers.getQualifiers().get(precisionQualifierKey) != null) {
+        columnInfo.setTypePrecision(
+            (long) typeQualifiers.getQualifiers().get(precisionQualifierKey).getI32Value());
+      }
     }
+    return columnInfo;
   }
 
   /**

src/main/java/com/databricks/jdbc/common/util/DatabricksTypeUtil.java

@@ -444,13 +444,13 @@ public static String inferDatabricksType(Object obj) {
     return type;
   }
 
-  public static TTypeId getThriftTypeFromTypeDesc(TTypeDesc typeDesc) {
+  public static TPrimitiveTypeEntry getTPrimitiveTypeOrDefault(TTypeDesc typeDesc) {
+    TPrimitiveTypeEntry defaultPrimitiveTypeEntry = new TPrimitiveTypeEntry(TTypeId.STRING_TYPE);
     return Optional.ofNullable(typeDesc)
         .map(TTypeDesc::getTypes)
         .map(t -> t.get(0))
         .map(TTypeEntry::getPrimitiveEntry)
-        .map(TPrimitiveTypeEntry::getType)
-        .orElse(TTypeId.STRING_TYPE);
+        .orElse(defaultPrimitiveTypeEntry);
   }
 
   public static ArrowType mapThriftToArrowType(TTypeId typeId) throws SQLException {