perf(refactor): optimizing JdbcBasePersistenceImpl.listEntities

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java

@@ -66,6 +66,7 @@
 import org.apache.polaris.core.storage.PolarisStorageIntegration;
 import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider;
 import org.apache.polaris.core.storage.StorageLocation;
+import org.apache.polaris.persistence.relational.jdbc.models.EntityNameLookupRecordConverter;
 import org.apache.polaris.persistence.relational.jdbc.models.ModelEntity;
 import org.apache.polaris.persistence.relational.jdbc.models.ModelGrantRecord;
 import org.apache.polaris.persistence.relational.jdbc.models.ModelPolicyMappingRecord;
@@ -425,38 +426,13 @@ public List<PolarisChangeTrackingVersions> lookupEntityVersions(
         .collect(Collectors.toList());
   }
 
-  @Nonnull
-  @Override
-  public Page<EntityNameLookupRecord> listEntities(
-      @Nonnull PolarisCallContext callCtx,
-      long catalogId,
-      long parentId,
-      @Nonnull PolarisEntityType entityType,
-      @Nonnull PolarisEntitySubType entitySubType,
-      @Nonnull PageToken pageToken) {
-    // TODO: only fetch the properties required for creating an EntityNameLookupRecord
-    return loadEntities(
-        callCtx,
-        catalogId,
-        parentId,
-        entityType,
-        entitySubType,
-        entity -> true,
-        EntityNameLookupRecord::new,
-        pageToken);
-  }
-
-  @Nonnull
-  @Override
-  public <T> Page<T> loadEntities(
-      @Nonnull PolarisCallContext callCtx,
+  private PreparedQuery buildEntityQuery(
       long catalogId,
       long parentId,
-      @Nonnull PolarisEntityType entityType,
-      @Nonnull PolarisEntitySubType entitySubType,
-      @Nonnull Predicate<PolarisBaseEntity> entityFilter,
-      @Nonnull Function<PolarisBaseEntity, T> transformer,
-      @Nonnull PageToken pageToken) {
+      PolarisEntityType entityType,
+      PolarisEntitySubType entitySubType,
+      PageToken pageToken,
+      List<String> queryProjections) {
     Map<String, Object> whereEquals =
         Map.of(
             "catalog_id",
@@ -467,17 +443,15 @@ public <T> Page<T> loadEntities(
             entityType.getCode(),
             "realm_id",
             realmId);
-    Map<String, Object> whereGreater;
 
     if (entitySubType != PolarisEntitySubType.ANY_SUBTYPE) {
       Map<String, Object> updatedWhereEquals = new HashMap<>(whereEquals);
       updatedWhereEquals.put("sub_type_code", entitySubType.getCode());
       whereEquals = updatedWhereEquals;
     }
 
-    // Limit can't be pushed down, due to client side filtering
-    // absence of transaction.
     String orderByColumnName = null;
+    Map<String, Object> whereGreater;
     if (pageToken.paginationRequested()) {
       orderByColumnName = ModelEntity.ID_COLUMN;
       whereGreater =
@@ -491,14 +465,58 @@ public <T> Page<T> loadEntities(
       whereGreater = Map.of();
     }
 
+    return QueryGenerator.generateSelectQuery(
+        queryProjections, ModelEntity.TABLE_NAME, whereEquals, whereGreater, orderByColumnName);
+  }
+
+  @Nonnull
+  @Override
+  public Page<EntityNameLookupRecord> listEntities(
+      @Nonnull PolarisCallContext callCtx,
+      long catalogId,
+      long parentId,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull PolarisEntitySubType entitySubType,
+      @Nonnull PageToken pageToken) {
+    try {
+      PreparedQuery query =
+          buildEntityQuery(
+              catalogId,
+              parentId,
+              entityType,
+              entitySubType,
+              pageToken,
+              ModelEntity.ENTITY_LOOKUP_COLUMNS);
+      AtomicReference<Page<EntityNameLookupRecord>> results = new AtomicReference<>();
+      datasourceOperations.executeSelectOverStream(
+          query,
+          new EntityNameLookupRecordConverter(),
+          stream -> {
+            results.set(
+                Page.mapped(pageToken, stream, Function.identity(), EntityIdToken::fromEntity));
+          });
+      return results.get();
+    } catch (SQLException e) {
+      throw new RuntimeException(
+          String.format("Failed to retrieve polaris entities due to %s", e.getMessage()), e);
+    }
+  }
+
+  @Nonnull
+  @Override
+  public <T> Page<T> loadEntities(
+      @Nonnull PolarisCallContext callCtx,
+      long catalogId,
+      long parentId,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull PolarisEntitySubType entitySubType,
+      @Nonnull Predicate<PolarisBaseEntity> entityFilter,
+      @Nonnull Function<PolarisBaseEntity, T> transformer,
+      @Nonnull PageToken pageToken) {
     try {
       PreparedQuery query =
-          QueryGenerator.generateSelectQuery(
-              ModelEntity.ALL_COLUMNS,
-              ModelEntity.TABLE_NAME,
-              whereEquals,
-              whereGreater,
-              orderByColumnName);
+          buildEntityQuery(
+              catalogId, parentId, entityType, entitySubType, pageToken, ModelEntity.ALL_COLUMNS);
       AtomicReference<Page<T>> results = new AtomicReference<>();
       datasourceOperations.executeSelectOverStream(
           query,

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/models/EntityNameLookupRecordConverter.java

@@ -0,0 +1,46 @@
+/*
+ * 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 org.apache.polaris.persistence.relational.jdbc.models;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.Map;
+import org.apache.polaris.core.entity.EntityNameLookupRecord;
+import org.apache.polaris.persistence.relational.jdbc.DatabaseType;
+
+public class EntityNameLookupRecordConverter implements Converter<EntityNameLookupRecord> {
+
+  @Override
+  public EntityNameLookupRecord fromResultSet(ResultSet rs) throws SQLException {
+    return new EntityNameLookupRecord(
+        rs.getLong("catalog_id"),
+        rs.getLong("id"),
+        rs.getLong("parent_id"),
+        rs.getString("name"),
+        rs.getInt("type_code"),
+        rs.getInt("sub_type_code"));
+  }
+
+  @Override
+  public Map<String, Object> toMap(DatabaseType databaseType) {
+    throw new UnsupportedOperationException(
+        "EntityNameLookupRecordConverter is read-only and does not support toMap operation");
+  }
+}

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/models/ModelEntity.java

@@ -54,6 +54,9 @@ public class ModelEntity implements Converter<PolarisBaseEntity> {
           "grant_records_version",
           "location_without_scheme");
 
+  public static final List<String> ENTITY_LOOKUP_COLUMNS =
+      List.of("id", "catalog_id", "parent_id", "type_code", "name", "sub_type_code");
+
   // the id of the catalog associated to that entity. use 0 if this entity is top-level
   // like a catalog
   private long catalogId;

polaris-core/src/main/java/org/apache/polaris/core/entity/EntityNameLookupRecord.java

@@ -22,7 +22,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import java.util.Objects;
 
-public class EntityNameLookupRecord {
+public class EntityNameLookupRecord implements Identifiable {
   // entity catalog id
   private final long catalogId;
 
@@ -41,10 +41,12 @@ public class EntityNameLookupRecord {
   // code representing the subtype of that entity
   private final int subTypeCode;
 
+  @Override
   public long getCatalogId() {
     return catalogId;
   }
 
+  @Override
   public long getId() {
     return id;
   }

polaris-core/src/main/java/org/apache/polaris/core/entity/Identifiable.java

@@ -0,0 +1,37 @@
+/*
+ * 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 org.apache.polaris.core.entity;
+
+/**
+ * An interface for entity types that have an id. These entities provide a * method `getId` and
+ * `getCatalogId` to identify the entity.
+ */
+public interface Identifiable {
+  /**
+   * @return the id of the catalog associated to that entity. NULL_ID if this entity is top-level
+   *     like a catalog
+   */
+  long getCatalogId();
+
+  /**
+   * @return the id of the entity
+   */
+  long getId();
+}

polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisEntityCore.java

@@ -27,7 +27,7 @@
  * of the entity tree. For some operations like updating the entity, change will mean any change,
  * i.e. entity version mismatch.
  */
-public class PolarisEntityCore {
+public class PolarisEntityCore implements Identifiable {
 
   // the id of the catalog associated to that entity. NULL_ID if this entity is top-level like
   // a catalog
@@ -116,6 +116,7 @@ public PolarisEntityCore build() {
     }
   }
 
+  @Override
   public long getId() {
     return id;
   }
@@ -136,6 +137,7 @@ public int getEntityVersion() {
     return entityVersion;
   }
 
+  @Override
   public long getCatalogId() {
     return catalogId;
   }

polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/EntityIdToken.java

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import jakarta.annotation.Nullable;
+import org.apache.polaris.core.entity.Identifiable;
 import org.apache.polaris.core.entity.PolarisBaseEntity;
 import org.apache.polaris.immutables.PolarisImmutable;
 
@@ -41,7 +42,7 @@ default String getT() {
     return ID;
   }
 
-  static @Nullable EntityIdToken fromEntity(PolarisBaseEntity entity) {
+  static @Nullable EntityIdToken fromEntity(@Nullable Identifiable entity) {
     if (entity == null) {
       return null;
     }