Adds usage stats for vectors (#45023)

mayya-sharipova · web-flow · commit 0c687650884e · 2019-07-31T12:32:41.000-04:00
Example of usage: _xpack/usage "vectors": { "available": true, "enabled": true, "dense_vector_fields_count" : 1, "sparse_vector_fields_count" : 1, "dense_vector_dims_avg_count" : 100 } Backport for #44512
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/vectors/VectorsFeatureSetUsage.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/vectors/VectorsFeatureSetUsage.java
@@ -7,18 +7,72 @@
 package org.elasticsearch.xpack.core.vectors;
 
 import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.XPackFeatureSet;
 import org.elasticsearch.xpack.core.XPackField;
 
 import java.io.IOException;
+import java.util.Objects;
 
 public class VectorsFeatureSetUsage extends XPackFeatureSet.Usage {
 
+    private final int numDenseVectorFields;
+    private final int numSparseVectorFields;
+    private final int avgDenseVectorDims;
+
     public VectorsFeatureSetUsage(StreamInput input) throws IOException {
         super(input);
+        numDenseVectorFields = input.readVInt();
+        numSparseVectorFields = input.readVInt();
+        avgDenseVectorDims = input.readVInt();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeVInt(numDenseVectorFields);
+        out.writeVInt(numSparseVectorFields);
+        out.writeVInt(avgDenseVectorDims);
     }
 
-    public VectorsFeatureSetUsage(boolean available, boolean enabled) {
+    public VectorsFeatureSetUsage(boolean available, boolean enabled, int numDenseVectorFields, int numSparseVectorFields,
+                                  int avgDenseVectorDims) {
         super(XPackField.VECTORS, available, enabled);
+        this.numDenseVectorFields = numDenseVectorFields;
+        this.numSparseVectorFields = numSparseVectorFields;
+        this.avgDenseVectorDims = avgDenseVectorDims;
+    }
+
+
+    @Override
+    protected void innerXContent(XContentBuilder builder, Params params) throws IOException {
+        super.innerXContent(builder, params);
+        builder.field("dense_vector_fields_count", numDenseVectorFields);
+        builder.field("sparse_vector_fields_count", numSparseVectorFields);
+        builder.field("dense_vector_dims_avg_count", avgDenseVectorDims);
+    }
+
+    public int numDenseVectorFields() {
+        return numDenseVectorFields;
+    }
+    public int numSparseVectorFields() {
+        return numSparseVectorFields;
+    }
+    public int avgDenseVectorDims() {
+        return avgDenseVectorDims;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(available, enabled, numDenseVectorFields, numSparseVectorFields, avgDenseVectorDims);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj instanceof VectorsFeatureSetUsage == false) return false;
+        VectorsFeatureSetUsage other = (VectorsFeatureSetUsage) obj;
+        return available == other.available && enabled == other.enabled && numDenseVectorFields == other.numDenseVectorFields
+            && numSparseVectorFields == other.numSparseVectorFields && avgDenseVectorDims == other.avgDenseVectorDims;
     }
 }
diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/vectors/VectorsFeatureSetUsageTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/vectors/VectorsFeatureSetUsageTests.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.core.vectors;
+
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.test.AbstractWireSerializingTestCase;
+
+import java.io.IOException;
+
+public class VectorsFeatureSetUsageTests extends AbstractWireSerializingTestCase<VectorsFeatureSetUsage> {
+
+    @Override
+    protected VectorsFeatureSetUsage createTestInstance() {
+        boolean available = randomBoolean();
+        boolean enabled = randomBoolean();
+        if (available && enabled) {
+            return new VectorsFeatureSetUsage(available, enabled, randomIntBetween(0, 100000), randomIntBetween(0, 100000),
+                randomIntBetween(0, 1024));
+        } else {
+            return new VectorsFeatureSetUsage(available, enabled, 0, 0, 0);
+        }
+    }
+
+    @Override
+    protected VectorsFeatureSetUsage mutateInstance(VectorsFeatureSetUsage instance) throws IOException {
+        boolean available = instance.available();
+        boolean enabled = instance.enabled();
+        int numDenseVectorFields = instance.numDenseVectorFields();
+        int numSparseVectorFields = instance.numSparseVectorFields();
+        int avgDenseVectorDims = instance.avgDenseVectorDims();
+
+        if (available == false || enabled == false) {
+            available = true;
+            enabled = true;
+        }
+        numDenseVectorFields = randomValueOtherThan(numDenseVectorFields, () -> randomIntBetween(0, 100000));
+        numSparseVectorFields = randomValueOtherThan(numSparseVectorFields, () -> randomIntBetween(0, 100000));
+        avgDenseVectorDims = randomValueOtherThan(avgDenseVectorDims, () -> randomIntBetween(0, 1024));
+        return new VectorsFeatureSetUsage(available, enabled, numDenseVectorFields, numSparseVectorFields, avgDenseVectorDims);
+    }
+
+    @Override
+    protected Writeable.Reader<VectorsFeatureSetUsage> instanceReader() {
+        return VectorsFeatureSetUsage::new;
+    }
+
+}
diff --git a/x-pack/plugin/src/test/resources/rest-api-spec/test/vectors/50_vector_stats.yml b/x-pack/plugin/src/test/resources/rest-api-spec/test/vectors/50_vector_stats.yml
@@ -0,0 +1,46 @@
+setup:
+  - skip:
+      features: headers
+      version: " - 7.3.99"
+      reason: "vector stats was added from 7.4"
+
+---
+"Usage stats on vector fields":
+  - do: {xpack.usage: {}}
+  - match: { vectors.available: true }
+  - match: { vectors.enabled: true }
+  - match: { vectors.dense_vector_fields_count: 0 }
+  - match: { vectors.sparse_vector_fields_count: 0 }
+  - match: { vectors.dense_vector_dims_avg_count: 0 }
+
+  - do:
+      indices.create:
+        index: test-index1
+        body:
+          mappings:
+            properties:
+              my_dense_vector1:
+                type: dense_vector
+                dims: 10
+              my_dense_vector2:
+                type: dense_vector
+                dims: 30
+
+  - do:
+      indices.create:
+        index: test-index2
+        body:
+          mappings:
+            properties:
+              my_dense_vector3:
+                type: dense_vector
+                dims: 20
+              my_sparse_vector1:
+                type: sparse_vector
+
+  - do: {xpack.usage: {}}
+  - match: { vectors.available: true }
+  - match: { vectors.enabled: true }
+  - match: { vectors.dense_vector_fields_count: 3 }
+  - match: { vectors.sparse_vector_fields_count: 1 }
+  - match: { vectors.dense_vector_dims_avg_count: 20 }
diff --git a/x-pack/plugin/vectors/src/main/java/org/elasticsearch/xpack/vectors/VectorsFeatureSet.java b/x-pack/plugin/vectors/src/main/java/org/elasticsearch/xpack/vectors/VectorsFeatureSet.java
@@ -6,25 +6,32 @@
 package org.elasticsearch.xpack.vectors;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.cluster.metadata.IndexMetaData;
+import org.elasticsearch.cluster.metadata.MappingMetaData;
+import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.xpack.core.XPackFeatureSet;
 import org.elasticsearch.xpack.core.XPackField;
 import org.elasticsearch.xpack.core.XPackSettings;
 import org.elasticsearch.xpack.core.vectors.VectorsFeatureSetUsage;
+import org.elasticsearch.xpack.vectors.mapper.DenseVectorFieldMapper;
+import org.elasticsearch.xpack.vectors.mapper.SparseVectorFieldMapper;
 
 import java.util.Map;
 
 public class VectorsFeatureSet implements XPackFeatureSet {
 
     private final boolean enabled;
     private final XPackLicenseState licenseState;
+    private final ClusterService clusterService;
 
     @Inject
-    public VectorsFeatureSet(Settings settings, XPackLicenseState licenseState) {
+    public VectorsFeatureSet(Settings settings, XPackLicenseState licenseState, ClusterService clusterService) {
         this.enabled = XPackSettings.VECTORS_ENABLED.get(settings);
         this.licenseState = licenseState;
+        this.clusterService = clusterService;
     }
 
     @Override
@@ -49,6 +56,40 @@ public Map<String, Object> nativeCodeInfo() {
 
     @Override
     public void usage(ActionListener<XPackFeatureSet.Usage> listener) {
-        listener.onResponse(new VectorsFeatureSetUsage(available(), enabled()));
+        boolean vectorsAvailable = available();
+        boolean vectorsEnabled = enabled();
+        int numDenseVectorFields = 0;
+        int numSparseVectorFields = 0;
+        int avgDenseVectorDims = 0;
+
+        if (vectorsAvailable && vectorsEnabled && clusterService.state() != null) {
+            for (IndexMetaData indexMetaData : clusterService.state().metaData()) {
+                MappingMetaData mappingMetaData = indexMetaData.mapping();
+                if (mappingMetaData != null) {
+                    Map<String, Object> mappings = mappingMetaData.getSourceAsMap();
+                    if (mappings.containsKey("properties")) {
+                        @SuppressWarnings("unchecked") Map<String, Map<String, Object>> fieldMappings =
+                            (Map<String, Map<String, Object>>) mappings.get("properties");
+                        for (Map<String, Object> typeDefinition : fieldMappings.values()) {
+                            String fieldType = (String) typeDefinition.get("type");
+                            if (fieldType != null) {
+                                if (fieldType.equals(DenseVectorFieldMapper.CONTENT_TYPE)) {
+                                    numDenseVectorFields++;
+                                    int dims = (Integer) typeDefinition.get("dims");
+                                    avgDenseVectorDims += dims;
+                                } else if (fieldType.equals(SparseVectorFieldMapper.CONTENT_TYPE)) {
+                                    numSparseVectorFields++;
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+            if (numDenseVectorFields > 0) {
+                avgDenseVectorDims = avgDenseVectorDims / numDenseVectorFields;
+            }
+        }
+        listener.onResponse(new VectorsFeatureSetUsage(vectorsAvailable, vectorsEnabled,
+            numDenseVectorFields, numSparseVectorFields, avgDenseVectorDims));
     }
 }
diff --git a/x-pack/plugin/vectors/src/test/java/org/elasticsearch/xpack/vectors/VectorsFeatureSetTests.java b/x-pack/plugin/vectors/src/test/java/org/elasticsearch/xpack/vectors/VectorsFeatureSetTests.java
@@ -5,43 +5,51 @@
  */
 package org.elasticsearch.xpack.vectors;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.cluster.ClusterName;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.IndexMetaData;
+import org.elasticsearch.cluster.metadata.MetaData;
+import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.core.XPackFeatureSet;
 import org.elasticsearch.xpack.core.vectors.VectorsFeatureSetUsage;
 import org.junit.Before;
+import org.mockito.Mockito;
 
-import static org.hamcrest.core.Is.is;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 public class VectorsFeatureSetTests extends ESTestCase {
 
     private XPackLicenseState licenseState;
+    private ClusterService clusterService;
 
     @Before
     public void init() {
         licenseState = mock(XPackLicenseState.class);
+        clusterService = mock(ClusterService.class);
     }
 
     public void testAvailable() throws Exception {
-        VectorsFeatureSet featureSet = new VectorsFeatureSet(Settings.EMPTY, licenseState);
+        VectorsFeatureSet featureSet = new VectorsFeatureSet(Settings.EMPTY, licenseState, clusterService);
         boolean available = randomBoolean();
         when(licenseState.isVectorsAllowed()).thenReturn(available);
-        assertThat(featureSet.available(), is(available));
+        assertEquals(available, featureSet.available());
 
         PlainActionFuture<XPackFeatureSet.Usage> future = new PlainActionFuture<>();
         featureSet.usage(future);
         XPackFeatureSet.Usage usage = future.get();
-        assertThat(usage.available(), is(available));
+        assertEquals(available, usage.available());
 
         BytesStreamOutput out = new BytesStreamOutput();
         usage.writeTo(out);
         XPackFeatureSet.Usage serializedUsage = new VectorsFeatureSetUsage(out.bytes().streamInput());
-        assertThat(serializedUsage.available(), is(available));
+        assertEquals(available, serializedUsage.available());
     }
 
     public void testEnabled() throws Exception {
@@ -54,17 +62,48 @@ public void testEnabled() throws Exception {
         } else {
             settings.put("xpack.vectors.enabled", enabled);
         }
-        VectorsFeatureSet featureSet = new VectorsFeatureSet(settings.build(), licenseState);
-        assertThat(featureSet.enabled(), is(enabled));
+        VectorsFeatureSet featureSet = new VectorsFeatureSet(settings.build(), licenseState, clusterService);
+        assertEquals(enabled, featureSet.enabled());
         PlainActionFuture<XPackFeatureSet.Usage> future = new PlainActionFuture<>();
         featureSet.usage(future);
         XPackFeatureSet.Usage usage = future.get();
-        assertThat(usage.enabled(), is(enabled));
+        assertEquals(enabled, usage.enabled());
 
         BytesStreamOutput out = new BytesStreamOutput();
         usage.writeTo(out);
         XPackFeatureSet.Usage serializedUsage = new VectorsFeatureSetUsage(out.bytes().streamInput());
-        assertThat(serializedUsage.enabled(), is(enabled));
+        assertEquals(enabled, serializedUsage.enabled());
     }
 
+    public void testUsageStats() throws Exception {
+        MetaData.Builder metaData = MetaData.builder();
+        IndexMetaData.Builder index1 = IndexMetaData.builder("test-index1")
+            .settings(settings(Version.CURRENT)).numberOfShards(1).numberOfReplicas(0)
+            .putMapping("_doc",
+                "{\"properties\":{\"my_dense_vector1\":{\"type\":\"dense_vector\",\"dims\": 10}," +
+                    "\"my_dense_vector2\":{\"type\":\"dense_vector\",\"dims\": 30} }}");
+        IndexMetaData.Builder index2 = IndexMetaData.builder("test-index2")
+            .settings(settings(Version.CURRENT)).numberOfShards(1).numberOfReplicas(0)
+            .putMapping("_doc",
+                "{\"properties\":{\"my_dense_vector3\":{\"type\":\"dense_vector\",\"dims\": 20}," +
+                    "\"my_sparse_vector1\":{\"type\":\"sparse_vector\"} }}");
+        metaData.put(index1);
+        metaData.put(index2);
+        ClusterState clusterState = ClusterState.builder(new ClusterName("_testcluster")).metaData(metaData).build();
+
+        Mockito.when(clusterService.state()).thenReturn(clusterState);
+        when(licenseState.isVectorsAllowed()).thenReturn(true);
+        Settings.Builder settings = Settings.builder();
+        settings.put("xpack.vectors.enabled", true);
+
+        PlainActionFuture<XPackFeatureSet.Usage> future = new PlainActionFuture<>();
+        VectorsFeatureSet vectorsFeatureSet = new VectorsFeatureSet(settings.build(), licenseState, clusterService);
+        vectorsFeatureSet.usage(future);
+        VectorsFeatureSetUsage vectorUsage = (VectorsFeatureSetUsage) future.get();
+        assertEquals(true, vectorUsage.enabled());
+        assertEquals(true, vectorUsage.available());
+        assertEquals(3, vectorUsage.numDenseVectorFields());
+        assertEquals(1, vectorUsage.numSparseVectorFields());
+        assertEquals(20, vectorUsage.avgDenseVectorDims());
+    }
 }
