This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/seatunnel.git
The following commit(s) were added to refs/heads/dev by this push:
new 25bcc53fd1 [Fix][Connector-V2][Lance] Write one fragment per flush
instead of one per buffered row (#12737) (#12738)
25bcc53fd1 is described below
commit 25bcc53fd1b2850b79127980de4f9bcb41e545cb
Author: tzy <[email protected]>
AuthorDate: Tue Oct 13 16:06:13 2026 +0000
[Fix][Connector-V2][Lance] Write one fragment per flush instead of one per
buffered row (#12737) (#12738)
---
docs/en/connectors/sink/Lance.md | 16 ++-
docs/zh/connectors/sink/Lance.md | 14 +-
.../seatunnel/lance/config/LanceSinkConfig.java | 7 +
.../seatunnel/lance/config/LanceSinkOptions.java | 7 +
.../seatunnel/lance/sink/LanceSinkFactory.java | 4 +
.../seatunnel/lance/sink/LanceSinkWriter.java | 21 ++-
.../seatunnel/lance/utils/FragmentConverter.java | 70 ++++++----
.../seatunnel/lance/LanceFactoryTest.java | 33 +++++
.../lance/sink/LanceSinkWriterBatchTest.java | 98 ++++++++++++++
.../lance/utils/FragmentConverterTest.java | 150 +++++++++++++++++++++
10 files changed, 377 insertions(+), 43 deletions(-)
diff --git a/docs/en/connectors/sink/Lance.md b/docs/en/connectors/sink/Lance.md
index e60f872ff9..9dfb8a484e 100644
--- a/docs/en/connectors/sink/Lance.md
+++ b/docs/en/connectors/sink/Lance.md
@@ -52,6 +52,7 @@ It provides a sink only; there is no Lance source connector.
| namespace_ids | list | no | [] |
Lance namespace path segments used when resolving the target table namespace.
|
| root_namespace_path | string | no | /tmp |
Root path used by the Lance namespace.
|
| table | string | no | test |
Target Lance table name. If set, it overrides the upstream table name.
|
+| lance.write.batch-size | int | no | 10000 |
Rows buffered in the writer before flushing one Lance fragment.
|
| lance.write.max-rows-per-file | int | no | 10 |
Maximum number of rows written to one Lance file.
|
| lance.write.max-rows-per-group | int | no | 20 |
Maximum number of rows written to one Lance row group.
|
| lance.write.max-bytes-per-file | long | no | 20480 |
Maximum number of bytes written to one Lance file.
|
@@ -91,6 +92,14 @@ to create and write files under this path.
The target Lance table name. When it is not set, the connector uses the
upstream table name if one is available. The default option value is `test`.
+### lance.write.batch-size
+
+Number of SeaTunnel rows buffered in memory before the writer converts them
into **one** Lance fragment and appends it. The default is `10000` (the writer
previously used a fixed buffer of `1000` rows), and the value must be at least
`1`.
+
+Lance stores data by column, so a fragment of `10000` rows can still be small.
Too few rows per fragment produces many small columnar files and increases read
and write amplification later. A larger value means fewer, larger fragments,
but each writer holds more rows in memory before flushing. Raise this value
when the job has enough memory, and lower it when memory is tight, for example
with wide rows or vector columns.
+
+The writer also flushes on checkpoint (`prepareCommit`) and when the sink
closes, even if the buffer has not reached this size.
+
### lance.write.mode
Controls how Lance writes the dataset. The default is `CREATE`. Use a value
@@ -144,7 +153,7 @@ Lance uses the Apache Arrow type system. The sink creates
the Lance schema from
:::tip
-The sink does not interpret `UPDATE` / `DELETE` row kinds as CDC operations —
every upstream row is appended to the Lance dataset using the configured
`lance.write.mode`. In streaming mode, the writer flushes the in-memory row
buffer to Lance on every checkpoint.
+The sink does not interpret `UPDATE` / `DELETE` row kinds as CDC operations —
every upstream row is appended to the Lance dataset using the configured
`lance.write.mode`. In streaming mode, the writer flushes the in-memory row
buffer to Lance when `lance.write.batch-size` is reached and on every
checkpoint.
:::
@@ -194,8 +203,8 @@ sink {
### Append Mode With Larger File Fragments
`APPEND` mode keeps the existing dataset and adds new rows. Increase
-`lance.write.max-rows-per-file` and `lance.write.max-bytes-per-file` to reduce
-the number of Lance fragments when appending large batches.
+`lance.write.batch-size` to reduce the number of Lance fragments when appending
+large batches and a larger in-memory buffer is acceptable.
```hocon
env {
@@ -223,6 +232,7 @@ sink {
namespace_id = "root"
table = "lance_sink_table"
lance.write.mode = "APPEND"
+ lance.write.batch-size = 100000
lance.write.max-rows-per-file = 100000
lance.write.max-rows-per-group = 5000
lance.write.max-bytes-per-file = 134217728
diff --git a/docs/zh/connectors/sink/Lance.md b/docs/zh/connectors/sink/Lance.md
index 95636f1d6c..0c1aec791d 100644
--- a/docs/zh/connectors/sink/Lance.md
+++ b/docs/zh/connectors/sink/Lance.md
@@ -50,6 +50,7 @@ Lance sink 用于把 SeaTunnel 数据写入 Lance 数据集。它可以根据上
| namespace_ids | list | 否 | [] |
解析目标表 namespace 时使用的 namespace 路径片段。 |
| root_namespace_path | string | 否 | /tmp |
Lance namespace 的根路径。 |
| table | string | 否 | test |
目标 Lance 表名。设置后会覆盖上游表名。 |
+| lance.write.batch-size | int | 否 | 10000 |
Writer 内存攒批行数，满批后落成一个 fragment。 |
| lance.write.max-rows-per-file | int | 否 | 10 |
单个 Lance 文件最多写入的行数。 |
| lance.write.max-rows-per-group | int | 否 | 20 |
单个 Lance row group 最多写入的行数。 |
| lance.write.max-bytes-per-file | long | 否 | 20480 |
单个 Lance 文件最多写入的字节数。 |
@@ -86,6 +87,14 @@ Lance namespace 的根目录。SeaTunnel 运行用户需要有权限在该目录
目标 Lance 表名。不设置时，如果上游存在表名，连接器会使用上游表名；该配置本身的默认值是 `test`。
+### lance.write.batch-size
+
+Writer 在内存中攒多少行后再转换成 **一个** Lance fragment 并 append。默认值是 `10000`（此前 Writer 固定攒
`1000` 行），取值必须不小于 `1`。
+
+Lance 是列存格式，`10000` 行一个 fragment 可能仍然偏小。每个 fragment
行数过少，会产生大量小列存文件，增加后续的读写放大。调大后 fragment 更少、更大，但每个 Writer flush
前在内存里攒的行也更多。资源充足时可以调大这个值，资源紧张时（例如宽表或向量列）应调小。
+
+即使未凑满该行数，checkpoint（`prepareCommit`）和 sink close 时也会 flush。
+
### lance.write.mode
控制 Lance 的写入方式。默认值是 `CREATE`。该值需要是 Lance `WriteParams.WriteMode`
支持的值：`CREATE`、`APPEND` 或 `OVERWRITE`。
@@ -136,7 +145,7 @@ Lance 使用 Apache Arrow 类型系统。sink 会根据上游 SeaTunnel 表结
:::tip
-Sink 不会按 `UPDATE` / `DELETE` 行类型执行 CDC 语义 —— 每条上游记录都会按 `lance.write.mode` 追加到
Lance 数据集中。在流式模式下，Writer 会在每个 checkpoint 把内存中的行缓冲写入 Lance。
+Sink 不会按 `UPDATE` / `DELETE` 行类型执行 CDC 语义 —— 每条上游记录都会按 `lance.write.mode` 追加到
Lance 数据集中。在流式模式下，Writer 会在达到 `lance.write.batch-size` 以及每个 checkpoint
时把内存中的行缓冲写入 Lance。
:::
@@ -185,7 +194,7 @@ sink {
### 使用 APPEND 模式并调大文件分片
-`APPEND` 模式会保留已有数据集并写入新行。把 `lance.write.max-rows-per-file` 和
`lance.write.max-bytes-per-file` 调大，可以减少追加大批量数据时产生的 Lance fragment 数量。
+`APPEND` 模式会保留已有数据集并写入新行。追加大批量数据、并且可以接受更大内存缓冲时，调大 `lance.write.batch-size`
可以减少产生的 Lance fragment 数量。
```hocon
env {
@@ -213,6 +222,7 @@ sink {
namespace_id = "root"
table = "lance_sink_table"
lance.write.mode = "APPEND"
+ lance.write.batch-size = 100000
lance.write.max-rows-per-file = 100000
lance.write.max-rows-per-group = 5000
lance.write.max-bytes-per-file = 134217728
diff --git
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkConfig.java
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkConfig.java
index a561de578d..d768f6528b 100644
---
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkConfig.java
+++
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkConfig.java
@@ -25,6 +25,8 @@ import java.util.Map;
public class LanceSinkConfig extends LanceCommonConfig {
+ private final Integer batchSize;
+
private final Integer maxRowsPerFile;
private final Integer maxRowsPerGroup;
@@ -42,6 +44,7 @@ public class LanceSinkConfig extends LanceCommonConfig {
public LanceSinkConfig(ReadonlyConfig pluginConfig) {
super(pluginConfig);
this.namespaceId =
pluginConfig.get(LanceCommonOptions.KEY_NAMESPACE_ID);
+ this.batchSize = pluginConfig.get(LanceSinkOptions.WRITE_BATCH_SIZE);
this.maxBytesPerFile =
pluginConfig.get(LanceSinkOptions.WRITE_MAX_BYTES_PER_FILE);
this.maxRowsPerGroup =
pluginConfig.get(LanceSinkOptions.WRITE_MAX_ROWS_PER_GROUP);
this.maxRowsPerFile =
pluginConfig.get(LanceSinkOptions.WRITE_MAX_ROWS_PER_FILE);
@@ -50,6 +53,10 @@ public class LanceSinkConfig extends LanceCommonConfig {
this.enableStableRowIds =
pluginConfig.get(LanceSinkOptions.WRITE_ENABLE_STABLE_ROW_IDS);
}
+ public Integer getBatchSize() {
+ return batchSize;
+ }
+
public Integer getMaxRowsPerFile() {
return maxRowsPerFile;
}
diff --git
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkOptions.java
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkOptions.java
index 7cfd1c3c20..d0d0ce5aac 100644
---
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkOptions.java
+++
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/config/LanceSinkOptions.java
@@ -25,6 +25,13 @@ import java.util.Map;
public class LanceSinkOptions extends LanceCommonOptions {
+ public static final Option<Integer> WRITE_BATCH_SIZE =
+ Options.key("lance.write.batch-size")
+ .intType()
+ .defaultValue(10000)
+ .withDescription(
+ "Rows buffered in the sink writer before flushing
one Lance fragment.");
+
public static final Option<Integer> WRITE_MAX_ROWS_PER_FILE =
Options.key("lance.write.max-rows-per-file")
.intType()
diff --git
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkFactory.java
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkFactory.java
index 699f802234..2743533b84 100644
---
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkFactory.java
+++
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkFactory.java
@@ -36,6 +36,7 @@ import org.apache.commons.collections4.CollectionUtils;
import com.google.auto.service.AutoService;
+import static
org.apache.seatunnel.api.configuration.util.Conditions.greaterThan;
import static org.apache.seatunnel.api.configuration.util.Conditions.notBlank;
@AutoService(Factory.class)
@@ -61,6 +62,9 @@ public class LanceSinkFactory implements TableSinkFactory {
LanceCommonOptions.KEY_NAMESPACE_TYPE,
notBlank(LanceCommonOptions.KEY_DATASET_PATH),
notBlank(LanceCommonOptions.KEY_NAMESPACE_TYPE))
+ .optional(
+ LanceSinkOptions.WRITE_BATCH_SIZE,
+ greaterThan(LanceSinkOptions.WRITE_BATCH_SIZE, 0))
.optional(
LanceCommonOptions.KEY_NAMESPACE_ID,
LanceSinkOptions.WRITE_MAX_ROWS_PER_FILE,
diff --git
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkWriter.java
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkWriter.java
index d4d9808dad..c5fb3aaa71 100644
---
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkWriter.java
+++
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkWriter.java
@@ -55,8 +55,6 @@ public class LanceSinkWriter
SupportMultiTableSinkWriter<Void>,
SupportSchemaEvolutionSinkWriter {
- private static final int DEFAULT_BATCH_SIZE = 1000;
-
private SeaTunnelRowType seaTunnelRowType;
private TableSchema sourceTableSchema;
private final LanceSinkConfig config;
@@ -79,8 +77,8 @@ public class LanceSinkWriter
this.sourceTableSchema = sourceTableSchema;
this.config = config;
this.catalog = catalog;
- this.batchSize = DEFAULT_BATCH_SIZE;
- this.batchBuffer = new ArrayList<>(batchSize);
+ this.batchSize = config.getBatchSize();
+ this.batchBuffer = new ArrayList<>();
this.allocator = new RootAllocator(Long.MAX_VALUE);
}
@@ -139,14 +137,13 @@ public class LanceSinkWriter
}
try {
- List<FragmentMetadata> allFragments = new ArrayList<>();
-
- for (SeaTunnelRow row : batchBuffer) {
- List<FragmentMetadata> fragmentMetadata =
- FragmentConverter.reconvert(
- row, seaTunnelRowType, schema, allocator,
config.getDatasetPath());
- allFragments.addAll(fragmentMetadata);
- }
+ List<FragmentMetadata> allFragments =
+ FragmentConverter.reconvert(
+ batchBuffer,
+ seaTunnelRowType,
+ schema,
+ allocator,
+ config.getDatasetPath());
if (!allFragments.isEmpty()) {
Transaction transaction =
diff --git
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverter.java
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverter.java
index 6525d0569a..95edd6e240 100644
---
a/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverter.java
+++
b/seatunnel-connectors-v2/connector-lance/src/main/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverter.java
@@ -37,6 +37,7 @@ import com.lancedb.lance.Fragment;
import com.lancedb.lance.FragmentMetadata;
import com.lancedb.lance.WriteParams;
+import java.util.Collections;
import java.util.List;
/** The converter for converting {@link Fragment} and {@link SeaTunnelRow} * */
@@ -45,24 +46,21 @@ public class FragmentConverter {
private FragmentConverter() {}
public static List<FragmentMetadata> reconvert(
- SeaTunnelRow seaTunnelRow,
+ List<SeaTunnelRow> rows,
SeaTunnelRowType seaTunnelRowType,
Schema schema,
BufferAllocator allocator,
String datasetPath) {
-
- List<FragmentMetadata> fragmentMetas;
+ if (rows == null || rows.isEmpty()) {
+ return Collections.emptyList();
+ }
try (VectorSchemaRoot root =
- convertToVectorSchemaRoot(seaTunnelRow, seaTunnelRowType,
schema, allocator)) {
- fragmentMetas =
- Fragment.create(
- datasetPath,
- allocator,
- root,
- new WriteParams.Builder()
- .withMaxRowsPerFile(Integer.MAX_VALUE)
- .build());
- return fragmentMetas;
+ convertToVectorSchemaRoot(rows, seaTunnelRowType, schema,
allocator)) {
+ return Fragment.create(
+ datasetPath,
+ allocator,
+ root,
+ new
WriteParams.Builder().withMaxRowsPerFile(Integer.MAX_VALUE).build());
}
}
@@ -71,24 +69,44 @@ public class FragmentConverter {
SeaTunnelRowType seaTunnelRowType,
Schema schema,
BufferAllocator allocator) {
+ return convertToVectorSchemaRoot(
+ Collections.singletonList(seaTunnelRow), seaTunnelRowType,
schema, allocator);
+ }
+
+ static VectorSchemaRoot convertToVectorSchemaRoot(
+ List<SeaTunnelRow> rows,
+ SeaTunnelRowType seaTunnelRowType,
+ Schema schema,
+ BufferAllocator allocator) {
VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
try {
+ List<Field> fields = schema.getFields();
+ int fieldCount = fields.size();
+ FieldVector[] vectors = new FieldVector[fieldCount];
+ int[] fieldIndexes = new int[fieldCount];
+ for (int i = 0; i < fieldCount; i++) {
+ vectors[i] = root.getVector(fields.get(i).getName());
+ fieldIndexes[i] =
seaTunnelRowType.indexOf(fields.get(i).getName(), false);
+ }
+
+ int rowCount = rows.size();
root.allocateNew();
- int rowIndex = 0;
- for (Field field : schema.getFields()) {
- FieldVector vector = root.getVector(field.getName());
- int fieldIndex = seaTunnelRowType.indexOf(field.getName(),
false);
- if (fieldIndex >= 0) {
- Object fieldValue = seaTunnelRow.getField(fieldIndex);
- setVectorValue(vector, field, fieldValue, rowIndex,
allocator);
- } else {
- // The persisted dataset can be ahead of the restored
checkpoint schema when a
- // failover occurs after the physical schema change but
before the checkpoint.
- // Keep the dataset layout and explicitly materialize the
missing runtime field.
- vector.setNull(rowIndex);
+ for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
+ SeaTunnelRow seaTunnelRow = rows.get(rowIndex);
+ for (int i = 0; i < fieldCount; i++) {
+ if (fieldIndexes[i] >= 0) {
+ Object fieldValue =
seaTunnelRow.getField(fieldIndexes[i]);
+ setVectorValue(vectors[i], fields.get(i), fieldValue,
rowIndex, allocator);
+ } else {
+ // The persisted dataset can be ahead of the restored
checkpoint schema when
+ // a failover occurs after the physical schema change
but before the
+ // checkpoint. Keep the dataset layout and explicitly
materialize the
+ // missing runtime field.
+ vectors[i].setNull(rowIndex);
+ }
}
}
- root.setRowCount(1);
+ root.setRowCount(rowCount);
return root;
} catch (RuntimeException | Error e) {
root.close();
diff --git
a/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/LanceFactoryTest.java
b/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/LanceFactoryTest.java
index 9d922ac97c..897b06ede8 100644
---
a/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/LanceFactoryTest.java
+++
b/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/LanceFactoryTest.java
@@ -22,6 +22,8 @@ import
org.apache.seatunnel.api.configuration.util.ConfigValidator;
import org.apache.seatunnel.api.configuration.util.OptionRule;
import org.apache.seatunnel.api.configuration.util.OptionValidationException;
import
org.apache.seatunnel.connectors.seatunnel.lance.config.LanceCommonOptions;
+import org.apache.seatunnel.connectors.seatunnel.lance.config.LanceSinkConfig;
+import org.apache.seatunnel.connectors.seatunnel.lance.config.LanceSinkOptions;
import org.apache.seatunnel.connectors.seatunnel.lance.sink.LanceSinkFactory;
import org.junit.jupiter.api.Assertions;
@@ -114,6 +116,37 @@ public class LanceFactoryTest {
Assertions.assertEquals(value,
config.get(LanceCommonOptions.KEY_NAMESPACE_ID));
}
+ @Test
+ void testBatchSizeDefaultsTo10000() {
+ LanceSinkConfig sinkConfig = new
LanceSinkConfig(ReadonlyConfig.fromMap(validConfig()));
+
+ Assertions.assertEquals(10000, sinkConfig.getBatchSize());
+ }
+
+ @Test
+ void testPositiveBatchSizeAccepted() {
+ Map<String, Object> config = validConfig();
+ config.put(LanceSinkOptions.WRITE_BATCH_SIZE.key(), 1);
+
+ ReadonlyConfig readonlyConfig = Assertions.assertDoesNotThrow(() ->
validate(config));
+
+ Assertions.assertEquals(1, new
LanceSinkConfig(readonlyConfig).getBatchSize().intValue());
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {0, -1})
+ void testNonPositiveBatchSizeRejected(int value) {
+ Map<String, Object> config = validConfig();
+ config.put(LanceSinkOptions.WRITE_BATCH_SIZE.key(), value);
+
+ OptionValidationException error =
+ Assertions.assertThrows(OptionValidationException.class, () ->
validate(config));
+
+ Assertions.assertTrue(
+
error.getMessage().contains(LanceSinkOptions.WRITE_BATCH_SIZE.key()),
+ error.getMessage());
+ }
+
private void assertInvalidOption(String key, String value) {
Map<String, Object> config = validConfig();
config.put(key, value);
diff --git
a/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkWriterBatchTest.java
b/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkWriterBatchTest.java
new file mode 100644
index 0000000000..fc94696c7f
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/sink/LanceSinkWriterBatchTest.java
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.seatunnel.connectors.seatunnel.lance.sink;
+
+import org.apache.seatunnel.api.configuration.ReadonlyConfig;
+import org.apache.seatunnel.api.table.catalog.PhysicalColumn;
+import org.apache.seatunnel.api.table.catalog.TableSchema;
+import org.apache.seatunnel.api.table.type.BasicType;
+import org.apache.seatunnel.api.table.type.SeaTunnelRow;
+import org.apache.seatunnel.connectors.seatunnel.lance.catalog.LanceCatalog;
+import
org.apache.seatunnel.connectors.seatunnel.lance.config.LanceCommonOptions;
+import org.apache.seatunnel.connectors.seatunnel.lance.config.LanceSinkConfig;
+import org.apache.seatunnel.connectors.seatunnel.lance.config.LanceSinkOptions;
+
+import org.apache.arrow.memory.RootAllocator;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledOnOs;
+import org.junit.jupiter.api.condition.OS;
+import org.junit.jupiter.api.io.TempDir;
+
+import com.lancedb.lance.Dataset;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+
+@DisabledOnOs(OS.WINDOWS)
+class LanceSinkWriterBatchTest {
+
+ @TempDir Path tempDir;
+
+ @Test
+ void flushWritesOneFragmentPerBatch() throws IOException {
+ String datasetPath =
tempDir.resolve("batch_fragment.lance").toString();
+ Map<String, Object> configs = new HashMap<>();
+ configs.put(LanceCommonOptions.KEY_DATASET_PATH.key(), datasetPath);
+ configs.put(LanceCommonOptions.KEY_NAMESPACE_TYPE.key(), "dir");
+ configs.put(LanceCommonOptions.KEY_ROOT_NAMESPACE_PATH.key(),
tempDir.toString());
+ configs.put(LanceSinkOptions.WRITE_BATCH_SIZE.key(), 2);
+ ReadonlyConfig config = ReadonlyConfig.fromMap(configs);
+ TableSchema tableSchema =
+ TableSchema.builder()
+ .column(
+ PhysicalColumn.of(
+ "c_string",
+ BasicType.STRING_TYPE,
+ (Long) null,
+ true,
+ null,
+ "c_string"))
+ .column(
+ PhysicalColumn.of(
+ "c_int",
+ BasicType.INT_TYPE,
+ (Long) null,
+ true,
+ null,
+ "c_int"))
+ .build();
+
+ LanceSinkWriter writer =
+ new LanceSinkWriter(
+ tableSchema.toPhysicalRowDataType(),
+ tableSchema,
+ new LanceSinkConfig(config),
+ new LanceCatalog("lance_namespace", config));
+ for (int i = 0; i < 5; i++) {
+ writer.write(new SeaTunnelRow(new Object[] {"row-" + i, i}));
+ }
+ writer.prepareCommit();
+ writer.close();
+
+ try (RootAllocator allocator = new RootAllocator();
+ Dataset dataset = Dataset.open(datasetPath, allocator)) {
+ Assertions.assertEquals(5, dataset.countRows());
+ // Two full batches of 2 rows, then the checkpoint flushes the
remaining row.
+ Assertions.assertEquals(3, dataset.getFragments().size());
+ }
+ }
+}
diff --git
a/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverterTest.java
b/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverterTest.java
index 55745907c7..3575f094de 100644
---
a/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverterTest.java
+++
b/seatunnel-connectors-v2/connector-lance/src/test/java/org/apache/seatunnel/connectors/seatunnel/lance/utils/FragmentConverterTest.java
@@ -17,6 +17,7 @@
package org.apache.seatunnel.connectors.seatunnel.lance.utils;
+import org.apache.seatunnel.api.table.type.ArrayType;
import org.apache.seatunnel.api.table.type.BasicType;
import org.apache.seatunnel.api.table.type.SeaTunnelDataType;
import org.apache.seatunnel.api.table.type.SeaTunnelRow;
@@ -25,15 +26,32 @@ import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.pojo.ArrowType;
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;
+import org.junit.jupiter.api.condition.DisabledOnOs;
+import org.junit.jupiter.api.condition.OS;
+import com.lancedb.lance.Dataset;
+import com.lancedb.lance.FragmentMetadata;
+import com.lancedb.lance.WriteParams;
+
+import java.io.IOException;
import java.nio.charset.StandardCharsets;
+import java.nio.file.FileVisitResult;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
class FragmentConverterTest {
@@ -66,4 +84,136 @@ class FragmentConverterTest {
Assertions.assertEquals(1,
root.getVector("c_bytes").getValueCount());
}
}
+
+ @Test
+ void shouldPackMultipleRowsIntoOneVectorSchemaRoot() {
+ SeaTunnelRowType rowType =
+ new SeaTunnelRowType(
+ new String[] {"c_string"},
+ new SeaTunnelDataType<?>[] {BasicType.STRING_TYPE});
+ Schema schema = new Schema(Arrays.asList(Field.nullable("c_string",
new ArrowType.Utf8())));
+
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
+ FragmentConverter.convertToVectorSchemaRoot(
+ Arrays.asList(
+ new SeaTunnelRow(new Object[]
{"row-0"}),
+ new SeaTunnelRow(new Object[]
{"row-1"}),
+ new SeaTunnelRow(new Object[]
{"row-2"})),
+ rowType,
+ schema,
+ allocator)) {
+ Assertions.assertEquals(3, root.getRowCount());
+ VarCharVector vector = (VarCharVector) root.getVector("c_string");
+ Assertions.assertEquals("row-0", new String(vector.get(0),
StandardCharsets.UTF_8));
+ Assertions.assertEquals("row-1", new String(vector.get(1),
StandardCharsets.UTF_8));
+ Assertions.assertEquals("row-2", new String(vector.get(2),
StandardCharsets.UTF_8));
+ }
+ }
+
+ @Test
+ void shouldPackMultipleArrayRowsIntoOneVectorSchemaRoot() {
+ SeaTunnelRowType rowType =
+ new SeaTunnelRowType(
+ new String[] {"c_floats"},
+ new SeaTunnelDataType<?>[]
{ArrayType.FLOAT_ARRAY_TYPE});
+ Field arrayField =
+ new Field(
+ "c_floats",
+ FieldType.nullable(new ArrowType.List()),
+ Collections.singletonList(
+ Field.nullable(
+ "item",
+ new ArrowType.FloatingPoint(
+
FloatingPointPrecision.SINGLE))));
+ Schema schema = new Schema(Collections.singletonList(arrayField));
+ List<Float> fiveValues = Arrays.asList(1.0f, 2.0f, 3.0f, 4.0f, 5.0f);
+ List<Float> oneValue = Collections.singletonList(6.0f);
+ List<Float> empty = Collections.emptyList();
+
+ try (RootAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
+ FragmentConverter.convertToVectorSchemaRoot(
+ Arrays.asList(
+ new SeaTunnelRow(new Object[]
{fiveValues}),
+ new SeaTunnelRow(new Object[]
{oneValue}),
+ new SeaTunnelRow(new Object[] {empty}),
+ new SeaTunnelRow(new Object[] {null})),
+ rowType,
+ schema,
+ allocator)) {
+ Assertions.assertEquals(4, root.getRowCount());
+ ListVector vector = (ListVector) root.getVector("c_floats");
+ Assertions.assertEquals(fiveValues, vector.getObject(0));
+ Assertions.assertEquals(oneValue, vector.getObject(1));
+ Assertions.assertFalse(vector.isNull(2));
+ Assertions.assertEquals(empty, vector.getObject(2));
+ Assertions.assertTrue(vector.isNull(3));
+ }
+ }
+
+ @Test
+ @DisabledOnOs(OS.WINDOWS)
+ void shouldWriteOneFragmentPerBatch() throws IOException {
+ SeaTunnelRowType rowType =
+ new SeaTunnelRowType(
+ new String[] {"c_string"},
+ new SeaTunnelDataType<?>[] {BasicType.STRING_TYPE});
+ Schema schema = new Schema(Arrays.asList(Field.nullable("c_string",
new ArrowType.Utf8())));
+ List<SeaTunnelRow> rows =
+ Arrays.asList(
+ new SeaTunnelRow(new Object[] {"row-0"}),
+ new SeaTunnelRow(new Object[] {"row-1"}),
+ new SeaTunnelRow(new Object[] {"row-2"}));
+
+ Path tempDir = Files.createTempDirectory("lance-batch-fragment");
+ try (RootAllocator allocator = new RootAllocator()) {
+ String datasetPath = tempDir.resolve("batch.lance").toString();
+ createEmptyDataset(allocator, datasetPath, schema);
+ List<FragmentMetadata> fragments =
+ FragmentConverter.reconvert(rows, rowType, schema,
allocator, datasetPath);
+ Assertions.assertEquals(1, fragments.size());
+ Assertions.assertEquals(3, fragments.get(0).getNumRows());
+ } finally {
+ deleteRecursively(tempDir);
+ }
+ }
+
+ private static void createEmptyDataset(
+ RootAllocator allocator, String datasetPath, Schema schema) {
+ try (Dataset ignored =
+ Dataset.create(
+ allocator,
+ datasetPath,
+ schema,
+ new
WriteParams.Builder().withMode(WriteParams.WriteMode.CREATE).build())) {
+ // An empty dataset gives Fragment.create a place to write data
files.
+ }
+ }
+
+ private static void deleteRecursively(Path root) throws IOException {
+ if (root == null || !Files.exists(root)) {
+ return;
+ }
+ Files.walkFileTree(
+ root,
+ new SimpleFileVisitor<Path>() {
+ @Override
+ public FileVisitResult visitFile(Path file,
BasicFileAttributes attrs)
+ throws IOException {
+ Files.deleteIfExists(file);
+ return FileVisitResult.CONTINUE;
+ }
+
+ @Override
+ public FileVisitResult postVisitDirectory(Path dir,
IOException exc)
+ throws IOException {
+ if (exc != null) {
+ throw exc;
+ }
+ Files.deleteIfExists(dir);
+ return FileVisitResult.CONTINUE;
+ }
+ });
+ }
}
