You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2024/04/15 05:16:46 UTC
(spark) branch master updated: [SPARK-47803][SQL] Support cast to variant
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new dd4c7fc6ae91 [SPARK-47803][SQL] Support cast to variant
dd4c7fc6ae91 is described below
commit dd4c7fc6ae91ec1f2a348cf4633bff10717cde22
Author: Chenhao Li <ch...@databricks.com>
AuthorDate: Mon Apr 15 13:16:30 2024 +0800
[SPARK-47803][SQL] Support cast to variant
### What changes were proposed in this pull request?
This PR allows casting another type into the variant type. The changes can be divided into two major parts:
- The `VariantBuilder` class is greatly refactored. Many of its APIs are exposed so that Spark can use them to build a variant value without JSON parsing.
- The actual implementation of the cast.
### Why are the changes needed?
It provides a convenient way to build variant values from other Spark values. Before this PR, `parse_json` is the only SQL function that can produce variant values. If users want to do so, they may have to use `parse_json(to_json(input))`, which is inefficient and disallowed if the input has a scalar type.
### Does this PR introduce _any_ user-facing change?
Yes. Casting to variant was previously not allowed but now allowd.
### How was this patch tested?
Unit tests.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #45989 from chenhao-db/cast_to_variant.
Authored-by: Chenhao Li <ch...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../src/main/resources/error/error-classes.json | 6 +
.../org/apache/spark/types/variant/Variant.java | 6 +-
.../apache/spark/types/variant/VariantBuilder.java | 423 ++++++++++++++-------
docs/sql-error-conditions.md | 6 +
.../spark/sql/catalyst/expressions/Cast.scala | 7 +
.../variant/VariantExpressionEvalUtils.scala | 71 +++-
.../variant/VariantExpressionSuite.scala | 26 ++
7 files changed, 406 insertions(+), 139 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-classes.json b/common/utils/src/main/resources/error/error-classes.json
index 2a00edb9a4df..e1c8c881f98f 100644
--- a/common/utils/src/main/resources/error/error-classes.json
+++ b/common/utils/src/main/resources/error/error-classes.json
@@ -4725,6 +4725,12 @@
],
"sqlState" : "22023"
},
+ "VARIANT_DUPLICATE_KEY" : {
+ "message" : [
+ "Failed to build variant because of a duplicate object key `<key>`."
+ ],
+ "sqlState" : "22023"
+ },
"VARIANT_SIZE_LIMIT" : {
"message" : [
"Cannot build variant bigger than <sizeLimit> in <functionName>.",
diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java b/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java
index 4aeb2c6e1435..a705daaf323b 100644
--- a/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java
+++ b/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java
@@ -41,12 +41,12 @@ import static org.apache.spark.types.variant.VariantUtil.*;
* define a new class to avoid depending on or modifying Spark.
*/
public final class Variant {
- private final byte[] value;
- private final byte[] metadata;
+ final byte[] value;
+ final byte[] metadata;
// The variant value doesn't use the whole `value` binary, but starts from its `pos` index and
// spans a size of `valueSize(value, pos)`. This design avoids frequent copies of the value binary
// when reading a sub-variant in the array/object element.
- private final int pos;
+ final int pos;
public Variant(byte[] value, byte[] metadata) {
this(value, metadata, 0);
diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
index 21a12cbe9d71..ea7a7674baf5 100644
--- a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
+++ b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
@@ -17,6 +17,11 @@
package org.apache.spark.types.variant;
+import org.apache.spark.QueryContext;
+import org.apache.spark.SparkRuntimeException;
+import scala.collection.immutable.Map;
+import scala.collection.immutable.Map$;
+
import java.io.IOException;
import java.math.BigDecimal;
import java.math.BigInteger;
@@ -61,7 +66,7 @@ public class VariantBuilder {
}
// Build the variant metadata from `dictionaryKeys` and return the variant result.
- private Variant result() {
+ public Variant result() {
int numKeys = dictionaryKeys.size();
// Use long to avoid overflow in accumulating lengths.
long dictionaryStringSize = 0;
@@ -100,6 +105,272 @@ public class VariantBuilder {
return new Variant(Arrays.copyOfRange(writeBuffer, 0, writePos), metadata);
}
+ public void appendString(String str) {
+ byte[] text = str.getBytes(StandardCharsets.UTF_8);
+ boolean longStr = text.length > MAX_SHORT_STR_SIZE;
+ checkCapacity((longStr ? 1 + U32_SIZE : 1) + text.length);
+ if (longStr) {
+ writeBuffer[writePos++] = primitiveHeader(LONG_STR);
+ writeLong(writeBuffer, writePos, text.length, U32_SIZE);
+ writePos += U32_SIZE;
+ } else {
+ writeBuffer[writePos++] = shortStrHeader(text.length);
+ }
+ System.arraycopy(text, 0, writeBuffer, writePos, text.length);
+ writePos += text.length;
+ }
+
+ public void appendNull() {
+ checkCapacity(1);
+ writeBuffer[writePos++] = primitiveHeader(NULL);
+ }
+
+ public void appendBoolean(boolean b) {
+ checkCapacity(1);
+ writeBuffer[writePos++] = primitiveHeader(b ? TRUE : FALSE);
+ }
+
+ // Append a long value to the variant builder. The actual used integer type depends on the value
+ // range of the long value.
+ public void appendLong(long l) {
+ checkCapacity(1 + 8);
+ if (l == (byte) l) {
+ writeBuffer[writePos++] = primitiveHeader(INT1);
+ writeLong(writeBuffer, writePos, l, 1);
+ writePos += 1;
+ } else if (l == (short) l) {
+ writeBuffer[writePos++] = primitiveHeader(INT2);
+ writeLong(writeBuffer, writePos, l, 2);
+ writePos += 2;
+ } else if (l == (int) l) {
+ writeBuffer[writePos++] = primitiveHeader(INT4);
+ writeLong(writeBuffer, writePos, l, 4);
+ writePos += 4;
+ } else {
+ writeBuffer[writePos++] = primitiveHeader(INT8);
+ writeLong(writeBuffer, writePos, l, 8);
+ writePos += 8;
+ }
+ }
+
+ public void appendDouble(double d) {
+ checkCapacity(1 + 8);
+ writeBuffer[writePos++] = primitiveHeader(DOUBLE);
+ writeLong(writeBuffer, writePos, Double.doubleToLongBits(d), 8);
+ writePos += 8;
+ }
+
+ // Append a decimal value to the variant builder. The caller should guarantee that its precision
+ // and scale fit into `MAX_DECIMAL16_PRECISION`.
+ public void appendDecimal(BigDecimal d) {
+ checkCapacity(2 + 16);
+ BigInteger unscaled = d.unscaledValue();
+ if (d.scale() <= MAX_DECIMAL4_PRECISION && d.precision() <= MAX_DECIMAL4_PRECISION) {
+ writeBuffer[writePos++] = primitiveHeader(DECIMAL4);
+ writeBuffer[writePos++] = (byte) d.scale();
+ writeLong(writeBuffer, writePos, unscaled.intValueExact(), 4);
+ writePos += 4;
+ } else if (d.scale() <= MAX_DECIMAL8_PRECISION && d.precision() <= MAX_DECIMAL8_PRECISION) {
+ writeBuffer[writePos++] = primitiveHeader(DECIMAL8);
+ writeBuffer[writePos++] = (byte) d.scale();
+ writeLong(writeBuffer, writePos, unscaled.longValueExact(), 8);
+ writePos += 8;
+ } else {
+ assert d.scale() <= MAX_DECIMAL16_PRECISION && d.precision() <= MAX_DECIMAL16_PRECISION;
+ writeBuffer[writePos++] = primitiveHeader(DECIMAL16);
+ writeBuffer[writePos++] = (byte) d.scale();
+ // `toByteArray` returns a big-endian representation. We need to copy it reversely and sign
+ // extend it to 16 bytes.
+ byte[] bytes = unscaled.toByteArray();
+ for (int i = 0; i < bytes.length; ++i) {
+ writeBuffer[writePos + i] = bytes[bytes.length - 1 - i];
+ }
+ byte sign = (byte) (bytes[0] < 0 ? -1 : 0);
+ for (int i = bytes.length; i < 16; ++i) {
+ writeBuffer[writePos + i] = sign;
+ }
+ writePos += 16;
+ }
+ }
+
+ public void appendDate(int daysSinceEpoch) {
+ checkCapacity(1 + 4);
+ writeBuffer[writePos++] = primitiveHeader(DATE);
+ writeLong(writeBuffer, writePos, daysSinceEpoch, 4);
+ writePos += 4;
+ }
+
+ public void appendTimestamp(long microsSinceEpoch) {
+ checkCapacity(1 + 8);
+ writeBuffer[writePos++] = primitiveHeader(TIMESTAMP);
+ writeLong(writeBuffer, writePos, microsSinceEpoch, 8);
+ writePos += 8;
+ }
+
+ public void appendTimestampNtz(long microsSinceEpoch) {
+ checkCapacity(1 + 8);
+ writeBuffer[writePos++] = primitiveHeader(TIMESTAMP_NTZ);
+ writeLong(writeBuffer, writePos, microsSinceEpoch, 8);
+ writePos += 8;
+ }
+
+ public void appendFloat(float f) {
+ checkCapacity(1 + 4);
+ writeBuffer[writePos++] = primitiveHeader(FLOAT);
+ writeLong(writeBuffer, writePos, Float.floatToIntBits(f), 8);
+ writePos += 4;
+ }
+
+ public void appendBinary(byte[] binary) {
+ checkCapacity(1 + U32_SIZE + binary.length);
+ writeBuffer[writePos++] = primitiveHeader(LONG_STR);
+ writeLong(writeBuffer, writePos, binary.length, U32_SIZE);
+ writePos += U32_SIZE;
+ System.arraycopy(binary, 0, writeBuffer, writePos, binary.length);
+ writePos += binary.length;
+ }
+
+ // Add a key to the variant dictionary. If the key already exists, the dictionary is not modified.
+ // In either case, return the id of the key.
+ public int addKey(String key) {
+ int id;
+ if (dictionary.containsKey(key)) {
+ id = dictionary.get(key);
+ } else {
+ id = dictionaryKeys.size();
+ dictionary.put(key, id);
+ dictionaryKeys.add(key.getBytes(StandardCharsets.UTF_8));
+ }
+ return id;
+ }
+
+ // Return the current write position of the variant builder. It is used together with
+ // `finishWritingObject` or `finishWritingArray`.
+ public int getWritePos() {
+ return writePos;
+ }
+
+ // Finish writing a variant object after all of its fields have already been written. The process
+ // is as follows:
+ // 1. The caller calls `getWritePos` before writing any fields to obtain the `start` parameter.
+ // 2. The caller appends all the object fields to the builder. In the meantime, it should maintain
+ // the `fields` parameter. Before appending each field, it should append an entry to `fields` to
+ // record the offset of the field. The offset is computed as `getWritePos() - start`.
+ // 3. The caller calls `finishWritingObject` to finish writing a variant object.
+ //
+ // This function is responsible to sort the fields by key and check for any duplicate field keys.
+ public void finishWritingObject(int start, ArrayList<FieldEntry> fields) {
+ int dataSize = writePos - start;
+ int size = fields.size();
+ Collections.sort(fields);
+ int maxId = size == 0 ? 0 : fields.get(0).id;
+ // Check for duplicate field keys. Only need to check adjacent key because they are sorted.
+ for (int i = 1; i < size; ++i) {
+ maxId = Math.max(maxId, fields.get(i).id);
+ String key = fields.get(i).key;
+ if (key.equals(fields.get(i - 1).key)) {
+ @SuppressWarnings("unchecked")
+ Map<String, String> parameters = Map$.MODULE$.<String, String>empty().updated("key", key);
+ throw new SparkRuntimeException("VARIANT_DUPLICATE_KEY", parameters,
+ null, new QueryContext[]{}, "");
+ }
+ }
+ boolean largeSize = size > U8_MAX;
+ int sizeBytes = largeSize ? U32_SIZE : 1;
+ int idSize = getIntegerSize(maxId);
+ int offsetSize = getIntegerSize(dataSize);
+ // The space for header byte, object size, id list, and offset list.
+ int headerSize = 1 + sizeBytes + size * idSize + (size + 1) * offsetSize;
+ checkCapacity(headerSize);
+ // Shift the just-written field data to make room for the object header section.
+ System.arraycopy(writeBuffer, start, writeBuffer, start + headerSize, dataSize);
+ writePos += headerSize;
+ writeBuffer[start] = objectHeader(largeSize, idSize, offsetSize);
+ writeLong(writeBuffer, start + 1, size, sizeBytes);
+ int idStart = start + 1 + sizeBytes;
+ int offsetStart = idStart + size * idSize;
+ for (int i = 0; i < size; ++i) {
+ writeLong(writeBuffer, idStart + i * idSize, fields.get(i).id, idSize);
+ writeLong(writeBuffer, offsetStart + i * offsetSize, fields.get(i).offset, offsetSize);
+ }
+ writeLong(writeBuffer, offsetStart + size * offsetSize, dataSize, offsetSize);
+ }
+
+ // Finish writing a variant array after all of its elements have already been written. The process
+ // is similar to that of `finishWritingObject`.
+ public void finishWritingArray(int start, ArrayList<Integer> offsets) {
+ int dataSize = writePos - start;
+ int size = offsets.size();
+ boolean largeSize = size > U8_MAX;
+ int sizeBytes = largeSize ? U32_SIZE : 1;
+ int offsetSize = getIntegerSize(dataSize);
+ // The space for header byte, object size, and offset list.
+ int headerSize = 1 + sizeBytes + (size + 1) * offsetSize;
+ checkCapacity(headerSize);
+ // Shift the just-written field data to make room for the header section.
+ System.arraycopy(writeBuffer, start, writeBuffer, start + headerSize, dataSize);
+ writePos += headerSize;
+ writeBuffer[start] = arrayHeader(largeSize, offsetSize);
+ writeLong(writeBuffer, start + 1, size, sizeBytes);
+ int offsetStart = start + 1 + sizeBytes;
+ for (int i = 0; i < size; ++i) {
+ writeLong(writeBuffer, offsetStart + i * offsetSize, offsets.get(i), offsetSize);
+ }
+ writeLong(writeBuffer, offsetStart + size * offsetSize, dataSize, offsetSize);
+ }
+
+ // Append a variant value to the variant builder. We need to insert the keys in the input variant
+ // into the current variant dictionary and rebuild it with new field ids. For scalar values in the
+ // input variant, we can directly copy the binary slice.
+ public void appendVariant(Variant v) {
+ appendVariantImpl(v.value, v.metadata, v.pos);
+ }
+
+ private void appendVariantImpl(byte[] value, byte[] metadata, int pos) {
+ checkIndex(pos, value.length);
+ int basicType = value[pos] & BASIC_TYPE_MASK;
+ switch (basicType) {
+ case OBJECT:
+ handleObject(value, pos, (size, idSize, offsetSize, idStart, offsetStart, dataStart) -> {
+ ArrayList<FieldEntry> fields = new ArrayList<>(size);
+ int start = writePos;
+ for (int i = 0; i < size; ++i) {
+ int id = readUnsigned(value, idStart + idSize * i, idSize);
+ int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+ int elementPos = dataStart + offset;
+ String key = getMetadataKey(metadata, id);
+ int newId = addKey(key);
+ fields.add(new FieldEntry(key, newId, writePos - start));
+ appendVariantImpl(value, metadata, elementPos);
+ }
+ finishWritingObject(start, fields);
+ return null;
+ });
+ break;
+ case ARRAY:
+ handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> {
+ ArrayList<Integer> offsets = new ArrayList<>(size);
+ int start = writePos;
+ for (int i = 0; i < size; ++i) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+ int elementPos = dataStart + offset;
+ offsets.add(writePos - start);
+ appendVariantImpl(value, metadata, elementPos);
+ }
+ finishWritingArray(start, offsets);
+ return null;
+ });
+ break;
+ default:
+ int size = valueSize(value, pos);
+ checkIndex(pos + size - 1, value.length);
+ checkCapacity(size);
+ System.arraycopy(value, pos, writeBuffer, writePos, size);
+ writePos += size;
+ break;
+ }
+ }
+
private void checkCapacity(int additional) {
int required = writePos + additional;
if (required > writeBuffer.length) {
@@ -117,12 +388,12 @@ public class VariantBuilder {
// Temporarily store the information of a field. We need to collect all fields in an JSON object,
// sort them by their keys, and build the variant object in sorted order.
- private static final class FieldEntry implements Comparable<FieldEntry> {
+ public static final class FieldEntry implements Comparable<FieldEntry> {
final String key;
final int id;
final int offset;
- FieldEntry(String key, int id, int offset) {
+ public FieldEntry(String key, int id, int offset) {
this.key = key;
this.id = id;
this.offset = offset;
@@ -143,117 +414,32 @@ public class VariantBuilder {
case START_OBJECT: {
ArrayList<FieldEntry> fields = new ArrayList<>();
int start = writePos;
- int maxId = 0;
while (parser.nextToken() != JsonToken.END_OBJECT) {
String key = parser.currentName();
parser.nextToken();
- int id;
- if (dictionary.containsKey(key)) {
- id = dictionary.get(key);
- } else {
- id = dictionaryKeys.size();
- dictionary.put(key, id);
- dictionaryKeys.add(key.getBytes(StandardCharsets.UTF_8));
- }
- maxId = Math.max(maxId, id);
- int offset = writePos - start;
- fields.add(new FieldEntry(key, id, offset));
+ int id = addKey(key);
+ fields.add(new FieldEntry(key, id, writePos - start));
buildJson(parser);
}
- int dataSize = writePos - start;
- int size = fields.size();
- Collections.sort(fields);
- // Check for duplicate field keys. Only need to check adjacent key because they are sorted.
- for (int i = 1; i < size; ++i) {
- String key = fields.get(i - 1).key;
- if (key.equals(fields.get(i).key)) {
- throw new JsonParseException(parser, "Duplicate key: " + key);
- }
- }
- boolean largeSize = size > U8_MAX;
- int sizeBytes = largeSize ? U32_SIZE : 1;
- int idSize = getIntegerSize(maxId);
- int offsetSize = getIntegerSize(dataSize);
- // The space for header byte, object size, id list, and offset list.
- int headerSize = 1 + sizeBytes + size * idSize + (size + 1) * offsetSize;
- checkCapacity(headerSize);
- // Shift the just-written field data to make room for the object header section.
- System.arraycopy(writeBuffer, start, writeBuffer, start + headerSize, dataSize);
- writePos += headerSize;
- writeBuffer[start] = objectHeader(largeSize, idSize, offsetSize);
- writeLong(writeBuffer, start + 1, size, sizeBytes);
- int idStart = start + 1 + sizeBytes;
- int offsetStart = idStart + size * idSize;
- for (int i = 0; i < size; ++i) {
- writeLong(writeBuffer, idStart + i * idSize, fields.get(i).id, idSize);
- writeLong(writeBuffer, offsetStart + i * offsetSize, fields.get(i).offset, offsetSize);
- }
- writeLong(writeBuffer, offsetStart + size * offsetSize, dataSize, offsetSize);
+ finishWritingObject(start, fields);
break;
}
case START_ARRAY: {
ArrayList<Integer> offsets = new ArrayList<>();
int start = writePos;
while (parser.nextToken() != JsonToken.END_ARRAY) {
- int offset = writePos - start;
- offsets.add(offset);
+ offsets.add(writePos - start);
buildJson(parser);
}
- int dataSize = writePos - start;
- int size = offsets.size();
- boolean largeSize = size > U8_MAX;
- int sizeBytes = largeSize ? U32_SIZE : 1;
- int offsetSize = getIntegerSize(dataSize);
- // The space for header byte, object size, and offset list.
- int headerSize = 1 + sizeBytes + (size + 1) * offsetSize;
- checkCapacity(headerSize);
- // Shift the just-written field data to make room for the header section.
- System.arraycopy(writeBuffer, start, writeBuffer, start + headerSize, dataSize);
- writePos += headerSize;
- writeBuffer[start] = arrayHeader(largeSize, offsetSize);
- writeLong(writeBuffer, start + 1, size, sizeBytes);
- int offsetStart = start + 1 + sizeBytes;
- for (int i = 0; i < size; ++i) {
- writeLong(writeBuffer, offsetStart + i * offsetSize, offsets.get(i), offsetSize);
- }
- writeLong(writeBuffer, offsetStart + size * offsetSize, dataSize, offsetSize);
+ finishWritingArray(start, offsets);
break;
}
case VALUE_STRING:
- byte[] text = parser.getText().getBytes(StandardCharsets.UTF_8);
- boolean longStr = text.length > MAX_SHORT_STR_SIZE;
- checkCapacity((longStr ? 1 + U32_SIZE : 1) + text.length);
- if (longStr) {
- writeBuffer[writePos++] = primitiveHeader(LONG_STR);
- writeLong(writeBuffer, writePos, text.length, U32_SIZE);
- writePos += U32_SIZE;
- } else {
- writeBuffer[writePos++] = shortStrHeader(text.length);
- }
- System.arraycopy(text, 0, writeBuffer, writePos, text.length);
- writePos += text.length;
+ appendString(parser.getText());
break;
case VALUE_NUMBER_INT:
try {
- long l = parser.getLongValue();
- checkCapacity(1 + 8);
- if (l == (byte) l) {
- writeBuffer[writePos++] = primitiveHeader(INT1);
- writeLong(writeBuffer, writePos, l, 1);
- writePos += 1;
- } else if (l == (short) l) {
- writeBuffer[writePos++] = primitiveHeader(INT2);
- writeLong(writeBuffer, writePos, l, 2);
- writePos += 2;
- } else if (l == (int) l) {
- writeBuffer[writePos++] = primitiveHeader(INT4);
- writeLong(writeBuffer, writePos, l, 4);
- writePos += 4;
- } else {
- writeBuffer[writePos++] = primitiveHeader(INT8);
- writeLong(writeBuffer, writePos, l, 8);
- writePos += 8;
- }
+ appendLong(parser.getLongValue());
} catch (InputCoercionException ignored) {
// If the value doesn't fit any integer type, parse it as decimal or floating instead.
parseFloatingPoint(parser);
@@ -263,16 +449,13 @@ public class VariantBuilder {
parseFloatingPoint(parser);
break;
case VALUE_TRUE:
- checkCapacity(1);
- writeBuffer[writePos++] = primitiveHeader(TRUE);
+ appendBoolean(true);
break;
case VALUE_FALSE:
- checkCapacity(1);
- writeBuffer[writePos++] = primitiveHeader(FALSE);
+ appendBoolean(false);
break;
case VALUE_NULL:
- checkCapacity(1);
- writeBuffer[writePos++] = primitiveHeader(NULL);
+ appendNull();
break;
default:
throw new JsonParseException(parser, "Unexpected token " + token);
@@ -290,10 +473,7 @@ public class VariantBuilder {
private void parseFloatingPoint(JsonParser parser) throws IOException {
if (!tryParseDecimal(parser.getText())) {
- checkCapacity(1 + 8);
- writeBuffer[writePos++] = primitiveHeader(DOUBLE);
- writeLong(writeBuffer, writePos, Double.doubleToLongBits(parser.getDoubleValue()), 8);
- writePos += 8;
+ appendDouble(parser.getDoubleValue());
}
}
@@ -308,36 +488,11 @@ public class VariantBuilder {
}
}
BigDecimal d = new BigDecimal(input);
- checkCapacity(2 + 16);
- BigInteger unscaled = d.unscaledValue();
- if (d.scale() <= MAX_DECIMAL4_PRECISION && d.precision() <= MAX_DECIMAL4_PRECISION) {
- writeBuffer[writePos++] = primitiveHeader(DECIMAL4);
- writeBuffer[writePos++] = (byte)d.scale();
- writeLong(writeBuffer, writePos, unscaled.intValueExact(), 4);
- writePos += 4;
- } else if (d.scale() <= MAX_DECIMAL8_PRECISION && d.precision() <= MAX_DECIMAL8_PRECISION) {
- writeBuffer[writePos++] = primitiveHeader(DECIMAL8);
- writeBuffer[writePos++] = (byte)d.scale();
- writeLong(writeBuffer, writePos, unscaled.longValueExact(), 8);
- writePos += 8;
- } else if (d.scale() <= MAX_DECIMAL16_PRECISION && d.precision() <= MAX_DECIMAL16_PRECISION) {
- writeBuffer[writePos++] = primitiveHeader(DECIMAL16);
- writeBuffer[writePos++] = (byte)d.scale();
- // `toByteArray` returns a big-endian representation. We need to copy it reversely and sign
- // extend it to 16 bytes.
- byte[] bytes = unscaled.toByteArray();
- for (int i = 0; i < bytes.length; ++i) {
- writeBuffer[writePos + i] = bytes[bytes.length - 1 - i];
- }
- byte sign = (byte) (bytes[0] < 0 ? -1 : 0);
- for (int i = bytes.length; i < 16; ++i) {
- writeBuffer[writePos + i] = sign;
- }
- writePos += 16;
- } else {
- return false;
+ if (d.scale() <= MAX_DECIMAL16_PRECISION && d.precision() <= MAX_DECIMAL16_PRECISION) {
+ appendDecimal(d);
+ return true;
}
- return true;
+ return false;
}
// The write buffer in building the variant value. Its first `writePos` bytes has been written.
diff --git a/docs/sql-error-conditions.md b/docs/sql-error-conditions.md
index 3f9ae72a50a6..36d7cf58f09f 100644
--- a/docs/sql-error-conditions.md
+++ b/docs/sql-error-conditions.md
@@ -2823,6 +2823,12 @@ To tolerate the error on drop use DROP VARIABLE IF EXISTS.
Cannot construct a Variant larger than 16 MiB. The maximum allowed size of a Variant value is 16 MiB.
+### VARIANT_DUPLICATE_KEY
+
+[SQLSTATE: 22023](sql-error-conditions-sqlstates.html#class-22-data-exception)
+
+Failed to build variant because of a duplicate object key ``<key>``.
+
### VARIANT_SIZE_LIMIT
[SQLSTATE: 22023](sql-error-conditions-sqlstates.html#class-22-data-exception)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
index 94cf7130d485..e252075c9c1c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
@@ -128,6 +128,7 @@ object Cast extends QueryErrorsBase {
case (TimestampType, _: NumericType) => true
case (VariantType, _) => variant.VariantGet.checkDataType(to)
+ case (_, VariantType) => variant.VariantGet.checkDataType(from)
case (ArrayType(fromType, fn), ArrayType(toType, tn)) =>
canAnsiCast(fromType, toType) && resolvableNullability(fn, tn)
@@ -236,6 +237,7 @@ object Cast extends QueryErrorsBase {
case (_: NumericType, _: NumericType) => true
case (VariantType, _) => variant.VariantGet.checkDataType(to)
+ case (_, VariantType) => variant.VariantGet.checkDataType(from)
case (ArrayType(fromType, fn), ArrayType(toType, tn)) =>
canCast(fromType, toType) &&
@@ -1119,6 +1121,7 @@ case class Cast(
} else {
to match {
case dt if dt == from => identity[Any]
+ case VariantType => input => variant.VariantExpressionEvalUtils.castToVariant(input, from)
case _: StringType => castToString(from)
case BinaryType => castToBinary(from)
case DateType => castToDate(from)
@@ -1223,6 +1226,10 @@ case class Cast(
$evPrim = (${CodeGenerator.boxedType(to)})$tmp;
}
"""
+ case VariantType =>
+ val cls = variant.VariantExpressionEvalUtils.getClass.getName.stripSuffix("$")
+ val fromArg = ctx.addReferenceObj("from", from)
+ (c, evPrim, evNull) => code"$evPrim = $cls.castToVariant($c, $fromArg);"
case _: StringType => (c, evPrim, _) => castToStringCode(from, ctx).apply(c, evPrim)
case BinaryType => castToBinaryCode(from)
case DateType => castToDateCode(from, ctx)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
index 74fae91f98a6..d55edcd662b9 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
@@ -19,9 +19,11 @@ package org.apache.spark.sql.catalyst.expressions.variant
import scala.util.control.NonFatal
-import org.apache.spark.sql.catalyst.util.BadRecordException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayData, BadRecordException, MapData}
import org.apache.spark.sql.errors.QueryExecutionErrors
-import org.apache.spark.types.variant.{VariantBuilder, VariantSizeLimitException, VariantUtil}
+import org.apache.spark.sql.types._
+import org.apache.spark.types.variant.{Variant, VariantBuilder, VariantSizeLimitException, VariantUtil}
import org.apache.spark.unsafe.types.{UTF8String, VariantVal}
/**
@@ -41,4 +43,69 @@ object VariantExpressionEvalUtils {
input.toString, BadRecordException(() => input, cause = e))
}
}
+
+ /** Cast a Spark value from `dataType` into the variant type. */
+ def castToVariant(input: Any, dataType: DataType): VariantVal = {
+ val builder = new VariantBuilder
+ buildVariant(builder, input, dataType)
+ val v = builder.result()
+ new VariantVal(v.getValue, v.getMetadata)
+ }
+
+ private def buildVariant(builder: VariantBuilder, input: Any, dataType: DataType): Unit = {
+ if (input == null) {
+ builder.appendNull()
+ return
+ }
+ dataType match {
+ case BooleanType => builder.appendBoolean(input.asInstanceOf[Boolean])
+ case ByteType => builder.appendLong(input.asInstanceOf[Byte])
+ case ShortType => builder.appendLong(input.asInstanceOf[Short])
+ case IntegerType => builder.appendLong(input.asInstanceOf[Int])
+ case LongType => builder.appendLong(input.asInstanceOf[Long])
+ case FloatType => builder.appendFloat(input.asInstanceOf[Float])
+ case DoubleType => builder.appendDouble(input.asInstanceOf[Double])
+ case StringType => builder.appendString(input.asInstanceOf[UTF8String].toString)
+ case BinaryType => builder.appendBinary(input.asInstanceOf[Array[Byte]])
+ case DateType => builder.appendDate(input.asInstanceOf[Int])
+ case TimestampType => builder.appendTimestamp(input.asInstanceOf[Long])
+ case TimestampNTZType => builder.appendTimestampNtz(input.asInstanceOf[Long])
+ case VariantType =>
+ val v = input.asInstanceOf[VariantVal]
+ builder.appendVariant(new Variant(v.getValue, v.getMetadata))
+ case ArrayType(elementType, _) =>
+ val data = input.asInstanceOf[ArrayData]
+ val start = builder.getWritePos
+ val offsets = new java.util.ArrayList[java.lang.Integer](data.numElements())
+ for (i <- 0 until data.numElements()) {
+ offsets.add(builder.getWritePos - start)
+ buildVariant(builder, data.get(i, elementType), elementType)
+ }
+ builder.finishWritingArray(start, offsets)
+ case MapType(StringType, valueType, _) =>
+ val data = input.asInstanceOf[MapData]
+ val keys = data.keyArray()
+ val values = data.valueArray()
+ val start = builder.getWritePos
+ val fields = new java.util.ArrayList[VariantBuilder.FieldEntry](data.numElements())
+ for (i <- 0 until data.numElements()) {
+ val key = keys.getUTF8String(i).toString
+ val id = builder.addKey(key)
+ fields.add(new VariantBuilder.FieldEntry(key, id, builder.getWritePos - start))
+ buildVariant(builder, values.get(i, valueType), valueType)
+ }
+ builder.finishWritingObject(start, fields)
+ case StructType(structFields) =>
+ val data = input.asInstanceOf[InternalRow]
+ val start = builder.getWritePos
+ val fields = new java.util.ArrayList[VariantBuilder.FieldEntry](structFields.length)
+ for (i <- 0 until structFields.length) {
+ val key = structFields(i).name
+ val id = builder.addKey(key)
+ fields.add(new VariantBuilder.FieldEntry(key, id, builder.getWritePos - start))
+ buildVariant(builder, data.get(i, structFields(i).dataType), structFields(i).dataType)
+ }
+ builder.finishWritingObject(start, fields)
+ }
+ }
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
index 24675518646d..9aa1dcd2ef95 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.catalyst.expressions.variant
import java.time.{LocalDateTime, ZoneId, ZoneOffset}
+import scala.reflect.runtime.universe.TypeTag
+
import org.apache.spark.{SparkFunSuite, SparkRuntimeException}
import org.apache.spark.sql.Row
import org.apache.spark.sql.catalyst.analysis.ResolveTimeZone
@@ -797,4 +799,28 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper {
checkCast(Array(primitiveHeader(BINARY), 5, 0, 0, 0, 72, 101, 108, 108, 111), StringType,
"Hello")
}
+
+ test("cast to variant") {
+ def check[T : TypeTag](input: T, expectedJson: String): Unit = {
+ val cast = Cast(Literal.create(input), VariantType, evalMode = EvalMode.ANSI)
+ checkEvaluation(StructsToJson(Map.empty, cast), expectedJson)
+ }
+
+ check(null.asInstanceOf[String], null)
+ for (input <- Seq[Any](false, true, 0.toByte, 1.toShort, 2, 3L, 4.0F, 5.0D)) {
+ check(input, input.toString)
+ }
+ check(Array(null, "a", "b", "c"), """[null,"a","b","c"]""")
+ check(Map("z" -> 1, "y" -> 2, "x" -> 3), """{"x":3,"y":2,"z":1}""")
+ check(Array(parseJson("""{"a": 1,"b": [1, 2, 3]}"""),
+ parseJson("""{"c": true,"d": {"e": "str"}}""")),
+ """[{"a":1,"b":[1,2,3]},{"c":true,"d":{"e":"str"}}]""")
+ val struct = Literal.create(
+ Row(
+ Seq("123", "true", "f"),
+ Map("a" -> "123", "b" -> "true", "c" -> "f"),
+ Row(0)),
+ StructType.fromDDL("c ARRAY<STRING>,b MAP<STRING, STRING>,a STRUCT<i: INT>"))
+ check(struct, """{"a":{"i":0},"b":{"a":"123","b":"true","c":"f"},"c":["123","true","f"]}""")
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org