This is an automated email from the ASF dual-hosted git repository.
jerryshao pushed a commit to branch branch-1.3
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/branch-1.3 by this push:
new 54486a58f7 [Cherry-pick to branch-1.3] [#12988] fix(lance): reject
nonempty Arrow input before table creation (#12989) (#13025)
54486a58f7 is described below
commit 54486a58f736383844ed0dea1d9634a3a820aaa5
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Wed Sep 9 17:27:08 2026 +0800
[Cherry-pick to branch-1.3] [#12988] fix(lance): reject nonempty Arrow
input before table creation (#12989) (#13025)
**Cherry-pick Information:**
- Original commit: ade9e1b8caecaed925293cd183018e776ba7eb8f
- Target branch: `branch-1.3`
- Status: ✅ Clean cherry-pick (no conflicts)
Co-authored-by: Qi Yu <[email protected]>
---
.../gravitino/GravitinoLanceTableOperations.java | 4 +-
.../gravitino/lance/common/utils/ArrowUtils.java | 54 ++++++++
.../lance/common/utils/TestArrowUtils.java | 153 +++++++++++++++++++++
.../lance/integration/test/LanceRESTServiceIT.java | 79 +++++++++++
4 files changed, 288 insertions(+), 2 deletions(-)
diff --git
a/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/ops/gravitino/GravitinoLanceTableOperations.java
b/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/ops/gravitino/GravitinoLanceTableOperations.java
index 50ba09df30..bd3a67fcc0 100644
---
a/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/ops/gravitino/GravitinoLanceTableOperations.java
+++
b/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/ops/gravitino/GravitinoLanceTableOperations.java
@@ -167,11 +167,11 @@ public class GravitinoLanceTableOperations implements
LanceTableOperations {
Preconditions.checkArgument(
nsId.levels() == 3, "Expected at 3-level namespace but got: %s",
nsId.levels());
- // Parser column information.
+ // Reject unsupported record batches before any metadata or storage
mutation.
List<Column> columns = Lists.newArrayList();
if (arrowStreamBody != null) {
org.apache.arrow.vector.types.pojo.Schema schema =
- ArrowUtils.parseArrowIpcStream(arrowStreamBody);
+ ArrowUtils.parseSchemaOnlyIpcStream(arrowStreamBody);
columns = extractColumns(schema);
}
diff --git
a/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/utils/ArrowUtils.java
b/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/utils/ArrowUtils.java
index 5d8508ee45..b72ed08f7e 100644
---
a/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/utils/ArrowUtils.java
+++
b/lance/lance-common/src/main/java/org/apache/gravitino/lance/common/utils/ArrowUtils.java
@@ -23,11 +23,17 @@ import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.channels.Channels;
+import org.apache.arrow.flatbuf.Message;
+import org.apache.arrow.flatbuf.MessageHeader;
+import org.apache.arrow.flatbuf.RecordBatch;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamReader;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
+import org.apache.arrow.vector.ipc.ReadChannel;
+import org.apache.arrow.vector.ipc.message.MessageMetadataResult;
+import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.Schema;
public class ArrowUtils {
@@ -57,16 +63,64 @@ public class ArrowUtils {
}
public static Schema parseArrowIpcStream(byte[] stream) {
+ return parseArrowIpcStream(stream, false);
+ }
+
+ /**
+ * Parses a schema-only Arrow IPC stream, rejecting record batches
containing rows.
+ *
+ * @param stream the Arrow IPC stream
+ * @return the stream schema
+ * @throws UnsupportedOperationException if any record batch contains rows
+ * @throws IllegalArgumentException if the stream cannot be parsed
+ */
+ public static Schema parseSchemaOnlyIpcStream(byte[] stream) {
+ return parseArrowIpcStream(stream, true);
+ }
+
+ private static Schema parseArrowIpcStream(byte[] stream, boolean
requireEmpty) {
Schema schema;
+ boolean containsRows = false;
try (BufferAllocator allocator = new RootAllocator();
ByteArrayInputStream bais = new ByteArrayInputStream(stream);
ArrowStreamReader reader = new ArrowStreamReader(bais, allocator)) {
schema = reader.getVectorSchemaRoot().getSchema();
+ if (requireEmpty) {
+ containsRows = containsRecordBatchRows(bais);
+ }
} catch (Exception e) {
throw new IllegalArgumentException("Failed to parse Arrow IPC stream",
e);
}
Preconditions.checkArgument(schema != null, "No schema found in Arrow IPC
stream");
+ if (containsRows) {
+ throw new UnsupportedOperationException(
+ "CreateTable only supports schema-only Arrow streams; "
+ + "write records through a Lance client or engine after
creation");
+ }
return schema;
}
+
+ private static boolean containsRecordBatchRows(ByteArrayInputStream input)
throws IOException {
+ // The schema reader has consumed the schema message. Inspect only
subsequent message headers;
+ // skipping bodies avoids allocating or decoding vectors, including
dictionary values.
+ try (ReadChannel channel = new ReadChannel(Channels.newChannel(input))) {
+ MessageMetadataResult metadata;
+ while ((metadata = MessageSerializer.readMessage(channel)) != null) {
+ Message message = metadata.getMessage();
+ if (message.headerType() == MessageHeader.RecordBatch) {
+ RecordBatch batch = (RecordBatch) message.header(new RecordBatch());
+ Preconditions.checkArgument(batch.length() >= 0, "Invalid Arrow
record batch row count");
+ if (batch.length() > 0) {
+ return true;
+ }
+ } else if (message.headerType() != MessageHeader.DictionaryBatch) {
+ throw new IOException("Unexpected Arrow message type: " +
message.headerType());
+ }
+ Preconditions.checkArgument(message.bodyLength() >= 0, "Invalid Arrow
message body length");
+ input.skipNBytes(message.bodyLength());
+ }
+ return false;
+ }
+ }
}
diff --git
a/lance/lance-common/src/test/java/org/apache/gravitino/lance/common/utils/TestArrowUtils.java
b/lance/lance-common/src/test/java/org/apache/gravitino/lance/common/utils/TestArrowUtils.java
index 43f0bf6ec6..7f1a7ec482 100644
---
a/lance/lance-common/src/test/java/org/apache/gravitino/lance/common/utils/TestArrowUtils.java
+++
b/lance/lance-common/src/test/java/org/apache/gravitino/lance/common/utils/TestArrowUtils.java
@@ -18,9 +18,25 @@
*/
package org.apache.gravitino.lance.common.utils;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.nio.channels.Channels;
import java.util.Arrays;
+import java.util.List;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import
org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider;
+import org.apache.arrow.vector.ipc.ArrowStreamWriter;
+import org.apache.arrow.vector.ipc.ReadChannel;
+import org.apache.arrow.vector.ipc.message.MessageMetadataResult;
+import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
@@ -39,4 +55,141 @@ public class TestArrowUtils {
Assertions.assertEquals(schema, parsedSchema);
}
+ /** Verifies schema-only streams and zero-row batches remain supported. */
+ @Test
+ public void testSchemaOnlyStreams() throws Exception {
+ Schema expected = new Schema(List.of(Field.nullable("id", new
ArrowType.Int(32, true))));
+ Assertions.assertEquals(expected,
ArrowUtils.parseSchemaOnlyIpcStream(streamWithRows()));
+ Assertions.assertEquals(expected,
ArrowUtils.parseSchemaOnlyIpcStream(streamWithRows(0, 0)));
+ }
+
+ /** Verifies that a non-empty batch is rejected, including after empty
batches. */
+ @Test
+ public void testRejectRecordBatchesWithRows() throws Exception {
+ for (byte[] stream : List.of(streamWithRows(1), streamWithRows(0, 1))) {
+ UnsupportedOperationException exception =
+ Assertions.assertThrows(
+ UnsupportedOperationException.class,
+ () -> ArrowUtils.parseSchemaOnlyIpcStream(stream));
+ Assertions.assertTrue(exception.getMessage().contains("schema-only"));
+ // Existing callers of the general schema parser retain their previous
behavior.
+ Assertions.assertEquals(1,
ArrowUtils.parseArrowIpcStream(stream).getFields().size());
+ }
+ }
+
+ /** Verifies malformed input is reported as invalid rather than as
unsupported data. */
+ @Test
+ public void testRejectMalformedSchemaOnlyStream() {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> ArrowUtils.parseSchemaOnlyIpcStream(new byte[] {1, 2, 3}));
+ }
+
+ /** Verifies non-empty batches are rejected from metadata without decoding
their bodies. */
+ @Test
+ public void testRejectRowsBeforeReadingBatchBody() throws Exception {
+ byte[] stream = streamWithRows(1);
+ try (ReadChannel channel =
+ new ReadChannel(Channels.newChannel(new
ByteArrayInputStream(stream)))) {
+ MessageSerializer.deserializeSchema(channel);
+ MessageMetadataResult batch = MessageSerializer.readMessage(channel);
+ Assertions.assertTrue(batch.getMessageBodyLength() > 0);
+ byte[] headersOnly = Arrays.copyOf(stream, (int) channel.bytesRead());
+ Assertions.assertThrows(
+ UnsupportedOperationException.class,
+ () -> ArrowUtils.parseSchemaOnlyIpcStream(headersOnly));
+ }
+ }
+
+ /** Verifies dictionary values are skipped and do not count as table rows. */
+ @Test
+ public void testDictionaryBatches() throws Exception {
+ for (int rows : new int[] {0, 1}) {
+ byte[] stream = dictionaryStreamWithRows(rows);
+ if (rows == 0) {
+ Assertions.assertEquals(
+ ArrowUtils.parseArrowIpcStream(stream),
ArrowUtils.parseSchemaOnlyIpcStream(stream));
+ } else {
+ Assertions.assertThrows(
+ UnsupportedOperationException.class, () ->
ArrowUtils.parseSchemaOnlyIpcStream(stream));
+ }
+ }
+ }
+
+ /** Verifies skipping a truncated dictionary body still reports malformed
input. */
+ @Test
+ public void testRejectTruncatedDictionaryBody() throws Exception {
+ byte[] stream = dictionaryStreamWithRows(0);
+ try (ReadChannel channel =
+ new ReadChannel(Channels.newChannel(new
ByteArrayInputStream(stream)))) {
+ MessageSerializer.deserializeSchema(channel);
+ MessageMetadataResult dictionary =
MessageSerializer.readMessage(channel);
+ Assertions.assertTrue(dictionary.getMessageBodyLength() > 0);
+ byte[] truncated = Arrays.copyOf(stream, (int) channel.bytesRead() + 1);
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () ->
ArrowUtils.parseSchemaOnlyIpcStream(truncated));
+ }
+ }
+
+ /** Verifies a schema message cannot appear where a record batch is
expected. */
+ @Test
+ public void testRejectUnexpectedMessage() throws Exception {
+ byte[] stream = streamWithRows();
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (ReadChannel channel =
+ new ReadChannel(Channels.newChannel(new
ByteArrayInputStream(stream)))) {
+ MessageSerializer.deserializeSchema(channel);
+ output.write(stream, 0, (int) channel.bytesRead());
+ output.write(stream);
+ }
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> ArrowUtils.parseSchemaOnlyIpcStream(output.toByteArray()));
+ }
+
+ private byte[] dictionaryStreamWithRows(int rows) throws Exception {
+ DictionaryEncoding encoding = new DictionaryEncoding(0, false, new
ArrowType.Int(32, true));
+ Schema schema =
+ new Schema(
+ List.of(new Field("id", new FieldType(true,
encoding.getIndexType(), encoding), null)));
+ try (RootAllocator allocator = new RootAllocator();
+ VarCharVector values = new VarCharVector("values", allocator);
+ VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
+ ByteArrayOutputStream output = new ByteArrayOutputStream()) {
+ values.allocateNew();
+ values.setSafe(0, new byte[] {42});
+ values.setValueCount(1);
+ MapDictionaryProvider dictionaries =
+ new MapDictionaryProvider(new Dictionary(values, encoding));
+ try (ArrowStreamWriter writer = new ArrowStreamWriter(root,
dictionaries, output)) {
+ root.allocateNew();
+ ((IntVector) root.getVector("id")).setSafe(0, 0);
+ root.setRowCount(rows);
+ writer.start();
+ writer.writeBatch();
+ writer.end();
+ }
+ return output.toByteArray();
+ }
+ }
+
+ private byte[] streamWithRows(int... batches) throws Exception {
+ Schema schema = new Schema(List.of(Field.nullable("id", new
ArrowType.Int(32, true))));
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ ArrowStreamWriter writer = new ArrowStreamWriter(root, null, output)) {
+ root.allocateNew();
+ writer.start();
+ for (int rows : batches) {
+ for (int i = 0; i < rows; i++) {
+ ((IntVector) root.getVector("id")).setSafe(i, i);
+ }
+ root.setRowCount(rows);
+ writer.writeBatch();
+ }
+ writer.end();
+ return output.toByteArray();
+ }
+ }
}
diff --git
a/lance/lance-rest-server/src/test/java/org/apache/gravitino/lance/integration/test/LanceRESTServiceIT.java
b/lance/lance-rest-server/src/test/java/org/apache/gravitino/lance/integration/test/LanceRESTServiceIT.java
index dcaa97c7f2..b4762fedf0 100644
---
a/lance/lance-rest-server/src/test/java/org/apache/gravitino/lance/integration/test/LanceRESTServiceIT.java
+++
b/lance/lance-rest-server/src/test/java/org/apache/gravitino/lance/integration/test/LanceRESTServiceIT.java
@@ -21,6 +21,7 @@ package org.apache.gravitino.lance.integration.test;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
+import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.charset.Charset;
@@ -34,6 +35,10 @@ import java.util.Objects;
import java.util.Set;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.commons.io.FileUtils;
@@ -52,6 +57,7 @@ import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.lance.Dataset;
import org.lance.namespace.LanceNamespace;
import org.lance.namespace.client.apache.ApiClient;
import org.lance.namespace.client.apache.ApiException;
@@ -449,6 +455,45 @@ public class LanceRESTServiceIT extends BaseIT {
assertLanceErrorCode(exception, ErrorCode.NAMESPACE_NOT_FOUND);
}
+ /** Verifies unsupported input cannot silently discard records or destroy an
existing table. */
+ @Test
+ void testCreateRejectsNonEmptyArrowWithoutSideEffects() throws IOException {
+ catalog = createCatalog(CATALOG_NAME);
+ createSchema();
+ byte[] data = arrowStreamWithEmptyThenNonEmptyBatch();
+ for (String mode : List.of("create", "exist_ok")) {
+ String name = "nonempty_" + mode;
+ Path location = tempDir.resolve(name);
+ assertNonEmptyCreateRejected(
+ List.of(CATALOG_NAME, SCHEMA_NAME, name), location.toString(), data,
mode);
+ Assertions.assertFalse(
+ catalog.asTableCatalog().tableExists(NameIdentifier.of(SCHEMA_NAME,
name)));
+ Assertions.assertFalse(Files.exists(location));
+ }
+
+ String original = "nonempty_overwrite";
+ List<String> ids = List.of(CATALOG_NAME, SCHEMA_NAME, original);
+ String location = tempDir.resolve(original).toString();
+ try (VectorSchemaRoot root =
+ VectorSchemaRoot.of(
+ new IntVector("id", allocator), new VarCharVector("value",
allocator))) {
+ createTable(
+ ids, location, Map.of(),
ArrowUtils.generateIpcStream(root.getSchema()), "create");
+ }
+ assertNonEmptyCreateRejected(ids, location, data, "overwrite");
+ DescribeTableRequest describe = new DescribeTableRequest();
+ describe.setId(ids);
+ Assertions.assertEquals(
+ List.of("id", "value"),
+ ns.describeTable(describe).getSchema().getFields().stream()
+ .map(JsonArrowField::getName)
+ .toList());
+ try (Dataset dataset = Dataset.open().uri(location).build()) {
+ Assertions.assertEquals(0, dataset.countRows());
+ Assertions.assertEquals(2, dataset.getSchema().getFields().size());
+ }
+ }
+
@Test
void testCreateTable() throws IOException {
catalog = createCatalog(CATALOG_NAME);
@@ -968,6 +1013,40 @@ public class LanceRESTServiceIT extends BaseIT {
Assertions.assertFalse(new File(anotherLocation).exists());
}
+ private void assertNonEmptyCreateRejected(
+ List<String> ids, String location, byte[] data, String mode) {
+ ApiException error =
+ Assertions.assertThrows(
+ ApiException.class,
+ () ->
+ createTableApi()
+ .createTable(
+ String.join(DELIMITER, ids),
+ data,
+ DELIMITER,
+ mode,
+ null,
+ null,
+ Map.of(LanceConstants.LANCE_TABLE_LOCATION_HEADER,
location)));
+ Assertions.assertEquals(406, error.getCode());
+ }
+
+ private byte[] arrowStreamWithEmptyThenNonEmptyBatch() throws IOException {
+ try (VectorSchemaRoot root = VectorSchemaRoot.of(new IntVector("id",
allocator));
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ ArrowStreamWriter writer = new ArrowStreamWriter(root, null, output)) {
+ root.allocateNew();
+ root.setRowCount(0);
+ writer.start();
+ writer.writeBatch();
+ ((IntVector) root.getVector("id")).setSafe(0, 42);
+ root.setRowCount(1);
+ writer.writeBatch();
+ writer.end();
+ return output.toByteArray();
+ }
+ }
+
private CreateTableResponse createTable(
List<String> ids,
String location,
