Copilot commented on code in PR #12989:
URL: https://github.com/apache/gravitino/pull/12989#discussion_r3956458440
##########
lance/lance-common/src/main/java/org/apache/gravitino/lance/common/utils/ArrowUtils.java:
##########
@@ -57,16 +57,46 @@ public static byte[] generateIpcStream(Schema arrowSchema)
throws IOException {
}
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) {
+ while (reader.loadNextBatch()) {
+ if (reader.getVectorSchemaRoot().getRowCount() > 0) {
+ containsRows = true;
+ break;
+ }
+ }
+ }
} catch (Exception e) {
Review Comment:
When `requireEmpty` is true, this loads and materializes potentially large
record batches just to detect whether any rows exist. If the goal is only to
reject non-empty streams, consider a lighter-weight check (e.g., stop at the
first batch with `rowCount > 0` is good, but you may still pay the cost of
decoding earlier empty batches). If Arrow provides a way to detect presence of
record batches / counts without fully loading vectors (or by using a cheaper
reader/config), prefer that to reduce CPU/memory overhead for large inputs.
##########
lance/lance-rest-server/src/test/java/org/apache/gravitino/lance/integration/test/LanceRESTServiceIT.java:
##########
@@ -968,6 +1013,40 @@ void testDeclareTable() {
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[] arrowStreamWithRecord() 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();
+ }
+ }
Review Comment:
This test helper creates a non-empty stream by writing an empty batch first,
then a non-empty batch. That’s useful coverage, but it makes the helper name a
bit misleading and it couples multiple behaviors into one generator. Consider
either renaming it to reflect the two-batch shape (empty + non-empty), or
adding a parameterized helper (like the unit tests do) so the integration test
can explicitly control whether empty batches precede the non-empty one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
