Redact at the Speed of AI: De‑identifying PII in the Lakehouse with ai_query

Summary

ai_query is a built-in SQL function in the Databricks Data and AI platform that calls a foundation model inline inside a query, enabling PII redaction across both structured tables and unstructured clinical notes without requiring custom model pipelines.
The core principle in this video is that the prompt is the policy: writing a precise redaction instruction as the prompt controls exactly which fields are replaced while preserving analytically important information like medical diagnoses.
A validated redaction prompt can be stored as a Unity Catalog User-Defined Function, making the approved logic reusable and consistent across all teams without each analyst needing to write or maintain their own prompt.

Redact at the Speed of AI: De‑identifying PII in the Lakehouse with ai_query

Teams need to unlock data responsibly. This lightning talk shows how to use AI Functions—specifically ai_query—to automatically redact PII across both structured tables and unstructured documents in Databricks.
We'll walk through how it works for tabular data, then show how the same pattern extends to free-text fields like clinical notes. We'll briefly touch on batching for throughput and checking output quality. A brief live demo runs against realistic unstructured healthcare notes, de-identifying patient information so analysts and ML teams can experiment without risk. You'll leave with reusable queries, sample prompts for both structured and unstructured redaction, and a practical starting point to operationalize PII redaction in your Lakehouse.
Talk By: Han Tran, Solutions Architect, Databricks ;

Chapters

FAQs

What is ai_query and how does it work for PII redaction?

ai_query is a built-in SQL function in the Databricks Data and AI platform that calls a foundation model directly inside a SQL query. For PII redaction, it takes a model name and a prompt as parameters, applying the model inline to each row—whether that row contains a structured field like a social security number or an unstructured clinical note.

Why do traditional NER models fall short for unstructured text PII redaction?

Named Entity Recognition models are typically trained on structured or semi-structured text and struggle with free-form writing like clinical notes, where PII can appear in unpredictable positions and formats. Foundation models accessed through ai_query handle unstructured text more flexibly because they understand context rather than relying on predefined entity patterns.

What does 'the prompt is the policy' mean in this PII redaction approach?

It means the redaction behavior is entirely defined by the instructions written in the ai_query prompt—there is no separate configuration file or rule engine. If you want to remove names and social security numbers but preserve diagnoses, you encode that rule in plain language in the prompt, making it easy to adjust without changing application code.

How can teams ensure consistent PII redaction across an organization?

Once a redaction prompt has been validated, it can be stored as a User-Defined Function in Unity Catalog. All teams then call the same UDF rather than writing their own prompts, ensuring that the approved redaction logic is applied consistently and that any updates to the policy are reflected everywhere automatically.

Full transcript

[00:08] if we could make PII reactions simpler. Have you ever thought about that? My name is Han Tran and I'll be walking you through the session today. All right. So, this session is for you if you have
[00:20] teams that want to use clinical notes, patient records, and customer tables, but it's locked behind PII and PHI. And for the purpose of this talk, I'm going to be interchangeably using PII and PHI.
[00:33] So most of the time, the analysis doesn't need to have sensitive fields. Your analysts don't need someone's social security number in order to build a dashboard. And your data science team
[00:45] does not need somebody's address to build a new model. Right? So this really creates three different problems. First, the sensitive data is blocked for your teams to unlock. And the second is these
[01:00] manual reviews can take a very long time. And traditional models like NE models are not very good at unstructured text. So free forming text. It might be good at tabular data, but for this we're
[01:13] going to be talking about both structured and unstructured data. So this session might not be you if you don't know if you're dealing with PII or PHI data yet. If you don't have an idea
[01:26] if you're using PII, then you would use UFC you need a catalog data classification first to identify whether or not you're using PII and PHI. So this session might not be for you, but you
[01:39] are here because you have PII or PHI data and you're trying to see is it really that simple to be using a simple SQL query to redact PII.
[01:52] Let's take a look. This is AI query. It is a builtin SQL function in data bricks that calls a foundation model in line right inside of your query. And there's two things that you need to change. The first thing is
[02:06] the model name. I'm choosing an open-source Metal Lama model, but you can choose any data bricks hosted foundation model that's available in data bricks today, or you can even bring
[02:18] your own model, and you can just have that be hosted inside of data bricks here. The second thing you need to change is the prompt. And before I get ahead of myself, this is really where
[02:30] the magic really happens. The prompt is the policy. And I'm going to be repeating this throughout the presentation, right? So the prompt here is going to be uh very important. So it
[02:43] is all the information you need. Uh think of it like this. You are giving AI query instructions like you how you would with another human being. If you're just saying, hey, redact everything that's PII, people that work
[02:56] with PII like yourselves might know what you're talking about, but the everyday average person might not know what PII is. So I'm going to be giving examples of what PII is to AI query. I'm going to
[03:08] give it examples of what to replace it with and make it very concrete and clear to AI query. So very simple, right? Two different things you need to change the model name and just the prompt.
[03:21] I want to talk about AI functions. So AI query is part of a larger family called AI functions. I'm bringing this up because remember we talked about
[03:34] structured and unstructured data in this presentation. We're also going to be talking about unstructured. So AI parse document is alluded in here because if you have a very specific task for AI
[03:46] function for example like summarize or parsing which is an OCR model that extracts everything from unstructured uh uh documentation like PDFs or PGs and to
[03:58] turn that into something more structured. So AI query can later than uh digest then you would use those specific task AI functions but for our purpose there is no AI function that can just redact
[04:12] PII data and that's where we're using the more um generalpurpose AI called AI query and then of course you can use you can build AI functions uh within agent bricks as well and so that's why that
[04:25] option is Let's take a look at some examples. I'm going to go first through a tabular example structure data. There's rows and
[04:38] there's columns. On the left you have patient record name, SSN, diagnosis, provider and on the right is what happens after AI query processes it. You
[04:51] can see that the PII is gone, but notice that the diagnosis stays the same. The type 2 diabetes is clinical information, not PII. So, the model understands the difference there. And that's the
[05:04] advantage over those uh predated models that we talked about in the beginning. You don't have to enumerate every single pattern. You tell the model what is PII and what is uh what is an example and it will take that and run with it.
[05:17] This is where it gets interesting. unstructured data. So, same AI query, same pattern, the same prompt I'm putting in, but instead of clean columns, we're passing a free text
[05:30] clinical note. And you can see that the names, date of births, MRNs, facility names are all embedded in natural language. But again, they are identified, redacted,
[05:43] but still the uh the the context there fatigue and elevated glucose is still there because your teams probably need that information to do additional analysis.
[05:58] And I mentioned this before, the prompt is the policy. When I was building this demo and this presentation, I learned many different types of iterations that we had to change. Right? First is don't just say remove personal information. I made that
[06:11] example earlier where you have to really explain it like you were explaining it to another human being who does not know what PII is. So you have to give very clear examples, be specific, define the replacements. You can't just say redact
[06:25] everything and then also preserve that meaning. I'm going to dive deeper into all of this during the demo here, but I just wanted to call out these three big um ideas that's going to come up later
[06:38] again in the demo. And at a high level, this is how I built it. There is going to be a data bricks apps component that is completely optional. I had just built this on a
[06:50] data bricks app to make it digestible for you to see the demo but that's not going to be needed in in terms of pushing this to production later on with AI query you will look at uh we'll be
[07:02] looking at a serverless SQL warehouse as well and that's going to be calling back those AI functions that I mentioned AI query which is our general purpose uh query function that we're focusing on today and then again AI parse document
[07:15] remember for unstructured text we need to call AI parse document first to grab the unstructured text and make it to something structured then feeds into AI query and then lastly I'm calling I'm
[07:28] calling an open source model metalama and again you can change out the models you don't have to use the metal lama model if you do not want to
[07:51] Okay. So, uh this is a data bricks application for those of you who have not seen a data bricks app before. Uh I'm going to go through the structured format first. So, this is structured. Uh this is a tabular format. We have the patient ID, first name, last name, date
[08:06] of birth, etc. And down here, because it's structured, I did not need to call AI parse document, right? We'll get to that in a little bit. But for this I'm just calling AI query directly. This is
[08:18] again the model name that I am calling. And let's look at the entire full prompt. I concatenated it. But let's see here. I'm saying you are a PII reduction system for healthcare data. And I'm
[08:30] giving it very clear examples of what I'm trying to identify with tags and also to replace it with what. Let's look at one example.
[08:45] So I'm saying person's names could be patients, family members, providers. Replace that with name and underneath that I have an example says Dr. Sarah Chen, you are going to replace this with doctor just name. And then down here I'm
[08:59] saying preserve all clinical context exactly. keep the symptoms, diagnosis, medications, dosages, um, all the same and any nonidentifying
[09:11] content only redact identifiers of the person. I'm going to hit run. And now what's happening in the background here is that it's calling AI query throughout this entire table. And
[09:24] by the way, this is a fictitious data set. I this is not real, right? I generated this using AI. Um but let's take a look at the before and after. So in the highlighted red is what it
[09:38] identified as PII and then over on the green side is it doing exactly what I wanted to do where there was an MRN number and it just
[09:51] replaced with MRN phone number etc. Now I want to take a quick pause and show you all something called the LLM judge. There are many ways we can judge the
[10:03] output of AI query and you should not trust AI blindly, right? We know that. I'm using AI query actually on itself. So I'm calling AI query again to be a
[10:16] strict judge on whether or not the rejection was accurate. You there's many different ways we can do this. We could have a human in the loop. You can have ML flow do this. There's many different iterations.
[10:29] But in production, I do recommend that you have a human in the loop. Probably build a dashboard out to see how accurate some of these outputs are. Don't just trust it blindly. But for this purpose, I was curious. I wanted to use AI query on itself to see how
[10:42] accurate the actual outputs are. So you can see I'm giving very clear examples and direction of what I wanted to do. And it said, hey, it was about 80% accurate.
[10:55] It thought that Asheville, North Carolina, and the zip code was PII. You and I know that's not PII, right? That's that's not it. So, this is where the human in the loop aspect is very
[11:07] important to be like, uh, actually, you're fine. Uh, we don't really care about the state or the zip code. And then false positives. It thought female was PII, but it's not.
[11:19] So, again, human in the loop aspect here is really great. Okay, let's go over now and look at the unstructured data. There were two types of files that
[11:31] I put in. One was a PNG and one was a PDF. Let me show you the PNG. I found this online just on Google, right? Uh this is not real, but I have some names.
[11:43] I have their MRN. I have their date of birth and I have the history of their illness. So, let's see what happens. when I go and call this function. And
[11:55] now this is where AI pars comes in because remember we had to take that unstructured data transform it to something that AI query can actually digest. So here this is the um the AI parsed
[12:10] document what it parsed out from that that PNG and let's run this redact again. This is the same prompt by the way. Same prompt as a structured uh data format. Nothing nothing different.
[12:24] All right, let's take a look at the before and after again. So patient name Gretle Sweet Tooth is now just name. The formatting is a little bit off, but uh you can see that it it almost translates right over. Date
[12:38] of birth is now completely do MRN is completely taken out. And you can see there that Gretle Sweet Tooth is a 57y old female on the left, but on the right side, it just says name is a 57y old
[12:51] female. Let's take a look at the second document here. The second document I have is a bit more extensive. It's a six-page PDF
[13:05] that is a lot more extensive. Um, for simplicity sake, I'm not going to zoom in to everything here, but you can see that there's a lot more details in this PDF.
[13:17] Let's see how well we can uh parse this out and call AI query. Again, this is the exact same prompt, nothing different, but instead
[13:30] now we have a six-page PDF that's a lot more intensive than that uh PNG that we saw earlier.
[13:43] All right. You can see there that the name is now replaced. What else has it replaced here? Miss Rogers is now just a name. And let's take a look at the LLM uh
[13:56] judge one more time. Same LLM judge prompt from structured. Nothing has changed, but it says it missed one PII. This is probably in a different page. And this is where I'm trying to iterate again. Human in the loop or using some sort of
[14:10] actual judge to really concretely look at the outputs of AI queries can be very important. I also have a notebook that I can show later on after this uh presentation. But
[14:22] for the sake of this, let's go back to our uh oops, let's go back to here. All right.
[14:36] A few things I want to mention here is that AI functions handles parallelization automatically for you. So if you have very large PDF documents or you have hundreds even thousands of documents, you don't need to break this
[14:48] up. I I feel like uh a lot of people might be like, "Oh, we need to batch this up." You don't need to do that. That's the power of AI functions. It will do that for you. You don't need to worry about the infrastructure. It's going to run on serverless for you if
[15:01] you want to. You don't need to worry about that. Again, you can use data bricks hosted models, but you can bring your own models as well. One thing I wanted to mention is when I did a dry run with one of my colleagues,
[15:14] he asked, does this mean that I have to write a chunk of text every single time I want to write a prompt? Isn't that kind of a waste of time? And it is. So the the solve be behind that and a key
[15:26] takeaway I want you to take away from this session is you can have a golden prompt that is approved by your teams written in something called a UDF which is userdefined function and you can call
[15:39] that userdefined function inside of that ai query prompt that we saw. So now you can have repeatable prompts that works for your organization. You wouldn't have to sit there and type out every single prompt. It's going to be a waste of
[15:51] time. So, just know that if you have that golden prompt, you can still fit that into the AI query prompt and scale as you go into production.
[16:04] All right, three key takeaways. We're almost done here. One, the pattern AI query plus a good prompt equals maybe good redaction. Uh the prompt is the redaction itself. Uh you are going to have to play with what works for your
[16:16] organization. You are gonna have to find what is that golden prompt. And once you find that, put that into that userdefined function all governed by Unity catalog. And then lastly, this workflow right the the actual
[16:29] architecture of what I uh sent or what I showed you and the discussion is going to be available after this. So with that, I hope you've learned how to use AI query for action. It was a pleasure
[16:42] talking to you all. Thank you very much.

Learn more about the Databricks Data and AI platform.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.