This is an automated email from the ASF dual-hosted git repository. shuwenwei pushed a commit to branch flink-iotdb-table-connector in repository https://gitbox.apache.org/repos/asf/iotdb-extras.git
commit 8db6280388bf858860a2892a521aa662a5d32969 Author: shuwenwei <[email protected]> AuthorDate: Wed Sep 16 17:47:39 2026 +0800 predicate pushdown --- .../flink/catalog/IoTDBCatalogClient.java | 16 +- .../flink/cfg/IoTDBRelationalOptions.java | 12 +- .../sink/serializer/IoTDBTabletSerializer.java | 20 +- .../serializer/RowDataIoTDBTabletSerializer.java | 7 +- .../deserializer/IoTDBDeserializationSchema.java} | 29 +- .../deserializer/RowDataDeserializationSchema.java | 114 +++++ .../enumerator/IoTDBSourceEnumeratorState.java | 42 ++ .../IoTDBSourceEnumeratorStateSerializer.java | 81 +++ .../source/pushdown/IoTDBExpressionVisitor.java | 568 +++++++++++++++++++++ .../flink/source/pushdown/IoTDBLiteralUtils.java | 115 +++++ .../flink/source/split/IoTDBSourceSplit.java | 80 +++ .../source/split/IoTDBSourceSplitSerializer.java | 83 +++ .../flink/utils/IoTDBIdentifierUtils.java | 39 ++ .../flink-iotdb-table-connector-flink1/pom.xml | 13 + .../relational/flink/catalog/IoTDBCatalog.java | 27 +- .../iotdb/relational/flink/sink/IoTDBSink.java | 5 +- .../relational/flink/sink/IoTDBSinkWriter.java | 134 ++++- .../iotdb/relational/flink/source/IoTDBSource.java | 99 ++++ .../flink/source/IoTDBSourceEnumerator.java | 170 ++++++ .../relational/flink/source/IoTDBSourceReader.java | 229 +++++++++ .../table/IoTDBRelationalDynamicTableFactory.java | 5 +- .../table/IoTDBRelationalDynamicTableSink.java | 3 +- .../table/IoTDBRelationalDynamicTableSource.java | 68 ++- .../flink/IoTDBRelationalLocalQueryManual.java | 95 ++++ 24 files changed, 1977 insertions(+), 77 deletions(-) diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalogClient.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalogClient.java index 7919dad..c7d3a5a 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalogClient.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalogClient.java @@ -23,6 +23,7 @@ import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.SessionDataSet; import org.apache.iotdb.isession.pool.ITableSessionPool; import org.apache.iotdb.relational.flink.cfg.IoTDBRelationalOptions; +import org.apache.iotdb.relational.flink.utils.IoTDBIdentifierUtils; import org.apache.iotdb.session.pool.TableSessionPoolBuilder; import org.apache.flink.table.catalog.exceptions.CatalogException; @@ -40,8 +41,8 @@ import java.util.Locale; /** * Common IoTDB catalog metadata access layer. * - * <p>This class implements metadata and basic DDL operations used by the Flink catalog. It talks - * to IoTDB through the table model session and the following SQL statements: + * <p>This class implements metadata and basic DDL operations used by the Flink catalog. It talks to + * IoTDB through the table model session and the following SQL statements: * * <pre> * SHOW DATABASES @@ -71,8 +72,7 @@ public class IoTDBCatalogClient implements AutoCloseable { } public List<String> listTables(String database) { - return querySingleColumn( - "SHOW TABLES FROM " + quoteIdentifier(database), COLUMN_TABLE_NAME); + return querySingleColumn("SHOW TABLES FROM " + quoteIdentifier(database), COLUMN_TABLE_NAME); } public boolean databaseExists(String database) { @@ -116,11 +116,7 @@ public class IoTDBCatalogClient implements AutoCloseable { } public void dropTable(String database, String table) { - executeNonQuery( - "DROP TABLE " - + quoteIdentifier(database) - + "." - + quoteIdentifier(table)); + executeNonQuery("DROP TABLE " + quoteIdentifier(database) + "." + quoteIdentifier(table)); } public TableSchema getTable(String database, String table) { @@ -235,7 +231,7 @@ public class IoTDBCatalogClient implements AutoCloseable { } private static String quoteIdentifier(String identifier) { - return "\"" + identifier.replace("\"", "\"\"") + "\""; + return IoTDBIdentifierUtils.quoteIdentifier(identifier); } private ITableSessionPool getSessionPool() { diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/cfg/IoTDBRelationalOptions.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/cfg/IoTDBRelationalOptions.java index 3a656c1..b056436 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/cfg/IoTDBRelationalOptions.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/cfg/IoTDBRelationalOptions.java @@ -129,17 +129,23 @@ public class IoTDBRelationalOptions implements Serializable { return defaultDatabase; } - /** @return the configured IoTDB time-column name, or {@code null} when unspecified. */ + /** + * @return the configured IoTDB time-column name, or {@code null} when unspecified. + */ public String getTimeColumn() { return timeColumn; } - /** @return configured IoTDB TAG column names. */ + /** + * @return configured IoTDB TAG column names. + */ public List<String> getTagColumns() { return tagColumns; } - /** @return configured IoTDB ATTRIBUTE column names. */ + /** + * @return configured IoTDB ATTRIBUTE column names. + */ public List<String> getAttributeColumns() { return attributeColumns; } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/IoTDBTabletSerializer.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/IoTDBTabletSerializer.java index 83441e3..cbb2e68 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/IoTDBTabletSerializer.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/IoTDBTabletSerializer.java @@ -25,10 +25,10 @@ import java.io.IOException; import java.io.Serializable; /** - * Converts one Flink input record into an IoTDB {@link Tablet}. + * Serializes Flink input records into a writer-owned IoTDB {@link Tablet}. * - * <p>The returned tablet may contain zero, one, or multiple rows. The sink writer owns batching and - * flushing, while the serializer only defines how an input record is represented as tablet rows. + * <p>The sink writer owns tablet creation, batching, and flushing. The serializer only defines how + * an input record is appended to the provided tablet. * * @param <IN> input record type */ @@ -38,11 +38,19 @@ public interface IoTDBTabletSerializer<IN> extends Serializable { default void open() throws Exception {} /** - * Serializes one input record. + * Creates the writer-owned tablet used for batching. * - * <p>TODO: define the exact owner of the returned tablet and the batch-size contract. + * <p>The returned tablet must have a positive maximum row number. The writer owns the returned + * instance and reuses it until it is flushed. */ - Tablet serialize(IN record) throws IOException; + Tablet createTablet(int maxRows) throws IOException; + + /** + * Serializes one input record directly into the writer-owned tablet. + * + * @return {@code true} if the tablet is full after serialization, otherwise {@code false} + */ + boolean serialize(IN record, Tablet tablet) throws IOException; /** Closes the serializer after the last record has been processed. */ default void close() throws Exception {} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/RowDataIoTDBTabletSerializer.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/RowDataIoTDBTabletSerializer.java index 0a30dda..0119cac 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/RowDataIoTDBTabletSerializer.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/RowDataIoTDBTabletSerializer.java @@ -45,7 +45,12 @@ public class RowDataIoTDBTabletSerializer implements IoTDBTabletSerializer<RowDa } @Override - public Tablet serialize(RowData record) throws IOException { + public Tablet createTablet(int maxRows) throws IOException { + throw new UnsupportedOperationException("Not implemented yet."); + } + + @Override + public boolean serialize(RowData record, Tablet tablet) throws IOException { throw new UnsupportedOperationException("Not implemented yet."); } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/IoTDBTabletSerializer.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/deserializer/IoTDBDeserializationSchema.java similarity index 50% copy from connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/IoTDBTabletSerializer.java copy to connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/deserializer/IoTDBDeserializationSchema.java index 83441e3..a9f762e 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/sink/serializer/IoTDBTabletSerializer.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/deserializer/IoTDBDeserializationSchema.java @@ -17,33 +17,24 @@ * under the License. */ -package org.apache.iotdb.relational.flink.sink.serializer; +package org.apache.iotdb.relational.flink.source.deserializer; -import org.apache.tsfile.write.record.Tablet; +import org.apache.iotdb.isession.SessionDataSet; import java.io.IOException; import java.io.Serializable; /** - * Converts one Flink input record into an IoTDB {@link Tablet}. + * Converts the current row of an IoTDB {@link SessionDataSet.DataIterator} into the output type of + * the Flink source. * - * <p>The returned tablet may contain zero, one, or multiple rows. The sink writer owns batching and - * flushing, while the serializer only defines how an input record is represented as tablet rows. + * <p>The caller owns iteration and must invoke {@link SessionDataSet.DataIterator#next()} before + * calling this method. Implementations must only read the current row and must not advance the + * iterator. * - * @param <IN> input record type + * @param <OUT> output type */ -public interface IoTDBTabletSerializer<IN> extends Serializable { +public interface IoTDBDeserializationSchema<OUT> extends Serializable { - /** Opens the serializer before the first record is serialized. */ - default void open() throws Exception {} - - /** - * Serializes one input record. - * - * <p>TODO: define the exact owner of the returned tablet and the batch-size contract. - */ - Tablet serialize(IN record) throws IOException; - - /** Closes the serializer after the last record has been processed. */ - default void close() throws Exception {} + OUT deserialize(SessionDataSet.DataIterator iterator) throws IOException; } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/deserializer/RowDataDeserializationSchema.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/deserializer/RowDataDeserializationSchema.java new file mode 100644 index 0000000..14a953a --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/deserializer/RowDataDeserializationSchema.java @@ -0,0 +1,114 @@ +/* + * 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.iotdb.relational.flink.source.deserializer; + +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.rpc.StatementExecutionException; + +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.types.RowKind; +import org.apache.tsfile.utils.Binary; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; + +/** + * Table API deserializer that converts the current {@link SessionDataSet.DataIterator} row into + * {@link RowData}. + */ +public class RowDataDeserializationSchema implements IoTDBDeserializationSchema<RowData> { + + private static final long serialVersionUID = 1L; + + private final DataType rowDataType; + private final List<DataType> fieldDataTypes; + + public RowDataDeserializationSchema(DataType rowDataType) { + if (rowDataType.getLogicalType().getTypeRoot() != LogicalTypeRoot.ROW) { + throw new IllegalArgumentException("RowDataDeserializationSchema requires a ROW data type."); + } + this.rowDataType = rowDataType; + this.fieldDataTypes = DataType.getFieldDataTypes(rowDataType); + } + + @Override + public RowData deserialize(SessionDataSet.DataIterator iterator) throws IOException { + GenericRowData row = new GenericRowData(RowKind.INSERT, fieldDataTypes.size()); + for (int i = 0; i < fieldDataTypes.size(); i++) { + row.setField(i, readField(iterator, i + 1, fieldDataTypes.get(i))); + } + return row; + } + + public DataType getRowDataType() { + return rowDataType; + } + + private static Object readField( + SessionDataSet.DataIterator iterator, int columnIndex, DataType dataType) throws IOException { + try { + if (iterator.isNull(columnIndex)) { + return null; + } + + LogicalTypeRoot typeRoot = dataType.getLogicalType().getTypeRoot(); + switch (typeRoot) { + case BOOLEAN: + return iterator.getBoolean(columnIndex); + case TINYINT: + return (byte) iterator.getInt(columnIndex); + case SMALLINT: + return (short) iterator.getInt(columnIndex); + case INTEGER: + return iterator.getInt(columnIndex); + case BIGINT: + return iterator.getLong(columnIndex); + case FLOAT: + return iterator.getFloat(columnIndex); + case DOUBLE: + return iterator.getDouble(columnIndex); + case CHAR: + case VARCHAR: + return StringData.fromString(iterator.getString(columnIndex)); + case BINARY: + case VARBINARY: + Binary binary = iterator.getBlob(columnIndex); + return binary == null ? null : Arrays.copyOf(binary.getValues(), binary.getLength()); + case DATE: + return (int) iterator.getDate(columnIndex).toEpochDay(); + case TIMESTAMP_WITHOUT_TIME_ZONE: + case TIMESTAMP_WITH_LOCAL_TIME_ZONE: + return TimestampData.fromTimestamp(iterator.getTimestamp(columnIndex)); + default: + throw new IOException( + "Unsupported Flink type at column " + columnIndex + ": " + dataType.getLogicalType()); + } + } catch (StatementExecutionException e) { + throw new IOException( + "Failed to read IoTDB column " + columnIndex + " as " + dataType.getLogicalType(), e); + } + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/enumerator/IoTDBSourceEnumeratorState.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/enumerator/IoTDBSourceEnumeratorState.java new file mode 100644 index 0000000..884ebb7 --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/enumerator/IoTDBSourceEnumeratorState.java @@ -0,0 +1,42 @@ +/* + * 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.iotdb.relational.flink.source.enumerator; + +import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplit; + +import java.io.Serializable; +import java.util.Collections; +import java.util.List; + +/** Checkpoint state of the IoTDB source enumerator. */ +public class IoTDBSourceEnumeratorState implements Serializable { + + private static final long serialVersionUID = 1L; + + private final List<IoTDBSourceSplit> remainingSplits; + + public IoTDBSourceEnumeratorState(List<IoTDBSourceSplit> remainingSplits) { + this.remainingSplits = remainingSplits; + } + + public List<IoTDBSourceSplit> getRemainingSplits() { + return remainingSplits == null ? Collections.emptyList() : remainingSplits; + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/enumerator/IoTDBSourceEnumeratorStateSerializer.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/enumerator/IoTDBSourceEnumeratorStateSerializer.java new file mode 100644 index 0000000..67b559e --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/enumerator/IoTDBSourceEnumeratorStateSerializer.java @@ -0,0 +1,81 @@ +/* + * 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.iotdb.relational.flink.source.enumerator; + +import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplit; +import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplitSerializer; + +import org.apache.flink.core.io.SimpleVersionedSerializer; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +/** Serializer for {@link IoTDBSourceEnumeratorState}. */ +public class IoTDBSourceEnumeratorStateSerializer + implements SimpleVersionedSerializer<IoTDBSourceEnumeratorState> { + + private static final int VERSION = 1; + + private final IoTDBSourceSplitSerializer splitSerializer = new IoTDBSourceSplitSerializer(); + + @Override + public int getVersion() { + return VERSION; + } + + @Override + public byte[] serialize(IoTDBSourceEnumeratorState state) throws IOException { + List<IoTDBSourceSplit> splits = state.getRemainingSplits(); + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (DataOutputStream output = new DataOutputStream(bytes)) { + output.writeInt(splits.size()); + for (IoTDBSourceSplit split : splits) { + byte[] serializedSplit = splitSerializer.serialize(split); + output.writeInt(serializedSplit.length); + output.write(serializedSplit); + } + } + return bytes.toByteArray(); + } + + @Override + public IoTDBSourceEnumeratorState deserialize(int version, byte[] serialized) throws IOException { + if (version != VERSION) { + throw new IOException( + "Unsupported IoTDB source enumerator state serializer version: " + version); + } + try (DataInputStream input = new DataInputStream(new ByteArrayInputStream(serialized))) { + int splitCount = input.readInt(); + List<IoTDBSourceSplit> splits = new ArrayList<>(splitCount); + for (int i = 0; i < splitCount; i++) { + int splitLength = input.readInt(); + byte[] serializedSplit = new byte[splitLength]; + input.readFully(serializedSplit); + splits.add(splitSerializer.deserialize(splitSerializer.getVersion(), serializedSplit)); + } + return new IoTDBSourceEnumeratorState(splits); + } + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBExpressionVisitor.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBExpressionVisitor.java new file mode 100644 index 0000000..b970986 --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBExpressionVisitor.java @@ -0,0 +1,568 @@ +/* + * 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.iotdb.relational.flink.source.pushdown; + +import org.apache.iotdb.relational.flink.utils.IoTDBIdentifierUtils; +import org.apache.iotdb.relational.flink.utils.IoTDBRelationalTypeUtils; + +import org.apache.flink.table.expressions.CallExpression; +import org.apache.flink.table.expressions.Expression; +import org.apache.flink.table.expressions.ExpressionVisitor; +import org.apache.flink.table.expressions.FieldReferenceExpression; +import org.apache.flink.table.expressions.ResolvedExpression; +import org.apache.flink.table.expressions.TypeLiteralExpression; +import org.apache.flink.table.expressions.ValueLiteralExpression; +import org.apache.flink.table.functions.BuiltInFunctionDefinition; +import org.apache.flink.table.functions.BuiltInFunctionDefinitions; +import org.apache.flink.table.functions.FunctionDefinition; +import org.apache.flink.table.functions.FunctionKind; + +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +/** Converts supported Flink scalar expressions into IoTDB SQL expressions. */ +public class IoTDBExpressionVisitor implements ExpressionVisitor<String> { + + private static final Map<String, String> FLINK_TO_IOTDB_FUNCTION_NAMES; + + static { + Map<String, String> functionNames = new HashMap<>(); + functionNames.put("abs", "abs"); + functionNames.put("acos", "acos"); + functionNames.put("asin", "asin"); + functionNames.put("atan", "atan"); + functionNames.put("ceil", "ceil"); + functionNames.put("charlength", "length"); + functionNames.put("concat", "concat"); + functionNames.put("cos", "cos"); + functionNames.put("cosh", "cosh"); + functionNames.put("degrees", "degrees"); + functionNames.put("e", "e"); + functionNames.put("exp", "exp"); + functionNames.put("floor", "floor"); + functionNames.put("greatest", "greatest"); + functionNames.put("least", "least"); + functionNames.put("length", "length"); + functionNames.put("ln", "ln"); + functionNames.put("log10", "log10"); + functionNames.put("lower", "lower"); + functionNames.put("lowercase", "lower"); + functionNames.put("ltrim", "ltrim"); + functionNames.put("pi", "pi"); + functionNames.put("radians", "radians"); + functionNames.put("regexp", "regexp_like"); + functionNames.put("replace", "replace"); + functionNames.put("round", "round"); + functionNames.put("rtrim", "rtrim"); + functionNames.put("sign", "sign"); + functionNames.put("sin", "sin"); + functionNames.put("sinh", "sinh"); + functionNames.put("sqrt", "sqrt"); + functionNames.put("substr", "substring"); + functionNames.put("substring", "substring"); + functionNames.put("tan", "tan"); + functionNames.put("tanh", "tanh"); + functionNames.put("trim", "trim"); + functionNames.put("upper", "upper"); + functionNames.put("uppercase", "upper"); + FLINK_TO_IOTDB_FUNCTION_NAMES = Collections.unmodifiableMap(functionNames); + } + + @Override + public String visit(CallExpression call) { + if (call == null || call.getFunctionDefinition().getKind() != FunctionKind.SCALAR) { + return null; + } + + try { + return visitGeneralScalarExpression(call); + } catch (RuntimeException e) { + return null; + } + } + + @Override + public String visit(ValueLiteralExpression valueLiteral) { + return IoTDBLiteralUtils.render(valueLiteral); + } + + @Override + public String visit(FieldReferenceExpression fieldReference) { + if (fieldReference == null || fieldReference.getInputIndex() != 0) { + return null; + } + return IoTDBIdentifierUtils.quoteIdentifier(fieldReference.getName()); + } + + @Override + public String visit(TypeLiteralExpression typeLiteral) { + return renderType(typeLiteral); + } + + @Override + public String visit(Expression expression) { + return null; + } + + private String visitGeneralScalarExpression(CallExpression call) { + if (!(call.getFunctionDefinition() instanceof BuiltInFunctionDefinition)) { + return null; + } + + List<ResolvedExpression> children = call.getResolvedChildren(); + String expressionName = getExpressionName(call); + + switch (expressionName) { + case "=": + return visitEqualTo(children); + case "<>": + return visitNotEqualTo(children); + case "<": + return visitLess(children); + case "<=": + return visitLessOrEqual(children); + case ">": + return visitGreater(children); + case ">=": + return visitGreaterOrEqual(children); + case "+": + case "-": + case "*": + case "/": + case "%": + return visitArithmeticBinary(expressionName, children); + case "-u": + return visitArithmeticUnary(children); + case "and": + return visitAnd(children); + case "or": + return visitOr(children); + case "not": + return visitNot(children); + case "is_null": + return visitIsNull(children); + case "is_not_null": + return visitIsNotNull(children); + case "like": + return visitLike(children); + case "in": + return visitIn(children); + case "between": + return visitBetween(children); + case "not_between": + return visitNotBetween(children); + case "cast": + return visitCast("CAST", children); + case "try_cast": + return visitCast("TRY_CAST", children); + case "if": + return visitIf(children); + case "coalesce": + return visitCoalesce(children); + case "position": + return visitPosition(children); + case "locate": + return visitLocate(children); + case "instr": + return visitInstr(children); + case "current_database": + return visitCurrentTime("CURRENT_DATABASE", children); + case "current_date": + return visitCurrentTime("CURRENT_DATE", children); + case "current_time": + return visitCurrentTime("CURRENT_TIME", children); + case "current_timestamp": + return visitCurrentTime("CURRENT_TIMESTAMP", children); + case "localtime": + return visitCurrentTime("LOCALTIME", children); + case "localtimestamp": + return visitCurrentTime("LOCALTIMESTAMP", children); + default: + return visitScalarFunction(expressionName, children); + } + } + + private String visitEqualTo(List<ResolvedExpression> children) { + return visitBinary("=", children); + } + + private String visitNotEqualTo(List<ResolvedExpression> children) { + return visitBinary("<>", children); + } + + private String visitLess(List<ResolvedExpression> children) { + return visitBinary("<", children); + } + + private String visitLessOrEqual(List<ResolvedExpression> children) { + return visitBinary("<=", children); + } + + private String visitGreater(List<ResolvedExpression> children) { + return visitBinary(">", children); + } + + private String visitGreaterOrEqual(List<ResolvedExpression> children) { + return visitBinary(">=", children); + } + + private String visitArithmeticBinary(String operator, List<ResolvedExpression> children) { + return visitBinary(operator, children); + } + + private String visitArithmeticUnary(List<ResolvedExpression> children) { + return visitPrefix("-", children); + } + + private String visitAnd(List<ResolvedExpression> children) { + return visitVariadic("AND", children); + } + + private String visitOr(List<ResolvedExpression> children) { + return visitVariadic("OR", children); + } + + private String visitNot(List<ResolvedExpression> children) { + return visitPrefix("NOT", children); + } + + private String visitIsNull(List<ResolvedExpression> children) { + return visitSuffix("IS NULL", children); + } + + private String visitIsNotNull(List<ResolvedExpression> children) { + return visitSuffix("IS NOT NULL", children); + } + + private String visitLike(List<ResolvedExpression> children) { + return visitBinary("LIKE", children); + } + + private String visitIn(List<ResolvedExpression> children) { + if (children == null || children.size() < 2) { + return null; + } + + String value = buildIoTDBExpressionSQL(children.get(0)); + if (value == null) { + return null; + } + + StringBuilder builder = new StringBuilder("(").append(value).append(" IN ("); + for (int i = 1; i < children.size(); i++) { + String child = buildIoTDBExpressionSQL(children.get(i)); + if (child == null) { + return null; + } + if (i > 1) { + builder.append(", "); + } + builder.append(child); + } + return builder.append("))").toString(); + } + + private String visitBetween(List<ResolvedExpression> children) { + return visitBetween("BETWEEN", children); + } + + private String visitNotBetween(List<ResolvedExpression> children) { + return visitBetween("NOT BETWEEN", children); + } + + private String visitBetween(String operator, List<ResolvedExpression> children) { + if (children == null || children.size() != 3) { + return null; + } + + String value = buildIoTDBExpressionSQL(children.get(0)); + String lowerBound = buildIoTDBExpressionSQL(children.get(1)); + String upperBound = buildIoTDBExpressionSQL(children.get(2)); + if (value == null || lowerBound == null || upperBound == null) { + return null; + } + return "(" + value + " " + operator + " " + lowerBound + " AND " + upperBound + ")"; + } + + private String visitCast(String keyword, List<ResolvedExpression> children) { + if (children == null || children.size() != 2) { + return null; + } + + String value = buildIoTDBExpressionSQL(children.get(0)); + String type = renderType(children.get(1)); + if (value == null || type == null) { + return null; + } + return keyword + "(" + value + " AS " + type + ")"; + } + + private String visitIf(List<ResolvedExpression> children) { + return visitFunction("if", children); + } + + private String visitCoalesce(List<ResolvedExpression> children) { + return visitFunction("coalesce", children); + } + + private String visitPosition(List<ResolvedExpression> children) { + if (children == null || children.size() != 2) { + return null; + } + return visitStringPosition(children.get(1), children.get(0)); + } + + private String visitLocate(List<ResolvedExpression> children) { + if (children == null || children.size() != 2) { + return null; + } + return visitStringPosition(children.get(1), children.get(0)); + } + + private String visitInstr(List<ResolvedExpression> children) { + if (children == null || children.size() != 2) { + return null; + } + return visitStringPosition(children.get(0), children.get(1)); + } + + private String visitStringPosition( + ResolvedExpression valueExpression, ResolvedExpression searchExpression) { + String value = buildIoTDBExpressionSQL(valueExpression); + String search = buildIoTDBExpressionSQL(searchExpression); + if (value == null || search == null) { + return null; + } + return "strpos(" + value + ", " + search + ")"; + } + + private String visitScalarFunction(String functionName, List<ResolvedExpression> children) { + String iotdbFunctionName = FLINK_TO_IOTDB_FUNCTION_NAMES.get(functionName); + return iotdbFunctionName == null ? null : visitFunction(iotdbFunctionName, children); + } + + private String visitFunction(String functionName, List<ResolvedExpression> children) { + StringBuilder builder = new StringBuilder(functionName).append('('); + if (children != null) { + for (int i = 0; i < children.size(); i++) { + String child = buildIoTDBExpressionSQL(children.get(i)); + if (child == null) { + return null; + } + if (i > 0) { + builder.append(", "); + } + builder.append(child); + } + } + return builder.append(')').toString(); + } + + private String visitCurrentTime(String keyword, List<ResolvedExpression> children) { + return children == null || children.isEmpty() ? keyword : null; + } + + private String visitBinary(String operator, List<ResolvedExpression> children) { + if (children == null || children.size() != 2) { + return null; + } + + String left = buildIoTDBExpressionSQL(children.get(0)); + String right = buildIoTDBExpressionSQL(children.get(1)); + if (left == null || right == null) { + return null; + } + return "(" + left + " " + operator + " " + right + ")"; + } + + private String visitVariadic(String operator, List<ResolvedExpression> children) { + if (children == null || children.size() < 2) { + return null; + } + + StringBuilder builder = new StringBuilder("("); + for (int i = 0; i < children.size(); i++) { + String child = buildIoTDBExpressionSQL(children.get(i)); + if (child == null) { + return null; + } + if (i > 0) { + builder.append(' ').append(operator).append(' '); + } + builder.append(child); + } + return builder.append(')').toString(); + } + + private String visitPrefix(String operator, List<ResolvedExpression> children) { + if (children == null || children.size() != 1) { + return null; + } + + String child = buildIoTDBExpressionSQL(children.get(0)); + return child == null ? null : "(" + operator + " " + child + ")"; + } + + private String visitSuffix(String operator, List<ResolvedExpression> children) { + if (children == null || children.size() != 1) { + return null; + } + + String child = buildIoTDBExpressionSQL(children.get(0)); + return child == null ? null : "(" + child + " " + operator + ")"; + } + + private String buildIoTDBExpressionSQL(ResolvedExpression expression) { + return expression == null ? null : expression.accept(this); + } + + private String renderType(ResolvedExpression expression) { + if (!(expression instanceof TypeLiteralExpression)) { + return null; + } + try { + return IoTDBRelationalTypeUtils.toIoTDBDataType(expression.getOutputDataType()).name(); + } catch (RuntimeException e) { + return null; + } + } + + private static String getExpressionName(CallExpression call) { + FunctionDefinition functionDefinition = call.getFunctionDefinition(); + + if (BuiltInFunctionDefinitions.EQUALS.equals(functionDefinition)) { + return "="; + } + if (BuiltInFunctionDefinitions.NOT_EQUALS.equals(functionDefinition)) { + return "<>"; + } + if (BuiltInFunctionDefinitions.LESS_THAN.equals(functionDefinition)) { + return "<"; + } + if (BuiltInFunctionDefinitions.LESS_THAN_OR_EQUAL.equals(functionDefinition)) { + return "<="; + } + if (BuiltInFunctionDefinitions.GREATER_THAN.equals(functionDefinition)) { + return ">"; + } + if (BuiltInFunctionDefinitions.GREATER_THAN_OR_EQUAL.equals(functionDefinition)) { + return ">="; + } + if (BuiltInFunctionDefinitions.PLUS.equals(functionDefinition)) { + return "+"; + } + if (BuiltInFunctionDefinitions.MINUS.equals(functionDefinition)) { + return "-"; + } + if (BuiltInFunctionDefinitions.TIMES.equals(functionDefinition)) { + return "*"; + } + if (BuiltInFunctionDefinitions.DIVIDE.equals(functionDefinition)) { + return "/"; + } + if (BuiltInFunctionDefinitions.MOD.equals(functionDefinition)) { + return "%"; + } + if (BuiltInFunctionDefinitions.MINUS_PREFIX.equals(functionDefinition)) { + return "-u"; + } + if (BuiltInFunctionDefinitions.AND.equals(functionDefinition)) { + return "and"; + } + if (BuiltInFunctionDefinitions.OR.equals(functionDefinition)) { + return "or"; + } + if (BuiltInFunctionDefinitions.NOT.equals(functionDefinition)) { + return "not"; + } + if (BuiltInFunctionDefinitions.IS_NULL.equals(functionDefinition)) { + return "is_null"; + } + if (BuiltInFunctionDefinitions.IS_NOT_NULL.equals(functionDefinition)) { + return "is_not_null"; + } + if (BuiltInFunctionDefinitions.LIKE.equals(functionDefinition)) { + return "like"; + } + if (BuiltInFunctionDefinitions.IN.equals(functionDefinition)) { + return "in"; + } + if (BuiltInFunctionDefinitions.BETWEEN.equals(functionDefinition)) { + return "between"; + } + if (BuiltInFunctionDefinitions.NOT_BETWEEN.equals(functionDefinition)) { + return "not_between"; + } + if (BuiltInFunctionDefinitions.CAST.equals(functionDefinition)) { + return "cast"; + } + if (BuiltInFunctionDefinitions.TRY_CAST.equals(functionDefinition)) { + return "try_cast"; + } + if (BuiltInFunctionDefinitions.IF_NULL.equals(functionDefinition)) { + return "coalesce"; + } + if (BuiltInFunctionDefinitions.IF.equals(functionDefinition)) { + return "if"; + } + if (BuiltInFunctionDefinitions.COALESCE.equals(functionDefinition)) { + return "coalesce"; + } + if (BuiltInFunctionDefinitions.CURRENT_DATABASE.equals(functionDefinition)) { + return "current_database"; + } + if (BuiltInFunctionDefinitions.CURRENT_DATE.equals(functionDefinition)) { + return "current_date"; + } + if (BuiltInFunctionDefinitions.CURRENT_TIME.equals(functionDefinition)) { + return "current_time"; + } + if (BuiltInFunctionDefinitions.CURRENT_TIMESTAMP.equals(functionDefinition)) { + return "current_timestamp"; + } + if (BuiltInFunctionDefinitions.LOCAL_TIME.equals(functionDefinition)) { + return "localtime"; + } + if (BuiltInFunctionDefinitions.LOCAL_TIMESTAMP.equals(functionDefinition)) { + return "localtimestamp"; + } + if (BuiltInFunctionDefinitions.NOW.equals(functionDefinition)) { + return "current_timestamp"; + } + + String functionName = normalizeFunctionName(call.getFunctionName()); + String mappedFunctionName = FLINK_TO_IOTDB_FUNCTION_NAMES.get(functionName); + return mappedFunctionName == null ? functionName : mappedFunctionName; + } + + private static String normalizeFunctionName(String functionName) { + if (functionName == null) { + return ""; + } + int separatorIndex = functionName.lastIndexOf('.'); + if (separatorIndex >= 0) { + functionName = functionName.substring(separatorIndex + 1); + } + return functionName.replace("`", "").toLowerCase(Locale.ROOT); + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBLiteralUtils.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBLiteralUtils.java new file mode 100644 index 0000000..066093f --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBLiteralUtils.java @@ -0,0 +1,115 @@ +/* + * 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.iotdb.relational.flink.source.pushdown; + +import org.apache.flink.table.expressions.ValueLiteralExpression; +import org.apache.flink.table.types.logical.LogicalTypeRoot; + +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.util.Optional; + +/** Renders Flink literal expressions as IoTDB SQL literals. */ +public final class IoTDBLiteralUtils { + + private IoTDBLiteralUtils() {} + + /** + * Renders a literal for pushdown. Returns {@code null} if the literal type or value cannot be + * represented in IoTDB SQL. + */ + public static String render(ValueLiteralExpression literal) { + if (literal == null || literal.isNull()) { + return null; + } + + try { + LogicalTypeRoot typeRoot = literal.getOutputDataType().getLogicalType().getTypeRoot(); + switch (typeRoot) { + case BOOLEAN: + return literal.getValueAs(Boolean.class).map(String::valueOf).orElse(null); + case TINYINT: + return value(literal.getValueAs(Byte.class)); + case SMALLINT: + return value(literal.getValueAs(Short.class)); + case INTEGER: + return value(literal.getValueAs(Integer.class)); + case BIGINT: + return value(literal.getValueAs(Long.class)); + case FLOAT: + return literal + .getValueAs(Float.class) + .filter(value -> !value.isNaN() && !value.isInfinite()) + .map(String::valueOf) + .orElse(null); + case DOUBLE: + return literal + .getValueAs(Double.class) + .filter(value -> !value.isNaN() && !value.isInfinite()) + .map(String::valueOf) + .orElse(null); + case CHAR: + case VARCHAR: + return literal.getValueAs(String.class).map(IoTDBLiteralUtils::quoteString).orElse(null); + case BINARY: + case VARBINARY: + return literal.getValueAs(byte[].class).map(IoTDBLiteralUtils::formatBinary).orElse(null); + case DATE: + return literal + .getValueAs(LocalDate.class) + .map(value -> "CAST('" + value + "' AS DATE)") + .orElse(null); + case TIMESTAMP_WITHOUT_TIME_ZONE: + return literal + .getValueAs(LocalDateTime.class) + .map(value -> "CAST('" + value + "' AS TIMESTAMP)") + .orElse(null); + case TIMESTAMP_WITH_LOCAL_TIME_ZONE: + return literal + .getValueAs(Instant.class) + .map(value -> "CAST('" + value + "' AS TIMESTAMP)") + .orElse(null); + default: + return null; + } + } catch (RuntimeException e) { + return null; + } + } + + private static String value(Optional<?> value) { + return value.map(Object::toString).orElse(null); + } + + private static String quoteString(String value) { + return "'" + value.replace("'", "''") + "'"; + } + + private static String formatBinary(byte[] value) { + StringBuilder builder = new StringBuilder(value.length * 2 + 3); + builder.append("X'"); + for (byte b : value) { + builder.append(String.format("%02X", b)); + } + builder.append("'"); + return builder.toString(); + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/split/IoTDBSourceSplit.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/split/IoTDBSourceSplit.java new file mode 100644 index 0000000..23dcf5f --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/split/IoTDBSourceSplit.java @@ -0,0 +1,80 @@ +/* + * 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.iotdb.relational.flink.source.split; + +import org.apache.flink.api.connector.source.SourceSplit; + +import java.io.Serializable; +import java.util.Objects; + +/** Read split for the IoTDB relational table source. */ +public class IoTDBSourceSplit implements SourceSplit, Serializable { + + private static final long serialVersionUID = 1L; + + private final String splitId; + private final String database; + private final String table; + private final String sql; + + public IoTDBSourceSplit(String splitId, String database, String table, String sql) { + this.splitId = splitId; + this.database = database; + this.table = table; + this.sql = sql; + } + + @Override + public String splitId() { + return splitId; + } + + public String getDatabase() { + return database; + } + + public String getTable() { + return table; + } + + public String getSql() { + return sql; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof IoTDBSourceSplit)) { + return false; + } + IoTDBSourceSplit that = (IoTDBSourceSplit) o; + return Objects.equals(splitId, that.splitId) + && Objects.equals(database, that.database) + && Objects.equals(table, that.table) + && Objects.equals(sql, that.sql); + } + + @Override + public int hashCode() { + return Objects.hash(splitId, database, table, sql); + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/split/IoTDBSourceSplitSerializer.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/split/IoTDBSourceSplitSerializer.java new file mode 100644 index 0000000..2cf453f --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/split/IoTDBSourceSplitSerializer.java @@ -0,0 +1,83 @@ +/* + * 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.iotdb.relational.flink.source.split; + +import org.apache.flink.core.io.SimpleVersionedSerializer; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; + +/** Serializer for {@link IoTDBSourceSplit}. */ +public class IoTDBSourceSplitSerializer implements SimpleVersionedSerializer<IoTDBSourceSplit> { + + private static final int VERSION = 1; + + @Override + public int getVersion() { + return VERSION; + } + + @Override + public byte[] serialize(IoTDBSourceSplit split) throws IOException { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (DataOutputStream output = new DataOutputStream(bytes)) { + writeString(output, split.splitId()); + writeString(output, split.getDatabase()); + writeString(output, split.getTable()); + writeString(output, split.getSql()); + } + return bytes.toByteArray(); + } + + @Override + public IoTDBSourceSplit deserialize(int version, byte[] serialized) throws IOException { + if (version != VERSION) { + throw new IOException("Unsupported IoTDB source split serializer version: " + version); + } + try (DataInputStream input = new DataInputStream(new ByteArrayInputStream(serialized))) { + return new IoTDBSourceSplit( + readString(input), readString(input), readString(input), readString(input)); + } + } + + private static void writeString(DataOutputStream output, String value) throws IOException { + if (value == null) { + output.writeInt(-1); + return; + } + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + output.writeInt(bytes.length); + output.write(bytes); + } + + private static String readString(DataInputStream input) throws IOException { + int length = input.readInt(); + if (length < 0) { + return null; + } + byte[] bytes = new byte[length]; + input.readFully(bytes); + return new String(bytes, StandardCharsets.UTF_8); + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/utils/IoTDBIdentifierUtils.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/utils/IoTDBIdentifierUtils.java new file mode 100644 index 0000000..a34a437 --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/utils/IoTDBIdentifierUtils.java @@ -0,0 +1,39 @@ +/* + * 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.iotdb.relational.flink.utils; + +/** Utilities for rendering IoTDB table-model identifiers in SQL. */ +public final class IoTDBIdentifierUtils { + + private IoTDBIdentifierUtils() {} + + /** + * Quotes a resolved logical identifier using IoTDB's double-quote syntax. + * + * <p>The input is treated as a logical name, not as SQL text. Backticks are preserved as part of + * the identifier and double quotes are escaped by doubling. + */ + public static String quoteIdentifier(String identifier) { + if (identifier == null) { + throw new IllegalArgumentException("Identifier must not be null."); + } + return "\"" + identifier.replace("\"", "\"\"") + "\""; + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/pom.xml b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/pom.xml index c1a292a..2733253 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/pom.xml +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/pom.xml @@ -57,5 +57,18 @@ <groupId>org.apache.flink</groupId> <artifactId>flink-connector-base</artifactId> </dependency> + <!-- Temporary dependencies for local manual query verification. Remove after debugging. --> + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-clients</artifactId> + <version>${flink.version}</version> + <scope>test</scope> + </dependency> + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-table-planner-loader</artifactId> + <version>${flink.version}</version> + <scope>test</scope> + </dependency> </dependencies> </project> diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalog.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalog.java index cf4a5b9..249bd73 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalog.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/catalog/IoTDBCatalog.java @@ -61,8 +61,8 @@ import java.util.Set; * Flink 1.x Catalog adapter for IoTDB relational tables. * * <p>Flink Catalog API adaptation belongs here. IoTDB metadata access belongs in {@link - * IoTDBCatalogClient}; the current implementation supports database/table discovery, table - * schema resolution, and basic database/table DDL. + * IoTDBCatalogClient}; the current implementation supports database/table discovery, table schema + * resolution, and basic database/table DDL. */ public class IoTDBCatalog extends AbstractCatalog { @@ -363,8 +363,7 @@ public class IoTDBCatalog extends AbstractCatalog { } String columnName = column.getName(); - TSDataType dataType = - IoTDBRelationalTypeUtils.toIoTDBDataType((DataType) abstractDataType); + TSDataType dataType = IoTDBRelationalTypeUtils.toIoTDBDataType((DataType) abstractDataType); columnNames.add(columnName); dataTypes.add(dataType); dataTypesByColumn.put(normalizeColumnName(columnName), dataType); @@ -397,7 +396,8 @@ public class IoTDBCatalog extends AbstractCatalog { } for (String columnName : value.split(",", -1)) { if (columnName.trim().isEmpty()) { - throw new CatalogException("Table option '" + optionName + "' contains an empty column name."); + throw new CatalogException( + "Table option '" + optionName + "' contains an empty column name."); } String normalizedColumnName = normalizeColumnName(columnName); if (!columnNames.add(normalizedColumnName)) { @@ -438,15 +438,16 @@ public class IoTDBCatalog extends AbstractCatalog { String columnName, String optionName, Map<String, TSDataType> dataTypesByColumn) { if (!dataTypesByColumn.containsKey(columnName)) { throw new CatalogException( - "Column '" + columnName + "' declared by table option '" + optionName + "' does not exist."); + "Column '" + + columnName + + "' declared by table option '" + + optionName + + "' does not exist."); } } private static ColumnCategory toColumnCategory( - String columnName, - String timeColumn, - Set<String> tagColumns, - Set<String> attributeColumns) { + String columnName, String timeColumn, Set<String> tagColumns, Set<String> attributeColumns) { String normalizedColumnName = normalizeColumnName(columnName); if (timeColumn.equals(normalizedColumnName)) { return ColumnCategory.TIME; @@ -475,8 +476,7 @@ public class IoTDBCatalog extends AbstractCatalog { for (int i = 0; i < columns.size(); i++) { IMeasurementSchema column = columns.get(i); schemaBuilder.column( - column.getMeasurementName(), - IoTDBRelationalTypeUtils.toFlinkDataType(column.getType())); + column.getMeasurementName(), IoTDBRelationalTypeUtils.toFlinkDataType(column.getType())); switch (categories.get(i)) { case TIME: timeColumn = column.getMeasurementName(); @@ -492,7 +492,8 @@ public class IoTDBCatalog extends AbstractCatalog { } } if (timeColumn == null) { - throw new CatalogException("IoTDB table has no TIME column: " + databaseName + "." + tableName); + throw new CatalogException( + "IoTDB table has no TIME column: " + databaseName + "." + tableName); } Map<String, String> tableOptions = new HashMap<>(); diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSink.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSink.java index 46101af..b676c4a 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSink.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSink.java @@ -37,7 +37,7 @@ import java.io.IOException; * <li>DataStream API can create {@code IoTDBSink<IN>} with a user-provided serializer. * </ul> * - * <p>TODO: implement batching and {@code ITableSession.insert(Tablet)} in the writer. + * <p>The writer owns the IoTDB session and batches serialized tablets before insertion. * * @param <IN> input record type */ @@ -48,8 +48,7 @@ public class IoTDBSink<IN> implements Sink<IN> { private final IoTDBRelationalOptions options; private final IoTDBTabletSerializer<IN> serializer; - public IoTDBSink( - IoTDBRelationalOptions options, IoTDBTabletSerializer<IN> serializer) { + public IoTDBSink(IoTDBRelationalOptions options, IoTDBTabletSerializer<IN> serializer) { this.options = options; this.serializer = serializer; } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSinkWriter.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSinkWriter.java index d0a8e30..681eacb 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSinkWriter.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/sink/IoTDBSinkWriter.java @@ -19,44 +19,164 @@ package org.apache.iotdb.relational.flink.sink; +import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.relational.flink.cfg.IoTDBRelationalOptions; import org.apache.iotdb.relational.flink.sink.serializer.IoTDBTabletSerializer; +import org.apache.iotdb.session.TableSessionBuilder; import org.apache.flink.api.connector.sink2.SinkWriter; +import org.apache.tsfile.write.record.Tablet; import java.io.IOException; /** * Sink writer of the IoTDB relational (table model) Flink connector. * - * <p>TODO: open one {@code ITableSession}, buffer serialized tablets, flush through {@code - * ITableSession.insert(Tablet)}, and close all resources. + * <p>The writer owns one IoTDB session and one buffered tablet. Serialized tablets are validated, + * merged into the buffer, and flushed through {@code ITableSession.insert(Tablet)}. * * @param <IN> input record type */ public class IoTDBSinkWriter<IN> implements SinkWriter<IN> { + private static final int DEFAULT_BATCH_SIZE = 1024; + private final IoTDBRelationalOptions options; private final IoTDBTabletSerializer<IN> serializer; - public IoTDBSinkWriter( - IoTDBRelationalOptions options, IoTDBTabletSerializer<IN> serializer) { + private ITableSession session; + private Tablet buffer; + private boolean closed; + + public IoTDBSinkWriter(IoTDBRelationalOptions options, IoTDBTabletSerializer<IN> serializer) + throws IOException { this.options = options; this.serializer = serializer; + open(); } @Override public void write(IN element, Context context) throws IOException, InterruptedException { - throw new UnsupportedOperationException("Not implemented yet."); + ensureOpen(); + + if (buffer.getRowSize() >= buffer.getMaxRowNumber()) { + flushBuffer(); + } + + int rowSizeBefore = buffer.getRowSize(); + boolean full = serializer.serialize(element, buffer); + int rowSizeAfter = buffer.getRowSize(); + + if (rowSizeAfter < rowSizeBefore || rowSizeAfter > buffer.getMaxRowNumber()) { + throw new IOException("Serializer produced an invalid tablet row count."); + } + if (full || rowSizeAfter >= buffer.getMaxRowNumber()) { + flushBuffer(); + } } @Override public void flush(boolean endOfInput) throws IOException, InterruptedException { - throw new UnsupportedOperationException("Not implemented yet."); + ensureOpen(); + flushBuffer(); } @Override public void close() throws Exception { - // TODO: flush remaining rows and close the IoTDB session. + if (closed) { + return; + } + closed = true; + + Exception failure = null; + try { + if (session != null) { + flushBuffer(); + } + } catch (Exception e) { + failure = e; + } + + try { + serializer.close(); + } catch (Exception e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + + try { + if (session != null) { + session.close(); + session = null; + } + } catch (Exception e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + + if (failure != null) { + throw failure; + } + } + + private void open() throws IOException { + try { + TableSessionBuilder builder = + new TableSessionBuilder() + .nodeUrls(options.getNodeUrls()) + .username(options.getUsername()) + .password(options.getPassword()); + if (options.getDatabase() != null) { + builder.database(options.getDatabase()); + } + session = builder.build(); + serializer.open(); + buffer = serializer.createTablet(DEFAULT_BATCH_SIZE); + if (buffer == null || buffer.getMaxRowNumber() <= 0) { + throw new IOException("Serializer created an invalid tablet buffer."); + } + } catch (Exception e) { + closeQuietly(); + throw new IOException("Failed to open IoTDB sink writer.", e); + } + } + + private void ensureOpen() throws IOException { + if (closed || session == null) { + throw new IOException("IoTDB sink writer is already closed."); + } + } + + private void flushBuffer() throws IOException { + if (buffer != null && buffer.getRowSize() > 0) { + insert(buffer); + buffer.reset(); + } + } + + private void insert(Tablet tablet) throws IOException { + try { + session.insert(tablet); + } catch (Exception e) { + throw new IOException("Failed to insert tablet into IoTDB.", e); + } + } + + private void closeQuietly() { + if (session != null) { + try { + session.close(); + } catch (Exception ignored) { + // Preserve the original open failure. + } finally { + session = null; + } + } } } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSource.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSource.java new file mode 100644 index 0000000..0c7e10e --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSource.java @@ -0,0 +1,99 @@ +/* + * 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.iotdb.relational.flink.source; + +import org.apache.iotdb.relational.flink.cfg.IoTDBRelationalOptions; +import org.apache.iotdb.relational.flink.source.deserializer.IoTDBDeserializationSchema; +import org.apache.iotdb.relational.flink.source.enumerator.IoTDBSourceEnumeratorState; +import org.apache.iotdb.relational.flink.source.enumerator.IoTDBSourceEnumeratorStateSerializer; +import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplit; +import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplitSerializer; + +import org.apache.flink.api.connector.source.Boundedness; +import org.apache.flink.api.connector.source.Source; +import org.apache.flink.api.connector.source.SourceReader; +import org.apache.flink.api.connector.source.SourceReaderContext; +import org.apache.flink.api.connector.source.SplitEnumerator; +import org.apache.flink.api.connector.source.SplitEnumeratorContext; +import org.apache.flink.core.io.SimpleVersionedSerializer; +import org.apache.flink.table.types.DataType; + +import java.util.ArrayList; +import java.util.List; + +/** + * Bounded FLIP-27 source for reading an IoTDB table-model table. + * + * <p>The first implementation uses a single split and emits insert-only rows. + * + * @param <OUT> source output type + */ +public class IoTDBSource<OUT> implements Source<OUT, IoTDBSourceSplit, IoTDBSourceEnumeratorState> { + + private static final long serialVersionUID = 1L; + + private final IoTDBRelationalOptions options; + private final DataType rowDataType; + private final IoTDBDeserializationSchema<OUT> deserializer; + private final List<String> filterQueries; + + public IoTDBSource( + IoTDBRelationalOptions options, + DataType rowDataType, + IoTDBDeserializationSchema<OUT> deserializer, + List<String> filterQueries) { + this.options = options; + this.rowDataType = rowDataType; + this.deserializer = deserializer; + this.filterQueries = filterQueries == null ? new ArrayList<>() : new ArrayList<>(filterQueries); + } + + @Override + public Boundedness getBoundedness() { + return Boundedness.BOUNDED; + } + + @Override + public SourceReader<OUT, IoTDBSourceSplit> createReader(SourceReaderContext readerContext) { + return new IoTDBSourceReader<>(readerContext, options, rowDataType, deserializer); + } + + @Override + public SplitEnumerator<IoTDBSourceSplit, IoTDBSourceEnumeratorState> createEnumerator( + SplitEnumeratorContext<IoTDBSourceSplit> enumContext) { + return new IoTDBSourceEnumerator(enumContext, options, rowDataType, filterQueries); + } + + @Override + public SplitEnumerator<IoTDBSourceSplit, IoTDBSourceEnumeratorState> restoreEnumerator( + SplitEnumeratorContext<IoTDBSourceSplit> enumContext, IoTDBSourceEnumeratorState checkpoint) { + return new IoTDBSourceEnumerator(enumContext, options, rowDataType, filterQueries, checkpoint); + } + + @Override + public SimpleVersionedSerializer<IoTDBSourceSplit> getSplitSerializer() { + return new IoTDBSourceSplitSerializer(); + } + + @Override + public SimpleVersionedSerializer<IoTDBSourceEnumeratorState> getEnumeratorCheckpointSerializer() { + return new IoTDBSourceEnumeratorStateSerializer(); + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceEnumerator.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceEnumerator.java new file mode 100644 index 0000000..00cb723 --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceEnumerator.java @@ -0,0 +1,170 @@ +/* + * 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.iotdb.relational.flink.source; + +import org.apache.iotdb.relational.flink.cfg.IoTDBRelationalOptions; +import org.apache.iotdb.relational.flink.source.enumerator.IoTDBSourceEnumeratorState; +import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplit; +import org.apache.iotdb.relational.flink.utils.IoTDBIdentifierUtils; + +import org.apache.flink.api.connector.source.SplitEnumerator; +import org.apache.flink.api.connector.source.SplitEnumeratorContext; +import org.apache.flink.table.types.DataType; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.UUID; + +/** Split enumerator for the bounded IoTDB table source. */ +public class IoTDBSourceEnumerator + implements SplitEnumerator<IoTDBSourceSplit, IoTDBSourceEnumeratorState> { + + private final SplitEnumeratorContext<IoTDBSourceSplit> context; + private final IoTDBRelationalOptions options; + private final DataType rowDataType; + private final List<String> filterQueries; + private final Deque<IoTDBSourceSplit> pendingSplits = new ArrayDeque<>(); + private final Deque<Integer> readersAwaitingSplit = new ArrayDeque<>(); + private final Set<Integer> assignedReaders = new HashSet<>(); + + private boolean allSplitsCreated; + private boolean closed; + + public IoTDBSourceEnumerator( + SplitEnumeratorContext<IoTDBSourceSplit> context, + IoTDBRelationalOptions options, + DataType rowDataType, + List<String> filterQueries) { + this(context, options, rowDataType, filterQueries, null); + } + + public IoTDBSourceEnumerator( + SplitEnumeratorContext<IoTDBSourceSplit> context, + IoTDBRelationalOptions options, + DataType rowDataType, + List<String> filterQueries, + @Nullable IoTDBSourceEnumeratorState checkpoint) { + this.context = context; + this.options = options; + this.rowDataType = rowDataType; + this.filterQueries = filterQueries == null ? new ArrayList<>() : new ArrayList<>(filterQueries); + if (checkpoint != null) { + pendingSplits.addAll(checkpoint.getRemainingSplits()); + allSplitsCreated = true; + } + } + + @Override + public void start() { + if (!allSplitsCreated) { + pendingSplits.add(createSingleSplit()); + allSplitsCreated = true; + } + assignPendingSplits(); + } + + @Override + public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) { + if (closed) { + return; + } + assignedReaders.remove(subtaskId); + readersAwaitingSplit.addLast(subtaskId); + assignPendingSplits(); + } + + @Override + public void addSplitsBack(List<IoTDBSourceSplit> splits, int subtaskId) { + assignedReaders.remove(subtaskId); + for (int i = splits.size() - 1; i >= 0; i--) { + pendingSplits.addFirst(splits.get(i)); + } + assignPendingSplits(); + } + + @Override + public void addReader(int subtaskId) { + // Splits are assigned when the reader explicitly requests one. + } + + @Override + public IoTDBSourceEnumeratorState snapshotState(long checkpointId) { + return new IoTDBSourceEnumeratorState(new ArrayList<>(pendingSplits)); + } + + @Override + public void close() throws IOException { + closed = true; + pendingSplits.clear(); + readersAwaitingSplit.clear(); + assignedReaders.clear(); + } + + private void assignPendingSplits() { + while (!pendingSplits.isEmpty() && !readersAwaitingSplit.isEmpty()) { + int subtaskId = readersAwaitingSplit.pollFirst(); + IoTDBSourceSplit split = pendingSplits.pollFirst(); + assignedReaders.add(subtaskId); + context.assignSplit(split, subtaskId); + } + + if (allSplitsCreated && pendingSplits.isEmpty() && assignedReaders.isEmpty()) { + while (!readersAwaitingSplit.isEmpty()) { + context.signalNoMoreSplits(readersAwaitingSplit.pollFirst()); + } + } + } + + private IoTDBSourceSplit createSingleSplit() { + List<String> fieldNames = DataType.getFieldNames(rowDataType); + StringBuilder columns = new StringBuilder(); + for (String fieldName : fieldNames) { + if (columns.length() > 0) { + columns.append(", "); + } + columns.append(quoteIdentifier(fieldName)); + } + if (columns.length() == 0) { + throw new IllegalArgumentException("IoTDB source requires at least one selected column."); + } + + String splitId = UUID.randomUUID().toString(); + StringBuilder sql = + new StringBuilder("SELECT ") + .append(columns) + .append(" FROM ") + .append(quoteIdentifier(options.getTable())); + if (!filterQueries.isEmpty()) { + sql.append(" WHERE ").append(String.join(" AND ", filterQueries)); + } + return new IoTDBSourceSplit(splitId, options.getDatabase(), options.getTable(), sql.toString()); + } + + private static String quoteIdentifier(String identifier) { + return IoTDBIdentifierUtils.quoteIdentifier(identifier); + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceReader.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceReader.java new file mode 100644 index 0000000..5bceafd --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceReader.java @@ -0,0 +1,229 @@ +/* + * 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.iotdb.relational.flink.source; + +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.relational.flink.cfg.IoTDBRelationalOptions; +import org.apache.iotdb.relational.flink.source.deserializer.IoTDBDeserializationSchema; +import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplit; +import org.apache.iotdb.session.TableSessionBuilder; + +import org.apache.flink.api.connector.source.ReaderOutput; +import org.apache.flink.api.connector.source.SourceReader; +import org.apache.flink.api.connector.source.SourceReaderContext; +import org.apache.flink.core.io.InputStatus; +import org.apache.flink.table.types.DataType; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +/** Source reader for the bounded IoTDB table source. */ +public class IoTDBSourceReader<OUT> implements SourceReader<OUT, IoTDBSourceSplit> { + + private final SourceReaderContext context; + private final IoTDBRelationalOptions options; + private final DataType rowDataType; + private final IoTDBDeserializationSchema<OUT> deserializer; + private final Deque<IoTDBSourceSplit> pendingSplits = new ArrayDeque<>(); + + private IoTDBSourceSplit currentSplit; + private ITableSession session; + private SessionDataSet dataSet; + private SessionDataSet.DataIterator iterator; + private boolean noMoreSplits; + private boolean closed; + private CompletableFuture<Void> availability = CompletableFuture.completedFuture(null); + + public IoTDBSourceReader( + SourceReaderContext context, + IoTDBRelationalOptions options, + DataType rowDataType, + IoTDBDeserializationSchema<OUT> deserializer) { + this.context = context; + this.options = options; + this.rowDataType = rowDataType; + this.deserializer = deserializer; + } + + @Override + public void start() { + if (pendingSplits.isEmpty() && currentSplit == null && !noMoreSplits) { + markUnavailable(); + context.sendSplitRequest(); + } + } + + @Override + public InputStatus pollNext(ReaderOutput<OUT> output) throws Exception { + if (closed) { + return InputStatus.END_OF_INPUT; + } + + while (true) { + if (iterator == null && !openNextSplit()) { + if (noMoreSplits) { + return InputStatus.END_OF_INPUT; + } + markUnavailable(); + context.sendSplitRequest(); + return InputStatus.NOTHING_AVAILABLE; + } + + if (iterator.next()) { + OUT record = deserializer.deserialize(iterator); + if (record != null) { + output.collect(record); + return InputStatus.MORE_AVAILABLE; + } + continue; + } + + finishCurrentSplit(); + } + } + + @Override + public List<IoTDBSourceSplit> snapshotState(long checkpointId) { + List<IoTDBSourceSplit> splits = new ArrayList<>(pendingSplits.size() + 1); + if (currentSplit != null) { + splits.add(currentSplit); + } + splits.addAll(pendingSplits); + return splits; + } + + @Override + public CompletableFuture<Void> isAvailable() { + return availability; + } + + @Override + public void addSplits(List<IoTDBSourceSplit> splits) { + if (closed || splits == null || splits.isEmpty()) { + return; + } + pendingSplits.addAll(splits); + markAvailable(); + } + + @Override + public void notifyNoMoreSplits() { + noMoreSplits = true; + markAvailable(); + } + + @Override + public void close() throws Exception { + if (closed) { + return; + } + closed = true; + closeCurrentSplit(); + pendingSplits.clear(); + markAvailable(); + } + + private boolean openNextSplit() throws Exception { + currentSplit = pendingSplits.pollFirst(); + if (currentSplit == null) { + return false; + } + + TableSessionBuilder builder = + new TableSessionBuilder() + .nodeUrls(options.getNodeUrls()) + .username(options.getUsername()) + .password(options.getPassword()); + if (currentSplit.getDatabase() != null) { + builder.database(currentSplit.getDatabase()); + } + try { + session = builder.build(); + dataSet = session.executeQueryStatement(currentSplit.getSql()); + iterator = dataSet.iterator(); + return true; + } catch (Exception e) { + try { + closeCurrentSplit(); + } catch (Exception closeException) { + e.addSuppressed(closeException); + } + currentSplit = null; + throw e; + } + } + + private void finishCurrentSplit() throws Exception { + closeCurrentSplit(); + currentSplit = null; + markAvailable(); + } + + private void closeCurrentSplit() throws Exception { + Exception failure = null; + try { + if (dataSet != null) { + dataSet.close(); + } + } catch (Exception e) { + failure = e; + } finally { + dataSet = null; + iterator = null; + } + + try { + if (session != null) { + session.close(); + } + } catch (Exception e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } finally { + session = null; + } + + if (failure != null) { + throw failure; + } + } + + private synchronized void markUnavailable() { + if (!availability.isDone()) { + return; + } + availability = new CompletableFuture<>(); + } + + private synchronized void markAvailable() { + CompletableFuture<Void> current = availability; + if (!current.isDone()) { + current.complete(null); + } + availability = CompletableFuture.completedFuture(null); + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableFactory.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableFactory.java index 5b7bfd9..a7a865f 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableFactory.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableFactory.java @@ -32,8 +32,8 @@ import org.apache.flink.table.factories.FactoryUtil; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; -import java.util.List; import java.util.HashSet; +import java.util.List; import java.util.Set; /** @@ -95,7 +95,8 @@ public class IoTDBRelationalDynamicTableFactory .withTable(config.get(IoTDBRelationalOptions.TABLE)) .withTimeColumn(config.get(IoTDBRelationalOptions.TIME_COLUMN)) .withTagColumns(parseColumnNames(config.get(IoTDBRelationalOptions.TAG_COLUMNS))) - .withAttributeColumns(parseColumnNames(config.get(IoTDBRelationalOptions.ATTRIBUTE_COLUMNS))) + .withAttributeColumns( + parseColumnNames(config.get(IoTDBRelationalOptions.ATTRIBUTE_COLUMNS))) .build(); } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSink.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSink.java index 82a3d59..6f55e63 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSink.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSink.java @@ -53,8 +53,7 @@ public class IoTDBRelationalDynamicTableSink implements DynamicTableSink { @Override public SinkRuntimeProvider getSinkRuntimeProvider(Context context) { - IoTDBTabletSerializer<RowData> serializer = - new RowDataIoTDBTabletSerializer(options, schema); + IoTDBTabletSerializer<RowData> serializer = new RowDataIoTDBTabletSerializer(options, schema); return SinkV2Provider.of(new IoTDBSink<>(options, serializer)); } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java index 59bda40..a81291d 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java @@ -20,48 +20,94 @@ package org.apache.iotdb.relational.flink.table; import org.apache.iotdb.relational.flink.cfg.IoTDBRelationalOptions; +import org.apache.iotdb.relational.flink.source.IoTDBSource; +import org.apache.iotdb.relational.flink.source.deserializer.RowDataDeserializationSchema; +import org.apache.iotdb.relational.flink.source.pushdown.IoTDBExpressionVisitor; import org.apache.flink.table.catalog.ResolvedSchema; import org.apache.flink.table.connector.ChangelogMode; import org.apache.flink.table.connector.source.DynamicTableSource; -import org.apache.flink.table.connector.source.LookupTableSource; import org.apache.flink.table.connector.source.ScanTableSource; +import org.apache.flink.table.connector.source.SourceProvider; +import org.apache.flink.table.connector.source.abilities.SupportsFilterPushDown; +import org.apache.flink.table.connector.source.abilities.SupportsLimitPushDown; +import org.apache.flink.table.expressions.ResolvedExpression; +import org.apache.flink.table.types.DataType; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; /** - * Dynamic table source of the IoTDB relational (table model) Flink connector, covering scan and - * lookup reading. + * Dynamic table source of the IoTDB relational (table model) Flink connector. * - * <p>Mirrors the structure of Doris' {@code DorisDynamicTableSource}. TODO: implement the scan and - * lookup runtime providers. + * <p>Only scan reads are exposed for now. Projection pushdown and lookup reads are intentionally + * disabled until their runtime behavior is implemented. */ -public class IoTDBRelationalDynamicTableSource implements ScanTableSource, LookupTableSource { +public class IoTDBRelationalDynamicTableSource + implements ScanTableSource, SupportsFilterPushDown, SupportsLimitPushDown { private final IoTDBRelationalOptions options; private final ResolvedSchema schema; + private DataType physicalRowDataType; + private final List<String> resolvedFilterQueries = new ArrayList<>(); + private long limit = -1L; public IoTDBRelationalDynamicTableSource(IoTDBRelationalOptions options, ResolvedSchema schema) { this.options = options; this.schema = schema; + this.physicalRowDataType = schema.toPhysicalRowDataType(); } @Override public ChangelogMode getChangelogMode() { - throw new UnsupportedOperationException("Not implemented yet."); + return ChangelogMode.insertOnly(); } @Override public ScanRuntimeProvider getScanRuntimeProvider(ScanContext scanContext) { - throw new UnsupportedOperationException("Not implemented yet."); + return SourceProvider.of( + new IoTDBSource<>( + options, + physicalRowDataType, + new RowDataDeserializationSchema(physicalRowDataType), + resolvedFilterQueries)); + } + + @Override + public Result applyFilters(List<ResolvedExpression> filters) { + if (filters == null || filters.isEmpty()) { + return Result.of(Collections.emptyList(), Collections.emptyList()); + } + + List<ResolvedExpression> acceptedFilters = new ArrayList<>(); + List<ResolvedExpression> remainingFilters = new ArrayList<>(); + IoTDBExpressionVisitor expressionVisitor = new IoTDBExpressionVisitor(); + for (ResolvedExpression filter : filters) { + String filterQuery = filter.accept(expressionVisitor); + if (filterQuery == null || filterQuery.trim().isEmpty()) { + remainingFilters.add(filter); + } else { + acceptedFilters.add(filter); + resolvedFilterQueries.add(filterQuery); + } + } + return Result.of(acceptedFilters, remainingFilters); } @Override - public LookupRuntimeProvider getLookupRuntimeProvider(LookupContext lookupContext) { - throw new UnsupportedOperationException("Not implemented yet."); + public void applyLimit(long limit) { + // TODO: push this limit into the single-split SQL when the optimization is enabled. + this.limit = limit; } @Override public DynamicTableSource copy() { - return new IoTDBRelationalDynamicTableSource(options, schema); + IoTDBRelationalDynamicTableSource copy = new IoTDBRelationalDynamicTableSource(options, schema); + copy.physicalRowDataType = physicalRowDataType; + copy.resolvedFilterQueries.addAll(resolvedFilterQueries); + copy.limit = limit; + return copy; } @Override diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/test/java/org/apache/iotdb/relational/flink/IoTDBRelationalLocalQueryManual.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/test/java/org/apache/iotdb/relational/flink/IoTDBRelationalLocalQueryManual.java new file mode 100644 index 0000000..6423d40 --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/test/java/org/apache/iotdb/relational/flink/IoTDBRelationalLocalQueryManual.java @@ -0,0 +1,95 @@ +/* + * 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.iotdb.relational.flink; + +import org.apache.flink.table.api.EnvironmentSettings; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; +import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; + +/** + * Temporary manual verification class. This is intentionally kept in test sources and should not be + * committed as a production test. + * + * <p>Run with system properties such as: + * + * <pre> + * -Diotdb.nodeUrls=127.0.0.1:6667 + * -Diotdb.user=root + * -Diotdb.password=root + * -Diotdb.database=test + * -Diotdb.table=sensor + * </pre> + */ +public class IoTDBRelationalLocalQueryManual { + + public static void main(String[] args) throws Exception { + String nodeUrls = System.getProperty("iotdb.nodeUrls", "127.0.0.1:6667"); + String user = System.getProperty("iotdb.user", "root"); + String password = System.getProperty("iotdb.password", "root"); + String database = System.getProperty("iotdb.database", "test"); + String table = System.getProperty("iotdb.table", "sensor"); + + TableEnvironment tableEnvironment = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + + tableEnvironment.executeSql("DROP TABLE IF EXISTS iotdb_source"); + + // Replace these columns with the actual columns and types in the local IoTDB table. + String ddl = + String.format( + "CREATE TABLE iotdb_source (\n" + + " `time` TIMESTAMP(3),\n" + + " `device_id` STRING,\n" + + " `temperature` DOUBLE\n" + + ") WITH (\n" + + " 'connector' = 'iotdb-relational',\n" + + " 'nodeUrls' = '%s',\n" + + " 'user' = '%s',\n" + + " 'password' = '%s',\n" + + " 'database' = '%s',\n" + + " 'table' = '%s'\n" + + ")", + nodeUrls, user, password, database, table); + + System.out.println("DDL:\n" + ddl); + tableEnvironment.executeSql(ddl); + + run(tableEnvironment, "SELECT * FROM iotdb_source LIMIT 5"); + run( + tableEnvironment, + "SELECT device_id, temperature FROM iotdb_source WHERE temperature > 0 LIMIT 5"); + run( + tableEnvironment, + "SELECT device_id, temperature FROM iotdb_source " + "WHERE temperature + 1 > 0 LIMIT 5"); + } + + private static void run(TableEnvironment tableEnvironment, String sql) throws Exception { + System.out.println("\n=== EXECUTE ==="); + System.out.println(sql); + + TableResult result = tableEnvironment.executeSql(sql); + try (CloseableIterator<Row> iterator = result.collect()) { + while (iterator.hasNext()) { + System.out.println(iterator.next()); + } + } + } +}
