diff --git a/superset/assets/javascripts/explorev2/stores/fields.js b/superset/assets/javascripts/explorev2/stores/fields.js index 7a8cf48ba31d6..a6a5d5faab4a1 100644 --- a/superset/assets/javascripts/explorev2/stores/fields.js +++ b/superset/assets/javascripts/explorev2/stores/fields.js @@ -376,6 +376,23 @@ export const fields = { 'domain_granularity. Should be larger or equal to Time Grain', }, + graph_color: { + type: 'SelectField', + label: 'Node Color', + default: [], + description: 'Choose a source node dimension to color the nodes.', + mapStateToProps: (state) => ({ + choices: (state.datasource) ? state.datasource.all_cols : [], + }), + }, + + graph_labels: { + type: 'CheckboxField', + label: 'Node Lables', + default: true, + description: 'Display lables on all node?', + }, + link_length: { type: 'SelectField', freeForm: true, diff --git a/superset/assets/javascripts/explorev2/stores/visTypes.js b/superset/assets/javascripts/explorev2/stores/visTypes.js index 1aea323b05840..8ba987f586858 100644 --- a/superset/assets/javascripts/explorev2/stores/visTypes.js +++ b/superset/assets/javascripts/explorev2/stores/visTypes.js @@ -544,6 +544,8 @@ const visTypes = { { label: 'Force Layout', fieldSetRows: [ + ['graph_color'], + ['graph_labels'], ['link_length'], ['charge'], ], diff --git a/superset/assets/visualizations/directed_force.js b/superset/assets/visualizations/directed_force.js index e2003f122eb15..626d9371771aa 100644 --- a/superset/assets/visualizations/directed_force.js +++ b/superset/assets/visualizations/directed_force.js @@ -10,6 +10,7 @@ const directedForceVis = function (slice, json) { const height = slice.height() - 25; const linkLength = json.form_data.link_length || 200; const charge = json.form_data.charge || -500; + const labels = json.form_data.graph_labels | 0; const links = json.data; const nodes = {}; @@ -46,6 +47,7 @@ const directedForceVis = function (slice, json) { } nodes[targetName].total += link.value; + nodes[targetName].color = link.color; }); /* eslint-disable no-use-before-define */ @@ -124,11 +126,11 @@ const directedForceVis = function (slice, json) { .select('circle') .transition() .style('stroke-width', 5); - d3.select(this) .select('text') .transition() - .style('font-size', 25); + .style('font-size', 25) + .style('opacity', 1); }) .on('mouseleave', function () { d3.select(this) @@ -138,7 +140,8 @@ const directedForceVis = function (slice, json) { d3.select(this) .select('text') .transition() - .style('font-size', 12); + .style('font-size', 12) + .style('opacity', labels); }) .call(force.drag); @@ -150,15 +153,21 @@ const directedForceVis = function (slice, json) { .domain(ext) .range([3, 30]); + const color = d3.scale.category20(); + node.append('circle') .attr('r', function (d) { return circleScale(Math.sqrt(d.total)); + }) + .style('fill', function (d) { + return color(d.color); }); // add the text node.append('text') .attr('x', 6) .attr('dy', '.35em') + .style('opacity', labels) .text(function (d) { return d.name; }); diff --git a/superset/data/__init__.py b/superset/data/__init__.py index e96ef36f35af9..fddbeab8c29c1 100644 --- a/superset/data/__init__.py +++ b/superset/data/__init__.py @@ -55,6 +55,7 @@ def load_energy(): if_exists='replace', chunksize=500, dtype={ + 'cluster': String(255), 'source': String(255), 'target': String(255), 'value': Float(), @@ -123,6 +124,7 @@ def load_energy(): ], "having": "", "link_length": "200", + "graph_color": "cluster", "metric": "sum__value", "row_limit": "5000", "slice_name": "Force", diff --git a/superset/data/energy.json.gz b/superset/data/energy.json.gz index 624d71db683a9..6c51dc3cac07c 100644 Binary files a/superset/data/energy.json.gz and b/superset/data/energy.json.gz differ diff --git a/superset/forms.py b/superset/forms.py index 150c0b7097bbe..d8f42d36f9541 100755 --- a/superset/forms.py +++ b/superset/forms.py @@ -321,6 +321,17 @@ def __init__(self, viz): "choices": self.choicify(datasource.groupby_column_names), "description": _("One or many fields to pivot as columns") }), + 'graph_color': (SelectField, { + "label": _("Node Color"), + "default": None, + "choices": self.choicify([None] + datasource.column_names), + "description": _("Choose a source node dimension to color the nodes"), + }), + 'graph_labels': (BetterBooleanField, { + "label": _("Node Labels"), + "default": True, + "description": _("Display labels on all node?") + }), 'all_columns': (SelectMultipleSortableField, { "label": _("Columns"), "choices": self.choicify(datasource.column_names), diff --git a/superset/viz.py b/superset/viz.py index eaf5a63740883..235869a188919 100755 --- a/superset/viz.py +++ b/superset/viz.py @@ -1759,6 +1759,8 @@ class DirectedForceViz(BaseViz): }, { 'label': _('Force Layout'), 'fields': ( + 'graph_color', + 'graph_labels', 'link_length', 'charge', ) @@ -1775,11 +1777,16 @@ def query_obj(self): if len(self.form_data['groupby']) != 2: raise Exception("Pick exactly 2 columns to 'Group By'") qry['metrics'] = [self.form_data['metric']] + if self.form_data['graph_color'] != 'None': + qry['groupby'] += [self.form_data['graph_color']] return qry def get_data(self): df = self.get_df() - df.columns = ['source', 'target', 'value'] + if self.form_data['graph_color'] == 'None': + df.columns = ['source', 'target', 'value'] + else: + df.columns = ['source', 'target', 'color', 'value'] return df.to_dict(orient='records')