Skip to content

Commit

Permalink
[Flink] add arrow datastream source (#496)
Browse files Browse the repository at this point in the history 
* add lakesoul arrow source

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* fix flink writer path

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* support inferring schema

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* tmp commit

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* fix regression test

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* fix reading tail recordbatch

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* fix stup dao sql

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* fix dynamic schema

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* optimize imports

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* split DataFileInfo with file_cols

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* optimize import

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* rebase on  main

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* fix regresion test

Signed-off-by: zenghua <huazeng@dmetasoul.com>

* fix regresion test

Signed-off-by: zenghua <huazeng@dmetasoul.com>

---------

Signed-off-by: zenghua <huazeng@dmetasoul.com>
Co-authored-by: zenghua <huazeng@dmetasoul.com>
  • Loading branch information
@Ceng23333
Ceng23333 and zenghua committed Jun 19, 2024
1 parent 2d6451b commit e15c8e7
Show file tree
Hide file tree
Showing 29 changed files with 893 additions and 238 deletions.
24 changes: 6 additions & 18 deletions lakesoul-common/src/main/java/com/dmetasoul/lakesoul/meta/dao/PartitionInfoDao.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ public List<PartitionInfo> findByTableIdAndParList(String tableId, List<String>
descPlaceholders = String.join(",", Collections.nCopies(partitionDescList.size(), "?"));
}
String sql = String.format(
"select m.table_id, t.partition_desc, m.version, m.commit_op, m.snapshot, m.expression, m.domain from (" +
"select m.table_id, t.partition_desc, m.version, m.commit_op, m.snapshot, m.timestamp m.expression, m.domain from (" +
"select table_id,partition_desc,max(version) from partition_info " +
"where table_id = ? and partition_desc in (%s) " +
"group by table_id,partition_desc) t " +
Expand All @@ -216,7 +216,7 @@ public List<PartitionInfo> findByTableIdAndParList(String tableId, List<String>
}
rs = pstmt.executeQuery();
while (rs.next()) {
rsList.add(partitionInfoFromResultSetWithoutTimestamp(rs));
rsList.add(partitionInfoFromResultSet(rs));
}
} catch (SQLException e) {
throw new RuntimeException(e);
Expand All @@ -236,7 +236,7 @@ public PartitionInfo selectLatestPartitionInfo(String tableId, String partitionD
return partitionInfoList.isEmpty() ? null : partitionInfoList.get(0);
}
String sql = String.format(
"select m.table_id, t.partition_desc, m.version, m.commit_op, m.snapshot, m.expression, m.domain from (" +
"select m.table_id, t.partition_desc, m.version, m.commit_op, m.snapshot, m.expression, m.timestamp, m.domain from (" +
"select table_id,partition_desc,max(version) from partition_info " +
"where table_id = '%s' and partition_desc = '%s' " + "group by table_id,partition_desc) t " +
"left join partition_info m on t.table_id = m.table_id and t.partition_desc = m" +
Expand Down Expand Up @@ -407,7 +407,7 @@ public List<PartitionInfo> getPartitionDescByTableId(String tableId) {
ResultSet rs = null;
List<PartitionInfo> rsList = new ArrayList<>();
String sql = String.format(
"select m.table_id, t.partition_desc, m.version, m.commit_op, m.snapshot, m.expression, m.domain from (" +
"select m.table_id, t.partition_desc, m.version, m.commit_op, m.snapshot, m.expression, m.timestamp, m.domain from (" +
"select table_id,partition_desc,max(version) from partition_info " + "where table_id = '%s' " +
"group by table_id,partition_desc) t " +
"left join partition_info m on t.table_id = m.table_id and t.partition_desc = m" +
Expand All @@ -418,7 +418,7 @@ public List<PartitionInfo> getPartitionDescByTableId(String tableId) {
pstmt = conn.prepareStatement(sql);
rs = pstmt.executeQuery();
while (rs.next()) {
rsList.add(partitionInfoFromResultSetWithoutTimestamp(rs));
rsList.add(partitionInfoFromResultSet(rs));
}
} catch (SQLException e) {
throw new RuntimeException(e);
Expand Down Expand Up @@ -583,7 +583,7 @@ private PartitionInfo getPartitionInfo(String sql) {
rs = pstmt.executeQuery();

while (rs.next()) {
partitionInfo = partitionInfoFromResultSetWithoutTimestamp(rs);
partitionInfo = partitionInfoFromResultSet(rs);
}
} catch (SQLException e) {
throw new RuntimeException(e);
Expand Down Expand Up @@ -623,16 +623,4 @@ public static PartitionInfo partitionInfoFromResultSet(ResultSet rs) throws SQLE
return partitionInfo.build();
}

public static PartitionInfo partitionInfoFromResultSetWithoutTimestamp(ResultSet rs) throws SQLException {
PartitionInfo.Builder partitionInfo = PartitionInfo.newBuilder()
.setTableId(rs.getString("table_id"))
.setPartitionDesc(rs.getString("partition_desc"))
.setVersion(rs.getInt("version"))
.setCommitOp(CommitOp.valueOf(rs.getString("commit_op")));
Array snapshotArray = rs.getArray("snapshot");
partitionInfo.addAllSnapshot(Arrays.stream((UUID[]) snapshotArray.getArray()).map(DBUtil::toProtoUuid).collect(Collectors.toList()));
String expr = rs.getString("expression");
partitionInfo.setExpression(expr == null ? "" : expr);
return partitionInfo.build();
}
}
8 changes: 4 additions & 4 deletions lakesoul-flink/src/main/java/org/apache/flink/lakesoul/sink/writer/NativeParquetWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ public class NativeParquetWriter implements InProgressFileWriter<RowData, String

long lastUpdateTime;

String path;
Path path;

private long totalRows = 0;

Expand All @@ -72,8 +72,8 @@ public NativeParquetWriter(RowType rowType,
nativeWriter.setRowGroupRowNumber(this.batchSize);
batch = VectorSchemaRoot.create(arrowSchema, nativeWriter.getAllocator());
arrowWriter = ArrowUtils.createRowDataArrowWriter(batch, rowType);
this.path = path.makeQualified(path.getFileSystem()).toString();
nativeWriter.addFile(this.path);
this.path = path.makeQualified(path.getFileSystem());
nativeWriter.addFile(this.path.toUri().toString());

FlinkUtil.setFSConfigs(conf, this.nativeWriter);
this.nativeWriter.initializeWriter();
Expand Down Expand Up @@ -184,7 +184,7 @@ public PendingFileRecoverable closeForCommit() throws IOException {
} catch (Exception e) {
throw new RuntimeException(e);
}
return new NativeWriterPendingFileRecoverable(this.path, this.creationTime);
return new NativeWriterPendingFileRecoverable(this.path.toString(), this.creationTime);
}

@Override
Expand Down
14 changes: 10 additions & 4 deletions ...ain/java/org/apache/flink/lakesoul/source/LakeSoulAllPartitionDynamicSplitEnumerator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import com.dmetasoul.lakesoul.lakesoul.io.substrait.SubstraitUtil;
import com.dmetasoul.lakesoul.meta.DataFileInfo;
import com.dmetasoul.lakesoul.meta.DataOperation;
import com.dmetasoul.lakesoul.meta.LakeSoulOptions;
import com.dmetasoul.lakesoul.meta.MetaVersion;
import com.dmetasoul.lakesoul.meta.entity.PartitionInfo;
import com.dmetasoul.lakesoul.meta.entity.TableInfo;
Expand All @@ -15,6 +16,7 @@
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.fs.Path;
import org.apache.flink.lakesoul.tool.FlinkUtil;
import org.apache.flink.shaded.guava30.com.google.common.collect.Maps;
Expand Down Expand Up @@ -141,35 +143,39 @@ private void processDiscoveredSplits(Collection<LakeSoulPartitionSplit> splits,

public Collection<LakeSoulPartitionSplit> enumerateSplits() {
List<PartitionInfo> allPartitionInfo = MetaVersion.getAllPartitionInfo(tableId);
LOG.info("allPartitionInfo={}", allPartitionInfo);
List<PartitionInfo> filteredPartition = SubstraitUtil.applyPartitionFilters(allPartitionInfo, partitionArrowSchema, partitionFilters);
LOG.info("filteredPartition={}, filter={}", filteredPartition, partitionFilters);


ArrayList<LakeSoulPartitionSplit> splits = new ArrayList<>(16);
for (PartitionInfo partitionInfo : filteredPartition) {
String partitionDesc = partitionInfo.getPartitionDesc();
long latestTimestamp = partitionInfo.getTimestamp();
long latestTimestamp = partitionInfo.getTimestamp() + 1;
this.nextStartTime = Math.max(latestTimestamp, this.nextStartTime);

DataFileInfo[] dataFileInfos;
if (partitionLatestTimestamp.containsKey(partitionDesc)) {
Long lastTimestamp = partitionLatestTimestamp.get(partitionDesc);
LOG.info("getIncrementalPartitionDataInfo, startTime={}, endTime={}", lastTimestamp, latestTimestamp);
dataFileInfos =
DataOperation.getIncrementalPartitionDataInfo(tableId, partitionDesc, lastTimestamp + 1, latestTimestamp, "incremental");
DataOperation.getIncrementalPartitionDataInfo(tableId, partitionDesc, lastTimestamp, latestTimestamp, "incremental");
} else {
dataFileInfos =
DataOperation.getIncrementalPartitionDataInfo(tableId, partitionDesc, startTime, latestTimestamp, "incremental");
}
if (dataFileInfos.length > 0) {
Map<String, Map<Integer, List<Path>>> splitByRangeAndHashPartition =
Map<Tuple2<String, String>, Map<Integer, List<Path>>> splitByRangeAndHashPartition =
FlinkUtil.splitDataInfosToRangeAndHashPartition(tableInfo, dataFileInfos);
for (Map.Entry<String, Map<Integer, List<Path>>> entry : splitByRangeAndHashPartition.entrySet()) {
for (Map.Entry<Tuple2<String, String>, Map<Integer, List<Path>>> entry : splitByRangeAndHashPartition.entrySet()) {
for (Map.Entry<Integer, List<Path>> split : entry.getValue().entrySet()) {
splits.add(new LakeSoulPartitionSplit(String.valueOf(split.hashCode()), split.getValue(), 0, split.getKey(), partitionDesc));
}
}
}
partitionLatestTimestamp.put(partitionDesc, latestTimestamp);
}
LOG.info("partitionLatestTimestamp={}", partitionLatestTimestamp);

return splits;
}
Expand Down
4 changes: 2 additions & 2 deletions lakesoul-flink/src/main/java/org/apache/flink/lakesoul/source/LakeSoulRecordEmitter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@
import org.apache.flink.connector.base.source.reader.RecordEmitter;
import org.apache.flink.table.data.RowData;

public class LakeSoulRecordEmitter implements RecordEmitter<RowData, RowData, LakeSoulPartitionSplit> {
public class LakeSoulRecordEmitter<OUT> implements RecordEmitter<OUT, OUT, LakeSoulPartitionSplit> {
@Override
public void emitRecord(RowData element, SourceOutput<RowData> output, LakeSoulPartitionSplit splitState) throws Exception {
public void emitRecord(OUT element, SourceOutput<OUT> output, LakeSoulPartitionSplit splitState) throws Exception {
output.collect(element);
splitState.incrementRecord();
}
Expand Down
69 changes: 69 additions & 0 deletions lakesoul-flink/src/main/java/org/apache/flink/lakesoul/source/LakeSoulRowDataSource.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
// SPDX-FileCopyrightText: 2023 LakeSoul Contributors
//
// SPDX-License-Identifier: Apache-2.0

package org.apache.flink.lakesoul.source;

import io.substrait.proto.Plan;
import org.apache.flink.api.connector.source.SourceReader;
import org.apache.flink.api.connector.source.SourceReaderContext;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.lakesoul.tool.LakeSoulSinkOptions;
import org.apache.flink.lakesoul.types.TableId;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.logical.RowType;

import javax.annotation.Nullable;
import java.util.List;
import java.util.Map;

public class LakeSoulRowDataSource extends LakeSoulSource<RowData> {

public LakeSoulRowDataSource(TableId tableId,
RowType tableRowType,
RowType projectedRowType,
RowType rowTypeWithPk,
boolean isBounded,
List<String> pkColumns,
List<String> partitionColumns,
Map<String, String> optionParams,
@Nullable List<Map<String, String>> remainingPartitions,
@Nullable Plan pushedFilter,
@Nullable Plan partitionFilters
) {
super(tableId,
tableRowType,
projectedRowType,
rowTypeWithPk,
isBounded,
pkColumns,
partitionColumns,
optionParams,
remainingPartitions,
pushedFilter,
partitionFilters
);
}


@Override
public SourceReader<RowData, LakeSoulPartitionSplit> createReader(SourceReaderContext readerContext) throws Exception {
Configuration conf = Configuration.fromMap(optionParams);
conf.addAll(readerContext.getConfiguration());
return new LakeSoulSourceReader(
() -> new LakeSoulSplitReader(
conf,
this.tableRowType,
this.projectedRowType,
this.rowTypeWithPk,
this.pkColumns,
this.isBounded,
this.optionParams.getOrDefault(LakeSoulSinkOptions.CDC_CHANGE_COLUMN, ""),
this.partitionColumns,
this.pushedFilter),
new LakeSoulRecordEmitter(),
readerContext.getConfiguration(),
readerContext);
}

}
Loading

0 comments on commit e15c8e7

Please sign in to comment.