From 7de14bf91b2c7e01f08b5d562353f34cf14162b3 Mon Sep 17 00:00:00 2001
From: x30027624 <xiangxinwei@huawei.com>
Date: Tue, 23 Jan 2024 15:07:12 +0800
Subject: [PATCH] fix_partition_scan_decimal64_error

---
 .../orc/OmniOrcColumnarBatchReader.java       |  2 +-
 .../vectorized/OmniColumnVector.java          | 58 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/datasources/orc/OmniOrcColumnarBatchReader.java b/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/datasources/orc/OmniOrcColumnarBatchReader.java
index f9e3913ad..75535c15c 100644
--- a/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/datasources/orc/OmniOrcColumnarBatchReader.java
+++ b/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/datasources/orc/OmniOrcColumnarBatchReader.java
@@ -177,7 +177,7 @@ public class OmniOrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch
       DataType dt = resultFields[i].dataType();
       if (requestedPartitionColIds[i] != -1) {
         OmniColumnVector partitionCol = new OmniColumnVector(capacity, dt, true);
-        ColumnVectorUtils.populate(partitionCol, partitionValues, requestedPartitionColIds[i]);
+        OmniColumnVector.populate(partitionCol, partitionValues, requestedPartitionColIds[i]);
         partitionCol.setIsConstant();
         templateWrappers[i] = partitionCol;
         orcVectorWrappers[i] = new OmniColumnVector(capacity, dt, false);;
diff --git a/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/vectorized/OmniColumnVector.java b/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/vectorized/OmniColumnVector.java
index 4a84c2a7c..800c9c6ba 100644
--- a/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/vectorized/OmniColumnVector.java
+++ b/omnioperator/omniop-spark-extension/java/src/main/java/org/apache/spark/sql/execution/vectorized/OmniColumnVector.java
@@ -20,9 +20,11 @@ package org.apache.spark.sql.execution.vectorized;
 
 import nova.hetu.omniruntime.vector.*;
 
+import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.util.RebaseDateTime;
 import org.apache.spark.sql.types.*;
 import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.types.CalendarInterval;
 import org.apache.spark.unsafe.types.UTF8String;
 
 import java.math.BigDecimal;
@@ -307,6 +309,62 @@ public class OmniColumnVector extends WritableColumnVector {
         }
     }
 
+    public static void populate(WritableColumnVector col, InternalRow row, int fieldIdx) {
+        int capacity = col.capacity;
+        DataType t = col.dataType();
+
+        if (row.isNullAt(fieldIdx)) {
+            col.putNulls(0, capacity);
+        } else {
+            if (t == DataTypes.BooleanType) {
+                col.putBooleans(0, capacity, row.getBoolean(fieldIdx));
+            } else if (t == DataTypes.BinaryType) {
+                col.putByteArray(0, row.getBinary(fieldIdx));
+            } else if (t == DataTypes.ByteType) {
+                col.putBytes(0, capacity, row.getByte(fieldIdx));
+            } else if (t == DataTypes.ShortType) {
+                col.putShorts(0, capacity, row.getShort(fieldIdx));
+            } else if (t == DataTypes.IntegerType) {
+                col.putInts(0, capacity, row.getInt(fieldIdx));
+            } else if (t == DataTypes.LongType) {
+                col.putLongs(0, capacity, row.getLong(fieldIdx));
+            } else if (t == DataTypes.FloatType) {
+                col.putFloats(0, capacity, row.getFloat(fieldIdx));
+            } else if (t == DataTypes.DoubleType) {
+                col.putDoubles(0, capacity, row.getDouble(fieldIdx));
+            } else if (t == DataTypes.StringType) {
+                UTF8String v = row.getUTF8String(fieldIdx);
+                byte[] bytes = v.getBytes();
+                for (int i = 0; i < capacity; i++) {
+                    col.putByteArray(i, bytes);
+                }
+            } else if (t instanceof DecimalType) {
+                DecimalType dt = (DecimalType)t;
+                Decimal d = row.getDecimal(fieldIdx, dt.precision(), dt.scale());
+                if (dt.precision() <= Decimal.MAX_LONG_DIGITS()) {
+                    col.putLongs(0, capacity, d.toUnscaledLong());
+                } else {
+                    final BigInteger integer = d.toJavaBigDecimal().unscaledValue();
+                    byte[] bytes = integer.toByteArray();
+                    for (int i = 0; i < capacity; i++) {
+                        col.putByteArray(i, bytes, 0, bytes.length);
+                    }
+                }
+            } else if (t instanceof CalendarIntervalType) {
+                CalendarInterval c = (CalendarInterval)row.get(fieldIdx, t);
+                col.getChild(0).putInts(0, capacity, c.months);
+                col.getChild(1).putLongs(0, capacity, c.microseconds);
+            } else if (t instanceof DateType) {
+                col.putInts(0, capacity, row.getInt(fieldIdx));
+            } else if (t instanceof TimestampType) {
+                col.putLongs(0, capacity, row.getLong(fieldIdx));
+            } else {
+                throw new RuntimeException(String.format("DataType %s is not supported" +
+                        " in column vectorized reader.", t.sql()));
+            }
+        }
+    }
+
     @Override
     public void putNotNulls(int rowId, int count) {}
 
-- 
Gitee