Added Attribute to for ArrowSchema

googleapis · Feb 8, 2022 · 3297e4d · 3297e4d
1 parent c1743c9
commit 3297e4d
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryResultSetImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryResultSetImpl.java
@@ -16,6 +16,7 @@
 
 package com.google.cloud.bigquery;
 
+import com.google.cloud.bigquery.storage.v1.ArrowSchema;
 import java.math.BigDecimal;
 import java.sql.Date;
 import java.sql.ResultSet;
@@ -28,19 +29,25 @@
 // TODO: This implementation deals with the JSON response. We can have respective implementations
 public class BigQueryResultSetImpl<T> implements BigQueryResultSet<T> {
 
-  private final Schema schema;
+  private Schema schema;
+  private ArrowSchema arrowSchema;
   private final long totalRows;
   private final BlockingQueue<T> buffer;
   private T cursor;
   private final ResultSetWrapper underlyingResultSet;
   private final BigQueryResultSetStats bigQueryResultSetStats;
 
   public BigQueryResultSetImpl(
-      Schema schema,
+      Object schema,
       long totalRows,
       BlockingQueue<T> buffer,
       BigQueryResultSetStats bigQueryResultSetStats) {
-    this.schema = schema;
+    if (schema instanceof Schema) {
+      this.schema = (Schema) schema;
+    } else if (schema instanceof ArrowSchema) {
+      this.arrowSchema = (ArrowSchema) schema;
+    }
+
     this.totalRows = totalRows;
     this.buffer = buffer;
     this.underlyingResultSet = new ResultSetWrapper();
@@ -52,6 +59,8 @@ public Schema getSchema() {
     return schema;
   }
 
+  // TODO(prasmish): Implement a method to return ArrowSchema
+
   @Override
   public long getTotalRows() {
     return totalRows;