From dcac111300e00c244b2fe42ebb5c7ceffb6dba8b Mon Sep 17 00:00:00 2001 From: Antoni Stachowski Date: Fri, 10 Jan 2025 10:07:53 +0100 Subject: [PATCH] SNOW-1876614: Fix validity checks in arrow struct vectors (#2022) --- .../ThreeFieldStructToTimestampTZConverter.java | 17 ++++++++++------- .../TwoFieldStructToTimestampLTZConverter.java | 8 ++++---- .../TwoFieldStructToTimestampNTZConverter.java | 8 ++++---- .../TwoFieldStructToTimestampTZConverter.java | 12 ++++++------ ...eeFieldStructToTimestampTZConverterTest.java | 2 ++ ...oFieldStructToTimestampLTZConverterTest.java | 2 ++ ...oFieldStructToTimestampNTZConverterTest.java | 2 ++ ...woFieldStructToTimestampTZConverterTest.java | 2 ++ 8 files changed, 32 insertions(+), 21 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverter.java index 929045dd1..364317d7d 100644 --- a/src/main/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverter.java @@ -45,7 +45,10 @@ public ThreeFieldStructToTimestampTZConverter( @Override public boolean isNull(int index) { - return epochs.isNull(index); + return structVector.isNull(index) + || epochs.isNull(index) + || fractions.isNull(index) + || timeZoneIndices.isNull(index); } @Override @@ -54,7 +57,7 @@ public String toString(int index) throws SFException { throw new SFException(ErrorCode.INTERNAL_ERROR, "missing timestamp TZ formatter"); } try { - Timestamp ts = epochs.isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true); + Timestamp ts = isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true); return ts == null ? null : context.getTimestampTZFormatter().format(ts, timeZone, context.getScale(columnIndex)); @@ -65,7 +68,7 @@ public String toString(int index) throws SFException { @Override public byte[] toBytes(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return null; } throw new SFException( @@ -79,7 +82,7 @@ public Object toObject(int index) throws SFException { @Override public Timestamp toTimestamp(int index, TimeZone tz) throws SFException { - return epochs.isNull(index) ? null : getTimestamp(index, tz, false); + return isNull(index) ? null : getTimestamp(index, tz, false); } private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) throws SFException { @@ -98,7 +101,7 @@ private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) thr @Override public Date toDate(int index, TimeZone tz, boolean dateFormat) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return null; } Timestamp ts = getTimestamp(index, TimeZone.getDefault(), false); @@ -116,7 +119,7 @@ public Time toTime(int index) throws SFException { @Override public boolean toBoolean(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return false; } Timestamp val = toTimestamp(index, TimeZone.getDefault()); @@ -127,7 +130,7 @@ public boolean toBoolean(int index) throws SFException { @Override public short toShort(int rowIndex) throws SFException { - if (epochs.isNull(rowIndex)) { + if (isNull(rowIndex)) { return 0; } throw new SFException( diff --git a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverter.java index 6e3904751..e2f7d7c89 100644 --- a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverter.java @@ -41,7 +41,7 @@ public TwoFieldStructToTimestampLTZConverter( @Override public boolean isNull(int index) { - return epochs.isNull(index); + return structVector.isNull(index) || epochs.isNull(index) || fractions.isNull(index); } @Override @@ -51,7 +51,7 @@ public String toString(int index) throws SFException { } try { - Timestamp ts = epochs.isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true); + Timestamp ts = isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true); return ts == null ? null @@ -82,7 +82,7 @@ private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) thr @Override public byte[] toBytes(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return null; } throw new SFException( @@ -111,7 +111,7 @@ public Time toTime(int index) throws SFException { @Override public boolean toBoolean(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return false; } Timestamp val = toTimestamp(index, TimeZone.getDefault()); diff --git a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverter.java index 30467169e..0e282b9ba 100644 --- a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverter.java @@ -42,7 +42,7 @@ public TwoFieldStructToTimestampNTZConverter( @Override public boolean isNull(int index) { - return epochs.isNull(index); + return structVector.isNull(index) || epochs.isNull(index) || fractions.isNull(index); } @Override @@ -51,7 +51,7 @@ public String toString(int index) throws SFException { throw new SFException(ErrorCode.INTERNAL_ERROR, "missing timestamp NTZ formatter"); } try { - Timestamp ts = epochs.isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true); + Timestamp ts = isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true); return ts == null ? null @@ -92,7 +92,7 @@ private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) thr @Override public byte[] toBytes(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return null; } throw new SFException( @@ -119,7 +119,7 @@ public Time toTime(int index) throws SFException { @Override public boolean toBoolean(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return false; } Timestamp val = toTimestamp(index, TimeZone.getDefault()); diff --git a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverter.java index 2068073a7..4abb7afa9 100644 --- a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverter.java @@ -37,7 +37,7 @@ public TwoFieldStructToTimestampTZConverter( @Override public boolean isNull(int index) { - return epochs.isNull(index); + return structVector.isNull(index) || epochs.isNull(index) || timeZoneIndices.isNull(index); } @Override @@ -59,7 +59,7 @@ public Object toObject(int index) throws SFException { @Override public Timestamp toTimestamp(int index, TimeZone tz) throws SFException { - return epochs.isNull(index) ? null : getTimestamp(index, tz); + return isNull(index) ? null : getTimestamp(index, tz); } private Timestamp getTimestamp(int index, TimeZone tz) throws SFException { @@ -76,7 +76,7 @@ private Timestamp getTimestamp(int index, TimeZone tz) throws SFException { @Override public Date toDate(int index, TimeZone tz, boolean dateFormat) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return null; } Timestamp ts = getTimestamp(index, TimeZone.getDefault()); @@ -94,7 +94,7 @@ public Time toTime(int index) throws SFException { @Override public boolean toBoolean(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return false; } Timestamp val = toTimestamp(index, TimeZone.getDefault()); @@ -105,7 +105,7 @@ public boolean toBoolean(int index) throws SFException { @Override public byte[] toBytes(int index) throws SFException { - if (epochs.isNull(index)) { + if (isNull(index)) { return null; } throw new SFException( @@ -114,7 +114,7 @@ public byte[] toBytes(int index) throws SFException { @Override public short toShort(int rowIndex) throws SFException { - if (epochs.isNull(rowIndex)) { + if (isNull(rowIndex)) { return 0; } throw new SFException( diff --git a/src/test/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverterTest.java b/src/test/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverterTest.java index 09cd4a587..0f1db77e6 100644 --- a/src/test/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverterTest.java +++ b/src/test/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverterTest.java @@ -169,9 +169,11 @@ public void testTimestampTZ( seconds.setSafe(j, testSecondsInt64[i]); nanos.setSafe(j, testNanos[i]); timeZoneIdx.setSafe(j, testTimeZoneIndices[i++]); + structVector.setIndexDefined(j); } j++; } + structVector.setValueCount(j); ArrowVectorConverter converter = new ThreeFieldStructToTimestampTZConverter(structVector, 0, this); diff --git a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverterTest.java b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverterTest.java index 4fd4f07f3..8297452b0 100644 --- a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverterTest.java +++ b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverterTest.java @@ -152,9 +152,11 @@ public void testTimestampLTZ( } else { epochs.setSafe(j, testSecondsInt64[i]); fractions.setSafe(j, testNanoSecs[i++]); + structVector.setIndexDefined(j); } j++; } + structVector.setValueCount(j); ArrowVectorConverter converter = new TwoFieldStructToTimestampLTZConverter(structVector, 0, this); diff --git a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverterTest.java b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverterTest.java index 3b84176e4..e4c495b7d 100644 --- a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverterTest.java +++ b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverterTest.java @@ -175,9 +175,11 @@ public void testTimestampNTZ( } else { epochs.setSafe(j, testSecondsInt64[i]); fractions.setSafe(j, testNanoSecs[i++]); + structVector.setIndexDefined(j); } j++; } + structVector.setValueCount(j); ArrowVectorConverter converter = new TwoFieldStructToTimestampNTZConverter(structVector, 0, this); diff --git a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverterTest.java b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverterTest.java index 767938d06..f3b7406a9 100644 --- a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverterTest.java +++ b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverterTest.java @@ -114,9 +114,11 @@ public void testTimestampTZ(String tz) throws SFException { } else { epoch.setSafe(j, testEpochesInt64[i]); timeZoneIdx.setSafe(j, testTimeZoneIndices[i++]); + structVector.setIndexDefined(j); } j++; } + structVector.setValueCount(j); ArrowVectorConverter converter = new TwoFieldStructToTimestampTZConverter(structVector, 0, this);