Merge branch 'apache:main' into GH-125

Author: Kaustav-Sarkar

.github/workflows/test.yml

@@ -89,7 +89,7 @@ jobs:
         include:
           - arch: AMD64
             jdk: 11
-            macos: 13
+            macos: 15-intel
           - arch: AArch64
             jdk: 11
             macos: latest

c/src/test/java/org/apache/arrow/c/RoundtripTest.java

@@ -35,7 +35,6 @@
 import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
-import org.apache.arrow.memory.util.hash.ArrowBufHasher;
 import org.apache.arrow.vector.BaseLargeVariableWidthVector;
 import org.apache.arrow.vector.BaseVariableWidthVector;
 import org.apache.arrow.vector.BigIntVector;
@@ -44,7 +43,6 @@
 import org.apache.arrow.vector.DateMilliVector;
 import org.apache.arrow.vector.DecimalVector;
 import org.apache.arrow.vector.DurationVector;
-import org.apache.arrow.vector.ExtensionTypeVector;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.FixedSizeBinaryVector;
 import org.apache.arrow.vector.Float2Vector;
@@ -74,6 +72,7 @@
 import org.apache.arrow.vector.UInt2Vector;
 import org.apache.arrow.vector.UInt4Vector;
 import org.apache.arrow.vector.UInt8Vector;
+import org.apache.arrow.vector.UuidVector;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.VarBinaryVector;
 import org.apache.arrow.vector.VarCharVector;
@@ -92,6 +91,7 @@
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.complex.UnionVector;
 import org.apache.arrow.vector.complex.impl.UnionMapWriter;
+import org.apache.arrow.vector.extension.UuidType;
 import org.apache.arrow.vector.holders.IntervalDayHolder;
 import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder;
 import org.apache.arrow.vector.holders.NullableUInt4Holder;
@@ -100,7 +100,6 @@
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
-import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -810,9 +809,8 @@ public void testEmptyRunEndEncodedVector() {
 
   @Test
   public void testExtensionTypeVector() {
-    ExtensionTypeRegistry.register(new UuidType());
     final Schema schema =
-        new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
+        new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE)));
     try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
       // Fill with data
       UUID u1 = UUID.randomUUID();
@@ -830,13 +828,12 @@ public void testExtensionTypeVector() {
       assertEquals(root.getSchema(), importedRoot.getSchema());
 
       final Field field = importedRoot.getSchema().getFields().get(0);
-      final UuidType expectedType = new UuidType();
       assertEquals(
           field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
-          expectedType.extensionName());
+          UuidType.INSTANCE.extensionName());
       assertEquals(
           field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA),
-          expectedType.serialize());
+          UuidType.INSTANCE.serialize());
 
       final UuidVector deserialized = (UuidVector) importedRoot.getFieldVectors().get(0);
       assertEquals(vector.getValueCount(), deserialized.getValueCount());
@@ -1115,72 +1112,4 @@ private VectorSchemaRoot createTestVSR() {
 
     return new VectorSchemaRoot(fields, vectors);
   }
-
-  static class UuidType extends ExtensionType {
-
-    @Override
-    public ArrowType storageType() {
-      return new ArrowType.FixedSizeBinary(16);
-    }
-
-    @Override
-    public String extensionName() {
-      return "uuid";
-    }
-
-    @Override
-    public boolean extensionEquals(ExtensionType other) {
-      return other instanceof UuidType;
-    }
-
-    @Override
-    public ArrowType deserialize(ArrowType storageType, String serializedData) {
-      if (!storageType.equals(storageType())) {
-        throw new UnsupportedOperationException(
-            "Cannot construct UuidType from underlying type " + storageType);
-      }
-      return new UuidType();
-    }
-
-    @Override
-    public String serialize() {
-      return "";
-    }
-
-    @Override
-    public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
-      return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16));
-    }
-  }
-
-  static class UuidVector extends ExtensionTypeVector<FixedSizeBinaryVector> {
-
-    public UuidVector(
-        String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) {
-      super(name, allocator, underlyingVector);
-    }
-
-    @Override
-    public UUID getObject(int index) {
-      final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index));
-      return new UUID(bb.getLong(), bb.getLong());
-    }
-
-    @Override
-    public int hashCode(int index) {
-      return hashCode(index, null);
-    }
-
-    @Override
-    public int hashCode(int index, ArrowBufHasher hasher) {
-      return getUnderlyingVector().hashCode(index, hasher);
-    }
-
-    public void set(int index, UUID uuid) {
-      ByteBuffer bb = ByteBuffer.allocate(16);
-      bb.putLong(uuid.getMostSignificantBits());
-      bb.putLong(uuid.getLeastSignificantBits());
-      getUnderlyingVector().set(index, bb.array());
-    }
-  }
 }

ci/scripts/jni_macos_build.sh

@@ -77,7 +77,9 @@ cmake \
 cmake --build "${build_dir}/cpp" --target install
 github_actions_group_end
 
-export JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/protobuf_ep-install"
+JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/_deps/protobuf-build"
+JAVA_JNI_CMAKE_ARGS+=" -DProtobuf_SRC_ROOT_FOLDER=${build_dir}/cpp/_deps/protobuf-src"
+export JAVA_JNI_CMAKE_ARGS
 "${source_dir}/ci/scripts/jni_build.sh" \
   "${source_dir}" \
   "${install_dir}" \

dataset/pom.xml

@@ -32,7 +32,7 @@ under the License.
 
   <properties>
     <arrow.cpp.build.dir>../../../cpp/release-build/</arrow.cpp.build.dir>
-    <parquet.version>1.15.2</parquet.version>
+    <parquet.version>1.16.0</parquet.version>
     <avro.version>1.12.0</avro.version>
   </properties>
 

flight/flight-core/pom.xml

@@ -134,7 +134,7 @@ under the License.
     <dependency>
       <groupId>com.google.api.grpc</groupId>
       <artifactId>proto-google-common-protos</artifactId>
-      <version>2.56.0</version>
+      <version>2.63.2</version>
       <scope>test</scope>
     </dependency>
     <dependency>

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java

@@ -79,7 +79,8 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p
   public void closeStatement(final StatementHandle statementHandle) {
     PreparedStatement preparedStatement =
         statementHandlePreparedStatementMap.remove(new StatementHandleKey(statementHandle));
-    // Testing if the prepared statement was created because the statement can be not created until
+    // Testing if the prepared statement was created because the statement can be
+    // not created until
     // this moment
     if (preparedStatement != null) {
       preparedStatement.close();
@@ -224,7 +225,8 @@ public ExecuteResult prepareAndExecute(
           MetaResultSet.create(handle.connectionId, handle.id, false, handle.signature, null);
       return new ExecuteResult(Collections.singletonList(metaResultSet));
     } catch (SQLTimeoutException e) {
-      // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and Runtime
+      // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and
+      // Runtime
       // Exceptions.
       throw new RuntimeException(e);
     } catch (SQLException e) {
@@ -253,6 +255,20 @@ public boolean syncResults(
     return false;
   }
 
+  @Override
+  public ConnectionProperties connectionSync(ConnectionHandle ch, ConnectionProperties connProps) {
+    final ConnectionProperties result = super.connectionSync(ch, connProps);
+    final String newCatalog = this.connProps.getCatalog();
+    if (newCatalog != null) {
+      try {
+        ((ArrowFlightConnection) connection).getClientHandler().setCatalog(newCatalog);
+      } catch (SQLException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    return result;
+  }
+
   void setDefaultConnectionProperties() {
     // TODO Double-check this.
     connProps
@@ -268,7 +284,8 @@ PreparedStatement getPreparedStatement(StatementHandle statementHandle) {
     return statementHandlePreparedStatementMap.get(new StatementHandleKey(statementHandle));
   }
 
-  // Helper used to look up prepared statement instances later. Avatica doesn't give us the
+  // Helper used to look up prepared statement instances later. Avatica doesn't
+  // give us the
   // signature in
   // an UPDATE code path so we can't directly use StatementHandle as a map key.
   private static final class StatementHandleKey {

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java

@@ -47,7 +47,6 @@
 import org.apache.arrow.flight.FlightStatusCode;
 import org.apache.arrow.flight.Location;
 import org.apache.arrow.flight.LocationSchemes;
-import org.apache.arrow.flight.SessionOptionValue;
 import org.apache.arrow.flight.SessionOptionValueFactory;
 import org.apache.arrow.flight.SetSessionOptionsRequest;
 import org.apache.arrow.flight.SetSessionOptionsResult;
@@ -147,20 +146,26 @@ public List<CloseableEndpointStreamPair> getStreams(final FlightInfo flightInfo)
     try {
       for (FlightEndpoint endpoint : flightInfo.getEndpoints()) {
         if (endpoint.getLocations().isEmpty()) {
-          // Create a stream using the current client only and do not close the client at the end.
+          // Create a stream using the current client only and do not close the client at
+          // the end.
           endpoints.add(
               new CloseableEndpointStreamPair(
                   sqlClient.getStream(endpoint.getTicket(), getOptions()), null));
         } else {
           // Clone the builder and then set the new endpoint on it.
 
-          // GH-38574: Currently a new FlightClient will be made for each partition that returns a
-          // non-empty Location then disposed of. It may be better to cache clients because a server
-          // may report the same Locations. It would also be good to identify when the reported
+          // GH-38574: Currently a new FlightClient will be made for each partition that
+          // returns a
+          // non-empty Location then disposed of. It may be better to cache clients
+          // because a server
+          // may report the same Locations. It would also be good to identify when the
+          // reported
           // location
-          // is the same as the original connection's Location and skip creating a FlightClient in
+          // is the same as the original connection's Location and skip creating a
+          // FlightClient in
           // that scenario.
-          // Also copy the cache to the client so we can share a cache. Cache needs to cache
+          // Also copy the cache to the client so we can share a cache. Cache needs to
+          // cache
           // negative attempts too.
           List<Exception> exceptions = new ArrayList<>();
           CloseableEndpointStreamPair stream = null;
@@ -337,7 +342,8 @@ private boolean isBenignCloseException(FlightRuntimeException fre) {
    */
   private void logSuppressedCloseException(
       FlightRuntimeException fre, String operationDescription) {
-    // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during shutdown
+    // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during
+    // shutdown
     LOGGER.debug("Suppressed error {}", operationDescription, fre);
   }
 
@@ -388,25 +394,40 @@ public interface PreparedStatement extends AutoCloseable {
   /** A connection is created with catalog set as a session option. */
   private void setSetCatalogInSessionIfPresent() {
     if (catalog.isPresent()) {
-      final SetSessionOptionsRequest setSessionOptionRequest =
-          new SetSessionOptionsRequest(
-              ImmutableMap.<String, SessionOptionValue>builder()
-                  .put(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog.get()))
-                  .build());
-      final SetSessionOptionsResult result =
-          sqlClient.setSessionOptions(setSessionOptionRequest, getOptions());
+      try {
+        setCatalog(catalog.get());
+      } catch (SQLException e) {
+        throw CallStatus.INVALID_ARGUMENT
+            .withDescription(e.getMessage())
+            .withCause(e)
+            .toRuntimeException();
+      }
+    }
+  }
 
+  /**
+   * Sets the catalog for the current session.
+   *
+   * @param catalog the catalog to set.
+   * @throws SQLException if an error occurs while setting the catalog.
+   */
+  public void setCatalog(final String catalog) throws SQLException {
+    final SetSessionOptionsRequest request =
+        new SetSessionOptionsRequest(
+            ImmutableMap.of(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog)));
+    try {
+      final SetSessionOptionsResult result = sqlClient.setSessionOptions(request, getOptions());
       if (result.hasErrors()) {
-        Map<String, SetSessionOptionsResult.Error> errors = result.getErrors();
-        for (Map.Entry<String, SetSessionOptionsResult.Error> error : errors.entrySet()) {
+        final Map<String, SetSessionOptionsResult.Error> errors = result.getErrors();
+        for (final Map.Entry<String, SetSessionOptionsResult.Error> error : errors.entrySet()) {
           LOGGER.warn(error.toString());
         }
-        throw CallStatus.INVALID_ARGUMENT
-            .withDescription(
-                String.format(
-                    "Cannot set session option for catalog = %s. Check log for details.", catalog))
-            .toRuntimeException();
+        throw new SQLException(
+            String.format(
+                "Cannot set session option for catalog = %s. Check log for details.", catalog));
       }
+    } catch (final FlightRuntimeException e) {
+      throw new SQLException(e);
     }
   }
 
@@ -654,7 +675,8 @@ public static final class Builder {
 
     @VisibleForTesting @Nullable Duration connectTimeout;
 
-    // These two middleware are for internal use within build() and should not be exposed by builder
+    // These two middleware are for internal use within build() and should not be
+    // exposed by builder
     // APIs.
     // Note that these middleware may not necessarily be registered.
     @VisibleForTesting
@@ -980,15 +1002,17 @@ public Location getLocation() {
      * @throws SQLException on error.
      */
     public ArrowFlightSqlClientHandler build() throws SQLException {
-      // Copy middleware so that the build method doesn't change the state of the builder fields
+      // Copy middleware so that the build method doesn't change the state of the
+      // builder fields
       // itself.
       Set<FlightClientMiddleware.Factory> buildTimeMiddlewareFactories =
           new HashSet<>(this.middlewareFactories);
       FlightClient client = null;
       boolean isUsingUserPasswordAuth = username != null && token == null;
 
       try {
-        // Token should take priority since some apps pass in a username/password even when a token
+        // Token should take priority since some apps pass in a username/password even
+        // when a token
         // is provided
         if (isUsingUserPasswordAuth) {
           buildTimeMiddlewareFactories.add(authFactory);
@@ -1047,8 +1071,10 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
                 allocator, channelBuilder.build(), clientBuilder.middleware());
         final ArrayList<CallOption> credentialOptions = new ArrayList<>();
         if (isUsingUserPasswordAuth) {
-          // If the authFactory has already been used for a handshake, use the existing token.
-          // This can occur if the authFactory is being re-used for a new connection spawned for
+          // If the authFactory has already been used for a handshake, use the existing
+          // token.
+          // This can occur if the authFactory is being re-used for a new connection
+          // spawned for
           // getStream().
           if (authFactory.getCredentialCallOption() != null) {
             credentialOptions.add(authFactory.getCredentialCallOption());