diff --git a/.gitignore b/.gitignore index d5188c1..b84ea99 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,18 @@ hs_err_pid* replay_pid* /cmake-build-debug/ *.lst + +############################## +## Maven +############################## +target/ + +############################## +## IntelliJ +############################## +out/ +.idea/ +.idea_modules/ +*.iml +*.ipr +*.iws \ No newline at end of file diff --git a/pom.xml b/pom.xml index 36203c4..83e5999 100644 --- a/pom.xml +++ b/pom.xml @@ -12,6 +12,7 @@ 22 22 UTF-8 + 18.1.0 @@ -23,6 +24,23 @@ 5.8.1 test + + org.apache.arrow + arrow-vector + ${arrow.version} + + + org.apache.arrow + arrow-memory-unsafe + ${arrow.version} + runtime + + + org.apache.arrow + arrow-compression + ${arrow.version} + runtime + @@ -38,6 +56,7 @@ **/*Test.java + --add-opens=java.base/java.nio=ALL-UNNAMED diff --git a/src/jni/chdb_jni.cpp b/src/jni/chdb_jni.cpp index c10bb53..bcd5726 100644 --- a/src/jni/chdb_jni.cpp +++ b/src/jni/chdb_jni.cpp @@ -48,19 +48,21 @@ local_result_v2 * queryToBuffer( } -JNIEXPORT jobject JNICALL Java_org_chdb_jdbc_ChdbJniUtil_executeQuery(JNIEnv *env, jclass clazz, jstring query) { +JNIEXPORT jobject JNICALL Java_org_chdb_jdbc_ChdbJniUtil_executeQuery(JNIEnv *env, jclass clazz, jstring query, jstring format) { // 1. Convert Java String to C++ string std::cout << "call func: ChdbJniUtil_executeQuery!" << std::endl; const char *queryStr = env->GetStringUTFChars(query, nullptr); + const char *formatStr = env->GetStringUTFChars(format, nullptr); + if (queryStr == nullptr) { std::cerr << "Error: Failed to convert Java string to C++ string" << std::endl; return nullptr; } // 2. Call the native query function - local_result_v2 *result = queryToBuffer(queryStr); + local_result_v2 *result = queryToBuffer(queryStr, formatStr); // 3. Release the Java string resources env->ReleaseStringUTFChars(query, queryStr); diff --git a/src/jni/chdb_jni.h b/src/jni/chdb_jni.h index c2ff837..5fb1126 100644 --- a/src/jni/chdb_jni.h +++ b/src/jni/chdb_jni.h @@ -7,8 +7,7 @@ extern "C" { #endif -JNIEXPORT jobject JNICALL Java_org_chdb_jdbc_ChdbJniUtil_executeQuery(JNIEnv *, jclass, jstring); -//JNIEXPORT jstring JNICALL Java_org_chdb_jdbc_ChdbJniUtil_executeQuery(JNIEnv *, jclass, jstring); +JNIEXPORT jobject JNICALL Java_org_chdb_jdbc_ChdbJniUtil_executeQuery(JNIEnv *, jclass, jstring, jstring); #ifdef __cplusplus } diff --git a/src/main/java/org/chdb/jdbc/ChdbJniUtil.java b/src/main/java/org/chdb/jdbc/ChdbJniUtil.java index b0b2211..2c6c7e3 100644 --- a/src/main/java/org/chdb/jdbc/ChdbJniUtil.java +++ b/src/main/java/org/chdb/jdbc/ChdbJniUtil.java @@ -5,8 +5,11 @@ public class ChdbJniUtil { System.loadLibrary("chdbjni"); } - public static native LocalResultV2 executeQuery(String query); -// public static native String executeQuery(String query); + public static LocalResultV2 executeQuery(String query) { + return executeQuery(query, "CSV"); + } + + public static native LocalResultV2 executeQuery(String query, String format); // public static void main(String[] args) { // String query = "SELECT 1"; diff --git a/src/main/java/org/chdb/jdbc/ChdbResultSet.java b/src/main/java/org/chdb/jdbc/ChdbResultSet.java index 32057b5..56d0cb9 100644 --- a/src/main/java/org/chdb/jdbc/ChdbResultSet.java +++ b/src/main/java/org/chdb/jdbc/ChdbResultSet.java @@ -1,33 +1,38 @@ package org.chdb.jdbc; -import java.io.*; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.chdb.jdbc.memory.ArrowMemoryManger; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.Reader; import java.math.BigDecimal; import java.net.URL; import java.nio.ByteBuffer; -import java.nio.charset.StandardCharsets; import java.sql.*; -import java.util.ArrayList; import java.util.Calendar; -import java.util.List; import java.util.Map; public class ChdbResultSet implements ResultSet { - private LocalResultV2 result; private int cursor = -1; - private List records; + private int loadedRows = 0; + private ArrowStreamReader arrowStreamReader; + private VectorSchemaRoot batch; + private int batchCursor; public ChdbResultSet(LocalResultV2 result) throws IOException { - this.result = result; - this.records = new ArrayList<>(); - parseData(result.getBuf()); + parseData(result.getBuf()); } private void parseData(ByteBuffer buffer) throws IOException { - BufferedReader reader = new BufferedReader(new InputStreamReader(new ByteArrayInputStream(toByteArray(buffer)), StandardCharsets.UTF_8)); - String line; - while ((line = reader.readLine()) != null) { - records.add(line); - } + ByteArrayInputStream inputStream = new ByteArrayInputStream(toByteArray(buffer)); + this.arrowStreamReader = new ArrowStreamReader(inputStream, ArrowMemoryManger.ROOT_ALLOCATOR); } private byte[] toByteArray(ByteBuffer buffer) { @@ -38,24 +43,54 @@ private byte[] toByteArray(ByteBuffer buffer) { @Override public boolean next() throws SQLException { - if (cursor < records.size() - 1) { + if (cursor < loadedRows - 1) { cursor++; + batchCursor++; return true; } - return false; + try { + if (arrowStreamReader.loadNextBatch()) { + AutoCloseables.close(batch); + batch = arrowStreamReader.getVectorSchemaRoot(); + loadedRows += batch.getRowCount(); + cursor++; + batchCursor = 0; + return true; + } else { + return false; + } + } catch (Exception e) { + throw new SQLException(e); + } } - private String getValue(int columnIndex) throws SQLException { - if (cursor < 0 || cursor >= records.size()) { - throw new SQLException("Cursor out of bounds"); + /** + * Checks to see whether the given index is a valid column number and throws + * an SQLException if it is not. The index is out of bounds + * if it is less than 1 or greater than the number of + * columns in this rowset. + *

+ * This method is called internally by the getXXX and + * updateXXX methods. + * + * @param idx the number of a column, must be between 1 + * and the number of rows in this rowset + * @throws SQLException if the given index is out of bounds + */ + private void checkColumnIndex(int idx) throws SQLException { + if (idx < 1 || idx > getColumnCount()) { + throw new SQLException("Column index " + idx + " is out of bound[1, " + getColumnCount() +"]"); } + } - return records.get(cursor); + private int getColumnCount() { + return batch.getSchema().getFields().size(); } @Override public String getString(int columnIndex) throws SQLException { - return getValue(columnIndex); + checkColumnIndex(columnIndex); + return batch.getVector(columnIndex - 1).getObject(batchCursor).toString(); } @Override @@ -75,12 +110,22 @@ public short getShort(int i) throws SQLException { @Override public int getInt(int columnIndex) throws SQLException { - return Integer.parseInt(getValue(columnIndex)); + checkColumnIndex(columnIndex); + FieldVector vector = batch.getVector(columnIndex - 1); + if (vector instanceof IntVector) { + return ((IntVector) vector).get(batchCursor); + } + throw new SQLException("Column [" + vector.getField() + "] is not int"); } @Override - public long getLong(int i) throws SQLException { - throw new SQLException("This method has not been implemented yet."); + public long getLong(int columnIndex) throws SQLException { + checkColumnIndex(columnIndex); + FieldVector vector = batch.getVector(columnIndex); + if (vector instanceof BigIntVector) { + return ((BigIntVector) vector).get(batchCursor); + } + throw new SQLException("Column [" + vector.getField() + "] is not long"); } @Override @@ -135,7 +180,11 @@ public InputStream getBinaryStream(int i) throws SQLException { @Override public void close() throws SQLException { - // No resources to close + try { + AutoCloseables.close(batch, arrowStreamReader); + } catch (Exception e) { + throw new SQLException(e); + } } @Override diff --git a/src/main/java/org/chdb/jdbc/ChdbStatement.java b/src/main/java/org/chdb/jdbc/ChdbStatement.java index 9cff3a0..f0388b5 100644 --- a/src/main/java/org/chdb/jdbc/ChdbStatement.java +++ b/src/main/java/org/chdb/jdbc/ChdbStatement.java @@ -12,7 +12,7 @@ public ChdbStatement(ChdbConnection connection) { @Override public ResultSet executeQuery(String sql) throws SQLException { - LocalResultV2 result = ChdbJniUtil.executeQuery(sql); + LocalResultV2 result = ChdbJniUtil.executeQuery(sql, "ArrowStream"); System.out.println("sql: " + sql); try { return new ChdbResultSet(result); diff --git a/src/main/java/org/chdb/jdbc/memory/ArrowMemoryManger.java b/src/main/java/org/chdb/jdbc/memory/ArrowMemoryManger.java new file mode 100644 index 0000000..43af7b3 --- /dev/null +++ b/src/main/java/org/chdb/jdbc/memory/ArrowMemoryManger.java @@ -0,0 +1,7 @@ +package org.chdb.jdbc.memory; + +import org.apache.arrow.memory.RootAllocator; + +public class ArrowMemoryManger { + public static RootAllocator ROOT_ALLOCATOR = new RootAllocator(); +} diff --git a/src/test/java/ChdbJdbcTest.java b/src/test/java/ChdbJdbcTest.java index 48f3f86..658d985 100644 --- a/src/test/java/ChdbJdbcTest.java +++ b/src/test/java/ChdbJdbcTest.java @@ -1,4 +1,4 @@ -import org.junit.jupiter.api.Test; +import org.chdb.jdbc.memory.ArrowMemoryManger; import java.sql.Connection; import java.sql.DriverManager; @@ -7,8 +7,10 @@ import static org.junit.jupiter.api.Assertions.*; -public class ChdbJdbcTest { +public class ChdbJdbcTest { + // TODO failed to run by junit, it will crash + // Note: adding --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED public static void main(String[] args) throws Exception { // Load the JDBC driver Class.forName("org.chdb.jdbc.ChdbDriver"); @@ -31,5 +33,6 @@ public static void main(String[] args) throws Exception { rs.close(); conn.close(); + assertEquals(0, ArrowMemoryManger.ROOT_ALLOCATOR.getAllocatedMemory()); } }