optimize memory allocation of validity buffers

TileDB-Inc · Jan 7, 2023 · fcf028c · fcf028c
1 parent 4539473
commit fcf028c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/src/main/java/io/trino/plugin/tiledb/TileDBMetadata.java b/src/main/java/io/trino/plugin/tiledb/TileDBMetadata.java
@@ -69,7 +69,6 @@
 import static io.trino.plugin.tiledb.TileDBColumnProperties.getExtent;
 import static io.trino.plugin.tiledb.TileDBColumnProperties.getFilterList;
 import static io.trino.plugin.tiledb.TileDBColumnProperties.getLowerBound;
-import static io.trino.plugin.tiledb.TileDBColumnProperties.getNullable;
 import static io.trino.plugin.tiledb.TileDBColumnProperties.getUpperBound;
 import static io.trino.plugin.tiledb.TileDBErrorCode.TILEDB_CREATE_TABLE_ERROR;
 import static io.trino.plugin.tiledb.TileDBModule.tileDBTypeFromTrinoType;
@@ -438,7 +437,7 @@ public TileDBOutputTableHandle beginCreateArray(ConnectorSession session, Connec
                     if (filterPairs.isPresent()) {
                         attribute.setFilterList(Util.createTileDBFilterList(localCtx, filterPairs.get()));
                     }
-                    attribute.setNullable(getNullable(columnProperties));
+                    attribute.setNullable(column.isNullable());
                     arraySchema.addAttribute(attribute);
                 }
 

diff --git a/src/main/java/io/trino/plugin/tiledb/TileDBPageSink.java b/src/main/java/io/trino/plugin/tiledb/TileDBPageSink.java
@@ -176,7 +176,7 @@ private void resetBuffers(Map<String, Pair<NativeArray, NativeArray>> buffers) t
         buffers.clear();
         query.resetBuffers();
 
-        validityMaps = new short[columnHandles.size()][maxBufferSize];
+        validityMaps = new short[columnHandles.size()][];
         // Loop through each column
         for (int channel = 0; channel < columnHandles.size(); channel++) {
             // Datatype
@@ -195,9 +195,16 @@ private void resetBuffers(Map<String, Pair<NativeArray, NativeArray>> buffers) t
             if (isVariableLength) {
                 offsets = new NativeArray(ctx, maxBufferSize, Datatype.TILEDB_UINT64);
             }
-            if (!columnHandle.getIsDimension()) {
-                Arrays.fill(validityMaps[channel], (short) 1); // all valid
+
+            // allocate validity buffers for nullable attributes only.
+            if (array.getSchema().hasAttribute(columnName)) {
+                Attribute attribute = array.getSchema().getAttribute(columnHandle.getColumnName());
+                if (attribute.getNullable()) {
+                    validityMaps[channel] = new short[maxBufferSize / type.getNativeSize()]; //allocate for the max amount of values
+                    Arrays.fill(validityMaps[channel], (short) 1); // all valid
+                }
             }
+
             buffers.put(columnName, new Pair<>(offsets, values));
         }
     }