Add milvusReader and milvusWriter to support read and write data for milvus

milvuswriter/pom.xml

@@ -0,0 +1,109 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <groupId>com.alibaba.datax</groupId>
+        <artifactId>datax-all</artifactId>
+        <version>0.0.1-SNAPSHOT</version>
+    </parent>
+
+    <artifactId>milvuswriter</artifactId>
+
+    <properties>
+        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+        <kotlin.code.style>official</kotlin.code.style>
+        <kotlin.compiler.jvmTarget>1.8</kotlin.compiler.jvmTarget>
+    </properties>
+    <dependencyManagement>
+        <dependencies>
+            <dependency>
+                <artifactId>guava</artifactId>
+                <groupId>com.google.guava</groupId>
+                <version>32.0.1-jre</version>
+            </dependency>
+        </dependencies>
+    </dependencyManagement>
+
+    <dependencies>
+        <dependency>
+            <groupId>io.milvus</groupId>
+            <artifactId>milvus-sdk-java</artifactId>
+            <version>2.4.8</version>
+        </dependency>
+        <dependency>
+            <groupId>org.jetbrains.kotlin</groupId>
+            <artifactId>kotlin-test-junit5</artifactId>
+            <version>2.0.0</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter</artifactId>
+            <version>5.10.0</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.jetbrains.kotlin</groupId>
+            <artifactId>kotlin-stdlib</artifactId>
+            <version>2.0.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.alibaba.datax</groupId>
+            <artifactId>datax-common</artifactId>
+            <version>0.0.1-SNAPSHOT</version>
+            <scope>compile</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.projectlombok</groupId>
+            <artifactId>lombok</artifactId>
+            <version>1.18.30</version>
+            <scope>provided</scope>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <resources>
+            <!--将resource目录也输出到target-->
+            <resource>
+                <directory>src/main/resources</directory>
+                <includes>
+                    <include>**/*.*</include>
+                </includes>
+                <filtering>true</filtering>
+            </resource>
+        </resources>
+        <plugins>
+            <!-- compiler plugin -->
+            <plugin>
+                <artifactId>maven-compiler-plugin</artifactId>
+                <configuration>
+                    <source>${jdk-version}</source>
+                    <target>${jdk-version}</target>
+                    <encoding>${project-sourceEncoding}</encoding>
+                </configuration>
+            </plugin>
+            <!-- assembly plugin -->
+            <plugin>
+                <artifactId>maven-assembly-plugin</artifactId>
+                <configuration>
+                    <descriptors>
+                        <descriptor>src/main/assembly/package.xml</descriptor>
+                    </descriptors>
+                    <finalName>datax</finalName>
+                </configuration>
+                <executions>
+                    <execution>
+                        <id>dwzip</id>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>single</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+    </build>
+
+</project>

milvuswriter/src/main/assembly/package.xml

@@ -0,0 +1,36 @@
+<assembly
+        xmlns="http://maven.apache.org/plugins/maven-assembly-plugin/assembly/1.1.0"
+        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+        xsi:schemaLocation="http://maven.apache.org/plugins/maven-assembly-plugin/assembly/1.1.0 http://maven.apache.org/xsd/assembly-1.1.0.xsd">
+    <id></id>
+
+    <formats>
+        <format>dir</format>
+    </formats>
+    <includeBaseDirectory>false</includeBaseDirectory>
+    <fileSets>
+        <fileSet>
+            <directory>src/main/resources</directory>
+            <includes>
+                <include>plugin.json</include>
+                <include>plugin_job_template.json</include>
+            </includes>
+            <outputDirectory>plugin/writer/milvuswriter</outputDirectory>
+        </fileSet>
+        <fileSet>
+            <directory>target/</directory>
+            <includes>
+                <include>milvuswriter-0.0.1-SNAPSHOT.jar</include>
+            </includes>
+            <outputDirectory>plugin/writer/milvuswriter</outputDirectory>
+        </fileSet>
+    </fileSets>
+
+    <dependencySets>
+        <dependencySet>
+            <useProjectArtifact>false</useProjectArtifact>
+            <outputDirectory>plugin/writer/milvuswriter/libs</outputDirectory>
+            <scope>runtime</scope>
+        </dependencySet>
+    </dependencySets>
+</assembly>

milvuswriter/src/main/java/com/alibaba/datax/plugin/writer/milvuswriter/BufferUtils.java

@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.alibaba.datax.plugin.writer.milvuswriter;
+
+import java.nio.Buffer;
+import java.nio.ByteBuffer;
+
+public class BufferUtils {
+
+    public static ByteBuffer toByteBuffer(Short[] shortArray) {
+        ByteBuffer byteBuffer = ByteBuffer.allocate(shortArray.length * 2);
+
+        for (Short value : shortArray) {
+            byteBuffer.putShort(value);
+        }
+
+        // Compatible compilation and running versions are not consistent
+        // Flip the buffer to prepare for reading
+        ((Buffer) byteBuffer).flip();
+
+        return byteBuffer;
+    }
+
+    public static Short[] toShortArray(ByteBuffer byteBuffer) {
+        Short[] shortArray = new Short[byteBuffer.capacity() / 2];
+
+        for (int i = 0; i < shortArray.length; i++) {
+            shortArray[i] = byteBuffer.getShort();
+        }
+
+        return shortArray;
+    }
+
+    public static ByteBuffer toByteBuffer(Float[] floatArray) {
+        ByteBuffer byteBuffer = ByteBuffer.allocate(floatArray.length * 4);
+
+        for (float value : floatArray) {
+            byteBuffer.putFloat(value);
+        }
+
+        ((Buffer) byteBuffer).flip();
+
+        return byteBuffer;
+    }
+
+    public static Float[] toFloatArray(ByteBuffer byteBuffer) {
+        Float[] floatArray = new Float[byteBuffer.capacity() / 4];
+
+        for (int i = 0; i < floatArray.length; i++) {
+            floatArray[i] = byteBuffer.getFloat();
+        }
+
+        return floatArray;
+    }
+
+    public static ByteBuffer toByteBuffer(Double[] doubleArray) {
+        ByteBuffer byteBuffer = ByteBuffer.allocate(doubleArray.length * 8);
+
+        for (double value : doubleArray) {
+            byteBuffer.putDouble(value);
+        }
+
+        ((Buffer) byteBuffer).flip();
+
+        return byteBuffer;
+    }
+
+    public static Double[] toDoubleArray(ByteBuffer byteBuffer) {
+        Double[] doubleArray = new Double[byteBuffer.capacity() / 8];
+
+        for (int i = 0; i < doubleArray.length; i++) {
+            doubleArray[i] = byteBuffer.getDouble();
+        }
+
+        return doubleArray;
+    }
+
+    public static ByteBuffer toByteBuffer(Integer[] intArray) {
+        ByteBuffer byteBuffer = ByteBuffer.allocate(intArray.length * 4);
+
+        for (int value : intArray) {
+            byteBuffer.putInt(value);
+        }
+
+        ((Buffer) byteBuffer).flip();
+
+        return byteBuffer;
+    }
+
+    public static Integer[] toIntArray(ByteBuffer byteBuffer) {
+        Integer[] intArray = new Integer[byteBuffer.capacity() / 4];
+
+        for (int i = 0; i < intArray.length; i++) {
+            intArray[i] = byteBuffer.getInt();
+        }
+
+        return intArray;
+    }
+}

milvuswriter/src/main/java/com/alibaba/datax/plugin/writer/milvuswriter/KeyConstant.java

@@ -0,0 +1,20 @@
+package com.alibaba.datax.plugin.writer.milvuswriter;
+
+public class KeyConstant {
+    public static final String URI = "uri";
+    public static final String TOKEN = "token";
+    public static final String DATABASE = "database";
+    public static final String COLLECTION = "collection";
+    public static final String AUTO_ID = "autoId";
+    public static final String ENABLE_DYNAMIC_SCHEMA = "enableDynamicSchema";
+    public static final String BATCH_SIZE = "batchSize";
+    public static final String COLUMN = "column";
+    public static final String COLUMN_TYPE = "type";
+    public static final String COLUMN_NAME = "name";
+    public static final String VECTOR_DIMENSION = "dimension";
+    public static final String IS_PRIMARY_KEY = "isPrimaryKey";
+//    "schemaCreateMode":"createWhenTableNotExit"/"Ignore"/"exception"
+    public static final String schemaCreateMode = "schemaCreateMode";
+    public static final String IS_PARTITION_KEY = "isPartitionKey";
+    public static final String MAX_LENGTH = "maxLength";
+}

milvuswriter/src/main/java/com/alibaba/datax/plugin/writer/milvuswriter/MilvusBufferWriter.java

@@ -0,0 +1,43 @@
+package com.alibaba.datax.plugin.writer.milvuswriter;
+
+import com.google.gson.JsonObject;
+import io.milvus.v2.client.MilvusClientV2;
+import io.milvus.v2.service.vector.request.UpsertReq;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@Slf4j
+public class MilvusBufferWriter {
+
+    private final MilvusClientV2 milvusClientV2;
+    private final String collection;
+    private final Integer batchSize;
+    private List<JsonObject> dataCache;
+
+    public MilvusBufferWriter(MilvusClientV2 milvusClientV2, String collection, Integer batchSize){
+        this.milvusClientV2 = milvusClientV2;
+        this.collection = collection;
+        this.batchSize = batchSize;
+        this.dataCache = new ArrayList<>();
+    }
+    public void write(JsonObject data){
+        dataCache.add(data);
+    }
+    public Boolean needCommit(){
+        return dataCache.size() >= batchSize;
+    }
+    public void commit(){
+        if(dataCache.isEmpty()){
+            log.info("dataCache is empty, skip commit");
+            return;
+        }
+        UpsertReq upsertReq = UpsertReq.builder()
+                .collectionName(collection)
+                .data(dataCache)
+                .build();
+        milvusClientV2.upsert(upsertReq);
+        dataCache = new ArrayList<>();
+    }
+}