Fail hard when there are multi-valued blocks

swallez · swallez · commit c7cd12d02655 · 2024-06-27T17:52:33.000+02:00
diff --git a/x-pack/plugin/esql/arrow/src/main/java/org/elasticsearch/xpack/esql/arrow/ArrowResponse.java b/x-pack/plugin/esql/arrow/src/main/java/org/elasticsearch/xpack/esql/arrow/ArrowResponse.java
@@ -64,7 +64,16 @@ public ArrowResponse(List<Column> columns, List<Page> pages) {
         currentSegment = new SchemaResponse(this);
         List<ResponseSegment> rest = new ArrayList<>(pages.size());
         for (int p = 0; p < pages.size(); p++) {
-            rest.add(new PageResponse(this, pages.get(p)));
+            var page = pages.get(p);
+            rest.add(new PageResponse(this, page));
+            // Multivalued fields are not supported yet.
+            for (int b = 0; b < page.getBlockCount(); b++) {
+                if (page.getBlock(b).mayHaveMultivaluedFields()) {
+                    throw new IllegalArgumentException(
+                        "ES|QL response field [" + columns.get(b).name + "] is multi-valued. This isn't supported yet by the Arrow format"
+                    );
+                }
+            }
         }
         rest.add(new EndResponse(this));
         segments = rest.iterator();
diff --git a/x-pack/plugin/esql/arrow/src/test/java/org/elasticsearch/xpack/esql/arrow/ArrowResponseTests.java b/x-pack/plugin/esql/arrow/src/test/java/org/elasticsearch/xpack/esql/arrow/ArrowResponseTests.java
@@ -260,6 +260,41 @@ public void testSingleBlock() throws IOException {
         compareEsqlAndArrow(testCase);
     }
 
+    /**
+     * Test that multivalued arrays are rejected
+     */
+    public void testMultivaluedField() throws IOException {
+        IntBlock.Builder builder = BLOCK_FACTORY.newIntBlockBuilder(0);
+        builder.appendInt(42);
+        builder.appendInt(43);
+        builder.beginPositionEntry();
+        builder.appendInt(44);
+        builder.appendInt(45);
+        builder.endPositionEntry();
+        builder.appendInt(46);
+        IntBlock block = builder.build();
+
+        // Consistency check
+        assertTrue(block.mayHaveMultivaluedFields());
+        assertEquals(0, block.getFirstValueIndex(0));
+        assertEquals(1, block.getValueCount(0));
+        assertEquals(1, block.getValueCount(1));
+
+        assertEquals(2, block.getValueCount(2));
+        assertEquals(2, block.getFirstValueIndex(2));
+        assertEquals(45, block.getInt(block.getFirstValueIndex(2) + 1));
+
+        assertEquals(4, block.getFirstValueIndex(3));
+
+        var column = TestColumn.create("some-field", "integer");
+        TestCase testCase = new TestCase(List.of(column), List.of(new TestPage(List.of(new TestBlock(column, block, Density.Dense)))));
+
+        IllegalArgumentException exc = assertThrows(IllegalArgumentException.class, () -> compareEsqlAndArrow(testCase));
+
+        assertEquals("ES|QL response field [some-field] is multi-valued. This isn't supported yet by the Arrow format", exc.getMessage());
+
+    }
+
     /**
      * Test a random set of types/columns/pages/densities
      */
