From 3d6cedb25fd81f2415b4f46b025f3624e8501550 Mon Sep 17 00:00:00 2001 From: Akshat Jain Date: Fri, 9 Aug 2024 11:39:53 +0530 Subject: [PATCH] Fix IndexOutOfBoundsException for MSQ window function queries with empty RAC (#16865) * Fix IndexOutOfBoundsException for MSQ window function queries with empty RAC --- .../rowsandcols/ConcatRowsAndColumns.java | 9 +++ .../rowsandcols/ConcatRowsAndColumnsTest.java | 57 +++++++++++++++++++ .../sql/calcite/DrillWindowQueryTest.java | 7 +++ .../empty_over_clause/single_empty_over_3.e | 0 .../empty_over_clause/single_empty_over_3.q | 4 ++ 5 files changed, 77 insertions(+) create mode 100644 sql/src/test/resources/drill/window/queries/druid_queries/empty_over_clause/single_empty_over_3.e create mode 100644 sql/src/test/resources/drill/window/queries/druid_queries/empty_over_clause/single_empty_over_3.q diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java index 523a982fbaf..c6ced60849d 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java @@ -19,6 +19,7 @@ package org.apache.druid.query.rowsandcols; +import com.google.common.base.Preconditions; import org.apache.druid.java.util.common.ISE; import org.apache.druid.query.rowsandcols.column.Column; import org.apache.druid.query.rowsandcols.column.ColumnAccessor; @@ -30,6 +31,7 @@ import javax.annotation.Nullable; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; +import java.util.Collections; import java.util.Comparator; import java.util.LinkedHashMap; import java.util.Map; @@ -53,6 +55,7 @@ public class ConcatRowsAndColumns implements RowsAndColumns ArrayList racBuffer ) { + Preconditions.checkNotNull(racBuffer, "racBuffer cannot be null"); this.racBuffer = racBuffer; int numRows = 0; @@ -76,6 +79,9 @@ public class ConcatRowsAndColumns implements RowsAndColumns @Override public Collection getColumnNames() { + if (racBuffer.isEmpty()) { + return Collections.emptySet(); + } return racBuffer.get(0).getColumnNames(); } @@ -92,6 +98,9 @@ public class ConcatRowsAndColumns implements RowsAndColumns if (columnCache.containsKey(name)) { return columnCache.get(name); } else { + if (racBuffer.isEmpty()) { + return null; + } final Column firstCol = racBuffer.get(0).findColumn(name); if (firstCol == null) { for (int i = 1; i < racBuffer.size(); ++i) { diff --git a/processing/src/test/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumnsTest.java b/processing/src/test/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumnsTest.java index b95f750632d..0bf0450114b 100644 --- a/processing/src/test/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumnsTest.java +++ b/processing/src/test/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumnsTest.java @@ -19,7 +19,13 @@ package org.apache.druid.query.rowsandcols; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.query.rowsandcols.column.IntArrayColumn; +import org.junit.Assert; +import org.junit.Test; + import java.util.ArrayList; +import java.util.Arrays; import java.util.function.Function; public class ConcatRowsAndColumnsTest extends RowsAndColumnsTestBase @@ -42,4 +48,55 @@ public class ConcatRowsAndColumnsTest extends RowsAndColumnsTestBase return new ConcatRowsAndColumns(theRac); }; + + @Test + public void testConstructorWithNullRacBuffer() + { + final NullPointerException e = Assert.assertThrows( + NullPointerException.class, + () -> new ConcatRowsAndColumns(null) + ); + Assert.assertEquals("racBuffer cannot be null", e.getMessage()); + } + + @Test + public void testFindColumn() + { + MapOfColumnsRowsAndColumns rac = MapOfColumnsRowsAndColumns.fromMap( + ImmutableMap.of( + "column1", new IntArrayColumn(new int[]{1, 2, 3, 4, 5, 6}), + "column2", new IntArrayColumn(new int[]{6, 5, 4, 3, 2, 1}) + ) + ); + ConcatRowsAndColumns apply = MAKER.apply(rac); + Assert.assertEquals(1, apply.findColumn("column1").toAccessor().getInt(0)); + Assert.assertEquals(6, apply.findColumn("column2").toAccessor().getInt(0)); + } + + @Test + public void testFindColumnWithEmptyRacBuffer() + { + ConcatRowsAndColumns concatRowsAndColumns = new ConcatRowsAndColumns(new ArrayList<>()); + Assert.assertNull(concatRowsAndColumns.findColumn("columnName")); + } + + @Test + public void testGetColumns() + { + MapOfColumnsRowsAndColumns rac = MapOfColumnsRowsAndColumns.fromMap( + ImmutableMap.of( + "column1", new IntArrayColumn(new int[]{0, 0, 0, 1, 1, 2, 4, 4, 4}), + "column2", new IntArrayColumn(new int[]{3, 54, 21, 1, 5, 54, 2, 3, 92}) + ) + ); + ConcatRowsAndColumns apply = MAKER.apply(rac); + Assert.assertEquals(Arrays.asList("column1", "column2"), new ArrayList<>(apply.getColumnNames())); + } + + @Test + public void testGetColumnsWithEmptyRacBuffer() + { + ConcatRowsAndColumns concatRowsAndColumns = new ConcatRowsAndColumns(new ArrayList<>()); + Assert.assertTrue(concatRowsAndColumns.getColumnNames().isEmpty()); + } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java index 4a2f0945087..9bf56d97d38 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java @@ -7622,6 +7622,13 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest windowQueryTest(); } + @DrillTest("druid_queries/empty_over_clause/single_empty_over_3") + @Test + public void test_empty_over_single_empty_over_3() + { + windowQueryTest(); + } + @DrillTest("druid_queries/empty_over_clause/multiple_empty_over_1") @Test public void test_empty_over_multiple_empty_over_1() diff --git a/sql/src/test/resources/drill/window/queries/druid_queries/empty_over_clause/single_empty_over_3.e b/sql/src/test/resources/drill/window/queries/druid_queries/empty_over_clause/single_empty_over_3.e new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sql/src/test/resources/drill/window/queries/druid_queries/empty_over_clause/single_empty_over_3.q b/sql/src/test/resources/drill/window/queries/druid_queries/empty_over_clause/single_empty_over_3.q new file mode 100644 index 00000000000..ac8c5e930d3 --- /dev/null +++ b/sql/src/test/resources/drill/window/queries/druid_queries/empty_over_clause/single_empty_over_3.q @@ -0,0 +1,4 @@ +select countryName, row_number() over () as c1 +from wikipedia +where countryName in ('non-existent-country') +group by countryName, cityName, channel