Make GenericRecord and ClickHouseBinaryFormatReader easy to iterate

ClickHouse · Dec 6, 2024 · e2242b6 · e2242b6
1 parent 5e304cb
commit e2242b6
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 11 deletions.
diff --git a/...v2/src/main/java/com/clickhouse/client/api/data_formats/ClickHouseBinaryFormatReader.java b/...v2/src/main/java/com/clickhouse/client/api/data_formats/ClickHouseBinaryFormatReader.java
@@ -18,6 +18,7 @@
 import java.time.ZonedDateTime;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
 public interface ClickHouseBinaryFormatReader extends AutoCloseable {
@@ -537,4 +538,22 @@ public interface ClickHouseBinaryFormatReader extends AutoCloseable {
     ClickHouseBitmap getClickHouseBitmap(String colName);
 
     ClickHouseBitmap getClickHouseBitmap(int index);
+
+    /**
+     * Get number of columns.
+     * @return
+     */
+    int getColumnCount();
+
+    /**
+     * Get column names.
+     * @return
+     */
+    Set<String> getColumnNames();
+
+    /**
+     * Get current row.
+     * @return
+     */
+    Map<String, Object> current();
 }
diff --git a/...main/java/com/clickhouse/client/api/data_formats/internal/AbstractBinaryFormatReader.java b/...main/java/com/clickhouse/client/api/data_formats/internal/AbstractBinaryFormatReader.java
@@ -32,12 +32,7 @@
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.time.temporal.ChronoUnit;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.TimeZone;
-import java.util.UUID;
+import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
@@ -737,4 +732,19 @@ public ClickHouseBitmap getClickHouseBitmap(int index) {
     public void close() throws Exception {
         input.close();
     }
+
+    @Override
+    public int getColumnCount() {
+        return currentRecord.size();
+    }
+
+    @Override
+    public Set<String> getColumnNames() {
+        return currentRecord.keySet();
+    }
+
+    @Override
+    public Map<String, Object> current() {
+        return currentRecord;
+    }
 }
diff --git a/...c/main/java/com/clickhouse/client/api/data_formats/internal/BinaryReaderBackedRecord.java b/...c/main/java/com/clickhouse/client/api/data_formats/internal/BinaryReaderBackedRecord.java
@@ -18,6 +18,8 @@
 import java.time.LocalDateTime;
 import java.time.ZonedDateTime;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
 public class BinaryReaderBackedRecord implements GenericRecord {
@@ -377,4 +379,19 @@ public ClickHouseBitmap getClickHouseBitmap(String colName) {
     public ClickHouseBitmap getClickHouseBitmap(int index) {
         return reader.readValue(index);
     }
+
+    @Override
+    public int getSize() {
+        return reader.getColumnCount();
+    }
+
+    @Override
+    public Set<String> getNames() {
+        return reader.getColumnNames();
+    }
+
+    @Override
+    public Map<String, Object> getAll() {
+        return reader.current();
+    }
 }
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/MapBackedRecord.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/MapBackedRecord.java
@@ -5,7 +5,6 @@
 import com.clickhouse.client.api.query.GenericRecord;
 import com.clickhouse.client.api.query.NullValueException;
 import com.clickhouse.data.ClickHouseColumn;
-import com.clickhouse.data.value.ClickHouseArrayValue;
 import com.clickhouse.data.value.ClickHouseBitmap;
 import com.clickhouse.data.value.ClickHouseGeoMultiPolygonValue;
 import com.clickhouse.data.value.ClickHouseGeoPointValue;
@@ -23,10 +22,7 @@
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.time.temporal.ChronoUnit;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.UUID;
+import java.util.*;
 import java.util.function.Function;
 
 public class MapBackedRecord implements GenericRecord {
@@ -522,4 +518,19 @@ public Object getObject(String colName) {
     public Object getObject(int index) {
         return readValue(index);
     }
+
+    @Override
+    public int getSize() {
+        return record.size();
+    }
+
+    @Override
+    public Set<String> getNames() {
+        return record.keySet();
+    }
+
+    @Override
+    public Map<String, Object> getAll() {
+        return Map.of();
+    }
 }
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/GenericRecord.java b/client-v2/src/main/java/com/clickhouse/client/api/query/GenericRecord.java
@@ -16,6 +16,8 @@
 import java.time.LocalDateTime;
 import java.time.ZonedDateTime;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
 public interface GenericRecord {
@@ -500,4 +502,22 @@ public interface GenericRecord {
     ClickHouseBitmap getClickHouseBitmap(String colName);
 
     ClickHouseBitmap getClickHouseBitmap(int index);
+
+    /**
+     * Returns the number of columns.
+     * @return number of columns
+     */
+    int getSize();
+
+    /**
+     * Returns the names of the columns.
+     * @return names of the columns
+     */
+    Set<String> getNames();
+
+    /**
+     * Returns all columns with values.
+     * @return all columns values
+     */
+    Map<String, Object> getAll();
 }