[lake/lance] Lance lake writer and committer implementation (apache#1441)

Author: xx789633

fluss-lake/fluss-lake-lance/pom.xml

@@ -33,7 +33,7 @@
     <packaging>jar</packaging>
 
     <properties>
-        <lance.version>0.26.1</lance.version>
+        <lance.version>0.33.0</lance.version>
     </properties>
 
     <dependencies>

fluss-lake/fluss-lake-lance/src/main/java/com/alibaba/fluss/lake/lance/LanceConfig.java

@@ -32,17 +32,9 @@ public class LanceConfig implements Serializable {
 
     private static final String block_size = "block_size";
     private static final String version = "version";
-    private static final String index_cache_size = "index_cache_size";
-    private static final String metadata_cache_size = "metadata_cache_size";
     private static final String max_row_per_file = "max_row_per_file";
     private static final String max_rows_per_group = "max_rows_per_group";
     private static final String max_bytes_per_file = "max_bytes_per_file";
-    private static final String ak = "access_key_id";
-    private static final String sk = "secret_access_key";
-    private static final String endpoint = "aws_endpoint";
-    private static final String region = "aws_region";
-    private static final String virtual_hosted_style = "virtual_hosted_style_request";
-    private static final String allow_http = "allow_http";
     private static final String batch_size = "batch_size";
     private static final String warehouse = "warehouse";
 
@@ -64,11 +56,17 @@ public LanceConfig(
     }
 
     public static LanceConfig from(
-            Map<String, String> properties, String databaseName, String tableName) {
-        if (!properties.containsKey(warehouse)) {
+            Map<String, String> clusterConf,
+            Map<String, String> tableCustomProperties,
+            String databaseName,
+            String tableName) {
+        if (!clusterConf.containsKey(warehouse)) {
             throw new IllegalArgumentException("Missing required option " + warehouse);
         }
-        return new LanceConfig(databaseName, tableName, properties.get(warehouse), properties);
+        Map<String, String> options = new HashMap<>();
+        options.putAll(clusterConf);
+        options.putAll(tableCustomProperties);
+        return new LanceConfig(databaseName, tableName, clusterConf.get(warehouse), options);
     }
 
     public static int getBatchSize(LanceConfig config) {
@@ -83,14 +81,6 @@ public Map<String, String> getOptions() {
         return options;
     }
 
-    public String getDatabaseName() {
-        return databaseName;
-    }
-
-    public String getTableName() {
-        return tableName;
-    }
-
     public String getDatasetUri() {
         return datasetUri;
     }
@@ -104,12 +94,6 @@ public static ReadOptions genReadOptionFromConfig(LanceConfig config) {
         if (maps.containsKey(version)) {
             builder.setVersion(Integer.parseInt(maps.get(version)));
         }
-        if (maps.containsKey(index_cache_size)) {
-            builder.setIndexCacheSize(Integer.parseInt(maps.get(index_cache_size)));
-        }
-        if (maps.containsKey(metadata_cache_size)) {
-            builder.setMetadataCacheSize(Integer.parseInt(maps.get(metadata_cache_size)));
-        }
         builder.setStorageOptions(genStorageOptions(config));
         return builder.build();
     }
@@ -130,24 +114,8 @@ public static WriteParams genWriteParamsFromConfig(LanceConfig config) {
         return builder.build();
     }
 
-    private static Map<String, String> genStorageOptions(LanceConfig config) {
-        Map<String, String> storageOptions = new HashMap<>();
-        Map<String, String> maps = config.getOptions();
-        if (maps.containsKey(ak) && maps.containsKey(sk) && maps.containsKey(endpoint)) {
-            storageOptions.put(ak, maps.get(ak));
-            storageOptions.put(sk, maps.get(sk));
-            storageOptions.put(endpoint, maps.get(endpoint));
-        }
-        if (maps.containsKey(region)) {
-            storageOptions.put(region, maps.get(region));
-        }
-        if (maps.containsKey(virtual_hosted_style)) {
-            storageOptions.put(virtual_hosted_style, maps.get(virtual_hosted_style));
-        }
-        if (maps.containsKey(allow_http)) {
-            storageOptions.put(allow_http, maps.get(allow_http));
-        }
-        return storageOptions;
+    public static Map<String, String> genStorageOptions(LanceConfig config) {
+        return config.getOptions();
     }
 
     @Override

fluss-lake/fluss-lake-lance/src/main/java/com/alibaba/fluss/lake/lance/LanceLakeCatalog.java

@@ -26,31 +26,14 @@
 import com.alibaba.fluss.metadata.TablePath;
 
 import com.lancedb.lance.WriteParams;
-import org.apache.arrow.vector.types.TimeUnit;
-import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
 
 import java.util.ArrayList;
 import java.util.List;
 
-import static com.alibaba.fluss.metadata.TableDescriptor.BUCKET_COLUMN_NAME;
-import static com.alibaba.fluss.metadata.TableDescriptor.OFFSET_COLUMN_NAME;
-import static com.alibaba.fluss.metadata.TableDescriptor.TIMESTAMP_COLUMN_NAME;
-
 /** A Lance implementation of {@link LakeCatalog}. */
 public class LanceLakeCatalog implements LakeCatalog {
-    private static final List<Field> SYSTEM_COLUMNS = new ArrayList<>();
-
-    static {
-        SYSTEM_COLUMNS.add(Field.nullable(BUCKET_COLUMN_NAME, new ArrowType.Int(32, true)));
-        SYSTEM_COLUMNS.add(Field.nullable(OFFSET_COLUMN_NAME, new ArrowType.Int(64, true)));
-        SYSTEM_COLUMNS.add(
-                Field.nullable(
-                        TIMESTAMP_COLUMN_NAME,
-                        new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)));
-    }
-
     private final Configuration options;
 
     public LanceLakeCatalog(Configuration config) {
@@ -67,16 +50,17 @@ public void createTable(TablePath tablePath, TableDescriptor tableDescriptor) {
 
         LanceConfig config =
                 LanceConfig.from(
-                        options.toMap(), tablePath.getDatabaseName(), tablePath.getTableName());
+                        options.toMap(),
+                        tableDescriptor.getCustomProperties(),
+                        tablePath.getDatabaseName(),
+                        tablePath.getTableName());
         WriteParams params = LanceConfig.genWriteParamsFromConfig(config);
 
         List<Field> fields = new ArrayList<>();
         // set schema
         fields.addAll(
                 LanceArrowUtils.toArrowSchema(tableDescriptor.getSchema().getRowType())
                         .getFields());
-        // add system metadata columns to schema
-        fields.addAll(SYSTEM_COLUMNS);
         try {
             LanceDatasetAdapter.createDataset(config.getDatasetUri(), new Schema(fields), params);
         } catch (RuntimeException e) {

fluss-lake/fluss-lake-lance/src/main/java/com/alibaba/fluss/lake/lance/LanceLakeStorage.java

@@ -20,6 +20,7 @@
 import com.alibaba.fluss.config.Configuration;
 import com.alibaba.fluss.lake.lakestorage.LakeStorage;
 import com.alibaba.fluss.lake.lance.tiering.LanceCommittable;
+import com.alibaba.fluss.lake.lance.tiering.LanceLakeTieringFactory;
 import com.alibaba.fluss.lake.lance.tiering.LanceWriteResult;
 import com.alibaba.fluss.lake.source.LakeSource;
 import com.alibaba.fluss.lake.writer.LakeTieringFactory;
@@ -35,7 +36,7 @@ public LanceLakeStorage(Configuration configuration) {
 
     @Override
     public LakeTieringFactory<LanceWriteResult, LanceCommittable> createLakeTieringFactory() {
-        throw new UnsupportedOperationException("Not implemented");
+        return new LanceLakeTieringFactory(config);
     }
 
     @Override

fluss-lake/fluss-lake-lance/src/main/java/com/alibaba/fluss/lake/lance/tiering/ArrowWriter.java

@@ -0,0 +1,74 @@
+/*
+ * 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 com.alibaba.fluss.lake.lance.tiering;
+
+import com.alibaba.fluss.lake.lance.utils.LanceArrowUtils;
+import com.alibaba.fluss.lake.lance.writers.ArrowFieldWriter;
+import com.alibaba.fluss.row.InternalRow;
+import com.alibaba.fluss.types.RowType;
+
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+
+/** An Arrow writer for {@link InternalRow}. */
+public class ArrowWriter {
+    private final VectorSchemaRoot root;
+
+    private final ArrowFieldWriter<InternalRow>[] fieldWriters;
+
+    private int recordsCount;
+
+    /**
+     * Writer which serializes the Fluss rows to Arrow record batches.
+     *
+     * @param fieldWriters An array of writers which are responsible for the serialization of each
+     *     column of the rows
+     * @param root Container that holds a set of vectors for the rows
+     */
+    public ArrowWriter(ArrowFieldWriter<InternalRow>[] fieldWriters, VectorSchemaRoot root) {
+        this.fieldWriters = fieldWriters;
+        this.root = root;
+    }
+
+    public static ArrowWriter create(VectorSchemaRoot root, RowType rowType) {
+        ArrowFieldWriter<InternalRow>[] fieldWriters =
+                new ArrowFieldWriter[root.getFieldVectors().size()];
+        for (int i = 0; i < fieldWriters.length; i++) {
+            FieldVector fieldVector = root.getVector(i);
+
+            fieldWriters[i] =
+                    LanceArrowUtils.createArrowFieldWriter(fieldVector, rowType.getTypeAt(i));
+        }
+        return new ArrowWriter(fieldWriters, root);
+    }
+
+    /** Writes the specified row which is serialized into Arrow format. */
+    public void writeRow(InternalRow row) {
+        for (int i = 0; i < fieldWriters.length; i++) {
+            fieldWriters[i].write(row, i, true);
+        }
+        recordsCount++;
+    }
+
+    public void finish() {
+        root.setRowCount(recordsCount);
+        for (ArrowFieldWriter<InternalRow> fieldWriter : fieldWriters) {
+            fieldWriter.finish();
+        }
+    }
+}

fluss-lake/fluss-lake-lance/src/main/java/com/alibaba/fluss/lake/lance/tiering/LanceArrowWriter.java

@@ -0,0 +1,130 @@
+/*
+ * 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 com.alibaba.fluss.lake.lance.tiering;
+
+import com.alibaba.fluss.record.LogRecord;
+import com.alibaba.fluss.types.RowType;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.ipc.ArrowReader;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+import java.io.IOException;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static com.alibaba.fluss.utils.Preconditions.checkArgument;
+import static com.alibaba.fluss.utils.Preconditions.checkNotNull;
+
+/** A custom arrow reader that supports writes Fluss internal rows while reading data in batches. */
+public class LanceArrowWriter extends ArrowReader {
+    private final Schema schema;
+    private final RowType rowType;
+    private final int batchSize;
+
+    private volatile boolean finished;
+
+    private final AtomicLong totalBytesRead = new AtomicLong();
+    private ArrowWriter arrowWriter = null;
+    private final AtomicInteger count = new AtomicInteger(0);
+    private final Semaphore writeToken;
+    private final Semaphore loadToken;
+
+    public LanceArrowWriter(
+            BufferAllocator allocator, Schema schema, int batchSize, RowType rowType) {
+        super(allocator);
+        checkNotNull(schema);
+        checkArgument(batchSize > 0);
+        this.schema = schema;
+        this.rowType = rowType;
+        this.batchSize = batchSize;
+        this.writeToken = new Semaphore(0);
+        this.loadToken = new Semaphore(0);
+    }
+
+    void write(LogRecord row) {
+        checkNotNull(row);
+        try {
+            // wait util prepareLoadNextBatch to release write token,
+            writeToken.acquire();
+            arrowWriter.writeRow(row.getRow());
+            if (count.incrementAndGet() == batchSize) {
+                // notify loadNextBatch to take the batch
+                loadToken.release();
+            }
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    void setFinished() {
+        loadToken.release();
+        finished = true;
+    }
+
+    @Override
+    public void prepareLoadNextBatch() throws IOException {
+        super.prepareLoadNextBatch();
+        arrowWriter = ArrowWriter.create(this.getVectorSchemaRoot(), rowType);
+        // release batch size token for write
+        writeToken.release(batchSize);
+    }
+
+    @Override
+    public boolean loadNextBatch() throws IOException {
+        prepareLoadNextBatch();
+        try {
+            if (finished && count.get() == 0) {
+                return false;
+            }
+            // wait util batch if full or finished
+            loadToken.acquire();
+            arrowWriter.finish();
+            if (!finished) {
+                count.set(0);
+                return true;
+            } else {
+                // true if it has some rows and return false if there is no record
+                if (count.get() > 0) {
+                    count.set(0);
+                    return true;
+                } else {
+                    return false;
+                }
+            }
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public long bytesRead() {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    protected synchronized void closeReadSource() throws IOException {
+        // Implement if needed
+    }
+
+    @Override
+    protected Schema readSchema() {
+        return this.schema;
+    }
+}