filter values in an array based on whether they are in a column of another table

[dschneiderch]

how can I filter the segment_id array column I create in t.mutate and only keep the matching segment_id's from tax ?

import ibis
from ibis import _

ibis.options.interactive = True

t = ibis.memtable(
    {
        "user": ["a", "b"],
        "attribute": [
            [
                {"id": "attr1", "sv": 1},
                {"id": "attr2", "sv": 2},
                {"id": "attr3", "sv": 3},
            ],
            [{"id": "attr1", "sv": 1}, {"id": "attr5", "sv": None}],
        ],
    }
)

tax = ibis.memtable({"segment_id": ["attr1"], "segment_name": ["attribute 1"]})

def create_segment_id(row):
    return row["id"] + "_" + row["sv"].cast("string")

t.mutate(
    segment_id=_.attribute.map(create_segment_id)
    ).select(_.user, _.segment_id
             ).mutate(segment_id=_.segment_id.join(","))
    

I tried using partial and passing in tax.segment_id but I get ValueError: The truth value of an Ibis expression is not defined

from functools import partial


def create_segment_id(row, taxonomy):
    segment_id = row["id"] + "_" + row["sv"].cast("string")
    return segment_id if taxonomy.contains(segment_id) else None


create_segment_id_partial = partial(
    create_segment_id, taxonomy=tax.segment_id
)

t.mutate(
    segment_id=_.attribute.map(create_segment_id_partial
).select(_.user, _.segment_id
             ).mutate(segment_id=_.segment_id.join(","))

[dschneiderch]

I can get what I want with unnest() and agg and inner join but that seems awfully wasteful. it would be better to filter at the array level before unnest()

t.mutate(
    segment_id=_.attribute.map(
        lambda row: row["id"].concat("_", row["sv"].cast("string"))
    )
).select(_.user, _.segment_id).unnest("segment_id"
                                      ).join(tax, _.segment_id == tax.segment_id
                                      ).select(_.user, _.segment_id
                                               ).group_by("user"
                                                    ).agg(v=_.segment_id.collect().join(","))

[gforsyth]

In the relational model (which is what any SQL-ish execution engine is using), you either have to join the two tables, or pass in the values from one table inline as some sort of scalar value. There's not really any way around that.

If you want to perform a comparison against a set of tax.segment_id values, you can pass it in as a scalar value, doing something like:

t.filter(t.segment_id.contains(tax.segment_id.as_scalar())

That will pass in the array of tax.segment_id and check for membership, but this will happen to every row in t.segment_id.

[dschneiderch]

I expect to pass entire column of tax.segment_id to every array. In this case anyway, i only expect tax.segment_id to have ~500 records while t will have millions so i'd expect this to be more performant than unnesting millions x tens of array values in each record. if you're familiar with Beam at all, I'd treat tax.segment_id as a side_input.

That said, 1. t.segment_id.contains(tax.segment_id.as_scalar()) seems backwards. tax.segment_id is a list of all allowed segment_ids so i think I'd want t.segment_id.isin(tax.segment_id.as_scalar() . Am I misunderstanding these functions?

using contains doesn't filter anything when I tried it

t.mutate(
    segment_id=_.attribute.map(
        lambda row: row["id"].concat("_", row["sv"].cast("string"))
    )
).select(_.user, _.segment_id).filter(
    _.segment_id.contains(tax.segment_id.as_scalar())
)

2. when i tried with isin i get this error

SignatureValidationError: InValues(┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ segment_id                     ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ array<string>                  │
├────────────────────────────────┤
│ ['attr1_1', 'attr2_2', ... +1] │
│ ['attr1_1', None]              │
└────────────────────────────────┘, ┌─────────┐
│ attr1_1 │
└─────────┘) has failed due to the following errors:
  `options`: ┌─────────┐
│ attr1_1 │
└─────────┘ is not a tuple of coercibles to a Value

Expected signature: InValues(value: Value, options: tuple[Value, ...])

I appreciate your help!

[gforsyth]

You're right that contains is backwards for what you're trying to do -- however, isin isn't an array-level method, it works on the column as a whole, so I don't think it's going to do what you want.

And I realized that as_scalar really does require the result to only be a single value. So I think perhaps you'll need an intermediate materialization if a join isn't practical.

It's a bit verbose, but you can evaluate the tax.segment_id as pyarrow and then make it a list, and that can be passed inline to intersect -- the general condition here is filter where arr1.intersect(arr2).length() > 0

[ins] In [53]: tax
Out[53]: 
┏━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓
┃ segment_id ┃ segment_name ┃
┡━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩
│ string     │ string       │
├────────────┼──────────────┤
│ attr1_2    │ attribute 1  │
│ attr2_2    │ attribute 2  │
└────────────┴──────────────┘

[ins] In [54]: t
Out[54]: 
┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ user   ┃ attribute                            ┃ segment_id                     ┃
┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ string │ array<struct<id: string, sv: int64>> │ array<string>                  │
├────────┼──────────────────────────────────────┼────────────────────────────────┤
│ a      │ [{...}, {...}, ... +1]               │ ['attr1_1', 'attr2_2', ... +1] │
│ b      │ [{...}, {...}]                       │ ['attr1_1', None]              │
└────────┴──────────────────────────────────────┴────────────────────────────────┘

[ins] In [55]: t.segment_id.intersect(tax.segment_id.to_pyarrow().to_pylist())
Out[55]: 
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ ArrayIntersect(segment_id, ('attr1_2', 'attr2_2')) ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ array<string>                                      │
├────────────────────────────────────────────────────┤
│ ['attr2_2']                                        │
│ []                                                 │
└────────────────────────────────────────────────────┘

[ins] In [56]: t.segment_id.intersect(tax.segment_id.to_pyarrow().to_pylist()).length
          ...: () > 0
Out[56]: 
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Greater(ArrayLength(ArrayIntersect(segment_id, ('attr1_2', 'attr2_2'))), 0) ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ boolean                                                                     │
├─────────────────────────────────────────────────────────────────────────────┤
│ True                                                                        │
│ False                                                                       │
└─────────────────────────────────────────────────────────────────────────────┘

[dschneiderch]

Ok! That did it. I pasted both compiled sql into bigquery and actually the unnest version used less slots at the expense of a little more bytes shuffled.
Thank you!

(
    standard_segments.filter(
        _.last_mapped_id_modified == "2024-10-08 06:00:00+00:00", _.type == "ramp_id"
    )
    .mutate(segment_id=_.attribute.map(create_segment_id))
    .select(_.value, _.segment_id)
    .group_by("value")
    .mutate(
        segment_id=_.segment_id.intersect(openx_standard.segment_id.to_pyarrow().to_pylist()).join(
            ","
        )
    )
    .filter(_.segment_id.notnull())
)